Transformers learn through gradual rank increase
Authors: Enric Boix-Adserà
Etai Littwin
Emmanuel Abbe
Samy Bengio
Joshua M. Susskind
Venue: NeurIPS 2023 poster (2023)
Keywords: transformers
low-rank bias
incremental learning

Abstract: We identify incremental learning dynamics in transformers, where the difference between trained and initial weights progressively increases in rank. We rigorously prove this occurs under the simplifying assumptions of diagonal weight matrices and small initialization. Our experiments support the theory and also show that phenomenon can occur in practice without the simplifying assumptions.

Transformers learn through gradual rank increase
Enric Boix-Adserà 1.2
Etai Littwin"!
Emmanuel Abbel 1.3 Samy Bengio' Joshua Susskind'
Apple MIT SEPLL
eboixêmit edu, emmanuel abbe@epf1.ch emmanuel..beeeppp..c.. (elittwin, bengio, juseggoolsusstindaappplpeee com
Abstract
We identify incremental learning dynamics in transformers, where the difference between trained and initial weights progressively increases in rank. We rigorously prove this occurs under the simplifying assumptions of diagonal weight matrices and small initialization. Our experiments support the theory and also show that phenomenon can occur in practice without the simplifying assumptions.
1 Introduction
The transformer architecture achieves state of the performance in various domains, yet still lack a solid theoretical understanding its training dynamics [VSP* 17, DCLTI9, LOG 19, DBK 20J. Nevertheless, the theoretical toolbox has matured over the last years and there are promising approaches. One important line of work examines the role that initialization scale plays on the trajectory taken by gradient descent [JGH18, COB18, GSJW19, MGW+ 20, JGS SS21, KC22]. When the weights are initialized small, it has been shown for simple networks that an incremental learning behaviour occurs, where functions of increasing complexity are learned in stages. This regime is known to be richer than the large-initialization regime', but the incremental learning dynamics difficult analyze, and Can we apply this analysis to transformers? Namely:
far understood only for extremely simple architectures.
Are there incremental learning dynamics when training transformer architecture?
An obstacle that past work on incremental learning has mainly studied linear networks [Ber22, ACHL19, MKAA21, LLL20, WGL+119, JGS+21, GSSD19, SKZ*23, PF23], with one paper studying nonlinear 2-layer fully-connected networks [BPVF22]. In contrast, transformers have nonlinear attention heads that do not fall under previous analyses: given X E RRax attention head computes
attention(X X; WK. Wo. Wv, Wo) = smax x
(I)
where WK. WQ, Wv, Wo dxdd are trainable matrices, and the softmax applied row-wise. transforner is even more complex, since it is formed by stacking alternating layers of attention heads and feedforward networks, along with residual connections.
Main finding Our main finding that transformers exhibit incremental learning dynamies, where the difference between the trained and initial weights incrementally increases in rank. Our results have a theoretical component and an experimental component.
In the large-initializaton regime, deep learning behaves as a kernel method [JGHI8, COB18], Various separations with kernels are known for smaller initialization: e.g., [GMMM19, ABM22, MKAS21].
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Figure 1: For attention head in ViT trained on (a) CIFAR-10, and (b) ImageNet, ized spectra of Wx k W Q at initial- ization (in red), and of the learned perturbations to Wx k W at differ- ent iterations (in green).
0.75
1100
2-20000
1100
plot the normal-
0.50
0.25
0.25
0.00
0.00
100
125
23
30
75
100
125
#figerwatue
#bgenwake
Theoretical contributions For our theory, we study simplification of the transformer architec- ture, where the attention head weights are diagonal matrices: i.e., in each attention head we have WK diag(wx), ). where Wk E R are trainable weights, and similarly for Wo. Wv and Wo. We rigorously establish the training dynamics of this architecture under gradient flow when the initialization is small. We prove that dynamics occur in discrete stages: (I) during most of each stage, the loss plateaus because the weights remain close to saddle point, and (2) at the end, the saddle point is quickly escaped and the rank of the weights increases by at most one.
This theoretical result on transformers follows from a general theorem characterizing the learning dynamics of networks JNN that depend on the product of parameters u, E IRP
(a u, v) h(æ;; u v)
(2)
where æ is the input, @ denotes the elementwise product, and f is smooth function.
Theorem 1.1 (Informal statement of incremental learning dynamics), Let fNN be a network of the form (2), and suppose that the weights are initialized very small: i.e., the entries of u, v are initialized on the order (aa) for some small a > O. Then the dynamics of gradient flow training effectively proceeds in discrete stages, each one lasting time @(log(1/a)). In each stage, the number of nonnegligible entries of u @ v increases by at most one.
transformer with diagonal weight matrices falls under this result when only train the attention head weights. For example, if the transformer has one attention head, then we can take t = [wx,, wy] E 2d and = (wq, wo] E 2d to be concatenations of the diagonal entries of the weights of the head; see Example 3.2 for more details and the extension to transformers with many heads. Then, using Theorem 1.1, we see that in each stage either WK k W diag( (w )diag (wQ) or Wy W = diag(w)) (ddggg)) increases in effective rank by at most one. 2
Experimental contributions In our experiments, we first validate our theoretical results, which require the simplifying assumptions of small initialization and diagonal weight matrices.
Then, we conduct experiments on vision and language transformers in settings closer to practice, without any of the assumptions required by our theoretical analysis. Perhaps surprisingly, we again observe incremental learning dynamics, even though the assumptions of the theory are not met. The difference between trained and initial weights has low rank, and the rank of this difference grows gradually during training; see Figure 1. The incremental nature of the dynamics is easier to see for ImageNet, since for CIFAR-10 the rank of the weight difference does not grow as much.
1.1 Related work
Relation to LoRA We note intriguing connection to the LoRA algorithm, where pretrained base model is cheaply fine-tuned by training a low-rank perturbation of the weights [LFLY18, AZG20, HSW+21J. The method is surprisingly powerful, and recently LoRA has been fundamental to allowing the open-source community to inexpensively fine- tune language models [PA23, TGZ 23]. On the other hand, in work we observe that the trained weights initial weights due to the training dynamics, without having to apply an explicit rank constraint as in LoRA. This raises exciting open question for future work: can explain and improve algorithms like LoRA by better understanding and quantifying the incremental dynamics of large transformers?
low-rank perturbation of the
We also remark that Theorem 1.1 is interesting in its own right and may have other applications beyond transformers. It qualitatively recovers the incremental dynamics result of [Ber22, PF23] when specialized to linear diagonal networks, i.c., when (a u, v) Σ.' U.U.

---

Low-rank bias in nonlinear networks For 2-layer networks, it is known that low-rank bias in the weights emerges if the target function depends on a low-dimensional subspace of the input [ABM22, ABM23, DIS22, BBSS22, MHPG 22J. The results of [ABM22, ABM23] are especially relevant, since they show that the rank of the weights increases in a sequential manner, determined by the "leap complexity" of the target function, which is reminiscent of our empirical observations on transformers. See also [FVB 22, TVS23] for more investigations of low-rank bias in 2-layer networks under different assumptions. For transformers, [YW23] report that empirically the trained weights (using default initialization) are not low-rank. This is consistent with our claim that the difference between initial and trained weights is low-rank, since the initial weights might not be low-rank.
Incremental learning dynamics Several works prove incremental learning behaviour in deep linear networks when the initialization is small. [GBLJ19] has shown that gradient descent dynamics on 2-layer linear network with L2 loss effectively solve a reduced-rank regression problem with gradually increasing rank. [GSSD19] prove a dynamical depth separation result, allowing for milder assumptions on initialization scale. [ACHL19, MKAA21] show implicit bias towards low rank in deep matrix and tensor factorization. [LLL20] show deep matrix factorization dynamics with small initialization are equivalent to a greedy low-rank learning (GLRL) algorithm. And [JGS+21 21] independently provides a similar description of the dynamics, but without requiring balanced initialization. Finally, [Ber22, JLLL* U.L+22 23, PF23] overcome a technical hurdle from previous analyses by proving incremental learning for the entire training trajectory, rather than just the first stage. In contrast to our result, these prior works apply only to linear networks with certain convex losses, whereas our result applies to nonlinear networks. In order to make our extension to nonlinear networks possible, we must make stronger assumptions on the training trajectory, which we verify hold empirically. As far as we are aware, one other work on incremental learning handles nonlinear networks: [BPVF22] proves that 2-layer network learns with two-stage incremental dynamic; but that result needs the stylized assumption that all data points orthogonal.
1.2 Paper organization
Sections 2, 3, and 4 contain theoretical preliminaries, definitions of the models to which our theory applies, and our main theoretical result on incremental dynamics. Section 5 provides experiments which verify and extend the theory. Section 6 discusses limitations and future directions.
2 Preliminaries
We consider training a network fuN(:: 0) parametrized by a vector of weights 0, to minimize a loss
((0)= Em,,y [[(y fNN( f))]
where the expectation is over samples (x. y) E IRd x Rd from training data distribution, and e: Rdy x Rdont + RR. Consider a solution (()) to the gradient flow
d0 = V
e(0) = 00, dt
where a >> 0 is parameter governing the initialization scale, that we will take small. For our theory, we henceforth require the following mild regularity assumption on the loss and data.
Assumption 2.1 (Regularity of data distribution and loss). The function E(y, () is continuously dibeediifierettiable in the arguments [y. c] Rd +dout, There exists C' > 0 such that almost surely the data is bounded by ||>||| I|, |gll s C.
The assumption on l is satisfied in typical cases such as the square and the cross-entropy losses. The data boundedness is often satisfied in practice (c.g., if the data is normalized).
We also the notation supp(a) == {i: A 0} to denote the support of vector
*Gradient flow training can be obtained as a limit of SGD or GD training as the learning rate tends to 0 (see. e.g., [Bac20J). It is popular testbed for studying learning dynamics (see e.g., ISMG13, ACHI8, RC20J), since is generally simpler to analyze than SGD.
3

---

3 Neural networks with diagonal weights
Our theory analyzes the training dynamics of networks that depend on products of diagonal weight matrices. We use O to denote elementwise vector product.
Definition 3.1. A network ÍNN is smooth with diagonal weights O (u, v) R229 if it is of the form
JNN ax; 0) = h!a;;; 2 O.))
where h Rd x R" - Root is continuously twice-differentiable in its arguments in Rd++p,
The assumption on h precludes the use of the ReLU function since it is not Otherwise the assumption fairly mild since any can be used to express an architecture of any depth as long as the nonlincarities used in ViT). We describe how to express a transformer with diagonal weights.
wiice-differentiabl,, which includes for example GeLUs (as
Example 3.2 (Transformer with diagonal weights). A transformer with L layers and H attention heads on each layer is defined inductively by Zo = X E Raxd and
* (Attention layer) Z, = -Z E attentionnZee.... W K W W w
(Feedforward layer) Z: Z: + o (Z:W
where W w W w dd are attention parameters, and W A: W RR are the feedforward parameters, and a is continuously twice-diiferentiable activation. Suppose that the attention parameters are diagonal matrices: i.e., W diag(ukk) E gdxd and similarly for the W W w matrices. Then by the definition of the attention layer (1), the final output of the transformer Z only depends on the attention parameters through the elementwise products f w @ whit and wy @ wo other words, we can write
z u
for vectors u = [w wy ((.))(L)(( E RRHHL and v = [w w (( E 2d2 and some smooth model h. Thus, if only the attention layers are trained, the diagonal transformer fits under Definition 3.1.
4 Incremental learning in networks with diagonal weights
We prove that the initialization scale a is small, then learning proceeds in incremental stages, where in each stage the effective sparsity of the weights increases by at most one. These stages implicitly defined by Algorithm below, which constructs a sequence of times 0 = To < Ti <... < Ts < ... and weight vectors 0O 01 .g...... 0+ ,.... R22p that define the stages. We prove the following:
Theorem 4.1 (Incremental dynamics at small initialization). Let fNN be a model with diagonal weights as Definition 3.1. For any stage i and time t E (Tk, Tk+1) the following holds under Assumptions 2.1, 4.3, 4.4 and 4.5. There is (n(()) > 0 such thatfor all a < an, there exists a unique solution [0, log(1/a)) + TRP to the gradient flow (3) and
lim Ott log(l/o)) 0² a-0
and each stage the sparsity increases by at most one: supp(@)+!) s0p S (ix).
Application: transformer with diagonal weights Before giving the intuition for this theorem and stating the assumptions formally, let us discuss its application the diagonal transformer model from Example 3.2. As corollary of Theorem 4.1, the gradient flow on diagonal transformer with small initialization will learn in stages, where in each stage there will at most one head [H] in one layer E [L] such that either the rank of L.i of w (W-+))
K (Wo.)) = diag(w) (wag(wx f. k )diag(we") &.i or the rank dynamics in the toy case of a single attention head trained in student-teacher setup.
diag(w("") diag((wf") increases by at most one. Figure 2, we illustrate these
Abusing notation, for 0 (u, v) E RP 12 write supp(0) supp(u) U Usupp(e).

---

Algorithm Incremental learning in networks with diagonal weights
1: 6", 0 + R2, To 6-
2: for stage number k = 0, 1, ., do
3: # (A) Pick new coordinate ig E [p] to activate.
4: For each i, define time (1) until active using (10).
5: Pick winning coordinate i& using (11)
6:
Calculate time T using (11) and break if oo
7: Update logarithmic weight approximation b using (12)
8: # (B) Train activated coordinates stationarity.
9: gh+1 +- limiting dynamics point from (13)
end for
Entries of ding(wg)ddag((wy) for alpha = o. 0.0001
Diagonal of Wo/Ms
Training loss vs. rescaled time, for various alpha
alpha 0.1
0.0
slpha 0.01
0.08
alpha
alpha 1e-08
1e:16
alpha 1o-32
-1.0
200
600
1000
/ lag()/apph)
0.04
Entries of diag(w,ldiag(wwo. for alpha 0.0001
Dlagonal of WoMe
0.02
0.00
time (olph))
time log(2/elphs)
a
(b)
(a) Loss versus rescaled time in the toy task of learning an attention head with diagonal weights, for various initialization scales a. The loss curves converge as + 0 to curve with stagewise loss plateaus and sharp decreases, as predicted by the theory; some stagewise learning behavior is already clear with a 0.01. (b) Each line shows the evolution of one of the entries of diag( ( (diag(wx)) and diag( wy )diag(w over rescaled time, demonstrating that the rank of these matrices increases incrementally; see Appendix A for experimental details and further results.
Figure 2:
4.1 Intuition for incremental learning dynamics
We develop informal intuition for Theorem 4.1 and fill out the definition of Algorithm 1. A model fNN with diagonal weights = (u, as in Definition 3.1 evolves under the gradient flow (3) as de
du
dt = u G (())
u 3 g())
where
dt
g( = En.y h( 3 DM(x;
Here De(g. .) E xlxu the derivative of in the second argument and Dh(a,,) E Ro XP is the derivative of h in the second argument. The first key observation is a conservation law that simplifies the dynamics. It can be viewed as the balancedness property for networks with linear activations [ACHI8, DHL18], specialized to the case of diagonal layers.
Lemma 4.2 (Conservation law). For any i [p] and any time t, we have
ut(() of (t) = (0) - =
(5)
Proof. This follows from (u v/) (0) (0) 0.
5

---

The conservation law reduces the degrees of freedom and means that we need only keep track of p parameters in total. Specifically, if we define w,((t := 24 (t) + (t))), then the vector w(t) (l) + vl)) evolves by
dw dt
3 g
Using the conservation law (5), keep track of the weights in terms of the initialization and (t):
02 -0) (tt) = [ (t) u 22
(0),
02 (0)
02 (0)
w(t) )
(w((t)
w(()
Therefore it suffices analyze the dynamics of w()).
4.1.1 Stage of dynamics
Stage 1A of dynamics: loss plateau for time (log((1/o) At initialization, 0(0) = 0 because the weights are initialized on the order of a which is small. This motivates the approximation (t)()) ~ g0)), under which the dynamics solve to:
(tt (0) 200)
(8)
Of course, this approximation is valid only while the weights are still close to the small initialization. The approximation breaks once one of the entries of ((t) reaches constant size. By combining (7) and (8), this happens at time = Ti log(1/a) a) for
T min ( EElp
Until this time, the network remains close to its initialization, and so we observe a loss plateau.
Stage IB of dynamics: nonlinear dynamics for time O(1) Subsequently, the loss decreases nonlinearly during (1) time-scale, which is vanishingly short relative to the time-scale of the loss plateau. To prove this, we make the non-degeneracy assumption that there is a unique coordinate to such that 1/1gto (0)l = Ti. Under this assumption, in stage IA all weights except for those at coordinate ig remain vanishingly small, on the order of oa (1). Concretely, for any small 0, there a time & (e) Ti log(1/a) and sign E {+1, -1} such that
Wio (t,) R e,
(., (t1)
and le (1) for all i i io.
Because all coordinates except for ig have vanishingly small ox. weights after stage IA, we may perform the following approximation of the dynamics. Zero out the weights at coordinates except for lu and consider the training dynamics starting 0 (cesg, sce;)) ). After O(1) time, should expect these dynamics to approach a stationary point. Although the evolution nonlinear, all entries remain zero except for the in entries, the stationary point is also sparse. Mathematically. there is time [ = 4 + (1)) x Ti log(1/a) such that
(f1) x3 (ae,,, soe,.) == 0¹
for some a E Ryo, where 01 stationary point of the loss. Despite the nonlincarity of the dynamics, the approximation be proved using Grönwall's inequality since i O(1) is constant time-scale.
To summarize, we have argued that the network approximately reaches stationary point that is 1-sparse, where only the weights at coordinate in are nonzero.
ensure that initialization u(0) and u(0) v00) nonnegative. This implies u(t) nonnegative. The fact that usg and Vi are roughly equal magnitude but might differ in sign the conservation law (5). See Appendix C.3 details.
Witthott loss of generality,
'The entries of and close magnitude (but may differ in sign) because of the conservation law (5).

---

4.1.2 Later stages
We inductively extend the argument to any number of stages k where each stage has (loo((//a)) time plateau, and then a O(1)- nonlinear evolution, with the sparsity of the weights increasing by at most one. The argument to analyze multiple stages is analogous, but we must also keep track the magnitude of the weights on the logarithmic scale, since these determine how much longer. . Inductively on k, suppose that there is some TA E R, b E IR and Oc RP and a time l Tx log(1/a) such that
loga (w((1)) Re bh and ((Ta) R0 0²,
where 0h is stationary point of the loss. Our inductive step shows that there is Tk+1 E IR such that during times t (4x, Th+1 log(1/a) - (1))) the weights remain close to the stationary point from the previous stage, i.e. ((t) l And a time f--1 T (og(1/a) we have
log,, (fx+1))) = b k+1 and ((f++1) = g++1
where g and b*+1 are defined below (summarized in Algorithm 1). Most notably, @h++1 is stationary point of the loss whose support grows by at most one compared to lh.
Stage (k- + 1)A, loss plateau for time ([log(1/a)) At the beginming of stage k + 1, the weights are close to the stationary point 0%, and so, similarly to stage 1A, linear dynamics are valid.
(t) w(xx) (1-Fx)
Using the conservation law (7), derive a "time until active" for each coordinate E [p]. which corresponds to the time for the weight at that coordinate to grow from 0a (1) to @(1) magnitude:
(bf + (0: (O*), 00,
if (0*) + 0
Au
(10)
if ( 0
The linear dynamics approximation (9) breaks down at time Te+1 log(1/a), where
Ta+1 T (x())) is arg min (1),
(11)
EEjp
which corresponds to the first time at the weights at coordinate grow from 0a (1) O(1) magnitude. And at times x Tx+1 log(1/a), on the logarithmic scale w is given by
(w)() R k++ = b² - - Aax))
(12)
Stage (k: 1)B of dynamies: nonlinear dynamics for time O(1) Subsequently, the weights evolve nonlinearly during O(1) time. In a similar way to the analysis of Stage IB, we show that at time f++1 tet1 + O1) x Tk+1 log(1/a), we have
((fx+1) so +++1 = lim lim 2s (t, c) c-0 100
(13)
where the dynamics a (t, e) 2 are initialized (0, l + (gn and evolve according the gradient flow (t.s) (oL(b"). This concludes the inductive step.
4.2 Assumptions for incremental dynamics
To make this intuition rigorous, we formalize below the assumptions required for Theorem 4.1. In Figure 3 and Appendix A, we provide experiments validating these assumptions on the toy model.
The first assumption is that the dynamics are non-degenerate, in the sense that two coordinates do not have weights that grow from o,, (1) to O(1) size at the same rescaled time. We also place technical condition to handle the corner case when coordinate leaves the support of the current stage's stationary point.
Assumption 4.3 (Nondegeneracy of dynamics in part (A)). The initialization satisfies |u,(0) * [o,( for all For stage k. either Tk+1 (11). Finally, for all i supp(0"-1)) suup((0)) we have g, (0*) o 0.
there is unique minimizer in min, (+(1e) in
7

---

Entries of W.W. perturb time log(1/alpha) 778.0
Entries of W.Wg perturb time log(1/alpha) 321.4
Grordimato value
Coardinat value
-1
778
779
780
781
782
703
784
322.5 325.0
327.5
330.0
332.5 335.0 337.5
time/ log(1/olpha)
loght/lpph))
Figure Validation of assumptions the toy model of learning a single attention head. (a) Assumption 4.4: weights perturbed at random time during training (solid lines) tend back to the near-stationary point (dashed lines). (b) Assumption 4.5: weights perturbed the beginning of stage (solid lines) have same nonlinear evolution as without perturbation (dashed lines). Details of these experiments and further validations are provided in Appendix A.
Next, we require that very small perturbations of the coordinates outside of supp(O") do not change dynamics. For this, suffices that 0 strict local minimum.
Assumption 4.4 (Stationary points are strict local minima). For stage k, there exist d 0 and c 0 such that for ū E (uu 5) supported on supp(u"), we have
C(ū, 8h C ū) ca 1/ ū
Finally, we require robust version the assumption that the limit (13) exists, asking for convergence neighhorhood of 0 even when the initialization is slightly noisy.
Assumption 4.5 (Noise-robustness of dynamics in part (B)). For any stage k with Tx+1 < oc and any 6 > 0, there are and r R250- - R such that the following holds. For any E B(u*, 8) n Rgoo supported on supp(u) S supp(u") {&x there exists unique solution 2b oo) + R of the gradient flow de (e((") initialized (0) (ū, s s @ u) and at times T(
(|x((t) g <e.
5
Experimental results
We run experiments that go beyond the toy diagonal attention head model (see Figures 2 and 3) to test the extent to which low-rank incremental learning occurs in popular models used in practice. We conduct experiments with vision transformers (ViT) [DBK+20] trained on the CIFAR-10/100 and ImageNet datasets, and with the GPT-2 language transformer [BMR+ 20] trained on the Wikitext-103 dataset. Full experiments are deferred to Appendix B.
Gradual rank increase in vision and language models We train practical transformer architectures on vision and language tasks using Adam and the cross-entropy loss. We train all layers (including the feedforward layers). To capture the low-rank bias with a non-vanishing initialization scale, we study the spectrum of the difference A WW W and AW W between the weights post-training and their initial values. Specifically, in Figure 4, we plot the stable rank of the differences A AWR W
and A Wy W The weight perturbation learned during the training process gradually increases in stable rank during training. and is ultimately low-rank when compared to the initial spectrum. Finally, for CIFAR- 10, we plot the spectrum of Á AW W against that of its initialized state in Figure for different self-attention heads, illustrating that the weight perturbation learned during the training process is extremely low-rank when compared to the initial spectrum, In Appendix B, we also study optimization with SGD, which shows similar gradual rank increase behavior.
Effect of initialization scale We probe the effect initialization scale on gradual rank increase dynamics for a ViT trained on CIFAR-10. 10. We use a ViT of depth 6, with 8 self-attention heads per layer (with layer normalization). We use an embedding and MI. dimension of demb 512, and head dimension of d2 128 (i.e WK. Wo. Wv. Wo E B gglmsxdx) We train the transformer

---

Layer
Layer 3
Layer
SMLAS
Stabb wa
Stabie
10000 20000 30000 40000 30000 Faaratinn
20000 43009
20000 20000 30000 40000
nterooons
Rierobons
(a) VIT, CIFAR-10
(b) VIT, CIFAR-100
(c) VIT, ImageNet
(d) GPT-2, Wikitext-103
Figure 4: Stable rank of A AWK 2 (blue) and A AW W (orange) on an arbitrary chosen layer throughout training for four different pairs of networks and tasks. The stable rank matrix W is defined || deviation (shaded area) are computed across all heads in each attention layer. Full details and results are in Appendix B.
and gives smooth approximation of the rank. Mcan and standard
Layer
Layer
Layar
Spectrum at Int Spectrum of perturbation
Spectrum Int
0.0
Spectrum at int Spectrum of perturbacion
1.0
Spectrum perturtation
0.6
0.5
(b)
0.5
0.2
0.0
0.0
0.01
50
100 #t sigemaalae
150
200
50
100 Atigonoalue
:50
200
50
100 #P Pigeenvaee
150
203
Layer
Layer
Layer
0.55
Spectrum at int
0.0
Spectrum ini
Spectrum at int Spectnum of perturbarion
0.6
Specinum of perturbation
0.6
Spectrum of perturiation
0.75
(d)
0.4
0.50
0.2
0.2
0.25
0.0
0.0
0.00
'0
100
"in
700
0
50
100
110
200
h
10
100
150
200
#Eigenvolue
#Eigenvolue
dEigenvalue
Figure 5: Spectrum of the weight perturbation AWR WT vs. initialization in a vision transformer trained on CIFAR-10, using Adam and default initialization scale, in random self-attention heads in different layers. The learned perturbation exhibits extreme low-rank bias post-training even default initialization scales. Analogous plots for CIFAR-100 and ImageNet are in Appendix B.
using Adam with the cross cross-entropy loss. We train all layers (including the feedforward layers) while varying the initialization scale of all layers by multiplying their initial values by a scale factor (we fix the scale of the initial token mapper). Figure 6 shows the evolution of the principal components of AW WT and A AWr.i W for a randomly-chosen self-attention head and layer throughout training, exhibiting incremental learning dynamics and a low-rank bias. Note that incremental learning and low-rank bias are increasingly evident with smaller initialization scales, as further demonstrated in Figure 7.
6
Discussion
We have identified incremental learning dynamics in transformers, proved them rigorously a simplified setting, and shown them experimentally in networks trained with practical hyperparameters.
Limitations There are clear limitations to our theory: the diagonal weights and small initialization assumptions. More subtly, the theory does not apply to losses with exponential-like tails because the weights may not converge to a finite value and so Assumption 4.4 is not met (this could possibly be addressed by adding regularization). Also, the architecture must be smooth, which precludes ReLUs- but allows for smoothed ReLUs such as the GeLUs used in ViT (DBK+20), Finally, the theory is for training with gradient flow, while other optimizers such as Adam are used in practice instead [KB 14]. Nevertheless, our experiments on VITs indicate that the incremental learning dynamics occur even when training with Adam.

---

scale 10-4
scale 10-2
scale
0.5-
magnttat
EN maggntt. u00
8.6
(a)
(b)
(c)
84
Mtdlu 1e-4
btale 1e-2
SCale
2.0
0.5
EV maceituus
8.5-
6.6
(d)
(f)
0.4
2000 2000 4000 Beral
20000 2000 4000
1000
2000
3000 4000
Burnstonn
berations
Figure 6: Training a vision transformer on CIFAR-10 using Adam, while varying the initialization scale (unit scale indicates default initialization). Plotted are the evolution of the eigenvalues of AWK W (a) (c) and AW, W (d) (f) random self-attention head in the second layer throughout training. Incremental learning dynamics and low-rank bias are evident for all scales, albeit more pronounced at smaller initialization scales.
Layer
Layer 3
Layer 5
10
Stable rm-k
Stable rank
(a)
(b)
6
(c)
10
Wenght
Weigh
Whighht Eialieetiinn
Figure 7: Stable rank of AWxW2 per initialization scale (Unit scale refers the default initializa- tion) in different self-attention heads post-training, at layers 1, 3, 5. At each layer, the stable rank mean and standard deviation are computed across heads per layer, for each initialization scale. All models were trained on CIFAR-10 10 using the Adam optimizer. Smaller initialization scales lead to lower-rank attention heads.
Future directions An interesting avenue of future research is to develop a theoretical understanding low-rank perturbation of
of the implicit bias in function space of transformers whose weights randomly initialized weights. Another promising direction is to examine the connection between our results on incremental dynamics and the LoRA method [HSW+21] with the goal of explaining and improving on this algorithm; see also the discussion in Section 1.1. Along this vein, a concurrent work [ZZC-23] independently observes gradual rank increase dynamics during transformer training and this inspires a low-rank training algorithm that obtains runtime and memory improvements over regular training. The results of [ZZC*23] 23] are complementary to ours, since they study the feedforward layers of the transformer, and their theory applies to linear networks in the standard initialization scale; contrast, we study the attention layers, and our theory applies to nonlinear networks with small initialization scale.
Acknowledgments
We would like to thank Vimal Thilak for his help in setting up the infrastructure for conducting experiments, and the anonymous reviewers for their helpful feedback.
References
[ABM22] Emmanuel Abbe, Enric Boix- and Theodor Misiakiewicz. The merged-staircant property: necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks, COLT, 2022.
10

---

[ABM23] Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. arXiv preprint arXiv:2302. 1055, 2023.
[ACH18] Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In International Conference on Machine Learning, pages 244-253. PMLR, 2018.
[ACHL19] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
[AZG20] Armen Aghajanyan, Luke Zettlemoyer, and Sonal Gupta. Intrinsic dimension- ality explains the effectiveness of language model fine-tuning. arXiv preprint 13Xiv:20/2.13255, 2020.
[Bac20] Francis Bach. Effortless optimization through gradient flows. Machine Learning Research Blog. https:/ffrancisbach. comfgradient-flows. 2020.
[BBSS22] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-inde index models with shallow neural networks. arXiv preprint arXiv:2210.1565. 2022.
[Ber22] Raphaël Berthier. Incremental learning in diagonal linear networks. arXiv preprint arXiv:2208. /4673, 2022.
[BR + 20] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[BPVF22] Etienne Boursier, Loucas Pillaud-Vivien, and Nicolas Flammarion. Gradient flow dynamics of shallow relu networks for square loss and orthogonal inputs. arXiv preprint arXiv:2206.009333 2022.
[COB18] Lénaïc Chizat, Edouard Oyallon, and Francis R. Bach. On lazy training in differentiable programming. In Neural Information Processing Systems, 2018
[DBK Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghami, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x 16 words: Transformers for image recognition scale. ArXiv, abs/2010. 2020.
[DCLT19] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre- training of deep bidirectional transformers for language understanding. ArXiv, abs/1810.04800 2019.
[DHL18] Simon Du, Wei Hu, and Jason Lee. Algorithmic regularization in learning deep homogeneous models: Layers automatically balanced. Advances in Neural Information Processing Systems, 31, 2018.
[DLS22] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference Learning Theory, pages 5413-5452. PMLR, 2022.
[FV++22 Spencer Frei, Gal Vardi, Peter L Bartlett, Nathan Srebro, and Wei Hu. Implicit bias in leaky relu networks trained on high-dimensional data. arXiv preprint arXiv:2210.07083 2022.
[GBLJ19] Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of discrete gradient dynamies in linear neural networks. Advances Neural Information Processing Systems, 32, 2019.
[GMMM19] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limita- tions of lazy training of two-layers neural network. Advances in Neural Information Processing Systems, 32, 2019.
11

---

[GSJW19] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy learning in deep neural networks: an empirical study. ArXiv, 0oas/9006.08034 2019.
[GSSD19] Daniel Gissin, Shai Shalev-Shward and Amit Daniely. The implicit bias of depth: How incremental learning drives generalization. arXiv preprint TrrXiv:1099 12051, 2019.
[HSW*21] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09683 2021.
[JGH18] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Con- vergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[JG++21] Arthur Jacot, Francois Gaston Ged, Berfin Simsek, Clément Hongler, and Franck Gabriel. Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. 2021.
[JLL [2] Jikai Jin, Zhiyuan Li, Kaifeng Lyu, Simon S Du, and Jason D Lee. Understanding incremental learning of gradient descent: A fine-grained analysis of matrix sensing. arXiv preprint arXiv:2301./ /550,, 2023.
[BB14] Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv: /4/i::/42..0 /etl:/422..990 22800 2014.
[KC22] Daesung Kim and Hye Won Chung. Rank-1 matrix completion with gradient descent and small random initialization. ArXiv, abs/2212.09396, 2022.
[LFLY18] Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes. arXiv preprint 1oXiv:/804.0SS 2018.
[LLI220] Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradi- ent descent for matrix factorization: Greedy low-rank learning, ArXiv, abs/2012.09839, 2020.
[LOG + 19] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. ArXiv, abs/1907.1 1692, 2019.
[MGW+20] Edward Moroshko, Suriya Gunasekar, Blake E. Woodworth, J. Lee, Nathan Srebro, and Daniel Soudry. Implicit bias in deep linear classification: Initialization scale vs training accuracy. ArXiv, http//0070067666 2020.
[MHPG*22] Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn diw-dimensional representations with sgd. arXiv preprint arXiv:2209. 248,:22209.1466,, 2022.
[MKAA21] Paolo Milanesi, Hachem Kadri, S. Ayache, and Thierry Artières. Implicit regularization in deep tensor factorization. 202/ International Joint Conference Neural Networks
(IJCNN, pages 1-8, 2021.
[MKAS21] Eran Malach, Pritish Kamath, Emmanuel Abbe, and Nathan Srebro. Quantifying the benefit of using differentiable learning over tangent kernels. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 7379-7389. PMLR, 18 24 Jul 2021.
[PA23] Dylan Patel and Afzal Ahmad. Google "we have no moat, and neither does openai", May 2023.
[PF23] Scott Pesme and Nicolas Flammarion. Saddle-to-saddle dynamics in diagonal linear networks. arXiv preprint arXiv:2304.00488 2023.
12

---

[RC20] Noam Razin and Nadav Cohen. Implicit regularization in deep learning may not be explainable by norms. Advances neural information processing systems, 33:21174- 21 187, 2020.
[3K James B Simon, Maksis Knutins, Liu Ziyin, Daniel Geisz, Abraham J Fetterman, and Joshua Albrecht. On the stepwise nature of self--uuprrvied learning. arXiv preprint arXhv:2303.15436 2023.
[SMG13] Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013.
[SS21] Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction. In Neural Information Processing Systems, 2021.
[TGZ+23] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B Hashimoto. Alpaca: A strong, replicable instruction-following model. Stanford Center for Research on Foundation Models. hitps://rrfm. stanford. edu/2023/03/1//appee
html, 2023.
[TVS23] Nadav Timor, Gal Vardi, and Ohad Shamir. Implicit regularization towards rank minimization in relu networks. In International Conference on Algorithmic Learning Theory, pages 1429-1459, PMLR, 2023.
[VSP 17] Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. ArXiv, abs/1706.03762, 2017.
[WGI++19 19] Blake E. Woodworth, Suriya Gunasekar, Lee, Edward Moroshko, Pedro H. Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. ArXiv, abs/2002.00977, 2019.
[YW23] Hao Yu and Jianxin Wu. Compressing transformers: Features are low-rank, but weights not! 2023.
[ZZC+23] Jiawei Zhao, Yifei Zhang, Beidi Chen, Florian Schäfer, and Anima Anandkumar. Inrank: Incremental low-rank learning. arXiv preprint arXiv:2306. / /250, 2023.
13

---

Contents
Introduction
1.1 Related work
1.2 Paper organization
2
Preliminaries
3
Neural networks with diagonal weights
Incremental learning in networks with diagonal weights
4.1
Intuition for incremental learning dynamics
4.1.1 Stage of dynamics
4.1.2
Later stages
4.2
Assumptions for incremental dynamics
5
Experimental results
Discussion
Experimental validation of the assumptions in Theorem 4.1
15
Further experiments on vision and language transformers
20
B.1 SGD-trained transformers
20
B.2 Adam-trained transformers
21
C Proof for dynamics of networks with diagonal parametrization (Theorem 4.1)
25
C.1 Assumptions:
25
C.2 Rescaling time for notational convenience
25
C.3 Simplifying problem without loss of generality
26
C,4 Tracking the sum of the weights
26
C55 Claimed invariants in proof of Theorem C.4
26
C.6 Dynamics from time & to time f++1 (Linear dynamics for (log 1/a)) unrescaled time).
27
C.6.1 Analysis in case where Tx+1 < 00
27
C.6.2 Analysis in case where Th+1 0
30
C.7 Dynamics from time 4 to time & (Nonlinear evolution for 0(1) unrescaled time)
30
C.8 Concluding the proof of Theorem C.4
32
D Technical lemmas
32
D.1 Relating the sum of the weights to the original weights using the conservation law
32
D.2 Sign of gradients on coordinates that leave support
33
D.3 Local lipschitzness and smoothness
33
14

---

A Experimental validation of the assumptions in Theorem 4.1
In Figures 2, and 9, we plot the evolution of the losses, of the entries of WK W diag( uxx )diag(w)Q) and of the entries of Wv W = diag(wy (d)) in the toy task of training an attention head (1) with diagonal weights. The model is trained with SGD the mean- squared error loss on 1000 random samples (X, y). Each random sample has X E R00x50 which sequence of 10 tokens, each of dimension 50, which is distributed as isotropic Gaussian. The label is given by a randomly-generaan In Figures 8, and 9, for {0.1, 0.01, 0.0001, 10-8, 10-16 16 10 32) we plot the evolution of the loss and of the weights when initialized at ((0) = adg, for some random Gaussian 0g Qualitatively, as a 0 we observe that the loss curve and the trajectories of the weights appear to converge to limiting stagewise dynamics, where there are plateaus followed by movement on short time-scales, as predicted by the theory.
teacher model that is also an attention head (1) with diagonal weights.
Validation of Assumption 4.3 (non-degeneracy of dynamics) A - 0, notice that the stages appear to separate and happen at distinct times. Furthermore, the extra technical condition on coordinates i E supp(O*) supp((k-11) in Assumption 4.3 is satisfied since no coordinates ever leave the support of 6
Validation of Assumption 4.4 (stationary points are striet local minima) In Figure 10 we consider the a = 10-32 trajectory, since this is closest to the dynamics in the a - 0 limit. We randomly select several epochs. Since the transitions between stages are vanishing fraction of the total training time, each of these randomly-selecter epochs is likely during a plateau, as we see in the figure. For each epoch perform the following experiment. For each nonnegligible coordinate of the weights (those where the weight is of magnitude greater than the threshold = 10 10-",, we perturb the weights by adding noise of standard deviation 0.05. We then run the training dynamics starting at this perturhed initialization for 1000 epochs. We observe that the training dynamics quickly converge to the original unperturbed initialization, indicating that the weights were close strict local minimum of the loss.
of dynamies) Figure 11 we perform the same experiment as in Figure 10, except that the epochs we select to perturb the weights are those where there is newly-nonnegligible coordinate (less than 10-5 magnitude the previous epoch, and more than 10 in magnitude in this epoch). We find that the nonlinear dynamics are robust and tend to the limiting endpoint even under a random Gaussian perturbation of standard deviation on each of the nonnegligible coordinates, supporting Assumption 4.5,
Validation of Assumption 4.5
15

---

Entries of diaglw,ldlag(we)) for alpha 0.1
Entries of diag(ww))lag(wz)) for alpha 0.01
Diagonal entries W.W. We
Diagonal entries of WoWWk
0.5
0.0
0.0
-1.0
-10
200
400
600
800
1000
200
400
600
800
1000
logt3/slpha)
Entries of dijggw.)ddaggwv.) for alpha 0.0001
Entries of digglw,ldiag(wk). for alpha le 08
Diagonal entree W W
Diogonal entrees of WoWk
0.0
0.0
-0.5
-05
-1.0
time Malphh)
time / logti//ipha)
Entries of dilgglil alpha 1e:16
Entries (lgglw,l for alpha 1e:32
Diagonal entries W.
Diagonal entries WoWe
0.5
0.5
-0.5
-05
-1.0
-1.0
200
400
600
800
1000
200
400
600
800
1000
time / logt1/apphs
time / leg(1/alpha)
Figure 8: Evolution diag( (Q) diag((w)))) entries over rescaled time initializing at various scalings a. Notice that + 0, the training trajectories tend to limiting trajectory. Each line corresponds to diagonal entry of diag ww diagddiaaa))) (w)).
16

---

Entries of d for alpha 0.1
Entries of diaglw,)ldigglw).. for alpha 0.01
Diagonal entries Wole
Diagonal entries of wwWe
time / log(2/epph)
time/ logl1/elphal
Entries dlag(w,.loiaglw). for 0.0001
Entries digggw for alpha
Diagonal entries of WyWoo
Diagonal entries W.
200
400
600
800
1000
200
400
600
800
1000
time / log(2/slphh)
Umme toggira Ipha]
Entries of diagtw.)dlagtwe) for alpha 1e-16
Entries diagtwy)diagWaal for alpha 1e-32
Doggnnll entries of WWo
Diogonal entries W-WW
200
400
600
800
1000
200
400
600
800
1000
log()/llphh)
time/ logt1/s
Figure 9: Evolution of diag( uy )dia)(daa)))) entries in the toy task of learning attention head with diagonal weights. Each line corresponds to the evolution of entry of diag wy )diag (wo) over rescaled time. Each plot corresponds to different initialization magnitude a. Notice that as + 0. training trajectories tend to limiting trajectory.
17

---

Entriies of WWWg perturb time /og(l/a)pha)= = 373.4
Entries of W,Wg perturb time /og/1/alpha) 373.4
0.5
Coardnate value
Goordinete valuee
0.0
0.5
374
375
376
377
378
379
374
375
376
377
378
379
time/ log(L/aipha)
time/ (bgt1/alphh)
Entries of W,Wgg perturb time 1og(\/alpha) 5412
Entries W.WS. perturb time log(2/alpha) 541.2
Coordnate viluee
Coardinate veluee
0.5
0.0
-0.5
-1
541
542
543
544
545
546
547
541
542
543
544
545
546
547
log)1/alpha)
(bol2/olpha)
Entries W.W.g perturo (L/(ppp)))
Entries W.W., perturb time log(2/alpha)
1.0
Coordimate valuem
0.5
Courdnate
0.0
-0.5
-10
752
793
784
780
791
783
log(L/slpha)
(pgl1/alphh)) (pha)
Entries w.w. perturb time [og(1/alphaa) 862.4
Entries WyW? perturb time log(1/alphai 862.4
1.0
Coordinate valum
0.5
Courdiatte
time log(1/slphh)
time logll//llphh (pha)
Entries of W,WWJ, parturb time lag(2//)pha) 942.2
Entries WoWT, perturb time log(2/alpha) 942.2
Coordinate value
Conrdinate value
0.0
-1.0
942
943
944
945
946
967
948
942
943
944
965
9466
947
948
time / log(L/sipha))
time logl1///plll (pha)
Figure 10: Evolution of weights of toy attention model under perturbation, validating Assumption 4.4. At 5 different random times during training. we perturb the nonnegligible weight coordinates and continue train with SGD. The evolution of each line) is compared to the original evolution without perturbation (dashed line). Observe that the training dynamics quickly brings each weight back to the unperturbed weight trajectory, indicating that the weights originally close to strict local minimum.
weights under the initial perturbation (solid
18

---

Entries WyWo. perturb time log(1/alpha) 202.5
Entrigs of WkW perturb log(1/alpha) 202.5
1.0
2.5
Coordimate valuee
2.0
1,5
Coordinate
-0.5
200
220
240
260
280
300
200
220
240
260
280
300
time
WyWg, perturb log(1/alpha)
W,W. perturb log(2/alpha) 321.4
0.00
Coontimate value
Cocrdinate
0.04
-0.06
322.5
325.0
327.5
332.5
335.0
337.5
322.5
325.0
330.0
332.5
335.0
337.5
time (og(2//appha)
log(/aappha)
Entries WyWg, perturb time log(]/slpha) 356.8
Entries of WeWy, perturb time log(2/alpha) 356.8
Canrdinate value
Conninate
0.5
367.5
377.5
365.0
time /log(//lppha
time / log(2/aphaa
Entries of W.WG, perturb time (ag(pha)pha) 405.4
Entries of Wr W.W. perturb time log(2/alpha) 405.4
1-
-0.5
-1.0
406
408
410
412
414
416
418
406
408
410
412
414
416
418
time [1ggp/appha)
time / (og(2/lphha
Entries W.W.. perturb time log([/atohaa) 486.5
Entries of W.W. perturb time log(1/aloha) 486.5
0.5
Coondinate
Courdinate
0.0
-0.3
-1.0
500
525
55.0
575
600
625
650
500
525
55.0
575
600
625
630
time (og(3//pphh)
log(3/wlpha)
Figure 11: Validating Assumption 4.5 with the same experiment as in Figure 10, except that the epochs for the perturbation chosen are those where there newly nonnegligible coordinate. Perturbed dynamics (solid lines) are again robust to perturbation and track the original dynamics (dashed lines).
19

---

B Further experiments on vision and language transformers
The practice of training transformer models often deviates substantially from the assumptions made in our theoretical analysis, and it is priori unclear to what extent gradual rank increase behaviour and low rank bias are manifested in setups more common in practical applications. To gauge the relevancy of our findings we conduct experiments on popular vision and language benchmarks, using algorithms and hyperparameters common the literature. We use the stable rank of matrix W W
smooth approximation of rank. We track the value of the stable rank for the given by [W13 different attention matrices throughout training. Although we do not expect our theoretical results to to hold precisely in practice, we find evidence of gradual increase in stable rank, leading a low rank bias in Figures 12, 13, 15, 17 and 19. In these experiments we use off-the-shelf vision transformers (ViT) [DBK+t2] trained popular vision benchmarks, as well as off-the-shelf GPT-2 trained popular language benchmark. We use no weight decay or dropout in our experiments. All models were initialized using the default initialization scale.
B.1 SGD-trained transformers
CIFAR-10/100 We trained a 6-layer ViT with heads per layer, embedding dimension 512, head dimension 128, and MLP dimension 512 and patch-size 4 for 500 epochs on CIFAR IO/CIFAR 100 with SGD and learning rate 3e-1 and warmup. See Figures 12 and 13. Each run took 2 hours on one A100 GPU.
Layer
4.5-
4.5
4.5-
4.3-
Slate me
2.3
2.2
2.5
15-
1.5
3.2
0
10003
23000
30300
40020
50000
18
30000
20003
33000
40300 50500
1.gg
10000
20000
30000
43630 50500
4.5-
45-
4.5
4.3-
2.3
2.25
225
15:
1.
10002
23000
20100 40000
20000
18
20000
20003
33000
40000
20000
10g
10000
20000 50000 42020 30590
Figure 12: CIFAR-IO, ViT trained with SGD: Stable rank of AW K W (blue) and AW, WT (orange) throughout training. Mean and standard deviation (shaded area) are computed across 8 heads per attention layer.
20

---

Layer
Layer
4.5
255
2.55-
2.-
14
15
13
20300
40002
Layer 3
Layer 4
Layer 5
2M
4.5
mmw
20,05
30-
3.0-
225
1.5
Heration
Figure 13: CIFAR-100, 100, VIT trained with SGD: Stable rank of AW W (blue) and AW, J (orange) throughout training. Mean and standard deviation (shaded area) are computed across heads per attention layer.
B.2 Adam-trained transformers
CIFAR-10/100 For the CIFAR-10/100 datasets we use a VIT with 6 layers, patchsize of 4, heads per self attention layer, an embedding and MLP dimension of 512, and a head dimension of 128. We train the model using the Adam optimizer for 500 epochs with base learning rate of lc-4, a cyclic learning rate decay with linear warmup schedule for 15 epochs and a batchsize of 512. Our results are summarized in Figures 14 and 15 for CIFAR-10 10, and Figures 16 and 17 for CIFAR-100.
layer
layer
layer
int
mit
ist
0.75
0.75
1::00
0.75
1::00
(a)
0.50
50:000
(b)
105000
(c)
R
0.50
12:000
0.7%
0.25
0.25
25
90
75
100
125
125
Wagee
#igenvvliue
KElgenaauu
layer 3
ayer
layer
1.00
1.00
1.00
int
(d)
0.50
1000
5000
(c)
0.50
5000
1000
(f)
0.30
1=5000
1000
0.00
#E /Eigenvalue
Figure 14: CIFAR-10, ViT trained with Adam: normalized spectrum different stages of training, (a) (c) Normalized spectrum of Wx W at initialization and AWR W during training for different attention heads at different layers. (d) (e) equivalent figures for Wy WT.
ImageNet For ImageNet, use the VIT-Base/16 from [DBK+20] 20] trained with Adam for 360 epochs with base learning rate of 3e-3, a cyclic learning rate decay with a linear warmup schedule for 15 epochs and batchsize of 4096. Our results are summarized in Figures 18 and 19 for ImageNet.
21

---

Layer
Layer
2.0
Stbblrek
Stabler tank
2.0
2.0
(c)
7.5-
1.5
2
1.0
20000
40/00
20000
40006
100CH JU100 50000 ROUDO 50000 Iterations
Fe Terwsons
Iteretions
Layer
Layer
Layer
Stable rark
Stable
(d)
10000 20000 30000 40000 50000
10000 20000 30000 40000 50000
10000 20000 30:00 40000 20000
Baraliers
Butaliono
Figure 15: CIFAR-10, ViT trained with Adam: Stable rank of AWR K W (blue) and AW, W (red) throughout training, Mean and standard deviation (shaded area) are computed across 8 heads per attention layer.
layer
layer
layer
1.00
Intt
Init
Init
0.75
0.25
1:100
0.75
(a)
0.50
(c)
0.25
0.25
0.25
25
%
75
100
125
#Bgerwake
AEignvsilue
AEigenwalue
layer 3
ayer
layer 5
1.00
1.00
1.00
5-1000
1000
b-1000
(d)
0.50
5-5000
(e)
0.50
5000
(f)
0.30
5-5000
0.00
0.00
0:00
dEgerrvalue
Figure 16: CIFAR-100, VIT trained with Adam: normalized spectrum at different stages of training. (a) (c) Normalized spectrum of WK W at initialization and AWR WT during training for different attention heads different layers. (d) (c) equivalent figures for Wy Wy W..
Layer
Layer
Layer 2
Saade 141k
Stable cank i
(a)
(c)
2
Iterations
Layer 3
Layer
Layer 5
10
Stable
(d)
20000 Earrations
40300
10000 20000
20000 Burrations
Figure 17: CIFAR-100, ViT trained with Adam: Stable rank of AWK W Q (bluc) and A AW, W (red) throughout training. Mean and standard deviation (shaded area) are computed across heads per attention layer.
22

---

layer
layer
layer
0.75
t-400
0.75
t:400
0.75
t:400
(a)
Fower
00.00
t:400c
tu4000
(c)
30.50
t=400c
0.23
0.00
0.00
0.00
25
50
75
300
125
25
50
75
100
125
25
50
75
200
125
EBgerwake
aElgorvelue
ABgenaake
layer
layer 6
layer 10
1.00
1.00
1.00
Int
Init
I11t
0.75
0.25
b-400
0:75
T-400
(d)
Rowe:
0.50
[-24000
0.50
1-20000
0.30
24000
0.25
0.25
0:25
25
50
75
1.00
125
25
50
70
1.00
125
#Egerwale
PEigensulue
#Eigonvulus
Figure 18: ImageNet, VIT trained with Adam: normalized spectrum at different stages of training. (a)
(c) Normalized spectrum of Wx W at initialization and AW K W during training for different attention heads at different layers. (d) (c) equivalent figures for Wv Wo-
Layer O
Layer
Layer 2
Stade nank
7.5
(a)
(c)
5.0
Iterations
Layer 3
Layer
Layer 5
10
Stbble mank
10
10
(d)
10000 30000 40000
2000o Rtaraliorn
10000 30000 40000 Hlraiions
Ble Starations
Layer
Layer
Layer 8
Statbe 1410
Scaidee mikk
Statle
40
20
10000 Eterations
10000 meration
10300 30000 Etarations
Layer 9
Layer 10
Layer 11
Babke man
10
(k)
stabe
20
10900 20030 90000 40000
10000 5000n
VIDRO IDerotinns
19300 20800 10h00 40000
Rtarators
Aterators
Figure 19: ImageNet, VIT trained with Adam: Stable rank of AW WT (blue) and A AW, W (red)
throughout training. Mean and standard deviation (shaded area) computed across 12 heads per attention layer.
23

---

Wikitext-103 The gradual rank increase phenomenon also occurs in the NLP setting with language transformers. We trained GPT-2 on Wikitext- 103 using the HuggingFace training script with Adam learning rate 3e-4, per-GPU batch-size 8, and block-length 256. We trained for 3 epochs on 2 A100 GPUs, which took 12 hours. See Figure 20.
Layer
20032
40920
00600
30005
20000
40300
63000
10000
20005
43000
51000
80003
floradinn
20032
40020
500c0 20002
20020
40100
60200
20020
0
20603
42000
5:000
80002
Layer
Layer
Layer
20002
40020
50000
20002
20020
40300
90200
30020
0
20003
40000
50000
80002
Layer 10
31
20000
40020
50000
55002
20320
40300 90200
10000
20103
42000 52000
80003
Figure 20: Wikitext- 103, GPT-2 trained with Adam: Stable rank of AW, WT and WW K versus training iteration. Stable rank of the perturbation increases gradually, but remains small throughout training.
24

---

C Proof for dynamics of networks with diagonal parametrization (Theorem 4.1)
C.1
Assumptions
Recall we have defined 0º ....... . ...... ....... E 2 as the sequence of weights such that 0 = 0 and 0
defined inductively as follows. Consider the dynamics of (t, R initialized at a/. (0,
@h + (ee sgn(g, (g. ((*) eee;,) and evolving according the gradient flow (t.) )+) dt
We assume that there is a limiting point @ of these dynamics as e taken small and the time is taken large:
lim lim 26 (t, 0
Under the above assumption that this sequence 0" ..... 0 .... is well-defined, we can derive a useful property of it for free. Namely, the conservation law (5) implies that @ u ④ is preserved. It follows that for each l have that 0 (u w²) satisfies |e*1 | entrywise, In other words, there is s E {+1,,11 1 satisfying
0h (u", t un") R22p
We also abuse notation and write supp = supp(u*) [p]. since the support of Qh on the first p coordinates matches its support on the last coordinates.
Having fixed this notation, now recall the main assumptions of the theorem.
Assumption C.1 (Nondegeneracy of dynamics in part (A); Assumption 4.3). The initialization
satisfies [u,(0) * [0)[ for all i. For stage k. either Tk+1 00 or there is unique minimizer é to min; (ix))) in (11). Finally, for all i E supp(0 supp(O") we have 9: (0*) / 0.
Assumption C.2 (Stationary points are strict local minima; Assumption 4.4). For stage k, there exist 5x > and Ca > 0 such that for á E B(u*, "
5) supported on supp(u*), we have
C(ū, C ū) 2 cxllu* =
Assumption C.3 Moise-robustness of dynamics in part (B); Assumption 4.5). For stage k either T 0o or the following holds. For any c O, there are > 0 and 7 Rso - R such that the following holds. For any E B(u*, 5) n IRS supported on supp(ū) S. supp(u") U {ix),
there exists unique solution vb oo) + Rº of the gradient flow di de
- VoL(v) initialized at
(0) ū, s k+1 @ ū) and times / ((ū.))
[[)(() g++ <e.
C.2 Rescaling time for notational convenience
For ease of notation, we rescale time
(0) = (o)) v. = axx(1)
dua log(1/0))) 2 g(u,, va);
dva
log(1 /a)ua 2 g(un, va).
(14)
dt
dt
We also define
(t) = (ta (t) E 2
Because of this time-rescaling, equivalently state Theorem 4.1 as:
Theorem C.4 (Restatement Theorem 4.1). Let K Zzo be such that Assumptions 4.3 4.4 hold for all & < K and Assumption 4.5 holds for all k < K. Then for any k < K and time t E (T Tet1) the following holds. There is (l)()) 0 such that for all a on, there exists a unique solution 0,, [0, - T&P to the gradient flow (14) and
lim (a(1) + 0¹.
where at each stage |supp(u*) supp(u 1
For shorthand, we also write
S = supp(u² and St [p] supppuu)..
25

---

C.3
Simplifying problem without loss of generality
For each coordinate i [p] we have (0)) + [a),( by the non-degeneracy Assumption 4.3.
So we can assume (a,((0) > (o))(0)] without loss of generality. Furthermore, we can assume the entrywise inequality
Un (0) 0
by otherwise training weights (t), (t) initialized at ú. (0) = sgn( (0) (0) and öa (0) =
sgn(v..0)))er v. (0), as (t) @ (t) (t) @ (t) all times.
Since (1) - (1) - (0) - _ (0) by the conservation law (5), it holds that [ua,((()
[t), throughout. So by continuity
(t) 0
throughout training.
C.4 Tracking the sum of the weights
We define
(1) (1) + (1).
The reason for this definition is that during training we have
dw = (og(1/a)wa, 3 g(o.),
(15)
dt
Notice that since that we have assumed :a. (0) > [0)] for each E [p] we have w. (0) > entrywise. So, by (15) for all 0,
wa (t) > 0.
suffices to track (() because we can relate the log-scale magnitude of (4) to the magnitudes of the corresponding coordinates in u. (t) and va (t) see technical Lemmas D.1 D.2 and D.3.
C.5 Claimed invariants in proof Theorem C.4
In order to prove Theorem C.4, we consider any gradient flow 0. [0, - R solving (14) where T* E (TK, 1)) For now, we focus only on proving properties this gradient flow, and defer its existence and uniqueness to Section C.8.
We show the following invariants inductively on the stage k. For any >> any stage k < K, there is ag == a& (c) > 0 such that for all a < ag the following holds. There are times to := (a e) and 4-1 = f (a, e), such that
EN [Ts e, Th + c),
(16)
- e, Txt1 + +c),
if < 00 if T Th+11 00
fe41 E {(T)).
(17)
and the weights approximate the greedy limit for all times E [Tx L++1l
(t. 0+ < E,
(18)
and the weights at times Ta and 4++1 are correctly estimated by the incremental learning dynamics on the logarithmic-scale
log, (wa (Tx)) - b²1 <6
(19)
and if Tx41 0o then
|| logg (w.((+)))) k++ <e.
(20)
Base case l = 0: Take Fa(a, c) = 0. Then statement (16) holds since Tú = 0. Notice that as a + we have that u (0), v0) - 0 u², and also log, wa (a) 1 = ba So statement (19) follows if we take ag small enough. In Section C.6 we show how to construct time t such that (18) and (20) hold.
Inductive step: Suppose that (16), (18), (19) and (20) hold for some iteration K. We prove them for iteration & + 1. In Section C.7 we construct time fg. In Section C.6 construct time tx+1.
26

---

C.6 Dynamics from time Eg to time k (Linear dynamics for ((log(1/a)) unrescaled time)
Let S K, and suppose that we know that for any Eg 0, there is ax({2) such that for all 0 < a ag. there is a time & = (a, Ex) satisfying
[T - [x < Ex
[10((Fx) 0%1 < Ek
1loo. (fx)) b² < Ék
C.6.1 Analysis in case where Tk+1 00
Consider first the case where Ta+1 < oc. We show that, for any cx+1 > 0, there is P > 0
such that for all p Pk++1 (Ek+1) there is ax+1 (p. Cx+1) 0 such that for all a < ax-1 there
is a time 1++1 4++1 (a, 1x++11 satisfying
Tk+11 1++1 < Ck+1
(21)
10 (t) 0 < 6x41 for all E (bktt4))
(22)
log. (a((t++1) bk+111 < 6x41
(23)
Uaii (Lx+1) E [p,3p],
(24)
sgn(Varte (Lx+1)) = 8i k+1
(25)
For any p, let Es = p (4p) and choose Es = Ea, Ex). Then define
L+1 La, Ex+1)
(26)
=infft E [7x. 00): (|ua, (t) us (fx)) || (t) Va (f&) > 4p).
Now we show that the weights 0, cannot move much from time & to L.41. The argument uses the local Lipschitzness of the loss L (from technical Lemma D.7), and the strictness of @ as stationary point (from Assumption 4.4).
Lemma C.5 (Stability of active variables during part (A) of dynamics). There is pk41 small enough and 04-1 (p) small enough depending p.such that for all < pk+1 and a < ('++1 and all t [k, 1x41.
0t, 1 < < == max 18 pK nk/k))))
(27)
where eu is the strict-minimum constant from Assumption 4.4 and K Rg is the Lipschitzness constant from Lemma D.7 for the ball of radius R4 ||@*||
Proof. Assume by contradiction that (27) is violated at some time t < tx+1- Let us choose the first such time
=" inffft [Zx, (e+1) [t") u²1l [1a)(t)) s + u p'}.
Define θ = =(ūz (i, by
(I*),
i E Sk
0, (aa(
E St. 0, Φ St
i; {0:0000
and
0.
E Sp
By the definition of f this satisfies
û un (l" || = || (1')]l S 4p+ +
| (7+) S 4p + ft < <5p,
|| vo (1*) || = || (()) < 4p || (Tx)]l S 4p + fs < <5p.
Also
||0 8 3 r *|| = Uu, (t*) - + Va.s, (1") - - 3 2 2 p' 10p 2 1//2.
27

---

Using (a) the strict minimum Assumption 4.4 with constant since ||0 - 1|-0 S p' and we take p' small enough,
(0) ((),(t")) 2 L()) - 4pK R 2 c( - 4pKR. + cx(2'²²
16
2 [tx)(tt)) - (4p Ex) cx(d)) ((Oa(ta))
16
This is contradiction because L is nondecreasing along the gradient flow.
Lemma C.6 (Log-scale approximation is correct during part (A)). There are functions Pk+1 (Sx+1) O and 0&+ (p, Cx+1) > such that for all p < Pk+1 and < Cx411 and for all E (x41) we have for constant C' depending on k,
(1w - b² + (1 [0)g( < PEX-1 + CCp [2)-
(28)
Furthermore, for all i E St and (Fx, (x+1) we have
sgu(g; (t)( = (g
(29)
Proof. By Lemma C.5 and Lemma D.7, there is constant C depending on 0 such that for all t (41).
lg)( 9(0*)) SCp'
For shorthand, write g(0") = g( + Cp'1 and g(0²) = g(0") - Cp'1. Since wa (t) > entrywise have assumed without loss of generality (see Section C.3), have the following entrywise inequalities
(t) ((),(t)) (t g(0") o (t)...
(30)
di Since the dynamics are given by dwa = log(1/a)g o Wa
log(1/a)g(0) Wa (T& (t(-ba)) log(1/a)g(0*) < w. (t) S wa (Fx) boo) Taking the logarithms with base E 10,
(1- (l)())((u) S
(T&)) loga (w) S (1- (())gg)))
The bound (28) follows since || log, (w, (Tx)) ba E < 15*+1
Finally, the claim (29) follows from (30) since sgn(g(O")) = sgn(g(0")) sgn(g((")) if we take P small enough.
First, we show that the weights must move significantly by time roughly Tk+1. This is because of the contribution of coordinate ig.
Lemma C.7 (+++1 is not much larger than Tk+1). Suppose that Tkt1 < cc. Then there are P&+1(ER+1) > O and holds.
such that for all p < pk+1 and a < ak+1 the following
t < Tk+1 + S641
Proof. Assume by contradiction that l441 < Tet1 + Ek+1 For all times E [T2, min(L++1: Tet1 1141)., by Lemma C.6,
logg (Wa, (t)) - bi + (t -t- - Fg ) (0))
Since know (1a (TR+1 - [x)) < E and by¹ (1
(0, it follows that (Wa.i. (Te+1 + Ex+1)) & (-1g,, (e0)) )( - (x+1) 19,, (0 )(<(+) - Ek+1)) + O(vp). that 1g,, (0*) )[{{+)](++1 Ek+1) + ((pp) 8 0 for some 8 > 0
By taking P small enough, that is independent of a, so
log (Wa.te (Tx+1 + (x+1)) & (-
So (ua,i (Tk+1 + Ex+1)l > by Lemma D.2. But by the construction of 4+1 this means that 4-1 Te41 + Ex+1-
D
28

---

Next, we show that until time 1 none of the coordinates in Sj move significantly, with the possible exception of coordinate ig.
Lemma C.8 (No coordinates in S \. {ix} move significantly during part (A)). Suppose Tk+1 < co. Then there are Pk++1 (5x+1) > O and C++1 (P. Ex++1) > O such that for all Pk+1 and a < Ok+1 the following holds. There is a constant > depending on such that for all E Sg \ {ig) and t [
(ua,((t) Ua (f),:(t))ll - va,i(fz) < a + E
Proof. The previous lemma combined with the inductive hypothesis gives
4++ E < (ta)
(fx) + 2G441 \ {ix)-
We analyze the movement of each coordinate i S% {is) by breaking into two cases:
Coordinate i of is such that bf (0, By Assumption 4.3, there is unique winning coordinate so byt (0k) E (c,2-c) c) for some constant >> 0 for all T E [0, fk+1 fu] S [0, (14) + 16+1 By Lemma C.6, loga (t), E (-c/2, 2 c/2) for all times E [F,, 1 1- So by Lemma D.I, (t)l. (Va,(()) S a/44
Coordinate i f ig such that 6f 0. By Lemma D.4, we must be in the corner case where i 5k-1 O Si (i.e., the coordinate was active in the previous stage but was dropped from the support in this stage).
By Lemma D.4, since U/ (g (0"))) < O for all t E (fx, tx11).
we have 9:(0*) 0. By Lemma C.6, this means
We break the analysis into two parts. Since 6/ = 0. the sign is of +1. The inductive hypothesis (t& 0²1l < Eg implies that (tx 211 < E and |vo, (Fx) 2 < & = sgn(va,((Te))
+1. Now let 8' min inf > La (a,(()) = 0}). Since ua,(()) > Va,i without loss of generality (see Section C.3), we have sgn(ua,:()) = sgn("a, = +1 for all t E [Fx,
For small enough 6. this means that sgn("n,i
So
du., (1) del,.(()
O for all t [Ex, So, for any E [Ex, t").
(t). - (fx)( (t,, - Co, (fx.) < εk
Also, since (wa,(t") = 1, by Lemma C.6 we have ts > c > 0 for some constant c independent of a. So for all t E [t", tx+1] we have by Tg: -0: (0*) constant > 0. So (t)), (t)) I o/4 for all t [t", 4 1. The conclusion follows by triangle inequality.
E (c, 2 - c) for some
Coordinate i + is such that 6/ 2. The analysis is analogous to the case 6% = 0, except that we have sy = instead and (e*)) by Lemma D.4.
Finally, use this conclude that t. Rs Tk+1 and that the weights coordinate is are the only weights that change significantly, and by an amount approximately p.
Lemma C.9 (Coordinate in wins the part (A) race at time L++1 RS Tx+11). Suppose that Te+1 < co. Then there are p&++ (5x+1) 0 and 0x+1 (p, (x+1) > O such that for all p Pk+t and < Ox+1 the following holds.
Lx+1 Tetil < Ek+1
Ua.in (L++1) E [p,
sgn(Va in (tr= )) s/o k+1
29

---

Proof. Let us analyze the case that 6,; E (0, 2). Notice that be+1 = b, - (ixi)) (#*)) {0,2) and that if 0 then 9t, (0%) > 0 and if it is then b gú. (0 < O. So by Lemma C.6, for all times t min((x++++ f+ - Ex+1)], we have Wa,ir (t) E (c, 2 - c) for some > So for
small enough a by Lemma D. Mu..i (t)l, Ma, (t) I a0/2, Combining this with Lemma C.8, we
see that for & E [Eg, min(f++ T k+1 Cx+1)] we have
uL (1) a (7x |I + (l) - una (Tx)ll < 2(a + +.L) < P,
for small enough a. So by definition of fx+1 we must have 4+11 Th+1 Sx+1- Combined with Lemma C.7, we conclude that [T&+11 Ex+il il < LK+1, which is the first claim of the lemma. Furthermore, by Lemma C.8,
Σ [ta,(( - uaxx((x) + (Va,((e411) va, S 2p(a + +x)) < p/2,
so by definition of Lo41 and triangle inequality we have |ua,i, (4x+1)) + |Va,i (L++1) 1) 4p p/2 7p/2. Also, since ua (LR+1) v2 (1x+1) = O(a²) we have uaae (tx+1) E [p. Finally,
8/2 = 2, then S1 k+11 = -1 and log, (Wa (lx+1)) > 1.5 so sgn("a (t)) < 0 by Lemma D.3; in = and (wa.i, (L++1) < 0.5 so sgn(Vo (4++1) 0.
analogously, i k++ = O, we have k+1
The case b/. E 20 can be proved similarly to the analysis in Lemma C.8, where one shows that during the first period of time the magnitudes of Uu, (1) and Jos, (1) decrease, until the sign of vi flips and they once again increase.
We have shown the claims (21), (22), (23) (24), and (25) for the time L++1- In fact, if we let 111 E [Eg, be the first time such that Ua (t) = p we still have (21), (22), (23) and (25) by the same analysis as above, and (24) can be replaced with the slightly more convenient
Ma (1e+1) = p.
C.6.2 Analysis in case where Tk+1 = 00
In this case that T, we just have to show that the weights remain close to 0h, We show that for any Ed+1 > 0, there is 2x+1 (E6+1) > 0 such that for all a < 2++1 and times t E [Tk + $6.1.
(1) o*1 < Sx+1-
We can use Lemmas C.5 and C.6, which were developed for the case of Tk+1 < co, but still hold for Tk+1 co. Lemma C.5 guarantees that the weights do not move much until time 4--1- and only need show that Lx+1 2 T when we take p small enough. For this, observe that g 0 for all i & Sk. because otherwise T < oc. Therefore Lemma C.6 guarantees that until time min (11) all weights are close to the original on the logarithmic scale. Namely,
loga (w. (t)) ba1l < P56+1 + Cp'(T* - [2)
Furthermore, by the non-degeneracy Assumption 4.3 we know that 6/t E (0,2) for all i & Sy by Lemma D.4. So if we take p small enough and @ small enough, we must have that fx+1 2 T
C.7 Dynamics from time 4. to time Ta (Nonlinear evolution for O(1) unrescaled time)
Suppose that we know for some & < K that for any 6g 0, there is pk((x) > 0 such that for all p 0h there is a (p,sa) > 0 such that for all a < (g there is a time 4 = (ar p,Ex) satisfying
Th txl
< Ek
(31)
110,
0 -1 < E
(32)
loga (wa (L+)) - b²1 < Ek
(33)
Ua (tx) = P
(34)
sgn Sgn(Vaiian1 (tx)) = 8 Pio-1"
(35)
30

---

Now we will show that for any Es > O, there is ā = āx((+)) a time & (a, Ex) satisfying
such that for all 0 < < ag. there
[Tk- tk Ek
(36)
110 (4x) 0* EK
(37)
|| log, (wo (fx))
b" < Ek
(38)
We give the construction for &g. For any desired accuracy E >> O in this stage, we will construct an accuracy E & (&+) = / > We will also construct p p(Ex) > O which sufficiently small,
and we will construct an cutoff for equal to Og 0 (Ex) x) > 0 which satisfies oş < @ (0: fg). The values for these parameters S and p and ag will be chosen in the following lemma, and will depend only on Ex-
Lemma C.10 (New local minimum reached in time O log(1/)))) For any Ek > 0, we can choose 0g (Tx) > small enough that, for any 0 < < 0g. there is & (a, Cx) for which conditions (36) to (38) hold.
Furthermore, there is constant C" independent of a such that (0.(1) E [1/C", C at all times t E [ tx].
Proof. Let 4 (a Ex) be given by the induction. Let us compare the dynamics starting at 0, (Lx) with the dynamies starting at ((L)) (ū(L+), o(4)) which is given by
(tx):
E Sk U {ik-1}
(ta) i Sk- U (ia-1}
ūz (4x)
and (.(bx) 10
0.
otherwise
otherwise
and run with
dθ logg((/a))V
dt
By Assumption 4.5 we know there exists a unique solution 0 [L: 00) - RR as long as we take fx small enough because (otp) = S U {ix-1} and [tx, - 0h-11| E.- Furthermore, by Assumption 4.5 if we take E small enough there must be time == T(EL, p) < such that
||@(t) 8* || e for t 2 t + T//og(L//a)
(39)
Define
f = Lx + T/ lo///////1/)))
So for a small enough, [Tp - Egl < 2e, < εR, proving (36).
We now compare (Fx) with (, and show that if we take small enough, then the dynamics of closely match the dynamics of (.(() for times 4. + O(1/ log(1/a)). The argument uses Gronwall's inequality. Let 1. inf(t 4x 110(1*) (t)l 1/3). For times E (4x by Lemma D.7 we have
a ((t) 00)
- (t)l = log(1/a)IVec((11) - &
- Koto log(L/a)]I((1) - 0, (t)
dt
where Kace) the smoothness constant from Lemma D.7. Note that since [IO())] < oc for large enough by (39), the trajectory of Ô must lie in compact set. Therefore, there must be a finite set of times $m. m E [t., (²) such that Le) 1/2) S U B(O(s)) 3/4). So letting C' max." Ko(e.) < oc for all times t [t,. we have
d
- (0(1) - 0.1l S Clog(1/a)l0(()) -(.
di
By Gronwall's inequality. for all times E (4,(") 1*),
(θ(()) - (t)l < [l0(t,) - 0 (tx
-x))
31

---

We know from Lemma C.8 that there is constant > 0 such that for any small enough 0 < a < ag. such that
[0(L)) - (tx))) a
we take a small enough that a" cxp(CT) < f < 1/3, we must have t" > t + T//log(///) and prove (37)
@ (c 22 (10(Fa) (rx < <
remains to show that (38) satisfied. Since [[((t) 0a (t)] < 1 for all E [4 Fx). it holds that the trajectory of (t) lies in compact set. So by Lemma D.7 we have (g) < C' for some (ww.
c'I (t),, we must have (W)l/(/)l (+))) E [1/C", C"] for some constant C" independent of a and all I E L4 [x]. Therefore, (38) follows from (33). A similar argument shows that (1)/0 (4,)
constant C' at all times t E [Ex fu]. Since log(i/a] dyc..
[1/""
C.8 Concluding the proof of Theorem C.4
We have shown that Theorem 4.1 true for solutions 0, [ R to the gradient flow, where T. E (Tk, Tx+1). To establish Theorem C.4 it remains only to show that for any T. E (TK, Tk+1) and small enough a such solution the gradient flow exists and is unique. To see this, note that in the inductive proof of the invariants we construct sequence of times 0 = to S 144
T., where we guarantee that any gradient flow solution B(0*, 0.: 10, Rp satisfies 0a E UkE B(O .10,,) for all t E Uke(o..... [f tx+ And also for & E (o.......KK11 Ltx. 1x+1). we have (t) E B(0, Cf for some constant C% independent of a by Lemma C.10. So 0₂ (t) B(0, Cx) for some constant Cк at all times E [o, T*). By the compact set BO0 Cx). So 0. [0, + I2 exists and is unique by the Cauchy- theorem.
S S ... S [ S LK+1
Lemma D.7, the loss gradient VoL(0) = (v @ g u o g()) is Lipschitz-continune
D Technical lemmas
D.1 Relating the sum of the weights to the original weights using the conservation law
Lemma D.I. If for some constant 0 < < 1 we have log, (wa,(t)) (c, c). then for small
enough a
max([ux,
(t)) (M),,,, < 0/2
Proof. Let (l) (1) (1). By the conservation law (5), wo.,())...(0
(O),(0) (0)2 (0). By the non-degeneracy of initialization (Assumption 4.3), the right-hand-side is O(a²). So if log, (wa,(tt) E (c,2 c) then for small enough we have loya (la (()) E (3c/4,2 3c/4). So [(a,(()] < [w,,((!) wa,(()) S e/2 and [ta,s(t) < (t) (t) S 1/2
Lemma D.2. for some constant 0 < c we have log, (t)) & (-c, c), then for small enough a,
| (1) > 1.
Proof. Define w. = u. Uo as in the proof of Lemma D.I. If log, (Wa,i (1)) -c then log, (/ (t))) > 2 - c/2 for small enough a, so ut(t) > a c a 2-0/2 > 1. Similarly, loga (wa, (1)) > 2 + e then log (t)), < -c/2 so |ui(a) > a cf2 - -a 3+tex11 1.
Lemma D.3. ffr some constant e > O, there is small enough ar such that if we have loga (wa, 1 + then sgn(va,.(t)) 0. Otherwise, if log, (t))
then sgn (gn(va,s(t) > 0.
32

---

Proof. Follows from va
2
] (wa wa) Recall that (a) > 0 and notice that wa(() > 0.
In the first case, Wa,((t) (t) a al and (t) a1-o/2 In the latter case (a, > a Q and
(o,((l) < att//2
D.2 Sign of gradients on coordinates that leave support
Lemma D.4. For any k and E St. bf {0,2) then we must have E supp(u supp(u"), and we must have gi(u²) if b and gi(0)) by 2. In particular,
(a(4a) > for all k.
Proof. This by induction on k and using the non-degeneracy Assumption 4.3.
D.3 Local lipschitzness and smoothness
We provide several technical lemmas on the local Lipschitzness and smoothness of E, h, and g. Lemma D.5. The function E(y, -) is locally Lipschitz and smooth in its second argument: for any R > 0, there exists Ka such that for any Ç.C' 5' E B(0, R) |e(y. C) f(y, (') Kull< 5''l [DE(y, 5) Di 4')ll -
almost surely over y. Here D E idouu is the derivative in the second argument.
Proof. Since f is continuously twice-differentiable for each y IRd, Rdout there is Ky.c < such that for all y E Byy s/)) and Ç' E B(C, k/ we have
DDDE ,')Il S Ky.c and [1 S Ky.s
where DE and e denote the first and second derivative in the second argument. So for all such y' Bly 1/Ky,s) and ç' E BCC, 1/Ky,c) we have Mely', e(y'; ("") Ky.cll'' 5"1l and - S Ky,cll'' -"
Cover the set {wy ||vl|| < C, I[ < R} with the balls Ly By compactness, there finite subcover (y), Ç))...,, (Yr: (r), so we can take Kn max Ky.c. < oo and the lemma holds since ||ыl| C' almost surely by Assumption 2.1.
Lemma D.6. The function h(x;;) is locally bounded, Lipschitz and smooth in its second argument: for any R > O there exists KR n such that for any ub, a/ B(0, R).
|1([;; ub)1| < K
[1(2;; x/) h(2: ab || < Kxll 4/5
[Dh((;; wb) Dh(æ; ab S KRllb - bb11.
almost surely over æ. Here Dh(x,-) Rdout RP the derivative in the second argument.
Proof. Analogous to proof of Lemma D.5, using continuous twice-differentiablliyy of and bound- edness of ||=]|]
O
Lemma D.7 (Local Lipschitzness of loss and loss derivative). When 0 (u. v) R and JNN(ar (x; 0) = li(;;; u @ u) the following holds for g0)) defined (4). For any R > there exists
K co such that for any 0, O' E B(0, Kn).
ll)(0) g s K R 0 0'|
( (n(( < Ko 0 8'1l
[C(0) - C(0') < KR 0'|
Proof. Let 0 = v), 0' (u', v'). This follows immediately from the local Lipschitzness and smoothness of and in Lemmas D.5 and D.6, as well as
(l)(0) (00')l F.,,y (Dh v) De(y, Max T)) Dh(a; n' 3 oo) De(y, h( ul ② x'))) T]I
33


Reviews:
[{"review_id": "MWyHVWcMFS", "reviewer": "Reviewer_ryvs", "summary": "The article \"Transformers learn through gradual rank increase\" considers the dynamics of training for the neural networks models with attention mechanism. The authors relate the training dynamics to a particular type of gradient flow. They show under 3 important assumptions:\ni. diagonal weight matrices\nii. initialization is small\niii. only one coordinate is in \"active\" regime at a time\nthat dynamics occur in discrete stages:\n(1) during most of each stage, the loss plateaus because the weights remain close to a saddle point\n(2) at the end, the saddle point is quickly escaped and the rank of the weights increases by at most one\n\nThe developed theory is illustrated with a series of experiments.", "questions": "1. Is the paper more about attention in general than specifically about transformers?\n\n2. What will happen if some weights activate simultaneously? How the theory would be affected?", "limitations": "Limitations are adequately discussed", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Nice and clean results that show us the dynamics of neural networks with attention mechanism\n\n2. Well connected with the previous research on training dynamics\n\n3. Paper is well written", "weaknesses": "Assumptions are quite restrictive. It would be nice to relax the assumptions, especially assumption #3.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive review of our paper and your questions about the theory. We are glad that you found the results interesting and found our paper to be well-written. Please find below our response to your questions:\n\n* Q1: \"Is the paper more about attention in general than specifically about transformers?\"\n    * Our paper has two novel contributions: an experimental contribution and a theoretical contribution. \n        * The experimental contribution is that we observe that the difference between trained and initial weights grows in rank. The experiments are conducted on ViT and GPT2 transformers (which have attention layers inside). So the experiments are about transformers in practice.\n        * Our theory on the training dynamics applies generally to nonlinear networks with diagonal parameterization. Transformers with diagonal weights are the important special case on which we focus. Our theory applies to these transformers if the attention layers are the only trainable parameters (see Example 3.2 in our paper for details). So the theory applies to transformers of any depth, where only the attention layers are being trained.\n* Q2: “What will happen if some weights activate simultaneously? How the theory would be affected?“\n    * If more than one weight can activate simultaneously, then it becomes more burdensome to write down the dynamics in Algorithm 1. Nevertheless, we believe that we might be able to prove the multiple-weight-activation dynamics are valid if we also modify Assumption 4.4 appropriately. Thankfully, multiple weights do not seem to activate at exactly the same time in practice (see our experiments in Appendix C), so this non-degeneracy assumption seems to be a valid simplifying assumption.\n* Weakness: “Assumptions are quite restrictive. It would be nice to relax the assumptions, especially assumption #3.”\n    * We wholeheartedly agree that it would be nice to relax the assumptions 4.2, 4.3, and 4.4, but we also recall that they are validated on our toy model in Appendix C. Assumption 4.2 is mainly the assumption that two weights do not activate simultaneously, so see our answer above for how we believe this could be relaxed. We do not know how to relax Assumptions 4.3 and 4.4, but strong assumptions along these lines seem inevitable because of the generality of the theorem. In the revision, following the suggestion of reviewer yLKb, we will move some figures from Appendix C to the main text because these help illustrate the assumptions and our experimental verifications of them on the toy attention model.\n\n"}, {"review_id": "7XXcNMG0oA", "reviewer": "Reviewer_ZvU1", "summary": "In this paper, the authors study the learning dynamics of transformers and argued that the difference between weights and their initial values increase in rank as the training progresses. Under small initialization, smoothness, non-degeneracy. convergence and robustness assumptions, the authors proof the incremental dynamics in diagonal-weighted transformers. Empirical results on toy models and VITs shows similar rank-progression dynamics on learned perturbations.", "questions": "* Why is the non-toy experiments only conducted on VIT? Is there any reason to choose only the vision transformers but not regular NLP transformers?\n\n* How small is a \"small initialization\" in practice?\n\n* Dynamic (3) seems to hold only for gradient descent training, how does it apply to other optimizers like SGD/Adam?\n\n* The learning dynamics theories in this paper is proved in continuous setting (learning rate $\\rightarrow$ 0, learning steps $\\rightarrow\\infty$). The real discretized GD training is just a forward finite difference approximation of the dynamics (3). Under what condition does the result hold in the later case?", "limitations": "The authors discussed about the limitations that the theory requires diagonal weights and small initialization. However, it is not discussed whether dynamics (3) could cover optimizers in practical cases. I don't think this paper would have potential negative societal impact.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "The authors make an interesting observation that the learned perturbations in transformers are of low rank and exhibit an rank-increasing dynamics. Theoretical analysis is given under a simplified setting and the empirical results seem to support the observation.\nI do think this is a topic worth exploring as the findings could support the recently popular low ranking fine-tuning studies and potentially help identifying more efficient low-rank training methods through better understanding the learning dynamics.", "weaknesses": "* The assumptions are very strong such that it is not clear if the theoretical results has real implications in practical cases. Not only does the theory only hold for attention layer only, diagonal weighted transformers, it also need strong assumptions on convergence, robustness. \n\n* The presentation of the paper could be improved. Section 4.1.1 and 4.1.2 are kind of confusing and does not help in understanding the reasoning behind the rank increasing dynamics. In addition, the author uses a lot of vague expressions like \"very small\", \"non-negligible\", \"good approximation\", when not providing quantitative reasoning like error/magnitude bounds.\n\n* Dynamic (3) seems to hold only for gradient descent training, which is hardly used in practice. It's not clear if the analysis holds for optimizers like SGD/Adam.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback. We answer all your questions and address all your comments on weaknesses below. We hope that our responses will be sufficient to clear up any concerns and confusion. We would be happy to answer any more questions if you have them.\n\n* Q: On NLP transformers\n    * Thanks for the suggestion. We have **added experiments on NLP transformers** (GPT-2 trained on Wikitext) which show the exact same behavior; this is to be expected since the architecture is quite similar; see attached document.\n* Q: “How small is a ”small initialization“ in practice?”\n    * We observe gradual rank increase dynamics at the initialization scales used in standard practice (see our experiments on ViTs and on GPT-2). Furthermore, Figures 2 and 4 show that the gradual rank increase dynamics are even more pronounced if we take a smaller initialization scale, which is consistent with our theory.\n    * For an indication that we might observe these dynamics at practical initialization scales, see Figures 6, 7, and 8. These figures show that in our toy model there is already some stage-wise learning behavior at initialization scale 0.1 which is a quite practical scale. We will emphasize this in the revision by moving Figure 6 to the main text.\n* Q: “Dynamic (3) seems to hold only for gradient descent training, how does it apply to other optimizers like SGD/Adam?”\n    * For our theory, we analyze gradient flow training, which can be obtained as a limit of SGD or GD training with learning rate → 0 (see e.g., [Bach ’20]). Gradient flow is generally simpler to analyze than SGD, and it is a popular testbed for studying learning dynamics (see e.g., [Saxe et al. ‘13], [Arora et al. ’18], [Razin and Cohen ’20], to name just a few examples). Analysis of Adam or constant-step-size SGD is certainly an interesting question for future work. However, it is beyond the theoretical scope of this paper as (1) it would significantly increase the complexity of the analysis, which is already involved; (2) our analysis of gradient flow is already a significant novel contribution in view of the existing literature. \n    * Our experiments show that gradual rank increase dynamics hold with Adam training. For the revision we have **added experiments on SGD-trained transformers** which show the same behavior; see attached document.\n* Q: “The learning dynamics theories... the later case?”\n    * There are several works exploring how to transfer guarantees from gradient flow to GD (see e.g., [De Sa et al. '22]). In our case, the loss function is smooth, so by Gronwall’s inequality our main theorem holds automatically for GD with sufficiently small step size. We will add a remark in the main text.\n* Weakness: “The assumptions are very strong....”\n    * On diagonal weights: Diagonal linear networks have recently been used actively as a toy model from which useful insights can be drawn and for which rigorous results can be potentially derived. This line of work has been active for several years, primarily in NeurIPS,  ICML and COLT (see the various references in the paper), and remains an active area with several open problems on some of the simplest models. Our result in this line of work reaches a new level by obtaining a formal result for a transformer-inspired model that maintains the softmax; this is an important component as diagonal networks with a single non-linearity are still far from being well understood theoretically.  \n    * On initialization scale: see our response to your questions above.\n    * On assumptions 4.2, 4.3, 4.4: These assumptions indeed hold for our toy model (an attention mechanism with diagonal weights). Please see Appendix C for experimental verification of these assumptions in our toy model.\n* Weakness: “The presentation of the paper could be improved.“\n    * Thank you for your feedback on the presentation of Sections 4.1.1 and 4.1.2. We were faced with a typical problem for theoretical papers: how to balance rigor and reader-friendliness in our proof sketch. There is a rigorous proof in Appendix A, so we decided to be less explicit on the error/magnitude bounds in the main text. We see now that this can be confusing for some readers. We will improve the presentation of Section 4 in the revision:\n        * As suggested by Reviewer yLKb, we will move Figure 6 to the main text. This illustrates Theorem 4.5 more clearly than Figure 3.\n        * As suggested by Reviewer yLKb, we will move some panels from Figures 9 and Figure 10 to the main text, since these illustrate assumptions 4.2,4.3,4.4 and our experimental verification.\n        * We will avoid phrases like “very small” and “negligible”, which referred to order $o_{\\\\alpha}(1)$ in Lines 150, 170, 188, 191, 196.\n        * We will improve the proof sketch in 4.1.1 and 4.1.2. This means adding more substance to the proof sketch in Section 4.1.1, and also shortening the proof sketch in Section 4.1.2 by removing technicalities. Concretely, in Section 4.1.1., we will explain how we use the conservation law to understand when the approximation to the dynamics is valid, since this is an important element.\n* References\n    * [Bach ’20] \"Effortless optimization through gradient flows.\" Machine Learning Research Blog. https://francisbach.com/gradient-flows (2020).\n    * [Saxe et al. ’13] \"Exact solutions to the nonlinear dynamics of learning in deep linear neural networks.\" arXiv preprint arXiv:1312.6120 (2013).\n    * [Arora et al. ’18] \"On the optimization of deep networks: Implicit acceleration by overparameterization.\" International Conference on Machine Learning. PMLR, 2018.\n    * [Razin and Cohen ’20] \"Implicit regularization in deep learning may not be explainable by norms.\" Advances in neural information processing systems 33 (2020): 21174-21187.\n    * [De Sa et al. ‘22] \"From Gradient Flow on Population Loss to Learning with Stochastic Gradient Descent.\" Advances in Neural Information Processing Systems 35 (2022): 30963-30976."}, {"review_id": "Xyw9FvmcJo", "reviewer": "Reviewer_1Eo1", "summary": "This paper conducts solid analysis and experiments to demonstrate the theory and proofs provided in the paper offer valuable insights into the incremental learning dynamics in transformers and how they can be better understood. ", "questions": "Please refers to the point listed in the Weakness part.", "limitations": "The proofs and conclusions of this study are dependent on the assumptions of the diagonal weights and small initialization. Such limitations have been discussed in an entire paragraph of Section 6.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The theory and proofs provided in the paper offer valuable insights into the incremental learning dynamics in transformers and how they can be better understood. \n2. The experiments conducted in the paper also provide evidence to support the proposed approach.", "weaknesses": "Besides LoRA, are there any other methods and techniques related to the proof and conclusions of this paper? A further discussion may help.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive evaluation of our paper. We were happy to read that you thought our contribution provided valuable insights and that our experiments gave good evidence. We answer your question below:\n\n* Q1: Besides LoRA, are there any other methods and techniques related to the proof and conclusions of this paper? A further discussion may help.\n  * Training dynamics of transformers are a very active area of research, and so there are indeed other methods and techniques related to the ideas in this paper that have appeared recently.\n      * A very relevant paper called “InRank: Incremental Low-Rank Learning” [Zhao et al. ’23] appeared on arXiv in June, after the NeurIPS deadline. Similarly to us, the authors observe gradual rank increase dynamics during transformer training. This inspires [Zhao et al. ’23] to run LoRA to train a transformer starting at random initialization, and they obtain significant runtime/memory improvements over regular training, with little loss in performance. The theory in [Zhao et al. ’23] is incomparable to ours — and actually it is quite complementary. Their theory studies linear networks with orthogonal weights (which are equivalent to linear diagonal networks because the modes evolve separately) in standard initialization scale regimes. In contrast, in our theory we study nonlinear networks with diagonal weights in small initialization scale regimes.\n      * A recent ICML paper called “On the stepwise nature of self-supervised learning” analyzes the gradient flow dynamics of a simple model and an SSL objective. They reach a similar conclusion where progress is made in discrete steps, where the rank of the learned embeddings increase gradually. The analysis conducted in that paper differs from ours in that they use a linear model as a theoretical testbed, where we consider a nonlinear model. \n\n\n* Limitations. We agree that our theory has these limitations. But, for completeness, we recall that our experiments on ViTs and GPT2 are a contribution, and these do not depend on diagonal weights and small initialization. \n* References\n    * Zhao, Jiawei, et al. \"InRank: Incremental Low-Rank Learning.\" arXiv preprint arXiv:2306.11250 (2023).\n    * Simon, James B., et al. \"On the stepwise nature of self-supervised learning.\" arXiv preprint arXiv:2303.15438 (2023).\n"}, {"review_id": "1P4DnJcExY", "reviewer": "Reviewer_yLKb", "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.\n", "questions": "1. Why is the rank at initialization bounded (at 128?) in the plots in Figure 5?\n2. 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?\n3. Are the same learning dynamics observed in ViT when training with SGD as opposed to Adam?\n", "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.\n", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n2. **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!). \n", "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.\n2. **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.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review and for your helpful suggestions on presentation and paper organization. We are happy that you found our theory and experiments convincing. We answer your questions below.\n\n* Q1: “Why is the rank at initialization bounded (at 128?) in the plots in Figure 5?“\n    * The rank is bounded because the head dimension is 128 and embedding dimension is 512, meaning that the matrices $W_Q, W_K, W_V, W_O$ are in $\\\\mathbb{R}^{512 \\\\times 128}$.\n* Q2: “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?”\n    * Yes, the top/bottom organization is correct, but it is not consistent between Figure 2 and Figure 3, which may have caused confusion. We will make it consistent in the revision. Both plots demonstrate that incremental learning dynamics occur for both key-query and value-output matrices.\n    * You make an interesting point that there seem to be some qualitative differences in the evolution of keys-queries and values-outputs matrices. See also Figures 12, 14 where the stable-rank of $\\\\Delta W_VW_O^T$ is higher than that of the $\\\\Delta W_Q W_K^T$ for CIFAR-10 and CIFAR-100. However, this trend does not hold for ImageNet (Figure 16). Our current theory unfortunately does not give any clue for why $\\\\Delta W_VW_O^T$ might have higher stable-rank than $\\\\Delta W_QW_K^T$, or vice-versa.\n* Q3: “Are the same learning dynamics observed in ViT when training with SGD as opposed to Adam?“\n    * Yes. Thanks for the suggestion. We have added experiments on SGD confirming this. See attached document.\n* Thank you also for your excellent suggestions on organization. We agree that moving Figure 6 and some panels from Figures 9 and 10 to the main text will help readability of Section 4, both for illustrating the theorem statement and illustrating and justifying the assumptions. We will implement this in the revision.\n\n"}]

Author rebuttal:
We thank the reviewers for their generally positive evaluations and for their helpful feedback that has helped us improve the paper in the revision. As suggested by the reviewers, we have **added new experiments**:
* on NLP transformers (GPT2 trained on Wikitext)
* and SGD-trained transformers (ViTs on CIFAR-10/100).

See attached file. These experiments further confirm the observations of the paper. We have responded to all reviewers’ questions individually, and we are happy to respond to any other questions they may have.

Meta-review: The paper received mixed reviews initially (3666). After rebuttal and discussion, the negative reviewer raised the score to 4, citing mainly remaining doubts on presentation. Even though doubts remain about the implications in practical cases, where the strong assumptions are not met (ZvU1), the AC agrees that the paper provides new insights on the incremental learning dynamics in transformers. Therefore, given the positive evaluation of 3 reviewers, and the minor reservations of the negative reviewer, the AC recommends acceptance of the paper. The authors should incorporate the discussion in their final version.
Decision: Accept (poster)