Source-linked AI summary
Weight decay induces low-rank attention layers
Seijin Kobayashi, Yassir Akram, Johannes Von Oswald
TL;DR
The paper addresses limited understanding of weight decay and L2-regularization in multiplicatively parameterized neural networks, especially attention layers. It proves that factorized L2 regularization rapidly becomes nuclear-norm regularization during training and verifies rank reduction in attention products, while finding performance benefits from decoupling attention weight decay.
Problem
The effect of weight decay and L2-regularization on multiplicatively parameterized neural networks is poorly understood, despite their use in attention layers.
Method
The paper analyzes factorized losses theoretically and empirically studies weight decay on attention products W_K^T W_Q and P W_V in neural networks and transformer models.
Results
The two regularizations become exponentially quickly identical during training, and weight decay measurably reduces the rank of attention products, including under online training.
Takeaways & Limitations
Decoupling and tuning weight decay for attention layers can improve language-model performance relative to uniform regularization.
Takeaways & Limitations
The study highlights that the effects of varying regularization strengths across model components remain difficult to understand and require more nuanced layer-specific analysis.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper studies poorly understood L2-regularization and weight decay in multiplicatively parameterized neural networks, especially attention layers. It shows theoretically that factorized weight regularization rapidly becomes nuclear-norm regularization, creating low-rank pressure, and validates this effect empirically in transformer-related settings.
- Attention motivation: Attention layers expose this mechanism because their key-query and value-projection computations contain products such as W_K^T W_Q and P W_V.These products arise from the key, query, value, and output-projection matrices used in attention.
- Theory: Factorized L2-regularization makes the Frobenius-norm upper bound on the nuclear norm exponentially tight during gradient-based optimization, directly inducing low-rank pressure.The result holds for arbitrary loss functions and explains why the effect appears before convergence.
- Theory: All local minima of Frobenius-regularized factorized losses coincide with local minima of the corresponding nuclear-norm-regularized loss.For W = AB⊤, the factorized formulation and nuclear-norm formulation share local minima under the stated regularization.
- Empirical scope: The study empirically tests the predicted rank-regularizing effect across settings including deep linear networks, language models, and Vision Transformers.The contribution summary also includes experiments with decoupled weight decay and evidence from pretrained foundation models.
2 Related Work
Prior work connected factorized regularization, optimization dynamics, and low-rank structure, but did not establish how quickly the effect emerges for arbitrary differentiable losses. This paper positions its contribution as a theoretical and empirical account of early low-rank pressure in practical transformer training.
- Matrix completion: Matrix-completion research established global-optimum equivalence between nuclear-norm regularization and L2-regularization on factorized matrices.The factorized formulation is differentiable everywhere and has been used for large-scale matrix completion.
- Optimization and low rank: Earlier studies linked small initialization, gradient-based optimization, and regularization to low rank in deep or factorized networks.These works include analyses of deep linear networks, homogeneous nonlinearities, and factorized matrix regularization [Jacot et al., 2022, Arora et al., 2019, Li et al., 2021, Jacot, 2023].
- Closest prior work: Related general-architecture analyses bounded rank at sufficiently large training steps, while Wang and Jacot studied equivalence between critical points of the two regularized formulations.These results differ in scope or mechanism from the present claim about rapid equivalence during optimization.
- Paper distinction: This paper claims that arbitrary differentiable losses exhibit exponentially fast convergence between the two regularizations, bringing low-rank pressure into early training.The authors connect this result to practical non-converged training of foundation models and to earlier empirical observations.
- Attention-specific context: Prior work reports that reducing the rank of attention matrices after LLM training can hurt downstream performance.This motivates treating attention rank effects separately from rank reduction in other components.
3 Theoretical results
The paper shows that factorized L2-regularized objectives are equivalent to nuclear-norm-regularized objectives at local minima, while their optimization trajectories become equivalent exponentially quickly under bounded gradient flow. A linear-network case study shows that regularization removes singular values below its strength, inducing rank reduction.
- 3.2 Equivalence of optimization solution: Local minima of the factorized L2-regularized loss correspond exactly to local minima of the nuclear-norm-regularized product, subject to the achievable-rank constraint and balanced factors.At stationary points, the factors satisfy A^T A = B^T B, making the Frobenius-based and nuclear-norm objectives coincide.
- 3.3 Optimization dynamic in the gradient flow limit: Gradient-flow training makes the discrepancy between the two regularizers converge exponentially to zero when the factor norms remain bounded.The result implies that nuclear-norm optimization pressure appears early in training, before stationary points are reached.
- 3.4 Case study: 2-layer linear network: For a sufficiently wide two-layer linear network, stable equilibria retain singular values above λ and eliminate those below λ, so regularization controls the product rank.Increasing λ adaptively prunes hidden units in an overcomplete network, producing a low-rank approximation without requiring a bottleneck.
- 3.4 Case study: 2-layer linear network: The low-rank effect depends on applying regularization to a matrix-product parametrization; regularizing a single product matrix leaves its rank unchanged as λ varies.Under simplifying assumptions, AdamW equilibrium points coincide with those of an L2-regularized loss with a different regularization strength.
4 Empirical results
Experiments across associative recall, language modeling, vision, and pretrained Llama models show that weight decay makes attention matrix products lower-rank and drives the L2 and nuclear-norm objectives toward agreement.
- Measurement: The experiments use a pseudo-rank defined as the smallest number of singular values capturing at least 95% of their total sum.This measure is applied to attention matrix products in the reported experiments.
- Large-scale language and vision experiments: Increasing AdamW weight decay consistently reduces the ranks of key-query and value-projection products across language-modeling and Vision Transformer experiments.The effect is observed while varying decay only in attention layers, with other-layer decay fixed.
- Associative recall: The discrepancy between factor norms and the nuclear-norm bound vanishes quickly during training, with convergence speed correlated with decay strength.This behavior is observed under AdamW despite the absence of the theoretical exponential-decay guarantee.
- Language modeling: In language modeling, attention decay has a performance sweet spot near 0.1, whereas excessive attention decay and its associated rank regularization can hurt performance.Increasing MLP decay beyond 0.1 is generally beneficial, contrasting with the effect of stronger attention decay.
- Pretrained foundation models: Pretrained Llama 2 attention heads show near-equality of the relevant row and column norms, while Frobenius and nuclear norms almost perfectly coincide.These observations are consistent with the stationary conditions and tight norm bound associated with rank-regularizing training.
5 Discussion
The discussion connects the theory to transformer attention, where weight decay regularizes key-query and value-projection products, and argues that layer-specific decay can improve performance. It also emphasizes that regularization effects remain difficult to understand across model components.
- Connection to attention: The empirical findings support a rank-regularizing effect on attention products even without convergence and provide evidence that Llama training is affected by it.The products discussed are W_K^T W_Q and P W_V.
- Practical implications: Decoupling attention weight decay from other parameters and tuning its strength can improve language-modeling performance.The authors contrast this with uniform regularization and suggest model- or layer-specific strategies.
- Open questions: The study concludes that optimization effects depend intricately on neural-network parametrization and component-specific regularization choices.The authors identify layer-specific regularization and disentangling in-weight from in-context learning as directions for further work.
A Compute budget
The reported experiments used four Nvidia RTX 4090 GPUs for two months, while the language models were trained on 16 A100 GPUs for four days.
- Compute budget: The total compute budget was estimated at 4 Nvidia RTX 4090 GPUs for two months and 16 A100 GPUs for four days for the language models.These figures summarize the reported computational resources.
B.1 Proof of Proposition 3.1
The proof analyzes stationary factorized parameters by expressing their matrix factors through singular-value decompositions and deriving the resulting stationarity equations.
- SVD reduction: The proof rewrites A and B using singular-value decompositions to analyze the product AB^T and its singular structure.The construction uses commuting orthogonal factors and block-diagonal structure relative to the singular values.
- Stationarity conditions: At a stationary point of the unregularized loss, the gradients with respect to A and B are expressed through the matrix gradient of L(AB^T).The resulting stationarity equations provide the conditions used in the proposition's proof.
B.3 Proof of Theorem 3.3
The proof establishes that local minima of the factorized L2-regularized loss correspond exactly to local minima of the nuclear-norm-regularized product loss, subject to a balancing condition.
- Forward implication: A local minimum of the nuclear-norm loss can be lifted to a factorized local minimum using balanced SVD factors A⊤A = B⊤B.For a nearby factor perturbation, continuity keeps the product nearby and the nuclear-norm objective no larger than the factorized objective.
- Reverse implication: A local minimum of the factorized loss implies that W = AB^⊤ is a local minimum of the nuclear-norm loss over matrices of the achievable rank.The proof assumes a lower nearby nuclear-norm objective and lifts its SVD factors continuously to nearby A′ and B′, contradicting local minimality.
B.4 Proof of Theorem 3.4
The proof of Theorem 3.4 begins by showing that factor imbalance decays exponentially, then bounding the remaining gap between the factor and nuclear norms.
- Proof strategy: The argument first proves exponential balancing of A and B, then bounds the discrepancy between ∥AB⊤∥_* and the factorized regularizer.This two-step structure connects the training dynamics to the claimed exponential agreement between the two regularized losses.
B.4.1 Exponential decay of A⊤A −B⊤B
The section analyzes why A⊤A − B⊤B becomes small during optimization, first under gradient flow and then under stochastic momentum dynamics with decoupled weight decay.
- Gradient-flow balancing: A⊤A − B⊤B converges exponentially to zero under gradient flow over the L2-regularized loss.The same exponential decay is extended through an alternative analysis to stochastic gradient flow with momentum and decoupled weight decay.
- Stochastic dynamics: Until factor imbalance reaches the stochastic-noise scale, A⊤A − B⊤B continues to decay exponentially under boundedness assumptions.The stochastic analysis assumes bounded factors and bounded gradients, yielding high-probability control of the momentum terms.
- Deep linear extension: For deeper linear networks, exponentially balancing adjacent factors leads to co-optimization of an Lp-Schatten norm with time constant λ/L.This extension assumes sufficiently small imbalance relative to the factor norms and generalizes the two-factor argument to depth L.
- Boundedness scope: The boundedness conditions exclude objectives that can decrease faster than weight regularization, such as L(AB⊤) = −∥AB⊤∥².With suitable initialization, both factors can diverge despite weight decay for this constructed negatively unbounded loss.
E Link between solutions of AdamW and L2-regularization
The section relates AdamW equilibria to L2-regularized solutions and illustrates that factorized AdamW still reaches nuclear-norm-minimizing solutions under equivalent hyperparameters.
- Optimizer equivalence: AdamW equilibrium points approximately match L2-regularized stationary points when λ|W∞| is much smaller than one, with λL2 = ϵλ.The exact equilibrium condition implies λ|W∞| < 1; the L2 correspondence uses the stronger small-product assumption.
- Toy experiment: In the toy experiment, AdamW and Adam with equivalent regularization hyperparameters converge to the same points, while factorization selects the L1 or nuclear-norm solution.For scalar factors, the nuclear norm reduces to the L1 norm of the underlying coordinates.
- ViT analysis: In pretrained ViT attention heads, factor norms are mostly matched and Frobenius and nuclear norms almost perfectly coincide.These observations indicate that optimization has approached the sufficient balancing condition even though the model is not stationary.
G Language modelling experimental details
The language-modelling experiments use a conventional GPT-2-style Transformer setup with standardized hyperparameters, while varying attention weight decay to examine rank reduction. Rank reduction correlates strongly with attention weight-decay strength, but many experiments lack multiple seeds because large-scale models make them impractical.
- Experimental setup: The experiments use a conventional GPT-2 Transformer with LayerNorm, MLPs, skip-connections, masked cross-entropy, and an output projection.The setup omits an input embedding layer and uses standardized values from the literature with consistent, untuned hyperparameters.
- Experimental setup: The language-modelling hyperparameter configuration is documented in Table 2.
- Language-modelling results: Rank reduction correlates strongly with weight-decay strength in attention matrices when training autoregressive Transformers on the Pile.The analysis averages WQK and PWV across heads in layers 7 and 9 while varying attention decay and holding other-layer decay fixed.
- Experimental limitations: Multiple-seed error bars are not reported for many experiments because large-scale and foundation-model computation makes repeated runs unrealistic.