mlsanityData to 7 Oct 2026

Implicit Regularization of Gradient Flow on One-Layer Softmax Attention

arXiv 2403.08699 ↗First posted 13 Mar 2024cs.LG cs.AI math.OCcited by 16

Not on the map. Why?

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.

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

Builds on

52 of 60 references.

References · 52

  1. 2017Attention Is All You Needcited by 59,932
  2. 2020Language Models are Few-Shot Learnerscited by 32,199
  3. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  4. 2023GPT-4 Technical Reportcited by 18,132
  5. 2020Rethinking Attention with Performerscited by 1,584
  6. 2020Linformer: Self-Attention with Linear Complexitycited by 1,462
  7. 2021Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNetcited by 1,096
  8. 2021Transformers in Vision: A Surveycited by 1,009
  9. 2020Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuningcited by 658
  10. 2017The Implicit Bias of Gradient Descent on Separable Datacited by 649
  11. 2019Implicit Regularization in Deep Matrix Factorizationcited by 402
  12. 2023Trained Transformers Learn Linear Models In-Contextcited by 331
38 more
  1. 2017Implicit Regularization in Matrix Factorizationcited by 318
  2. 2019Gradient Descent Maximizes the Margin of Homogeneous Neural Networkscited by 308
  3. 2020Kernel and Rich Regimes in Overparametrized Modelscited by 302
  4. 2018Characterizing Implicit Bias in Terms of Optimization Geometrycited by 259
  5. 2018Implicit Bias of Gradient Descent on Linear Convolutional Networkscited by 251
  6. 2018Gradient descent aligns the layers of deep linear networkscited by 182
  7. 2020Directional convergence and alignment in deep learningcited by 172
  8. 2020Towards Resolving the Implicit Bias of Gradient Descent for Matrix Factorization: Greedy Low-Rank Learningcited by 148
  9. 2019Implicit Regularization of Discrete Gradient Dynamics in Linear Neural Networkscited by 122
  10. 2024How Transformers Learn Causal Structure with Gradient Descentcited by 114
  11. 2023In-Context Convergence of Transformerscited by 112
  12. 2018Convergence of Gradient Descent on Separable Datacited by 108
  13. 2021What Happens after SGD Reaches Zero Loss? --A Mathematical Frameworkcited by 106
  14. 2023Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformercited by 101
  15. 2022Vision Transformers provably learn spatial structurecited by 92
  16. 2020Shape Matters: Understanding the Implicit Bias of the Noise Covariancecited by 88
  17. 2023How Do Transformers Learn Topic Structure: Towards a Mechanistic Understandingcited by 85
  18. 2019Implicit Regularization for Optimal Sparse Recoverycited by 78
  19. 2021Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstructioncited by 77
  20. 2023A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexitycited by 77
  21. 2020A Unifying View on Implicit Bias in Training Linear Neural Networkscited by 74
  22. 2024Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimalitycited by 71
  23. 2020Implicit Bias in Deep Linear Classification: Initialization Scale vs Training Accuracycited by 70
  24. 2021On the Implicit Bias of Initialization Shape: Beyond Infinitesimal Mirror Descentcited by 69
  25. 2023Transformers as Support Vector Machinescited by 67
  26. 2023Max-Margin Token Selection in Attention Mechanismcited by 57
  27. 2023Saddle-to-Saddle Dynamics in Diagonal Linear Networkscited by 54
  28. 2019Lexicographic and Depth-Sensitive Margins in Homogeneous and Non-Homogeneous Deep Modelscited by 53
  29. 2023On the Role of Attention in Prompt-tuningcited by 49
  30. 2022Implicit Bias of Gradient Descent on Reparametrized Models: On Equivalence to Mirror Descentcited by 38
  31. 2023On the Optimization and Generalization of Multi-head Attentioncited by 37
  32. 2023(S)GD over Diagonal Linear Networks: Implicit Regularisation, Large Stepsizes and Edge of Stabilitycited by 35
  33. 2021Fast Margin Maximization via Dual Accelerationcited by 33
  34. 2024Implicit Bias and Fast Convergence Rates for Self-attentioncited by 31
  35. 2021Implicit Sparse Regularization: The Impact of Depth and Early Stoppingcited by 30
  36. 2018Stochastic Gradient/Mirror Descent: Minimax Optimality and Implicit Regularizationcited by 29
  37. 2021More is Less: Inducing Sparsity via Overparameterizationcited by 22
  38. 2020Understanding Implicit Regularization in Over-Parameterized Single Index Modelcited by 12

50 most cited of 52 All 52 (JSON)

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 6; 2025: 7; 2026: 3.

16 papers.

Sort:

By area:

  1. 2025Training Dynamics of In-Context Learning in Linear Attentioncited by 44
  2. 2024From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiencycited by 28
  3. 2024Exploring the Frontiers of Softmax: Provable Optimization, Applications in Diffusion Model, and Beyondcited by 15
  4. 2024Do LLMs dream of elephants (when told not to)? Latent concept association and associative memory in transformerscited by 14
  5. 2024Mask in the Mirror: Implicit Sparsificationcited by 12
  6. 2025Gating is Weighting: Understanding Gated Linear Attention through In-context Learningcited by 11
  7. 2025Mirror, Mirror of the Flow: How Does Regularization Shape Implicit Bias?cited by 10
  8. 2025How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Biascited by 10
  9. 2025Generalization or Hallucination? Understanding Out-of-Context Reasoning in Transformerscited by 10
  10. 2024Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selectioncited by 9
  11. 2024Implicit Regularization of Sharpness-Aware Minimization for Scale-Invariant Problemscited by 9
  12. 2026Breaking the Reversal Curse in Autoregressive Language Models via Identity Bridgecited by 2
4 more
  1. 2025Revisiting Transformers through the Lens of Low Entropy and Dynamic Sparsitycited by 1
  2. 2026Provably Learning Attention with Queriescited by 1
  3. 2025Attention with Trained Embeddings Provably Selects Important Tokenscited by 0
  4. 2026Transformer as an Euler Discretization of Score-based Variational Flowcited by 0