SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Consequently, researchers have proposed initialization and scaling schemes which remedy this issue [1, 71, 4, 19]. Layer-scale with initialization 0 is an example of one such scheme—at initialization the transformer is an identity function. While  $\gamma_{1},\gamma_{2}$  are typically initialized as vectors of  $10^{-4}$  or  $10^{-6}$ , we use 0 for simplicity.
Figure 2 (right) demonstrates that the layer-scale intervention is successful at controlling the average magnitude output. Without the intervention, the average feature magnitude  $\mathbb{E}[\mathrm{abs}(x_k)]$  becomes high for later blocks. Previous work [17] has shown that large feature magnitudes result in issues for low precision training.
Results for simulated fp8 training are shown in Figure 2 (left) for ViT-Large. We find that all fp8 runs diverge except for when we use layer-scale initialized to zero. Concretely, Figure 2 compares i) the baseline which uses bfloat16 training, ii) using fp8 with tensor-wise quantization and no further modifications, which slowly diverges, iii) adding gradient clipping to ii), which also diverges, iv) adding KQ layernorm [14] to ii), which also diverges, and v) using zero-init layerscale, which trains without diverging. While there is a difference still between fp8 and bfloat16 training, this is primarily because of layerscale. Moreover, we believe that with hyperparameter tuning layerscale would match standard training in terms of accuracy.
Figure 3: Loss spikes increase with model size for fixed learning rate and batch size. Reducing AdamW  $\beta_{2}$  from its default in PyTorch of 0.999 mitigates loss spikes. Reducing  $\beta_{2}$  too much slows training.
3 Stability
We now switch focus from accelerating learning by reducing precision to addressing instabilities which can arise during training. Section 3.1 reviews preliminaries and related work while Section 3.2 details the experimental setup. Next, Section 3.3 examines trends for training instability, finding loss spikes to increase with model scale but decrease with lower AdamW  $\beta_{2}$ . Then, Section 3.4 finds that loss spikes arise in our setting due to an out-of-date AdamW second moment estimator leading Section 3.5 to adopt and tests a fix developed in the context of AdaFactor [54]. Appendix Section G connects this section on stability with the previous section on low-precision training.
3.1 Preliminaries and related work
Loss spikes can emerge when scaling up models [8, 25, 14, 68, 70, 54, 73]. These instabilities may slow learning, or even destabilize training completely. Various solutions have been proposed, including freezing the embedding layer [8], adding additional layer normalization [14, 25], or reparametrizing the weights [68].
In our work we investigate instabilities which arise during CLIP training. Unlike the instabilities observed in [14, 68] which lead to a slow divergence, we study fast loss spikes. Our results indicate that these spikes arise when the second moment estimator is out of date for early layers.
While our analysis and methods build directly on Shazeer and Stern [54] (AdaFactor), there are important differences. In contrast with Shazeer and Stern [54], who only observe instabilities without warmup, we observe instabilities despite a long warmup period. Moreover, in contrast with Shazeer and Stern [54] we find that an out-of-date second moment estimator is primarily an issue for the (patch) embedding layer, and measure how well loss spikes are predicted by this event. Finally, we note that researchers have moved away from AdaFactor in its original formulation for large-scale training [47, 11, 69], finding AdaFactor to under-perform AdamW [47]. We believe this is due to the factored second moment or absence of first moment. This is why our focus is AdamW [37] which is the de facto standard optimizer for transformers.
After the initial version of this paper we became aware of Cohen et al. [12] which offers a general and principled treatment of fast loss spikes, and which we recommend to readers. Moreover, we direct the readers attention to the concurrent work of [41].
3.2 Experimental setup
As in Section 2, we train ViT CLIP models on LAION [53] using OpenCLIP [29] and evaluate them zero-shot on ImageNet. Since we are not interested in final performance and instead interested in studying instability—even for very large models—we use a short run which allows us to conduct multiple experiments. Concretly, we use patch-dropout 0.5 [35] and 20k iterations. The first 5k iterations are linear warmup while the remainder are cosine decay [36]. We follow the CLIP paper [46] in that i) we do not use gradient clipping unless otherwise mentioned, though we do clip the logit_scale parameter, and ii) we add a layer-norm after the patch embedding and before the main transformer. Unless otherwise mentioned, experiments use batch size 16384 (per-gpu batch size of 256), learning rate 2e-3 and weight decay 0.2. We initially tried adding a layer-norm before the patch embedding as in [34], but removed this as we found it to hurt performance at CLIP ViT-Huge scale.
Figure 4: The learning signal can change so that the AdamW second moment estimator  $u_{t}$  is out-of-date and underestimates the squared gradients  $g_{t}^{2}$ . This can be detected if the aggregate quantity  $\mathrm{RMS}_t = \sqrt{\mathbb{E}[g_t^2 / u_t]}$  is far from 1. This figure observes a predictive relationship between the event of an RMS spike and a loss spike—we observe a spike in  $\mathrm{RMS}_t$  1-8 iterations before a loss spike. For lower  $\beta_{2}$ ,  $\mathrm{RMS}_t$  does not deviate far from 1. This result looks at  $\mathrm{RMS}_t$  for the patch embedding layer only. This predictive relationship is further examined in Figures 15 to 20 of Appendix E.
3.3 Loss spikes increase with model size, batch size, and learning rate
We begin our studying of loss spikes by observing how their presence varies when changing model size, batch size, and learning rate. The following sections build on these observations—in particular the finding that lowering the AdamW  $\beta_{2}$  hyperparameter removes spikes entirely.
We find that loss spikes increase when increasing model size, batch size, or learning rate. The first result is shown in Figure 3 while the second two analogous results are in Appendix Figures 12 and 13. Importantly, these figures show that loss spikes can be avoided by reducing the  $\beta_{2}$  hyperparameter for in AdamW. On the other hand, if  $\beta_{2}$  is reduced too much then learning is slowed which results in worse performance [50].
3.4 On  $\beta_{2}$  and an out-of-date second moment estimator
Based on the observation in the previous section that lowering  $\beta_{2}$  reduces spikes, this section traces the cause of loss spikes to an out-of-date second moment estimator in the patch embedding layer.
Overview. Adaptive optimizers such as AdaGrad [22], Adam [33], or AdaFactor [54] scale the update differently for each individual parameter. This is often conceptualized a per-parameter learning rate. For instance, in Adam/AdamW, per-parameter updates are scaled by the inverse root of the exponential moving average of squared gradients (see the code for AdamW in Algorithm 2, ignoring for now the modifications in pink which we discuss in Section 3.5).
This adaptivity can be a very useful tool for accelerating training, but can also cause issues when the learning signal changes. Concretely, exponential moving averages can become out of date causing updates to be scaled by a value that is too large. This issue is discussed in Section 5 of Shazeer and Stern [54], and we summarize below.
As in Algorithm 2, let  $u_{t} = \{u_{t,j}\}_{j=1}^{n}$  denote the exponential moving average (EMA) of squared gradients  $g_{t}^{2} = \{g_{t,j}^{2}\}_{j=1}^{n}$  for neural network parameters  $\theta \in \mathbb{R}^n$ . Ignoring the bias correction term<sup>1</sup>, at each iteration  $t$ ,  $u_{t}$  is updated as  $\beta_{2}u_{t-1} + (1 - \beta_{2})g_{t}^{2}$  where  $\beta_{2}$  is referred to as the decay for the EMA. Then, the update is scaled by  $1 / (\sqrt{u_{t}} + \epsilon)$ , where  $\epsilon$  is a small value added numerical stability. Often the ratio  $v_{t} / (\sqrt{u_{t}} + \epsilon)$  is thought of as signal-to-noise ratio of the gradient over time.
However, this method can break down when the learning signal changes and  $u_{t}$  ceases to be a good estimator for the running average of  $g_{t}^{2}$ . Consider the case where the gradient magnitudes have been historically very small for some parameters so  $1 / (\sqrt{u_t} +\epsilon)$  is large for those parameters. If, then, at iteration  $t$  those parameters suddenly receive a larger gradient signal the update can be catastrophically big. We refer to the scenario as the stuck-in-the-past scenario.
Overall, if  $\beta_{2}$  is too small then convergence may be slowed [50]. If  $\beta_{2}$  is too large then  $u_{t}$  can become out-of-date and no longer a good estimator for  $g_{t}^{2}$ , resulting in per-parameter scaling that is too large.
In practice, the EMA is debiased with a correction term. Algorithm 2 follows AdaFactor section 7.1 in applying the correction term to  $\beta_{1},\beta_{2}$ . Adam is often written with the correction term applied to  $v_{t},u_{t}$  but they are equivalent [54].
Figure 5: Adding update clipping to AdamW mitigates loss spikes and outperforms other interventions such as gradient clipping with norm 1. Code for the AdamW-AdaFactor hybrid we recommend of AdamW + update clipping is in Algorithm 2. The left plot shows loss curves for  $\beta_{2} = 0.99$  while the right displays accuracy abating over  $\beta_{2}$ .
Measurement. We now discuss measurement of the aforementioned stuck-in-the-past scenario and search for a predictive relationship between this event and a loss spike. We follow Shazeer and Stern [54] and measure the following root-mean-square quantity,  $\mathsf{RMS}_t = \sqrt{\mathbb{E}\left[g_t^2 / u_t\right]}$ . If  $u_{t}$  is a good estimator for  $g_{t}^{2}$  then the aggregate quantity  $\mathsf{RMS}_t$  will be around 1. The stuck-in-the-past scenario described above corresponds to an  $\mathsf{RMS}_t \gg 1$ .
As illustrated in Figure 3 we observe instability for high  $\beta_{2}$  in our experiments even though we have  $5\mathrm{k}$  iterations of warm-up. While Shazeer and Stern [54] first recognize the out-of-date second moment estimator issue, in their experimental setting they only observe instability without warm-up.
We now aim to establish a predictive relationship between the stuck-in-the-past scenario and loss spikes. We present initial results in Figure 4, where we examine  $\mathrm{RMS}_t$  for the visual transformer patch embedding layer, visual.conv1.weight. This means that the expectation is computed over parameters in visual.conv1.weight only. This figure illustrates a few important findings: i) loss spikes tend to follow 1-8 iterations after an RMS spike, ii) loss spikes slow learning as recovery time is required, and iii),  $\mathrm{RMS}_t$  stays around 1 for lower  $\beta_{2}$ .
As this is just one example, we further elaborate on the predictive relationship between an RMS spike in the embedding layer in Section E through Figures 15, 16, 17, 18, 19, and 20. For analysis purposes, we define a heuristic to characterize loss and RMS spikes in visual.conv1.weight. We then show that 28 out of 30 detected loss spikes follow an RMS spike by 1-8 iterations, while the probability that a loss spike follows an RMS spike by chance is only  $1\%$ . Moreover, we find that the same predictive relationship does not exist for the RMS in other transformer layers.
3.5 StableAdamW: AdamW with update clipping from AdaFactor
This Section develops and tests StableAdamW (Algorithm 2), an AdamW-Adafactor hybrid.
To stabilize training, the AdaFactor optimizer divides the learning rate for iteration  $t$  by  $1 / \max(\mathsf{RMS}_t, 1)$ . They refer to this as update clipping. The effect is to slow training when  $u_t$  is no longer a good estimator for  $g_t^2$ .
As discussed in Section 3.4, our stability issues can be traced to an out-of-date  $u_{t}$  which is what led Shazeer and Stern [54] to update clipping, even though their stability issues are also solved with warm-up. Therefore, we port update clipping to the standard AdamW optimizer with  $d = 1$  and refer to the resulting AdamW-Adafactor hybrid as StableAdamW (Algorithm 2). A modification we make is to compute and divide learning rate by  $\max(\mathsf{RMS}_t, 1)$  independently for each tensor, which is for implementation convenience. This means that the expectation will be computed independently for each layer to produce a different  $\mathsf{RMS}_t$ .
We now test how StableAdamW compares with other stability interventions such as gradient clipping or lowering  $\beta_{2}$ . These results, presented in Figure 5 find that StableAdamW (i.e., AdamW + update clipping) outperforms these aforementioned interventions for CLIP ViT-Huge. While gradient clipping and update clipping both remove instability, update clipping performs better in terms of zero-shot ImageNet accuracy. With update or gradient clipping, higher  $\beta_{2}$  such as 0.99 tends to
2They actually introduce a hyperparameter  $d$  and use  $1 / \max (\mathsf{RMS}_t / d,1)$  , but recommend setting  $d = 1$  3We clip at global norm 1. We observed instability when trying 2 instead of 1. We did not tune this further, but note that 1.0 is standard in, e.g., PaLM [11], and Scaling Vision Transformers [69].
perform better. Appendix F provides further commentary and implementation considerations for StableAdamW.
4 Limitations, broader impacts, and conclusion
We believe the main limitation of our work is that it is non-exhaustive. For instance, we only simulate float8 training and our experiments focus solely on CLIP-style training. In terms of broader impact, our work may enable additional CLIP models, whose broader impact is examined extensively by Section 7 of Radford et al. [46]. Finally, we believe that our findings on accelerating and stabilizing large multi-modal model training will be broadly useful to the community.
For insightful discussions we thank Romain Beaumont, Yair Carmon, Mehdi Cherti, Brian Cheung, Alex Fang, Gabriel Ilharco, Jenia Jitsev, LAION, Sarah Pratt, Christoph Schuhmann, Ross Whightman, and Sho Yaida. We thank Emad Mostaque and stability.ai for compute resources.
This work is in part supported by NSF IIS 1652052, IIS 17303166, DARPA N66001-19-2-4031, DARPA W911NF-15-1-0543 and gifts from Allen Institute for Artificial Intelligence.
[1] Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence at large depth. In Uncertainty in Artificial Intelligence, pages 1352-1361. PMLR, 2021.
[2] Haoli Bai, Wei Zhang, Lu Hou, Lifeng Shang, Jing Jin, Xin Jiang, Qun Liu, Michael R. Lyu, and Irwin King. Binarybert: Pushing the limit of bert quantization. ArXiv, abs/2012.15701, 2021.
[3] Charlie Blake, Douglas Orr, and Carlo Luschi. Unit scaling: Out-of-the-box low-precision training. arXiv preprint arXiv:2303.11257, 2023.
[4] Andy Brock, Soham De, Samuel L Smith, and Karen Simonyan. High-performance large-scale image recognition without normalization. In International Conference on Machine Learning, pages 1059-1071. PMLR, 2021.
[5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems (NeurIPS), 2020. https://arxiv.org/abs/2005.14165.
[6] Léopold Cambier, Anahita Bhiwandiwalla, Ting Gong, Oguz H. Elibol, Mehran Nekuii, and Hanlin Tang. Shifting and squeezed 8-bit floating point format for low-precision training of deep neural networks. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL https://openreview.net/forum?id=Bkxe2AVtPS.
[7] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, et al. Symbolic discovery of optimization algorithms. arXiv preprint arXiv:2302.06675, 2023.
[8] Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9640-9649, 2021.
[9] Mehdi Cherti, Romain Beaumont, Ross Wightman, Mitchell Wortsman, Gabriel Ilharco, Cade Gordon, Christoph Schuhmann, Ludwig Schmidt, and Jenia Jitsev. Reproducible scaling laws for contrastive language-image learning. arXiv preprint arXiv:2212.07143, 2022.
[10] Minsk Cho, Keivan A Vahid, Saurabh Adya, and Mohammad Rastegari. Dkm: Differentiable k-means clustering layer for neural network compression. arXiv preprint arXiv:2108.12659, 2021.
[11] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
[12] Jeremy M Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardozo, Zachary Nado, George E Dahl, et al. Adaptive gradient methods at the edge of stability. arXiv preprint arXiv:2207.14484, 2022.
[64] Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning, pages 23965-23998. PMLR, 2022.
[65] Guangxuan Xiao, Ji Lin, Mickael Seznec, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. arXiv preprint arXiv:2211.10438, 2022.
[66] Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong He. Zeroquant: Efficient and affordable post-training quantization for large-scale transformers. arXiv preprint arXiv:2206.01861, 2022.
[67] Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang, Hanyu Lai, Ming Ding, Zhuoyi Yang, Yifan Xu, Wendi Zheng, Xiao Xia, et al. Glm-130b: An open bilingual pre-trained model. arXiv preprint arXiv:2210.02414, 2022.
[68] Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Josh Susskind. Stabilizing transformer training by preventing attention entropy collapse. arXiv preprint arXiv:2303.06296, 2023.
[69] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers, 2021. https://arxiv.org/abs/2106.04560.
[70] Xiaohua Zhai, Basil Mustafa, Alexander Kolesnikov, and Lucas Beyer. Sigmoid loss for language image pre-training. arXiv preprint arXiv:2303.15343, 2023.
[71] Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. arXiv preprint arXiv:1901.09321, 2019.
[72] Susan Zhang. Open pretrained transformers lecture, 2023. https://www.youtube.com/watch?v=p9IxoSkvZ-M.
[73] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022.
[74] Wei Zhang, Lu Hou, Yichun Yin, Lifeng Shang, Xiao Chen, Xin Jiang, and Qun Liu. Ternarybert: Distillation-aware ultra-low bit bert. In EMNLP, 2020.
[75] Changsheng Zhao, Ting Hua, Yilin Shen, Qian Lou, and Hongxia Jin. Automatic mixed-precision quantization search of bert. arXiv preprint arXiv:2112.14938, 2021.
[76] Kang Zhao, Sida Huang, Pan Pan, Yinghan Li, Yingya Zhang, Zhenyu Gu, and Yinghui Xu. Distribution adaptive int8 quantization for training cnns. In Proceedings of the Thirty-Fifth AAAI Conference on Artificial Intelligence, 2021.
[77] Chenzhuo Zhu, Song Han, Huizi Mao, and William J. Dally. Trained ternary quantization. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https://openreview.net/forum?id=S1_pAu9x1.
[78] Feng Zhu, Ruihao Gong, Fengwei Yu, Xianglong Liu, Yanfei Wang, Zhelong Li, Xiuqi Yang, and Junjie Yan. Towards unified int8 training for convolutional neural network. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1969-1979, 2020.
Figure 6: (Left) Individually profiling operations which constitute a forward and backward pass in a linear layer for i) SwitchBack using triton kernels and ii) an fp16 baseline using torch/matmul. Times are averaged over a linear layer from dim to  $4\cdot$  dim and a linear layer from  $4\cdot$  dim to dim—representative of the linear layers in a transformer MLP. (Right) The  $\%$  speedup of SwitchBack over a standard fp16 linear layer when all operations in Figure 6 (left) are summed.
A Additional Related Work on Quantization
The literature on training neural networks in low-bit precision is vast. The main differentiating factor of our work is that we train relatively large models – in fact, we train the largest 8-bit vision transformers to date.
The literature agrees that quantization of very large networks is more difficult than for smaller networks [18, 17, 65, 24]. As such, we divide our related work into three parts: (1) large-scale low-precision neural network (larger than BERT-large), and (2) low-precision training of smaller networks.
Large-scale Low-precision Neural Networks. Our work is currently the only work that does low-precision (8-bit and below) training of very large networks with more than 230M parameters. Other related work studies inference at scale. SmoothQuant [65], ZeroQuant [66], NuQmm [42], and LLM.int8() [17] study inference with Int8 matrix multiplication. Another line of work studies large models inference with more than 250M parameters by considering 16-bit inputs and k-bit weights [16, 24, 67].
Small Scale Low-precision Training Training of small-scale low-precision neural networks can take many shapes and forms, such as quantization for integer only devices, quantization for mobile device, or quantization to accelerate training. One way to break up these directions is through the data type used and the neural network trained. One major direction is to quantize convolutional neural networks often for fast and memory efficient usage on edge devices [77, 6, 13, 78, 23, 76, 30]. Further work in this area is discussed in the survey by [45]. Another line of work is centered around 8-bit float data types which can be used to accelerate training of neural networks [21, 56, 62, 38, 3]. Lastly, a common application is to finetune (similar to training) BERT models to particular datasets. This not only decreases the model footprint and increases inference speed but adjusts the model to new data [2, 32, 74, 55, 75].
B Achieving speed-ups with SwitchBack
We now test the speedups offered by SwitchBack by first examining individual operations and then end-to-end training.
We profile all of the operations which constitute a forward and backward pass for a single linear layer in Figure 6 (left) for both SwitchBack and the baseline. For SwitchBack we profile our custom triton kernels and for the baseline we profile torch.matmul. Overall, we observe that int8 multiplies occupy just over half the time as standard fp16 matmuls, and that quantize operations are roughly an order of magnitude less time than a matmul. Note that our int8 matmuls are fused with the dequantize operation.
Figure 6 (right) displays the  $\%$  speedup of SwitchBack over a standard fp16 layer when all operations in Figure 6 (left) are summed. Overall, the advantage of SwitchBack is greater for larger dim and batch_size * sequence_length. Overall, the speedup ranges from  $5\%$  to  $35\%$ . We see a bump at  $\dim = 1280$  because standard PyTorch matmuls do not have optimized kernels for matrices of this