Title: Weight decay induces low-rank attention layers
Venue: NeurIPS 2024
Accepted: yes
Keywords: weight decay

attention
Abstract: The effect of regularizers such as weight decay when training deep neural networks is not well understood. We study the influence of weight decay as well as $L2$-regularization when training neural network models in which parameter matrices interact multiplicatively. This combination is of particular interest as this parametrization is common in attention layers, the workhorse of transformers. Here, key-query, as well as value-projection parameter matrices, are multiplied directly with each other: $W_K^TW_Q$ and $PW_V$. 
We extend previous results and show on one hand that any local minimum of a $L2$-regularized loss of the form $L(AB^\top) + \lambda (\|A\|^2 + \|B\|^2)$ coincides with a minimum of the nuclear norm-regularized loss $L(AB^\top) + \lambda\|AB^\top\|_*$, and on the other hand that the 2 losses become identical exponentially quickly during training. We thus complement existing works linking $L2$-regularization with low-rank regularization, and in particular, explain why such regularization on the matrix product affects early stages of training.
Based on these theoretical insights, we verify empirically that the key-query and value-projection matrix products $W_K^TW_Q, PW_V$ within attention layers, when optimized with weight decay, as usually done in vision tasks and language modelling, indeed induce a significant reduction in the rank of $W_K^TW_Q$ and $PW_V$, even in fully online training.
We find that, in accordance with existing work, inducing low rank in attention matrix products can damage language model performance, and observe advantages when decoupling weight decay in attention layers from the rest of the parameters.

Weight decay induces low-rank attention layers

Abstract

The effect of regularizers such as weight decay when training deep neural networks is not well understood. We study the influence of weight decay as well as L 2 -regularization when training neural network models in which parameter matrices interact multiplicatively. This combination is of particular interest as this parametrization is common in attention layers, the workhorse of transformers. Here, key-query, as well as value-projection parameter matrices, are multiplied directly with each other: W T K W Q and PW V . We extend previous results and show on one hand that any local minimum of a L 2 -regularized loss of the form L ( AB ⊤ ) + λ ( ∥ A ∥ 2 + ∥ B ∥ 2 ) coincides with a minimum of the nuclear normregularized loss L ( AB ⊤ ) + λ ∥ AB ⊤ ∥ ∗ , and on the other hand that the 2 losses become identical exponentially quickly during training. We thus complement existing works linking L 2 -regularization with low-rank regularization, and in particular, explain why such regularization on the matrix product affects early stages of training. Based on these theoretical insights, we verify empirically that the key-query and value-projection matrix products W T K W Q , P W V within attention layers, when optimized with weight decay, as usually done in vision tasks and language modelling, indeed induce a significant reduction in the rank of W T K W Q and PW V , even in fully online training. We find that, in accordance with existing work, inducing low rank in attention matrix products can damage language model performance, and observe advantages when decoupling weight decay in attention layers from the rest of the parameters.

1 Introduction

The influence of L 2 -regularization, as well as weight decay regularization when training deep neural network models remains poorly understood and is still a subject of active research [van Laarhoven, 2017, Zhang et al., 2021, 2019, Loshchilov and Hutter, 2019, Zhang et al., 2021, Xie et al., 2023, Andriushchenko et al., 2023]. Given a model parametrized by matrix W , the standard motivation of adding λ 2 ∥ W ∥ 2 to the optimization loss L ( W ) comes from framing learning the model weights W as maximum a posteriori (MAP) estimation and choosing a Gaussian prior with zero mean [Mackay, 1995, Krogh and Hertz, 1991].

Previous works have studied the effect of regularization on the rank of weight matrices when training a model with gradient-based optimization [Ziyin and Wang, 2023, Arora et al., 2019, Li et al., 2021, Razin and Cohen, 2020, Gunasekar et al., 2017]. Here, we focus on the effect of L 2 -regularization on models using a factorized parametrization, where some weight matrices are parametrized as products

of (often lower rank) matrices, W = AB ⊤ . This parametrization is used heavily in attention layers inside transformers [Vaswani et al., 2017] which we will focus on in the following.

Indeed, at the heart of the Transformer architecture is the attention operation which updates the T tokens concatenated into a matrix E ∈ R d m × T inside the network according to

where ϕ is typically a softmax operation applied column-wise and M is typically a causal mask. The matrices W V , W K , W Q ∈ R d k × d m are respectively the value, key, and query matrices that linearly transform E into some typically smaller space of dimension d k [Phuong and Hutter, 2022], which can potentially subsume bias terms by appending a constant 1 to the tokens. The weight matrix P ∈ R d m × d k projects the weighted sum of value vectors back into the original token dimension. Therefore (multi-head) attention layers indeed consist of parameter matrix products i.e. W QK = W ⊤ K W Q as well as W V P = PW V , regardless of the choice of ϕ , or the presence or absence of causal masks.

When optimizing neural network models with this particular parametrization in conjunction with L 2 -regularization, and for any such two weight matrices A and B (e.g. the P and W V for a given layer and a given head), we can rewrite the loss as:

where θ accounts for all the remaining parameters. We will see in the following that optimizing such losses has in practice implications on regularizing the rank of W = AB ⊤ . In fact, while it is classically known that the summed Frobenius norm 1 2 ( ∥ A ∥ 2 + ∥ B ∥ 2 ) is a tight upper bound on the nuclear norm ∥ AB ⊤ ∥ ∗ [Srebro and Shraibman, 2005, Tibshirani, 2021], we theoretically show in the following that gradient-based optimization of the above objective result in the upper bound becoming tight exponentially quickly, for arbitrary loss, and thus directly optimizes for the nuclear norm which is known to induce low rank.

We highlight the relevance of this study since high weight decay is commonly used when training Transformer models. For example, GPT-3 [Brown et al., 2020], LLaMa [Touvron et al., 2023], LLaMa 2 [Touvron et al., 2023] and ViT [Dosovitskiy et al., 2021] report a weight decay strength of λ = 0 . 1 . Interestingly, this is even true when fine-tuning, for example with low-rank adaptation (LoRA) [Hu et al., 2021a].

We summarize our contributions below:

We show that for models with factorized parametrization, all local minima of any loss regularized by the Frobenius norm of A,B coincide with local minima of the same loss regularized by the nuclear norm of W . We further show theoretically that the discrepancy between the 2 regularizations vanishes exponentially quickly during training, thus implying that training such models with weight regularization can be subjected to low rank inducing pressure long before convergence.

We empirically validate our result on various experimental settings, including when optimization with decoupled weight decay [Loshchilov and Hutter, 2019], on models ranging from deep linear networks to language models as well as Vision Transformers Dosovitskiy et al. [2021]. Intriguingly, we observe that this inductive bias of factorized parametrization with weight decay seems to hurt the performance on some tasks, raising the question of whether it is a feature or a bug.

We provide evidence suggesting that this rank-regularizing effect in fact seems to affect the pretraining of popular pre-trained foundation models such as LLAMA 2 [Touvron et al., 2023] and Vision Transformer [Wu et al., 2020], by analyzing their pre-trained weights.

2 Related Work

The setting we study is closely related to a setting extensively studied in the Matrix Completion literature [Srebro and Shraibman, 2005, Sun and Luo, 2016, Candes and Tao, 2009], where the goal is to recover an unknown low-rank matrix for which only a subset of its entries are specified. Nuclear norm regularization is often used as a convex relaxation of the problem [Hu et al., 2021b], and

its equivalence at the global optimum with the L 2 -regularization on factorized matrix [Srebro and Shraibman, 2005], which has the advantage of being differentiable everywhere, has been exploited as a popular approach for large-scale matrix completion. Extensive prior work has focused in this setting on the theoretical guarantee of the factorization formulation to recover the underlying low-rank matrix correctly [Sun and Luo, 2016, Candes and Tao, 2009]. Similarly, similar loss landscape analyses were performed in the context of unconstrained features models [Zhu et al., 2021]. In contrast, our analysis does not rely on assumptions about the data, the loss (other than its differentiability) or convergence.

In a different line of work, recent efforts have focused on the effect of gradient-based optimization of deep networks on the parametrized matrix. For example, small weight initialization in this setting was shown to induce low rank in deep linear networks [Jacot et al., 2022, Arora et al., 2019, Li et al., 2021]. More recently, [Jacot, 2023] has shown the representation cost of deep networks with homogeneous nonlinearity converges to a notion of rank over nonlinear functions. More related to our work, equivalence between L 2 regularization applied on factorized matrices and a low-rank inducing L p -Schatten norm on the matrix they parametrize has been shown in several prior works [Dai et al., 2021, Tibshirani, 2021]. This is particularly relevant as L 2 regularization can be applied explicitly or implicitly, such as when training deep networks with homogeneous activation coupled with e.g. the cross entropy loss [Jacot et al., 2022, Arora et al., 2019]. Crucially, however, these existing works characterize the low-rank inducing bias on neural networks that globally minimize L 2 regularization while fitting training data.

Recently, [Galanti et al., 2023] have studied the effect of SGD with L2-regularization on a general architecture. Similarly to our work, they consider a general differentiable loss, but bound the rank of matrices at sufficiently large training steps, employing a theoretical argument that crucially does not leverage low-rank inducing norms due in part to the generality of the architecture they consider. [Wang and Jacot, 2023] have studied the same effect in the context of deep fully connected linear networks, showing that SGD strengthens the already existing low-rank bias induced by L2-regularization, albeit on matrix completion problems. Similarly to our work, they draw for the first time, to the best of our knowledge, an equivalence between the critical points of L2-regularized loss on the factorized matrix and Nuclear norm regularized loss on the parametrized matrix.

In contrast to these past works, we show both theoretically and empirically that for any arbitrary differentiable loss, the two regularizations become exponentially quickly identical during gradientbased optimization, and thus, that the low-rank inducing effect comes into play very early in during training. This brings a theoretical understanding to empirical observations made in previous works [Khodak et al., 2022], and is particularly relevant for many practical settings, in which learning does not converge, such as foundation model trained online, as is commonly done for large language models (LLMs) and large vision models.

Finally, given the significance of self-attention models, there has been work trying to understand the implicit inductive biases of some of their design choices. [Bhojanapalli et al., 2020] shows, in particular, the head size heuristic commonly used causes a low-rank bottleneck and limits the expressive power of the multi-head attention layer. Recent work has shown indeed that reducing the rank of attention matrices post-training of LLMs can hurt downstream performance [Sharma et al., 2023]. Our empirical work complements these observations and sheds light on the potentially damaging effect of the implicit rank-reducing effect of weight decay in the context of Attention layers, an unintended side effect contrary to the matrix completion setting.

3 Theoretical results

3.1 Preliminaries

We begin by reviewing the definition of the nuclear norm of a matrix and its upper bound when applied to a factorized matrix. We denote by ∥ · ∥ the Frobenius norm when applied on matrices.

3.1.1 Nuclear norm

The nuclear norm (also known as trace norm) of a real-valued matrix W , denoted by ∥ W ∥ ∗ , is defined as √

When using the singular value decomposition (SVD) of W , W = USV ⊤ , denoting ( s i ) i the singular values, we can see that

i.e. the nuclear norm is the sum of the singular values of W .

The nuclear norm is often used in the low-rank regularization literature [Hu et al., 2021b] as it intuitively is a convex relaxation of the rank, and regularizing it typically induces low rank by injecting sparsity in the singular values.

3.1.2 Upper bound of the nuclear norm of a factorized matrix

Let two matrices A,B such that W = AB ⊤ . Then, using the Cauchy-Schwarz inequality, we have that

3.1.3 Considered losses

We will consider L 2 losses of the format

As a consequence of the above inequality, the L 2 -regularized objective (8) is an upper bound of the nuclear norm-regularized objective.

The meticulous reader should spot that those objectives don't account for the remaining parameters θ as in (2), while those parameters also evolve through learning. In fact, one can convince oneself that this can be safely ignored without loss of generality. The reader is referred to appendix D for more details about this point.

3.2 Equivalence of optimization solution

In the following, we will first show that in fact, any objective of the form in (8) will coincide at any stationary point with the nuclear-norm regularized loss in (9), thus introducing a low-rank inducing bias in the solution found. We assume A,B to have a bottleneck, i.e. to have the number of rows greater or equal to the number of columns, as is usual in attention layers. All proofs can be found in Appendix B.

We start by providing a sufficient condition under which the averaged Frobenius norm of two matrices would correspond to the nuclear norm of their product.

Proposition 3.1. Let A,B be matrices such that A ⊤ A = B ⊤ B . Then, denoting AB ⊤ = USV ⊤ the SVD of AB ⊤ , there exist an orthogonal matrix O such that A = U ( √ S 0 ) O ⊤ and B = √

This condition states that the scalar product of any two columns of A should match the scalar product of corresponding columns of B . We will show next that at any stationary point of the objective L L 2 , that condition is fulfilled. We assume the loss L is differentiable and λ > 0 .

Lemma 3.2. At any stationary point A,B of L L 2 we have that A ⊤ A = B ⊤ B .

and their L ⋆ counterpart

The above Lemma, together with Proposition 3.1, implies that at a stationary point A,B , L L 2 ( A,B ) and L ∗ ( AB ⊤ ) coincide. However, this is not enough to claim that finding a (local) minimum of L L 2 will in fact find a (local) minimum of L ∗ . We now provide a result which shows that this claim is true.

Theorem 3.3. A,B is a local minimum of L L 2 , if and only if 1) W = AB ⊤ is a local minimum of L ∗ , constrained to matrices of rank r where r is the maximum rank achievable by AB ⊤ , and 2) A ⊤ A = B ⊤ B .

A more general formulation of the above results, albeit without a bottleneck dimension, was recently shown in [Wang and Jacot, 2023] (c.f. Theorem 3.1) where it was applied in the matrix completion context. We restate and prove it here for completion in the context of self-attention and transformer models.

The Theorem states that there is in fact a one-to-one mapping between the local minima of L L 2 ( A,B ) and (the equivalence class of) local minima of L ∗ ( AB ⊤ ) , for a general unregularized loss L .

In particular, if one wishes to optimize L ∗ for some matrix W , potentially under rank constraint, one can reparametrize W as a product of two matrices A,B and optimize the differentiable objective L L 2 on A,B without introducing bad minima, and obtain rank-regularized solutions. In principle, one can still converge to a bad minimum for a general loss, but this is not due to the reparametrization.

On the other hand, the theorem shows that naively optimizing the L 2 -regularized loss with a factorized parametrization will (often inadvertently) result in actually finding solutions that exactly minimize the nuclear-norm regularized loss, introducing unintended low-rank inducing bias to the solution.

Note that however, the two parametrizations may result in different optimization, and thus different solutions, even if the loss landscape shares the same local minima.

3.3 Optimization dynamic in the gradient flow limit

The above result establishes equivalence of the local minima of the two losses. Our next result shows that the two losses will in fact coincide exponentially quickly during training.

Theorem 3.4. Consider the gradient flow limit over the loss L L 2 . If ∥ A ∥ , ∥ B ∥ remain bounded during training, then we have that ∣ ∣ L L 2 ( A,B ) -L ∗ ( AB ⊤ ) ∣ ∣ converges exponentially to 0 .

In order to prove the theorem, we first show that during gradient flow optimization, the condition from Proposition 3.1 becomes true exponentially quickly. This is then followed by a new bound bounding the gap between ∥ AB ⊤ ∥ ∗ and 1 2 ( ∥ A ∥ 2 + ∥ B ∥ 2 ) by the norm of A ⊤ A -B ⊤ B . For completeness, we also provide a general result bounding the analogous gap for a L -layer deep linear network.

We provide in the appendix a similar result when considering gradient flow with noise, as well as with momentum and decoupled weight decay. We furthermore provide in appendix B.5 a discussion about the soundness condition.

The above result complements Theorem 3.3 by showing that optimizing L L 2 will result in cooptimizing L ∗ very quickly during training, long before stationary points are found. The theorem also confirms previous empirical observations [Khodak et al., 2022].

3.4 Case study: 2-layer linear network

To illustrate the low-rank inducing bias of the factorized parametrization coupled with weight decay, we will study in the following the optimization within a 2-layer linear network and characterize the network at equilibrium. Such a network corresponds in fact to a drastically simplified softmax attention layer with T = 1 . The derivations are similar to those used when studying deep linear networks [Ziyin et al., 2022, Saxe et al., 2013] and the redundant parameterization studied in [Ziyin and Wang, 2023].

Consider the following model

where B ⊤ ∈ R d 2 × d 1 , A ∈ R d 3 × d 2 . For simplicity of presentation, we assume d 3 = d 1 = d 1 , 3 , but the result can be easily extended to the general case. Given D data points ( x i , y i ) 1 ≤ i ≤ D , in matrix

0°1

1.0

1.50

IS

AB

100

100

0.8

0.8

1.25

S2

S3

2

B

+

10-3

109

106

司

三

B

A

102

10-4

9°0

0.4

0.2

0°0

入

0.6

0.4

0.75

1.00

0.50

S5

二

1012

10-6

0.2

0.25

10-15

0

500

t

1000

10-8

0

500

t

1000

0.0

0.0

0.5

入

1.0

0.00

100

101

t

102

103

form, the L 2 -regularized mean squared error can be expressed as

where X = ( x i ) i ∈ R d 1 × D , Y = ( y i ) i ∈ R d 3 × D , and λ > 0 .

Using full batch gradient flow, the differential equation governing the parameter dynamic becomes

where Σ Y X = Y X ⊤ , Σ XX = XX ⊤ , and τ is a constant controlling the learning rate.

To further simplify the above equations, we follow [Saxe et al., 2013] and assume Σ XX = I , an assumption which holds exactly for whitened input data. Finally, without loss of generality, we perform a change of basis such that Σ Y X = S where S is the diagonal matrix which diagonal consists of the singular values ( s i ) i ∈ [1 ..d 1 , 3 ] of Y X ⊤ .

At equilibrium, we thus have the following set of equations

Denoting by a i , b i the i -th row of A,B , and assuming the ( s i ) i ∈ [1 ..d 1 , 3 ] are all non-zero and distinct, we have the following conditions at equilibrium (cf Appendix C)

̸

In particular, this implies, for any i , λ ∥ a i ∥ 2 = ( s i - ∥ a i ∥ 2 ) ∥ a i ∥ 2 . Clearly, if λ ≥ s i , then the equation can only be true if a i = 0 . If on the other hand λ < s i , either a i = 0 or ∥ a i ∥ 2 = s i -λ satisfy the equilibrium condition, with the former being an unstable equilibrium point if the number of hidden units d 2 is greater than the number of elements in { i ∈ [1 ..d 1 , 3 ] | s i > λ } .

To highlight the result, let us consider the case where the hidden layer has enough capacity, i.e. d 2 ≥ d 1 , 3 . In that case, the result tells us that at a stable equilibrium, AB ⊤ will drop all singular values s that are less than λ , while keeping those that are larger. In other words, it performs a sort of low rank approximation of the input-output correlation matrix where the rank is controlled by λ . A related result was already obtained in the analyses of [Saxe et al., 2013] who studied the exact

solutions of 11 without the regularization term but introducing a bottleneck in the hidden layer, i.e. d 2 < d 1 , 3 . Remarkably here, regularization achieves a similar effect even in an overcomplete network, where increasing λ gradually prunes the hidden neurons to ignore the smallest variations of the data, i.e. reducing d 2 adaptively. We confirm these results empirically in Figure 1.

Importantly, this result is only obtained because the regularization is applied to the parametrization involving a matrix multiplication. If AB ⊤ were replaced by a single matrix W ∈ R d 1 , 3 × d 1 , 3 , then the equilibrium condition would be W = 1 1+ λ S , whose rank remains constant w.r.t. the regularization strength.

3.5 Weight decay with Adam optimizer

While the regularized loss is a convenient setting for studying what happens to the parameters at equilibrium, in the vast majority of practical settings, decoupled weight decay [Loshchilov and Hutter, 2019], simply referred to as weight decay in the following, is used instead optimizing a regularized loss. A popular choice of optimizer for deep neural networks, including those with self-attention layers, is AdamW [Loshchilov and Hutter, 2019], which update the weights by using the Adam optimizer on the non-regularized loss while simultaneously applying weight decay.

While it is non-trivial to analyze the equilibrium points of AdamW in general, we show in Appendix E that under some simplifying assumptions, they coincide with those of a L 2 -regularized loss with a different regularization strength.

4 Empirical results

0.6

0°0 =

入=0.01

W

N

1^M

10

2N

一

20

入=0.3

入 = 0.03

入=0.1

MI

rank

0.4

5

10

P

+

0

25

50

75

100

0

0

25

50

75

100

→2

0

0

25

50

75

100

t

t

t

The primary objective of our experimental analysis is to empirically validate the theoretical findings in more practical settings. Specifically, we aim to investigate the effect of decoupled weight decay, adaptive optimizers, as well as noisy gradient and lack of exact convergence to stationary points on the theoretical findings.

The second objective is to establish that the theory is relevant in the training of large neural network models. Due to the large computational costs we chose to avoid re-training large scale models but trained small-scale language models as well as a Vision Transformer without changing common hyperparameters. We aim to demonstrate that their typical training is affected by the rank-regularizing effect predicted by our theory. Finally, we investigate pre-trained weights of the relevant foundation models to show that they are consistent with rank-regularizing training.

To quantitatively measure the rank of matrices in the context of our experiments with attention layers, we use the following definition of pseudo rank : Let W be a weight matrix with singular values σ 1 , σ 2 , ..., σ n , ordered such that σ 1 ≥ σ 2 ≥ ... ≥ σ n . The pseudo rank (referred to simply as rank in the following) of W is defined as k n where k is the smallest number such that:

In simpler terms, it represents the fraction of the largest singular values required to capture at least 95% of the sum of all singular values of the matrix W .

rank

0.90

0.85

0.80

0.75

A = 0.1

A = 0.0

A = 0.01

A = 0.03

入=0.3

0.80

0.90

0.75

0.85

入 = 0.0

入=0.01

入 = 0.03

A = 0.1

入=0.3

rank

WoK

0.9

0.8

0.7

0.6

0.5

0.4

rank

0.9

0.8

0.7

0.6

0.5

A= 0.0

A= 0.01

A = 0.1

入= 0.25

0.70

0.70

0.3

0.4

0.65

0.65

0.2

0.3

0.60

0.1

0.2

0

20000

40000

60000

80000

0

20000

40000

60000

80000

0

F

10

15

20

25

30

0

5

10

15

20

25

30

Steps

Steps

Epoch

Epoch

4.1 Associative recall task

In this simple memory task, a model is presented with a sequence of paired tokens [ x 1 , y 1 , . . . , x T , y T , x T +1 ] . Specifically, the task is parameterized by an integer N , representing the number of unique tokens that can be mapped to N corresponding tokens. The sequence presented to the model therefore consists of 2 N +1 tokens (with T = N ), and the final token is repeated and appears in the sequence before, i.e. x T +1 = x j for some j ∈ [0 , . . . , T ] . The model is trained to remember the correct association observed in-context and predict y j . This task has been attributed and proposed as a proxy for language modelling [Fu et al., 2023, Poli et al., 2023].

We train a 2-layer self-attention only Transformer with AdamW optimizer on minibatches of size 128, for N = 20 . To simulate additional noise, we perturb 5% of the labels with random labelling [Zhang et al., 2021].

Figure 4 shows that even in this setting, the stationary condition of a L 2 -regularized loss in Lemma 3.2 is approached, and the gap between the nuclear norm and the Frobenius norm in (5) vanishes, thus confirming that AdamW in fact also optimizes for the nuclear norm. Furthermore, the convergence speed is perfectly correlated with the weight decay strength. The results furthermore show that AdamW leads indeed to a consistent decrease in the rank in both parameter weight products as the decay strength increases. This aligns with the effect of optimizing the nuclear norm of these matrices.

4.2 Language Modelling

In order to validate our theoretical findings in larger scale experiments, we now present results when training standard small scale Transformer models, with 125 million parameters, on the Pile [Gao et al., 2020] - a common language modeling dataset. All design decisions such as the Transformer architecture as well as the optimizer and training schedule are identical to the ones proposed in the GPT-3 paper [Brown et al., 2020], which are now used in various other studies e.g. [Fu et al., 2023, von Oswald et al., 2023]. Details can be found in the Appendix G.

First, we confirm again that increasing weight decay with AdamW drastically reduces the rank of W T K W Q as well as PW V , on average across depth and heads, of the trained models (c.f. Figure 3.

Appendix Figure 7). Furthermore, we observe that while increasing the weight decay strength of MLP beyond 0.1 is generally beneficial, see Table 1, doing the same for attention matrices starts slightly hurting performance. Results are averaged over 3 seeds. We observe a sweet spot around weight decay strength of 0.1 applied to self-attention weights, indicating that some rank regularization is beneficial for this task. Nevertheless, too much weight decay and therefore rank regularization seems to be detrimental. Finally, applying weight decay to the MLP weights seems to more important with a generally higher effect on performance. We leave a more nuanced investigation of decoupling the weight decay strength of matrices affected by our theory from the rest of the parameters for future research.

4.3 Vision Transformers

Next, we focus on computer vision tasks and train a Vision Transformer on the ImageNet dataset [Deng et al., 2009] for 24 hours, following the exact training protocol of [Irandoust et al., 2022]. We follow the previous section and vary the decay strength only in the attention layers, while keeping every other hyperparameter fixed. We observe a similar effect of the decay strength on the ranks of the matrices W QK , W V P (c.f. Fig 3).

4.4 Pretrained foundation models

Finally, we turn to pre-trained foundation models, and provide some evidence that their training is also impacted by the rank-regularizing effect of weight decay. Specifically, following Proposition B.4, it is sufficient to observe that the matrices W Q W ⊤ Q resp. P ⊤ P are close to W K W ⊤ K resp. W V W ⊤ V . Because the matrices W Q , W K , W V , P ⊤ are typically wide rectangular matrices, the off-diagonal elements of W Q W ⊤ Q , etc, are mostly 0 . For L L 2 to approximately correspond to L ∗ , it thus suffices that the diagonal elements of W Q W ⊤ Q resp. P ⊤ P are close to those of W K W ⊤ K resp. W V W ⊤ V .

6

2.5

009

300

30

5

2.0

500

250

M

3

4

P

P

1.5

1.0

Ml

300

400

200

150

20

layer

2

200

100

10

1

0.5

100

50

0

2

4

6

0.0

0

WWT

1

0

0

Wol + WkI12

200

400

600

0

0

100

200

300

Q

5 Discussion

Our results provide further insights into the interplay between L 2 -regularization and weight decay regularization and the optimization of models that consist of parameter matrix products. This is of particular interest since attention layers in transformer exhibit this parametrization as key-query, as well as value-projection parameter matrices, are multiplied directly with each other: W T K W Q and PW V . Our empirical findings strongly support our theoretical predictions about the impact of weight decay on the rank of attention layers and clearly show a rank-regularizing effect even without convergence. We provide evidence that the training of some foundation models such as Llama are in fact in practice affected by the same regularization.

Furthermore, we find that decoupling weight decay in the attention weights and tuning its weight decay strength can improve performance, for example in our language modelling experiments. These findings complement the recent observation that reducing the rank of language model MLP matrices post-training improves their reasoning performance, while doing the same for attention layer matrices mostly hurt it [Sharma et al., 2023]. In particular, our findings suggest that the conventional practice of applying uniform regularization strategies across all layers may not be optimal for other deep learning architectures as well. This finding opens up new avenues for model- or layer-specific regularization strategies that could significantly enhance the performance of these models.

Our findings once more highlight the complexity of understanding optimization techniques in conjunction with particular neural network models, particularly transformers. For example, the difficulty of understanding the effect when varying regularization strengths on different components of these models underscores the need for a more nuanced theoretical understanding of layer-specific regularization. We are particularly excited about further research that aims to disentangle the role of weight decay in in-weight vs. in-context learning within MLPs and self-attention layers, building on [Singh et al., 2023]. In conclusion, while our findings mark a step forward in understanding and improving the usage of weight decay when training deep neural networks, in particular transformers, our study shed light on the intricate interplay of neural network regularization and its parametrization.

Acknowledgments

Seijin Kobayashi, Yassir Akram and Johannes von Oswald deeply thank João Sacramento and Angelika Steger for their support and guidance. The authors also thank Robert Obryk, Moritz Firsching, Luca Versari, Nicolas Zucchet, Alexander Meulemans, Simon Schug, Blaise Agüera y Arcas, Alexander Mordvintsev, Ettore Randazzo and Eyvind Niklasson for fruitful discussions.

References

Twan van Laarhoven. L2 regularization versus batch and weight normalization, 2017.

Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Commun. ACM , 64(3):107-115, feb 2021. ISSN 0001-0782. doi: 10.1145/3446776. URL https://doi.org/10.1145/3446776 .

Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger Grosse. Three mechanisms of weight decay regularization. In International Conference on Learning Representations , 2019. URL https://openreview.net/forum?id=B1lz-3Rct7 .

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2019. URL https://openreview.net/forum?id= Bkg6RiCqY7 .

Zeke Xie, zhiqiang xu, Jingzhao Zhang, Issei Sato, and Masashi Sugiyama. On the overlooked pitfalls of weight decay and how to mitigate them: A gradient-norm perspective. In Thirty-seventh Conference on Neural Information Processing Systems , 2023. URL https://openreview.net/ forum?id=vnGcubtzR1 .

Maksym Andriushchenko, Francesco D'Angelo, Aditya Varre, and Nicolas Flammarion. Why do we need weight decay in modern deep learning?, 2023.

David J C Mackay. Probable networks and plausible predictions - a review of practical bayesian methods for supervised neural networks. Network: Computation in Neural Systems , 6(3):469-505, 1995.

Anders Krogh and John A. Hertz. A simple weight decay can improve generalization. In Neural Information Processing Systems , 1991. URL https://api.semanticscholar.org/CorpusID: 10137788 .

Liu Ziyin and Zihao Wang. spred: Solving l1 penalty with SGD. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 43407-43422. PMLR, 23-29 Jul 2023.

Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. CoRR , abs/1905.13655, 2019. URL http://arxiv.org/abs/1905.13655 .

Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=AHOs7Sm5H7R .

Noam Razin and Nadav Cohen. Implicit regularization in deep learning may not be explainable by norms, 2020.

Suriya Gunasekar, Blake Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nathan Srebro. Implicit regularization in matrix factorization, 2017.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2017.

Mary Phuong and Marcus Hutter. Formal algorithms for transformers, 2022.

Nathan Srebro and Adi Shraibman. Rank, trace-norm and max-norm. In Peter Auer and Ron Meir, editors, Learning Theory , pages 545-560, Berlin, Heidelberg, 2005. Springer Berlin Heidelberg. ISBN 978-3-540-31892-7.

Ryan J Tibshirani. Equivalences between sparse models and neural networks. Working Notes. URL https://www. stat. cmu. edu/ryantibs/papers/sparsitynn. pdf , 2021.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. arXiv preprint arXiv:2005.14165 , 2020.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.

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. In International Conference on Learning Representations , 2021. URL https://openreview. net/forum?id=YicbFdNTTy .

Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. CoRR , abs/2106.09685, 2021a. URL https://arxiv.org/abs/2106.09685 .

Bichen Wu, Chenfeng Xu, Xiaoliang Dai, Alvin Wan, Peizhao Zhang, Zhicheng Yan, Masayoshi Tomizuka, Joseph Gonzalez, Kurt Keutzer, and Peter Vajda. Visual transformers: Token-based image representation and processing for computer vision, 2020.

Ruoyu Sun and Zhi-Quan Luo. Guaranteed matrix completion via non-convex factorization. IEEE Transactions on Information Theory , 62(11):6535-6579, November 2016. ISSN 1557-9654. doi: 10.1109/tit.2016.2598574. URL http://dx.doi.org/10.1109/TIT.2016.2598574 .

Emmanuel J. Candes and Terence Tao. The power of convex relaxation: Near-optimal matrix completion, 2009.

Zhanxuan Hu, Feiping Nie, Rong Wang, and Xuelong Li. Low rank regularization: A review. Neural Networks , 136:218-232, 2021b. ISSN 0893-6080. doi: https://doi.org/10. 1016/j.neunet.2020.09.021. URL https://www.sciencedirect.com/science/article/ pii/S089360802030352X .

Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. A geometric analysis of neural collapse with unconstrained features, 2021.

Arthur Jacot, François Ged, Berfin ¸ Sim¸ sek, Clément Hongler, and Franck Gabriel. Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity, 2022.

Arthur Jacot. Implicit bias of large depth networks: a notion of rank for nonlinear functions, 2023.

Zhen Dai, Mina Karzand, and Nathan Srebro. Representation costs of linear neural networks: Analysis and design. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 26884-26896. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper_files/paper/ 2021/file/e22cb9d6bbb4c290a94e4fff4d68a831-Paper.pdf .

Tomer Galanti, Zachary S. Siegel, Aparna Gupte, and Tomaso Poggio. Characterizing the implicit bias of regularized sgd in rank minimization, 2023.

Zihan Wang and Arthur Jacot. Implicit bias of sgd in l 2 -regularized linear dnns: One-way jumps from high to low rank, 2023.

Mikhail Khodak, Neil Tenenholtz, Lester Mackey, and Nicolò Fusi. Initialization and regularization of factorized neural layers, 2022.

Srinadh Bhojanapalli, Chulhee Yun, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar. Low-rank bottleneck in multi-head attention models, 2020.

Pratyusha Sharma, Jordan T. Ash, and Dipendra Misra. The truth is in there: Improving reasoning in language models with layer-selective rank reduction, 2023.

Liu Ziyin, Botao Li, and Xiangming Meng. Exact solutions of a deep linear network. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id=X6bp8ri8dV .

Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks, 2013. URL https://arxiv.org/abs/1312.6120 .

Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: an 800GB dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027 , 2020.

Saghar Irandoust, Thibaut Durand, Yunduz Rakhmangulova, Wenjie Zi, and Hossein Hajimirsadeghi. Training a vision transformer from scratch in less than 24 hours with 1 gpu, 2022.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models, 2023.

Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models, 2023.

Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Blaise Agüera y Arcas, Max Vladymyrov, Razvan Pascanu, and João Sacramento. Uncovering mesa-optimization algorithms in transformers, 2023.

Aaditya K. Singh, Stephanie C. Y. Chan, Ted Moskovitz, Erin Grant, Andrew M. Saxe, and Felix Hill. The transient nature of emergent in-context learning in transformers, 2023.

Paul Lévy. Sur certains processus stochastiques homogènes. Compositio Mathematica , 7:283-339, 1940. URL http://www.numdam.org/item/CM_1940__7__283_0/ .

Robert T Powers and Erling Størmer. Free states of the canonical anticommutation relations. Communications in Mathematical Physics , 16(1):1-33, 1970.

Diederik P. Kingma and Jimmy Ba. Adam: a method for stochastic optimization. In International Conference on Learning Representations , 2015.

Richard H. R. Hahnloser, Rahul Sarpeshkar, Misha A. Mahowald, Rodney J. Douglas, and H. Sebastian Seung. Digital selection and analogue amplification coexist in a cortex-inspired silicon circuit. Nature , 405(6789):947-951, 2000.

A Compute budget

We estimate the total compute budget to 4 Nvidia RTX 4090 for two months. The LLMs were punctually trained on a cluster of 16xA100 GPUs for 4 days.

B Proofs of theoretical results

B.1 Proof of Proposition 3.1

Proof. Using singular value decomposition, we write A = U A Σ A V ⊤ A and B = U B Σ B V ⊤ B , with Σ A = ( S A 0 ) and Σ B = ( S B 0 ) . Substituting in the equation A ⊤ A = B ⊤ B , we get

By positivity and uniqueness of singular values, necessarily, Σ A = Σ B = ( S 0 ) . Furthermore, by rearranging the above equation, we get S 2 V ⊤ A V B = V ⊤ A V B S 2 , i.e. that V ⊤ A V B commutes with S . We rewrite A as

Redefining U A as U A ( V ⊤ A V B 0 0 I ) , setting Σ = Σ A , and V = V B , we can write A = U A Σ V ⊤ and B = U B Σ V ⊤ .

In particular, AB ⊤ = U A ΣΣ ⊤ U ⊤ B , which is a valid SVD of AB ⊤ . It remains to show that, if AB ⊤ is diagonal, then there exists an orthogonal matrix O such that A = Σ O ⊤ and B = Σ O ⊤ .

Let us assume the diagonality, i.e. AB ⊤ = ΣΣ ⊤ . Then, we have

i.e. that U A , U B commute with ΣΣ ⊤ , and thus that they are block diagonal. Furthermore, ΣΣ ⊤ U A U ⊤ B = ΣΣ ⊤ . They can then be written as

where U, U ′ A , U ′ B are orthogonal matrices, and the block of U corresponds to the non zero singular values of ΣΣ ⊤ .

We can then rewrite A,B as A = Σ ( U 0 0 I ) V ⊤ and B = Σ ( U 0 0 I ) V ⊤ , which conclude the proof by setting O = ( U 0 0 I ) V ⊤ .

Finally, AB ⊤ = U A ΣΣ ⊤ U ⊤ B , and therefore ∥ AB ⊤ ∥ ∗ = ∥ U A ΣΣ ⊤ U ⊤ B ∥ ∗ = Tr(ΣΣ ⊤ ) = 1 2 ( ∥ A ∥ 2 + ∥ B ∥ 2 ) .

B.2 Proof of Lemma 3.2

Proof. Let A,B a stationary point of the unregularized loss L in L L 2 . One can show that the gradient of L ( W = AB ⊤ ) with respect to A (resp. B ) is

where ∂L ∂W ∣ ∣ W = AB ⊤ is a matrix, which we denote by -G . Differentiating L , at the stationary point, the following equations must then be satisfied

In particular, A ⊤ A = 1 λ A ⊤ GB = 1 λ ( G ⊤ A ) ⊤ B = B ⊤ B .

B.3 Proof of Theorem 3.3

Proof. ( ⇐ ) We start by proving the backward implication, by contradiction. Let M a local minimum of L ∗ , and A,B such that M = AB ⊤ and A ⊤ A = B ⊤ B . Then by Proposition 3.1, L L 2 ( A,B ) = L ∗ ( M ) . Assume A,B is not a local minimum of L L 2 , i.e. there exists an infinitesimally perturbed matrices A ′ , B ′ such that L L 2 ( A ′ , B ′ ) < L L 2 ( A,B ) . By continuity of matrix multiplication, M ′ = A ′ B ′⊤ is an infinitesimally perturbed matrix M . Since L ∗ ( M ′ ) ≤ L L 2 ( A ′ , B ′ ) < L L 2 ( A,B ) = L ∗ ( M ) , we get a contradiction.

( ⇒ ) Assume now that A,B is a local minimum of L L 2 , and that W = AB ⊤ is not a local minimum of L ∗ constrained to rank r matrices. Then, we can construct a sequence ( W n ) n of rank r matrices such that lim n →∞ W n = W , and for all n , L ∗ ( W n ) < L ∗ ( W ) . For all n , let W n = U n S n V ⊤ n the SVD of W n . By continuity of the mapping from a matrix to its singular values, lim n →∞ S n = S , where S are the singular values of W . Because the set of orthogonal matrices is compact, there exists a subsequence of (( U n , V n )) n which converges to some orthogonal matrices ( U, V ) . Without loss of generality, we redefine the sequence to this converging subsequence. By continuity of matrix multiplication, necessarily USV ⊤ = W . USV ⊤ is a valid SVD of W . Since by local minimality of A,B , following Lemma 3.2 and Proposition 3.1, we get that A = U Σ O ⊤ and B = V Σ O ⊤ √

V n Σ n O ⊤ , where Σ n = ( S n 0 ) . Then, lim n →∞ ( A n , B n ) = ( A,B ) and yet, because for all n , A ⊤ n A n = B ⊤ n B n and A n B ⊤ n = W n , we have L L 2 ( A n , B n ) = L ∗ ( W n ) < L ∗ ( W ) = L L 2 ( A,B ) . This is a contradiction.

where Σ = ( S 0 ) and O is some orthogonal matrix. Let for all n , A n = U n Σ n O ⊤ and B n = √

B.4 Proof of Theorem 3.4

In order to prove the theorem, we first show that during optimization, the condition from Proposition 3.1 becomes true exponentially quickly. This is then followed by a new bound bounding the gap between ∥ AB ⊤ ∥ ∗ and 1 2 ( ∥ A ∥ 2 + ∥ B ∥ 2 ) by the norm of A ⊤ A -B ⊤ B .

B.4.1 Exponential decay of A ⊤ A -B ⊤ B

We provide the result for the vanilla gradient flow limit, but also provide an alternative proof for the stochastic gradient flow with momentum and decoupled weight decay, to illustrate that the exponential decay would hold in many practical settings. We note that the gradient flow limit is a good approximation for a small learning rate in the discrete dynamic.

Lemma B.1. In the gradient flow limit over the loss L L 2 , A ⊤ A -B ⊤ B will converge exponentially to 0 .

Proof. For any i , we denote by a i , b i the i -th column of A and B . The columns follow the following differential equations:

where G = -∂L ∂W ∣ ∣ W = AB ⊤ , and τ is some time constant controlling the learning rate. Given a pair i, j , we can now look at the dynamic of a i ⊤ a j -b i ⊤ b j :

Therefore, we have A ⊤ A -B ⊤ B = Qe -2 λt τ , where Q is A ⊤ A -B ⊤ B at initialization, and in particular, every entry of A ⊤ A -B ⊤ B converge to 0 exponentially.

We now provide a similar result, in the gradient flow regime but with momentum, as well as decoupled weight decay - a tractable approximation to AdamW as shows the following proposition:

Proposition B.2. We consider the following dynamics approximating stochastic gradient flow with weight decay:

where µ, η, σ > 0 and W A and W B are independent matrix Wiener processes. Initial condition are H A 0 = 0 and H B 0 = 0 . H A (resp. H B ) is the momentum gradient with respect to A (resp. B ). Then,

Proof. We have such that

Note that the second term is an abuse of notation. The derivation of the integral form of H B t , A t and B t follows the same logic.

For the last two equations, we get

which gives

Before analyzing the implications of this proposition, let us state a lemma that will allow us to bound the probability of a Brownian motion diverging:

Lemma B.3. [Lévy, 1940] Let ( B t ) be a 1D Wiener process. Then, for t, L > 0 , P [max s ∈ [0 ,t ] B s > L ] = 2 P [ B t > L ] .

Let us now analyse the consequences of Proposition B.2. We will assume that A t and B t remain L2-bounded by M > 0 (which is true for all converging dynamic modulo steady state noise), and that G t remains L2-bounded by K (either using a Lipschitzian loss, or using clipping).

First observe that, using lemma B.3, for ε > 0 , with probability 1 -ε , the term √ µ 2 σW A 1 -e -2 µt and the correspond B will remain bounded by σ √ µnd ln 4 nd ε .

This way, H A and H B will with probability 1 -ε remain bounded by KM + σ √ µnd ln 4 nd ε . With that same probability, the term

will remain bounded by 4 ηM λ ( KM ++ σ √ µnd ln 4 nd ε ) . This is the same order of magnitude as the stochastic term. Until A ⊤ A -B ⊤ B is of that order, it exponentially decays.

B.4.2 Upper bound of ∥ ∥ ∥ AB ⊤ ∥ ∗ -1 2 ( ∥ A ∥ 2 + ∥ B ∥ 2 ) ∥ ∥

Finally, we provide the following general result, bounding the gap between ∥ AB ⊤ ∥ ∗ and 1 2 ( ∥ A ∥ 2 + ∥ B ∥ 2 ) by the norm of A ⊤ A -B ⊤ B .

Proposition B.4. For any matrices A,B , we have

In particular,

Proof. Let Q := A ⊤ A -B ⊤ B . Using singular value decomposition, we write A = U A Σ A V ⊤ A and B = U B Σ B V ⊤ B , with Σ A = ( S A 0 ) and Σ B = ( S B 0 ) . Substituting in the previous equation, we get

where V B ∆ V ⊤ B := √ V B S 2 B V ⊤ B + Q -√ V B S 2 B V ⊤ B . By the Powers-Stormer inequality [Powers and Størmer, 1970], we have ∥ ∆ ∥ 2 F = ∥ V B ∆ V ⊤ B ∥ 2 F ⩽ ∥ Q ∥ ∗ . From there, we rewrite A as

Consequently, ∥ AB ⊤ ∥ ∗ = ∥ S 2 B +∆ S B ∥ ∗ and

Similarily, we have ∣ ∣ ∥ AB ⊤ ∥ ∗ -∥ A ∥ 2 F ∣ ∣ ⩽ √ ∥ Q ∥ ∗ ∥ A ∥ ∗ .

The second inequality is obtained by using the triangle inequality.

B.4.3 Upper bound of ∥ ∥ ∥ ∥ ∏ l A l ∥ 2 /L 2 /L -1 L ∑ ∥ A l ∥ 2 F ∥ ∥ ∥

We here provide a more general version of Proposition 3.1.

Proposition B.5. Let q ≥ r > 0 , A 1 ∈ M q,r , A l ∈ M r,r for l ∈ [2 ..L -1] , A L ∈ M q,r and A = ∏ L l =1 A l . Let ε > 0 . We assume that the sequence ( A l ) l ∈ [1 ..L ] is ε -balanced, i.e. that for l ∈ [1 ..L -1] ,

i.e.

Furthermore, assume that

Then for k ∈ [1 ..L ] , we have

In particular,

Proof. We will assume for this proof that the A l are square matrices for all l ∈ [1 ..L ] . This proof holds as it is for the more general case, with more cumbersome notations, and with the trick used in equation 33. For l ∈ [1 ..L ] , we denote by A l = U l Σ l V ⊤ l be the SVD of A l . Let k ∈ [1 ..L ] . We first show by recurrence that we can write

for an appropriate choice of orthogonal matrices ( O l ) l ∈ [0 ..L ] and symmetric matrices (∆ l ) l ∈ [1 ..L ] , and such that : O k -1 = U k , O k = V k and ∥ ∆ l ∥ F ⩽ √ ε | l -k | for l ∈ [1 ..k ] . The recurrence is symmetric for l ⩽ k and for l ⩾ k . We will prove it in the latter case. For l = k , the statement holds. Let l ∈ [ k..L -1] . We assume we can write A l = O l -1 (Σ k +∆ l ) O l ⊤ , with ∥ ∆ l ∥ F ⩽ √ ε ( l -k ) . Let Q := A ⊤ l A l -A l +1 A ⊤ l +1 . We have

Similar to the previous proof, we can write

We set O l +1 = V l +1 U ⊤ l +1 O l , ∆ l +1 = ∆ l +∆ to get the desired result. We verify that ∥ ∆ l +1 ∥ F ⩽ √ ε ( l -k +1) .

Now that we have proven our lemma, let us observe that

Notice that ∥ ∥ Σ L k ∥ ∥ 2 /L 2 /L = ∥ A k ∥ 2 F . Using the triangular inequality of A →∥ A ∥ p p for 0 < p ≤ 1 , we have

with ∆ symmetric verifying

We can rewrite with

Furthermore, using the fact that ∥ AB ∥ 2 /L ≤ ∥ A ∥ 2 /L ∥ B ∥ F , we have

where in the penultimate line, we used (1 + x ) n -1 ≤ enx whenever x < 1 n , and using the assumption of the proposition.

Ultimately, we thus obtain

which concludes the proof for the first inequality. The second inequality is obtained by using the triangular inequality.

Note that the above proposition can be used to argue that optimizing a deep linear network of depth L will very quickly co-optimize the L p -Schatten norm of the product, with an exponential time constant λ/L . In fact, even in a deep linear network, it can be shown that the matrices become exponentially balanced using the same proof as in Lemma B.1. For the assumption ε ≤ 1 L 4 min l ∥ A l ∥ L ∗ to hold, it suffices that the matrix norm remains lower-bounded by a strictly positive value, a reasonable assumption for loss functions of interest. Note however that this assumption is used only to obtain the convenient upper bound expression, but the exponential decay is trivially true even without it.

B.5 On the boundness condition of A and B

We now examine the assumption that both A and B are bounded. It can be shown that, under stochastic dynamics with momentum and for certain loss types, A and B will remain bounded with high probability. Below, we present two examples of sufficient conditions on the loss function to ensure this boundedness. Although one could formulate an expanding set of such sufficient conditions to cover a broader class of losses that lead to bounded parameters, there are still practical scenarios where certain losses might not satisfy these conditions. In these cases, practitioners often achieve stable dynamics through hyperparameter tuning, and our theorem remains applicable, as it is designed to accommodate these scenarios rather than exclude them. Therefore, the boundedness assumption is broadly relevant in practice, particularly in training setups that use weight decay, and we leverage this assumption to discuss the rank-regularization effect that impacts such training processes.

B.5.1 Sufficient conditions

Henceforth, we will refer by θ = ( A,B ) .

Sufficient condition 1: Gradient flow with lower bounded loss function We consider a the following gradient flow dynamics with weight decay:

η is the learning rate hyperparameter, and λ the weight decay strength.

A sufficient condition on the loss is that it is lower bounded, which is the case for most common losses.

Indeed, the above dynamic is the gradient flow dynamic of the loss L ′ ( θ ) := L ( θ ) + λ ∥ θ ∥ 2 . Given that L ′ is also lower bounded, and that L ′ ( θ t ) is a monotonically decreasing function of time, L ′ ( θ t ) must converge to a constant real value, i.e. L ( θ t ) + λ ∥ θ t ∥ 2 → t →∞ c for some c . If ∥ θ ∥ is unbounded from above, then necessarily L is unbounded from below, which is a contradiction.

Sufficient condition 2: Gradient flow with momentum with Lipschitz gradient We consider a the following gradient flow dynamics with momentum and decoupled weight decay:

where G is the exponential average of the gradient of θ . µ is the momentum hyperparameter, η the learning rate, and λ the weight decay strength.

A sufficient condition on the gradient is to be min(1 , ηλ µ ) -Lipschitz with respect to the parameters θ sufficiently far, i.e. for ∥ θ ∥ > P for a given P :

The momentum makes the analysis of the dynamics more complicated. However, defining F =

[ ] one can rewrite the equations as:

The derivative of the squared norm of F verifies:

The dynamics of F are flow dynamics. Whenever ∥ F ∥ reaches P , its norm is decreasing. ∥ F ∥ can thus never exceed P . As ∥ F ∥ is an upperbound on ∥ θ ∥ , the same holds for ∥ θ ∥ . We also observe that the Lipschitz condition doesn't need to hold for all θ > P . In fact, it suffices that it holds for any borderless submanifold of codim -1 (for example the sphere of radius M) that contains the initialization point.

Note on stochasticity To deal with stochasticity, we consider similar equations:

which become:

The integral form is:

The process can diverge because of the stochasticity. However, similarly to proposition B.2, we can fix for any δ > 0 an upper bound on W that holds with probability at least 1 -δ ; this allows to bound the contribution of the stochasticity. We deal with the gradient component similar to the non-stochastic proof.

B.5.2 Pathological examples where the boundedness does not hold

An example of a loss for which the parameters will diverge is when we allow losses to be negative, and diverge to minus infinity "stronger" than the weight regularization term.

An obvious, albeit constructed such loss is L ( AB ⊤ ) = -∥ AB ⊤ ∥ 2 . Then, the gradient of L w.r.t. A (resp. B ) is ∇ A L = -AB ⊤ B (resp. ∇ B L = -BA ⊤ A ). Even with the decay term, one can see that if A,B are initialized to be e.g. orthogonal matrices scaled by some α > λ , both A,B will diverge to infinity.

Such negatively unbounded objective functions to be minimized may be found in e.g. the reinforcement learning setting, when using undiscounted returns.

C Equilibrium condition of 2-layer Linear network

We assume λ > 0 and that the singular values of Y ⊤ X are all non-zero and distinct.

We start with the following set of equations:

Clearly, (34) implies and thus that B ⊤ B = A ⊤ A .

Furthermore, (34) also implies

Using B ⊤ B = A ⊤ A , we get that AA ⊤ S = SBB ⊤ .

Denoting by a i , b i the i -th row of A,B , for any i, j ∈ [1 ..d 1 , 3 ] 2 , we have s i a ⊤ i a j = s j b ⊤ i b j and s j a ⊤ j a i = s i b ⊤ j b i . Thus, ∥ a i ∥ 2 = ∥ b i ∥ 2 , and a ⊤ i a j = b ⊤ i b j = 0 , since s i , s j are distinct and positive. Taken together, AA ⊤ is a diagonal matrix, which we denote by D . We have D = diag (( ∥ a i ∥ 2 ) i ∈ [1 ..d 1 , 3 ] ) .

̸

Using the positivity of s i and λ , one can see that necessarily, a i = b i .

In particular, (38) implies λAB ⊤ = D ( S -AB ⊤ ) , i.e. ( λI + D ) AB ⊤ = DS . Because the entries of D are positive, ( λI + D ) is invertible, and thus AB ⊤ = ( λI + D ) -1 DS . In other words, the off-diagonal entries of AB ⊤ are zero, i.e. a ⊤ i b j = 0 for all i = j , i, j ∈ [1 ..d 1 , 3 ] 2 .

In particular, for a given i , we have

D On the link between the full loss and the restricted loss

In the theoretical section, we study a pruned version of the total losses:

where the remaining parameters are not accounted for. This in fact still accounts for the general case. Indeed:

stationary points of L L 2 will be also stationary in ( A,B ) , hence Lemma 3.2 still holds. In fact, the latter condition suffices.

For theorem 3.3, the same proof B.3 shows that ( A,B,θ ) is a local minima of L L 2 iff 1) ( W = AB ⊤ .θ ) is a local minima of L ∗ constrained to matrices W of rank r and 2) A ⊤ A = B ⊤ B

in the proof B.4.1 of theorem 3.4, the gradient G of the first lemma now hides a dependence in θ , but the proof still hold as is. In the second lemma, the gradient G t depends on the time, and is indirectly hiding a dependence on θ .

E Link between solutions of AdamW and L 2 -regularization

Consider the following dynamic induced by AdamW, with λ > 0 :

where η represents the learning rate and β 1 , β 2 , ε are the common hyperparameters of Adam, W t is the parameter at time t , and where the various operations are applied element-wise.

Note that the term -λW t stems from weight decay.

If the dynamic converges, then necessarily, G ∞ = ∇ W L ( W ∞ ) , B ∞ = ( ∇ W L ( W ∞ )) 2 , and thus λW ∞ = -∇ W L ( W ∞ ) |∇ W L ( W ∞ ) | + ϵ . Clearly, this implies that λ | W ∞ | < 1 . If we further assume that λ | W ∞ |≪ 1 , then the condition becomes ϵλW ∞ ≈ -∇ W L ( W ∞ ) , which is the equilibrium point of a L 2 -regularized loss with regularization strength ϵλ 2 . Thus, the stationary points of the AdamW optimizer can in practice correspond to stationary points of L 2 -regularized loss, and thus the same low-rank inducing solutions can be found.

Weshow in Fig. 5 a toy experiments illustrating the equivalence in the solutions found by AdamW with decay strength λ WD and hyperparameter ϵ , with those found by Adam with L 2 -regularization with regularization strength λ L 2 = λ WD ϵ . In particular, we illustrate how a factorized parametrization in this setting will still result in solutions that minimize the nuclear norm, even when trained with AdamW.

F Pretrained foundation models: ViT

We provide in Figure 6 the ViT counterpart of Figure 4.

0.06

0.04

0.02

0.00

0.02

0.00

0.05

W1

0.06

0.04

0.02

0.00

-0.02

0.00

0.05

W1

S

min L2

min L1

入L2, ^d = 1e  07

入L2, d =1e 06

入2,f = 1e  05

入L2,^d =0.0001

入L2,^d = 0.001

入L2, d =0.01

入L2, ^dE = 0.1

入L2,ud = 1.0

12

10

12

500

400

8

10

350

10

TK

6

P

8

400

M

300

8

4

T

P

6

4

300

200

250

elayer

4

2

2

200

150

2

8.0

2.5

5.0

7.5

10.0

0

0

WWT

5

10

200

lWol2 + WkI12

300

400

500

100

100

PI2 + WI12

200

300

400

G Language modelling experimental details

Here, we present details of our language modeling experiments, employing standardized values from the literature and consistent, untuned hyperparameters across all trials. Unless specified otherwise, we utilize the conventional GPT-2 transformer architecture with LayerNorm (Ba et al., 2016), incorporating MLPs between self-attention layers and applying skip-connections after each layer. Training is conducted using a standard (autoregressively) masked cross-entropy loss, omitting an input embedding layer but incorporating an output projection before computing logits. Further details can be found in Table 2.

0.90

A = 0.0

A = 0.01

0.90

A = 0.0

入 = 0.01

0.90

A=0.0

入 = 0.01

0.90

A = 0.0

A = 0.01

0.85

0.80

A= 0.1

入=0.3

A = 0.03

0.85

0.80

入=0.3

A = 0.03

入=0.1

0.85

A =0.03

入=0.1

入=0.3

0.85

入=0.3

A = 0.03

入=0.1

rank

0.75

rank

0.75

rank

W

0.80

rank

0.80

0.70

0.65

0.70

0.75

0.75

0.65

0.60

0.70

0.70

0.60

0

20000

40000

60000

80000

0.55

0

20000

40000

60000

80000

0

20000

40000

60000

80000

0

20000

40000

60000

80000

Steps

Steps

Steps

Steps

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We provide both theoretical and empirical results supporting the main claim of the paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The discussion highlights some of the limitation of this work.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All statements are provided with proofs.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We hope to provide all hyperparameters and experimental details in the appendix and provide code to reproduce most of the experiments.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: We aim to collect the code as soon as possible in a Git repository.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide all experimental details in the main text, as well as the appendix.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer:[No]

Justification: Many of the experiments were performed on large scale models or foundation models, rendering the computation of multiple seeds unrealistic.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide an estimate of compute used.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The paper should be considered a theory and/or conceptual paper. We discussed implication for robust machine learning in the main text, and can not anticipate that the presented results can not conform in any aspect with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The paper shows the importance of choosing a proper weight decay while training transformer.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: No data and models realese.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the effect of weight decay on losses where the trained parameters include two matrices that are multiplies, and specifically the bias of such losses towards low rank. It is shown theoretically that under certain conditions, local minima of the L_2 regularized loss coincide with minima of the L_* loss, regularized by nuclear norm, which is an analog to low rank. Also, it is shown that with gradient flow, the distance between the minima of the L_2 and L_* losses converges exponentially fast to 0.\nSeveral experiments are given, mostly on transformers, that verify these findings empirically.\n\nSTRENGTHS:\nI think that overall this is a nice paper that should be accepted. The results, although pretty simple mathematically, show a nice observation on the bias of certain parameterizations towards low rank. The experiments seem thorough and cover both toy examples with smaller models, and larger-scale transformers.\n\nWEAKNESSES:\nMy biggest issue with the paper is its framing. The theoretical results seem to apply more for compressed sensing or matrix compleletion, where the loss is L(A^\\top B) for some differentiable L. However, the paper (and also the title) seem to discuss mostly transformers. It might be popular to write papers about transformers these days, but this specific paper looks more like a paper about implicit bias in general, with experiments on transformers.\n\nSecond, the discussion about linear models focuses on the underspecified regime since the matrix S is invertible (i.e. the input dimension is larger than the number of samples). This is not the standard regime for linear regressions. I would suggest, to make this discussion more relevant to what happens in practice, is to discuss kernel regimes (e.g. NTK), where it is more standard where the kernel’s dimensions is large than the number of samples.\n\nQUESTIONS:\n- What happens if L in Eq. (2) is non-differentiable? In practice, L it is non-diferentiable for transformers since it incorporates an MLP with a ReLU activation.\n\n- In Figure 3 - what is the performance of each setting? It is not clear for the more realistic values of wight decay (i.e. those with the best performance) what is the rank.\n\n- Line 212 - How do the authors define an unstable equilibrium?\n\nLIMITATIONS:\nThe authors discuss the limitations of their results.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-23"}

comments: SUMMARY:
This paper studies the effect of weight decay on the product of matrices, which appear in the attention layers. This paper shows that weight decay will have an effect of reducing the rank hence hurting the generalization. The theoretical results are verified in extensive experiments.

STRENGTHS:
1. Understanding the attention layers is an important question hence is of great interest.
2. The experiments are extensive and well support the theoretical results.

WEAKNESSES:
The novelty of this paper is quite weak:
- Some theoretical results have already been shown by Wang and Jacot [2023].  Specifically, as mentioned by the authors, Theorem 3.1 in Wang and Jacot [2023] is a general version of Theorem 3.3.
- Theorem 3.4 is for gradient flow which seems far from the algorithms used in practice. 

- The empirical observation that low rank weights hurt the generalization also has been shown in [Sharma et al., 2023]. 

Please correct me if I miss any key contributions in this paper.

QUESTIONS:
1. About Theorem 3.4, how to see that "... long before stationary points are found"?

LIMITATIONS:
The contribution of this paper, as outlined in the weaknesses section, does not seem substantial enough.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-06-29

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper explores how applying weight decay to matrices affects the rank of their product. They show that L2 regularization of the operands is equivalent to regularizing the nuclear norm (sum of singular values) of the product which could result in a lower rank. The attention block of transformers contains several plain matrix products that this theory applies to. The authors experiment with different transformer variants, showing that applying weight decay to these matrices results in a loss of rank which is detrimental to performance. They suggest not applying weight decay to these specific matrices, while keeping it for the rest.\n\nSTRENGTHS:\n+ Interesting topic of high relevance to the community, weight decay is almost universally used for transformer training and improved guidelines and understanding could result in practical performance gains.\n+ The theory looks sound with simple experiments that support the main conclusion.\n+ The empirical evaluation on real transformers supports the theory.\n+ The paper is well written overall although a bit dense at times. The figures are sufficiently clear although a bit small and not in a vector format.\n\nWEAKNESSES:\n- The theory relies on analyzing the converged solution with gradient flow. I’m not sure how well this corresponds to real training (it would be nice to discuss this).\n- The experiments could be stronger to eliminate some potential confounding explanations (see details below).\n- (minor) Missing a couple of related work on weight decay (see details below).\n\nQUESTIONS:\n**Experiments**:\n- The ViT experiments show a much larger loss of rank but the performance impact of this is not quantified, why not?\n- For the GPT experiments the loss of rank seems very small as you point out. I’m not convinced that this is the root cause for the performance loss, rather then some temperature effects for the softmax.\n- Many related works (see below) explore weight decay in terms of effective learning rates. They would suggest that when changing the weight decay the learning rate should be adjusted to compensate. Otherwise it is unclear if the performance degradation truly results from changes in the rank and weight decay or changes in the effective learning rate. Showing that these rank effects occur from the effective learning rate as well would strengthen the results and relate it to existing lines of work.\n- Suggested experiment: Disentangle the softmax temperature effects (from the magnitude of the weight matrices / activations) from the rank effects. Maybe you could repeat the GPT or ViT experiment using a scale-invariant softmax from [1], reporting the rank and performance again. Since this softmax alternative does not depend on the scale of the inputs, it would eliminate this confounding effect.\n- Suggested experiment: Disentangle the effective learning rate effects from the rank effects. When changing the weight decay, keep the weight decay * learning rate product constant. This will keep the effective learning rate in the steady state constant as described by [5] for AdamW, but also change the resulting weight magnitude, so it should be combined with something like the scale-invariant softmax to eliminate that effect.\n\n**Related work**:\n- Overall you provide a good overview of prior weight decay literature but are missing one important line of work that explores weight decay in terms of the effective learning rate. Here are a couple of notable works from this line: [2] [3] [4] [5]\n\n**Other Questions**:\n- In my experience high (effective) learning rates cause a loss of rank (at least certain measures like the stable rank) even in standard matrices (not just the products). Is this something you observe in your experiments (e.g. a loss of rank in A and B that you apply weight decay to, not just the product AB)?\n\n[1]: Li, Zhiyuan, Srinadh Bhojanapalli, Manzil Zaheer, Sashank Reddi, and Sanjiv Kumar. \"Robust training of neural networks using scale invariant architectures.\" In International Conference on Machine Learning, pp. 12656-12684. PMLR, 2022.\n\n[2]: Li, Zhiyuan, and Sanjeev Arora. \"An exponential learning rate schedule for deep learning.\" ICLR 2020.\n\n[3]: Wan, Ruosi, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. \"Spherical motion dynamics: Learning dynamics of normalized neural network using sgd and weight decay.\" Advances in Neural Information Processing Systems 34 (2021): 6380-6391.\n\n[4]: Li, Zhiyuan, Kaifeng Lyu, and Sanjeev Arora. \"Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate.\" Advances in Neural Information Processing Systems 33 (2020): 14544-14555.\n\n[5]: Kosson, Atli, Bettina Messmer, and Martin Jaggi. \"Rotational equilibrium: How weight decay balances learning across neural networks.\" In International Conference on Machine Learning, 2024.\n\nLIMITATIONS:\nThe paper does not really discuss limitations in any depth. I would not expect any particular adverse societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors investigate the landscape of two different optimization problems. For a general objective function $L$ defined on a matrix space, they consider two regularized objectives $\\mathcal L_*(B, A) = L(AB^\\top) + \\lambda ||AB^\\top||_*$, $\\mathcal L_2(B, A) = L(AB^\\top) + \\frac{\\lambda}{2} (||A||_F^2 + ||B||_F^2)$. \n\nThe authors prove that these two objective functions share the same set of critical points up to equivalence. Moreover, along integral lines of $\\nabla \\mathcal L_2$,  the difference $|\\mathcal L_2-\\mathcal L_*|$ decays exponentially if both $||A||$ and $||B||$ remain bounded, effectively showing an implicit bias towards low-rank solutions. These results naturally apply to transformer training, for which the authors present some numerical results.\n\nSTRENGTHS:\nThe work is very well-presented, and the contribution is self-contained. The theoretical analysis covers cases of practical interest. Presenting the results through gradient flows enhances readability.\n\nWEAKNESSES:\n1. Lack of an additional lemma addressing the effect of time discretization in the gradient flow scenario.\n2. The stochastic case is studied through SDEs, which may differ significantly from the exact setting encountered in deep learning.\n\nQUESTIONS:\n1. The balance condition is often enforced directly at initialization and can be proven [1] to be conserved during the continuous-time flow. By the use of proposition 3.1, this would imply that the gradient flow of $\\mathcal L_2$ is **exactly the same gradient flow** of $\\mathcal L_*$. This kind of condition is for example satisfied by spectral initialization, which is commonly used. I believe a comment about this would be interesting to see in the manuscript.\n2. Do the authors see any way to address the two weaknesses? I believe further clarifications in this direction in the manuscript would make the result sound more solid.\n3. I suggest to organize better the references section, to make them all uniform in style.\n\n[1] S.S.Du, W. Hu, J.D. Lee. \"Algorithmic Regularization in Learning Deep Homogeneous\nModels: Layers are Automatically Balanced\", NeurIPS 2018.\n\nLIMITATIONS:\nI have no suggestions for the authors.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors study the role of the weight decay in training of matrices especially when they appear in multiplicative form. First estabilshing the equivalence between the L2 regularization and nuclear norm regularized loss at stationary points and local minima, they establish how even training on the L2 loss invariably leads to the latter and hence low-rank solutions at the minima. Using this main observation, they empirically validate their claim for various practical LLMs and provide interesting insights about their behavior.\n\nSTRENGTHS:\nThe authors consider a well-motivated problem and characterize it mathematically with substantive evidence. Given the impact of transformers and the underlying need to study them fundamentally, such analyses are indeed quite insightful and useful. In particular, a rigorous theoretical anallysis of how low-rank solutions emerge is indeed a very interesting observation.\n\nWEAKNESSES:\nWhile overall I enjoyed reading the paper, I wish the authors did a slightly better job at the following two things which could have made it even more solid:\n\n1. The writing and exposition could be greatly improved and at various places especially concerning mathematical statements, it feels a bit loose though the details are provided in appendix. For example, Theorem 3.4 concerns the gradient flow analysis but neither the flow equation nor the precise statement for the \"L2-nuclear norm\" loss gap going to zero is written mathematically. Since this is an important result, it would be better to present enough mathematical details and expound upon them later. Similarly in Theorem 3.3, regarding the rank-r. It's not clear what this \"r\" is supposed to be. At the optimum or it's something fixed before hand. Things like this should be fixed. Also there are a couple of typos such as Eq.(8) should be a full stop, line 176 it reads better with \"two losses\" rather \"2 losses\" (same comment applies at various places like line 173\" etc.\n\n2. While the L2 loss in equation (2) concerns a pair of matrices (A,B) with the loss of the form L(AB^\\top), given that the main motivation of the authors stems from transformers, it would have been nice to have a discussion about how their analysis extends to Eq.(1), let's say. Here there are two such multiplicative terms with $PW_V$ and $E^\\top W_k^\\top W_Q E$. In this case, can your results about low-rank minima still carry over? Or not? Some insights about this?\n\nQUESTIONS:\nPlease refer to the above weakness part for questions.\n\n\n-------------------------\n\nMinor note: In a recent work, the authors observe a similar low-rank structure when training transformers with Markov chains. They reconstruct the formula for these low-rank matrices and indeed it seems your low-rank optimal conditions seem to be satisfied there (Appendix C I think). Thought I would share with the authors if they find interesting (not to compare as such): https://arxiv.org/abs/2402.04161\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}