Kati8/ICLR_2023_2025 huggingface dataset:
authors: Mitchell Wortsman
authors: Peter J Liu
authors: Lechao Xiao
authors: Katie E Everett
authors: Alexander A Alemi
authors: Ben Adlam
authors: John D Co-Reyes
authors: Izzeddin Gur
authors: Abhishek Kumar
authors: Roman Novak
authors: Jeffrey Pennington
authors: Jascha Sohl-Dickstein
authors: Kelvin Xu
authors: Jaehoon Lee
authors: Justin Gilmer
authors: Simon Kornblith
reviews: 
Comments: 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.
Strengths: Large scale transformers are expensive, important and suffer from instabilities. Providing a small-scale proxy model is impactful.

The paper is well written and the experiments are cleanly described. 

The 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.

Section 3.3 reads a little anecdotal to me. A more systematic study would be better.
Recommendation: 8: accept, good paper
Reviewer confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness/correctness: 4 excellent
Clarity: 4 excellent
reviews: 
Comments: 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. 

The 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.

The 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?

One insight which seems useful is that default eps=1e-8 in AdamW might appear too large and cause updates to vanish.
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?
Note: the concerns have been addressed.
Recommendation: 8: accept, good paper
Reviewer 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.
Soundness/correctness: 3 good
Clarity: 3 good
reviews: 
Comments: 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.

Two 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.

The 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.
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.

The 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.

I 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.

As 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.
Recommendation: 8: accept, good paper
Reviewer confidence: 2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness/correctness: 3 good
Clarity: 3 good