Research map / Architectures and efficiency / RNNs and normalization
Weight decay induces low-rank attention layers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows theoretically that L2-regularized factorized products (as in attention W_K^T W_Q and P W_V) equal nuclear-norm regularization at local minima and converge to it exponentially fast, and empirically that weight decay induces low-rank attention matrix products in small LMs, ViTs, and pretrained LLaMA 2.
Problem
Weight decay (L2 regularization) is widely used when training transformers, but its effect on the rank of factorized attention parameter products (W_K^T W_Q and P W_V) is poorly understood.
Approach
For any differentiable loss on a factorized product AB^T, local minima of the L2-regularized loss coincide with local minima of the nuclear-norm-regularized loss, and the two objectives become identical exponentially quickly during gradient flow, so weight decay on attention factors implicitly induces low rank even before convergence.
Key results
- Weight decay of 0.1 on attention with AdamW substantially reduces pseudo rank of W_K^T W_Q and PW_V in 125M LMs trained on the Pile for 10B tokens.
- On the Pile, 125M Transformer test perplexity is best (11.52) with MLP weight decay 0.25 and SA weight decay 0.1; increasing attention weight decay beyond 0.1 slightly hurts performance.
- In pretrained LLaMA 2 7B, squared row norms of W_Q and W_K (and W_V and P) are nearly matched, and nuclear and Frobenius norms almost perfectly coincide across heads and layers.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase self-attention weight decay from λ=0.0 to λ=0.1 | 125M Transformer · Pile test set | perplexity: 11.72$^{\pm0.02}$ -> 11.52$^{\pm0.03}$ | Helped |
| Increase self-attention weight decay from λ=0.1 to λ=0.25 | 125M Transformer · Pile test set | perplexity: 11.52$^{\pm0.03}$ -> 11.58 $^{\pm0.03}$ | Hurt |
| Increase MLP weight decay from λ=0.0 to λ=0.25 | 125M Transformer · Pile test set | perplexity: 11.92$^{\pm0.02}$ -> 11.52$^{\pm0.03}$ | Helped |
| Increase attention-layer weight decay | Transformer attention products W_K^T W_Q and PW_V · Pile language modeling | pseudo rank | Hurt |
| Tune and decouple attention-layer weight decay from other parameters | 125M Transformer · Pile language modeling | perplexity | Mixed |
Surprising findings
- Increasing attention weight decay beyond 0.1 starts to hurt language modelling performance, while increasing MLP weight decay beyond 0.1 is generally beneficial.
- Regularization on the factorized product prunes hidden units adaptively even in overcomplete linear networks, unlike regularizing a single matrix.
Limitations
- Theory for AdamW holds only approximately; exponential convergence is proven for gradient flow and momentum-based SDE approximations.
- Language model experiments are small scale (125M parameters) and LLaMA 2 analysis is only of pretrained weights without retraining.
Experimental setup
Datasets: The Pile ImageNet associative recall synthetic task
Benchmarks: The Pile test perplexity ImageNet
Metrics: pseudo rank (95% of nuclear norm mass) test perplexity nuclear vs Frobenius norm gap accuracy
Baselines: L2-regularized factorized matrices nuclear norm regularization
Builds on: AdamW decoupled weight decay GPT-3 training recipe Vision Transformer training protocol
Closest papers by meaning
- Why Do We Need Weight Decay in Modern Deep Learning? (NeurIPS 2024 · similarity 0.90)
- Mind the Gap: a Spectral Analysis of Rank Collapse and Signal Propagation in Attention Layers (ICML 2025 · similarity 0.88)
- From Attention to Activation: Unraveling the Enigmas of Large Language Models (ICLR 2025 · similarity 0.87)
- From Low Rank Gradient Subspace Stabilization to Low-Rank Weights: Observations, Theories, and Applications (ICML 2025 · similarity 0.86)
- Low Rank Gradients and Where to Find Them (NeurIPS 2025 · similarity 0.86)
- Wide Neural Networks Trained with Weight Decay Provably Exhibit Neural Collapse (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.