Title: Transformers learn through gradual rank increase
Venue: NeurIPS 2023
Accepted: yes
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

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 art performance in various domains, yet we still lack a solid theoretical understanding of its training dynamics [VSP + 17, DCLT19, LOG + 19, DBK + 20]. Nevertheless, the theoretical toolbox has matured over the last years and there are promising new 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 + 21, 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 1 , but the incremental learning dynamics are difficult to analyze, and are so far understood only for extremely simple architectures. Can we apply this analysis to transformers? Namely:

Are there incremental learning dynamics when training a transformer architecture?

An obstacle is that past work on incremental learning has mainly studied linear networks [Ber22, ACHL19, MKAA21, LLL20, WGL + 19, 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 ∈ R n × d , an attention head computes

where W K , W Q , W V , W O ∈ R d × d ′ are trainable matrices, and the softmax is applied row-wise. A transformer 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 is that transformers exhibit incremental learning dynamics, where the difference between the trained and initial weights incrementally increases in rank . Our results have a theoretical component and an experimental component.

(a)

(b)

Theoretical contributions For our theory, we study a simplification of the transformer architecture, where the attention head weights are diagonal matrices: i.e., in each attention head we have W K = diag( w K ) , where w K ∈ R d are trainable weights, and similarly for W Q , W V and W O . 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: (1) during most of each stage, the loss plateaus because the weights remain close to a 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 f NN that depend on the product of parameters u , v ∈ R p as

where x is the input, ⊙ denotes the elementwise product, and h is a smooth function.

Theorem 1.1 (Informal statement of incremental learning dynamics) . Let f NN 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 Θ( α ) for some small α > 0 . Then the dynamics of gradient flow training effectively proceeds in discrete stages, each one lasting time Θ(log(1 /α )) . In each stage, the number of nonnegligible entries of u ⊙ v increases by at most one.

A transformer with diagonal weight matrices falls under this result when we only train the attention head weights. For example, if the transformer has one attention head, then we can take u = [ w K , w V ] ∈ R 2 d and v = [ w Q , w O ] ∈ R 2 d 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 W K W ⊤ Q = diag( w K )diag( w Q ) or W V W ⊤ O = diag( w V )diag( w O ) 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 an intriguing connection to the LoRA algorithm, where a pretrained base model is cheaply fine-tuned by training a low-rank perturbation of the weights [LFLY18, AZG20, HSW + 21]. 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 our work we observe that the trained weights are a low-rank perturbation of the initial weights due to the training dynamics, without having to apply an explicit rank constraint as in LoRA. This raises an exciting open question for future work: can we explain and improve algorithms like LoRA by better understanding and quantifying the incremental dynamics of large transformers?

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, DLS22, BBSS22, MHPG + 22]. 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 a 2-layer linear network with L 2 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] independently provides a similar description of the dynamics, but without requiring balanced initialization. Finally, [Ber22, JLL + 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 a 2-layer network learns with a two-stage incremental dynamic; but that result needs the stylized assumption that all data points are 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 f NN ( · ; θ ) parametrized by a vector of weights θ , to minimize a loss

where the expectation is over samples ( x , y ) ∈ R d x × R d y from a training data distribution, and ℓ : R d y × R d out → R . Consider a solution θ ( t ) to the gradient flow 3

where α > 0 is a 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 ℓ ( y , ζ ) is continuously twice-differentiable in the arguments [ y , ζ ] ∈ R d y + d out . There exists C > 0 such that almost surely the data is bounded by ∥ x ∥ , ∥ y ∥ ≤ C .

The assumption on ℓ is satisfied in typical cases such as the square and the cross-entropy losses. The data boundedness is often satisfied in practice (e.g., if the data is normalized).

̸

We also use the notation supp( a ) := { i : a i = 0 } to denote the support of a vector a .

3 Neural networks with diagonal weights

Our theory analyzes the training dynamics of networks that depend on products of diagonal weight matrices. We use ⊙ to denote elementwise vector product.

Definition 3.1. A network f NN is smooth with diagonal weights θ = ( u , v ) ∈ R 2 p if it is of the form

where h : R d x × R p → R d out is continuously twice-differentiable in its arguments in R d x + p .

The assumption on h precludes the use of the ReLU function since it is not continuously-differentiable. Otherwise the assumption is fairly mild since any h can be used to express an architecture of any depth as long as the nonlinearities are twice-differentiable, which includes for example GeLUs (as used in ViT). We describe how to express a transformer with diagonal weights.

Example 3.2 (Transformer with diagonal weights) . A transformer with L layers and H attention heads on each layer is defined inductively by Z 0 = X ∈ R n × d and

(Attention layer) ˜ Z ℓ = Z ℓ -1 + ∑ H i =1 attention( Z ℓ -1 ; W ℓ,i K , W ℓ,i Q , W ℓ,i V , W ℓ,i O )

(Feedforward layer) Z ℓ = ˜ Z ℓ + σ ( ˜ Z ℓ W ℓ A )( W ℓ B ) ⊤ ,

where W ℓ,i K , W ℓ,i Q , W ℓ,i V , W ℓ,i O ∈ R d × d ′ are attention parameters, and W ℓ A , W ℓ B ∈ R d × d ′ are the feedforward parameters, and σ is a continuously twice-differentiable activation. Suppose that the attention parameters are diagonal matrices: i.e., W ℓ,i K = diag( w ℓ,i K ) ∈ R d × d , and similarly for the W ℓ,i Q , W ℓ,i V , W ℓ,i O matrices. Then by the definition of the attention layer (1), the final output of the transformer Z L only depends on the attention parameters through the elementwise products w ℓ,i K ⊙ w ℓ,i Q and w ℓ,i V ⊙ w ℓ,i O . In other words, we can write

for vectors u = [ w ℓ,i K , w ℓ,i V ] ( ℓ,i ) ∈ [ L ] × [ H ] ∈ R 2 dHL and v = [ w ℓ,i Q , w ℓ,i O ] ( ℓ,i ) ∈ [ L ] × [ H ] ∈ R 2 dHL , 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 if the initialization scale α 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 are implicitly defined by Algorithm 1 below, which constructs a sequence of times 0 = T 0 < T 1 < · · · < T k < · · · and weight vectors θ 0 , θ 1 , . . . , θ k , . . . ∈ R 2 p that define the stages. We prove the following:

Theorem 4.1 (Incremental dynamics at small initialization) . Let f NN be a model with diagonal weights as in Definition 3.1. For any stage k and time t ∈ ( T k , T k +1 ) the following holds under Assumptions 2.1, 4.3, 4.4 and 4.5. There is α 0 ( t ) > 0 such that for all α < α 0 , there exists a unique solution θ : [0 , t log(1 /α )] → R p to the gradient flow (3) and

and at each stage the sparsity increases by at most one: supp( θ k +1 ) \ supp( θ k ) ⊆ { i k } . 4

Application: transformer with diagonal weights Before giving the intuition for this theorem and stating the assumptions formally, let us discuss its application to the diagonal transformer model from Example 3.2. As a corollary of Theorem 4.1, the gradient flow on a diagonal transformer with small initialization will learn in stages, where in each stage there will be at most one head i ∈ [ H ] in one layer ℓ ∈ [ L ] such that either the rank of W ℓ,i K ( W ℓ,i Q ) ⊤ = diag( w ℓ,i K )diag( w ℓ,i Q ) or the rank of W ℓ,i V ( W ℓ,i O ) ⊤ = diag( w ℓ,i V )diag( w ℓ,i O ) increases by at most one. In Figure 2, we illustrate these dynamics in the toy case of a single attention head trained in a student-teacher setup.

Algorithm 1 Incremental learning in networks with diagonal weights

1: b 0 , θ 0 ← 0 ∈ R p , T 0 ← 0

2: for stage number k = 0 , 1 , 2 , . . . do

3: # (A) Pick new coordinate i k ∈ [ p ] to activate.

4: For each i , define time ∆ k ( i ) until active using (10).

5: Pick winning coordinate i k using (11)

6: Calculate time T k +1 using (11) and break if ∞

0.5

0.0

7: Update logarithmic weight approximation b k +1 using (12)

8: # (B) Train activated coordinates to stationarity.

9: θ k +1 ← limiting dynamics point from (13)

10: end for

-0.5

-1.0

3

Diagonal en ries of WQWK

4.1 Intuition for incremental learning dynamics

We develop an informal intuition for Theorem 4.1 and fill out the definition of Algorithm 1. A model f NN with diagonal weights θ = ( u , v ) as in Definition 3.1 evolves under the gradient flow (3) as

Here Dℓ ( y , · ) ∈ R 1 × d out is the derivative of ℓ in the second argument and Dh ( x , · ) ∈ R d out × p 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 [ACH18, DHL18], specialized to the case of diagonal layers.

Lemma 4.2 (Conservation law) . For any i ∈ [ p ] and any time t , we have

0

200

)diag(wK

400

), for alpha = 0.1

600

time / log(1/alpha)

800

1000

0

200

400

600

800

1000

time / log(1/alpha)

0.00

0.02

0.04

0.06

0.08

training loss

Training loss vs. rescaled time, for various alpha

alpha = 0.1

alpha = 0.01

alpha = 0.0001

alpha = 1e-08

alpha = 1e-16

alpha = 1e-32

0

200

400

600

800

1000

ime / log(1/alpha)

-1.0

-0.5

0.0

0.5

Diagonal en ries of WQWK

)diag(wK

), for alpha = 0.0001

0

200

400

600

800

1000

time / log(1/alpha)

-1.0

-0.5

0.0

0.5

Diagonal en ries of WQWK

)diag(wK

), for alpha = 1e-16

0

200

400

600

800

1000

time / log(1/alpha)

-1

0

1

2

3

Diagonal entries of WVWO

Entries of diag(wV

)diag(wO

), for alpha = 0.1

0

200

400

600

800

1000

time / log(1/alpha)

-1

0

1

2

3

Diagonal entries of WVWO

Entries of diag(wV

)diag(wO

), for alpha = 0.0001

Entries of diag(wV

)diag(wO

), for alpha = 1e-16

(a)

(b)

Diagonal en ries of WQWK

Diagonal en ries of WQWK

Diagonal entries of WVWO

Diagonal en ries of WQWK

Diagonal entries of WVWO

time / log(1/alpha)

0

200

400

600

800

1000

0.5

0.0

-0.5

-1.0

ime / log(1/alph

0

200

400

Entries of diag(wV

3

)diag(wK

),

600

)diag(wK

)diag(wO

0.5

2

0.0

1

-0.5

0

-1

-1.0

0

200

0

200

400

), f

), f

600

time / log(1/alpha

400

600

ime / log(1/alph

Entries of diag(wV

3

2

)diag(wO

)diag(wK

0.5

0.0

1

0

-0.5

-1

3

2

1

0

-1

-1.0

0

200

400

time / log(1/alph

0

200

400

600

time / log(1/alpha

Entries of diag(wV

), f

), fo

600

)diag(wO

0

200

400

), fo

600

time / log(1/alpha

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 i ( t ) := u i ( t ) + v i ( t ) , then the vector w ( t ) = u ( t ) + v ( t ) evolves by

Using the conservation law (5), we can keep track of the weights in terms of the initialization and w ( t ) :

Therefore it suffices to analyze the dynamics of w ( t ) .

4.1.1 Stage 1 of dynamics

Stage 1A of dynamics: loss plateau for time Θ(log(1 /α )) At initialization, θ (0) ≈ 0 because the weights are initialized on the order of α which is small. This motivates the approximation g ( θ ( t )) ≈ g ( 0 ) , under which the dynamics solve to:

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 t ≈ T 1 · log(1 /α ) for

Until this time, the network remains close to its initialization, and so we observe a loss plateau.

Stage 1B of dynamics: nonlinear dynamics for time O (1) Subsequently, the loss decreases nonlinearly during a O (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 i 0 such that 1 / | g i 0 ( 0 ) | = T 1 . Under this assumption, in stage 1A all weights except for those at coordinate i 0 remain vanishingly small, on the order of o α (1) . Concretely, for any small ϵ > 0 , there is a time t 1 ( ϵ ) ≈ T 1 · log(1 /α ) and sign s ∈ { +1 , -1 } such that 5

̸

Because all coordinates except for i 0 have vanishingly small o α (1) weights after stage 1A, we may perform the following approximation of the dynamics. Zero out the weights at coordinates except for i 0 , and consider the training dynamics starting at ˜ θ = ( ϵ e i 0 , sϵ e i 0 ) . After O (1) time, we should expect these dynamics to approach a stationary point. Although the evolution is nonlinear, all entries remain zero except for the i 0 entries, so the stationary point is also sparse. Mathematically, there is a time ¯ t 1 = t 1 + O (1) ≈ T 1 · log(1 /α ) such that

for some a ∈ R > 0 , where θ 1 is a stationary point of the loss. 6 Despite the nonlinearity of the dynamics, the approximation can be proved using Gr¨ onwall's inequality since ¯ t 1 -t 1 = O (1) is a 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 i 0 are nonzero.

4.1.2 Later stages

We inductively extend the argument to any number of stages k , where each stage has a Θ(log(1 /α )) -time plateau, and then a O (1) -time 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 of the magnitude of the weights on the logarithmic scale, since these determine how much longer . Inductively on k , suppose that there is some T k ∈ R , b k ∈ R p and θ k ∈ R 2 p and a time ¯ t k ≈ T k · log(1 /α ) such that

where θ k is a stationary point of the loss. Our inductive step shows that there is T k +1 ∈ R such that during times t ∈ ( ¯ t k , T k +1 · log(1 /α ) -Ω(1)) the weights remain close to the stationary point from the previous stage, i.e., θ ( t ) ≈ θ k . And at a time ¯ t k +1 ≈ T k +1 · log(1 /α ) we have

where θ k +1 and b k +1 are defined below (summarized in Algorithm 1). Most notably, θ k +1 is a stationary point of the loss whose support grows by at most one compared to θ k .

Stage ( k +1) A, loss plateau for time Θ(log(1 /α )) At the beginning of stage k +1 , the weights are close to the stationary point θ k , and so, similarly to stage 1A, linear dynamics are valid.

Using the conservation law (7), we derive a 'time until active' for each coordinate i ∈ [ p ] , which corresponds to the time for the weight at that coordinate to grow from o α (1) to Θ(1) magnitude:

̸

The linear dynamics approximation (9) breaks down at a time t ≈ T k +1 · log(1 /α ) , where

which corresponds to the first time at the weights at a coordinate grow from o α (1) to Θ(1) magnitude. And at times t ≈ T k +1 · log(1 /α ) , on the logarithmic scale w is given by

Stage ( k + 1) B of dynamics: nonlinear dynamics for time O (1) Subsequently, the weights evolve nonlinearly during O (1) time. In a similar way to the analysis of Stage 1B, we show that at a time ¯ t k +1 = t k +1 + O (1) ≈ T k +1 · log(1 /α ) , we have

where the dynamics ψ k ( t, ϵ ) ∈ R 2 p are initialized at ψ k (0 , ϵ ) = θ k +( ϵ e i k , sgn( g i ( θ k )) ϵ e i k ) and evolve according to the gradient flow d ψ k ( t,ϵ ) dt = -∇ θ L ( ψ k ) . 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 Θ(1) size at the same rescaled time. We also place a technical condition to handle the corner case when a coordinate leaves the support of the current stage's stationary point.

̸

Assumption 4.3 (Nondegeneracy of dynamics in part (A)) . The initialization satisfies | u i (0) | = | v i (0) | for all i . For stage k , either T k +1 = ∞ or there is a unique minimizer i k to min i ∆ k ( i k ) in (11). Finally, for all i ∈ supp( θ k -1 ) \ supp( θ k ) we have g i ( θ k ) = 0 .

̸

Coordinate val

1

0

-1

541

time / log(1/alpha)

time / log(1/al ha)

time / log(1/al ha)

200

220

240

260

280

300

0.0

-0.5

-1.0

200

220

542

543

544

545

546

547

0.0

1.5

1.0

0.5

Coo dinate val

-0.5

0.0

-1.0

541

542

543

544

545

546

547

778

779

780

781

782

783

784

time / log(1/al ha)

-1

0

1

2

Coordinate value

Entries of WVW T

O

,  erturb time / log(1/al ha) = 778.0

778

779

780

781

782

783

784

time / log(1/al ha)

-1.0

-0.5

0.0

0.5

1.0

Coordinate value

Entries of WKW T

Q

,  erturb time / log(1/al ha) = 778.0

322.5

325.0

327.5

330.0

332.5

335.0

337.5

time / log(1/alpha)

0

1

2

Coo dinate value

Ent ies of WVW T

O

, pe tu b time / log(1/alpha) = 321.4

time / log(1/alpha)

240

260

, pe tu b time / log

327.5

time / log(1/alpha)

0.00

Ent ies of WKW T

Coo dinate value

-0.02

-0.04

-0.06

322.5

325.0

330.0

332.

Ent ies of WKW T

Q

, pe tu b time / log

357.5 360.0 362.5 365.0 367.5 370.

time / log(1/alpha)

Ent ies of WKW T

2 O 0.5 Q 2 3 O 0.5 Assumption 4.4 (Stationary points are strict local minima) . For stage k , there exist δ k > 0 and c k > 0 such that for ˜ u ∈ B ( u k , δ ) supported on supp( u k ) , we have

-1

-1.0

-1

863 864 865 866 867 868 time / log(1/al ha) 863 864 865 866 867 868 time / log(1/al ha) Entries of WVW T ,  erturb time / log(1/al ha) = 942.2 Entries of WKW T ,  erturb time / log(1/al ha) = 942.2 357.5 360.0 362.5 365.0 367.5 370.0 372.5 375.0 377.5 time / log(1/alpha) Ent ies of WVW T , pe tu b time / log(1/alpha) = 405.4 Next, we require that very small perturbations of the coordinates outside of supp( θ k ) do not change the dynamics. For this, it suffices that θ k be a strict local minimum.

Coordinate value

1

Coo dinate value

0.0

-0.5

0

1

0

-0.5

942 943 944 945 946 947 948 -1 942 943 944 945 946 947 948 -1.0 406 408 410 412 414 416 418 -1 -1.0 Finally, we require a robust version of the assumption that the limit (13) exists, asking for convergence to a neighborhood of θ k +1 even when the initialization is slightly noisy.

406

408

Ent ies of WKW T

Q

0.0

-0.5

time / log(1/al ha) time / log(1/al ha) time / log(1/alpha) 0 2 Coo dinate value Ent ies of WVW T O , pe tu b time / log(1/alpha) = 486.5 0.5 Coo dinate value Assumption 4.5 (Noise-robustness of dynamics in part (B)) . For any stage k with T k +1 < ∞ and any ϵ > 0 , there are δ > 0 and τ : R > 0 → R such that the following holds. For any ˜ u ∈ B ( u k , δ ) ∩ R p ≥ 0 supported on supp(˜ u ) ⊆ supp( u k ) ∪ { i k } , there exists a unique solution ψ : [0 , ∞ ) → R p of the gradient flow d ψ dt = -∇ θ L ( ψ ) initialized at ψ (0) = (˜ u , s k +1 ⊙ ˜ u ) , and at times t ≥ τ ( ˜ ψ i k ) , ∥ ψ ( t ) -θ k +1 ∥ < ϵ .

-1

-1.0

Q

, pe tu b time / log

410

time / log(1/alpha)

412

414

, pe tu b time / log

550

575

600

time / log(1/alpha)

Q

5 Experimental results

time / log(1/alpha)

500

525

550

575

600

625

650

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 Wetrain 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 ∆ W K W ⊤ Q and ∆ W V W ⊤ O between the weights post-training and their initial values. Specifically, in Figure 4, we plot the stable rank of the differences ∆ W K W ⊤ Q and ∆ W V W ⊤ O . 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 ∆ W K W ⊤ Q against that of its initialized state in Figure 5 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 of initialization scale on gradual rank increase dynamics for a ViT trained on CIFAR-10. We use a ViT of depth 6, with 8 self-attention heads per layer (with layer normalization). We use an embedding and MLP dimension of d emb = 512 , and a head dimension of d h = 128 (i.e W K , W Q , W V , W O ∈ R d emb × d h ). We train the transformer

Coordinate val

Coo dinate val

500

525

0

20000

40000

60000

80000

Iteration

1

2

3

4

5

6

Stable rank

Layer 3

WQW T

K

WVW T

O

(a) ViT, CIFAR-10

(b) ViT, CIFAR-100

(c) ViT, ImageNet

(d) GPT-2, Wikitext-103

Layer 3

Layer 3

Layer 3

4

Stable rank

Stable rank

10

Stable rank

10

3

2

5

5

1

0

1000020000300004000050000

0

20000

40000

O

1000020000

30000

40000

Iterations

Iterations

Itterations

(a)

(b)

(c)

(d)

(e)

(f)

Layer 1

Layer 3

Layer 5

Spectrum atinit

Spectrum at init

0.8

Spectrum atinit

1.0

Spectrum of perturbation

1.0

Spectrum of perturbation

0.6

Spectrum of perturbation

Power

0.4

0.2

0.0

0.0

0.0

0

50

100

150

200

0

50

100

150

200

0

50

100

150

200

#Eigenvalue

#Eigenvalue

#Eigenvalue

Layer 4

Layer5

Layer 6

0.8

Spectrum at init

0.8

Spectrum at init

1.00

Spectrum at init

Power

0.6

Spectrum of perturbation

Power

0.6

Spectrum of perturbation

Power

0.75

Spectrum of perturbation

0.4

0.4

0.50

0.2

0.2

0.25

0.0

0.0

0.00

0

50

100

150

200

0

50

100

150

200

0

50

100

150

200

#Eigenvalue

#Eigenvalue

#Eigenvalue

using Adam with the 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 ∆ W K W ⊤ Q and ∆ W V W ⊤ O 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 in 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 [KB14]. Nevertheless, our experiments on ViTs indicate that the incremental learning dynamics occur even when training with Adam.

(a)

(b)

(c)

(d)

(e)

(f)

Scale=1e-4

Scale=1e-2

scale =1

EV magnitude

3

EV magnitude

4

EVmagnitude

0.8

2

3

0.6

2

0.4

1

1

0.2

0

0

0.0

0

10002000300040005000

0

10002000300040005000

0

10002000300040005000

Iterations

Iterations

Iterations

Scale=1e-4

scale =1e-2

scale =1

EVmagnitude

2.0

EV magnitude

EV magnitude

0.8

1.5

0.6

1.0

1.0

0.4

0.5

0.5

0.2

0.0

0.0

0.0

0

10002000300040005000

0

10002000300040005000

0

10002000300040005000

Iterations

Iterations

Iterations

(a)

(b)

(c)

Layer 1

Layer 3

Layer 5

10

6

15

Stablerank

rank

8

lerank

4

Stable

6

Stable

10

2

4

5

2

1e-2

1e-1

0.5

1

1.5

2

le-21e-1

0.5

1

1.5

2

le-21e-1

0.5

1

1.5

2

Weight scale at initialization

Weight scale at initialization

Weight scale at initialization

Future directions An interesting avenue of future research is to develop a theoretical understanding of the implicit bias in function space of transformers whose weights are a low-rank perturbation of 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] 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; in 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-Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks, COLT, 2022.

[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.11055 , 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 dimensionality explains the effectiveness of language model fine-tuning. arXiv preprint arXiv:2012.13255 , 2020.

[Bac20] Francis Bach. Effortless optimization through gradient flows. Machine Learning Research Blog. https://francisbach. com/gradient-flows , 2020.

[BBSS22] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. arXiv preprint arXiv:2210.15651 , 2022.

[Ber22] Rapha¨ el Berthier. Incremental learning in diagonal linear networks. arXiv preprint arXiv:2208.14673 , 2022.

[BMR + 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.00939 , 2022.

[COB18] L´ ena¨ ıc Chizat, Edouard Oyallon, and Francis R. Bach. On lazy training in differentiable programming. In Neural Information Processing Systems , 2018.

[DBK + 20] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. ArXiv , abs/2010.11929, 2020.

[DCLT19] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pretraining of deep bidirectional transformers for language understanding. ArXiv , abs/1810.04805, 2019.

[DHL18] Simon S Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homogeneous models: Layers are 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 on Learning Theory , pages 5413-5452. PMLR, 2022.

[FVB + 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.07082 , 2022.

[GBLJ19] Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of discrete gradient dynamics in linear neural networks. Advances in Neural Information Processing Systems , 32, 2019.

[GMMM19] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations of lazy training of two-layers neural network. Advances in Neural Information Processing Systems , 32, 2019.

[GSJW19] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy learning in deep neural networks: an empirical study. ArXiv , abs/1906.08034, 2019.

[GSSD19] Daniel Gissin, Shai Shalev-Shwartz, and Amit Daniely. The implicit bias of depth: How incremental learning drives generalization. arXiv preprint arXiv:1909.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.09685 , 2021.

[JGH18] Arthur Jacot, Franck Gabriel, and Cl´ ement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems , 31, 2018.

[JGS + 21] Arthur Jacot, Francois Gaston Ged, Berfin Simsek, Cl´ ement Hongler, and Franck Gabriel. Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. 2021.

[JLL + 23] 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.11500 , 2023.

[KB14] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 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 arXiv:1804.08838 , 2018.

[LLL20] Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient 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.11692, 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 , abs/2007.06738, 2020.

[MHPG + 22] Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with sgd. arXiv preprint arXiv:2209.14863 , 2022.

[MKAA21] Paolo Milanesi, Hachem Kadri, S. Ayache, and Thierry Arti` eres. Implicit regularization in deep tensor factorization. 2021 International Joint Conference on 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.

[RC20] Noam Razin and Nadav Cohen. Implicit regularization in deep learning may not be explainable by norms. Advances in neural information processing systems , 33:2117421187, 2020.

[SKZ + 23] James B Simon, Maksis Knutins, Liu Ziyin, Daniel Geisz, Abraham J Fetterman, and Joshua Albrecht. On the stepwise nature of self-supervised learning. arXiv preprint arXiv:2303.15438 , 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¨ oger 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. https://crfm. stanford. edu/2023/03/13/alpaca. 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.

[WGL + 19] Blake E. Woodworth, Suriya Gunasekar, J. Lee, Edward Moroshko, Pedro H. P. Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. ArXiv , abs/2002.09277, 2019.

[YW23] Hao Yu and Jianxin Wu. Compressing transformers: Features are low-rank, but weights are not! 2023.

[ZZC + 23] Jiawei Zhao, Yifei Zhang, Beidi Chen, Florian Sch¨ afer, and Anima Anandkumar. Inrank: Incremental low-rank learning. arXiv preprint arXiv:2306.11250 , 2023.

Contents

A Experimental validation of the assumptions in Theorem 4.1

In Figures 2, 8, and 9, we plot the evolution of the losses, of the entries of W K W ⊤ Q = diag( w K )diag( w Q ) , and of the entries of W V W ⊤ O = diag( w V )diag( w O ) in the toy task of training an attention head (1) with diagonal weights. The model is trained with SGD on the meansquared error loss on 1000 random samples ( X , y ) . Each random sample has X ∈ R 10 × 50 , which a sequence of 10 tokens, each of dimension 50, which is distributed as isotropic Gaussian. The label y is given by a randomly-generated teacher model that is also an attention head (1) with diagonal weights. In Figures 2, 8, and 9, for α ∈ { 0 . 1 , 0 . 01 , 0 . 0001 , 10 -8 , 10 -16 , 10 -32 } we plot the evolution of the loss and of the weights when initialized at θ (0) = α θ 0 , for some random Gaussian θ 0 . Qualitatively, as α → 0 we observe that the loss curve and the trajectories of the weights appear to converge to a limiting stagewise dynamics, where there are plateaus followed by movement on short time-scales, as predicted by the theory.

Validation of Assumption 4.3 (non-degeneracy of dynamics) As α → 0 , notice that the stages appear to separate and happen at distinct times. Furthermore, the extra technical condition on coordinates i ∈ supp( θ k ) \ supp( θ k -1 ) in Assumption 4.3 is satisfied since no coordinates ever leave the support of θ k .

Validation of Assumption 4.4 (stationary points are strict local minima) In Figure 10 we consider the α = 10 -32 trajectory, since this is closest to the dynamics in the α → 0 limit. We randomly select several epochs. Since the transitions between stages are a vanishing fraction of the total training time, each of these randomly-selected 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 -5 ), we perturb the weights by adding noise of standard deviation 0 . 05 . We then run the training dynamics starting at this perturbed initialization for 1000 epochs. We observe that the training dynamics quickly converge to the original unperturbed initialization, indicating that the weights were close to a strict local minimum of the loss.

Validation of Assumption 4.5 (noise-robustness of dynamics) In 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 a newly-nonnegligible coordinate (less than 10 -5 in magnitude in the previous epoch, and more than 10 -5 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 10 -2 on each of the nonnegligible coordinates, supporting Assumption 4.5.

0

200

400

600

800

1000

time / log(1/alpha)

-1.0

-0.5

0.0

0.5

Diagonal en ries of WQWK

)diag(wK

), for alpha = 0.1

0

200

400

600

800

1000

ime / log(1/alpha)

-1.0

-0.5

0.0

0.5

Diagonal en ries of WQWK

)diag(wK

), for alpha = 0.01

0

200

400

600

800

1000

ime / log(1/alpha)

-1.0

-0.5

0.0

0.5

Diagonal en ries of WQWK

)diag(wK

), for alpha = 0.0001

0

200

400

600

800

1000

ime / log(1/alpha)

-1.0

-0.5

0.0

0.5

Diagonal en ries of WQWK

)diag(wK

), for alpha = 1e-08

0

200

400

600

800

1000

time / log(1/alpha)

-1.0

-0.5

0.0

0.5

Diagonal en ries of WQWK

)diag(wK

), for alpha = 1e-16

0

200

400

600

800

1000

time / log(1/alpha)

-1.0

-0.5

0.0

0.5

Diagonal en ries of WQWK

)diag(wK

), for alpha = 1e-32

0

200

400

600

800

1000

time / log(1/alpha)

-1

0

1

2

3

Diagonal entries of WVWO

Entries of diag(wV

)diag(wO

), for alpha = 0.1

0

200

400

600

800

1000

time / log(1/alpha)

-1

0

1

2

3

Diagonal entries of WVWO

Entries of diag(wV

)diag(wO

), for alpha = 0.01

0

200

400

600

800

1000

time / log(1/alpha)

-1

0

1

2

3

Diagonal entries of WVWO

Entries of diag(wV

)diag(wO

), for alpha = 0.0001

0

200

400

600

800

1000

time / log(1/alpha)

-1

0

1

2

3

Diagonal entries of WVWO

Entries of diag(wV

)diag(wO

), for alpha = 1e-08

0

200

400

600

800

1000

time / log(1/alpha)

-1

0

1

2

3

Diagonal entries of WVWO

Entries of diag(wV

)diag(wO

), for alpha = 1e-16

0

200

400

600

800

1000

time / log(1/alpha)

-1

0

1

2

3

Diagonal entries of WVWO

Entries of diag(wV

)diag(wO

), for alpha = 1e-32

374

375

376

377

378

379

time / log(1/al ha)

-1

0

1

2

3

Coordinate value

Entries of WVW T

O

,  erturb time / log(1/al ha) = 373.4

374

375

376

377

378

379

time / log(1/al ha)

-1.0

-0.5

0.0

0.5

Coordinate value

Entries of WKW T

Q

,  erturb time / log(1/al ha) = 373.4

541

542

543

544

545

546

547

time / log(1/al ha)

-1

0

1

2

3

Coordinate value

Entries of WVW T

O

,  erturb time / log(1/al ha) = 541.2

541

542

543

544

545

546

547

time / log(1/al ha)

-1.0

-0.5

0.0

0.5

Coordinate value

Entries of WKW T

Q

,  erturb time / log(1/al ha) = 541.2

778

779

780

781

782

783

784

time / log(1/al ha)

-1

0

1

2

Coordinate value

Entries of WVW T

O

,  erturb time / log(1/al ha) = 778.0

778

779

780

781

782

783

784

time / log(1/al ha)

-1.0

-0.5

0.0

0.5

1.0

Coordinate value

Entries of WKW T

Q

,  erturb time / log(1/al ha) = 778.0

863

864

865

866

867

868

time / log(1/al ha)

-1

0

1

2

3

Coordinate value

Entries of WVW T

O

,  erturb time / log(1/al ha) = 862.4

863

864

865

866

867

868

time / log(1/al ha)

-1.0

-0.5

0.0

0.5

1.0

Coordinate value

Entries of WKW T

Q

,  erturb time / log(1/al ha) = 862.4

942

943

944

945

946

947

948

time / log(1/al ha)

-1

0

1

2

Coordinate value

Entries of WVW T

O

,  erturb time / log(1/al ha) = 942.2

942

943

944

945

946

947

948

time / log(1/al ha)

-1.0

-0.5

0.0

0.5

Coordinate value

Entries of WKW T

Q

,  erturb time / log(1/al ha) = 942.2

200

220

240

260

280

300

time / log(1/alpha)

0.0

0.5

1.0

1.5

2.0

2.5

Coo dinate value

Ent ies of WVW T

O

, pe tu b time / log(1/alpha) = 202.5

200

220

240

260

280

300

time / log(1/alpha)

-1.0

-0.5

0.0

0.5

1.0

Coo dinate value

1e-5

Ent ies of WKW T

Q

, pe tu b time / log(1/alpha) = 202.5

322.5

325.0

327.5

330.0

332.5

335.0

337.5

time / log(1/alpha)

0

1

2

Coo dinate value

Ent ies of WVW T

O

, pe tu b time / log(1/alpha) = 321.4

322.5

325.0

327.5

330.0

332.5

335.0

337.5

time / log(1/alpha)

-0.06

-0.04

-0.02

0.00

Coo dinate value

Ent ies of WKW T

Q

, pe tu b time / log(1/alpha) = 321.4

357.5 360.0 362.5 365.0 367.5 370.0 372.5 375.0 377.5

time / log(1/alpha)

-1

0

1

2

3

Coo dinate value

Ent ies of WVW T

O

, pe tu b time / log(1/alpha) = 356.8

357.5 360.0 362.5 365.0 367.5 370.0 372.5 375.0 377.5

time / log(1/alpha)

-0.5

0.0

0.5

Coo dinate value

Ent ies of WKW T

Q

, pe tu b time / log(1/alpha) = 356.8

406

408

410

412

414

416

418

time / log(1/alpha)

-1

0

1

2

3

Coo dinate value

Ent ies of WVW T

O

, pe tu b time / log(1/alpha) = 405.4

406

408

410

412

414

416

418

time / log(1/alpha)

-1.0

-0.5

0.0

0.5

Coo dinate value

Ent ies of WKW T

Q

, pe tu b time / log(1/alpha) = 405.4

500

525

550

575

600

625

650

time / log(1/alpha)

-1

0

1

2

Coo dinate value

Ent ies of WVW T

O

, pe tu b time / log(1/alpha) = 486.5

500

525

550

575

600

625

650

time / log(1/alpha)

-1.0

-0.5

0.0

0.5

Coo dinate value

Ent ies of WKW T

Q

, pe tu b time / log(1/alpha) = 486.5

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 a priori unclear to what extent gradual rank increase behaviour and a 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 in the literature. We use the stable rank of a matrix W given by ∥ W ∥ 2 F ∥ W ∥ 2 2 as a smooth approximation of rank. We track the value of the stable rank for the 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 to a low rank bias in Figures 12, 13, 15, 17 and 19. In these experiments we use off-the-shelf vision transformers (ViT) [DBK + 20] trained on popular vision benchmarks, as well as off-the-shelf GPT-2 trained on a 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 8 heads per layer, embedding dimension 512, head dimension 128, and MLP dimension 512 and patch-size 4 for 500 epochs on CIFAR10/CIFAR100 with SGD and learning rate 3e-1 and warmup. See Figures 12 and 13. Each run took 2 hours on one A100 GPU.

0

10000

20000

30000

40000

50000

Iteration

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

Stable rank

Layer 0

WQW T

K

WVW T

O

0

10000

20000

30000

40000

50000

Iteration

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

Stable rank

Layer 1

WQW T

K

WVW T

O

0

10000

20000

30000

40000

50000

Iteration

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

Stable rank

Layer 2

WQW T

K

WVW T

O

0

10000

20000

30000

40000

50000

Iteration

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

Stable rank

Layer 3

WQW T

K

WVW T

O

0

10000

20000

30000

40000

50000

Iteration

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

Stable rank

Layer 4

WQW T

K

WVW T

O

0

10000

20000

30000

40000

50000

Iteration

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

Stable rank

Layer 5

WQW T

K

WVW T

O

0

10000

20000

30000

40000

50000

Iteration

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

Stable rank

Layer 0

WQW T

K

WVW T

O

0

10000

20000

30000

40000

50000

Iteration

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

Stable rank

Layer 1

WQW T

K

WVW T

O

0

10000

20000

30000

40000

50000

Iteration

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

Stable rank

Layer 2

WQW T

K

WVW T

O

0

10000

20000

30000

40000

50000

Iteration

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

Stable rank

Layer 3

WQW T

K

WVW T

O

0

10000

20000

30000

40000

50000

Iteration

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

Stable rank

Layer 4

WQW T

K

WVW T

O

0

10000

20000

30000

40000

50000

Iteration

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

Stable rank

Layer 5

WQW T

K

WVW T

O

B.2 Adam-trained transformers

CIFAR-10/100 For the CIFAR-10/100 datasets we use a VIT with 6 layers, patchsize of 4, 8 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 a base learning rate of 1e-4, a cyclic learning rate decay with a linear warmup schedule for 15 epochs and a batchsize of 512. Our results are summarized in Figures 14 and 15 for CIFAR-10, and Figures 16 and 17 for CIFAR-100.

(a)

(b)

(c)

(d)

(e)

(f)

layer4

layer5

1.00

1.00

1.00

init

init

init

0.75

0.75

0.75

ower

t=1000

ower

t=1000

Power

t=1000

0.50

0.50

0.50

t=5000

t=5000

t=5000

0.25

0.25

0.25

0.00

0.00

0.00

0

25

50

75

100

125

0

25

50

75

100

125

0

25

50

75

100

125

#Eigenvalue

#Eigenvalue

#Eigenvalue

layer 4

layer5

1.00

1.00

1.00

init

init

init

0.75

0.75

0.75

Power

t=1000

Power

t=1000

ower

t=1000

0.50

t=5000

0.50

t=5000

0.50

t=5000

0.25

0.25

0.25

0.00

0.00

0.00

0

25

50

75

100

125

0

25

50

75

100

125

0

25

50

75

100

125

#Eigenvalue

#Eigenvalue

#Eigenvalue

ImageNet For ImageNet, we use the VIT-Base/16 from [DBK + 20] trained with Adam for 360 epochs with a base learning rate of 3e-3, a cyclic learning rate decay with a linear warmup schedule for 15 epochs and a batchsize of 4096. Our results are summarized in Figures 18 and 19 for ImageNet.

(a)

(b)

(c)

(d)

(e)

(f)

Layer 0

Layer 1

Layer 2

2.5

3.0

Stable rank

Stable rank

2.5

Stablerank

3

2.0

1.5

1.5

2

1.0

1.0

1

0

20000

40000

0

20000

40000

0

100002000030000 4000050000

Iterations

Iterations

Iterations

Layer3

Layer 4

Layer5

5

Stable rank

4

Stable rank

4

Stable rank

4

3

3

3

2

2

2

1

1

1

1000020000300004000050000

0

1000020000300004000050000

0

1000020000300004000050000

Iterations

Iterations

Iterations

(a)

(b)

(c)

(d)

(e)

(f)

layer 4

layer5

1.00

1.00

1.00

init

init

init

0.75

0.75

0.75

t=1000

ower

t=1000

Power

t=1000

0.50

t=5000

0.50

t=5000

0.50

t=5000

0.25

0.25

0.25

0.00

0.00

0.00

0

25

50

75

100

125

0

25

50

75

100

125

0

25

50

75

100

125

#Eigenvalue

#Eigenvalue

#Eigenvalue

layer 4

layer5

1.00

1.00

1.00

init

init

init

0.75

0.75

0.75

Power

t=1000

Power

t=1000

Power

t=1000

0.50

t=5000

0.50

t=5000

0.50

t=5000

0.25

0.25

0.25

0.00

0.00

0.00

0

25

50

75

100

125

0

25

50

75

100

125

0

25

50

75

100

125

#Eigenvalue

#Eigenvalue

#Eigenvalue

(a)

(b)

(c)

(d)

(e)

(f)

Layer 0

Layer 1

Layer 2

Stablerank

Stablerank

8

Stable rank

8

6

4

4

4

2

2

2

O

100002000030000 4000050000

0

1000020000300004000050000

0

1000020000300004000050000

Iterations

Iterations

Iterations

Layer3

Layer 4

Layer 5

15

Stablerank

10

Stable rank

10

Stable rank

10

5

5

5

20000

40000

0

1000020000300004000050000

0

20000

40000

Iterations

Iterations

Iterations

(a)

(b)

(c)

(d)

(e)

(f)

layer6

layer10

1.00

1.00

1.00

init

init

init

0.75

t=400

0.75

t=400

0.75

t=400

t=4000

ower

t=4000

Power

t=4000

0.50

t=24000

0.50

t=24000

0.50

t=24000

0.25

0.25

0.25

0.00

0.00

0.00

0

25

50

75

100

125

0

25

50

75

100

125

0

25

50

75

100

125

#Eigenvalue

#Eigenvalue

#Eigenvalue

layer6

layer10

1.00

1.00

1.00

init

init

init

0.75

t=400

0.75

t=400

0.75

t=400

Power

t=4000

Power

t=4000

ower

t=4000

0.50

t=24000

0.50

t=24000

0.50

t=24000

0.25

0.25

0.25

0.00

0.00

0.00

0

25

50

75

100

125

0

25

50

75

100

125

0

25

50

75

100

125

#Eigenvalue

#Eigenvalue

#Eigenvalue

(a)

(b)

(c)

(d)

(e)

(f)

(g)

(h)

(i)

(j)

(k)

(l)

Layer 0

Layer 1

Layer 2

5

10.0

Stable rank

4

Stablerank

7.5

3

2

4

5.0

2

2.5

1

10000

20000

30000

40000

0

10000

20000

30000

40000

0

100002000030000

40000

Itterations

Itterations

Itterations

Layer 3

Layer 4

Layer 5

Stable rank

10

Stablerank

10

Stable rank

10

5

5

0

10000200003000040000

0

10000200003000040000

0

10000200003000040000

Itterations

Itterations

Itterations

Layer6

Layer 7

Layer8

15

15

20

Stable rank

Stable rank

Stable rank

15. 0 5

10

0

10000200003000040000

0

10000200003000040000

0

10000200003000040000

Itterations

Itterations

Itterations

Layer 9

Layer 10

Layer 11

20

Stable rank

15

Stable rank

20

Stable rank

20

101

10

10

0

0

10000200003000040000

0

10000200003000040000

0

10000200003000040000

Itterations

Itterations

Itterations

0

20000

40000

60000

80000

Iteration

1

2

3

4

5

6

Stable rank

Layer 0

WQW T

K

WVW T

O

0

20000

40000

60000

80000

Iteration

1

2

3

4

5

6

Stable rank

Layer 1

WQW T

K

WVW T

O

0

20000

40000

60000

80000

Iteration

1

2

3

4

5

6

Stable rank

Layer 2

WQW T

K

WVW T

O

0

20000

40000

60000

80000

Iteration

1

2

3

4

5

6

Stable rank

Layer 3

WQW T

K

WVW T

O

0

20000

40000

60000

80000

Iteration

1

2

3

4

5

6

Stable rank

Layer 4

WQW T

K

WVW T

O

0

20000

40000

60000

80000

Iteration

1

2

3

4

5

6

Stable rank

Layer 5

WQW T

K

WVW T

O

0

20000

40000

60000

80000

Iteration

1

2

3

4

5

6

Stable rank

Layer 6

WQW T

K

WVW T

O

0

20000

40000

60000

80000

Iteration

1

2

3

4

5

6

Stable rank

Layer 7

WQW T

K

WVW T

O

0

20000

40000

60000

80000

Iteration

1

2

3

4

5

6

Stable rank

Layer 8

WQW T

K

WVW T

O

0

20000

40000

60000

80000

Iteration

1

2

3

4

5

6

Stable rank

Layer 9

WQW T

K

WVW T

O

0

20000

40000

60000

80000

Iteration

1

2

3

4

5

6

Stable rank

Layer 10

WQW T

K

WVW T

O

0

20000

40000

60000

80000

Iteration

1

2

3

4

5

6

Stable rank

Layer 11

WQW T

K

WVW T

O

C Proof for dynamics of networks with diagonal parametrization (Theorem 4.1)

C.1 Assumptions

Recall we have defined θ 0 , . . . , θ k , . . . ∈ R 2 p as the sequence of weights such that θ 0 = 0 and θ k +1 is defined inductively as follows. Consider the dynamics of ψ k ( t, ϵ ) ∈ R 2 p initialized at ψ k (0 , ϵ ) = θ k +( ϵ e i k , sgn( g i ( θ k )) ϵ e i k ) and evolving according to the gradient flow d ψ k ( t,ϵ ) dt = -∇ θ L ( ψ k ) . We assume that there is a limiting point θ k +1 of these dynamics as ϵ is taken small and the time is taken large:

Under the above assumption that this sequence θ 0 , . . . , θ k , . . . is well-defined, we can derive a useful property of it for free. Namely, the conservation law (5) implies that u ⊙ u -v ⊙ v is preserved. It follows that for each k we have that θ k = ( u k , v k ) satisfies | u k | = | v k | entrywise. In other words, there is s k ∈ { +1 , -1 } p satisfying

We also abuse notation and write supp( θ k ) := supp( u k ) ⊆ [ p ] , since the support of θ k on the first p coordinates matches its support on the last p coordinates.

Having fixed this notation, we now recall the main assumptions of the theorem.

̸

Assumption C.1 (Nondegeneracy of dynamics in part (A); Assumption 4.3) . The initialization satisfies | u i (0) | = | v i (0) | for all i . For stage k , either T k +1 = ∞ or there is a unique minimizer i k to min i ∆ k ( i k ) in (11). Finally, for all i ∈ supp( θ k -1 ) \ supp( θ k ) we have g i ( θ k ) = 0 .

̸

Assumption C.2 (Stationary points are strict local minima; Assumption 4.4) . For stage k , there exist δ k > 0 and c k > 0 such that for ˜ u ∈ B ( u k , δ ) supported on supp( u k ) , we have

Assumption C.3 (Noise-robustness of dynamics in part (B); Assumption 4.5) . For stage k , either T k +1 = ∞ or the following holds. For any ϵ > 0 , there are δ > 0 and τ : R > 0 → R such that the following holds. For any ˜ u ∈ B ( u k , δ ) ∩ R p ≥ 0 supported on supp(˜ u ) ⊆ supp( u k ) ∪ { i k } , there exists a unique solution ψ : [0 , ∞ ) → R p of the gradient flow d ψ dt = -∇ θ L ( ψ ) initialized at ψ (0) = (˜ u , s k +1 ⊙ ˜ u ) , and at times t ≥ τ (˜ u i k ) ,

C.2 Rescaling time for notational convenience

For ease of notation, we rescale time

We also define

Because of this time-rescaling, we equivalently state Theorem 4.1 as:

Theorem C.4 (Restatement of Theorem 4.1) . Let K ∈ Z ≥ 0 be such that Assumptions 4.3 4.4 hold for all k ≤ K and Assumption 4.5 holds for all k < K . Then for any k ≤ K and time t ∈ ( T k , T k +1 ) the following holds. There is α 0 ( t ) > 0 such that for all α < α 0 , there exists a unique solution θ α : [0 , t ] → R p to the gradient flow (14) and

where at each stage | supp( u k ) \ supp( u k -1 ) | ≤ 1 .

For shorthand, we also write

C.3 Simplifying problem without loss of generality

̸

For each coordinate i ∈ [ p ] we have | u α,i (0) | = | v α,i (0) | by the non-degeneracy Assumption 4.3. So we can assume | u α,i (0) | > | v α,i (0) | without loss of generality. Furthermore, we can assume the entrywise inequality

by otherwise training weights ˜ u α ( t ) , ˜ v α ( t ) initialized at ˜ u α (0) = sgn( u α (0)) u α (0) and ˜ v α (0) = sgn( v α (0)) v α (0) , as ˜ u α ( t ) ⊙ ˜ v α ( t ) = u α ( t ) ⊙ v α ( t ) at all times.

Since u 2 α,i ( t ) -v 2 α,i ( t ) = u 2 α,i (0) -v 2 α,i (0) by the conservation law (5), it holds that | u α,i ( t ) | > | v α,i ( t ) | throughout. So by continuity

throughout training.

C.4 Tracking the sum of the weights

We define and if T k +1 < ∞ then

The reason for this definition is that during training we have

Notice that since that we have assumed u α,i (0) > | v α,i (0) | for each i ∈ [ p ] we have w α (0) > 0 entrywise. So, by (15) for all t > 0 ,

It suffices to track w α ( t ) because we can relate the log-scale magnitude of w α ( t ) to the magnitudes of the corresponding coordinates in u α ( t ) and v α ( t ) - see technical Lemmas D.1 D.2 and D.3.

C.5 Claimed invariants in proof of Theorem C.4

In order to prove Theorem C.4, we consider any gradient flow θ α : [0 , T ∗ ] → R p solving (14) where T ∗ ∈ ( T K , T K +1 ) . For now, we focus only on proving properties of 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 ϵ > 0 , any stage k ≤ K , there is α k := α k ( ϵ ) > 0 such that for all α < α k the following holds. There are times ¯ t k := ¯ t k ( α, ϵ ) and t k +1 := t k +1 ( α, ϵ ) , such that

and the weights approximate the greedy limit for all times t ∈ [ ¯ t , t ]

and the weights at times ¯ t k and t k +1 are correctly estimated by the incremental learning dynamics on the logarithmic-scale

Base case k = 0 : Take ¯ t 0 ( α, ϵ ) = 0 . Then statement (16) holds since T 0 = 0 . Notice that as α → 0 we have that u α (0) , v α (0) → 0 = u 0 , and also log α w α (0) → 1 = b 0 . So statement (19) follows if we take α 0 small enough. In Section C.6 we show how to construct time t 1 such that (18) and (20) hold.

Inductive step : Suppose that (16), (18), (19) and (20) hold for some iteration k < K . We prove them for iteration k +1 . In Section C.7 we construct time ¯ t k . In Section C.6 we construct time t k +1 .

C.6 Dynamics from time ¯ t k to time t k +1 (Linear dynamics for O (log(1 /α )) unrescaled time)

Let k ≤ K , and suppose that we know that for any ¯ ϵ k > 0 , there is ¯ α k (¯ ϵ k ) > 0 such that for all 0 < α < ¯ α k , there is a time ¯ t k = ¯ t k ( α, ¯ ϵ k ) satisfying

C.6.1 Analysis in case where T k +1 < ∞

Consider first the case where T k +1 < ∞ . We show that, for any ϵ k +1 > 0 , there is ρ k +1 ( ϵ k +1 ) > 0 such that for all 0 < ρ < ρ k +1 (¯ ϵ k +1 ) there is α k +1 ( ρ, ϵ k +1 ) > 0 such that for all α < α k +1 , there is a time t k +1 = t k +1 ( α, ρ, ϵ k +1 ) satisfying

For any ρ, α , let ¯ ϵ k = ρϵ k +1 / (4 p ) and choose ¯ t k = ¯ t k ( α, ¯ ϵ k ) . Then define

Now we show that the weights θ α ( t ) cannot move much from time ¯ t k to t k +1 . The argument uses the local Lipschitzness of the loss L (from technical Lemma D.7), and the strictness of θ k as a stationary point (from Assumption 4.4).

Lemma C.5 (Stability of active variables during part (A) of dynamics) . There is ρ k +1 small enough and α k +1 ( ρ ) small enough depending on ρ ,such that for all ρ < ρ k +1 and α < α k +1 and all t ∈ [ ¯ t k , t k +1 ) ,

where c k is the strict-minimum constant from Assumption 4.4 and K R k is the Lipschitzness constant from Lemma D.7 for the ball of radius R k = ∥ θ k ∥ +1 .

Proof. Assume by contradiction that (27) is violated at some time t < t k +1 . Let us choose the first such time

Define ˜ θ = (˜ u , ˜ v ) by

By the definition of t k +1 , this satisfies

Also

Using (a) the strict minimum Assumption 4.4 with constant c k , since ∥ ˜ θ -θ k ∥ ≤ ρ ′ and we take ρ ′ small enough,

This is a contradiction because L is nondecreasing along the gradient flow.

Lemma C.6 (Log-scale approximation is correct during part (A)) . There are functions ρ k +1 ( ϵ k +1 ) > 0 and α k +1 ( ρ, ϵ k +1 ) > 0 such that for all ρ < ρ k +1 and α < α k +1 , and for all t ∈ ( ¯ t k , t k +1 ) we have for a constant C depending on k ,

Furthermore, for all i ∈ S c k and t ∈ ( ¯ t k , t k +1 ) we have

Proof. By Lemma C.5 and Lemma D.7, there is a constant C depending on θ k such that for all t ∈ ( ¯ t k , t k +1 ) ,

For shorthand, write ¯ g ( θ k ) = g ( θ k ) + Cρ ′ 1 and g ( θ k ) = g ( θ k ) -Cρ ′ 1 . Since w α ( t ) > 0 entrywise as we have assumed without loss of generality (see Section C.3), we have the following entrywise inequalities

Since the dynamics are given by d w α dt = log(1 /α ) g ( w α ) ⊙ w α ,

Taking the logarithms with base α ∈ (0 , 1) ,

The bound (28) follows since ∥ log α ( w α ( ¯ t k )) -b k ∥ < ¯ ϵ k < ρϵ k +1 .

Finally, the claim (29) follows from (30) since sgn(¯ g ( θ k )) = sgn( g ( θ k )) = sgn( g ( θ k )) if we take ρ small enough.

First, we show that the weights must move significantly by time roughly T k +1 . This is because of the contribution of coordinate i k .

Lemma C.7 ( t k +1 is not much larger than T k +1 ) . Suppose that T k +1 < ∞ . Then there are ρ k +1 ( ϵ k +1 ) > 0 and α k +1 ( ρ, ϵ k +1 ) > 0 such that for all ρ < ρ k +1 and α < α k +1 , the following holds.

Proof. Assume by contradiction that t k +1 < T k +1 + ϵ k +1 . For all times t ∈ [ ¯ t k , min( t k +1 , T k +1 + ϵ k +1 )] , by Lemma C.6,

Since we know | ∆ k ( i k ) -( T k +1 -¯ t k ) | < ¯ ϵ k and b k i -∆ k ( i k ) g i k ( θ k ) ∈ { 0 , 2 } , it follows that

By taking ρ small enough, we see that | g i k ( θ k ) | ( ϵ k +1 -¯ ϵ k +1 ) + O ( √ ρ ) > δ > 0 for some δ > 0 that is independent of α , so

So | u α,i k ( T k +1 + ϵ k +1 ) | > 1 by Lemma D.2. But by the construction of t k +1 this means that t k +1 < T k +1 + ϵ k +1 .

Next, we show that until time t k +1 , none of the coordinates in S c k move significantly, with the possible exception of coordinate i k .

Lemma C.8 (No coordinates in S c k \ { i k } move significantly during part (A)) . Suppose T k +1 < ∞ . Then there are ρ k +1 ( ϵ k +1 ) > 0 and α k +1 ( ρ, ϵ k +1 ) > 0 such that for all ρ < ρ k +1 and α < α k +1 , the following holds. There is a constant c > 0 depending on k such that for all i ∈ S c k \ { i k } and t ∈ [ ¯ t k , t k +1 ] ,

Proof. The previous lemma combined with the inductive hypothesis gives

We analyze the movement of each coordinate i ∈ S c k \ { i k } by breaking into two cases:

̸

Coordinate i = i k such that b k i ∈ (0 , 2) . By Assumption 4.3, there is a unique winning coordinate so b k i -τg i ( θ k ) ∈ ( c, 2 -c ) for some constant c > 0 for all τ ∈ [0 , t k +1 -¯ t k ] ⊆ [0 , ∆ k ( i k ) + 2 ϵ k +1 ] . By Lemma C.6, log α ( w α,i ( t )) ∈ ( -c/ 2 , 2 -c/ 2) for all times t ∈ [ ¯ t k , t k +1 ] . So by Lemma D.1, | u α,i ( t ) | , | v α,i ( t ) | ≤ α c/ 4 .

̸

Coordinate i = i k such that b k i = 0 . By Lemma D.4, we must be in the corner case where i ∈ S k -1 ∩ S c k (i.e., the coordinate was active in the previous stage but was dropped from the support in this stage).

By Lemma D.4, since b k i = 0 we have g i ( θ k ) < 0 . By Lemma C.6, this means sgn( g i ( θ α ( t ))) = sgn( g i ( θ k )) < 0 for all t ∈ ( ¯ t k , t k +1 ) .

We break the analysis into two parts. Since b k i = 0 , the sign is s k i = +1 . The inductive hypothesis ∥ θ α ( ¯ t k ) -θ k ∥ < ¯ ϵ k implies that | u α,i ( ¯ t k ) -z k i | < ¯ ϵ k and | v α,i ( ¯ t k ) -z k i | < ¯ ϵ k . For small enough ¯ ϵ k this means that sgn( u α,i ( ¯ t k )) = sgn( v α,i ( ¯ t k )) = +1 . Now let t ∗ = min( t k +1 , inf { t > ¯ t k : v α,i ( t ) = 0 } ) . Since u α,i ( t ) > v α,i ( t ) without loss of generality (see Section C.3), we have sgn( u α,i ( t )) = sgn( v α,i ( t )) = +1 for all t ∈ [ ¯ t k , t ∗ ] . So du α,i ( t ) dt , dv α,i ( t ) dt < 0 for all t ∈ [ ¯ t k , t ∗ ] . So, for any t ∈ [ ¯ t k , t ∗ ] ,

Also, since log α ( w α,i ( t ∗ )) ≈ 1 , by Lemma C.6 we have t ∗ > c > 0 for some constant c independent of α . So for all t ∈ [ t ∗ , t k +1 ] we have b k i -τg i ( θ k ) ∈ ( c, 2 -c ) for some constant c > 0 . So | u α,i ( t ) | , | v α,i ( t ) | ≤ α c/ 4 for all t ∈ [ t ∗ , t k +1 ] . The conclusion follows by triangle inequality.

̸

Coordinate i = i k such that b k i = 2 . The analysis is analogous to the case b k i = 0 , except that we have s k i = -1 instead and g i ( θ k ) > 0 by Lemma D.4.

Finally, we use this conclude that t k +1 ≈ T k +1 and that the weights at coordinate i k are the only weights that change significantly, and by an amount approximately ρ .

Lemma C.9 (Coordinate i k wins the part (A) race at time t k +1 ≈ T k +1 ) . Suppose that T k +1 < ∞ . Then there are ρ k +1 ( ϵ k +1 ) > 0 and α k +1 ( ρ, ϵ k +1 ) > 0 such that for all ρ < ρ k +1 and α < α k +1 , the following holds.

Proof. Let us analyze the case that b k i k ∈ (0 , 2) . Notice that b k +1 i k = b k i k -∆ k ( i k ) g i k ( θ k ) ∈ { 0 , 2 } and that if b k +1 i = 0 then g i k ( θ k ) > 0 and if it is 2 then b k +1 i k = g i k ( θ k ) < 0 . So by Lemma C.6, for all times t ∈ [ ¯ t k , min( t k +1 , T k +1 -ϵ k +1 )] , we have w α,i k ( t ) ∈ ( c, 2 -c ) for some c > 0 . So for small enough α by Lemma D.1, | u α,i k ( t ) | , | v α,i k ( t ) | ≤ α c/ 2 . Combining this with Lemma C.8, we see that for t ∈ [ ¯ t k , min( t k +1 , T k +1 -ϵ k +1 )] we have

for small enough α . So by definition of t k +1 we must have t k +1 > T k +1 -ϵ k +1 . Combined with Lemma C.7, we conclude that | T k +1 -t k +1 | < ϵ k +1 , which is the first claim of the lemma. Furthermore, by Lemma C.8,

so by definition of t k +1 and triangle inequality we have | u α,i k ( t k +1 ) | + | v α,i k ( t k +1 ) | ≥ 4 ρ -ρ/ 2 = 7 ρ/ 2 . Also, since u 2 α,i k ( t k +1 ) -v 2 α,i k ( t k +1 ) = Θ( α 2 ) we have u α,i k ( t k +1 ) ∈ [ ρ, 3 ρ ] . Finally, if b k +1 i k = 2 , then s k +1 i k = -1 and log α ( w α,i k ( t k +1 )) > 1 . 5 so sgn( v α,i k ( t )) < 0 by Lemma D.3; analogously, if b k +1 i k = 0 , we have s k +1 i k = 1 and log α ( w α,i k ( t k +1 ) < 0 . 5 so sgn( v α,i k ( t k +1 ) > 0 .

The case b k i k ∈ { 0 , 2 } can be proved similarly to the analysis in Lemma C.8, where one shows that during the first period of time the magnitudes of | u i k ( t ) | and | v i k ( t ) | decrease, until the sign of v i k flips and they once again increase.

We have shown the claims (21), (22), (23) (24), and (25) for the time t k +1 . In fact, if we let t ′ k +1 ∈ [ ¯ t k , ∞ ) be the first time t such that u α,i k ( t ) = ρ we still have (21), (22), (23) and (25) by the same analysis as above, and (24) can be replaced with the slightly more convenient

C.6.2 Analysis in case where T k +1 = ∞

In this case that T k +1 , we just have to show that the weights remain close to θ k . We show that for any ϵ k +1 > 0 , there is α k +1 ( ϵ k +1 ) > 0 such that for all α < α k +1 and times t ∈ [ T k + ϵ k +1 , T ∗ ] ,

We can use Lemmas C.5 and C.6, which were developed for the case of T k +1 < ∞ , but still hold for T k +1 = ∞ . Lemma C.5 guarantees that the weights do not move much until time t k +1 , and so we only need to show that t k +1 ≥ T ∗ when we take ρ small enough. For this, observe that g i ( θ k ) = 0 for all i ̸∈ S k , because otherwise T k +1 < ∞ . Therefore Lemma C.6 guarantees that until time min( T ∗ , t k +1 ) all weights are close to the original on the logarithmic scale. Namely,

Furthermore, by the non-degeneracy Assumption 4.3 we know that b k i ∈ (0 , 2) for all i ̸∈ S k by Lemma D.4. So if we take ρ small enough and α k +1 small enough, we must have that t k +1 ≥ T ∗ .

C.7 Dynamics from time t k to time ¯ t k (Nonlinear evolution for O (1) unrescaled time)

Suppose that we know for some k ≤ K that for any ϵ k > 0 , there is ρ k ( ϵ k ) > 0 such that for all ρ < ρ k there is α k ( ρ, ϵ k ) > 0 such that for all α < α k , there is a time t k = t k ( α, ρ, ϵ k ) satisfying

Now we will show that for any ¯ ϵ k > 0 , there is ¯ α k = ¯ α k (¯ ϵ k ) > 0 such that for all 0 < α < ¯ α k , there is a time ¯ t k = ¯ t k ( α, ¯ ϵ k ) satisfying

We give the construction for ¯ t k . For any desired accuracy ¯ ϵ k > 0 in this stage, we will construct an accuracy ϵ k = ϵ k (¯ ϵ k ) = ¯ ϵ k / 3 > 0 . Wewill also construct a ρ = ρ ( ϵ k ) > 0 which is sufficiently small, and we will construct an cutoff for α equal to ¯ α k = ¯ α k +1 (¯ ϵ k ) > 0 which satisfies ¯ α k < α k ( ρ, ϵ k ) . The values for these parameters ϵ k and ρ and ¯ α k will be chosen in the following lemma, and will depend only on ¯ ϵ k .

Lemma C.10 (New local minimum reached in time O (1 / log(1 /α )) ) . For any ¯ ϵ k > 0 , we can choose ¯ α k = ¯ α k (¯ ϵ k ) > 0 small enough so that, for any 0 < α < ¯ α k , there is ¯ t k = ¯ t k ( α, ¯ ϵ k ) for which conditions (36) to (38) hold.

Furthermore, there is a constant C ′′ independent of α such that | θ α ( t ) | / | θ α ( t k ) | ∈ [1 /C ′′ , C ′′ ] 2 p at all times t ∈ [ t k , ¯ t k ] .

Proof. Let t k = t k ( α, ρ, ϵ k ) be given by the induction. Let us compare the dynamics starting at θ α ( t k ) with the dynamics starting at ˜ θ ( t k ) = (˜ u ( t k ) , ˜ v ( t k )) which is given by

and run with

Define

By Assumption 4.5 we know there exists a unique solution ˜ θ : [ t k , ∞ ) → R p as long as we take ϵ k small enough because supp( ˜ θ ( t k )) = S k -1 ∪ { i k -1 } and ∥ ˜ θ i ( t k ) -θ k -1 ∥ < ϵ k . Furthermore, by Assumption 4.5 if we take ϵ k small enough there must be a time τ := τ (¯ ϵ k , ρ ) < ∞ such that

So for α small enough, | T k -¯ t k | < 2 ϵ k < ¯ ϵ k , proving (36).

We now compare θ α ( ¯ t k ) with ˜ θ ( ¯ t k ) , and show that if we take α small enough, then the dynamics of ˜ θ closely match the dynamics of θ α ( t ) for times t k + O (1 / log(1 /α )) . The argument uses Gronwall's inequality. Let t ∗ = inf { t > t k : ∥ ˜ θ ( t ∗ ) -θ α ( t ) ∥ > 1 / 3 } . For times t ∈ [ t k , t ∗ ) by Lemma D.7 we have

where K ˜ θ ( t ) is the smoothness constant from Lemma D.7. Note that since ∥ ˜ θ ( t ) ∥ < ∞ for large enough t by (39), the trajectory of ˜ θ must lie in a compact set. Therefore, there must be a finite set of times s 1 , . . . , s m ∈ [ t k , t ∗ ) such that ∪ t ∈ [ t k ,t ∗ ) B ( ˜ θ ( t ) , 1 / 2) ⊆ ∪ i m =1 B ( ˜ θ ( s i ) , 3 / 4) . So letting C = max i m =1 K ˜ θ ( s i ) < ∞ for all times t ∈ [ t k , t ∗ ) we have

By Gronwall's inequality, for all times t ∈ [ t k , t ∗ ) ,

We know from Lemma C.8 that there is a constant c > 0 such that for any small enough 0 < α < α k , such that

If we take α small enough that α c exp( Cτ ) < ¯ ϵ k / 2 < 1 / 3 , we must have t ∗ > t k + τ/ log(1 /α ) and so we prove (37)

It remains to show that (38) is satisfied. Since ∥ ˜ θ ( t ) -θ α ( t ) ∥ < 1 / 3 for all t ∈ [ t k , ¯ t k ] , it holds that the trajectory of θ α ( t ) lies in a compact set. So by Lemma D.7 we have ∥ g ( θ α ( t )) ∥ < C ′ for some constant C ′ at all times t ∈ [ t k , ¯ t k ] . Since 1 log(1 /α ) | dw α,i dt | = | w α,i ( t ) || g i ( w α ( t )) | < C ′ | w α,i ( t ) | , we must have | w α,i ( t ) | / | w α,i ( t k ) | ∈ [1 /C ′′ , C ′′ ] for some constant C ′′ independent of α and all t ∈ [ t k , ¯ t k ] . Therefore, (38) follows from (33). A similar argument shows that | θ α ( t ) / θ α ( t k ) | ∈ [1 /C ′′ , C ′′ ] 2 p .

C.8 Concluding the proof of Theorem C.4

We have shown that Theorem 4.1 is true for solutions θ α : [0 , T ∗ ] → R 2 p to the gradient flow, where T ∗ ∈ ( T K , T K +1 ) . To establish Theorem C.4 it remains only to show that for any T ∗ ∈ ( T K , T K +1 ) and small enough α such a solution to the gradient flow exists and is unique. To see this, note that in the inductive proof of the invariants we construct a sequence of times 0 = ¯ t 0 ≤ t 1 ≤ ¯ t 1 ≤ · · · ≤ ¯ t K ≤ t K +1 > T ∗ , where we guarantee that any gradient flow solution θ α : [0 , t k +1 ] → R p satisfies θ α ∈ ∪ k ∈{ 0 ,...,K } B ( θ k , 1) for all t ∈ ∪ k ∈{ 0 ,...,K } [ ¯ t k , t k +1 ] . And also for t ∈ ∪ k ∈{ 0 ,...,K -1 } [ t k , ¯ t k +1 ] , we have θ α ( t ) ∈ B (0 , C ′′ k θ k ) for some constant C ′′ k independent of α by Lemma C.10. So θ α ( t ) ∈ B (0 , C K ) for some constant C K at all times t ∈ [0 , T ∗ ] . By Lemma D.7, the loss gradient ∇ θ L ( θ ) = ( v ⊙ g ( θ ) , u ⊙ g ( θ )) is Lipschitz-continuous on the compact set B (0 , C K ) . So θ α : [0 , T ∗ ] → R p exists and is unique by the Cauchy-Lipschitz theorem.

D Technical lemmas

D.1 Relating the sum of the weights to the original weights using the conservation law

Lemma D.1. If for some constant 0 < c < 1 we have log α ( w α,i ( t )) ∈ ( c, 2 -c ) , then for small enough α

Proof. Let ˜ w α ( t ) = u α ( t ) -v α ( t ) . By the conservation law (5), w α,i ( t ) ˜ w α,i ( t ) = w α,i (0) ˜ w α,i (0) = u α,i (0) 2 -v α,i (0) 2 . By the non-degeneracy of initialization (Assumption 4.3), the right-hand-side is Θ( α 2 ) . So if log α ( w α,i ( t )) ∈ ( c, 2 -c ) then for small enough α , we have log α ( | ˜ w α,i ( t ) | ) ∈ (3 c/ 4 , 2 -3 c/ 4) . So | u α,i ( t ) | ≤ | w α,i ( t ) + ˜ w α,i ( t ) | ≤ α c/ 2 and | v α,i ( t ) | ≤ | w α,i ( t ) -˜ w α,i ( t ) | ≤ α c/ 2 .

Lemma D.2. If for some constant 0 < c we have log α ( w α,i ( t )) ̸∈ ( -c, 2+ c ) , then for small enough α ,

Proof. Define ˜ w α = u α -v α as in the proof of Lemma D.1. If log α ( w α,i ( t )) < -c then log α ( | ˜ w α,i ( t ) | ) > 2 -c/ 2 for small enough α , so u i ( t ) > α -c -α 2 -c/ 2 > 1 . Similarly, if log α ( w α,i ( t )) > 2 + c then log α ( | ˜ w α,i ( t ) | ) < -c/ 2 so | u i ( α ) | > α -c/ 2 -α 2+ c > 1 .

LemmaD.3. If for some constant c > 0 , there is small enough α such that if we have log α ( w α,i ( t )) > 1 + c then sgn( v α,i ( t )) < 0 . Otherwise, if log α ( w α,i ( t )) < 1 -c then sgn( v α,i ( t )) > 0 .

Proof. Follows from v α = 1 2 ( w α -˜ w α ) . Recall that w α ( t ) > 0 and notice that ˜ w α ( t ) > 0 . In the first case, w α,i ( t ) < α 1+ c and ˜ w α,i ( t ) > α 1 -c/ 2 . In the latter case w α,i ( t ) > α 1 -c and ˜ w α,i ( t ) < α 1+ c/ 2 .

D.2 Sign of gradients on coordinates that leave support

Lemma D.4. For any k ≥ 1 and i ∈ S c k , if b k i ∈ { 0 , 2 } then we must have i ∈ supp( u k -1 ) \ supp( u k ) , and we must have g i ( u k ) < 0 if b k i = 0 and g i ( θ k ) > 0 if b k i = 2 . In particular, ∆ k ( i k ) > 0 for all k .

Proof. This is 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 ℓ , h , and g .

Lemma D.5. The function ℓ ( y , · ) is locally Lipschitz and smooth in its second argument: for any R > 0 , there exists K R such that for any ζ , ζ ′ ∈ B (0 , R )

almost surely over y . Here Dℓ ( y , · ) ⊤ ∈ R d out is the derivative in the second argument.

Proof. Since ℓ is continuously twice-differentiable, for each y ∈ R d y , ζ ∈ R d out there is K y , ζ < ∞ such that for all y ∈ B ( y , 1 /K y , ζ ) and ζ ′ ∈ B ( ζ , 1 /K y , ζ ) we have

where Dℓ and D 2 ℓ denote the first and second derivative in the second argument. So for all such y ′ ∈ B ( y , 1 /K y , ζ ) and ζ ′ , ζ ′′ ∈ B ( ζ , 1 /K y , ζ ) we have

Cover the set { ( y , ζ ) : ∥ y ∥ ≤ C, ∥ ζ ∥ ≤ R } with the balls ∪ y B ( y , 1 /K y , ζ ) . By compactness, there is a finite subcover ( y 1 , ζ 1 ) , . . . , ( y r , ζ r ) , so we can take K R = max i ∈ [ r ] K y i , ζ i < ∞ and the lemma holds since ∥ y ∥ ≤ 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 > 0 there exists K R such that for any ψ , ψ ′ ∈ B (0 , R ) ,

almost surely over x . Here Dh ( x , · ) ∈ R d out × R p is the derivative in the second argument.

Proof. Analogous to proof of Lemma D.5, using continuous twice-differentiability of h and boundedness of ∥ x ∥ .

Lemma D.7 (Local Lipschitzness of loss and loss derivative) . When θ = ( u , v ) ∈ R 2 p and f NN ( x ; θ ) = h ( x ; u ⊙ u ) the following holds for g ( θ ) defined in (4) . For any R > 0 , there exists K R < ∞ such that for any θ , θ ′ ∈ B (0 , K R ) ,

Proof. Let θ = ( u , v ) , θ ′ = ( u ′ , v ′ ) . This follows immediately from the local Lipschitzness and smoothness of h and ℓ in Lemmas D.5 and D.6, as well as

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n1. **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\nWEAKNESSES:\n1. **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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

comments: 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.

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

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-09"}

comments: 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:
i. diagonal weight matrices
ii. initialization is small
iii. only one coordinate is in "active" regime at a time
that dynamics occur in discrete stages:
(1) during most of each stage, the loss plateaus because the weights remain close to a saddle point
(2) at the end, the saddle point is quickly escaped and the rank of the weights increases by at most one

The developed theory is illustrated with a series of experiments.

STRENGTHS:
1. Nice and clean results that show us the dynamics of neural networks with attention mechanism

2. Well connected with the previous research on training dynamics

3. Paper is well written

WEAKNESSES:
Assumptions are quite restrictive. It would be nice to relax the assumptions, especially assumption #3.

QUESTIONS:
1. Is the paper more about attention in general than specifically about transformers?

2. What will happen if some weights activate simultaneously? How the theory would be affected?

LIMITATIONS:
Limitations are adequately discussed
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-26