Paper: Small-scale proxies for large-scale Transformer training instabilities (ICLR-2024)

Reviewer comment (Experiments):
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, 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, 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.

Author response:
We absolutely agree with these limitations. We have revised Section 5 to include a subsection which highlights both of these limitations.

As you highlight in your review, “the experimental work in this paper was well executed and carefully controlled”. Unfortunately, we believe that in order to provide concrete rules of thumb or study the effect of data, a large number of further experiments would be required for us to be confident in the findings and maintain the overall standard of careful experimentation. Therefore, we believe that these directions are promising topics for further research, but beyond the scope of our investigation.

Reviewer comment (Experiments):
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, performing experiments on additional data modalities would be a large overhead in effort, and the current findings are convincing enough.

Author response:
We absolutely agree with these limitations. We have revised Section 5 to include a subsection which highlights both of these limitations.

As you highlight in your review, “the experimental work in this paper was well executed and carefully controlled”. Unfortunately, we believe that in order to provide concrete rules of thumb or study the effect of data, a large number of further experiments would be required for us to be confident in the findings and maintain the overall standard of careful experimentation. Therefore, we believe that these directions are promising topics for further research, but beyond the scope of our investigation.

Reviewer comment (Writing):
The caption of Figure 1 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.

Author response:
We apologize that this was unclear. In this context, we use optimal to refer to $\ell^*$ in Section 2.2, which is the minimum loss achieved in the learning rate sweep. To address this, we have revised the caption to read “LR sensitivity measures the expected deviation from the minimum achieved loss when varying learning rate across three orders of magnitude”.

Reviewer comment (Evaluation):
In the introduction, the authors state: "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 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?

Author response:
We believe it is an open question why depth increases LR sensitivity at a faster rate than width in our experiments. As you suggest, a different initialization or parameterization when scaling depth may resolve this issue. This is an interesting question for future research.

Reviewer comment (Writing):
In section 3.3, the authors write "We now examine whether it is possible to predict the logit growth instability before it occurs." 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. Stating this more clearly in the main text would be beneficial.

Author response:
Thank you for the suggestion. We agree and have changed the main text to state this more clearly as you’ve recommended. In our revised version, Section 3.3 begins with: “A central question when studying instabilities is whether they can be predicted using small-scale proxy experiments. We now examine whether it is possible to predict the logit growth instability before it occurs using previous runs with smaller models.”.

Reviewer comment (Experiments):
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?

Author response:
Indeed, we use AdamW in our experiments. We choose AdamW because it is a popular optimizer used for many large scale runs including GPT-3 and LLaMA. We believe that some of the Transformer training instabilities we study are present when using other optimizer variants. For instance, the output logit divergence instability (Section 3.1.2) is reported by Chowdhery et al., 2022 for the PaLM model, which uses the “parameter scaling” feature from Adafactor. On the other hand, we expect that some instabilities we study are absent when using other optimizers. For instance, the epsilon instability from Section 3.4 is not expected when using SGD, since SGD has no epsilon hyperparameter.