Revisiting Checkpoint Averaging for Neural Machine Translation (2022)

Abstract: Checkpoint averaging is a simple and effective method to boost the performance of converged neural machine translation models. The calculation is cheap to perform and the fact that the translation improvement almost comes for free, makes it widely adopted in neural machine translation research. Despite the popularity, the method itself simply takes the mean of the model parameters from several checkpoints, the selection of which is mostly based on empirical recipes without many justifications. In this work, we revisit the concept of checkpoint averaging and consider several extensions. Specifically, we experiment with ideas such as using different checkpoint selection strategies, calculating weighted average instead of simple mean, making use of gradient information and fine-tuning the interpolation weights on development data. Our results confirm the necessity of applying checkpoint averaging for optimal performance, but also suggest that the landscape between the converged checkpoints is rather flat and not much further improvement compared to simple averaging is to be obtained.

## introduction checkpoint averaging is a simple method to improve model performance at low computational cost. the procedure is straightforward: select some model checkpoints, average the model parameters, and obtain a better model. because of its simplicity and effectiveness, it is widely used in neural machine translation (nmt), e.g. in the original transformer paper @xcite , in systems participating in public machine translation (mt) evaluations such as conference on machine translation (wmt) @xcite and the international conference on spoken language translation (iwslt) @xcite : @xcite ; @xcite ; @xcite ; @xcite ; @xcite ; @xcite ; @xcite ; di @xcite ; @xcite , and in numerous mt research papers @xcite @xcite @xcite . apart from nmt, checkpoint averaging also finds applications in transformer-based automatic speech recognition models @xcite @xcite @xcite . despite the popularity of the method, the recipes in each work are rather empirical and do not differ much except in how many and exactly which checkpoints are averaged. in this work, we revisit the concept of checkpoint averaging and consider several extensions. we examine the straightforward hyperparameters like the number of checkpoints to average, the checkpoint selection strategy and the mean calculation itself. because the gradient information is often available at the time of checkpointing, we also explore the idea of using this piece of information. additionally, we experiment with the idea of fine-tuning the interpolation weights of the checkpoints on development data. as reported in countless works, we confirm that the translation performance improvement can be robustly obtained with checkpoint averaging. however, our results suggest that the landscape between the converged checkpoints is rather flat, and it is hard to squeeze out further performance improvements with advanced tricks. ## related work the idea of combining multiple models for more stable and potentially better prediction is not new in statistical learning @xcite . in nmt, ensembling, more specifically, ensembling systems with different architectures is shown to be helpful @xcite @xcite . in contrary, checkpoint averaging uses checkpoints from the same training run with the same neural network (nn) architecture. compared to ensembling, checkpoint averaging is cheaper to calculate and does not require one to store and query multiple models at test time. the distinction can also be made from the perspective of the interpolation space, i.e. model parameter space for checkpoint averaging, and posterior probability space for ensembling. as a trade-off, the performance boost from checkpoint averaging is typically smaller than ensembling @xcite . in the literature, @xcite study the use of checkpoints from the same training run for ensembling; smith (2017) proposes cyclic learning rate schedules to improve accuracy and convergence; @xcite propose to use a cyclic learning rate to obtain snapshots of the same model during training and ensemble them in the probability space; @xcite perform model parameter averaging on-the-fly during training and argue for better generalization in this way; @xcite discuss empirical findings related to checkpoint averaging for nmt; @xcite and @xcite maintain an exponential moving average during model training; @xcite propose a boosting algorithm and ensemble checkpoints in the probability space; @xcite exploit the fisher information matrix to calculate weighted average of model parameters. here, we are interested in the interpolation happening in the model parameter space, and therefore restrain ourselves from further discussing topics like ensembling or continuing training on the development data. ## methodology in this section, we discuss extensions to checkpoint averaging considered in this work. an intuitive illustration is shown in fig. 1 . ## experiments we re-implement transformer @xcite using pytorch @xcite and experiment on iwslt14 german-, russian-, and spanish-to-english (de-en, ru-en, es-en), and wmt16 english-to-romanian, wmt14 englishto-german, wmt19 chinese-to-english (en-ro, ende, zh-en) datasets. due to limited length, we only present representative results on de-en in this section. results on other language pairs can be found in the appendix and the trends are similar to that reported in this section. note that, in the experiments below, the test bleu scores are under consideration. however, we argue that it is not critical because checkpoint averaging is a vetted trick to boost system performance and our goal is to better understand the parameter space and not to obtain "the state-of-the-art" in some public scoreboard. in fig. 2 , we plot the bleu @xcite ) scores versus increasing k, where the previous k checkpoints starting from the best checkpoint (in terms of devppl) are selected. as can be seen, initial bleu improvements are obtained but as worse and worse checkpoints are included, the bleu score drops as expected. in fig. 3 , ranking all checkpoints by their de-vppl, the top-k checkpoints are selected for averaging. notice that up to @xmath0, the devppl is still around 5, whereas in the last-k case, significantly worse checkpoints (the early checkpoints) are already included in the interpolation. it can be seen that the final bleu score is much less sensitive to the choice of k in this case. of course the final performance also relies on the checkpointing settings (e.g. the checkpointing frequency) but it is clear from the comparison that one should prefer to include checkpoints with better devppl. in fig. 4 , we plot the bleu scores against the temperature τ in eq.2. here, we select last-k checkpoints as in fig. 2 to artificially include some bad-performing checkpoints. two sanity checks can be done here. when τ is very small, uniform weights are used and the performance is close to the vanilla last-40 case. when τ is very large, one-hot weights are used and the performance is close to that of the best checkpoint. we observe that using the devppl-dependent weights results in similar performance increase compared to the vanilla case, meaning that the checkpoint selections can be automated by selecting a proper τ . next, we study how the system performance changes with the step size used in the one-shot gradient update (fig. 1b and eq.3). as shown in fig. 5 , we interpolate three systems selecting top-k checkpoints with @xmath1, @xmath2, respectively. here, temperature τ = 100. in line with the results in fig. 2 and fig. 3 , the models with @xmath3. however, as the step size η increases, the bleu score quickly drops as the averaged model diverges further away from the initial mean. it is clear from the figure that nothing is gained in terms of bleu during the η scan. in other words, these results suggest a very flat surface along the direction of averaged gradients. ## -3 given the results so far, it is clear that although a small boost of bleu score can be robustly obtained in various checkpoint averaging settings, it is hard to squeeze out any further improvement with the extensions considered here. we therefore perform a grid search over the interpolation weights w k with @xmath4, to examine the landscape between the checkpoints. shown in fig. 7 , is the intersection of w 1 + w 2 + w 3 = 1, 0 ≤ w @xmath5. from the figure, except when really close to the vertices, i.e. (w 1 , w 2 , w 3 ) = (1, 0, 0) or (0, 1, 0) or (0, 0, 1), the surface is rather flat with small fluctuations here and there. considered together with the previous results, this suggests that the gradient direction in the flat area may be unreliable and not much improvement is to be gained by further tuning the interpolation weights. of course one could argue that in higher dimensions the surface could look different by moving off of the @xmath6, but we think it is unlikely to be helpful as fig. 5 is a counter-evidence at hand. ## conclusion we consider checkpoint averaging, a simple and effective method in neural machine translation to boost system performance. specifically, we examine different checkpoint selection strategies, calcu-late weighted average, make use of gradient information and optimize the interpolation weights. we confirm the robust improvements from checkpoint averaging and that the checkpoint selection can be automated with the weighted average scheme. however, by closely looking at the landscape between the checkpoints, we find the surface to be rather flat and conclude that tuning in the space of the interpolation weights may not be a meaningful direction to squeeze out further improvements.