Weight decay induces low-rank attention layers

We thank the reviewer for their positive feedback and constructive criticism. Below, we address your concerns point by point.

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.

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

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.

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

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.

This is indeed an important question. We point out that modern transformers, and in particular LLMs, use the GeLU or SwiGLU activation, making them differentiable.

However, 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.

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

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.

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

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

We will include in the final manuscript the performance on ViT, and add the above clarification point.

Line 212 - How do the authors define an unstable equilibrium?

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

We thank the reviewer again for their review and appreciation of our work. We remain available if you have any further questions.

We thank the reviewer for their feedback. We address your concerns point by point below:

Some theoretical results have already been shown

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

Theorem 3.4 is for gradient flow which seems far from the algorithms used in practice.

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

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

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

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

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

In 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. Furthermore, 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.

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

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

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

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

We sincerely thank the reviewer for providing us with valuable feedback. We address your concerns and questions point by point below.

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

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

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

[1] Dynamic of Stochastic Gradient Descent with State-Dependent Noise. Qi Meng et. al Archive 2020

(minor) Missing a couple of related work on weight decay (see details below).

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

The ViT experiments show a much larger loss of rank but the performance impact of this is not quantified, why not?

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

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

Disentangling the softmax temperature effects from the rank effects.

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

Now, 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.

Disentangling the effective learning rate effects from the rank effects.

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

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

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

We thank the reviewer for their constructive feedback and interest. Below, we address your concerns point by point.

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.

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

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

For 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$.

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

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

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.

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

However, 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.

Obviously, 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.

I suggest organizing the references section better to make them all uniform in style.

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

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

We thank the reviewer for their encouraging feedback and constructive criticism. We address your concerns point by point below.

The writing and exposition could be greatly improved

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

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$?

This 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. For 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$.

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

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

Thank you for the reminder - we could indeed not see the first message. Below, the requested discussion on the boundedness assumption, with some sufficient conditions on the loss when it is provably satisfied. We will include this in the appendix of the final version of the manuscript. We are happy to discuss if you have any further questions.

Sufficient conditions

Henceforth, we will refer by $\theta=(A,B)$.

Sufficient condition 1: Gradient flow with lower bounded loss function

We consider a the following gradient flow dynamics with weight decay:

$

\dot{\theta} &= -\eta (\nabla_{\theta}L + \lambda \theta)

$

$\eta$ is the learning rate hyperparameter, and $\lambda$ the weight decay strength.

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

Indeed, the above dynamic is the gradient flow dynamic of the loss $L'(\theta) := L(\theta) + \lambda \|\theta\|^2$. Given that $L'$ is also lower bounded, and that $L'(\theta_t)$ is a monotonically decreasing function of time, $L'(\theta_t)$ must converge to a constant real value, i.e. $L(\theta_t) + \lambda \|\theta_t\|^2 \rightarrow_{t\rightarrow\infty} c$ for some $c$. If $\|\theta\|$ is unbounded from above, then necessarily $L$ is unbounded from below, which is a contradiction.

Sufficient condition 2: Gradient flow with momentum with Lipschitz gradient

We consider a the following gradient flow dynamics with momentum and decoupled weight decay:

$\dot{G} = \mu (\nabla_{\theta}L - G)$

$\dot{\theta} = -\eta (G + \lambda \theta)$

where $G$ is the exponential average of the gradient of $\theta$. $\mu$ is the momentum hyperparameter, $\eta$ the learning rate, and $\lambda$ the weight decay strength.

A sufficient condition on the gradient is to be $\min (1, \frac{\eta\lambda}{\mu})$-Lipschitz with respect to the parameters $\theta$ sufficiently far, i.e. for $\|\theta\| > P$ for a given $P$:

The momentum makes the analysis of the dynamics more complicated. However, defining $F = ( \theta, G )^\top$, and $M =((\eta\lambda, 0)^\top, (\eta, \mu)^\top)$ and $U = ( 0, \\ \nabla_{\theta}L )^\top$ one can rewrite the equations as:

The derivative of the squared norm of $F$ verifies:

$\frac{d}{dt}\|F\|^2 = Tr \dot{F} F^T$

$= -Tr{MFF^{\top}} + \mu Tr U^{\top}F$

$\leqslant -\min (\eta\lambda, \mu) \|F\|^2 + \mu Tr U^\top F$

$\leqslant -\min (\eta\lambda, \mu) \|F\|^2 + \mu \|\nabla_{\theta}L\|\|F\|$

$= \mu\|F\|\left(\|\nabla_{\theta}L\| - \min (1, \frac{\eta\lambda}{\mu} )\|F\|\right)$

$\leqslant \mu\|F\|\left(\|\nabla_{\theta}L\| - \min (1, \frac{\eta\lambda}{\mu}) \|\theta\|\right)$

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

Note on stochasticity

To deal with stochasticity, we consider similar equations:

$dG = \mu (\nabla_{\theta}L dt + dW - G dt)$

$d\theta = -\eta (G + \lambda \theta) dt$

which become: $dF = -MFdt + \mu Udt + \mu dW$ The integral form is: $F = F(0) + \mu e^{-tM} \int e^{sM}Uds + \mu e^{-tM} \int e^{sM}dW$ The process can diverge because of the stochasticity. However, similarly to proposition B.2, we can fix for any $\delta > 0$ an upper bound on $W$ that holds with probability at least $1-\delta$; this allows to bound the contribution of the stochasticity. We deal with the gradient component similar to the non-stochastic proof.

Dear Authors,

As I am not sure whether the above comment triggered a notification, I am posting this further comment to trigger one.

Best,

AC