mlsanityData to 7 Oct 2026

Max-Margin Token Selection in Attention Mechanism

arXiv 2306.13596 ↗First posted 23 Jun 2023cs.LG cs.AI cs.CLcited by 57

Not on the map. Why?

Abstract (arXiv)

Attention mechanism is a central component of the transformer architecture which led to the phenomenal success of large language models. However, the theoretical principles underlying the attention mechanism are poorly understood, especially its nonconvex optimization dynamics. In this work, we explore the seminal softmax-attention model f(boldsymbolX)=langle boldsymbolXv, textttsoftmax(boldsymbolXWp)rangle, where boldsymbolX is the token sequence and (boldsymbolv,boldsymbolW,boldsymbolp) are trainable parameters. We prove that running gradient descent on boldsymbolp, or equivalently boldsymbolW, converges in direction to a max-margin solution that separates textitlocally-optimal tokens from non-optimal ones. This clearly formalizes attention as an optimal token selection mechanism. Remarkably, our results are applicable to general data and precisely characterize textitoptimality of tokens in terms of the value embeddings boldsymbolXv and problem geometry. We also provide a broader regularization path analysis that establishes the margin maximizing nature of attention even for nonlinear prediction heads. When optimizing boldsymbolv and boldsymbolp simultaneously with logistic loss, we identify conditions under which the regularization paths directionally converge to their respective hard-margin SVM solutions where boldsymbolv separates the input features based on their labels. Interestingly, the SVM formulation of boldsymbolp is influenced by the support vector geometry of boldsymbolv. Finally, we verify our theoretical findings via numerical experiments and provide insights.

Key facts. Cited by 57 papers in the map (data through 7 Oct 2026). Builds on 85 of its 105 references that are in the map. First posted on arXiv 23 Jun 2023.

Builds on

85 of 105 references.

References · 85

  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. 2020Language Models are Few-Shot Learnerscited by 32,199
  5. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  6. 2023GPT-4 Technical Reportcited by 18,132
  7. 2022Chain-of-Thought Prompting Elicits Reasoning in Large Language Modelscited by 12,302
  8. 2021Evaluating Large Language Models Trained on Codecited by 8,053
  9. 2022PaLM: Scaling Language Modeling with Pathwayscited by 5,304
  10. 2020Training data-efficient image transformers & distillation through attentioncited by 4,560
  11. 2014Neural Machine Translation by Jointly Learning to Align and Translatecited by 4,530
  12. 2021Zero-Shot Text-to-Image Generationcited by 4,276
38 more
  1. 2021On the Opportunities and Risks of Foundation Modelscited by 3,961
  2. 2021Prefix-Tuning: Optimizing Continuous Prompts for Generationcited by 3,878
  3. 2021The Power of Scale for Parameter-Efficient Prompt Tuningcited by 3,769
  4. 2021Finetuned Language Models Are Zero-Shot Learnerscited by 3,747
  5. 2018Neural Tangent Kernel: Convergence and Generalization in Neural Networkscited by 2,597
  6. 2023PaLM-E: An Embodied Multimodal Language Modelcited by 2,037
  7. 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,528
  8. 2019BERT4Rec: Sequential Recommendation with Bidirectional Encoder Representations from Transformercited by 952
  9. 2022A Generalist Agentcited by 802
  10. 2018A Convergence Theory for Deep Learning via Over-Parameterizationcited by 784
  11. 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 700
  12. 2022What learning algorithm is in-context learning? Investigations with linear modelscited by 671
  13. 2017The Implicit Bias of Gradient Descent on Separable Datacited by 649
  14. 2018Universal Transformerscited by 613
  15. 2018Gradient Descent Finds Global Minima of Deep Neural Networkscited by 572
  16. 2017Deep Interest Network for Click-Through Rate Predictioncited by 541
  17. 2021Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depthcited by 411
  18. 2019Are Transformers universal approximators of sequence-to-sequence functions?cited by 404
  19. 2019Implicit Regularization in Deep Matrix Factorizationcited by 402
  20. 2017A Deep Reinforced Model for Abstractive Summarizationcited by 364
  21. 2017A Structured Self-attentive Sentence Embeddingcited by 318
  22. 2019Gradient Descent Maximizes the Margin of Homogeneous Neural Networkscited by 308
  23. 2020Kernel and Rich Regimes in Overparametrized Modelscited by 302
  24. 2023Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspectivecited by 293
  25. 2020Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Losscited by 267
  26. 2018Characterizing Implicit Bias in Terms of Optimization Geometrycited by 259
  27. 2023Looped Transformers as Programmable Computerscited by 224
  28. 2023Transformers as Algorithms: Generalization and Stability in In-context Learningcited by 220
  29. 2016A Decomposable Attention Model for Natural Language Inferencecited by 214
  30. 2022Online Decision Transformercited by 194
  31. 2018Gradient descent aligns the layers of deep linear networkscited by 182
  32. 2019Towards moderate overparameterization: global convergence guarantees for training shallow neural networkscited by 176
  33. 2020Directional convergence and alignment in deep learningcited by 172
  34. 2020The Lipschitz Constant of Self-Attentioncited by 171
  35. 2019Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networkscited by 162
  36. 2016Long Short-Term Memory-Networks for Machine Readingcited by 161
  37. 2021All Tokens Matter: Token Labeling for Training Better Vision Transformerscited by 144
  38. 2021Inductive Biases and Variable Creation in Self-Attention Mechanismscited by 144

50 most cited of 85 All 85 (JSON)

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2023: 4; 2024: 24; 2025: 16; 2026: 13.

57 papers.

Sort:

By area:

  1. 2024How Transformers Learn Causal Structure with Gradient Descentcited by 114
  2. 2023JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attentioncited by 52
  3. 2023Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regressioncited by 47
  4. 2025Training Dynamics of In-Context Learning in Linear Attentioncited by 44
  5. 2023What can a Single Attention Layer Learn? A Study Through the Random Features Lenscited by 39
  6. 2024Attention with Markov: A Framework for Principled Analysis of Transformers via Markov Chainscited by 39
  7. 2024How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?cited by 35
  8. 2024From Self-Attention to Markov Models: Unveiling the Dynamics of Generative Transformerscited by 34
  9. 2025Implicit Bias of Spectral Descent and Muon on Multiclass Separable Datacited by 33
  10. 2024Implicit Bias and Fast Convergence Rates for Self-attentioncited by 31
  11. 2024A phase transition between positional and semantic learning in a solvable model of dot-product attentioncited by 30
  12. 2024Linear Transformers are Versatile In-Context Learnerscited by 28
38 more
  1. 2025Cognitive Memory in Large Language Modelscited by 28
  2. 2024Self-attention Networks Localize When QK-eigenspectrum Concentratescited by 19
  3. 2023Enhancing Graph Transformers with Hierarchical Distance Structural Encodingcited by 16
  4. 2024Implicit Regularization of Gradient Flow on One-Layer Softmax Attentioncited by 16
  5. 2024Selective Attention: Enhancing Transformer through Principled Context Controlcited by 16
  6. 2024What Improves the Generalization of Graph Transformers? A Theoretical Dive into the Self-attention and Positional Encodingcited by 14
  7. 2024Trained Transformer Classifiers Generalize and Exhibit Benign Overfitting In-Contextcited by 14
  8. 2024Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysiscited by 14
  9. 2024Transformers Learn Low Sensitivity Functions: Investigations and Implicationscited by 13
  10. 2024Non-asymptotic Convergence of Training Transformers for Next-token Predictioncited by 13
  11. 2025The Rich and the Simple: On the Implicit Bias of Adam and SGDcited by 12
  12. 2025Gating is Weighting: Understanding Gated Linear Attention through In-context Learningcited by 11
  13. 2025How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Biascited by 10
  14. 2025Generalization or Hallucination? Understanding Out-of-Context Reasoning in Transformerscited by 10
  15. 2024Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformerscited by 9
  16. 2024Theoretical Insights into Fine-Tuning Attention Mechanism: Generalization and Optimizationcited by 9
  17. 2024Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selectioncited by 9
  18. 2024Learning on Transformers is Provable Low-Rank and Sparse: A One-layer Analysiscited by 8
  19. 2024Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysiscited by 8
  20. 2024How Transformers Get Rich: Approximation and Dynamics Analysiscited by 7
  21. 2025Sequential keypoint density estimator: an overlooked baseline of skeleton-based video anomaly detectioncited by 7
  22. 2024On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recoverycited by 6
  23. 2024Can Implicit Bias Imply Adversarial Robustness?cited by 5
  24. 2024On the Power of Convolution Augmented Transformercited by 5
  25. 2025On the Convergence of Gradient Descent on Learning Transformers with Residual Connectionscited by 5
  26. 2026Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioningcited by 4
  27. 2025Dynamics of Spontaneous Topic Changes in Next Token Prediction with Self-Attentioncited by 2
  28. 2025Diagonalizing the Softmax: Hadamard Initialization for Tractable Cross-Entropy Dynamicscited by 2
  29. 2026Unlabeled Data Can Provably Enhance In-Context Learning of Transformerscited by 2
  30. 2026Breaking the Reversal Curse in Autoregressive Language Models via Identity Bridgecited by 2
  31. 2025How Transformers Learn In-Context Recall Tasks? Optimality, Training Dynamics and Generalizationcited by 1
  32. 2025Advancing Complex Wide-Area Scene Understanding with Hierarchical Coresets Selectioncited by 1
  33. 2025Transformers are almost optimal metalearners for linear classificationcited by 1
  34. 2025Towards Understanding Transformers in Learning Random Walkscited by 1
  35. 2026Provably Learning Attention with Queriescited by 1
  36. 2025Attention with Trained Embeddings Provably Selects Important Tokenscited by 0
  37. 2026Finite-Time Analysis of Gradient Descent for Shallow Transformerscited by 0
  38. 2026Gaussian Match-and-Copy: A Minimalist Benchmark for Studying Transformer Inductioncited by 0

50 most cited of 57 All 57 (JSON)