Summary
This paper presents theoretical justifications and empirical evidence that transformers demonstrate incremental learning dynamics in the low-initialization regime. The authors consider a very restricted diagonal attention model along with a range of restrictive assumptions to theoretically characterize a number of features of learning dynamics in single-layer transformers (a single attention layer). The assumptions, which include nondegeneracy of dynamics (4.2), existence of stationary points that are strict local minima (4.3), and robustness of gradient flow to perturbations (4.4). These assumptions, in the diagonal model, allow the authors to prove that learning evolves through discrete, incremental stages of gradual rank increase. These assumptions and theoretical predictions are experimentally verified using a toy learning scenario matching the assumptions, and results are given on a full multi-layer, multi-head vision transformer which demonstrate similar learning dynamics despite being significantly far from the restrictive assumptions of the theory.
Strengths
1. **Convincing Theory**. The theoretical framework and proof of incremental learning dynamics is fairly easy to follow and, although based on a very restricted model and stringent assumptions, I think provides a number of useful tools of understanding and studying the learning dynamics of more complex models.
2. **Fairly Convincing Empirical Confirmation**. The empirical validation of incremental dynamics in the single attention layer model (Figure 3) are interesting -- in both the diagonal and full model the empirically observed behavior follows the predicted incremental "activation" of directions. However I find it odd that the evolution of singular values is so markedly different between the value-output and key-query matrices (Figures 3a-d and 3b-e). The value-output eigenvalues are significantly more sparse and their activations more "instantaneous". This difference is also roughly observable in the ViT results in Figure 2 (although reversed!).
Weaknesses
1. **Organization**. As will many theoretical papers, some of the most interesting observations are relegated to the Supplementary Material -- which I appreciate is inevitable. Reorganizing the preliminaries and theoretical development of the main paper I admit is delicate, but I think that the stable rank analysis in Figure 4 and illustration of low-rank bias in Figure 5 add little to the empirical results section of the main paper. Instead, I think this space could have been used better providing experimental evidence of the assumptions 4.2, 4.3, and 4.4 (is there an assumption 4.1 that I am missing?). Some of the plots in Figures 9 and 10 of the Supplementary Material, or even the *very* clear illustration of stepped behavior and dependence on $\alpha$ from Figure 6 would be excellent illustrations to include in the main paper. I find the plots in rescaled training time much easier to interpret.
2. **ViT Training Regime**. All of the empirical results on ViT were generated using Adam, which I fear might introduce its own dynamics into the learning process due to its modulation of gradients using the diagonal empirical Fisher matrix. It would be interesting to see results generated using vanilla SGD to verify that the observed learning dynamics are due to the low-initialization regime.
Questions
1. Why is the rank at initialization bounded (at 128?) in the plots in Figure 5?
2. Is the top/bottom organization of plots indicated in the captions of Figures 2 and 3 correct? If so, can you explain the difference between behavior between the key-query and value-output eigenvalues?
3. Are the same learning dynamics observed in ViT when training with SGD as opposed to Adam?
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
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.
Limitations
The authors provide a discussion of the limitations of the results presented in the paper, acknowledging that the restrictive assumptions and requirements of the theory are quite far from actual practice. They also connect their to recent work on low-rank model adaptation which could potentially exploit the theoretical and empirically observed incremental learning dynamics of transformers.