SlowGuess/NeurIPS_2023_OCR huggingface dataset:
A
B
Figure 4: Ablation Experiments (A) Average GLUE score and (B) Individual GLUE tasks. BERT-base: standard BERT-base (110M parameters) with attention dropout=0.1 and feedforward dropout=0.1, vocab size set to 30522, MLM=15% (all Hugging Face standard configurations). BERT+drpt=0: standard BERT-base, except that the attention in the dropout layer is set to 0 instead of the default 0.1. BERT+GLU: standard BERT-base, with GLU for the feedforward component of the encoder block. BERT+lpLN: standard BERT-base, except with low precision LayerNorm (bfloat16). BERT+mlm30: standard BERT-base, except with a masked language modeling masking ratio of 30%. MosaicBERT: the complete MosaicBERT-Base including GLU (where the dimension of the intermediate layer is 3072 resulting in 137M total parameters), ALiBi, low precision LayerNorm, unpadding, MLM 30%, vocab size 30528 (a multiple of 64) and the attention dropout=0. MosaicBERT-FlashAttn+drpt=0.1: MosaicBERT-Base without Flash Attention and with the attention dropout set to 0.1.
BERT-Large has 24 repeated transformer layers (BERT-Base has 12), and as a result consumes much more memory and takes longer to train. We found that MosaicBERT-Large reached an average GLUE score of 83.2 in 15.85 hours, while BERT-Large took 23.35 hours (Figure 3). MosaicBERT-Large therefore had a  $1.47 \times$  speedup over BERT-Large in this training regime.
While MosaicBERT-Base is optimal for a constrained budget, the MosaicBERT-Large average GLUE score eventually surpasses MosaicBERT-Base after 25 hours of training on a  $8 \times \mathrm{A}100 - 80\mathrm{GB}$  node, and reaches an average score of 85.5 in roughly 50 hours. In our experiments, the MosaicBERT-Large architecture and pretraining was the same as MosaicBERT-Base, outside of the number of attention heads, number of hidden layers, and intermediate size of the feedforward units.
A striking result from Figures 3 and S2 is that MosaicBERT-Base is Pareto optimal relative to BERT-large in this training regime, and is Pareto optimal to MosaicBERT-Large during the first half of training. MosaicBERT-Base takes only 25 hours to complete 2 epochs, while MosaicBERT-Large takes close to 70 to complete 2 epochs. A potential takeaway from this is that MosaicBERT-Large only surpasses Base performance in the large data regime. For certain tasks such as QQP, SST-2, and MRPC MosaicBERT-Base achieves a maximum accuracy on par with the maximum accuracy of MosaicBERT-Large, for far fewer pretraining hours. When building encoders for specific domains and tasks, bigger is not always better.
3.4 Ablation Experiments
When deciding on the final setup for MosaicBERT, we tried to make educated guesses about what combination of modifications could work well together to both increase FLOP utilization and maintain or improve accuracy. Our approach was inspired by Amdahl's Law, i.e. the idea that optimizing a portion of a system responsible for  $\mathrm{N\%}$  of the costs can lead to, at most, an  $\mathrm{N\%}$  speedup [48], and chose each of the modifications such that they did not target the same exact part of the of the stack.
For example, low precision LayerNorm and GLU affect different parts of the encoder block, and so potentially compose well.
In the following ablation experiments we benchmark the individual effects of various architecture and training choices on pretraining wall clock time. We plot downstream GLUE accuracy as a function of measured pretraining wall clock time in Figure 4.
The patterns in Figure 4A-B shed light on the individual effects of various architectures (e.g. BERT+GLU, BERT+low precision LayerNorm) and training configurations (e.g. BERT + 30% masking ratio). On average, all methods seem to provide a slight accuracy boost to BERT-base. Increasing the masking ratio to 30% leads to a slight accuracy boost while not affecting the WCT, while turning off dropout in the attention layer (BERT+drpt=0) leads to a slight improvement in both accuracy and WCT. Low precision LayerNorm (BERT+lpLN) leads to a significant speedup (i.e. a shift to the left). Gated Linear Units (BERT+GLU) add more parameters to the model and lead to a significant slowdown while providing an accuracy boost. As a point of reference, we also benchmark the full MosaicBERT as well as MosaicBERT without FlashAttention and with attention dropout set to 0.1 (the standard BERT-base configuration).
In these experiments, all BERT-Base models here were pretrained on C4 for 70,000 steps with batch size 4096, and microbatch size 256 on  $8 \times \mathrm{A}100$  80GB GPUs. All models were initialized with the same seed and shared all other hyperparameters including the bert-base uncased tokenizer, the learning rate schedule, AdamW as the optimizer, etc. Final pretraining checkpoints were then finetuned on GLUE following the details in the appendix. The points represented in these GLUE plots are final finetuning checkpoints.
These plots highlight the importance of benchmarking with Pareto curves, as it is not possible to tell from these plots alone whether training BERT-base for 2 more hours leads to better performance than BERT+GLU, for example.
Various studies rightly focus on a single method or modification that improves throughput, such as FlashAttention [11] and unpadding [66]. In this study, we incorporate many of these techniques to investigate whether they combine advantageously.
RoBERTa ("Robustly optimized BERT approach") is the most influential work in this regard [37]. In this study, they kept the exact BERT architecture but changed the training recipe by removing the next sentence prediction objective, training for longer on much larger datasets, and changing the batch size, among other things. They showed that the original BERT was significantly undertrained - while the original BERT trained on 16GB worth of text, the top accuracy RoBERTa (Large) was trained on 160GB of data for 500,000 steps with batch size 8192. Many of the training choices in RoBERTa have become standard practice; our training recipe therefore more closely resembles RoBERTa than the original BERT. See Appendix for a more detailed comparison of RoBERTa and MosaicBERT.
Improvements in transformer architecture and GPU hardware have caused the cost of pretraining BERT models to decline precipitously. The recent paper "Cramming: Training a Language Model on a Single GPU in One Day" [18] exemplifies this trend. The goal of this rigorous study was to train the best BERT in 24 hours on a single GPU. Similar to us, they tweaked the BERT architecture to incorporate FlashAttention and Gated Linear Units (but without increasing the dimensionality of the hidden block). Unlike MosaicBERT, they used scaled sinusoidal positional embeddings [59] as well as Pre-LayerNorm (applying LayerNorm before the attention and feedforward layers) and did not change the pretraining masking rate. With this setup, they were able to train their modified BERT-Base to an average GLUE score of 80.4 in 24 hours on a single A6000 GPU (i.e. 24 GPU hours, as per Table 3 of [18]). Our study is similar in spirit, but asks what is the fastest architecture for pretraining, and expands on this in greater detail by showing that MosaicBERT is Pareto optimal relative to BERT.
While we believe that our training optimization choices for MosaicBERT go a long way to improve BERT training efficiency, there are likely further modifications that could lead to an increase in accuracy with no effect on throughput, such as replacing LayerNorm with RMSNorm [67] and GeGLU with SwiGLU [50, 42]. Removing all biases from the linear layers can also potentially lead to slight speed ups in training time without a significant hit to accuracy. "NarrowBERT" [36]
suggested a clever change to the way encoder blocks are stacked so that computation is not wasted on unmasked tokens. Another recent study showed that dynamically masking the masking ratio during pretraining leads to downstream accuracy gains in BERT [1]. Note, however, that accounts such as [18] and [29] document many examples where certain modifications to the BERT architecture and training recipe fail to improve accuracy. We also note here that there is an ongoing debate about the benefits of RoPE (rotary positional embeddings) [55] versus AliBi [44]. Finally, using a more modern tokenizer will likely have an effect on downstream accuracy [18].
Approaches such as knowledge distillation [5] might additionally push the Pareto frontier of BERT models during pretraining and finetuning. Progressive layer stacking is one example technique; by initializing a larger BERT model with the weights from smaller pretrained BERT models, pretraining can be reliably accelerated [21, 29]. As the field is learning how to optimize stable model training at the billion parameter scale [8, 13], we expect some of these innovations to cycle back to smaller models such as BERT-Base. For example, it has been hypothesized that incorporating more Layer-Norm modules at the QK matrix output (i.e. QK-LayerNorm [24]) can lead to improved stability during training; combining this with a more aggressive learning rate schedule could lead to faster convergence.
In this study we show how combinations of architecture choices can improve BERT pretraining speed and accuracy. We built MosaicBERT to enable ML researchers and engineers to pretrain BERT models from scratch on their own data and build better models for their specific domains without facing time and cost restrictions. We ultimately hope that by making BERT training faster and cheaper, our work contributes to the trend in NLP research away from finetuning generic models and towards of pretraining custom encoders on domain specific data. A large body of research has highlighted the success of BERT-style models pretrained on a specific domains such as biomedicine [3, 33, 22, 17], math [51], chemistry [2, 25], financial communications [49], and code [56]. This seems to hold true even in the age of LLMs [6, 65].
A stable webpage for this work can be found at mosaicbert.github.io. Code for pretraining and finetuning MosaicBERT can be found in the MosaicML examples repository https://github.com/mosaicml/examples. The exact code for this study was pinned to v0.0.4 of the MosaicML mosaicml/examples repository https://github.com/mosaicml/examples/tree/v0.0.4/examples/bert. All pretraining and finetuning was done in PyTorch 1.13 using the MosaicMLComposer library https://github.com/mosaicml/composer. Model weights for MosaicBERT-Base can be found on the HuggingFace hub https://huggingface.co/mosaicml/mosaic-bert-base.
Code for pretraining and finetuning transformers such as MPT and derivative versions of MosaicBERT can be found in the MosaicML llm-foundry repository https://github.com/mosaicml/llm-foundry.
The authors would like to thank Erica Yuen, Vitaliy Chiley, Connor Jennings, and Aaron Gokaslan for their feedback on the code and manuscript. We would also like to thank Ishana Shastri and Vlad Ivanchuk for some of their early work as interns that led to MosaicBERT. The majority of this work appeared as a blogpost with the title "MosaicBERT: Pretraining BERT from Scratch for $20" in March 2023. $^{13}$  This work also builds on the MosaicML MLPerf v2.1 results for BERT-Large [31].
13mosaicml.com/blog/mosaicbert
[56] J. Tabassum, M. Maddela, W. Xu, and A. Ritter. Code and named entity recognition in stackoverflow. arXiv preprint arXiv:2005.01634, 2020.
[57] P. Tillet, H.-T. Kung, and D. Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10–19, 2019.
[58] T. H. Trinh and Q. V. Le. A simple method for commonsense reasoning. arXiv preprint arXiv:1806.02847, 2018.
[59] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[60] A. Wang, A. Singh, J. Michael, F. Hill, O. Levy, and S. R. Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461, 2018.
[61] L. Wang, N. Yang, X. Huang, B. Jiao, L. Yang, D. Jiang, R. Majumder, and F. Wei. Text embeddings by weakly-supervised contrastive pre-training. arXiv preprint arXiv:2212.03533, 2022.
[62] A. Warstadt, A. Singh, and S. R. Bowman. Neural network acceptability judgments. Transactions of the Association for Computational Linguistics, 7:625-641, 2019.
[63] A. Wettig, T. Gao, Z. Zhong, and D. Chen. Should you mask  $15\%$  in masked language modeling? arXiv preprint arXiv:2202.08005, 2022.
[64] A. Williams, N. Nangia, and S. R. Bowman. A broad-coverage challenge corpus for sentence understanding through inference. arXiv preprint arXiv:1704.05426, 2017.
[65] S. Wu, O. Irsoy, S. Lu, V. Dabravolski, M. Dredze, S. Gehrmann, P. Kambadur, D. Rosenberg, and G. Mann. Bloomberggpt: A large language model for finance. arXiv preprint arXiv:2303.17564, 2023.
[66] J. Zeng, M. Li, Z. Wu, J. Liu, Y. Liu, D. Yu, and Y. Ma. Boosting distributed training performance of the unpadded bert model. arXiv preprint arXiv:2208.08124, 2022.
[67] B. Zhang and R. Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[68] Y. Zhu, R. Kiros, R. Zemel, R. Salakhutdinov, R. Urtasun, A. Torralba, and S. Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In Proceedings of the IEEE international conference on computer vision, pages 19-27, 2015.
A Pretraining Hyperparameters
We detail our pretraining hyperparameter values in Table S1 below.
Model	Optimizer	LR	beta	eps	WD	MB	Warmup	Final LR	MLM
BERT-Base	Decoupled AdamW	5e-4	[0.9,0.98]	1e-6	1e-5	512	6%	0.02LR	0.15
MosaicBERT-Base	Decoupled AdamW	5e-4	[0.9,0.98]	1e-6	1e-5	512	6%	0.02LR	0.3
BERT-Large	Decoupled AdamW	2e-4	[0.9,0.98]	1e-6	1e-5	256	6%	0.02LR	0.15
MosaicBERT-Large	Decoupled AdamW	2e-4	[0.9,0.98]	1e-6	1e-5	256	6%	0.02LR	0.3
Table S1: Pretraining hyperparameters for BERT-Base, MosaicBERT-Base, BERT-Large and MosaicBERT-Large. LR = learning rate, WD = weight decay, MB = Microbatch and MLM = Masked Language Modeling ratio. All pretraining runs were done on 8× A100-80 GB GPUs.
All pretraining runs followed a warmup + linear decay learning rate schedule. Evaluation was done every 2000 batches, and checkpoints were saved every 3500 batches - both of which slow down training slightly.
B Finetuning Hyperparameters
We used the hyperparameters in Table S2 for finetuning all BERT and MosaicBERT models. All finetuning datasets used a maximum sequence length of 256 tokens. We found that these values worked well across all tasks for BERT-Base, MosaicBERT-Base, and MosaicBERT-Large; BERT-Large however was somewhat less performant on QQP for some pretraining seeds.
Task	learning rate	beta	epsilon	weight decay	epochs	seeds
MNLI	5e-5	[0.9, 0.98]	1e-6	5e-6	3	1
QNLI	1e-5	[0.9, 0.98]	1e-6	1e-6	10	1
QQP	3e-5,1.5e-5	[0.9,0.988]	1e-6	3e-6	5	1
RTE	1e-5	[0.9, 0.98]	1e-6	1e-5	3	5
CoLA	5e-5	[0.9, 0.98]	1e-6	5e-6	10	4
SST-2	3e-5	[0.9,0.988]	1e-6	3e-6	3	3
MRPC	8e-5	[0.9, 0.98]	1e-6	8e-6	10	5
STS	3e-5	[0.9, 0.98]	1e-6	3e-6	10	5
Table S2: Finetuning hyperparameters for BERT and MosaicBERT across Base and Large. All finetuning was done with sequence length 256. Note that finetuning BERT-Large on QQP with learning rate 3e-5 was highly unstable; we therefore lowered the learning rate in this instance to 1.5e-5 for BERT-Large (but not MosaicBERT-Large).
As can be seen from Figure S2, we found the finetuning of BERT-Large to be highly variable. Finetuning runs that did not converge, or that led to outlier final values were dropped. For this reason, not all checkpoints are represented in the BERT-Large average GLUE plot in Figure 3.
We did not do a careful hyperparameter sweep for BERT-Large finetuning, and instead tried to keep values consistent across BERT-Base and Large as well as MosaicBERT-Base and Large.
C Pareto Curves for Benchmarking Neural Network Training
In order to make meaningful improvements in neural network training efficiency, ML practitioners must be able to compare between different choices of network architectures, hyperparameters, and training algorithms. One straightforward way to do this is to characterize the tradeoff between
Figure S1: Pareto curves for BERT-Base and MosaicBERT-Base for runs trained for 70,000 and 178,000 steps with batch size 4096 and sequence length 128. Same data as Figure 1B and 3.
accuracy and training time with a "tradeoff curve" or a "Pareto curve." Pareto curves can be generated by varying the length of training for each model configuration; longer training runs take more time but tend to reach higher quality. For a fixed model and task configuration, this method of generating tradeoff curves is an estimate of the theoretical Pareto frontier, i.e. the set of all of the best possible tradeoffs between training time and accuracy, where any further attempt to improve one of these metrics worsens the other. See [40, 43] for more details.
We therefore consider one model Pareto optimal relative to another if it has superior accuracy across different training budgets while keeping everything fixed. Many studies advertise novel architecture approaches without measuring wall clock time, or simply show an improvement for a single training duration. In this study we show that MosaicBERT-Base is Pareto optimal relative to our BERT-Base baseline for both short training durations and long training durations (Figures 1, 2, 3). We also show that BERT-Large and MosaicBERT-Large are not Pareto optimal relative to BERT-Base and MosaicBERT-Base for training durations under 30 hours (Figure 3).
The Pareto plots in Figures 1, 2, 3 and S2 are constructed by taking points from the same runs (with the same learning schedule), instead of having one run per pretraining budget on the plot. We did this for cost reasons; it would have been too expensive to do every run separately.
For ResNets, it has been established that the most accurate way to generate Pareto curves is to run completely separate training runs for each point on the Pareto curve, where each point also has their own learning rate schedule [43]. However, to the best of our knowledge, this has not been clearly established for BERT models. Can Pareto curves for BERT simply be constructed by evaluating checkpoints of a single training run?
Since we ran BERT-Base and MosaicBERT-Base experiments for two separate learning rate schedules (70,000 steps and 178,000 steps with batch size 4096), we can plot the Pareto curves for all of these runs on top of each other (Figure S1). Somewhat surprisingly, the 70k and 178k curves for BERT-Base (and MosaicBERT-Base) mostly lie on top of each other. This is strong evidence that at least for BERT models, a reasonable Pareto curve can be constructed using checkpoints from a single training run.
We note that the Pareto curve approach discussed here is also related to "tokens per parameter" (TPP) considerations that are often discussed in LLM pretraining (e.g. see [15]).
Figure S2: MosaicBERT-Base and Large accuracy (dev) vs. pretraining speed Pareto curves for individual GLUE benchmarks. All models were trained on roughly 2 epochs of C4 (178,000 steps with batch size 4096 with maximum sequence length of 128). Error bars represent  $95\%$  confidence interval for  $n = 2 - 3$  pretraining seeds. Dashed black line represents BERT-Large accuracy on the GLUE (dev) data as reported in [37]. The average GLUE score can be found in Figure 3.
D MosaicBERT-Base Experiments with Increasing Sequence Length
Batch size, maximum sequence length, and number of training steps are important hyperparameters that affect both downstream GLUE accuracy as well as wall clock time. While increasing the maximum number of tokens in a sequence leads to a direct increase in wall clock time, it also exposes the model to more data. For example, increasing sequence length from 128 to 256 while keeping all other hyperparameters constant leads to  $2 \times$  the amount of total training tokens (assuming no padding). However, this will lead to a commensurate increase in pretraining wall clock time. Given our Pareto-based approach, we can therefore ask: is it better to train with fewer steps and longer sequence length, or more steps and shorter sequence length?
Here we find that training for more steps (i.e. 178,000) on shorter sequence length (i.e. 128 tokens) is Pareto optimal for GLUE scores.
As part of our experimental benchmarking, we trained MosaicBERT-Base models with varying maximum sequence lengths ranging from 256 tokens to 2048 tokens for 70,000 steps with batch size 4096. We then plotted the downstream GLUE accuracy in Figures S3 and S4. Besides maximum sequence length, all other pretraining and finetuning hyperparameters are the same as MosaicBERT-Base in the main text (and Tables S1 and S2). Each point represents  $n = 1 - 2$  pretraining checkpoints finetuned with multiple seeds according to Table S2. The curve representing MosaicBERT-Base trained with sequence length 128 and 178,000 steps is the same as in Figure 3.
When looking at average GLUE accuracy as a function of pretraining steps in Figure S3A, longer sequence lengths of 256, 512, 1024 and 2048 nominally do better than a maximum sequence length of 128 tokens when trained for 70,000 steps. However, plotting steps on the x-axis can be misleading, as increasing sequence length leads to increased wall clock time (ceteris paribus).
Instead, it is much more informative to plot wall clock time on the x-axis, as in Figure S3B. Here, the order of the curves is flipped and the curve representing maximum sequence length of 128 tokens is mostly Pareto optimal (red curve). Interestingly, training with batch size and sequence length of  $\{4096,128\}$  for 178,000 steps leads to the same final average GLUE accuracy and wall clock time as training with  $\{4096,512\}$  for 70,000 steps (purple curve in Figure S3B). Note that the x-axis here is
B
Figure S3: MosaicBERT-Base Pareto curves for maximum sequence lengths ranging from 128 to 2048 tokens, with all other hyperparameters fixed. (A) Average GLUE score as a function of pretraining steps. (B) Same data as in A, but with wall clock time on the x-axis. The curve representing a maximum sequence length of 128 tokens is mostly Pareto optimal (red curve). Interestingly, training with batch size and sequence length of \{4096,128\} for 178,000 steps (red curve) leads to the same final average GLUE accuracy and wall clock time as training with \{4096,512\} for 70,000 steps (purple curve). Top dashed line represents BERT-Large average (dev) GLUE score from Devlin et al. 2018, while bottom dashed line represents the BERT-Base average (dev) GLUE score from Devlin et al. 2018 [14].
plotted on a logarithmic scale, as training with sequence length 2048 for 70,000 steps takes more than 100 hours on  $8 \times \mathrm{A}100$  80GB GPUs.
Figure S4: MosaicBERT-Base Pareto curves for maximum sequence lengths ranging from 128 to 2048 tokens, with all other hyperparameters fixed. Individual GLUE scores are plotted as a function of pretraining wall clock time. All curves used batch size 4096 and were pretrained for 70,000 steps, except for the curve representing sequence length 128 tokens, which was trained for 178,000 steps (and is the same as MosaicBERT-Base in Figure 3).
Figure S4 plots the GLUE scores for individual tasks as a function of pretraining wall clock time. The curve representing maximum sequence length of 128 tokens is mostly Pareto optimal (red curve). However, QNLI and STSB seem to benefit from longer maximum sequence lengths.