# Multi-armed bandits for resource efficient, online optimization of language model pre-training: the use case of dynamic masking

arXiv: https://arxiv.org/abs/2203.13151 · HTML: https://mlsanity.com/paper/2203.13151/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 67 references that are in the map.
- First posted on arXiv 24 Mar 2022.

## Builds on (most cited first, 30 of 40)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [BioBERT: a pre-trained biomedical language representation model for biomedical text mining](https://mlsanity.com/paper/1901.08746/) (2019)
- [Don't Stop Pretraining: Adapt Language Models to Domains and Tasks](https://mlsanity.com/paper/2004.10964/) (2020)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [SciBERT: A Pretrained Language Model for Scientific Text](https://mlsanity.com/paper/1903.10676/) (2019)
- [Domain-Specific Language Model Pretraining for Biomedical Natural Language Processing](https://mlsanity.com/paper/2007.15779/) (2020)
- [Publicly Available Clinical BERT Embeddings](https://mlsanity.com/paper/1904.03323/) (2019)
- [Ethical and social risks of harm from Language Models](https://mlsanity.com/paper/2112.04359/) (2021)
- [SpanBERT: Improving Pre-training by Representing and Predicting Spans](https://mlsanity.com/paper/1907.10529/) (2019)
- [A Tutorial on Bayesian Optimization](https://mlsanity.com/paper/1807.02811/) (2018)
- [Hyperband: A Novel Bandit-Based Approach to Hyperparameter Optimization](https://mlsanity.com/paper/1603.06560/) (2016)
- [GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration](https://mlsanity.com/paper/1809.11165/) (2018)
- [Carbon Emissions and Large Neural Network Training](https://mlsanity.com/paper/2104.10350/) (2021)
- [Defending Against Neural Fake News](https://mlsanity.com/paper/1905.12616/) (2019)
- [Deep Exploration via Bootstrapped DQN](https://mlsanity.com/paper/1602.04621/) (2016)
- [Introduction to Multi-Armed Bandits](https://mlsanity.com/paper/1904.07272/) (2019)
- [A Tutorial on Thompson Sampling](https://mlsanity.com/paper/1707.02038/) (2017)
- [ERNIE 2.0: A Continual Pre-training Framework for Language Understanding](https://mlsanity.com/paper/1907.12412/) (2019)
- [Predictive Entropy Search for Efficient Global Optimization of Black-box Functions](https://mlsanity.com/paper/1406.2541/) (2014)
- [Kernel Interpolation for Scalable Structured Gaussian Processes (KISS-GP)](https://mlsanity.com/paper/1503.01057/) (2015)
- [Bayesian Optimization is Superior to Random Search for Machine Learning Hyperparameter Tuning: Analysis of the Black-Box Optimization Challenge 2020](https://mlsanity.com/paper/2104.10201/) (2021)
- [Fast Bayesian Optimization of Machine Learning Hyperparameters on Large Datasets](https://mlsanity.com/paper/1605.07079/) (2016)
- [AMMUS : A Survey of Transformer-based Pretrained Models in Natural Language Processing](https://mlsanity.com/paper/2108.05542/) (2021)
- [Bayesian optimization for materials design](https://mlsanity.com/paper/1506.01349/) (2015)
- [Parallel and Distributed Thompson Sampling for Large-scale Accelerated Exploration of Chemical Space](https://mlsanity.com/paper/1706.01825/) (2017)

## Built on by (most cited first, 2 of 2)

- [Birdie: Advancing State Space Models with Reward-Driven Objectives and Curricula](https://mlsanity.com/paper/2411.01030/) (2024)
- [A Component-Based Survey of Interactions between Large Language Models and Multi-Armed Bandits](https://mlsanity.com/paper/2601.12945/) (2026)

## Abstract (arXiv)

We design and evaluate a Bayesian optimization framework for resource efficient pre-training of Transformer-based language models (TLMs). TLM pre-training requires high computational resources and introduces many unresolved design choices, such as selecting its pre-training hyperparameters. We propose a multi-armed bandit framework for the sequential selection of TLM pre-training hyperparameters, aimed at optimizing language model performance, in a resource efficient manner. We design a Thompson sampling algorithm, with a surrogate Gaussian process reward model of the Masked Language Model (MLM) pre-training objective, for its sequential minimization. Instead of MLM pre-training with fixed masking probabilities, the proposed Gaussian process-based Thompson sampling (GP-TS) accelerates pre-training by sequentially selecting masking hyperparameters that improve performance. We empirically demonstrate how GP-TS pre-trains language models efficiently, i.e., it achieves lower MLM loss in fewer epochs, across a variety of settings. In addition, GP-TS pre-trained TLMs attain competitive downstream performance, while avoiding expensive hyperparameter grid search. GP-TS provides an interactive framework for efficient and optimized TLM pre-training that, by circumventing costly hyperparameter selection, enables substantial computational savings.