SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(b) Validation losses for training the RHO loss proxy/holdout model across three datasets.
We split all datasets into  $20\%$  proxy model pre-training,  $\sim 1\%$  proxy model validation, and  $\sim 79\%$  target model pre-training set (later, during target model pre-training, we further split the remaining  $79\%$  into train and validation set). For BCWK and C4, we train for 12 hours; for MP, we train for 6 hours only since it is much smaller and we want to avoid over-fitting. Other than varying the dataset splits and training budgets, we use the same hyper-parameters from Section 3.
Figure 10b shows the validation loss during proxy model pre-training; none of the models over-fit, and all of them reach reasonable losses. For reference, one may compare their values with the baseline in Figure 9, which shows the validation losses using data from the same dataset sources.
C.5 Lion: Learning Rate (LR) And Weight Decay (WD)
The authors of Lion provide the following guidelines [13, page 14]:
"The default values for  $\beta_{1}$  and  $\beta_{2}$  in AdamW are set as 0.9 and 0.999, respectively, with an  $\epsilon$  of  $1e - 8$ , while in Lion, the default values for  $\beta_{1}$  and  $\beta_{2}$  are discovered through the program search process and set as 0.9 and 0.99, respectively."
We adopt these default  $\beta_{1},\beta_{2}$  hyper-parameters. Next, we look at LR and WD (also from [13]).
Architecture	LR	WD
BERT-Base-like	{1e-4, 3e-4, 5e-4, 7e-4}	{0.03, 0.05, 0.07, 0.1}
T5-Base	{5e-4, 7.5e-4, 1e-3, 2.5e-3, 5e-3, 2e-2}	{0.0}
Table 5: Grid Search Space for Lion.
(b) T5
Figure 11: Lion grid search for both the BERT and T5 model. Training loss over time during pre-training.
"Based on our experience, a suitable learning rate for Lion is typically 3-10x smaller than that for AdamW. Note that the initial value, peak value, and end value of the learning rate should be changed simultaneously with the same ratio compared to AdamW. We do not modify other training settings such as the learning rate schedule, gradient and update clipping. Since the effective weight decay is  $\mathrm{lr} * \lambda$ ; update +=w *  $\lambda$ ; update *=  $\mathrm{lr}$ , the value of  $\lambda$  used for Lion is 3-10x larger than that for AdamW in order to maintain a similar strength."
To recap (details in Section 3), for AdamW, we use base LRs of 1e-3 and 0.02, respectively. For BERT, we use a WD of 0.01, while we disable it for T5 respectively.
Hence, following the above guidelines, we define the grid search space as described in Table 5.
For BERT, we determine a LR of 7e-4 and a weight decay of 0.1 to be best, as illustrated in Figure 11a. For T5, we do not use any weight decay (following Raffel et al. [77], Shazeer [82]) and find an LR of 7.5e-4 to yield the best performance, as shown in Figure 11b. For all optimizers, we follow Nawrot [69] to integrate RMS-LR scaling to facilitate convergence.
C.6 Sophia: Learning Rate (LR), Weight Decay (WD) and  $\rho$
The official code repository  $^5$  suggests the following:
<sup>5</sup>https://github.com/Liuhong99/Sophia/blob/main/README.md
Figure 12: Inserting Sophia as Drop-in Replacement (FP16) for BERT resulted in NaN Losses.
Figure 13: Lowering LR to 1e-4 of Sophia for BERT Slows Down Convergence (FP16).
"Choose lr to be about the same as the learning rate that you would use for AdamW. Some partial ongoing results indicate that lr can be made even larger, possibly leading to a faster convergence. Consider choosing  $\rho$  in [0.01, 0.1].  $\rho$  seems transferable across different model sizes. We choose rho = 0.03 in 125M SophiaG. The (Ir, rho) for 355M, SophiaG is chosen to be (5e - 4, 0.05) (more aggressive and therefore, even faster!) Slightly increasing weight decay to 0.2 seems also helpful."
First, we followed the above guidelines, replaced AdamW with Sophia, and simply tuned  $\rho \in \{0.01, 0.03, 0.05, 0.1\}$ . All of these runs led to NaN losses, which we illustrate in Figure 12. Next, we lowered the learning rate to 1e-4, which mitigated the instabilities but resulted in much slower convergence, as shown in Figure 13.
We then decided to switch from FP16 to BF16, which improved training stability. Further, we manually tuned the LR, WD, and  $\rho$  values, as shown in Figure 14. Since we noticed a strong negative correlation between  $\rho$  and the training loss, we decided to stick to  $\rho = 0.01$ . The best performance was achieved with an LR of 4e-4 and a WD of 0.015. We follow Nawrot [69] to integrate RMS-LR scaling to facilitate convergence, as for all optimizers.
For T5, we vary the learning rate within the range of  $\{2e - 2, 5e - 3, 2.5e - 3, 1e - 3, 7.5e - 4, 5e - 4\}$ , and  $\rho$  in  $\{5e - 2, 1e - 2\}$ . Figure 15 show the results; similar to BERT, we observe better performance with a smaller  $\rho$ . The best performance comes with  $\rho = 1e - 2$  and LR of 1e-3.
We acknowledge that there are additional hyper-parameters like  $\epsilon$  and  $k$  that we did not tune because we followed the author's suggestions. Future work may investigate their effects on the training speed-ups too.
(a) Weight decay  $= 0.01$
(b) Weight decay  $= 0.15$
Figure 14: Sophia grid search for BERT (BF16). Training loss over time during pre-training.
(a)  $\mathrm{Rho} = 1\mathrm{e} - 2$
Figure 15: Sophia grid search for the T5 model. Training loss over time during pre-training.
(b)  $\mathrm{Rho} = 5\mathrm{e} - 2$
D.1 Training BERT, Evaluating on GLUE
We notice that on specific GLUE datasets, the efficiency methods have little effect on performance (MNLI, SST-2, QNLI, QQP, MRPC). Across all datasets, CoLA appears to have the largest gains from efficient training. Averaged across all datasets, efficient training methods have little effect. We report the exact numbers in Table 1 (left) for layer stacking and layer dropping. In these experiments, layer dropping barely outperforms the baseline given a 6-hour RST budget. For a 12-hour budget, both stacking and layer dropping inch above the baseline, and stacking only produces more accurate results than the baseline and layer dropping for a 24-hour budget. In our study, we also compare the performance of the efficient methods on the T5-base model evaluated on the SNI benchmark and report the results in Table 2. From our observations, layer stacking is particularly noteworthy, demonstrating superior performance in a 6-hour training period, significantly outperforming the other methods. However, as the training budget is increased to 12 hours, the gap between baseline and layer stacking starts to diminish. In a 24-hour training scenario, baseline exhibits a marginally better performance than layer stacking, highlighting the efficacy of the baseline method given sufficient training time. Notably, both baseline and layer stacking consistently outperform layer dropping across all the time budgets.
D.2 Layer dropping on small datasets
The layer dropping paper [107] trains on the Bookcorpus+Wikipedia dataset for 186 epochs, while the original BERT paper trains for only 40 epochs [22]. They do not report results for the baseline trained
Figure 16: BERT models evaluated on GLUE. The black vertical error bars indicate the standard deviation over three seeds. The black horizontal line shows the average performance of the baseline. For clarity, the individual tasks are plotted against the left-hand axis, while the average accuracy is plotted against the right-hand axis.
(a) Ablation for layer dropping, investigating its lack of performance. It prevents overfitting when performing multiple epochs over the dataset.
with a schedule based on fewer epochs. Given the possibility of overfitting due to the high number of epochs layer dropping's similarity to dropout [89], and dropout's known efficacy for language model pre-training with repeated data [104], we raise the question of whether layer dropping acts as a regularizer.
Note that given the abundance of pre-training data for language models, even the largest and most-expensively-trained models [34, 14, 92] are typically trained within a one-epoch regime, which has been shown to improve performance over training for multiple epochs on a smaller dataset [51, 77]. Hence, we exhaust the compute budget before passing through all available training data more than once. In contrast, Zhang & He [107] use a smaller dataset and thus complete 186 epochs during training. Thus, their training setup likely corresponds to an overfitting regime.
To investigate this, we repeat the experiment with a truncated training dataset, resulting in the experiment performing roughly 180 epochs. The result is shown in Appendix D.1. The plot confirms our suspicion: when the baseline training procedure overfits, layer dropping can help mitigate this, preventing the validation loss from increasing as training time increases.