A Hyperparameter Optimization Toolkit for Neural Machine Translation Research (2023)

Abstract: Hyperparameter optimization is an important but often overlooked process in the research of deep learning technologies. To obtain a good model, one must carefully tune hyperparameters that determine the architecture and training algorithm. Insufficient tuning may result in poor results, while inequitable tuning may lead to exaggerated differences between models. We present a hyperparameter optimization toolkit for neural machine translation (NMT) to help researchers focus their time on the creative rather than the mundane. The toolkit is implemented as a wrapper on top of the open-source Sockeye NMT software. Using the Asynchronous Successive Halving Algorithm (ASHA), we demonstrate that it is possible to discover near-optimal models under a computational budget with little effort. Code: https://github.com/kevinduh/sockeye-recipes3Video demo: https://cs.jhu.edu/kevinduh/j/demo.mp4

## introduction deep learning models are difficult to train. although they achieve impressive results on many tasks, non-trivial amounts of effort are required for selecting appropriate hyperparameters, such as the number of layers, vocabulary size, embedding dimension, and optimization algorithm. this trialand-error process is necessary for each task, domain, or language. further, the rapid development of new neural network architectures implies that this hyperparameter optimization process will only become more expensive. currently, hyperparameter optimization tends to be performed manually by researchers in an ad hoc fashion, using scripts put together independently. the lack of open-source support tools means that the level of rigor in hyperparameter optimization may vary widely. this poses two risks: the importance of documenting the hyperparameter optimization process in research has already been widely recognized and is included as an item under the "responsible nlp checklist" foot_1 required for paper submissions in the field. to support these efforts, we believe it will be beneficial to develop open-source tools to improve the hyperparameter optimization process itself. this paper presents a hyperparameter optimization toolkit for nmt research. it enables researchers to easily explore the hyperparameter space of various nmt models based on the pytorch codebase of aws sockeye framework @xcite . one simply specifies (1) the desired set of hyperparameter options to search, (2) the compute resource constraints, and (3) the training data paths, then the toolkit will plan and execute an automatic hyperparameter optimization and return the best model discovered. the toolkit implements the asynchronous successive halving algorithm (asha) @xcite , which is well-suited for commodity off-the-shelf distributed grids. in the following, we first give an overview of the toolkit (section 2) and hyperparameter optimization algorithm (section 3). then, the case study in section 4 illustrates how the toolkit can help a researcher search over thousands of hyperparameter configurations with ease. finally, section 5 discusses our design choices, hopefully serving as a reference for those who want to implement similar toolkits for different nlp software. ## usage overview our hyperparameter optimization toolkit is named sockeye-recipes3, since it cooks up different models by training models with the aws sockeye nmt framework, version 3. an overview is shown in figure 1 . for concreteness, let us suppose the researcher in figure 1 wants to run a rigorous hyperparameter optimization to obtain a strong transformer baseline for a new dataset. step 1: the researcher designs a hyperparameter search space for his/her model. step 2: sockeye-recipes3 samples from the full hyperparameter space to generate a set of bash files called hpm files. each hpm file represents a specific hyperparameter configuration and encapsulates all the information needed to train a model. this includes not only hyperparameter settings but also paths to training and validation data. for example, config1.hpm might train a model with: the set of hpm files represents all the hyperparameter configurations to be explored by the hyperparameter optimization algorithm. rather than randomly sampling a subspace, one can also generate the full cartesian product or manually edit some hpm files based on prior knowledge. depending on the researcher's usage scenario, this set typically numbers from tens to thousands. step 3: once the researcher is ready, he/she starts the asha program with resource specifications such as the number of concurrent gpus to use and the number of checkpoints per training run. this python code dispatches the training processes as standard sockeye jobs to a distributed grid. foot_2 asha will attempt to efficiently train as many models as possible given the computational constraints. it is a bandit learning method that automatically learns when to stop a not-so-promising training run in order to allocate resources to other hyperparameter configurations. details are in section 3. ## case study goal: to illustrate how sockeye-recipes3 works in practice, we show a case study on building a strong transformer baseline for a new telugu-to-english dataset. our initial training set consists of 900k lines of bitext obtained from public sources via the opus portal @xcite . this is augmented with 7 million lines of back-translated data obtained by running a reverse system (englishto-telugu nmt trained on 900k) on web-scraped news from the leipzig corpus @xcite . 3000 lines are held out from the initial training set to serve as the validation set. given this setup, our goal is to run hyperparameter optimization on a standard transformer architecture to obtain the best possible model according to validation bleu. this model can serve as a strong baseline for any future nmt experiment based on the same dataset. since this is a lowresource language pair that is relatively unexplored in the research community, we opt to search a large hyperparameter space. hyperparameter space: our space.yaml file is defined according to the options listed in ## related work asha and variants can be viewed as bandit algorithms that balance exploration (trying new configurations) with exploitation (training the current configurations for longer). they obtain efficiency algorithm 1 asha pseudocode while budget remains do for all @xmath0, g, conda-env) ▷ sockeye train() end for pause for m minutes end while by early stopping. another class of methods are "blackbox" optimizers, e.g. bayesian optimization and evolutionary methods @xcite : they treat hyperparameters as input features, observed accuracy as output targets, and train a proxy model to predict new hyperparameters that are worth sampling. these two classes of methods can be combined @xcite ; this is potentially future work. several benchmarks provide comparisons of state-of-the-art @xcite . neural architecture search @xcite is related to hyperparameter optimization but focuses more on fine-grained choices (e.g. changing skip connections at different layers). this is an active area of research, but out-of-scope for our purpose of improving nmt experimentation. there are some existing toolkits like vizier @xcite and ray tune @xcite , which are suitable for those wanting general rather than application-specific solutions. ## conclusions there is a progression of toolkit development that enables researchers to do better work. deep learning toolkits like pytorch and tensorflow made it easy to exploit gpu hardware. applicationspecific toolkits like sockeye and fairseq build on top of that, and enabled researchers to quickly prototype new ideas. further on top, we believe that hyperparameter optimization toolkits and experiment management toolkits in general will further help advance the speed and rigor of research. we presented sockeye-recipes3, an open-source hyperparameter optimization toolkit for nmt research. our hope is this will relieve some of the mundane aspects of manual hyperparameter tuning so that researchers can focus on more creative activities. a rigorous and automated hyperparameter optimization process will also lead to more trustworthy experiment results.