mlsanityData to 7 Oct 2026

Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection

arXiv 2410.14581 ↗First posted 18 Oct 2024cs.LG cs.AI cs.CLcited by 9

Not on the map. Why?

Abstract (arXiv)

Attention mechanisms have revolutionized several domains of artificial intelligence, such as natural language processing and computer vision, by enabling models to selectively focus on relevant parts of the input data. While recent work has characterized the optimization dynamics of gradient descent (GD) in attention-based models and the structural properties of its preferred solutions, less is known about more general optimization algorithms such as mirror descent (MD). In this paper, we investigate the convergence properties and implicit biases of a family of MD algorithms tailored for softmax attention mechanisms, with the potential function chosen as the p-th power of the ℓp-norm. Specifically, we show that these algorithms converge in direction to a generalized hard-margin SVM with an ℓp-norm objective when applied to a classification problem using a softmax attention model. Notably, our theoretical results reveal that the convergence rate is comparable to that of traditional GD in simpler models, despite the highly nonlinear and nonconvex nature of the present problem. Additionally, we delve into the joint optimization dynamics of the key-query matrix and the decoder, establishing conditions under which this complex joint optimization converges to their respective hard-margin SVM solutions. Lastly, our numerical experiments on real data demonstrate that MD algorithms improve generalization over standard GD and excel in optimal token selection.

Key facts. Cited by 9 papers in the map (data through 7 Oct 2026). Builds on 57 of its 65 references that are in the map. First posted on arXiv 18 Oct 2024.

Builds on

57 of 65 references.

References · 57

  1. 2017Attention Is All You Needcited by 59,932
  2. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,204
  3. 2021Learning Transferable Visual Models From Natural Language Supervisioncited by 33,456
  4. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  5. 2023GPT-4 Technical Reportcited by 18,132
  6. 2021Evaluating Large Language Models Trained on Codecited by 8,053
  7. 2014Neural Machine Translation by Jointly Learning to Align and Translatecited by 4,530
  8. 2021Zero-Shot Text-to-Image Generationcited by 4,276
  9. 2023PaLM-E: An Embodied Multimodal Language Modelcited by 2,037
  10. 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,528
  11. 2017The Implicit Bias of Gradient Descent on Separable Datacited by 649
  12. 2014Convex Optimization: Algorithms and Complexitycited by 569
38 more
  1. 2019Implicit Regularization in Deep Matrix Factorizationcited by 402
  2. 2019Gradient Descent Maximizes the Margin of Homogeneous Neural Networkscited by 308
  3. 2020Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Losscited by 267
  4. 2018Characterizing Implicit Bias in Terms of Optimization Geometrycited by 259
  5. 2020Directional convergence and alignment in deep learningcited by 172
  6. 2020Towards Resolving the Implicit Bias of Gradient Descent for Matrix Factorization: Greedy Low-Rank Learningcited by 148
  7. 2023In-Context Convergence of Transformerscited by 112
  8. 2018Convergence of Gradient Descent on Separable Datacited by 108
  9. 2023Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformercited by 101
  10. 2022Vision Transformers provably learn spatial structurecited by 92
  11. 2022On the Implicit Bias in Deep-Learning Algorithmscited by 84
  12. 2018Risk and parameter convergence of logistic regressioncited by 81
  13. 2023A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexitycited by 77
  14. 2024Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimalitycited by 71
  15. 2023Transformers as Support Vector Machinescited by 67
  16. 2023Max-Margin Token Selection in Attention Mechanismcited by 57
  17. 2023JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attentioncited by 52
  18. 2023On the Role of Attention in Prompt-tuningcited by 49
  19. 2023Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regressioncited by 47
  20. 2022Implicit Bias in Leaky ReLU Networks Trained on High-Dimensional Datacited by 46
  21. 2020Implicit Regularization in ReLU Networks with the Square Losscited by 45
  22. 2019Stochastic Mirror Descent on Overparameterized Nonlinear Models: Convergence, Implicit Regularization, and Generalizationcited by 40
  23. 2024An Information-Theoretic Analysis of In-Context Learningcited by 39
  24. 2024Attention with Markov: A Framework for Principled Analysis of Transformers via Markov Chainscited by 39
  25. 2023On the Optimization and Generalization of Multi-head Attentioncited by 37
  26. 2024From Self-Attention to Markov Models: Unveiling the Dynamics of Generative Transformerscited by 34
  27. 2024Mechanics of Next Token Prediction with Self-Attentioncited by 34
  28. 2021Fast Margin Maximization via Dual Accelerationcited by 33
  29. 2024Implicit Bias and Fast Convergence Rates for Self-attentioncited by 31
  30. 2022Unraveling Attention via Convex Duality: Analysis and Interpretations of Vision Transformerscited by 30
  31. 2018Stochastic Gradient/Mirror Descent: Minimax Optimality and Implicit Regularizationcited by 29
  32. 2022Mirror Descent Maximizes Generalized Margin and Can Be Implemented Efficientlycited by 25
  33. 2024Implicit Geometry of Next-token Prediction: From Language Sparsity Patterns to Model Representationscited by 24
  34. 2024In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitznesscited by 22
  35. 2024Provably learning a multi-head attention layercited by 21
  36. 2024Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannotcited by 21
  37. 2024Self-attention Networks Localize When QK-eigenspectrum Concentratescited by 19
  38. 2024Implicit Regularization of Gradient Flow on One-Layer Softmax Attentioncited by 16

50 most cited of 57 All 57 (JSON)

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2025: 7; 2026: 2.

9 papers.

  1. 2025Training Dynamics of In-Context Learning in Linear Attentioncited by 44
  2. 2025Implicit Bias of Spectral Descent and Muon on Multiclass Separable Datacited by 33
  3. 2025Gating is Weighting: Understanding Gated Linear Attention through In-context Learningcited by 11
  4. 2025Mirror, Mirror of the Flow: How Does Regularization Shape Implicit Bias?cited by 10
  5. 2025Dynamics of Spontaneous Topic Changes in Next Token Prediction with Self-Attentioncited by 2
  6. 2025Diagonalizing the Softmax: Hadamard Initialization for Tractable Cross-Entropy Dynamicscited by 2
  7. 2025Attention with Trained Embeddings Provably Selects Important Tokenscited by 0
  8. 2026Implicit Bias of Mirror Flow in Homogeneous Neural Networks: Sparse and Dense Feature Learningcited by 0
  9. 2026The Implicit Bias of Depth: From Neural Collapse to Softmax Codescited by 0