SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mindermann et al. [66] argue that solely prioritizing high training loss results in two types of examples that are unwanted: (i) mislabeled and ambiguous data, as commonly found in noisy, web-crawled data; and (ii) outliers, which are less likely to appear at test time. The authors propose down-weighting such data via a selection objective called Reducible Holdout (RHO) loss, shown in Algorithm 4.
The authors acknowledge that their method comes with three overhead costs: (1) pre-training a proxy model using holdout data, (2) one extra forward pass over the entire training set to store the proxy model's training losses, and (3) additional forward passes for the batch selection. In our evaluations we never count the costs of (1) and (2) because, as Mindermann et al. [66] point out, these costs can be amortized over many training runs. For (1), we pre-train a model for 6 hours of RST on held-out data, which is enough to reach reasonable performances (Table 1).
Despite Mindermann et al. [66] motivating their method for a scenario where additional workers are available to perform batch selection, we wondered if it might still provide performance gains even if this is not the case. Mindermann et al. [66] do not implement or evaluate an algorithm where extra workers are used to select batches. Because of this, we evaluate RHO loss under two protocols: in the main results, we count the batch selection costs (3) against the training budget, while in Section 5.4, we provide a second set of results where we ignore the cost of batch selection.
We assume that the effects of selecting better training batches should be largely agnostic to whether we pre-train a BERT or T5 model. Hence, instead of training both architectures, we decide to pre-train only BERT models and instead vary the datasets and budgets as follows.
For the first set of experiments, we fix the budget to 12 hours and consider three different pre-training datasets: (i) C4 [77], consisting only of webpage text which, despite being regularly used for pre
Algorithm 5 Sophia [58]
1: Input:  $\theta_{1}$ , learning rate  $\{\eta_t\}_{t=1}^T$ , hyperparameters  $\{\lambda, \beta_1, \beta_2, \epsilon, \rho, k\}$ , and GNB-Estimator  
2: Set  $m_0 = 0$ ,  $v_0 = 0$ ,  $h_{1-k} = 0$   
3: for  $t = 1$  to  $T$  do  
4: Compute minibach loss  $\mathcal{L}_t(\theta_t)$ .  
5: Compute  $g_t = \nabla \mathcal{L}_t(\theta_t)$ .  
6:  $m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t$   
7: if  $t \mod k = 1$  then  
8: Compute  $\hat{h}_t = \mathrm{GNB}(\theta_t)$ .  
9:  $h_t = \beta_2 h_{t-k} + (1 - \beta_2) \hat{h}_t$   
10: else  
11:  $h_t = h_{t-1}$   
12:  $\theta_t = \theta_t - \eta_t \lambda \theta_t$  (weight decay)  
13:  $\theta_{t+1} = \theta_t - \eta_t \cdot \mathrm{clip}(m_t / \max\{h_t, \epsilon\}, \rho)$
Algorithm 6 Lion [13]
1: Input:  $\theta_{1}$ , learning rate  $\{\eta_t\}_{t=1}^T$ , hyperparameters  $\{\lambda, \beta_1, \beta_2\}$   
2: Set  $m_0 = 0$   
3: for  $t = 1$  to  $T$  do  
4: Compute minibach loss  $\mathcal{L}_t(\pmb{\theta}_t)$ .  
5: Compute  $g_t = \nabla \mathcal{L}_t(\pmb{\theta}_t)$ .  
6:  $u_t = \text{sign}(\beta_1 m_{t-1} + (1 - \beta_1) g_t)$   
7:  $m_t = \beta_2 m_{t-1} + (1 - \beta_2) g_t$   
8:  $\pmb{\theta}_{t+1} = \pmb{\theta}_t - \eta_t u_t$
training, is known to have quality issues [53], (ii) Bookcorpus and Wikipedia [22], which contain polished, book(-like) text and MiniPile [41], a subset of the diverse Pile pre-training corpus [28], containing code, mathematics, books, webpages, and other scientific articles.
Validation loss. To rank the pre-training performances, we compare the validation loss, which is directly comparable as we use the same inputs for all methods (which is not the case for the training data). This is shown in Figure 3, which shows the validation loss every 3 hours throughout training. We find that both batch selection methods do not improve over the baseline.
Downstream performance. Next, we investigate downstream performances, fix the C4 corpus as the pre-training corpus, and vary the budgets (6, 12, and 24 hours). Figures 2 and 16 show the results on SuperGLUE and GLUE. We observe very small differences between the methods and the baseline. On average (Table 1) no batch selection method significantly outperforms the baseline.
5.4 Ablation: Batch Selection For Free
We want to disentangle whether batch selection methods fail to improve over the baseline because their gains do not compensate for their computational overhead. In the previous section (Section 5.3) and the main results (Table 1),
we accounted for this overhead. Therefore, in those experiments, selective backprop and RHO loss effectively update the model for fewer iterations than the baseline within the fixed budgets. Here, we re-run a small number of experiments where batch selection is free, and so we train these models with the same number of iterations as the baseline.
Figure 3: Validation losses for different datasets. Results for batch selection methods (selective backprop and RHO loss) for a 12-hour RST budget.
Specifically, we run an experiment for 12 hours using the C4 corpus and GLUE downstream tasks. For selective backprop, we choose  $\beta = 1$ , which resulted in  $\sim 1.7x$  of the wall-clock time; while for RHO loss, we choose a mega-batch size that is  $10x$  larger (15360) than the mini-batch size (1536), following Mindermann et al. [66], which led to  $\sim 5.3x$  of the original pre-training time. Table 4 shows the results. We see that RHO loss reaches a slightly better final validation loss but performs worse on the GLUE downstream tasks than baseline and selective backprop, which does not improve over the baseline.
(b) T5
Figure 5: Training losses, efficient optimizer methods (Lion, and Sophia). Results are shown for RST budgets of 6 hours  $(\dots)$ , 12 hours  $(\cdot \cdot \cdot)$ , and 24 hours  $(—)$ , on C4.
6 Case Study 3: Efficient Optimizers
Recently, two efficient optimizers were proposed to replace the ubiquitous Adam(W) optimizers [48, 62] for language models: Lion [13] and Sophia [58].
Lion [13] is an optimizer symbolically discovered in the vast program space of first-order optimization primitives. As such, it does not follow any theory-grounded principle that would justify its acceleration property a priori; however, Chen et al. [13] report empirical speed-ups over AdamW in many settings.
Sophia [58] is a scalable stochastic second-order optimizer primarily designed for and evaluated on language model pre-training. The authors claim that Sophia achieves a 2x speed-up compared with AdamW in the number of steps, total compute, and wall-clock time. The authors study two Hessian estimators, but as of this writing, only open-source the code for the empirically better one (Gauss-Newton-Bartlett), which we use here.
The baseline in this section simply refers to AdamW [62].
Mixed Precision Training with BF16 vs. FP16 A common practice for training language models is to use mixed precision training [65]. In initial experiments with BERT, we observed several numerical instabilities (NaN training losses) during hyper-parameter search after inserting Lion and Sophia into our training pipelines as drop-in replacements. Our default mixed-precision mode was FP16, and as noted by other practitioners of Sophia [57], BF16 can sometimes be more stable. Hence; for the optimizer comparisons with BERT, we report results in BF16; including the baseline (although we notice that the baseline's training curves are essentially identical across both modes). For the T5 model, we use the TF32 precision format.
RMS-Scaling for T5 Pre-training T5 pre-training [77, 82] typically employs the Adafactor optimizer [83], which relies on tensorwise learning rate scaling determined by the tensor's root mean square (RMS). This scaling has been identified as critical for convergence during pre-training when using AdamW [69]. Initial tests with Lion and Sophia without additional adjustments led to divergence for higher learning rates or suboptimal performance. This mirrored the behavior of AdamW without RMS scaling. To address this, we incorporated RMS scaling into Lion and Sophia for a subsequent round of experiments, which we outline in Appendix C.6.
In the case of BERT downstream performances (Table 1), we find that Lion and Sophia perform about the same as the baseline. Figure 6 shows the results in more detail. We note that baseline (AdamW) consistently ranks the highest and has a comparatively low standard deviation over random seeds.
Figure 4: BERT Validation loss when we do not count Sophia's Hessian update steps.
Analogous to the batch selection for free ablation in Section 5.4, we also experiment with Sophia while not counting for Hessian update steps and running for the same number of iterations, as shown in Figure 4. Surprisingly, we still do not observe any speedup.
(a) GLUE
Figure 6: BERT BF16 models evaluated on (Super-)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.
(b) SuperGLUE
In this work, we closely examined efficient training algorithms which promised to deliver training speed-ups. First, we clarify that quantifying a training speed-up without specifying an explicit training budget can be misleading and that some previous works missed this. To normalize the wall clock time across different hardware, we introduce the reference system time measure. Then, we put three classes of algorithms to the test in various pre-training settings of BERT and T5 models. We found only a few settings where some of the considered algorithms improved over the baseline.
The authors would like to thank Edoardo Ponti for his feedback on an earlier version of this manuscript. JK and OK acknowledge support from the Engineering and Physical Sciences Research Council with grant number EP/S021566/1. OK was supported by G-Research. PN was supported by the UKRI Centre for Doctoral Training in Natural Language Processing, funded by the UKRI (grant EP/S022481/1) and the University of Edinburgh, School of Informatics and School of Philosophy, Psychology & Language Sciences. PM was partially funded by the European Union's Horizon 2020 research and innovation programme under grant agreement no. 875160, ELIAI (The Edinburgh Laboratory for Integrated Artificial Intelligence) EPSRC (grant no. EP/W002876/1), an industry grant from Cisco, and a donation from Accenture LLP; and is grateful to NVIDIA for the GPU donations. This work was supported by the Edinburgh International Data Facility (EIDF) and the Data-Driven Innovation Programme at the University of Edinburgh.
[102] Xie, S. M., Santurkar, S., Ma, T., and Liang, P. Data selection for language models via importance resampling, 2023. URL https://arxiv.org/abs/2302.03169.
[103] Xing, C., Arpit, D., Tsirigotis, C., and Bengio, Y. A walk with sgd. arXiv preprint arXiv:1802.08770, 2018.
[104] Xue, F., Fu, Y., Zhou, W., Zheng, Z., and You, Y. To Repeat or Not To Repeat: Insights from Scaling LLM under Token-Crisis, May 2023. URL http://arxiv.org/abs/2305.13230.arXiv:2305.13230 [cs].
[105] Yao, X., Zheng, Y., Yang, X., and Yang, Z. Nlp from scratch without large-scale pretraining: A simple and efficient framework. In International Conference on Machine Learning, pp. 25438-25451. PMLR, 2022.
[106] Zhang, C., Bengio, S., and Singer, Y. Are all layers created equal? arXiv preprint arXiv:1902.01996, 2019.
[107] Zhang, M. and He, Y. Accelerating training of transformer-based language models with progressive layer dropping. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in neural information processing systems, volume 33, pp. 14011-14023. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper_files/paper/2020/file/a1140a3d0df1c81e24ae954d935e8926-Paper.pdf.
[108] Zhang, M. R., Lucas, J., Hinton, G., and Ba, J. Lookahead Optimizer: k steps forward, 1 step back, December 2019. URL http://arxiv.org/abs/1907.08610. arXiv:1907.08610 [cs, stat].
[109] Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., Mihaylov, T., Ott, M., Shleifer, S., Shuster, K., Simig, D., Koura, P. S., Sridhar, A., Wang, T., and Zettlemoyer, L. OPT: Open Pre-trained Transformer Language Models, June 2022. URL http://arxiv.org/abs/2205.01068. arXiv:2205.01068 [cs].
[110] Zhu, C., Ni, R., Xu, Z., Kong, K., Huang, W. R., and Goldstein, T. Gradinit: Learning to initialize neural networks for stable and efficient training. Advances in Neural Information Processing Systems, 34:16410-16422, 2021.
[111] Zhu, Y., Kiros, R., Zemel, R. S., Salakhutdinov, R., Urtasun, R., Torralba, A., and Fidler, S. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In ICCV, pp. 19-27. IEEE Computer Society, 2015.
[112] Zhuang, B., Liu, J., Pan, Z., He, H., Weng, Y., and Shen, C. A survey on efficient training of transformers. arXiv preprint arXiv:2302.01107, 2023.
[113] Zhuang, J., Tang, T., Ding, Y., Tatikonda, S. C., Dvornek, N., Papademetris, X., and Duncan, J. AdaBelief Optimizer: Adapting Stepsizes by the Belief in Observed Gradients. In Advances in Neural Information Processing Systems, volume 33, pp. 18795-18806. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/bitstream/d9d4f495e875a2e075a1a4a6e1b9770f-Abstract.html.
[114] zhuofan xia, Pan, X., Jin, X., He, Y., Xue', H., Song, S., and Huang, G. Budgeted training for vision transformer. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=sVzBN-D1JRi.
A Limitations and Future Work
Firstly, while we ran extensive ablations within the 6-, 12-, and 24-hour training regimes, it is possible that our results do not generalize to much longer ones. We justify this choice for two reasons. Firstly, all downstream task performances observed are only a few percentage points worse than state-of-the-art performances; e.g., our 24-hour BERT model reaches  $\sim 79\% /58.5\%$  on GLUE/SuperGLUE, while fine-tuning the original BERT-base model reaches  $80.9\% /60.8\%$  , respectively. Similarly, our 24h-T5 model reaches  $39.5\%$  on SNI, while using the original checkpoint, reaches  $41.0\%$  [69]. Secondly, we argue that an efficient training method should work well specifically in settings with limited budgets.
Next, as illustrated in Appendix B, there is an abundance of efficient training algorithms, and rigorously evaluating all of them is prohibitively expensive. Hence, one limitation of this work remains that we did not consider other efficiency-promoting algorithms, and we hope to explore more in future work.
Another limitation is our sole focus on language model pre-training. Investigating approaches for (i) efficient fine-tuning of (large) language models or (ii) pre-training on other data-intensive modalities such as images and video remains promising too. We expect that our experimental protocol utilizing RSTs will benefit future works doing so.
B.1 Efficient Training Algorithms
There is an abundance of proposals for efficient training; surveys of these can be found in [8, 85, 43]. We roughly categorize them into architecture-centric, data-centric, optimization-centric, and others.
Architecture-centric strategies. These decide how to avoid forward/backward passes of specific weights in the network. The idea of gradually growing a network to accelerate training, as in layer stacking, dates back to the 90s [25, 55]. There are a number of follow-ups to the layer stacking paper [30], including automated stacking using elastic supernets [56] loss- and training-dynamics preserved stacking [86] and stacking small pre-trained models to train larger models [96].
Methods akin to layer dropping include FreezeOut [9], LayerDrop[26] (which focuses on network sparsity), and AutoFreeze [60]. In these methods, forward/backward passes for specific layers are skipped based on either a pre-determined schedule [9, 26] or based on statistics from prior forward/backward passes [60].
While not the focus of this work, a similar motivation can also be found in dynamic sparse training methods, which aim to identify relevant sub-networks during training and promote the pruning ability of the network [67, 21, 24]. However, these approaches do typically not lead to training speedups since unstructured sparsity is not GPU-friendly, which is why we do not consider them here [59].
Data-centric strategies. Besides batch selection methods, which subsample relevant data from a mega-batch throughout training as discussed in Section 5, there are a number of other approaches aiming to either order or subsample training data. We did not consider them in our experiments since they either have a different motivation than training efficiency or are not directly suitable for a drop-in modification of the training procedure. However, we briefly summarize three classes of such work. Firstly, one of the oldest classes of data-selection strategies for training is curriculum learning (for a recent survey, see Wang et al. [97]). We do not consider curriculum learning here as it has already been extensively evaluated [100], and because it was initially motivated to improve generalization rather than achieve efficiency gains. Secondly, task-specific retrieval [32, 105] use task-specific data to retrieve similar data from a large unsupervised corpus. Different from the task-agnostic batch selection methods we consider here, these methods are specifically designed to improve downstream task performance. Lastly, another line of work tries to subsample the entire dataset decoupled from the training procedure by various scoring heuristics [16, 31, 74, 54, 88, 1, 102, 101, 61], and we believe that further investigating such can be a promising direction for future work.
<sup>4</sup>https://huggingface.co/bert-base-uncased
(b) T5
Figure 7: Layer stacking grid search: we tune the intervals at which the model is doubled. The notation "stacking  $(a, b)$ " signifies that the model's size was doubled once at  $a$  hours and then again at  $b$  hours, measured using RST.
Optimization-centric strategies. Lots of optimizers have been proposed with the goal of speeding up convergence [23, 108, 113]; yet, Schmidt et al. [80] report that none of them consistently outperforms Adam(W) [48, 62] when put to a rigorous test. Instead of modifying the training procedure, another line of work observed intermediate performance speedups by averaging weights along the training trajectory post-hoc training [40]; however, such gains arise primarily through effectively lowering the learning rates without directly intervening in the training process [78, 79], which is different to the in this work considered methods which do intervene.
Others. These include ways to improve the faster computation of Transformer building blocks [91, 18, 70], allocate compute conditioned on the input [84, 71, 2], network initialization [7, 110].
B.2 Efficient Training Meta-Studies
Budget-centric recipes. A different line of work investigates budget-centric training recipes, for example, for academic settings with multiple GPUs [38, 114], or hard time constraints such as training on a single GPU for one day [36, 29, 69]. Our work adopts some of the recipes proposed by Geiping & Goldstein [29], Nawrot [69] and aims to complement them by investigate (other) speedup techniques.
Empirical Meta-Studies of Training Transformer-based models. Narang et al. [68] study Transformer modifications across implementations and applications, finding that most do not meaningfully improve performances. Similarly, Tay et al. [90] study the scaling properties of various Transformer-based architectures (some of which are designed for efficiency), concluding that the original Transformer proposed by Vaswani et al. [93] has the best scaling behavior. Dehghani et al. [20] discuss how common cost indicators of machine learning models have different pros and cons and how they can contradict each other. Further, they show how training time can be gamed. Our work is heavily inspired by Dehghani et al. [20] and aims to complement it in two ways: (1) by proposing to normalize WCT across different systems (using RST) and (2) by a thorough experimental study in the case of pre-training Transformer-based language models.
Closest to our work is the concurrent work by Dahl et al. [17], who exhaustively discuss various pitfalls of benchmarking neural network optimizers. Among other considerations, they propose to benchmark algorithms within a fixed runtime budget, a practice we agree with and use in our work too. We believe our work additionally complements theirs by other methods beyond optimizers.
(b) T5
Figure 8: Grid search performed on the hyperparameter  $\gamma_{f}$  for layer dropping in a 24-hour budget setting. For the T5 model with layer dropping we also test smaller learning rate  $(1e - 2)$ , because we observe instabilities with the original one  $(2e - 2)$ . We mark runs with learning rate  $= 2e - 2$  with a *. The upper plots depict the training loss for each method, while the lower plots showcase the average number of active layers.
C Hyper-Parameter Search
C.1 Layer stacking: When To Stack
Figure 7 illustrates that layer stacking has relatively low sensitivity to different stacking RST hour times, namely  $\{(2.4, 7.2), (3, 7.2), (7.2, 9.6)\}$ , with  $\{(2.4, 7.2), (3, 7.2)\}$  yielding similar performance and (7.2, 9.6) slightly underperforming. For our experiments in Section 4.3, for both BERT and T5, we choose (3, 7.2), as it maintains the same  $\frac{\text{stacking step}}{\text{all training steps}}$  ratio proposed by Gong et al. [30].
C.2 Layer dropping: How to Drop
In Figure 8, we observe that different choices of  $\gamma_{f}$  yield comparable performance for both the BERT and T5 models. However, the layer dropping training curves for the T5 model are notably spikier than those of the baseline, suggesting that the layer dropping method demonstrates less stability during training with this architecture. As a result, we also tested a smaller learning rate  $(1e - 2)$ , in contrast to the original training rate  $(2e - 2)$ , which ultimately yielded better results. The parameters selected for our experiments in Section 4.3 were  $\gamma_{f} = 20$ ,  $lr = 1e - 2$  for T5 and  $\gamma_{f} = 100$  for BERT.
C.3 Selective backprop: Selectivity Scale
We tune the selectivity scale  $\beta$ , where  $\beta \in \{1,2,3\}$ . Jiang et al. [39] use  $33\%$  and  $50\%$  selectivity in their experiments, which approximately corresponds to  $\beta = \{1,2\}$ , respectively. We find that the larger the  $\beta$  value, the worse the pre-training performance. Note that the higher the  $\beta$  value, the more forward passes selective backprop needs to perform in order to collect enough samples for a backward pass, which decreases the total number of parameter update steps within the RST budget. For the experiments in Section 5.3, we chose  $\beta = 1$ , as it consistently achieves the best performance.
C.4 RHO loss: Mega-batch Size
RHO loss requires one additional hyper-parameter, the size of the mega-batch, from which the mini-batch then gets subsampled from. We tune this hyper-parameter in Figure 10a and similar to Appendix C.3, we find that the larger this size gets, the worse the validation loss. We started tuning it based on BCWK, and given the clear hierarchy we observed; we decided not to tune it on other datasets and simply set it to  $2\mathrm{x}$  (3072).
Another implicit set of hyper-parameters is how to pre-train the proxy/irreducible loss models. Here, we follow suggestions by Mindermann et al. [66] and choose the same architecture as the target model.
(a) Bookcorpus & Wikipedia
(c) MiniPile
Figure 9: Selective backprop grid search: we tune the  $\beta$  hyperparameter. Each plot shows the validation loss over time during training for the given dataset.
(a) RHO loss grid search: we tune the mega batch size hyper-parameter on BCWK as multiples of the mini-batch size (1536):  $\{2\mathrm{x}$  3072),  $3\mathrm{x}$  4608), 4x (6144), 10x (15360)}.