Weight decay induces low-rank attention layers
Authors: Seijin Kobayashi
Yassir Akram
Johannes von Oswald
Venue: NeurIPS 2024 poster (2024)
Keywords: weight decay
attention
Primary area: optimization_for_deep_networks

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
Seijin Kobayashi" 122 Yassir Akram"! Johannes von Oswald?
Department of Computer Science, ETH Zürich Google, Paradigms of Intelligence Team
{seijink, jvoswald)@google com, yakram@ethz.ch
Abstract
The effect of regularizers such as weight decay when training deep neural net- works is not well understood. We study the influence of weight decay well as L2-regularization when training neural network models in which parameter matrices interact multiplicatively. This combination is of particular interest this parametrization is common in attention layers, the workhorse of transformers. Here, key-query, well directly with each other: W WQQ and PWv. We extend previous results and show on one hand that any local minimum of L2-regularized loss of the form L(AB") + A(IIAAI2 + |1 coincides with a minimum of the nuclear norm- regularized loss L(AB become identical exponentially quickly during training. We thus complement existing works linking 1.2-reglarrization with low-rank regularization, and in par- ticular, explain why such regularization the matrix product affects early stages of training. Based on these theoretical insights, we verify empirically that the key-query and projecttoonttionn matrix products W W WQ. PWv 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 WTWa Wo and PWv, 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.
calue-projection parameter matrices, multiplied
>||AB ||,, and on the other hand that the 2 losses
Introduction
The influence of L2-regularization, well as weight decay regularization when training deep neural
network models remains poorly understood and is still subject of active research [van Laarhoven
2017 ZZang al. 1 2019 Loshchilov and Hutter 2019 Zhang al. 202T Xie al. 2023
Andrushchennoo et al.) 2003]. Given model parametrized by matrix W, the standard motivation of
adding 2 to the optimization loss L(W) comes from framing learning the model weights W
maximum posteriori (MAP) estimation and choosing 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-
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 12-regularization on models using a factorized parametrization, where some weight matrices are parametrized as products
* Equal contribution, order determined randomly
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

of (often lower rank) matrices. W = -B"". 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 matrix E E IR xT inside the network according to
E - E + PWv Wx O M)
where é is typically a softmax operation applied column-wise and M is typically causal mask. The matrices Wr, Wx, Wq E Rd xd.. are respectively the value, key, and query matrices that linearly transform E into some typically smaller space of dimension dg [Phuong and Hutter. [022], which can potentially subsume bias terms by appending a constant to the tokens. The weight matrix P E Ri Xdx 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. Wok = Wx We as well as Wvp = PWv, 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 L2- L2-regularization, and for any such two weight matrices A and B (e.g. the P and Wy for given layer and a given head), we can rewrite the loss as:
> (11411 + [B]12), L12(A, = L(AB", 0) +
(2)
2
where 0 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 (11A1] + B a tight upper bound on the nuclear norm |AB ||+ [Srebro and Shraibman 2005 Tibshirani 2021), we theoretically show 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 commonly used when training Transformer models. For example, GPT-3 [Brown et al., [020]. LLaMa [Touvron et al.] [023] LLaMa 2 [Touvron et al.] 2023) and ViT (Dosovitskiy et al.] 2021 report a weight decay strength of X = 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 opti- mization with decoupled weight decay [Loshchilov and Hutter 2019]. models ranging from deep linear networks to language models as well as Vision Transformers Dosovitskiy et al. [222]. 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 feature or bug.
We provide evidence suggesting that this rank aankrreularrizing effect in fact seems to affect the pretraining of popular pre-trained foundation models such as LLAMA [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 the Matrix Completion
literature [Srebro and Shraibman 2005 Sun and Luo 2016 Candes and Tao 2009). where the goal to recover an unknown low-rank matrix for which only a subset of its entries are specified. Nuclear
norm regularization is often used as convex relaxation of the problem [Hu et al. 2021b), and

---

its equivalence at the global optimum with the L2-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 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., [022] Arora et al. 2019 Li
iet al. [02T]. More recently, Pacot [023] 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 1.2 regularization applied on factorized matrices and low-rank inducing Ty-SSchttten norm on the matrix they parametrize has been shown several prior works (Dai et al. 2021 Tibshirani] 2021). This is particularly relevant as L2 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 L2 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 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 12-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 gradient- based optimization, and thus, that the low-rank inducing effect comes into play very early in during training. This brings theoretical understanding to empirical observations made in previous works [Khodak al., 20022. 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. [2200] 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 L.LMs can hurt downstream performance [Sharma et al. [003). 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 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
|||| Tr Tr(VWWT)
(3)
3

---

When using the singular value decomposition (SVD) of W, W = -USV, denoting (s;); the singular values, we can see that
[ TIVVSVVVVUUT Tr(5) Σ s6
(4)
Le. the nuclear norm is the sum of the singular values of W.
The nuclear norm often used in the low-rank regularization literature (Hu et al., 2021b] as intuitively is a convex relaxation of the rank, and regularizing it typically induces low rank by injecting sparsity the singular values.
3.1.2 Upper bound of the nuclear norm of factorized matrix
Let two matrices A, such that ABT Then, using the Cauchy-Schwarz inequality, we have that
|W]|- Tr(s) TT(U VB
(5)
s THUT BBB
(6)
= ||AI|BB| S (11111 + [|
(7)
3.1.3 Considered losses
We will consider 7.2 losses of the format
= L(AB)) + > (11411 + [|B|)))
(8)
and their L, counterpart
C.(AB = L(ABT) + A||AB ||..
(9)
As consequence of the above inequality, the I.2rregllarized objective ( 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 generality. The reader is referred to appendix for more details about this point.
3.2 Equívalence of optimization solution
In the following, we will first show that in fact, any objective of the form in ( will coincide at any stationary point with the nuclear-norm regularized loss in (2). thus introducing low-rank inducing bias in the solution found. We assume A, B to have a bottleneck, i.c. to have the number of rows greater equal the number of columns, as usual in attention layers. All proofs can be found in Appendix B
We start by providing sufficient condition under which the averaged Frobenius norm of two matrices would correspond to the nuclear norm their product.
Proposition 3.1. Let A, be matrices such that A A BT B. Then, denoting ABT USV
the SVD of ABT, there exist orthogonal matrix O such that U
VS
0 )0
O and B
VS 0 OT.
)OT In particular, ||ABT ||- (1 + |[B|)T)
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 LL2- that condition is fulfilled. We assume the loss L is differrentabbee and X > 0.
Lemma 3.2. At any stationary point A, B of Lia we have that A A B B.

---

The above Lemma, together with Proposition 3.1 implies that at stationary point A, B, L12 [A, B) and coincide. However, this is not enough to claim that finding a (local) minimum of LL will in fact find a (local) minimum of C.. We now provide result which shows that this claim true.
Theorem 3.3. A, B is a local minimum of L12, if and only if 1) W AB is a local minimum of L., constrained to matrices of rank where r. is the maximum rank achievable by AB", aad 2) A B.
more general formulation of the above results, albeit without a bottleneck dimension, was recently shown in (Wang and Jacot, [023] (c.f. Theorem 3.1) where it was applied in the matrix completion context. We restate and prove here for completion the context of self-attention and transformer models.
The Theorem states that there is in fact one-to-one mapping between the local minima of L1 A B) and (the equivalence class of) local minima of L, (AB" ) for general unregularized loss L.
In particular, 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 Li2 on A, B without introducing bad minima, and obtain regulaegularized solutions. In principle, one can still converge to bad minimum for general loss, but this is not due to the reparametrization.
On the other hand, the theorem shows that naively optimizing the T.2-regularized loss with factorized parametrization will (often inadvertenll) result in actually finding solutions that exactly minimize the nuclear- regularized loss, introducing unintended low-rank inducing bias 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 L1.- [|Al]] [B]] remain bounded
during training, then we have that [I) (AB )} converges exponentially to O.
In order to prove the theorem, we first show that during gradient flow optimization, the condition from Proposition h.tiiioonnnnmeecee.eee becomes true exponentially quickly. This is then followed by a new bound bounding the gap between ||AB |] and & AI| + [|]] 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 LL2 will result in optimizing L. very quickly during training, long before stationary points are found. The theorem also confirms previous empirical observations (Khodak 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, will study in the following the optimization within 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 [Zyiin et al. 2022 Saxe et al. 2013) and the redundant parameterization studied in (Ziyin
and Wang 2023
Consider the following model
f(ABT) AB
(10)
where BT E Rdd XRd For simplicity of presentation, we assume dg = d, d1,3. 3, but the result can be easily extended to the general case. Given D data points (z,, 31)1555 D. in matrix
5

---

100
LAB
0.5
120
10
0.6
0.30
0.55
301 IBCC
Figure 1: Optimization by gradient descent of two 5-by-5 matrices A, B on the T:2-regllarized loss |IABT - D1 [[] where D) = diag(0.2, 0.4, 0.6, 0.8, 1), with various regularization strength X. denotes the number of optimization steps. Left: difference between the nuclear norm ||B ||- with the Frobenius norm ||| + 2 BI throughout optimization. For all cases, other than 0, the trajectory converges exponentially quickly to predicted by our theory. Center left: Norm of the discrepancy between A A and B' B over training steps. As predicted the discrepancy exponentially vanishes, with a time constant proportional the X. Center right: Singular values of the matrix AB' at l = 1000, for various regularization strength X. As predicted, 8; decays linearly with X, until À 2 8i, which point the singular value vanishes. Right: Singular values of the matrix AB during optimization, for x 0.4.
form, the 12-regularized mean squared error can be expressed as
L _ |Y AB x
X B1] + ||13)
(11)
2
where X = (I:);; E R4IxD,Y Y = (wi); E Rd and À > 0. Using full batch gradient flow, the differential equation governing the parameter dynamic becomes
B A AB Exx) A TB
(12)
T (Σyx - AB Exx) Exx)x AA
(13)
where Σyx YXT, Σxx = XXT, and is 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 basis such that Ey where S the diagonal matrix which diagonal consists of the singular values (s; (si))e(L.dd,.) of Y
At equilibrium, we thus have the following set of equations
XB A T(S A
(14)
AA = (S - )
(15)
Denoting by a; the i-th row of A, B, and assuming the (s;)ig(i.at.a) are non-zero and distinct, we have the following conditions at equilibrium (cf Appendix C
Vi [1..d13), a; = b;
(16)
Vi,j E [1.. s.t. i # j, a. bj 0.
(17)
In particular, this implics, for any i, Xl|a;11 (s: || 1 ||.1. Clearly, if A 8,. then the equation can only be true if a, = 0. If on the other hand < s,, either a; = or ||a;||"= satisfy the equilibrium condition, with the former being an unstable equilibrium point if the number of hidden units d2 is greater than the number of elements in {i [1.dd,s] 8, > A}.
To highlight the result, let us consider the case where the hidden layer has enough capacity, i.e. d2 2 d1,3- In that case, the result tells us that stable equilibrium, AB will drop all singular values that are less than X, 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 controlled by A related result was already obtained in the analyses of [Saxe et al., 2013 who studied the exact

---

solutions of [1 without the regularization term but introducing bottleneck in the hidden layer, Le. d2 < d1,3. Remarkably here, regularization achieves similar effect even an overcomplete network, where increasing À gradually pranes the hidden neurons to ignore the smallest variations of the data, i.e. reducing d2 adaptively. We confirm these results empirically in Figure !
Importantly, this result is only obtained because the regularization is applied to the parametrization involving a matrix multiplication. If AB" were replaced by single matrix W 2 then the cquilibrium condition would be W T+X 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 as weight decay in the following, is used instead optimizing regularized loss. A popular choice of optimizer for deep neural networks, including those with self-attention layers, is A [Loshchilov and Hutter 2019], which update the weights by using the Adam optimizer the non-regularized loss while simultaneously applying weight decay.
While it is non-trivial to analyze the equilibrium points of Adam in general, we show Appendix that under some simplifying assumptions, they coincide with those of 1.2-regularized loss with different regularization strength.
4 Empirical results
0.8
[P] p W.W
rank Ww
10
20
10
HPP
25
50
75
100
0
25 781 70
1000
25
(4)
100
Figure 2: Left: The rank of weight matrix product PWy of the first layer of a 2-layer Transformer trained on the associative recall task, during training, with AdamW, for various decay strengths. To better account for the effect of weight decay on the attention layers, only the decay strength applied to attention layers is varied, while the strength for all other layers fixed at 0.1. We observe that rank reduction correlates strongly with weight decay strength. Center: Norm of the discrepancy predicted, the difference seems converge
between PT P and Wy WT, during training. when À > 0 towards the end of training. While for AdamW we no longer have the guarantee of an exponential decay, we see that the discrepancy nonetheless vanishes quickly, with a time constant which perfectly correlates with the decay strength. Right: The difference of the nuclear norm of Wvp with the Frobenius norm upper bounding it. As the discrepancy between P and Wy W decreases, the difference approaches 0, and thus the bound becomes tight. The optimization of Lr.2 thus gradually switches to that of L., explaining the rank regularization. Qualitative findings are identical when studying W WQ-
The primary objective of our experimental analysis is to empirically validate the theoretical findings in more practical settings. Specifically, aim 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 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 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 consistent with ranl-regularizing training.
7

---

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 weight matrix with singular values 11, 22...., nn ordered such that 2 Ø2 2 2 nn- The pseudo rank (refered to simply as rank in the following) of W is defined as fLr where k is the smallest number such that:
Ea_10 2i 2 0.95.
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.
3.00
378
03
0.55
0.3
3.52
03
0
20002
63300
30000
00002
20002
10833
00000
30000
23
15
Epcch
Figure 3: Left, center left: The rank of weight matrix products K Wo and PWy averaged across heads of layer 5 of an autoregressive transformers trained on the Pile [Gao et al. 2020]. Center right, right: The rank of weight matrix products W and PWv averaged over all heads and all layers of a Vision Transformer trained following [Irandoust et al. 2022] on the ImageNet dataset (Deng ct al. 2009]. In both settings, the decay strength applied to attention layers is varied, while keeping the strength for all other layers fixed. In all cases, we observe again that rank reduction correlates strongly with weight decay strength when optimizing with AdamW. The weight decay strength of 0.1 commonly used to pretrain some known large foundation models in fact noticeably reduces the rank of the generated matrices compared when weight decay is turned off.
4.1 Associative recall task
In this simple memory task, a model is presented with sequence of paired tokens [x1, "1....,, IT, yr, ZT+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 2N + 1 tokens (with T = _N), and the final token is repeated and appears in the sequence before, i.e. rT+1 x., for some remember the correct association observed in co-nentex and predict y,- This task has been attributed and proposed as proxy for language modelling (Fu et al. 2023 Poli et al.) 2023].
[0....,T ,T. The model is trained to
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 Zhhng et al. 202T].
Figure shows that even in this setting, the stationary condition of a L2-regularized loss in Lemma 3.2 approached, and the gap between the nuclear norm and the Frobenius norm (5) vanishes, thus confirming that Adam W 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 Adam W leads indeed to 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, 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 al. 20223. Details can be found in the Appendix G
First, we confirm again that increasing weight decay with Adam drastically reduces the rank of W We well as PWv, average across depth and heads, of the trained models (c.f. Figure 3
8

---

Table 1: Test set perplexity of 125 million Transformer models trained on the Pile for 10 billion tokens with A and different weight decays for the self-attention (SA) and the feed-forward (MLP) weights. + standard error of the mean computed over 5 seeds.
SA-À = 0.0
SA-À = 0.01
SA-À = 0.025
SA-À 0.1
SA-À = 0.25
MLP-À 0.0
12.00
12.01 #0.05
11.98 +0.00
11.92 +0.02
12.02 =0.01
MLP-A = 0.01
11.94-00.02
11.95
11.9440..0
11.89=00.02
11.97-40.03
MLP-A = 0.025
11.89:0.001
11.90-10.04
11.90
11.80
11.92-10.02
MLP- =011
11.72=00.2
11.71
11.68=00.3
11.67=0.03
11.70##0.2
MLP-A = 0.25
11.63+#0.02
11.65#0.04
11.62200.4
11.52=00.3
11.58:+0.03
Appendix Figure 7. Furthermore, we observe that while increasing the weight decay strength of MLP beyond 0.1 is generally beneficial, see Table[] doing the same for attention matrices starts slightly hurting performance. Results are averaged over 3 seeds. We observe 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 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 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 similar effect of the decay strength the ranks of the matrices WQK, Wyr (c.f. Fig 3).
4.4 Pretrained foundation models
Finally, we turn to pre- pre-trained foundation models, and provide some evidence that their training is also impacted by the rank regglaauua effect of weight decay. Specifically, following Proposition B.4 it is sufficient to observe that the matrices W W resp. PT P are close to Wx W K resp. Wy Wy Because the matrices Wo, WK, Wv, PT are typically wide rectangular matrices, the off-diagonal elements of W W that the diagonal elements of WoW resp. PT are close to those of Wx W resp. Wv W.
Q. etc, are mostly For L approximately correspond L., it thus suffices
Figure 4 shows that this is mostly the case, for all layers of the model. For each layer and head, we further find that the gap from (5) is indeed mostly tight, consistent with rank regularizing training.
25
50
PRT:
layer
00
200
36
2003
36
Wo
1
11PF= + 11
Figure 4: Analyses of attention layers in the pretrained LLAMA 2 model with 7 Billion parameters TTovrron al. 2023]. The leftmost (resp. center left) shows the squared norm of every row of Wo (resp. Wv), ), for the first head of each layer, against the norm of the corresponding row of Wx (resp. column of P). The condition Wx W WoW would require these norms be equal, which in fact is mostly true. While the model has not reached a stationary point, this indicates the optimization has advanced enough for this sufficient condition for L. be identical to L12 to emerge. In fact, the center right (resp. rightmost) plot show the scatter plot mapping the Frobenius norm against the nuclear norm for all heads across all layers. The two norms almost perfectly coincide.

---

Discussion
Our results provide further insights into the interplay between L2-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-proyj poouee-pooettonn parameter matrices, are multiplied directly with each other: W Wo and pWr. Our empirical findings strongly support our theoretical predictions about the impact of weight decay 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, 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 (Sharma 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 con- junction 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 more nuanced theoretical understanding of layer-specific regular- ization. We are particularly excited about further research that aims disentangle the role of weight decay in in- -n-weihht vs. in-context learning within MLPs and ceff-attention layers, building [Singh] al. 2022]. In conclusion, while our findings mark a step forward 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 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 Arcas, Alexander Mordvintsev, Ettore Randazzo and Eyvind Niklasson for fruitful discussions.
10

---

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. 10.1145/344677. URL htttss:///oi wwg/10..//... 344677
Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger Grosse. Three mechanisms of weight decay regularization. In International Conference Learning Representations, 2019. URL
https://openreviexr http://ew.nee//roommdd-i-ii---
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
Bkg6RiCqY7
Zeke Xie, zhiqiang xu, Jingzhao Zhang, Issei Sato, and Masashi Sugiyama. On the overlooked pitfalls of weight decay and how to mitigate them: gradient-norm perspective. In Thirty-seventh Conference Neural Information Processing Systems, 2023. URL https::://ppenreiivv..o net/ "ooun?id=vnGuubtaa
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 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. https:/:: httplppgge..pp....emmmaccccccorr wwt.s/sacicccccccc..oooooooooaaa 10137788
Liu Ziyin and Zihao Wang. spred: Solving 11 penalty with SGD. 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, absp/195. 2019. URL httt:://ataxiyy org/ abs/aa/s...1//////// 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. htt..itttpss//peeneeeeee..eee
Noam Razin and Naday 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-1
Ryan Tibshirani. Equivalences between sparse models and neural networks. Working Notes. URI. https:///ww.. stat. cmu. chiu/yenttis/ppperssppriinnnn pdf. 2021.
11

---

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, 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 14015 2020.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cris- tian 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, lliyan 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 Internutional Conference on Learning Representations, 2021. URL htttps://operreviee het.fforunPdd-YicoFaNTTY
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, /bss/210..0965,, 2021a. la. URL
https://arxiv org abs/ /bs//111 /206. 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) 6579, November 2016. ISSN 1557-9654. doi: 10.1 10.11099h1.2001625598777 URL http: doi. www....ogg1.... 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 re- view. Neural Networks, 136:218-232, 2021b. ISSN 0893-6080. doi: https://doi.or//10 http////wwwnnnnceeeettooom http://rrr.....ooocccooceeerrraa
0t16j]mmeueet2222000000002 021. pii p11/$0893608020303322
URL
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 Şimşek, 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: notion of rank for nonlinear functions, 2023.
Zhen Dai, Mina Karzand, and Nathan Srebro. Representation costs of linear neural networks: Analy- sis and design. In M. Ranzato, A. Beygelzimer, Dauphin, P.S. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 26884 26896. Cur- ran Associates, Inc., 2021, URL https::/p// B http:////ppoceeeegg.eeeeeeee neurips cwtpsper.pppp.ee files/ files/paper/ o0211111122220bb0000098444468633
pdf
12

---

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 12-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 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 deep linear network. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances Neural Information Processing Systems, 2022. URL BtRLlattte://pppeeeeeee..e http://ppwooooooooooooooooooddd/
Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutioon to the nonlinear dynamics of learning in deep linear neural networks, 2013. URL https:://arivv.org//aaiiii 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.000 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 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 computer vision and pattern recognition, pages 248-255 Icce, 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 Arcas, Max Vladymyrov, Razvan Pascanu, and João Sacramento. Uncovering mesa-optimization algorithms transformers, 2023.
Aaditya K. Singh, Stephanic 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. TtRt./ttp:// //wp...///w//////////////.//.. numdam.
Robert T Powers and Erling Størmer. Free states of the canonical anticommutation relations. Com- munications 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. Sebas- tian Seung. Digital selection and analogue amplification coexist in cortex-inspired silicon circuit. Nature, 40766789):44-995 2000.
13

---

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 UVEAVA and B UB with
SA and Σg = ( SB
ΣA =
Substituting in the equation A A B B, get
0
0
VA = V
By positivity and uniqueness of singular values, necessarily, ΣA = ΣB = (50) ))) Furthermore, by
rearranging the above equation, we get V V VBS i.. that V VB commutes with We rewrite A as
A= UA VBV = =UA (
(5V/VV )
SV
VB
= -UA VVV
V = =UA
VV VB 0
1 >>VVVV
VT 1 ) setting E = EA. and V = VB, we can write A = UΛEVT
Redefining UA as UA and B =U
In particular, AB UAEETU which is a valid SVD of ABT. remains show that, AB diagonal, then there exists orthogonal matrix O such that = ΣO and B ΣOT.
Let us assume the diagonality, i.e. ABT ΣΣT Then, we have
(2BT)² UAEE Ey TU EN = UBΣΣ T
ie. that UA Up commute with ΣΣT, and thus that they are block diagonal. Furthermore, ΣY UnUB = ΣΣT They can then be written as
U
U
0
U%
UB
0
U/
where U, U,,, U/, are orthogonal matrices, and the block of U corresponds to the non zero singular values of ΣST
We can then rewrite A, B as A E U o 9) VT and B ==
(0 U 0 ) VT. which conclude
the proof by setting O U 0
1 VT.
Finally, ABT UAΣΣ UT, and therefore ||B ||. JIAΣE UT TREET [11
O
14

---

B.2 Proof of Lemma 3.2
Proof. Let A, B stationary point of the unregularized loss in LL.. One can show that the gradient of L(W = ABT) with respect to A (resp. B) is
aL
d^L
B
(18)
aW
W=ABT
0gL
(19)
OW
W-ABT
aL where SW - TR is matrix, which we denote by G. Differentiaingg L, at the stationary point, the following equations must then be satisfied
XA GB
(20)
AB G A
(21)
In particular, A-A= AA A) B B
B.3 Proof of Theorem 3.3
Proof. (=) We start by proving the backward implication, by contradiction. Let M local minimum of C., and A, B such that M = AB and A AA- B" B. Then by Proposition 3.1 Li B) (M). Assume A, B not local minimum of Lr.2, i.e. there exists an infinitesimally perturbed matrices A', B3 such that (A2,A', B¹) < Lra(A, (A, B). By continuity of matrix multiplication, M' A is an infinitesimally perturbed matrix M. Since L.(M) < L12(A', B') < L12(A, = f. (M), get contradiction.
(+) Assume now that A, B is a local minimum of L12, and that W = AB is not local minimum of L. constrained rank matrices. Then, can construct a sequence (W,).. n. of rank matrices C. (W). For all let W Sn VT the
such that lim,-roc Wn, = W, and for all n, L. (W.) SVD of Wn. By continuity the mapping from a matrix to its singular values, lim,- Sn = S, where S are the singular values of W. Becaaus the set of orthogonal matrices is compact, there exists a subsequence of ((U,, V,)), which converges to some orthogonal matrices (U, V). Without loss of generality, we redefine the sequence 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 VEOT
VS and O is some orthogonal matrix. Let for all A, = UnE., and B,
where Σ=
VSS, 0
Then, lim (Ans Br) (A, B) and yet, because for all AT A, B B, and A, Wn. have L12(A, Ba) (WW,) < L.(W) L12(A, B). This is contradiction.
V, OT, where Σ,
B.4 Proof of Theorem
In order to prove the theorem, we first show that during optimization, the condition from Proposi- tion 3.1 becomes true exponentially quickly. This is then followed by a new bound bounding the gap between |AB ||. and (1[AA]2 ? by the norm of A A T B B.
B.4.1
Exponential decay of A B
We provide the result for the vanilla gradient flow limit, but also provide 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 good approximation for small learning rate in the discrete dynamic.
Lemma B.I. In the gradient flow limit over the loss L12. AT A BT B will converge exponentially 0.
15

---

Proof. For any i, we denote by a", b' the i-th column of A and B. The columns follow the following differential equations:
Tà' = Gb - Xa' T G a' Xb"
(22)
(23)
AL where aW W=ABT and is some time constant controlling the learning rate. Given a pair i, we can now look at the dynamic of a" aT - -
d la' a bt 64) =( + a aTT - TT - -bT
dt
a AaTTaa a²
+ Gb Aa
(a) Gb NE AT b')
Gb Gb \b b)
-22a a be 82)
Therefore, we have A A - B B = Qe where Q is A A-AA-A BT particular, every entry of AT B B converge to exponentially.
B at initialization, and in
We now provide a similar result, in the gradient flow regime but with momentum, as well as decoupled weight decay- tractable approximation to AdamW as shows the following proposition: Proposition B.2. We consider the following dynamics approximating stochastic gradient flow with weight decay:
dHA B,G + adW, dt) dt)
dH. p(G. A,dt + odW H. dt)
dAA, HH, + dAt)dt
dB. = - + dBt)dt
where μ. a >
and W and W are independent matrix Wiener processes. Initial condition are Hº and H = 0. H (resp. HB) ) is the momentum gradient with respect to A (resp. B). Then,
H A =
plt pitts) Bdds +
W
A
(24)
00
2
1-e-2pa
H =p (' cull--o)
G A,ds +
WBB
(25)
2
1-e-2u
A: =E nAe Ao
E -)(--)) H ds
(26)
B₁ = --=-mii -
nAE Bu -1
-) H ds
(27)
0
AT Ar BT B, = 2nd (AT Ao BB Bb)
(28)
17 /--
2 A, A H HBT B, B. ds.
)
(29)
Proof. We have
d(c = = (uH, dt )dH,)) = + adW;A)
such that
HA -p H6 + 11 eH' odd
2o02 W G, +
16

---

Note that the second term is an abuse of notation. The derivation of the integral form of H₂?. A and B, follows the same logic.
For the last two equations, we get
d(A, A, B, =dA- A, + d dB, B, - B, dB.
=- (H, + +AAi) T A. A (H₁ + AAi))
-n + XBi) T B. + B. (H) + +ABi))
--21A(A A, B, B.) T(11 Ar A 11 11 B. B, 11, jdt
-- =-2nAA A. - B, + Qidt
+
which gives
( -2nd((-o)) Quds
A A. - BT B, = 2-2 BJ Bo)
Before analyzing the implications of this proposition, let us state lemma that will allow us bound the probability of a Brownian motion diverging:
Lemma B.3. Lévy 1940 Let (B.) be a 1D Wiener process. Then, for t, L >> 0, P(maxxe|o.1) B. > L] = B, DJ.
Let analyse the consequences of Proposition B.2 We will assume that A, and B, remain 1.2-bounded by M > 0 (which is true for all converging dynamic modulo steady state noise), and that G2 remains L2-bounded by (either using Lipschitzian loss, or using clipping).
First observe that, using lemma
B.3 for > 0, with probability - e, the term v /ooW 1-e-2e and
the correspond B will remain bounded by a judd And
This way, A and H will with probability - remain bounded by KM + a und 4nd With that same probability, the term
BB. L. 20A( A, + A H B. H
will remain bounded by 4'M KM + to - und 4nd
This is the same order of magnitude as the stochastic term. Until A A - T B is of that order, it exponentially decays.
B.4.2 Upper bound of || ||. 1/2 311 + ||B|])||
Finally, provide the following general result, bounding the gap between |AB ||. and (11AA]
2) by the norm A- TT B.
Proposition B.4. For any matrices A, B, have
AB 11. - -|A|| & & IATA A BTB |
In particular,
|AB 11 |A + |B||
2
| |B|-
VVIA |AAT
B B1I
17

---

Proof. Let Q= = A A Using singular value decomposition, we write A = UAEAVA and
B = BBBBBBV with AA = SΛ and MB
SB
Substituting in the previous equation,
we get
VISVV = VBSVV +Q
i.e.
VASAVA V VB +Q VB A)V,
(30)
where VBAV B VBSVV VT + VVBSBV V By the Powers-Stormer inequality
Powers
and Störmer 1970 we have || VV B I QQ
From there, we rewrite A as
SAVI VB
A= A= VBV = =UA
(31)
0
30
UA
VA Vn(Sn + +4) 0
VTT
(32)
VB
0
(1) A ))
= UA
1
ΣB +
VB
(33)
Consequently, ||A
||. ||5
AS ||
and
|||ABT] ||~ |B
|
||+ < S ||||=|SB|||.
|QQ| B||
Similarily, we have ||.AB" ||+ - ||||1
&
||
The second inequality is obtained by using the triangle incquality.
2/L B.4.3 Upper bound of II A A1l
1 Σ |
We here provide a more general version of Proposition 3.1
Proposition B.5. Let q 2 O, Ai E Magr, At E Mr. for l [2..L- 11, Ar E Mar and - Ar. Let > O. We assume that the sequence (Ai)IE(L..)) is e-balanced, i.e. that for E -1.
|| A A441 E
Furthermore, assume that
1
6 S
min |
1
Then for k E [I..L), we have
2/L
L
A
||
5 <r1AAn1
- 1/2
l=1
2/L
In particular,
L
2/L
L Σ
A
HAN
11/
2/L
1=1
18

---

Proof. We will assume for this proof that the A are square matrices for all E [I..L]. This proof holds is for the more general case, with more cumbersome notations, and with the trick used in equation 33 For 1 [1..T], we denote by A, = U be the SVD of A,. Let & E [1..A.. We first show by recurrence that we can write
Ar O1-1 A
for I [I..]]
for appropriate choice of orthogonal matrices (Ot) Lj and symmetric matrices (Ar) 4 and such that Ok-1 kk: O₂ V& and || & Vell - k| for 1 E [1..k]. The recurrence is symmetric for & and for e We will prove the latter case. For = the statement holds. Let / E [k..L 1J. We assume we can write A Ot-1 (Ex Oi) T with ||r]|| vE(l - k).
Let Q Al A: A1-14/41 We have
Og + A.) U141 Q.
Similar to the previous proof, we can write
Og(Σx + + U44UU+1UHH
with A symmetric verifying
|A < 10 E
We can rewrite
A111 = E441 -O00, U141 1441 Vita
= =O{EE + A A)O, V141V
We set O+41 Vi41U,,00, A411 A to get the desired result. We verify that || &
Vē(1 k
Now that we have proven our lemma, let us observe that
1
II
A Oo [I Ai) O2
with
II
[[ + A1)
2/L
2/L
Notice that |
||A |3.. Using the triangular inequality of A + |All for < S 1, we
have
2/L
2/I.
A
||
II ( + A1)
-112
2/1
2/L
1/1L
2/1.
E + Ar) - Σ
2/L
19

---

Furthermore, using the fact that ||AB 2/rr. S ||A412/r.11B11 Bl we have
L
L
I
+ A1)
1 11222 1'
2/L
L/22
1 2
1.//1 || [/2(L-11 E
(1ER)]) 1.12 L621126
|E 11 12))
S1 12112
(1+
E
-1)
S/"111E411 12 eT.
V |
S1 L(L-1)) 11/2 eL222/2
where in the penultimate line, we used (1 + x)" -1 S enx whenever x < 1 n. and using the assumption of the proposition. Ultimately, we thus obtain
2/L
L
II
-|Ax11?
< 1
2/L
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 argue that optimizing a deep linear network of depth L will very quickly co-optimize the Lp Schhatttenn norm of the product, with an exponential time constant X/L. In fact, even in 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 I _ min; || to hold, suffices that the matrix norm remains nower-bounder by a strictly positive value, 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 expanding set of such sufficient conditions to cover 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 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-regularizztion effect that impacts such training processes.
B.5.1 Sufficient conditions
Henceforth, we will refer by 0 (A, B).
20

---

Sufficient condition 1: Gradient flow with lower bounded loss function We consider the following gradient flow dynamics with weight decay:
θ -V 20)
?) is the learning rate hyperparameter, and the weight decay strength.
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'(0) == L(0) + A|1011 Given that L' is also lower bounded, and that L'(0,) is monotonically decreasing function of time, (',)))
must converge to a constant real value, i.c. L(O,) + A| -tc+00 c for some c. If ||0|| unbounded from above, then necessarily L is unbounded from below, which is contradiction.
Sufficient condition 2: Gradient flow with momentum with Lipschitz gradient We consider the following gradient flow dynamics with momentum and decoupled weight decay:
G = (VoL - G) 0 = (/ + A0)
where G the exponential average of the gradient of A. yu is the momentum hyperparameter, the learning rate, and À the weight decay strength.
A sufficient condition the gradient is to be min(1, 2d Lipischttzz with respect to the parameters 0 sufficiently far, i.e. for ||O]| > P for given P:
The momentum makes the analysis of the dynamics more complicated. However, defining 0 and M G
17> 17 and U 0
rewrite the equations as:
11
VoL.
F - MF HU
The derivative of the squared norm of F verifies:
d
F1|2 = Tr FFT
dt
Tr MFFF + F
< (ni F
< m) [oLI|FFI
nd |[[]] [|VaL|| min(1,
= =||
nd ||||| [Voo min(1,
< FIl
The dynamics of are flow dynamics. Whenever || ||| reaches P, its norm is decreasing. [| can thus never exceed P. As [|] is an upperbound on [|@]], the same holds for ||0]]- 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 -I (for example the sphere of radius M) that contains the initialization point.
Note on stochasticity
To deal with stochasticity, we consider similar equations:
dG Ldl + dW Gdt)
dB (n 20)
21

---

which become:
dF MFdt + pti pdW
The integral form is:
f fcM + pe pe f
F(0 + μe
ue- tM
The process can diverge because of the stochasticity. However, similarly to proposition B.2, we can fix for any 5 > 0 an upper bound on W that holds with probability at least 1 - 8: 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 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(ABT) - ||ABTI (resp. B) is V VAL = -AB' B (resp. VBL - BAT A). Even with the decay term, one can see that A, are initialized to be e.g. orthogonal matrices scaled by some a X, both A, B will diverge to infinity.
1|2 Then, the gradient of L w.r.t. A
Such negatively unbounded objective functions be minimized may be found in e.g. the reinforce- ment learning setting, when using undiscounted returns.
C Equilibrium condition of 2-layer Linear network
We assume À and that the singular values of Y X are all non-zero and distinct.
We start with the following set of equations:
AB = A A - )
(34)
AA (S ABT) )
(35)
Clearly, 34 implies
BB AT(S BB
(36)
AAT AT(S AB
(37)
and thus that B B ATA. A.
Furthermore, 34 also implies
XAB = AA - A AA
(38)
XABT = =(S -S- AB >-BBBBBT = SBB AB BBT
(39)
Using BTB= AT A, we get that AATS= = =SBT.
Denoting by a; by the i-th row of A, B, for any i,j E [1..d1.a] x 2 we have 8;a, a; = szb by and 5y03 a; = s;b, b, Thus, ||a,l| = |b,11. and a 0j = bībb, = 0, since S; are distinct and positive. Taken together, AAT is diagonal matrix, which we denote by D. We have D ((a))
In particular, (38) implies XABT = D(S AB ABT), i.. (AI + D)ABB- = DS. Because the entries of are positive, (VI + D) invertible, and thus AB (M + D)- other words, the off-diagonal entries of AB are zero, i.e. a b; for all + j. i, E [1..d1,]]².
In particular, for a given i, we have
Xa by = ||a:112(s: a. bz)
(40)
\l|a;|12= 2 (8: a. b.)a, by
(41)
Using the positivity of s; and X, one can see that necessarily, a; b..
22

---

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:
L (A, B,0) I(ABT.) (1I + [|]]
(AB 8) L(ABT) A|| Il..
where the remaining parameters not accounted for. This in fact still accounts for the general case. Indeed:
stationary points of L2 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 LL2 iff 1) (W = AB 0) local minima of L. constrained to matrices W of rank r and 2) A B
in the |...f of theorem 3.4 the gradient G of the first lemma now hides a dependence in 0, but the proof still hold as is. In the second lemma, the gradient G, depends on the time, and is indirectly hiding dependence 0.
E
Link between solutions of Adam W and I2-regularrization
Consider the following dynamic induced by AdamW, with > 0:
G 81 G. + (1 B1) VwW(W)
B, + 32 B (1 32) VwL(W))²
G, + G1/(1 84)
B, + B1((1 B5)
GGl/ (VB++e) t + AW,
W411 + W. -
where 1) represents the learning rate and Bi, 32, E are the common hyperparameters of Adam, W, is the parameter at time and where the various operations are applied element-wise.
Note that the term -AW, stems from weight decay.
the dynamic converges, then necessarily, Goo = Vw Bx = (VwL(W and thus XW. 2w Clearly, this implies that A|W < 1. If we further assume that X|W 1, then the condition becomes eAW - Www.CWW <())), which is the equilibrium point of L2-regularized loss with regularization strength aA Thus, the stationary points of the AdamW optimizer can in practice correspond to stationary points of L2-regularized loss, and thus the same low-rank inducing solutions can be found.
V ((W>)
We show in F toy experiments illustrating the equivalence in the solutions found by Adam W with decay strength Åw D and hyperparameter with those found by Adam with I.2regularizztion with regularization strength λ72 = DCDC. 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.
Pretrained foundation models: VIT
We provide in Figure the VIT counterpart of Figure
23

---

0.06
0,06
0.04
0,05
App. Auaf
103
0.02
Aga Aust
0.0001
0.00
0.00
-0.02
-0.02
Aca Auct
Ace.
0.05
Figure 5: Trajectory of wi, W2 in the 2D plane when optimizing the underlying parameter for various hyperparameters. At every coordinate in the plane, the loss is defined as the squared distance the surface in orange. The red (resp. blue) cross represents the points on S minimizing the L2-norm (resp. L1-norm). Left: W1, W2 directly parametrized and optimized by Adam with decoupled weight decay (in solid line) or Adam with L.2regularization (in dotted line). As conjectured, the convergence point of AdamW given the hyperparameter and decay strength A corresponds that of the equilibrium point of the 1.2-regularized loss with regularization strength λ = de. Right: product two scalars, i.c. wi agbi, W2 a2b2. where a1 b1, a2, b2 are now optimized by AdamW or Adam with L2 regularization. Again, the two optimizers find the same convergence point for equivalent hyperparameters. However, the solution found now corresponds to those of the loss regularized by the L1-norm of wh W2, (corresponding to the nuclear norm for scalars) predicted.
W1, u2 are parameterized
ID
[| Mor
-20
dager
00
200
400
300
20
19-100
31
Figure 6: Analyses of attention layers in the pre-trained Vision Transformer (Wu et al. 2020 available huggingface under the id center left) shows the squared norm of every row of WQ (resp. Wy), for the first head of each layer, against the norm of the corresponding row of WK (resp. column of P). The condition W WxW k W would require these norms to be equal, which in fact is mostly true. While the model has not reached stationary point, this indicates the optimization has advanced enough for this sufficient condition for L. to be identical to L12 to emerge. In fact, the centre right (resp. rightmost) plot shows the seatter plot mapping the Frobenius norm against the nuclear norm for all heads across all layers. The two norms almost perfectly coincide.
The leftmost (resp.
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-ate self-attention layers and applying skip-connections after each layer. Training is conducted using standard (autoregressively) masked cross cross-entropy loss, omitting an input embedding layer but incorporating an output projection before computing logits. Further details can be found in Table 2
24

---

Table 2: Hyperparameters for language modelling experiments.
Hyperparameter
Value
Dataset
The pile [Gao al. 2020)
Tokenizer
GPT-2 tokenizer we append special "EOS" token between every sequence
Context size
756
Vocabulary size
50257
Vocabulary dim
756
Optimizer
Adam Kingma and Ba 2015 with e 1e 1,=",,80 = 0.9, ,3 0.95
Weight decay
See main text
Batchsize
Gradient clipping
Global norm of 1.
Positional encodings
We add standard positional encodings.
Dropout
We use embedding dropout of 0.1 right after adding positional encodings. 12 layers, heads, key size 64, token size 756, no input- but outuut-mmbedding W ~ (0, with a = 0.02 and bias parameter to zero. We scale all weight matrices before skip connection with 2/X with N the number of layers.
Architecture details
Weight init
Learning rate sched-
Linear warm-up up starting from 1e to 1e- the first 8000 training steps, cosine annealing 10% of the learning
uler
after warm-up for the end of training Widening factor hidden dimension 756 with ReI.U
MI.P size
non-linearities Hahnloser et al. 20X00
3.40
0.90
399
1w000
1.000
0.83
14510
0.0
15282 03381
J5882 10330
/2003 53281
2.880 10833 10531
Figure 7: The rank of weight matrix products W WQ and PWy averaged across heads of layer 7 (left and outer left) and layer (right and outer right) of an autoregressive transformers trained on the Pile [Gao et al. 20200] For both layers, the decay strength applied to attention layers is varied, while keeping the strength for all other layers fixed. In all cases, we observe again that rank reduction correlates strongly with weight decay strength when optimizing with Adam W.
NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
We provide both theoretical and empirical results supporting the main claim of the paper.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
25

---

The abstract and/or introduction should clearly state the claims made, including the contributions made 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.
is fine to include aspirational goals as motivation as long as is clear that these goals 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 separate "Limitations" section 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 few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 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, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers penalize honesty concerning limitations.
will be specifically instructed
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All statements are provided with proofs.
Guidelines:
The answer NA means that the paper does include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
26

---

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 ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We hope to provide 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 cross- referenced.
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 provide 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 instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [No]
Justification: aim to collect the code as soon as possible in Git repository. Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (kttps://nips.ce
for more details.
While encourage the release of code and data, 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 new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ https/cccppbbcgguuueeeeee Sttpl////bgddde//ddeeuubbbbbbiii ssi asionnolicyy for more details.
The authors should provide instructions 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. only 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) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We provide all experimental details in the main text, well as the appendix.
27

---

7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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 provided either with the code, in appendix, or as supplemental material.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide 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 compute required for each of individual experimental rums well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make 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 http::////errrps..cccc///////taa
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 NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
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 proper weight decay while training transformer.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
28

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate to point that 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 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
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)?
Question:
Answer: [NA]
Justification: No data and models realese.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 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 encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
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?
Question:
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 used and, 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 use in package should be provided. For popular datasets, paperswithcone com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
29

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information not available online, the authors are encouraged to reach out to the asset's 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/moodll 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 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 crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According 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 crowdsoureing nor research with human subjects.
Depending 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 NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
30


Reviews:
[{"review_id": "ElCt9vEmpe", "reviewer": "Reviewer_veeJ", "summary": "In 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.", "questions": "Please 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", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The 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.", "weaknesses": "While 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their encouraging feedback and constructive criticism. We address your concerns point by point below.\n\n> The writing and exposition could be greatly improved\n\nWe really appreciate all the feedback about the writing. We fully agree with all the above points, all of which are now taken into account in the revised version. \n\n> Can the results about low-rank minima still carry over when the multiplicative terms are $PW_V$ and $E^\\top W_K^\\top W_Q E$?\n\nThis is indeed a crucial point which could be clearer. We thank the reviewer for bringing it to our attention. For our theory to hold, the paramount condition is that two matrices, $A$ and $B$, enter the unregularized loss $L$ only as their multiplied form, $AB$. That unregularized loss may depend on other parameters and inputs, which obviously may interact with $AB$. But as long as this condition holds (as is the case for the attention matrices), the dependence of $L$ on other quantities can be safely ignored without loss of generality.\nFor any such two matrices $A$ and $B$, we can write the loss as $\\mathcal{L}: (A, B, \\Theta) \\mapsto L(AB^\\top, \\Theta)$, where $\\Theta$ denotes all the remaining parameters. Stationary points of $\\mathcal{L}$ are also stationary in $A$ and $B$ so Lemma 3.2 still holds. For Theorem 3.3, the exact same proof would allow us to show that $A, B, \\Theta$ is a local minimum of $\\mathcal{L}\\_{L2}$, if and only if $AB, \\Theta $ is a local minimum of $\\mathcal{L}_\\star$ (constrained to ...). For Theorem 3.4, in the warm-up proof on line 534, $G$ hides a dependence on $\\Theta$, but the result still holds as $G$ cancels out. The same applies to the gradient flow proof on line 542, where $G_t$ now depends on time, and thus indirectly on $\\Theta$.\n\nWe realize we did not provide enough explanation of why the loss can be assumed to depend solely on $A$ and $B$ without loss of generality. In the updated paper, we now very clearly motivate, using the specific example of transformers, why studying our loss in this manner is the right thing to do.\n\n---\n\nWe thank the reviewer again for these suggestions and questions, which we believe help improve the paper. We remain at your disposal if you have any further questions, and hope our clarification convinces the reviewer for a strong acceptance of the paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6zrHwfIkdN", "reviewer": "Reviewer_fDM3", "summary": "The 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.", "questions": "1. 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.", "limitations": "I have no suggestions for the authors.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback and interest. Below, we address your concerns point by point.\n\n> Lack of an additional lemma addressing the effect of time discretization in the gradient flow scenario. The stochastic case is studied through SDEs, which may differ significantly from the exact setting encountered in deep learning.\n\nWe thank you for bringing up this point. We would like to briefly comment on how our result could be extended to the discrete dynamic setting and why we omitted it in the manuscript.\n\nFor the non-stochastic version (Lemma B.1), the proof remains unchanged. One essentially obtains an exponential decrease in $(1-\\lambda\\eta)^2$, where $\\eta$ is the learning rate. Note that for $\\eta \\to 0$, we retrieve the $2\\lambda\\eta$ factor from the continuous case.\n\nFor a stochastic version, assuming we model the minibatch noise similarly, one would get an exponential decrease of $A^\\top A - B^\\top B$ until it becomes of the order of $\\sigma M$, where $\\sigma^2$ is the variance modeling the stochasticity of SGD, and $M$ is an upper bound on $A$ and $B$.\n\nThings get one order more complicated if we want to model discrete time, stochasticity, and momentum. The discreteness of time adds additional interaction terms between the noise terms and the matrices $A$ and $B$, that we now need to take into account. While this is totally within reach, it would add a lot of technicalities that don’t provide any additional insights or intuition. The proofs would be barely readable, and even the phrasing of a precise proposition would be much more complicated. As you rightfully pointed out, the continuous SDE offers a nice approximation of real dynamics, as $\\eta$ is often chosen to be small, while still capturing all the intuition.\n\nWe agree, however, that the quality of these approximations to the practical training dynamics is nonetheless an important point to discuss, and will add a short paragraph in the final manuscript.\n\n> 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 $L2$ is exactly the same as the gradient flow of $L^\\star$. 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.\n\nThis is an intriguing point, and we thank the reviewer for bringing it to our attention. There is indeed a deep connection with spectral initialization, which will now have a dedicated paragraph in the discussion for the final manuscript.\n\nHowever, we stress that $\\mathcal{L}\\_\\star(A_t,B_t)=\\mathcal{L}\\_{L2}(A_t,B_t)$ for all $t$ during optimization with respect to $\\mathcal{L}_{L2}$ does not imply the equality of the gradients of the two losses. We can theoretically show that the equality of the gradients would in fact hold whenever the singular values of $AB$ are all equal, but this does not hold in general.\n\nObviously, however, the equality $\\mathcal{L}\\_\\star(A_t,B_t)=\\mathcal{L}\\_{L2}(A_t,B_t)$ implies $\\mathcal{L}\\_\\star$ is co-optimized from the beginning, thus the rank of $AB$ is regularized, and given our theoretical result, optimization will find a local minimum of $\\mathcal{L}\\_\\star$. We stress that while optimizing directly $\\mathcal{L}\\_\\star$ will also find a local optimum, it may take a different trajectory and find a different optimum. The study of this difference may be an interesting future work.\n\n> I suggest organizing the references section better to make them all uniform in style.\n\nWe fully agree with this point, and we have now fixed it. We thank the reviewer for bringing this to our attention. We will keep on improving our manuscript for the final version.\n\n---\n\nWe thank the reviewer again for bringing intriguing connections to our attention, which will help improve the discussion. We hope we have addressed your concerns and remain at your disposal for any further questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dFWkPb8jNv", "reviewer": "Reviewer_hhKp", "summary": "This 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.", "questions": "**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.", "limitations": "The paper does not really discuss limitations in any depth. I would not expect any particular adverse societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ 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.", "weaknesses": "- 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for providing us with valuable feedback. We address your concerns and questions point by point below.\n\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\nWe clarify that we have three main theoretical contributions: the correspondence of the local optima of the two regularized losses; the matrix inequalities and the exponential decay of the difference between the two losses during optimization. The first two results are independent of the optimization method and, therefore, are relevant for real training. Specifically, if the solution found is an approximation of a local (or global) minimum of the L2-regularized loss, it is also an approximation of the nuclear-norm regularized loss.\n\nAs for the result on the training dynamics, we provide results for both the simple gradient flow regime and stochastic gradient flow with momentum. There are indeed several apparent sources of discrepancy between these dynamics and real training dynamics, e.g., continuous vs. discrete dynamics, how to model the minibatch noise exactly, and how to approximate the theoretically intractable AdamW dynamics. While we believe some of these approximations are benign (see our 2nd response to reviewer fDM3), others may not be (see, e.g., [1]). Ultimately, we had to strike the right balance between theoretical tractability and proximity to real training, such that relevant insights may be extracted without overcomplicating proofs, and we believe our empirical verifications validate the modeling. We agree that these are nonetheless important points to discuss, and will add these points as a limitation of the current theory in the final manuscript.\n\n[1] Dynamic of Stochastic Gradient Descent with State-Dependent Noise. Qi Meng et. al Archive 2020\n\n> (minor) Missing a couple of related work on weight decay (see details below).\n\nThank you for providing these references. The relationship between weight decay and the effective learning rate is very relevant in our setting, as you pointed out. We appreciate the reviewer for bringing this to our attention. We will incorporate these references into the final version of the manuscript.\n\n> The ViT experiments show a much larger loss of rank but the performance impact of this is not quantified, why not?\n\nThis is indeed a good point that needs clarification. In our experiments, we found that removing the weight decay from the attention layers in ViTs did not result in a significant gain in performance. In fact, removing weight decay even slightly degraded performance, indicating that for visual tasks, low rank doesn't seem to hurt as much as it does for language tasks.\n\nThis highlights a subtlety in our message that we now clarified in the revised version. The practical takeaway of our work should not be that practitioners should always turn off weight decay in attention layers. In fact, we speculate that when increasing, for example, the key-query dimension, the rank regularization of weight decay may become beneficial, even in language modeling tasks. Our work aims to uncover the confounding, low-rank-inducing inductive bias of weight decay coupled with attention layers and demonstrate the relevance of this inductive bias in real applications, including popular foundation models. The benefit of this bias is problem- and architecture-dependent. By taking an off-the-shelf model and hyperparameters and showing that performance can be improved by turning off weight decay in its attention layers, we aim to highlight the need for practitioners to take this inductive bias seriously.\n\n> Disentangling the softmax temperature effects from the rank effects. \n\nThis is an excellent point. During our experimentation, we tried selectively turning off the weight decay on the key-query matrices only, as well as on the value-projection matrices only. Both of these changes, in fact, improved upon the baseline where weight decay was left at 0.1, yielding about half of the improvement achieved by turning weight decay off in all attention matrices. This suggests that low rank in both \\( W_{KQ} \\) and \\( W_{PV} \\) is beneficial and that the benefits are cumulative. \n\nNow, since turning the weight decay off in the value-projection matrices does not affect the effective temperature of the softmax attention, we believe the performance improvement could be (at least partially) disentangled from it.\n\n> Disentangling the effective learning rate effects from the rank effects. \n\nThis is also an excellent point. We thank the reviewer for bringing this confounding effect to our attention. We conducted the following experiment: to understand whether reducing the effective learning rate of attention layers can account for the performance improvement, we reused the off-the-shelf hyperparameters (where weight decay is set to 0.1 on all attention layers) and modulated the learning rate of the value-projection matrices by 1, 0.1, and 0.01. We left the key-query matrices untouched to disentangle the softmax temperature effect, as you suggested. Our early results suggest that reducing the learning rate in this manner results in significantly worse performance than the baseline, with a decrease of about 1% for the 0.1 modulation and 3% for the 0.01 modulation.\n\nWhile it is difficult to truly disentangle the various confounding effects, we hope these new pieces of evidence are sufficient to convince the reviewer that the improved performance is, at least partially, due to the reduction in rank.\n\n\nWe thank the reviewer again for their inputs, particularly for their great ideas for analyzing confounding factors, which we also will add to our discussion. We believe that the various new clarifications and the emphasis on the new controls strengthen our paper. We hope the reviewer agrees, and will raise their score accordingly.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BbaQ1V3Pay", "reviewer": "Reviewer_vYXP", "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.", "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.", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "1. Understanding the attention layers is an important question hence is of great interest.\n2. The experiments are extensive and well support the theoretical results.", "weaknesses": "The novelty of this paper is quite weak:\n- 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.\n- Theorem 3.4 is for gradient flow which seems far from the algorithms used in practice. \n\n- The empirical observation that low rank weights hurt the generalization also has been shown in [Sharma et al., 2023]. \n\nPlease correct me if I miss any key contributions in this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback. We address your concerns point by point below:\n\n> Some theoretical results have already been shown \n\nWe stress that while Wang and Jacot [2023] show a similar result, their result does not apply to transformers where $AB^\\top$ is not full rank due to an architectural bottleneck. More importantly, their result only describes what happens when training is converged to an equilibrium point. We stress that this setting is practically never achieved. Our theoretical contribution goes beyond this and considers what happens *during* optimization, which is theoretically far more complex and meaningful than studying behaviour at stationary points. We therefore stress that our results establishes that rank regularization does happen already early in training. This not only provides a theoretical explanation of past empirical observations (e.g., Khodak et al., [2022]), but is especially relevant for understanding (online) training large foundation models such as LLMs, where optimization typically is never brought to completion. Our empirical results, and in particular the analyses we performed on the pretrained foundation model weights, clearly demonstrate this relevance. We hope that this clarifies the novelty and relevance of our work.\n\n> Theorem 3.4 is for gradient flow which seems far from the algorithms used in practice.\n\nWe note that we provide in the appendix a similar result when considering gradient flow with noise, as well as with momentum and decoupled weight decay. \n\nAs for the continuous dynamics, we believe this to be a benign approximation. Indeed, for the non-stochastic version (Lemma B.1), the proof remains unchanged. One obtains an exponential decrease in $(1-\\lambda\\eta)^2$, where $\\eta$ is the learning rate. Note that for $\\eta \\to 0$, we retrieve the $2\\lambda\\eta$ factor from the continuous case. For the stochastic version, one would get an exponential decrease of $A^\\top A - B^\\top B$ until it becomes of the order of $\\sigma M$, where $\\sigma^2$ is the variance modeling the stochasticity of SGD, and $M$ is an upper bound on $A$ and $B$. Things get more complicated if we want to model discrete time, stochasticity, and momentum. The discreteness of time adds additional interaction terms between the noise terms and the matrices $A$ and $B$ that needs to be taken into account. The continuous SDE offers a nice approximation of real dynamics, as $\\eta$ is often chosen to be small, while still capturing all the intuition. \n\nWe agree however that the quality of these approximations to the practical training dynamics is nonetheless an important point to discuss, and will add a short paragraph in the final manuscript.\n\n\n> The empirical observation that low rank weights hurt the generalization also has been shown in [Sharma et al., 2023].\n\nWe respectfully disagree with this assessment. Sharma et al. [2023] show that *after pretraining*, when surgically reducing the rank of the MLP weight matrices, the performance of various LLMs improves on *downstream tasks*. As part of their hyperparameter search, they found no setting in which rank reduction of (even subsets of) attention matrices improved performance, and show some evidence that it may even hurt on, e.g., the CounterFact dataset. \n\nIn contrast, to the best of our knowledge, we are the first to show that the low-rank regularization induced by weight decay on attention layers *during optimization* can hurt the perplexity of LLMs. \nFurthermore, our empirical results demonstrate, by showing the equality of the entries of e.g. $W_KW_K^T$ and $W_QW_Q^T$   (Fig. 4 and Proposition 3.1), that the attention weights of popular foundation models, such as Llama 2 and ViTs, are being rank regularized through the mechanism we describe. \n\nThese are, to the best of our knowledge, very different insights from Sharma et al., 2023. We welcome the reviewer to clarify and point out results in Sharma et al., 2023 where the two aforementioned points were observed empirically.\n\n> About Theorem 3.4, how to see that \"... long before stationary points are found\"?\n\nThis is a goot point and deserves some additional clarification in the main text. What we theoretically show is that under reasonable conditions, the timescale at which the rank regularization can be observed is independent from the rest of the optimization (with a characteristic time equalling $ \\lambda $ ). This means that for long enough optimization, such as that of a foundation model, there is ample time for the co-optimization of the two regularizations to happen. Our analyses on pretrained model weights in Fig. 4 once again support that view.\n\n----\n\nWe hope we could clarify our contribution to the reviewer and convince them that the paper deserves acceptance. Please let us know if you have further questions or concerns.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2khkEKyozi", "reviewer": "Reviewer_eBg4", "summary": "The 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.", "questions": "- 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?", "limitations": "The authors discuss the limitations of their results.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I 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.", "weaknesses": "My 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback and constructive criticism. Below, we address your concerns point by point.\n\n> The theoretical results seem to apply more for compressed sensing or matrix completion, where the loss is $L(A^\\top B)$ for some differentiable $L$. However, the paper (and also the title) seem to discuss mostly transformers. \n\nThis is a valid point. Yet, we stress that the motivation of our research is to understand the inductive bias of weight decay in attention layers. Weight decay is virtually always used by default, and attention layers are almost ubiquitous in modern architectures. Therefore, we believe shedding light on this inductive bias is of great relevance to the community, and justifies the current framing despite our results ultimately being useful in other fields.\n\n> The discussion about linear models focuses on the underspecified regime since the matrix $S$ is invertible. This is not the standard regime for linear regressions.\n\nWe may have misunderstood the reviewer here, but we would like to clarify that our setting considers the case where $D$ (number of data points) is larger than $d$ (dimension of the input), since we assume the matrix $XX^\\top$ is the identity. As you point out, this is the standard regime for linear regression. Furthermore, even in the case where $d < D$, as long as $XX^\\top$ and $YX^\\top$ are still co-diagonalizable (a convenient assumption also made by Saxe et al. [2013]), we believe the derivation to be very similar, and that the weight regularization would simply set the weights corresponding to the degenerate dimensions to zero. \n\n> What happens if $L$ in Eq. (2) is non-differentiable? In practice, $L$ is non-differentiable for transformers since it incorporates an MLP with a ReLU activation.\n\nThis is indeed an important question. We point out that modern transformers, and in particular LLMs, use the GeLU or SwiGLU activation, making them differentiable.\n\nHowever, when the model is merely almost everywhere differentiable, we do not know whether the correspondence of the local minima between $L_\\star$ and $L_{L2}$, and in particular whether Lemma 3.2, still holds. Note that if the Lemma does hold, given that our proof in Theorem 3.3 does not use differentiability, the correspondence result would hold. \n\nAs for our result on the optimization dynamics (Theorem 3.4), our theory uses the gradient flow limit to study the problem. The flow is, of course, undefined at non-differentiable points. However, the Brownian noise accounting for the minibatch noise would allow escape from any non-differentiable point, so one could still define it soundly (e.g., setting the gradient to zero or to any subgradient at those points). Thus, almost everywhere differentiability is enough for the dynamics to be well described by our approximation.\n\n> In Figure 3 - what is the performance of each setting? It is not clear for the more realistic values of weight decay (i.e., those with the best performance) what is the rank.\n\nThe performance of the LLM experiment for various values of weight decay can be found in Table 1. For ViT, interestingly, we found that removing the weight decay from the attention layers did not result in a significant gain in performance. In fact, removing weight decay even slightly degraded performance, indicating that for visual tasks, low rank doesn't seem to hurt as much as it does for language tasks.\n\nThis highlights a subtlety in our message that we believe was not clear in the original manuscript. The takeaway of our work should not be that practitioners should always turn off weight decay from attention layers. Our work aims to uncover the confounding, low-rank-inducing inductive bias of weight decay coupled with attention layers and demonstrate the relevance of this inductive bias in real applications, including popular foundation models. The benefit of this bias is problem- and architecture-dependent. By taking an off-the-shelf model and hyperparameters and showing that performance can be improved by turning off weight decay in its attention layers, we aim to highlight the need for practitioners to take this inductive bias seriously.\n\nWe will include in the final manuscript the performance on ViT, and add the above clarification point.\n\n> Line 212 - How do the authors define an unstable equilibrium?\n\nIn this setting, we mean by \"unstable equilibrium\" a stationary point that is not a local minimum, i.e., the gradient is null but the Hessian is not positive semi-definite.\n\n----\n\nWe thank the reviewer again for their review and appreciation of our work. We remain available if you have any further questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We sincerely thank all the reviewers for their efforts in evaluating our work. Please find your personalized responses addressing your specific concerns. We welcome any further questions and remain open to continued discussions.

Meta-review: This paper studies the rank sparsity inducing effects of the weight decay regularizer applied to factor matrices, with a particular focus on transformer architectures. It also proves the convergence of the gradient flow in this case. 
Theorem 3.3 shows that the local minima of (2)  are the same as those of (8) (the only difference with existing works [1,2,3] is the arbitrary loss $\mathcal{L}$, the techniques are by now classic). Theorem 3.4 shows that if the matrices $A$ and $B$ remain bounded during training, the loss converges to zero exponentially during gradient flow training.  Interestingly, the authors argue that such low-rank effects are undesirable in transformer architectures, and demonstrate that this type of regularization harms performance in many cases.


 The reviewers are mostly positive about the paper. One exception is reviewer vYXP, which is concerned by the similarity to existing literature including [1]. I share this concern generally, as similar results have been shown in [1,2,3]. However, I like the fact that this paper considers a completely arbitrary loss function over the matrix, which makes the results slightly more general and elegant. In addition, I find some of the proofs to be of independent interest, such as the quantitative result in Proposition B.4. I did also have concerns about the "assumption" that $A,B$ remain bounded during training in Theorem 3.4, however, I really appreciate that the authors have been very reactive during their discussion with me, and provided **additional results** which show examples of loss functions where this assumption can be shown to hold, as well as examples where it does not. Overall, this paper makes an interesting addition to the existing literature. 






**Comments for the camera-ready**: 

Please incorporate the interesting analysis you provided in your comment (https://openreview.net/forum?id=oDeqjIM9Sk&noteId=w5ZHaJbJ1D) in the revision. I believe this greatly improves the quality of the work since you need a little push for the originality to be just above the borderline. In addition, I also wonder whether you could provide a generalization of Proposition B.14 and  Theorems 3.3 and 3.4 to the case of deep matrix completion (Schatten p quasi norm regularization), though I understand it would be mostly a tangential question since that is not relevant to transformers. 



 Please fix the citations: none of the **venues** are showing up in your bibliography. For instance,  the paper [1]  is published at ICLR 2023 (notable, top 25 percent). There are other  closely related works which are not even cited, e.g. [3].  It might also be advisable to expand the discussion of the related works by clarifying the difference in perspective between the present paper and the rest of the literature: in most literature, be it on neural networks, matrix completion or deep matrix completion, a the rank-sparsity inducing effect is a positive thing, which results in favorable generalization guarantees. In the present paper, the emphasis is on the way it can harm representation power in LLMs.


**References**

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

[2] Zhen Dai, Mina Karzand, Nathan Srebro, "Representation Costs of Linear Neural Networks: Analysis and Design". NeurIPS 2021. 

[3] Implicit bias of SGD in L2-regularized linear DNNs: One-way jumps from high to low rank. Zihan Wang, Arthur Jacot.
Decision: Accept (poster)