Summary
In this work, the authors examine sources of training instabilities in transformer models through a detailed experimental study.
They 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.
They 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.
In particular, they focus on two instabilities observed in practice, namely the the growth of logits in attention layers, and the divergence of output logits.
They show that increasing the learning rate at training time can reproduce these instabilities for smaller models.
Further, 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.
A range of ablation studies across learning rates, interventions and model size yield a number of practical insights on training stability.
Strengths
Overall, I think the experimental work in this paper was well executed and carefully controlled.
The 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:
__Reproducing instabilities on small models:__
The authors successfully reproduce training instabilities on smaller transformers, by increasing the learning rate.
They show that as model size increases (e.g. figure 1 and figure 6), training instabilities occur at smaller learning rates.
Furthermore, 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.
This 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.
In 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.
__Verifying the effectiveness of qk-layernorm and z-regularisation:__
The 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.
This 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.
__Extrapolating instabilities:__
The authors demonstrate that the hyperparameter regimes which result in instabilities can be predicted by looking at the maximum attention logits from other runs.
In 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.
__Overall thoroughness of ablations:__
I found that the ablations performed in this work were very thorough and supported the claims made in the main text very well.
The documentation of the various parameter settings used in the experiments are also clearly documented.
__Motivation and clarity:__
Overall, 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:
__Absence of concrete rules of thumb:__
One 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.
Specifically, 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.
Given the thoroughness of the ablations, this is a relatively minor point.
However, 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.
__Limitation to C4 data:__
To my understanding, all experiments in this work involve the C4 dataset, which is textual.
While 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.
However, I appreciate that performing experiments on additional data modalities would be a large overhead in effort, and the current findings to be convincing enough.
Questions
- __Figure 1:__
The caption says "LR sensitivity measures the expected deviation from optimal."
What do the authors mean by "optimal" in this context?
Is the meaning of "optimal" coming from the discussion in section 2.2?
Some clarification on this in the main text would be good.
- __Introduction comment:__
"One interesting finding is that scaling depth increases LR sensitivity at a faster rate than scaling width."
One 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.
As 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.
Can the authors comment on why this occurs?
- __Point on phrasing:__
In section 3.3 the authors write "We now examine whether it is possible to predict the logit growth instability before it occurs."
I 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.
By 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.
I think stating this more clearly in the main text would be beneficial.
- __Effect of different optimisers:__
To my understanding, all experiments in this paper use AdamW.
Can the authors comment on whether they expect their findings to extend to other commonly used optimisers?
Rating
8: accept, good paper
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.