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

Abstract: 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 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.

## introduction in the field of natural language processing (nlp), models for learning unsupervised representations from unlabeled text based on transformer architectures @xcite are the state-of-the-art on a variety of tasks @xcite . transformer-based language models (tlms) like bert @xcite , roberta @xcite , and their linage of advanced models @xcite , rely on the combination of an unsupervised pre-training of the model, and a subsequent task-specific fine-tuning procedure. tlms are pre-trained over large unlabeled text data using self-supervision, to learn the relationships between different sentences or words of the input. once the tlm is pre-trained over large volumes of data, it can be used in various downstream tasks, by fine-tuning task-specific model layers. with pre-training, tlms learn language representations that are useful across downstream tasks, minimizing the need and burden of retraining the entire model from scratch, again, for each task. extensive pre-training can lead to downstream performance improvements, i.e., it is worth learning complex tlms in huge natural language corpora before fine-tuning them for particular tasks. many have replicated the pre-train-then-finetune strategy in different domains, e.g., pre-training bert with scientific @xcite and biomedical corpora @xcite @xcite ; or in-house, industryspecific tlms @xcite . in addition, continual pre-training -taking a model pre-trained with general corpora to continue pre-training it with in-domain data-is of great value, yielding significant downstream gains @xcite . even if conceptually simple and empirically powerful, pre-training is challenging and expensive. beyond the significant resources needed to pre-train the original bert model by @xcite , the improvements of roberta @xcite relied on orders of magnitude higher computational resources @xcite . the relationship between tlm architecture, training corpus, pre-training hyperparameters, and evaluation metrics is complex and obscure. therefore, previously overlooked pre-training design choices, e.g., pre-training hyperparameter selection, result in significant performance differences. with this work, we aim to improve the pretraining procedure of tlms, by sequentially selecting hyperparameters that result in a more efficient and superior pre-training performance. we hypothesize that an interactive selection of pre-training hyperparameters can accelerate and improve pre-training, i.e., we can achieve a better metric value in fewer epochs. it is critical not only to achieve superior performance, but to reduce the computational cost, steering clear from timeand resource-expensive procedures. increased efficiency in tlm pre-training is paramount amidst concerns pertaining to the carbon footprint of large language models @xcite ; and specifically, the significant impact of hyperparameter selection on resource utilization and power consumption @xcite . our tlm pre-training use-case is random dynamic masking of masked language models (mlms) -in contrast to rule or task-based mlm dynamic masking solutions proposed in the literature @xcite . even though @xcite showed the benefits of random dynamic masking, the search for optimal masking hyperparameters is often carried out based on heuristic techniques and grid-based search. in machine learning (ml), hyperparameter selection is commonly addressed as a black-box optimization problem, which can be solved using evolutionary algorithms @xcite , entropy search methods @xcite , and bayesian optimization (bo) @xcite . in particular, bo can tackle the problem of optimizing an unknown objective function with possibly noisy evaluations @xcite , and of speeding up resource allocation to promising hyperparameter configurations @xcite . aligned with the recent successes of @xcite in hyperparameter selection via bo, we propose a bo framework for sequential tuning of mlm pre-training hyperparameters. our framework is different from bo techniques that speed up hyperparameter set evaluations, such as hyperband @xcite , which is a pure-exploration adaptive resource al-location algorithm for allocating resources among configurations in the non-stochastic setting. we here cast the tlm pre-training procedure as a sequential decision process, in which at each interaction, a reinforcement learning agent selects an action (e.g., pre-training hyperparameters) to maximize cumulative rewards (e.g., the pre-training metric of interest). to accommodate the black-box nature of the pre-training objective function, we fit a probabilistic surrogate model to the empirical evaluations of the pre-training metric, and propose a bandit-based technique for its sequential optimization. in the mlm dynamic masking use case, the bandit actions are the dynamic masking probabilities; and the mlm performance, the unknown function the bandit is trying to maximize, based on estimates computed in the validation set. contrary to dynamic masking techniques that decide which subsets of tokens to mask via combinatorial optimization and dynamic programming @xcite ; we target online, sequential selection of masking hyperparameters for accelerated and improved pre-training. in contrast to proposals that adapt the language model's masking policy to a particular task of interest @xcite , we devise a generic online optimization framework that, by sequential selection of mlm design choices, provides fast and superior tlm pre-training performance, when pre-training -from-scratch and continually-across diverse corpora. the contributions of this work are: to the best of our knowledge, this work is the first to address online optimization of tlm pretraining with bandit-based bo, and to showcase its performance and resource efficiency benefits. the manuscript is organized as follows: section 2 provides the background on bayesian optimization, multi-armed bandits and tlm pretraining; section 3 describes the proposed gp-ts method for tlm pre-training optimization; with its empirical performance evaluated in section 4. concluding remarks are provided in section 5. the multi-armed bandit (mab) is an abstraction for problems that require learning while simultaneously maximizing attained rewards, i.e., balancing the exploration-exploitation tradeoff @xcite . a mab is a sequential decision process that requires decision-making under uncertainty @xcite . at each interaction @xmath0, • • • , t , a bandit agent chooses an action a @xmath1, and it observes stochastic reward r t drawn from an unknown distribution of the selected arm, a t , often characterized parametrically, r t ∼ p(•|a t , θ). the mab agent's goal is to maximize (ex-pected) cumulative rewards, r @xmath2,t , with each arm's expected reward denoted as µ @xmath3, θ}. the challenge is on the lack of knowledge about the reward generating mechanism, which demands learning its properties (e.g., its parameters), as it interacts with the environment. a plethora of mab algorithms have been proposed and analyzed over the years, from computing optimal strategies @xcite and greedy approaches @xcite , to upper confidence interval @xcite and thompson sampling @xcite algorithms. for models in the exponential family, the latter have been empirically and theoretically proven to perform competitively @xcite @xcite @xcite , and extensions have been proposed that model observed rewards via ensembles of models (lu and roy, 2017), gaussian mixture models (urteaga and wiggins, 2018), gaussian processes @xcite , and neural networks @xcite . in the context of bo in general, and mabs in particular, reward uncertainty quantification is critical. gaussian processes @xcite provide not only adequate bayesian uncertainty estimates, but a flexible solution for surrogate models that encode smoothness assumptions of the payoff function @xcite @xcite . we resort to a gaussian process reward model in the proposed bandit-based bo framework for tlm pre-training. ## proposed bandit-based framework we cast tlm pre-training as a sequential decision process to be solved by a multi-armed bandit agent that interactively optimizes the analytically unknown pre-training loss, based on its sequentially observed empirical evaluations. we define pre-training steps, i.e., a fixed number of stochastic gradient updates u in the training set, as bandit interactions @xmath4, • • • , t . the goal is to minimize the tlm pre-training objective l(•|ψ) given tunable hyperparameters ψ, with (stochastic) evaluations of the loss function in the validation set. pre-training hyperparameters at interaction t, ψ t , are the bandit's arms, i.e., a @xmath5. for mlm pre-training with dynamic masking, at each bandit interaction, the agent selects hyperparameters ψ (the proportion of tokens to mask and their masking probabilities), pre-trains the tlm for certain stochastic updates to minimize the mlm loss, and evaluates its performance in the validation subset, as per equation ( 5 ). due to the black-box nature of the pre-training objective, for which only stochastic evaluations are available, we formulate a surrogate reward function (leveraging empirical mlm validation loss estimates) for the bandit to maximize, as it sequentially selects which arm to play. ## conclusion we present a multi-armed bandit-based bayesian optimization framework for the sequential selection of pre-training hyperparameters towards optimized transformer-based language model performance. we develop and evaluate an interactive, gaussian process-based thompson sampling (gp-ts) framework for accelerated language model pre-training. we model noisy evaluations of the pre-training objective (e.g., the mlm loss) as drawn from a surrogate gaussian process that the bandit agent aims to minimize. we provide empirical evidence of how gp-ts, when applied to mlm dynamic masking, attains superior and accelerated (both from-scratch and continual) pre-training performance, along with excellent in-domain downstream metric values. while @xcite randomly select -with fixed probability-which input tokens to mask, we show that sequentially adapting the masking hyperparameters with gp-ts results in enhanced and efficient pre-training. notably, gp-ts interactively selects hyperparameters that result in top performing models faster, enabling significant resource efficiency, of critical importance in practice. building upon our formulation and the provided evidence, we envision follow-up work investigating the proposed method's ability to successfully pretrain large-scale models in general purpose corpora, as well as for optimizing domain-specific models.