# Implicit Regularization of Gradient Flow on One-Layer Softmax Attention

arXiv: https://arxiv.org/abs/2403.08699 · HTML: https://mlsanity.com/paper/2403.08699/

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- Builds on 52 of its 60 references that are in the map.
- First posted on arXiv 13 Mar 2024.

## Builds on (most cited first, 30 of 52)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet](https://mlsanity.com/paper/2101.11986/) (2021)
- [Transformers in Vision: A Survey](https://mlsanity.com/paper/2101.01169/) (2021)
- [Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning](https://mlsanity.com/paper/2012.13255/) (2020)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Implicit Regularization in Deep Matrix Factorization](https://mlsanity.com/paper/1905.13655/) (2019)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Implicit Regularization in Matrix Factorization](https://mlsanity.com/paper/1705.09280/) (2017)
- [Gradient Descent Maximizes the Margin of Homogeneous Neural Networks](https://mlsanity.com/paper/1906.05890/) (2019)
- [Kernel and Rich Regimes in Overparametrized Models](https://mlsanity.com/paper/2002.09277/) (2020)
- [Characterizing Implicit Bias in Terms of Optimization Geometry](https://mlsanity.com/paper/1802.08246/) (2018)
- [Implicit Bias of Gradient Descent on Linear Convolutional Networks](https://mlsanity.com/paper/1806.00468/) (2018)
- [Gradient descent aligns the layers of deep linear networks](https://mlsanity.com/paper/1810.02032/) (2018)
- [Directional convergence and alignment in deep learning](https://mlsanity.com/paper/2006.06657/) (2020)
- [Towards Resolving the Implicit Bias of Gradient Descent for Matrix Factorization: Greedy Low-Rank Learning](https://mlsanity.com/paper/2012.09839/) (2020)
- [Implicit Regularization of Discrete Gradient Dynamics in Linear Neural Networks](https://mlsanity.com/paper/1904.13262/) (2019)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Convergence of Gradient Descent on Separable Data](https://mlsanity.com/paper/1803.01905/) (2018)
- [What Happens after SGD Reaches Zero Loss? --A Mathematical Framework](https://mlsanity.com/paper/2110.06914/) (2021)
- [Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer](https://mlsanity.com/paper/2305.16380/) (2023)
- [Vision Transformers provably learn spatial structure](https://mlsanity.com/paper/2210.09221/) (2022)
- [Shape Matters: Understanding the Implicit Bias of the Noise Covariance](https://mlsanity.com/paper/2006.08680/) (2020)
- [How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding](https://mlsanity.com/paper/2303.04245/) (2023)
- [Implicit Regularization for Optimal Sparse Recovery](https://mlsanity.com/paper/1909.05122/) (2019)

## Built on by (most cited first, 19 of 19)

- [Training Dynamics of In-Context Learning in Linear Attention](https://mlsanity.com/paper/2501.16265/) (2025)
- [From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency](https://mlsanity.com/paper/2410.05459/) (2024)
- [Exploring the Frontiers of Softmax: Provable Optimization, Applications in Diffusion Model, and Beyond](https://mlsanity.com/paper/2405.03251/) (2024)
- [Do LLMs dream of elephants (when told not to)? Latent concept association and associative memory in transformers](https://mlsanity.com/paper/2406.18400/) (2024)
- [Mask in the Mirror: Implicit Sparsification](https://mlsanity.com/paper/2408.09966/) (2024)
- [Gating is Weighting: Understanding Gated Linear Attention through In-context Learning](https://mlsanity.com/paper/2504.04308/) (2025)
- [Mirror, Mirror of the Flow: How Does Regularization Shape Implicit Bias?](https://mlsanity.com/paper/2504.12883/) (2025)
- [How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias](https://mlsanity.com/paper/2505.00926/) (2025)
- [Implicit Regularization of Sharpness-Aware Minimization for Scale-Invariant Problems](https://mlsanity.com/paper/2410.14802/) (2024)
- [Generalization or Hallucination? Understanding Out-of-Context Reasoning in Transformers](https://mlsanity.com/paper/2506.10887/) (2025)
- [Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection](https://mlsanity.com/paper/2410.14581/) (2024)
- [Phase Transitions in Attention: A Bayesian Theory of Copy Head Emergence](https://mlsanity.com/paper/2606.12058/) (2026)
- [Breaking the Reversal Curse in Autoregressive Language Models via Identity Bridge](https://mlsanity.com/paper/2602.02470/) (2026)
- [Revisiting Transformers through the Lens of Low Entropy and Dynamic Sparsity](https://mlsanity.com/paper/2504.18929/) (2025)
- [Attention with Trained Embeddings Provably Selects Important Tokens](https://mlsanity.com/paper/2505.17282/) (2025)
- [Provably Learning Attention with Queries](https://mlsanity.com/paper/2601.16873/) (2026)
- [Focus and Dilution: The Multi-stage Learning Process of Attention](https://mlsanity.com/paper/2605.01199/) (2026)
- [Transformer as an Euler Discretization of Score-based Variational Flow](https://mlsanity.com/paper/2604.23740/) (2026)
- [Federated LoRA Fine-Tuning for LLMs via Collaborative Alignment](https://mlsanity.com/paper/2605.21217/) (2026)

## Abstract (arXiv)

We study gradient flow on the exponential loss for a classification problem with a one-layer softmax attention model, where the key and query weight matrices are trained separately. Under a separability assumption on the data, we show that when gradient flow achieves the minimal loss value, it further implicitly minimizes the nuclear norm of the product of the key and query weight matrices. Such implicit regularization can be described by a Support Vector Machine (SVM) problem with respect to the attention weights. This finding contrasts with prior results showing that the gradient descent induces an implicit regularization on the Frobenius norm on the product weight matrix when the key and query matrices are combined into a single weight matrix for training. For diagonal key and query matrices, our analysis builds upon the reparameterization technique and exploits approximate KKT conditions of the SVM associated with the classification data. Moreover, the results are extended to general weights configurations given proper alignment of the weight matrices' singular spaces with the data features at initialization.