"Weight decay induces low-rank attention layers"

Your transformer's attention layers might be unnecessarily constrained by weight decay

Weight decay secretly forces attention layers to operate in lower dimensions

https://arxiv.org/abs/2410.23819

🔍 Original Problem:

Weight decay, commonly used in training transformers (λ=0.1 in GPT-3, LLaMA, LLaMA 2, ViT), has poorly understood effects on attention layers. The impact of L2-regularization on matrix products in attention mechanisms needs investigation.

-----

🛠️ Solution in this Paper:

The researchers mathematically prove that L2-regularization on factorized matrices (like attention layers) induces low-rank solutions. They show the equivalence between L2 and nuclear norm regularization during training, explaining why this affects model performance early in training. The solution involves analyzing the optimization dynamics and providing theoretical guarantees for this equivalence.

-----

💡 Key Insights:

→ Weight decay in attention layers unintentionally creates low-rank solutions

→ The effect happens exponentially quickly during training, long before convergence

→ This rank reduction can hurt model performance in language tasks

→ Decoupling weight decay in attention layers from other parameters may improve results

-----

📊 Results:

→ Analysis of pre-trained LLaMA 2 and Vision Transformer weights shows clear signs of rank reduction

→ Empirical validation across deep linear networks to LLMs confirms theoretical predictions

→ The rank-reducing effect correlates directly with weight decay strength

→ Performance improves when decoupling attention layer weight decay from other parameters