Small-scale proxies for large-scale Transformer training instabilities
Authors: Mitchell Wortsman
Peter J Liu
Lechao Xiao
Katie E Everett
Alexander A Alemi
Ben Adlam
John D Co-Reyes
Izzeddin Gur
Abhishek Kumar
Roman Novak
Jeffrey Pennington
Jascha Sohl-Dickstein
Kelvin Xu
Jaehoon Lee
Justin Gilmer
Simon Kornblith
Venue: ICLR 2024 oral (2024)
Keywords: Small Transformers
Training
Stability
Primary area: optimization
TL;DR: We seek ways to reproduce, study and predict training instability with smaller models.

Abstract: Teams that have trained large Transformer-based models have reported training instabilities at large scale that did not appear when training with the same hyperparameters at smaller scales. Although the causes of such instabilities are of scientific interest, the amount of resources required to reproduce them has made investigation difficult. In this work, we seek ways to reproduce and study training instability at smaller scales. First, we focus on two sources of training instability described in previous work: the growth of logits in attention layers (Dehghani et al., 2023) and divergence of the output logits from the log probabilities (Chowdhery et al., 2022). By measuring the relationship between learning rate and loss across scales, we show that these instabilities also appear in small models when training at high learning rates, and that mitigations previously employed at large scales are equally effective in this regime. This prompts us to investigate the extent to which other known optimizer and model interventions influence the sensitivity of the final loss to changes in the learning rate. To this end, we study methods such as warm-up, weight decay, and the MuParam (Yang et al., 2022), and combine techniques to train small models that achieve similar losses across orders of magnitude of learning rate variation. Finally, to conclude our exploration we study two cases where instabilities can be predicted before they emerge by examining the scaling behavior of model characteristics such as activation and gradient norms.

Published as a conference paper at ICLR 2024
SMALL-SCALE PROXIES FOR LARGE-SCALE TRANS- FORMER TRAINING INSTABILITIES
Mitchell Wortsman' Peter J. Liu Lechao Xiao Katie Everett
John D. Co-Reyes Izzeddin Gur Abhishek Kumar Roman Novak Jeffrey Pennington Jascha Sohl-dickstein Kelvin Xu Jachoon Lee" Justin Gilmer" Simon Kornblith
Alex Alemi Ben Adlam
Google DeepMind
ABSTRACT
Teams that have trained large Transformer-basedd models have reported training instabilities at large scale that did not appear when training with the same hy- perparameters smaller scales. Although the causes of such instabilities are of scientific interest, the amount of resources required to reproduce them has made investigation difficult. this work, we seek ways to reproduce and study training instability at smaller scales. First, we focus on two sources of training instabil- ity described in previous work: the growth of logits in attention layers (Dehghani et al., 2023) and divergence of the output logits from the log probabilities (Chowd- hery et al., 2022). By measuring the relationship between learning rate and loss across scales, training at high learning rates, and that mitigations previously employed at large scales equally effective in this regime. This prompts to investigate the extent to which other known optimizer and model interventions influence the sensitivity of the final loss to changes in the learning rate. To this end, we study methods such as warm-up, weight decay, and the μParam (Yang et al., 2022), and combine techniques to train small models that achieve similar losses across orders of mag- nitude of learning rate variation. Finally, to conclude our exploration we study two cases where instabilities can be predicted before they emerge by examining the scaling behavior of model activation and gradient norms.
show that these instabilities also appear in small models when
INTRODUCTION
Scaling up transformers has led remarkable progress from chat models image generation. How- ever, not every training run is successful. When training large Transformers, researchers have re- ported instabilities which slow destabilize learning (Chowdhery et al., 2022; Dehghani et al., 2023; Zhang et al., 2022; Molybog et al., 2023; Cohen et al., 2022). As the resources required for large runs continue to grow, it is important to examine the ways that Transformer training can fail.
In this report we reproduce, study, and predict training instability in Transformer models. We find that measuring the relationship between learning rate and loss across scales is a useful tool identify instability (c.g.. Figure 1). Therefore, we introduce learning rate (LR) sensitivity, which serves as a useful summary statistic for learning rate vs. loss curves. LR sensitivity measures the deviation from optimal performance when varying LR across orders of magnitude.
We show that two sources of instability, which have previously been described scale, can be reproduced in small Transformers.' This enables their study without access to large resource pools. In particular, we examine the growth of logits in attention layers (Dehghani et al., 2023; Gilmer all; Zhai et al., 2023a) and divergence of the output logits from the log probabilities (Chowdhery al., 2022). As evident from the learning rate vs. loss curves and by inspecting model characteristics, both instabilities appear high learning rates in small models. Moreover, interventions which have previously been employed at scale are also successful in this regime (e.g.. Figure 1). These
"qqual contribution. Authors are now at Anthropic. We focus instabilities which lead slow divergence, not loss spikes (see Section 4).

---

Published as a conference paper at ICLR 2024
qk-layemorm (Dehghani et al., 2023) and z-loss regularization (Chowdhery et al., 2022)-reduee reduce LR sensitivity and enable successful training across three orders magnitude of learning rate variation.
interventions
These observations raise the question of how other known optimizer and model interventions affect the shape of the learning rate vs. loss curves across scales. Therefore, we study the effect of techniques such as warm-up. weight decay, and pPaarm (Yang et al., 2022) in this context. When employing la-layernorm and z- loss regularization, these other techniques usu- ally have little impact on the range of learning rates which models can be stably trained, but do affect the sensitivity to learning rate within this range. In line with previous work, we find that longer warm-up reduces learning rate sensitivity, as does the independent scaling of learning rate and weight decay recommended by Loshchilov Hutter (2019). One interest- ing finding is that scaling depth increases LR sensitivity at faster rate than scaling width.
4.25
ok.layemerm True
akllayemorm False
9.4
199+07
0.2e+07
8.5
Le
3.75
3. 3.0e+008
Final eval lass
3.50
3.25
2,50
10
10
The remainder of our investigation centers the scaling behavior for model characteristics such as activation and gradient norms. Using the attention logit growth instability as ex- ample, we show that it possible to predict an
Learning rate
LR sensitivity
10a
10-1
instability before it emerges. This is in contrast to prior works on scaling which primarily focus on scaling trends related to loss (Kaplan et al., 2020; Hoffmann et al., 2022).
10
10
Number of parameters
Figure 1: Qk-layernorm (Dehghani et al., 2023) enables stable training across three orders of mag- nitude of learning rate (LR) variation. (Top) For transformers with N parameters, we plot the ef- fect of learning rate on final evaluation loss. (Bot- tom) We use LR sensitivity to summarize plot. LR sensitivity measures the expected devia- tion from the minimum achieved loss when vary- ing learning rate across three orders of magnitude. layleyoroorm reduces LR sensitivity, but LR sen- sitivity still increases with model scale.
We conclude by using the scaling behavior of model characteristics to search for instabili- ties that currently not well documented. Our investigation shows that gradient norms de- crease with both scale and learning rate, such that the default AdamW W (Loshchilov & Hutter, 2019) epsilon hyperparameter
top
large. This causes updates that are too small. We connect this phenomenon and the attention logit growth instability to parameter norm growth (Merrill al., 2020; Lee, 2023).
Overall, we believe our work presents new scientific opportunities for studying training stability without access to large resource pools.
2 EXPERIMENTAL METHODOLOGY
This section details our experimental set-up (Section 2. and useful tools employed by our analysis: (i measuring the relationship between learning rate and loss across scales (Section 2.2) and (ii) examining scaling trends for model characteristics (Section 2.3).
2. EXPERIMENTAL SET-UP
We train small Transformer models (Vaswani et al., 2017) with similar experimental set-up GPT- 2 (Radford et al., 2019) implemented Flax (Heek et al., 2023): the models are decoder-only (Liu al., 2018) and trained with auto-regressive loss (refer Section A for infrastructure
2

---

Published as a conference paper at ICLR 2024
details). While we experimentally manipulate many of the following hyperparameters, this section provides their default values, which we use unless otherwise specified.
By default, gradient clipping at global norm 1. The default warmup is 5e3 steps, and the default number of total steps is le5. We 2016) schedule for the remainder, with minimum learning rate 1e-5. We use an independent weight decay of le-4 and auxiliary z-loss (Chowdhery al., 2022) with coefficient le-4. Sections 3.2.2 and 3.1.2 respectively provide additional information and ablations on decoupled weight decay and z-loss. We use pre-normalization (Radford et al., 2019) Transformers with qk-layernorm (Dehghani et al., 2023) (see Section 3.1.1 for information). We do not use any biases following Chowdhery et al. (2022), and the layernorm (Ba et al., 2016) e remains at the default value in Flax (Heek al., 2023) of le-6. We jointly scale up the embedding size, depth, and number of heads when scaling parameters. We do not use weight tying of the first and last layer (Press & Wolf, 2017), and when reporting the number of parameters we exclude the embedding and head (as in Kaplan et al. (2020)). We use rotary positional embeddings (Su et al., 2021), and for training data we use C4 (Raffel et al., 2020a). Letting d refer to the model dimension (i.e., the embedding size), the feed-forward component of the Transformer is an MLP with hidden dimension of 4d and gelu (Hendrycks & Gimpel, 2016) activations. As in Vaswani et al. (2017) we use factor 1/vā scaling in the self- attention. The embedding initialization is the default in Flax, which is normally distributed with standard deviation 1/Vd. The remainder of the weights are initialized with a truncated normal distribution with inverse root fan-in standard deviation (Glorot & Bengio, 2010). The default batch size is 256, where each batch element has a sequence length of 512 tokens. Sequences are packed so that no padding is required. Finally, we use the vocabulary from Raffel et al. (2020b) which has size 32101 and uses a SentencePiece (Kudo & Richardson, 2018) tokenizer. We train on TPUs (Jouppi et al., 2017) in 6 precision using Flax (Heek et al., 2023) and JAX (Bradbury et al., 2018).
AdamW (Loshchilov & Hutter, 2019) with 31 0.9, 32 0.95, 1c-8, and
linear schedule for warmup and and a cosine-decay (Loshchilov & Hutter,
2.2 LR vs. LOSS CURVES AND LEARNING RATE SENSITIVITY
To investigate how model instability emerges with scale, it is useful to plot the relationship between learning rate (LR) and loss for models of different sizes. For instance, an instability is often char- acterized by an explosion in the loss at high learning rates. LR vs. loss curves can reveal how the lowest unstable learning rate changes
function of model size.
To summarize LR vs. loss curves, we use LR sensitivity. LR sensitivity measures the deviation in final validation loss from optimal when sweeping LR across three orders of magnitude. If model fails to train at high learning rates, then LR sensitivity will be high. There are cases where LR vs. loss curves and LR sensitivity are no longer meaningful, for instance if intervention changes the meaning of learning rate-see Appendix B for a detailed discussion.
Let 0 A(n)) denote the model weights 0 obtained when training with learning rate n, and let (()) denote the validation loss when using weights 6. For learning rate range [a, b], let E denote the loss obtained with the best learning rate, i.c., e- min é (A(r)). Moreover, let Eo denote loss at initialization. Then, LR sensitivity is defined E, [min (e (AA (o) C*).
Unless otherwise mentioned, we use the learning rate range 3e-4 to 3e-1 with AdamW (Loshchilov & Hutter, 2019) to measure LR sensitivity, where LR refers the maximum value in cosine decay schedule with warm-up (Loshchilov & Hutter, 2016). We consider LRs in (3e-4, le-3, 3e-3, le-2, 3e-2, le-1, 13 when computing the minimum and expectation.
2.3 SCALING TRENDS FOR MODEL CHARACTERISTICS
To study instability, we also find it useful to examine scaling trends for model characteristics such as gradient or activation norms. This method helpful for predicting instabilities and contrasts with previous work on scaling, which primarily focuses on trends relating model scale and loss (Kaplan et al., 2020; Hoffmann et al., 2022).
RESULTS
This section presents our results on training stability for small Transformers. Equipped with LR sensitivity (Section 2.2), we study two known instabilities and their corresponding mitigation at
3

---

Published as a conference paper at ICLR 2024
small scale (Section 3.1). This raises the question of how other model and optimizer interventions effect sensitivity of final loss to learning rate, which we investigate in Section 3.2. Finally, we examine whether instabilities can be reliably predicted before they emerge: Section 3.3 predicts when the logit growth instability may cause divergence a larger model, while Section 3.4 aims to find other issues that may occur when scaling up with default hyperparameters.
3.1 REPRODUCING TWO KNOWN INSTABILITIES AT SMALL SCALE
Here, we examine two instabilities that have previously been described at scale: the growth of logits in attention layers (Dehghani et al., 2023; Gilmer et al.; Zhai et al., 2023a) and divergence of the output logits from the log probabilities (Chowdhery et al., 2022). By examining LR vs. loss curves, we show that these instabilities can be reproduced in small models by using high learning rates and that mitigations employed at scale are effective in this regime.
3.1.1 ATTENTION LOGIT GROWTH
Researchers have previously documented that Transformer training fails when the attention logits become large (Dehghani et al., 2023; Zhai et al., 2023a). In Dehghani et al. (2023), this issue emerged when training ViT model (Dosovitskiy et al., 2021) with 22 billion parameters.
In the self-attention layer Transformer (Vaswani et al., 2017). queries qi and keys k are com- bined compute the attention logits 2,j = (q: k:y) /dh, where d is the head dimension. Next, the attention logits are passed through a softmax to produce attention weights, which are used to com- bine values Ug. Dehghani et al. (2023) observed that the attention logits z became large, which they refered to as attention logit growth. As result, the attention weights collapse to one-hot vectors, which was named attention entropy collapse by Zhai al. (2023a). To resolve this issue, Dehghani et al. (2023) proposed ql-layernorm, which applies LayerNorm (Ba et al., 2016) to the queries and keys before computing the attention logits.
In our experiments, we find that models need not be large to exhibit instability related to attention logit growth. As shown in Figure 1, the maximum learning rate which small models can be trained increases when using qk-layernorm. Without lk-layernorm, the learning rate at which models di- verge becomes smaller with increasing model size. By contrast, models with qk-layernorm exhibit considerably lower LR sensitivity and train to low loss at high learning rates. As a highlight, qk- layernorm allows training a model with 1.2B parameters at learning rate 0.3. Both with and without gk-layemnrm, LR sensitivity increases with scale.
Figure G.I displays the loss and max attention logit for two model scales that differ by three orders of magnitude. In both cases, the loss diverges without gk-layernorm. Our results in Appendix Figure G.7 suggest that attention logit growth is due to growth in the queries and keys, not due to increase in their alignment. Finally, Appendix C connects this instability to the quadratic dependence of attention logits on parameter norms.
3.1.2 OUTPUT LOGIT DIVERGENCE
Another instability reported by researchers training large models is divergence in the output logits from the log probabilities (Chowdhery et al., 2022). Just as before, we reproduce this instability with small models at large learning rates, and the proposed mitigation ameliorates the issue. Overall, Figure summarizes the effect.
Let y denote the model's output logits, which used to compute class probabilities p; via softmax = e /Z where Z E, evi This instability occurs when the logits diverge and become very negative, as illustrated in Figure G.2 for a 2.4M parameter model at learning rate 0.1. In contrast to the attention logit growth instability, this divergence occurs towards the end of training. The mitigation proposed by Chowdhery et al. (2022) is to encourage log Z to remain close to zero. They add an auxiliary loss log² Z, referred to as z-loss, with coefficient le-4.
As illustrated in Figures 2, we find that instability related to output logit divergence occurs in models with no weight decay regardless of scale, and z-loss resolves this instability. Weight decay also mitigates this instability for the larger models we test.
4

---

Published as a conference paper at ICLR 2024
4.4
2-loss confficient 1e-4
2,4e+06
50
2.4+
Warmup -50
9.4e+
Weight decay 1e-4
1,9e-07
Warmup 5000
1.9+ 07
weight decay
4,20+07
Warmup 10000
07
8,5e+07
Warmup 25000
0.5e+ 07
4.0
1
3.8
3.0e+
3.8
Final eval loss
3.6
Final
2.4
3.4
3.2
3.2
Learning
Learning rate
LR sensitivity
sensitivity
10a
10-1
6x38
10-1
Number of parameters
Number of parameters
Figure 2: The effect of the output logit diver- gence instability (Chowdhery al., 2022) and the z-loss mitigation (Chowdhery et al., 2022) (Section 3.1.2). Models in this experiment have qk-layernorm (Dehghani et al., 2023).
Figure The effect warm-up length for dif- ferent model sizes. Longer warm-up reduces LR sensitivity and loss, especially for the larger models test. Models in this experiment qk-layernorm (Dehghani et al., 2023).
3.2 MEASURING THE EFFECT OF OTHER KNOWN INTERVENTIONS
The previous section used the relationship between learning rate and loss as a useful tool for ex- amining known instabilities and their mitigation. This raises the question of how other known model and optimizer interventions affect the shape of LR vs. loss curves across scales. In particular, can LR sensitivity help identify additional issues resolutions when scaling? This section aims to answer this question for common techniques such as warm-up, weight decay, and μParam (Yang al., 2022). Additional interventions and hyperparamters are tested Appendix Section D.
3.2. WARM-UP
As illustrated by Figure longer warm-up period reduces LR sensitivity. This is most clear for the larger models, which are not stable at LR 3e-1 without long warm-up. The number total steps is fixed le5 in this experiment, and all models gk-layernorm. The importance of warm-up for stability has previously been highlighted (Gilmer et al., 2021; Shazeer Stern, 2018; Liu et al., 2019), although these works do not measure scaling behavior.
3.2.2 INDEPENDENT WEIGHT DECAY
Parameterizing weight decay independently of learning rate reduces LR sensitivity, as illustrated in Figure While this was recommended by Loshchilov & Hutter (2019), not common practice in the default AdamW implementations of PyTorch (Paszke al., 2019) Optax (Babuschkin et al., 2020). We explain the differences below.
For parameters 0, let A v/ (VT+= e) denote the AdamW update without learning rate weight decay. For weight decay coefficient X, max learning rate 77, and schedule s; E [0, 1], Loshchilov & Hutter (2019) recommend the update 0 / st(nA AO), which refer to independent decay. On the other hand, the default implementation in PoTch or Optax applies the update t- ( X0), i.e., 1) scales both terms.
5

---

Published as a conference paper at ICLR 2024
Independent decay True
N 2.4++06
Scaling wiath
+ dim 256
layors
Independent decay
layet
== 1,9e+00
dim 768
layers 12
42++07
layers
4.0
8.5e+007
dim 2048
layers 48
15
3.0e+
Final eval evallos loss
Fnall eval loss
3.6
22
2.0
10-4
10-3
10-2
10
10
10
10-2
10-2
10F
Learning rate
Learning rate
10"1
LR sensitivity
sensitivity
Scaling
Scaling depth
107
100 parameters
Number parameters
Figure 4: Independently scaling LR without also scaling weight decay reduces LR sensitiv- ity. While this was recommended by Loshchilov & Hutter (2019), it is not common practice in the default A implementations in popular libraries. Refer to Section 3.2.2 for more detail.
Figure 5: Independently scaling depth increases LR sensitivity at faster rate than scaling width, though also produces a model with lower loss at the largest scale we test. Refer to Ap- pendix Figure G.8 for this experiment without qk-layernorm.
When reporting LRR sensitivity without independent decay in Figure we report the minimum LR sensitivity over ranges [1e-4, le-1] and [3e-4, 3e-1] because the former is sometimes better centered on the minimum. The default setting in this paper to use independent decay. When using indepen- dent decay we set X=le-4, and without independent decay we set A=0.1. A sweep weight decay values conducted Figure G.16.
3.2.3 SCALING WIDTH VS. DEPTH
have so far consistently observed that increasing the number of parameters increases LR sensi- tivity. We now examine which part of scaling is most responsible.
Our results, illustrated by Figure indicate that scaling depth increases LR sensitivity faster rate than scaling width. However, at the largest scale we test, independently scaling depth produces a model with lower validation loss. A validation loss comparison between width scaling, depth scaling, and joint scaling is in Appendix Figure G.9. The standard practice of joint scaling performs best at the largest scale and also has a more reliable scaling prediction when extrapolating.
When scaling depth, heads scaled proportionally with width, so that the head dimension remains the same.
= 512, and when scaling width,
6 layers. The number of
Figure G.8 repeats this experiment without qk-layermorm, finding that the attention logit growth instability occurs more frequently at scale regardless of whether width or depth are scaled.
3.2.4 PPARAM
Yang & Hu (2021) introduced the pParam method for parameterizing a neural network. As a prod- uct, the optimal LR remains consistent when scaling model width (Yang et al., 2022). This section tests the effect μPaarm on LR sensitivity, and examines whether 4Param alleviates the need for gk-layernorm (Dehghani al., 2023).

---

Published as a conference paper at ICLR 2024
LR 0.0001
LR = 0.0003
LR 0.003
LR 0.01
LR = 0.03
LR 0.1 PX
Max attention logit
10"
10*
10'
10" 10"
10'
10* 10'
10' 10" 10"
10' 10 10'
10'
10" 10"
10/
10" 10"
Num params
Num params
Num params
Num params
Num params
Num params
ok-layernom True
Successful Instability
Quadratic fit
Additional experiment va validate predicted instability
layenammmmm False
Extrapolated trend
Figure 6: Predicting the attention logit growth instability via scaling behavior of model character- isties. We extrapolate to predict that a larger model will become unstable at LR 1e-2, and run an experiment to confirm the prediction. Refer Section 3.3 for more information.
As illustrated by Figure G.3, jdParam does succeed in stabilizing the optimal LR at the scale we test. However, pPaara does not improve loss reduce LR sensitivity our experiments. Appendix Figure G.4 repeats this experiment without gk-llayrnorm. Our results indicate that μlParm does not alleviate the need for this intervention at high learning rates. We note that from practical perspective, reducing LR sensitivity is not important if the optimal LR does not change.
We refer to the variant of μPaarm that we use in these experiments as pPaarm (simple) because it maintains only the core feature of jPParmm We add additional features from Yang et al. (2022) in Figure G.II without measurable improvement the largest scale we test. For μPaarm (sim- ple) we make the following changes from our standard baseline: scale the LR for linear layers by base-fand ffnn-nnn.. For paaarm (full) there are three additional changes: (i) initialize the head with standard deviation v waase-fn-innfaawe ffnn-in; (ii) change the 1/vd, scaling factor in attention layers to 1 1/d, where d, is the head dimension; and (iii) initialize the query projection weights with zeros. For base-fan-in we use the fan-in values for the smallest model we test, which has width 256. We ablate change (ii) in isolation in Figure G.12. initial experiments (iii) had no noticeable effect.
3.3 PREDICTING ATTENTION LOGIT GROWTH INSTABILITY FROM SCALING BEHAVIOR OF MODEL CHARACTERISTICS
A central question when studying instabilities is whether they can be predicted using small-scale proxy experiments. We now examine whether is possible to predict the logit growth instability before it occurs using previous runs with smaller models. We track the attention logit maximums across model scales and fit a curve to the data. We use this to predict that a 4.8B parameter model will be unstable at LR le-2 without qk-layernorm and run experiment to confirm this prediction.
Figure plots the number of parameters vs. max attention logit at different learning rate values. 2 At each LR, fit quadratic to predict how the max attention logit will change with model scale.
We first noticed that all points with attention logits above 1e4 diverged. Morcover, the quadratic fit predicted that for LR 1c-2 the next model scale would also cross that value. Based on this prediction, we trained new 4.8B parameter model at LR le-2. This model diverged as predicted. Not only do predict the divergence, but our fit closely extrapolates to predict the value of max attention logit.
Refer to Appendix E for a preliminary experiment on whether we could have predicted that instabil- ity arises when the max attention logit exceeds le4 without manipulating LR and model size.
3.4 SEARCHING FOR NEW INSTABILITIES VIA SCALING TRENDS OF MODEL CHARACTERISTICS
This section examines whether the scaling behavior of model characteristics can be used to predict new issues with the default model and hyperparameter settings.
We use block 0, which typically has the largest logits, and consider the value at step 2c3. Much earlier than 2c3 uninformative, and much later the unstable points had long past diverged.

---

Published as a conference paper at ICLR 2024
10-4
10
Grad RMS
10-7
10
10
10-6
10-4
10-0
10-7
10
10.8
10
107
10a
10a
10
10-2
10
5
10
15
20
Num params
LR
Block index
0.0003
0.3
9.4e6
1.2.9
9.4e6
1.2e9
LR
Num params
Num params (at LR 3e-1)
Figure 7: Predicting potential instability from the scaling behavior of model characteristics. The gradient root mean square (RMS) decreases with num params (left) and learning rate (middle). These trends indicate that hyperparameter adjustment may be required to successfully scale further, as the RMS is approaching the default Adam hyperparameter. If the gradient RMS becomes too small without adjusting
weight decay, a layer may collapse. The gradient RMS in the left and middle plot is reported for the first MLP layer of block O, but we observe similar trends for other layers (e.g.. Appendix Figure G.18). Gradient RMS across different blocks is also reported (right). Gradient and update RMS are averaged over the final 500 steps, refer to Appendix Figure G.19 for the data during training.
In Figure 7 we examine scaling trends for the gradient root mean square RMS(g) = / E, | This figure reports the RMS for the first layer of the MLP, though we observe similar trends for other layers (Appendix Figure G.18).
As models get larger, the value that grad RMS approaches is cause for concern. At the largest scale and learning rate we test, grad RMS is around the default AdamW e hyperparameter. Recall that the unscaled AdamW update is A v/ (vū + e), where and u are the first and second gradient moment EMA, respectively. If the grad RMS is on the same order as e, then A will decrease in magnitude as illustrated by Figure G.6, and parameters will not receive learning signals intended.
An obvious mitigation for this issue is to simply lower the A e hyperparameter from its default of 1e-8. We conduct this experiment for a 4.8B parameter model at LR 0.3 and present the results in Figure G.5. Decreasing to le-15 15 improves loss and mitigates a collapse in grad RMS. We believe this improvement will only increase scale. On the other hand, increasing instability (shown in Figure G.10).
1e-6 results in an
Figure G.6 expands on this result by illustrating the grad and update RMS throughout training the largest scale and learning rate we test. When the grad RMS reaches the update RMS becomes small. Figure G.19 presents data from an analogous experiment at many different scales and LRs, demonstrating that this issue is most apparent for the larger models and LRs we test.
Although we identified the instability above by empirically measuring the scaling behavior of the gradients, mechanistic explanation exists. As learning rate increases, so does the parameter RMS. larger parameter RMS leads to a larger RMS for the features output by each Transformer block. Then, the overall output RMS in turn increases with depth due to residual connections. The overall effect is that for larger networks and learning rates, the Transformer output RMS entering the final layernorm will grow. Since the layernorm gradients are scaled by the inverse of their input RMS, the gradient received by the Transformer will shrink. Refer to Appendix C for a detailed discussion.
RELATED WORK
This paper mainly focuses on the effect of known interventions and instabilities, and so related work has been primarily discussed when relevant. This includes the attention growth instability observed by Dehghani et al. (2023); Zhai et al. (2023a), and the final logit divergence issue encountered by Chowdhery et al. (2022). However, we highlight similar experimental methods in previous work. For instance, Yang et al. (2022) also measure the relationship between LR and loss across seales, but their focus is on centering the optimum (see Section 3.2.4). In addition, Zhai et al. (2023a) elicit instability in base models by doubling learning rate, and Dettmers et al. (2022) measure the presence of outlier features function of scale.
8

---

Published as a conference paper at ICLR 2024
There are also important instabilities and related topies have not directly discussed so far. For instance, we have primarily focused on instabilities that lead to a slow divergence, and we now summarize research on fast loss spikes. This instability is characterized by quick increase in the loss that often eventually recovers.
The Edge of Stability and fast spikes. The conventional understanding of gradient descent pre- dicts that loss instability only occurs when the learning rate exceeds 2/ 2/>max (H), where H is the Hessian. However recent investigations into large batch neural network training dynamics have re- vealed more complicated picture via edge of stability (BoS) (Cohen et al., 2021). When training neural networks with large batch SGD, the loss curvature constantly evolves via the interaction of two processes: progressive sharpening and self stabilization. Progressive sharpening is the empirical observation that when L.R < 2/>mux ()) the curvature gradually increases until the stability thresh- old is violated. When the learning rate becomes too large relative to the curvature, faxt loss spikes occur and the parameters oscillate into a region with smaller Xmax (H) where stable training and progressive sharpening resumes. The latter process where instability results in smaller Amax (H) is self-stabilization,, a theoretical model of which is given in Damian et al. (2022). Gradually shrinking Amax (H) via self stabilization
shown be primary mechanism behind the success of learn- ing rate warmup in Gilmer et al. (2021), who closely studied the connections between curvature, initialization, architecture and max trainable learning rates.
Cohen et al. (2022) further analyze edge of stability of dynamics with adaptive optimizers, showing that progressive sharpening interacts with both the self-stabilization process and the adaptive opti- mizer state. This interaction results in the preconditioned sharpness Amax (P-1)) /1) oscillating around an optimizer specific threshold (38/LR in the case of Adam with =0.9). Adaptive EoS (AEoS) can also result in periodic loss spikes when progressive sharpening pushes the preconditioned sharpness above the stability threshold, however the optimizer hyperparameters play a role. In particular, when LR>38/> (P-1 H). two mechanisms are now in play to resolve the step size being too either H can shrink or p-1 can shrink (or both). Cohen et al. (2022) found that when 33 is large, H tends to shrink and fast loss spikes result during the process, resembling the self stabilization process ob- served with gradient descent. However when 32 is small, p-1 tends to shrink, no loss spikes are observed, and Amax (H) tends to gradually increase throughout training.
is noteworthy that the adaptive edge of stability process (and the role of B2) studied in Cohen et al. (2022) offers more complete understanding for loss spikes studied in body of literature (Shazeer & Stern, 2018; Chowdhery et al., 2022; Molybog et al., 2023; Wortsman et al., 2023a; Zhai et al., 2023b; Chen et al., 2021). For example, Shazeer & Stern (2018) argue that during training of Transformers with adaptive optimizers the optimizer update can become too big resulting in a loss spike followed by recovery. This is sometimes attributed to the adaptive optimizer state becoming "stale", which is consistent with the observation the reducing B3 resolves the loss spikes (Shazeer & Stern, 2018; Wortsman et al., 2023a; Zhai et al., 2023b). This is perhaps the same observation as Cohen et al. (2022) that reducing 32 allows p-1 to change quicker to adjust to the process of progressive sharpening. AEoS also offers an explanation for the periodic loss spikes observed when training large transformer models (Molybog et al., 2023).
Parameter-free methods and more parameterizations. While our work has studied sensitivity to learning rate, there is also research that aims to eliminate the need to specify learning rate (Ivgi et al., 2023; Defazio & Mishchenko, 2023). Based on their analysis, Ivgi et al. (2023) set the step size for iteration to the maximum distance from the initialization divided by the root sum of historical gradient squares. Moreover, while our work investigated jPaarm, there are additional parameterizations for which it would be interesting to explore LR vs. loss (Dinan et al., 2023; Yaida, 2022; Bordelon & Pehlevan, 2023; Jacot et al., 2018).
CONCLUSION
This paper demonstrates that useful insights on instability can be gained from small Transformers. Our results indicate that: (1) instabilities previously reported at scale can be reproduced in small- scale proxy models, facilitating their study without access to large resource pools; 2) instabilities previously reported at scale can be predicted before they emerge by extrapolating g from experiments with small-scale proxy models; and 3) new instabilities can be found using small-scale proxy mod- els.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
We thank George Dahl for thorough comments and suggestions, and Hugo Larochelle and Rif A. Saurous for helpful discussion. Also, we thank the members of the Google DeepMind PAGI team for their support of this effort, Noah Fiedel, Noah Constant, Aaron Parisi, Alex Rizkowsky, Avi Singh, Azade Nova, Bernd Bohnet, Daniel Freeman, Gamaleldin Elsayed, Hanie Sedghi, Isabelle Simpson, James Harrison, Jiri Hron, Kathleen Kenealy, Kevin Swersky. Kshiteej Mahajan, Laura Culp, Max Bileschi, Merrie Morris, Rosanne Liu, Yundi Qian, Sharad Vikram, Tris Warkentin.
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.064500 2016.
Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacei, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, lurüi Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring. Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojevié, Wojciech Stokowiee, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http: //gi //github. com/ deepmieepnnnooo//
Blake Bordelon and Cengiz Pehlevan. Dynamics of finite width kernel and prediction fluctuations in field neural networks. arXiv preprint arXiv:2304.03444 2023.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github.com com/ google e/jax.
X. Chen, S. Xie, and K. He. An empirical study of training self-supervised vision transform- ers. In 202/ IEEE/CVF International Conference on Computer Vision (ICCV), pp. 9620-9629 Los Alamitos, CA, USA, oct 2021. IEEE Computer Society. doi: 10.1 10110099/CCV 00950. URL https: /oo. weseecooptterrscciee..crooo wwtp:/org11..0 w19p://ettoooo...1111/////// DTC.10/91TCVV4822... 2021.0095
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.022. 02311, 2022.
Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent neural networks typically occurs at the edge of stability. arXiv preprint 11Xiv:2103.00065, 2021.
Jeremy M Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E Dahl, al. Adaptive gradient methods at the edge of stability. arXiv preprint arXiv:2207.14488 2022.
Alex Damian, Eshaan Nichani, and Jason D Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. arXiv preprint arXiv:2209. 2022.
Aaron Defazio and Konstantin Mishchenko. Leaming-rate--rre learning by d-adaptation. arXiv preprint XrXiv:2301.07753 2023.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. arXiv preprint arXiv22302.05442, 2023.
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 O): 8-bit matrix multiplication for transformers scale. arXiv preprint arXiv:2088.07339, 2022.
10

---

Published as a conference paper at ICLR 2024
Emily Dinan, Sho Yaida, and Susan Zhang. Effective theory of transformers at initialization. arXiv preprint arXiv:2304.020::2 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold. Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 66 words: Transformers for image recognition at scale, In International Conference on Learning Representations (ICLR), 2021. https:
//arxiv. wttp.//ggaaaa...ooo./// 11929.
Colin Gaffney, Dinghua Li, Ruoxin Sang, Ayush Jain, and Haitang Hu. Orbax, 2023. URL http: //github. com/ gowglo/oolllooooo /orbax.
Justin Gilmer, Andrea Schioppa, and Jeremy Cohen. Intriguing properties of transformer training instabilities. To appear.
Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Car- doze, George Dahl, Zachary Nado, and Orhan Firat. A loss curvature perspective on training 2021.
instability in deep learning. arXiv preprint arXiv:2110.044:
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference artificial intelligence and statistics, pp. 249-256. JMLR Workshop and Conference Proceedings. 2010.
Google. Grain feeding jax models, 2023. URL htttt:///ggth//t com/ htte://oolleeeee //rainn
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2023. URL
http: http://github: com/ gwwggo/oogelleoooooooooooooo flax.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint
1.X.016066.8415,, 2016.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optinet large language models. arXiv preprint arXiv:2203. 15556, 2022.
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In Inter- national Conference Machine Learning, pp. 9099-9117. PMLR, 2022.
Maor Ivgi, Oliver Hinder, and Yair Carmon. Dog is sgd's best friend: A parameter-free dynamic step size schedule. arXiv preprint 2222.2220.2. 2023.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gener- alization in neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2018. https: www.wviwww.gaaaaaaaaaaaaeooeeeee ht/p///s//1166 100.6 07572.
Norman P Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al. In-datacenter performance analysis of tensor processing unit. In Proceedings of the 44th annual international symposium on computer architecture, pp. 1-12, 2017.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836/ 2020.
Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. arXiv preprint arXiv:1808.0660 2018.
Jachoon Lec. A random walk model of transformer parameter growth. 2023.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. arXiv preprint 1908::990..0:0. 2019.
11

---

Published as a conference paper at ICLR 2024
Peter J. Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. In International Confer- Learning Representations, 2018. URL https: seppereview.net/forrreeeeeeeeeee Hyg0vbWC
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. In Interna- tional Conference on Learning Representations (ICLR), 2016. https: https///axivvvvo a 1608.00 03983.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Con- ference Learning Representations (ICLR), 2019. https: http::///pppeeeeeeeeeeee new.eee.eeeeeeeeeeeeeeeeeeeeeeee forum? id=Bkg 6RiCQY77
William Merrill, Vivek Ramanujan, Yoav Goldberg, Roy Schwartz, and Noah Smith. Effects of parameter norm growth during transformer training: Inductive bias from gradient descent. arXiv preprint #rXhv:2010009697, 2020.
Igor Molybog, Peter Albert, Moya Chen, Zachary Vetoto, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, et al. A theory on adam instability in large- scale machine learning. arXiv preprint arXiv:2304.0983 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradhury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, high-performance deep learning library. In Advances in Neural Information Processing Systems (NeurIPS), 2019. https: //tp://aaviv////aa www.gi/iirrr.ooooooooo.ooo abs/ 1912. 01703.
Pytorch: An imperative style,
Ofrr Press and Lior Wolf. Using the output embedding to improve language models. In Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 2, Short Papers, pp. 157-163, Valencia, Spain, April 2017. Association for Computational Linguisties. URL httpt::t////mm /a//a/aaaanntcoooooooooooooo org/ /E17 2017772
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Lan- guage Models are Unsupervised Multitask Learners, 2019. https: https://openal: com/blog/
better-language madde-aeaaaggggeemdel///////s// models/.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Leurning Research, 2020a. http: // jmmr.. org papers/ v21/20-00 007/20000444 4. html.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to- text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020b. URL http: //jml r. www.:o////pp..//ooo.oeo
/20- 074, html.
Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Min- jia Zhang, Dong Li, and Yuxiong He. [ZeRO-Offloady Democratizing Sallion-Scalll model training. In 202/ USENIX Annual Technical Conference (USENIX ATC 2/), pp. 551-564, 2021.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596 4596-4604. PMLR, 2018
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. arXiv preprint arXiv: 2 100.009644 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Fukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- fion processing systems, 30, 2017.
Mitchell Wortsman, Tim Dettmers, Luke Zettlemoyer, Ari Morcos, Ali Farhadi, and Ludwig Schmidt. Stable and owwriiiiio training for large-scale rision-languee models. arXiv preprint arXiv:2304. 13013, 2023a.
12

---

Published as a conference paper at ICLR 2024
Mitchell Wortsman, Jachoon Lee, Justin Gilmer, and Simon Kornblith. Replacing softmax with relu in vision transformers. arXiv preprint arXiv:2309.08886, 2023b.
Sho Yaida. Meta-principled family of hyperparameter scaling strategies. MrXiv:2210.04909 2022.
arXiv preprint
Greg Yang and Edward Hu. Tensor programs iv: Feature learning in infinite-wide wiith neural networks. In International Conference on Machine Learning, pp. 11727-11737. PMLR, 2021.
Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ry- der, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466, 2022.
Shuangfei Zhai, Tatiana Likhomanenko, Etai Litwwi, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Josh Susskind. Stabilizing transformer training by preventing attention entropy collapse. arXiv preprint arXiv:2303.06296, 2023a.
Xiaohua Zhai, Basil Mustafa, Alexander Kolesnikov, and Lucas Beyer. Sigmoid loss for language image pre-training. arXiv preprint arXiv:2303. 15343, 2023b.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christo- pher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint 2r5.vv2205.01068, 2022.
13

---

Published as a conference paper at ICLR 2024
A ADDITIONAL INFRASTRUCTURE DETAILS
This Section provides more details on the training infrastructure, which is built on Flax (Heek et al., 2023), Jax (Bradbury et al., 2018), and TPUs (Jouppi et al., 2017). To enable larger model training, we shard the model and optimizer states as in FSDP (Ren et al., 2021), then specify these shadings when compiling with JIT. We use Orbax (Gaffney at., 2023) for checkpointing, and Grain (Google, 2023) for deterministic data loading. When loading data, sequences are packed so that no padding required-if sequence is less tokens than the context length hyperparameter, then an end of sequence token is appended, followed by the beginning of a new sequence.
B WHEN IS LEARNING RATE SENSITIVITY A USEFUL METRIC
There are cases where LR sensitivity (defined in Section 2.2) is no longer a useful metric. This section details these scenarios and justifies the use of LR sensitivity for the interventions in this paper.
Interventions which change the meaning of learning rate. When an intervention changes the meaning of learning rate then comparing LR sensitivity not useful. A clear example of this would be taking the square root of the LR before passing it to the optimizer, but there are more subtle cases to be cautious of when using LR sensitivity.
In general, we avoid manipulations where the meaning of LRR meaningfully changes. In some cases, we have good empirical evidence that the meaning of the learning rate has not changed when inter- vening. For instance, the LR vs. loss curves are indistinguishable up to some critical learning rate when using gk-layernone (Figure 1), adding z-loss (Figure 2), or changing warm-up.
In other cases, such as when testing yPaarm (Section 3. 3.2.4), we believe that LR sensitivity is useful despite a per-layer modification of LR. This because the per-layer LR is manipulated linearly, and this modification does not change for different points on the LR vs loss curve.
Shifting of the optimal LR. The definition of LR sensitivity in Section 2.2 does not account for the optimal LR shifting when specifying the LR range (a, b]. In practice we recommend shifting the three order of magnitude range [a, bj to correspond with this shift. For instance, we shift the range main experiments (e.g.,
in Section 3.2.2, as discussed in more detail in the section. However, Figure 1) do not test at a large enough scale to necessitate this shift.
LR sensitivity is invariant invariant to the scale of the loss. If the network consistently achieves random performance across learning rates, then LR sensitivity will be zero. We do not offer solution to this, and instead recommend that LR sensitivity should always be examined in combination with the LR vs. loss curves as we do this paper. meant as a useful summary of the LR vs. loss curves, not metric to optimize in isolation.
loss. Another limitation of the LR sensitivity metric is that it is
C PARAMETER AND OUTPUT NORM GROWTH
This section discusses the growth of the parameter norm during Transformer training as previously studied by Merrill et al. (2020); Lee (2023), and relates this phenomenon to the attention logit growth and Adam epsilon instabilities (Sections 1.1.1 and 3.4, respectively).
An observation of Lee (2023) is that, when using an adaptive optimizer, the movement of parameters can be approximated by random walk. We show parameter root mean square (RMS) throughout training in Figure C.1², which appears to follow predictable trend. This aligns with the afore- mentioned observation, and is further supported by Figure C.2. Figure C.2 displays parameter RMS function model scale and learning rate, averaged over the last 500 training steps. As before, parameter RMS is determined primarily by learning rate. As parameter RMS grows, we would ex- pect output RMS to also grow. This is validated by Figure C.3, which shows that the RMS of the Transformer block output is mainly determined by learning rate, and follows a very similar trend to parameter RMS.
We show parameter RMS for the first MI. layer in various blocks, hut expect other layers to exhibit similar trends.
14

---

Published as a conference paper at ICLR 2024
There are two interesting takeaways. First, this observation helps explain why the attention output logits become large at high learning rates as observed by Dehghani et al. (2023) and Section 3.1.1. This is the only feature in the network test whose magnitude depends quadratically on parameter RMS. For inputs X with unit RMS, a typical matrix multiply XW with parameters W will result in features Y where RMS(Y) is linear function of RMS(().. On the other hand, the attention logit entries are computed via {.X x so depend quadratically on RMS(W). They are therefore the first to become large when the parameter norm grows, Next, this helps explain the decreasing trend in gradient scale observed in Section 3.4 (Figure 7). In pre- rre-normalization (Radford et al., 2019) Transformer (Vaswani et al., 2017) there an output layemorm layer (Ba et al., 2016) after the last Transformer block and before the final linear layer. The gradient from this output layernorm layer is scaled by the reciprocal of the input RMS. In addition to growing with LR, this RMS growing with depth because of the residual connections (Figure C.3). As the RMS leaving the last Transformer block grows, the gradient received shrinks.
For completeness now compute the layernorm gradient to input We assume the input mean zero and the layernorm has no bias for simplicity. Let
x
LayerNorm(x) = a
a.
(1)
v
E.
[x
+
m
where m = E, [ +
Then
0e Σ
Oe Ozk
0x,
3zk dzg
08
a
Σ de
1
2
02j
m 1/2
Ozy
2
33/2
n
2zj
(3)
1/2
1
(Kj Ozj
de
1/2 1/m-/2
nmm
Σ
Ozk
Ol
(4)
k
OE
jxj Ozg
(V,, a 2))
(5)
1/2
Tum
Equivalently,
1
22 x)
Vz a
V
nmm 1/2 0
(6)
m//2
D ADDITIONAL MODEL AND OPTIMIZER INTERVENTIONS
This section recreates the plots from Section 3.2 with additional interventions or hyperparameter changes.
Changing the number of training steps from le5 to 5e4 or 2e5 does not meaningfully change LR sensitivity (Figure 1.14).
We try applying gk-layernorm across the whole model dimension instead of individually per-head with shared paramters. As illustrated in Figure G.13, the latter performs better. We use per-head qk-layernorm as the default all other experiments.
Increasing the batch size from 256 512 1024 does not meaningfully change LR sensi- tivity (Figure G.15, each batch element contains 512 tokens). When increasing batch size we decrease the number of training steps so that the amount of data seen is constant. We believe a similar effect would be observed instead we held the number of steps constant because changing the number of steps has no impact on LR sensitivity at batch size 256 (Figure G.14).
The effect of changing the weight decay from le-4 is illustrated in Figure G.16. Increasing decay appears to slightly shift the optimal LR right.
We find that the logit growth instability is not due to the softmax in the self-attention layer, as it still occurs with pointwise variant of attention (Figure G.17).
15

---

Published as a conference paper at ICLR 2024
Block
Block
Block 11
Block =
10a
10
10
109
Block output RMS
107
10
10
107
207
10
10
107
201
10
10
10
100
10
10
10
10
10
10
10
10
25000 50000 750001 D0000
10
25000 50000 75000100000
10
25000
75000100600
10
25000 50000 55000000000
Step
Step
Step
Step
First MLP weight RMS
20"
10
101
10"
10-4
10-1
10-4
10-1
10
10
10
10"
75000
75000100000
Step
Step
Step
Step
LR 0.0001
LR 0.001
- 0.01
LR 0.1
Figure C.I: (Top) The root square (RMS) of the Transformer block outputs throughout train- ing. (Bottom) The RMS the MLP weights throughout training, for the first of the two layers the MLP. Recall RMS(X)= = [x,] RMS is mostly determined by LR, with higher LR correspond- ing with higher RMS for block outputs and MLP weights. Different curves the same learning rate correspond to different model scales. This experiment uses a decoupled weight decay values of le-4.
10
10
101
MLP weight RMS
10a
10
10".
10-1
10-1
10-2
10-2
107
10a
10a
10
3
2
5
Num params
LR
Block index
0.0003
0.3
9.4e6
1.2c9
9.4e6
1.2e9 Num params (at LR 3e-1)
Num params
Figure C.2: The root mean square (RMS) the MLP weights are roughly consistent with scale (left) but increase reliably with learning rate (center). At high learning rates, parameters later in the network can be affected by AdamW epsilon instability discussed Section 3.4 (right), This experiment considers the first of the two MLP layers in the block, and data for the first two plots are from block two. RMS is averaged over the final 500 training steps, where RMS(X) [X7]
16

---

Published as a conference paper at ICLR 2024
Block output RMS
10
10a
10a
10a
10
10
10a
10a
10a
10a
10a
107
10a
10a
10
10
10
10
15
20
25
Num params
LR
Block index
0.0003
0.3
9.4e6
1.2e9
9.4e6
1.2e9
LR
Num params
Num params (at LR 3e-1)
Figure C.3: The root mean square (RMS) of the Transformer block outputs are roughly consistent with scale (left) but increase with learning rate (center). RMS increases deeper in the transformer because of the residual connections, which is shown for very high learning rates (right). The first two plots are for block index two, and RMS is averaged over the final 500 training steps. Recall RMS(X) = v [X7
ADDITIONAL EXPERIMENTS ON PREDICTING THE ATTENTION LOGIT GROWTH INSTABILITY
One question unresolved by our analysis in Sec- tion 3.3 whether could have predicted that in- stability arises when the max attention logit exceeds le4 without manipulating learning rate and model size. We take initial steps towards an answer by transplanting different values of max attention logit into small network with 10M parameters. For different constants K we pass the queries and keys
Fixed scale
Zero layer baseline
6.0
Final loss
5.0
through (()) V 2/V [ before computing the attention logits. Results are illustrated in Fig- ure E.I. Loss deteriorates around к =le3, and by lee the loss exceeds that of a zero-layer bigram model consisting of the Transformer we use without any self-attention or MLP layers.
4.0
10a
10
102
10a
10"
Scaling constant
Figure E.1: Enforcing a max attention logit of approximately in small model to de- termine which value of x inhibits learning.
F AUTHOR CONTRIBUTIONS
Mitchell Wortsman led the project, ran the experi- ments and produced the figures, contributed substan- tially the infrastructure for experimentation, the framing and direction, and the writing.
Peter Liu led the infrastructure and creation of NanoDO for experimentation, provided key insights and advice on multiple technical areas, and contributed the writing.
Lechao Xiao and Katie Everett contributed to the infrastructure used for experimentation, provided key insight related to parameterization, and contributed to the writing.
Alex Alemi, Ben Adlam, John D. Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl-dickstein, and Kelvin Xu were active participants in weekly brainstorming meetings which motivated, influenced, and elucidated technical concepts pertaining to this work.
17

---

Published as a conference paper at ICLR 2024
Jachoon Lee and Justin Gilmer were senior authors advising the project, contributed substan- tially to the framing and direction, provided key insight and advice on multiple technical areas, and contributed to the writing. Jachoon led the connection with output norm growth. Justin proposed to plot loss as a function of learning rate for different model sizes, and performed initial experiments demonstrating that attention logit growth could be reproduced at high learning rates in small models.
Simon Kornblith was the lead advisor on the project, contributing substantially to the framing, di- rection, infrastructure, and writing. Simon initially brainstormed the project with Mitchell, and was Mitchell's host for the summer internship during which this research was conducted, providing substantial technical support.
G ADDITIONAL FIGURES
This Section contains the additional Figures referenced in the main text and appendix.
Num 9.4c6, LR
LO55
4.4
Output logit mean
-20
4.2
Num perams = 4.0e9, LR 0.01
-40
attenbon logit
-60
LooS
-100
50000 Step
100000
50000
100000
Step
2-loss coefficient le-4
Weight decay 1e-4
2000
4000
6000
1000
10000
2000
2000
6000 Stop
1000 15600
No z-loss
No weight decay
Stop
ok lyyaammmoommmmmmmm
da-layemorm
Figure G.2: An example the output logit vergence instability (Chowdhery et al., 2022) (Section 3.1.2) in a 2.4M parameter Trans- former at learning rate 0.1.
Figure G.I: The attention logit growth instabil- ity (Dehghani et al., 2023; Zhai et al., 2023a) ap- pears in small models at high learning rates. The mitigation of applying gk-layernorm proposed by Dehghan et al. (2023) is equally effective in the small-scale regime. The max attention logit reported for layer 0, which we typically ob- serve to have the largest logit values.
18

---

Published as a conference paper at ICLR 2024
Standard (simple]
2.4e+ 06
Standard MuParam (simple)
2.40
4.2
N 9.4e+
N 199-07
N 1.9e+
N 4. 4.2e+07
0.5e+07
4.00
1.50008
N 1.5e+000
3.0e+08
3.75
1,20409
3.75
N 1.2e 09
Rial eva loss
3.50
2.25
3.25
3,00
3.00
2.75
2.75
2.50
10-4
10
10a4
10:4
2:50
10
10"
Learning rate
Learning rate
sensitivity
Standard
LR sensitivity
Standard
MuParam (simple)
1.0
(slm pla)
3-21-
2.
107
10
10'
10' Number parameters
10a
parameters
Figure G.3: Measuring the effect of μPPara on LR sensitivity for models with qk- layernorm (Dehghani et al., 2023). μParam succeeds stabilizing the optimal L.R, though it does not improve loss or reduce LR sensitivity. For more information refer to Section 3.2.4.
Figure G.4: The effect of pParam on LR sen- sitivity for models without qk-layemorr (De- hghani et al., 2023). paParam succeeds in stabi- lizing the optimal LR, but does not alleviate the need for qk-layernorm. For more information refer to Section 3.2.4.
3.6
10 8
3.4
10-9
Grad RMS
Loss
3.2
10 100
3.0
10 11
10
2.8
10 -1
2.6
0
50000
100000
0
10
20
Step
Block index
Eps 1e-8
Eps 1e-15
Figure G.5: Decreasing AdamW e from default value of 1e-8 to le-15 improves loss for 4.8B parameter model at LR 0.3. When increasing to 1c-6, loss diverged. Grad RMS is averaged over the final 500 steps for the first layer in the MLP: refer to Figure G.6 for data throughout training.
19

---

Published as a conference paper at ICLR 2024
Block 0
Block
Block 16
Block 22
10a1
10"
10-1
10"
Grad RMS
10
10
10-7
10-1
10
10-9
10-3
10-9
10
10 10
10:
20
10"44
10"14
10:11
25000 yo000
730000 00000
25000 s0000 750001
25000 50000 73000100000
25000 50000 75001100000
Step
Step
Step
Step
Unscaled update RMS
10-1
10-1
10-1
10-1
10-2
10-2
10-2
10-1
10:4
10'
10"4
10'
50000 0000000000
750001
75000108000
50000
Step
Step
Step
Step
Eps 1e-6 (diverged) ed)
Eps = 1e-8 (default)
Eps 1e-15
Figure G.6: The top row displays the root mean square (RMS) of the gradient for the first MLP layer at different blocks throughout the network. When the grad RMS drops below the AdamW hyperparameter, the magnitude of the update decreases, as illustrated by the bottom row. Experiment conducted with 4.8B parameter model trained with LR 0.3. The experiment with 1e-6 was stopped when loss diverged.
Num params 9.4e6, LR 0.1, block 0
Max attention logit
10a
Mean query nom
10*
Loss
10a
10a
10a
10'
102
2000
4000
6000 Step
8000
10000
2000
4000
6000
8000
10000
2000
4000
6000 Step
8000
10000
Step
Query key cossim mean
1.0
1.0
1.0
0.8
0.8
attention
0.8
0.6
0.6
0.6
0.4
0.2
0.2
2000
4000
6000
8000
10000
2000
4000
6000 Step
8000
10000
2000
4000
6000 Step
8000
10000
Step
ok-layemorm True
k- laylayymm = False
Figure G.7: The logit growth instability (Dehghani et al., 2023; Zhai et al., 2023a) occurs when the norm of the query and keys increases, not due to an increase in their cosine similarity.
20

---

Published as a conference paper at ICLR 2024
Joint scaling
Scaling depth
Scaling width
Final evall loss
34x10'
32x
4.2
20x
Scaling width
dim 256
layers
Scaling depth
dim 512
layers 6
4.0
dim 760
layers 12
dim
107
10a
dim 2048
layers 48
Num params
layers
Figure G.9: Jointly scaling width and depth leads to lower loss than independently scaling depth or width at the largest scale we test. also leads to more reliable scaling prediction when extrapolating from models with less than 1e8 parameters. Best loss reported in sweep over learning rates.
Rinal eval loss
3.6
3.4
2.0
10
10
10
10
10"
10--
Learning
10-3
Figure G.8: The effect of scaling width vs. scaling depth without qk-layernorm (Dehghani et al., 2023).
Grad
10-1
2000
4000
2000
4000
Step
Step
Eps 1e-6
Figure G.10: Increasing the AdamW from its default value of 1e-8 to 1e-6 causes loss di- vergence for 4.8B parameter model at LR 0.3. Grad RMS is for the first layer in the MLP.
21

---

Published as a conference paper at ICLR 2024
gk-layernorm
qk-layemorm
False
4.4
4.4
Standard
N 2.ee006
Standard
2.4
MuParam (simple)
N 9.4
MuParam (simple)
9.44+06
4.2
MUParam (Intermedlate)
+ 1.9e
4.2
MuParam (Intermedlate)
1.9e407
MuParam
4.2e+07
(full)
1.2e+07
4.0
N 8.
4.0
8.5e+07
N 1.5e+08
1.5e+08
N 3.0e+08
N 3.0e+08
3.8
3.8
Final eval loss
Final eval loss
3.6
3.6
3.4
3.4
3.2
3.2
3.0
3.0
2.8
2.8
10
10-1
10
10-1
10-1
10-1
10
10"
Learning rate
Learning rate
LR sens/tivity
LR sensitivity
10-1
10-1
10-7
10-1
107
10a
107
Number of parameters
Number of parameters
Figure G. 11: Comparing μPaaam (full), which implements jPaaram as described in Yang et (2022) with and without gk-layernorm, with μPaarm (simple) and μPaarm (intermediate). There are four changes in pParam (full), (i) Scale the LR for linear layers by base-fan-in/ fai--aaann-ff-in,, ffan-in, (ii) initialize the head with standard deviation / maase--an-in/fa-iin (iii) change the 1/v di scaling factor in attention layers to 1/d, where d, is the head dimension, and (iv) initialize the query projection weights with zeros. jPParm (intermediate) consists of (i) and (ii), while pParam (simple) is only (i). With pParam (full) and gk-layernorm, the model trains without diverging LR 1. However at the best LR there is no measurable improvement over jParam (simple) at the largest scale we test.
22

---

Published as a conference paper at ICLR 2024
gk-layernorm True
qk-layernorm False
4.4
1/sq/thhead dim) scaling
N 2.4e+06
1/sqrt(head dim) scaling
N 244+06
1/thead dim) scaling
+ N = 9.4e+06
1/(head dim) scaling
N 9.4e+06
N 1.9e+07
N 1.9e407
4.2e+07
N 4.2e+07
8.5e+07
8.5e+07
N 08
1.5e+08
N 3,0e+08
N 3.0e+08
Final eval loss
Final loss
3.0
10
Learning rate
Learning rate
Figure G.12: Measuring the effect of changing the 1/vd, term attention to 1/dh, where dy, is head dimension, Vaswani al. (2017) use 1/ id while Yang al. (2022) use 1dz.-
Steps 50000
N 2
dim layernommm
2.4e+06
4.25
4.2
Model-dim ek layemorm
9,4e106
Steps 200000
N 1 9e+
1.9+007
4.2++
N
42e+07
4,00
0.Se+ 07
8,5e+07
3.0e
1.e+08
3.8
1,0c+08
Final lvaal loss
3.75
Final eval loss
3.6
3.50
3.4
3.00
3.0
2.75
2.0
10a1
10a
10-2
10
10
Learning rate
10a3
10
sensitivity
Learning rate
ax30
Figure G.13: We achieve slightly better perfor- mance when applying qk-layermorm individu- ally per-head instead of across the model dimen- sion. The per-head variant has only head-dim learnable parameters instead of model-dim pa- rameters. We use the hea- variant as the de- fault this paper, and we never use biases.
20"
Number parameters
Figure G.14: Changing the number of total training steps from le5 to 5e4 2e5 does have large effect of the shape of the learning rate vs. loss curves the scales test.
23

---

Published as a conference paper at ICLR 2024
qk-layemorm True
gk-layemorr False
Batch 256
2.4e+06
Batch size 256
2.4e+06
4.25
9.4e+06
4.25
Batch size 1024
1.9c+07
Batch size 1024
1.9c+07
+
4.2e+D7
N 4.2e+07
N
8.
4.00
N 8.5e+07
1.5e+08
1.5e+08
3.75
Final eval loss
3.50
3.50
Final
3.00
3.00
2.75
2.75
10
10
10
10
10
10
10a
Learning rate
Leaming rate
LR sensitivity
10a
LR sensitivity
10a
10-1
10-1
107
10a
107
109
Number parameters
Number parameters
Figure G.15: Increasing the batch size from 256 to 512 1024 does not have large effect on the shape of the learning rate vs. loss curves at the scales we test. Each batch element contains 512 tokens, and 256 the default.
24

---

Published as a conference paper at ICLR 2024
gk-layemom True
qk-layeroom False
Weight 3e 04
2,4
Weight decay 3e 04
N = 2.4++06
4.25
Weight decay - 1e
0.4
4.25
Weight decay 1e 04
N 9.4c+06
Weight decay 3e-05
+
1.9e+07
Weight decay 3e-05
1.9e+07
Weight decay
4.2e+07
Weight decay
N 0.5e+07
N 8.5e+07
N 1.
N = 1.5e+08
* 3.0c+08
N 3.0e+08
3.75
Final eval eval loss
3.75
3.50
3.50
Final
3.25
3.25
3.00
3,00
10
10
10
10
10a
10-4
10-2
10
16
101
Learning rate
Learning rate
LR sensitivity
10a
LR sensitivity
10a
10-1
10-1
10-2
10-
10"
10a
10'
10a
Number of parameters
Number of parameters
Figure G.16: The effect of weight decay on LR sensitivity. We use independent weight decay as described in Section 3.2.2 and recommended by (Loshchilov & Hutter, 2019).
gk-layemorm = True
qk-layernorm False
attention
2.e+066
Softmax attention
2.4e+06
4.2
Pointwise attention
9.4+066
4.2
Pointwise attention
9.4e+06
1.9e+07
N 1.9e+07
+
4.2e+07
N 4.2e+07
8.5e+07
8,5e+07
15e+08
3.8
3.0c+08
3.6
3.0c+08
Final eval loss
Final eval loss
3.6
3.2
3.2
2.8
10
10
10
10-1 10-1
10
10
10-1
Learning rate
Learning rate
Figure G.17: The logit growth instability occurs even without softmax. For the pointwise variant of attention here, we replace softmax with squared-relu as described by (Hua et al., 2022). As recommended in (Wortsman et al., 2023b) we add a scaling factor which depends sequence length. In this case, we use inverse square root.
25

---

Published as a conference paper at ICLR 2024
10
10-4
10-7
Grad RMS
10-5
10-4
10-4
10-7
10-4
10-10
10
107
10a
10a
10
10
10
10
15
20
Num params
LR
Block index
9.4e6
9.4e6
LR
Num params
Num params (at LR 3e-1)
Figure G.18: Recreating Figure 7 with the kernel projection instead of the first MLP layer.
LR 0.0003
LR 0.001
0.01
=R 0.3
Loss
30
3x
25000 $8000
75000 300000
25000 50000
75000 100000
25000 50000
75000 180000
25000 50000
75000 106000
Step
Step
Step
Step
Grad RMS
1.0
10
10
10
10
10
10-4
10-4
10-4
25000
50000 Step
75000 300000
25000 50000
75000 100000
25000 50000
75000 100000
25,00 5,000
75000 100000
Step
Step
Step
Unscaled update RMS
10-1
10-3
10-2
10
50000 Step
75800 100000
25000
100000
50h00
Step
Step
Step
- 9.4e+006
N 1.9
4,26107
8.5
1.5e
3.0e108
1.2e
Figure G.19; For various learning rates and model sizes we display the gradient root mean square (RMS), and the unscaled update RMS. The unscaled udpate is the update returned by the optimizer before scaling by learning rate. The gradient and update are shown here for the first MLP layer of the Transformer. The update RMS falls when the grad RMS approaches the AdamW of 1c-8.
26


Reviews:
[{"review_id": "DkWeWfZsih", "reviewer": "Reviewer_TVMA", "summary": "In this work, the authors examine sources of training instabilities in transformer models through a detailed experimental study.\nThey motivate their study with the fact that instabilities observed in large transformer models are difficult to study and mitigate because of the large computational costs of these runs.\nThey therefore examine these and show that they can be reproduced in smaller models, which can be trained faster and can be used to design mitigations for the instabilities which will hopefully translate to larger architectures.\nIn particular, they focus on two instabilities observed in practice, namely the the growth of logits in attention layers, and the divergence of output logits.\nThey show that increasing the learning rate at training time can reproduce these instabilities for smaller models.\nFurther, they show that commonly used mitigation approaches, such as qk-layernorm and z-regularisation can help with instabilities induced by large learning rates, and also examine the effect of a range of other optimiser and model interventions to the sensitivity of the training procedure on the learning rate.\nA range of ablation studies across learning rates, interventions and model size yield a number of practical insights on training stability.", "questions": "- __Figure 1:__\nThe caption says \"LR sensitivity measures the expected deviation from optimal.\"\nWhat do the authors mean by \"optimal\" in this context?\nIs the meaning of \"optimal\" coming from the discussion in section 2.2?\nSome clarification on this in the main text would be good.\n\n- __Introduction comment:__\n\"One interesting finding is that scaling depth increases LR sensitivity at a faster rate than scaling width.\"\nOne factor at play with this finding may be the fact that in standard initialisation schemes, changing the width of the network affects the initialisation scale of the weights, whereas increasing the depth does not.\nAs a result, it is reasonable to expect that changing the width does not impact stability as much as depth, because the change in width is somewhat accounted for by the adaptive initialisation.\nCan the authors comment on why this occurs?\n\n- __Point on phrasing:__\nIn section 3.3 the authors write \"We now examine whether it is possible to predict the logit growth instability before it occurs.\"\nI think this phrasing is a little ambiguous because it may be interpreted as predicting whether a logit growth instability will occur in an ongoing run, based on the data collected in the current run.\nBy contrast, to my understanding, the authors are using previous runs with different hyperparameters, to determine whether a particular hyperparameter setting will cause an instability or not.\nI think stating this more clearly in the main text would be beneficial.\n\n- __Effect of different optimisers:__\nTo my understanding, all experiments in this paper use AdamW.\nCan the authors comment on whether they expect their findings to extend to other commonly used optimisers?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall, I think the experimental work in this paper was well executed and carefully controlled.\nThe strengths of the paper, in my view, include a number of useful findings and insights, as well as the overall high quality of the ablations and the paper itself:\n\n__Reproducing instabilities on small models:__\nThe authors successfully reproduce training instabilities on smaller transformers, by increasing the learning rate.\nThey show that as model size increases (e.g. figure 1 and figure 6), training instabilities occur at smaller learning rates.\nFurthermore, the authors show that two existing instabilities that are observed in large transformers (i.e. the growth of logits in attention layers and the divergence of output logits) can be reproduced in smaller models.\nThis is convincing evidence that the authors' findings on interventions made on smaller models are likely to translate to larger ones, since the mechanism of the instabilities is common across different scales.\nIn addition, pointing out this relationship is interesting and also potentially useful towards the development of large transformer models, as it provides strong evidence for adjusting the learning rate as a function of model size.\n\n\n__Verifying the effectiveness of qk-layernorm and z-regularisation:__\nThe authors showed that using qk-layernorm (figure 1) and/or z-regularisation (figure 2) significantly helps mitigate instabilities, reducing sensitivity to the learning rate across a range of model sizes, and increases the range of stable learning rates.\nThis suggests that qk-layernorm and z-regularisation are good candidates for mitigating instabilities in small models, and likely also sufficient for mitigating these effects in large transformers as well.\n\n\n__Extrapolating instabilities:__\nThe authors demonstrate that the hyperparameter regimes which result in instabilities can be predicted by looking at the maximum attention logits from other runs.\nIn particular, in figure 6, they show that for a model with no qk-layernorm, both the value of the maximum attention logit as well as the occurrence of an instability can be predicted by extrapolating from smaller runs and different learning rates.\n\n__Overall thoroughness of ablations:__\nI found that the ablations performed in this work were very thorough and supported the claims made in the main text very well.\nThe documentation of the various parameter settings used in the experiments are also clearly documented.\n\n__Motivation and clarity:__\nOverall, I also found the paper to be well motivated and clear, and the figures to be insightful and informative.", "weaknesses": "I did not find significant flaws in the paper, I thought that two possible weakness are the following:\n\n__Absence of concrete rules of thumb:__\nOne weaker point in the paper is that it does not provide concrete rules of thumb for setting the relevant hyperparameters of transformer models and their training loops.\nSpecifically, I think that the paper goes a long way reproducing instabilities and performing detailed ablations, but does not provide concrete advice (i.e. general recipes) for hyperparameter settings.\nGiven the thoroughness of the ablations, this is a relatively minor point.\nHowever, I think that a short discussion of how a practitioner could use the insights in this paper to fix training instabilities and extract better model performance (by utilising smaller scale runs), would be useful.\n\n__Limitation to C4 data:__\nTo my understanding, all experiments in this work involve the C4 dataset, which is textual.\nWhile it is most likely that the authors' findings generalise to other datasets, it is not fully clear that the scalings shown in this paper would be encountered in other data modalities.\nHowever, I appreciate that performing experiments on additional data modalities would be a large overhead in effort, and the current findings to be convincing enough.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cUz4owUsKP", "reviewer": "Reviewer_SBqh", "summary": "This article studies optimization instabilities observed in the training of large transformer-based models. The central contribution is the reproduction and analysis as small scale of instabilities that were previously observed on large-scale models. This allows to study the stability of those models without needing the large computing power required for large-scale training.\n\nTwo central kind of instabilities are studied by the authors: the growth of logic in attention layers, and the divergence of output logits of the model. In both cases, it is experimentally shown that those instability can be reproduced on small models when using a large learning rate, and that the mitigation techniques that were developed for large models are equally effective in this context. The core tool used for this analysis is introduced to be the measure of the sensibility of the model performance to the learning rate used for the optimization, and the experimental results show that those mitigations tend to reduce that sensibility, stabilizing the training.\n\nThe authors finally extend their analysis to study the impact of several other interventions that have been proposed, such as the parameterization of the trainable weights, the integration of weight decay in the optimizer, the scaling of the model size and the use of warm-up periods.", "questions": "I don't have more questions.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This is an extensive and detailed experimental study of the stability of transformer models with regard to the training learning rate and the various mitigation methods that have been considered.\n\nThe experimental setup is described with abundance of details, the conducted experiments are well motivated and presented, and the analysis tools (as the LR sensibility) allows a synthetic and clear summary of the impact of the parameters & methods evaluated.\n\nI believe this article has the potential to provide a wealth of useful information and heuristics for practitioners working with such models.", "weaknesses": "While I am not extremely familiar with the large-transformer-models community, I am under the impression that the pool of persons effectively concerned by this work is very small. As the authors note, training such large models is very computationally expensive, and currently only very few groups have the means to train such models.\n\nAs a result, I wonder if this subject might be in practice rather niche, in terms of how much of the community could actually use it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "O7IcOYLnKy", "reviewer": "Reviewer_Fq7A", "summary": "An experimental paper. The authors' main point seems to be about the attention and output logits in transformer yielding instabilities. This may be a valid point, although it is not very how to mitigate this problem, and it is hard to be completely convinced that this is *the* reason for instability of large transformers. That being said, some of the experiments are valuable and help us a little bit to understand some issues that may arise in the training of transformers. There is an emphasis on considering the learning rate size. \n\nThe suggested experimental evidence supporting this claim is \"val loss vs learning rate\" curves. However, (1) there is no surprise in training divergence when lr becomes too large, and (2) I do not see any experimental evidence that divergence is indeed caused by the considered instabilities and not by something else.\n\nThe paper also studies how \"learning rate sensitivity\" is affected by certain design choices. Learning rate sensitivity is defined as the average of \"excess val loss\" over learning rate range. However, the choice of particularly this metric does not seem well-motivated. Do authors use uniform distribution over lr? If yes, why not uniform over log(lr)? Why not simply use maximal stable lr?\n\nOne insight which seems useful is that default eps=1e-8 in AdamW might appear too large and cause updates to vanish.", "questions": "Are we really sure that the reason for large transformers not training well is logit divergence? What are other possible problems? What do we learn in the end from your analysis? Is it clear that such problems don't arise in other architectures?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Understanding stability of transformer training is an important problem. The hypothesis that instabilities may be related to attention logits is not without interest. The numerical experiments seem to be very carefully made, and overall they bring some value. I thank the authors for the clarifications.", "weaknesses": "The structure of the paper is a little weird (the conclusion is very short and contains no useful information, the discussion of existing results is just put at the end without much being done from it, the main points seem to be made in the figures. ). The way the logits in the attention mechanism pose problem is not made super clear or intuitive (obviously, it's a little hard to prove something, but at least some intuition would be appreciated). For instance, we learn that high enough learning rate will pose problem at some point, but that's the kind of things that is not surprising. Does this validate the whole hypothesis?\nNote: the concerns have been addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FalYCeXy6N", "reviewer": "Reviewer_hm5n", "summary": "The paper studies instabilities of transformers on a smaller scale. Specifically, the authors performs ablation experiments over learning rates of small transformers, and finds that techniques that are known to improve stability for large transformers also improve the stability of small transformers when using high learning rate. Among other things, the authors show 1) that qk normalization enables higher LR, 2) that the z-loss enables higher LR 3) LR warmup makes the model less LR sensitive, 4) Independently parametrizing WD and LR makes the model less LR sensitive, 5) model LR sensitivity grows faster with depth than width.", "questions": "Should LR sensitivity be normalized somehow? The optimal loss scales with model size, so the delta in eval loss between models of different scales are not really comparable.\n\nWill the code be open sourced?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Large scale transformers are expensive, important and suffer from instabilities. Providing a small-scale proxy model is impactful.\n\nThe paper is well written and the experiments are cleanly described. \n\nThe observations on independent weight decay and the scaling of the gradient RMS are relatively novel.", "weaknesses": "A significant part of the paper is dedicated to replicating observations made in large transformers to small transformers. The utility of this is a little unclear. While it demonstrates that a small model with high LR could serve as a proxy for a larger model, it doesn’t demonstrate any new insights regarding large models. It would be more impactful if the authors would make previously unknown observations at a small scale, and then show that they hold at a larger scale.\n\nSection 3.3 reads a little anecdotal to me. A more systematic study would be better.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (oral)