# Max-Margin Token Selection in Attention Mechanism

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

## Key facts

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

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [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)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Prefix-Tuning: Optimizing Continuous Prompts for Generation](https://mlsanity.com/paper/2101.00190/) (2021)
- [The Power of Scale for Parameter-Efficient Prompt Tuning](https://mlsanity.com/paper/2104.08691/) (2021)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [PaLM-E: An Embodied Multimodal Language Model](https://mlsanity.com/paper/2303.03378/) (2023)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [BERT4Rec: Sequential Recommendation with Bidirectional Encoder Representations from Transformer](https://mlsanity.com/paper/1904.06690/) (2019)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [Gradient Descent Finds Global Minima of Deep Neural Networks](https://mlsanity.com/paper/1811.03804/) (2018)
- [Deep Interest Network for Click-Through Rate Prediction](https://mlsanity.com/paper/1706.06978/) (2017)
- [Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth](https://mlsanity.com/paper/2103.03404/) (2021)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)

## Built on by (most cited first, 30 of 58)

- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention](https://mlsanity.com/paper/2310.00535/) (2023)
- [Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression](https://mlsanity.com/paper/2310.17086/) (2023)
- [Training Dynamics of In-Context Learning in Linear Attention](https://mlsanity.com/paper/2501.16265/) (2025)
- [Implicit Bias of Spectral Descent and Muon on Multiclass Separable Data](https://mlsanity.com/paper/2502.04664/) (2025)
- [What can a Single Attention Layer Learn? A Study Through the Random Features Lens](https://mlsanity.com/paper/2307.11353/) (2023)
- [Attention with Markov: A Framework for Principled Analysis of Transformers via Markov Chains](https://mlsanity.com/paper/2402.04161/) (2024)
- [How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?](https://mlsanity.com/paper/2402.15607/) (2024)
- [From Self-Attention to Markov Models: Unveiling the Dynamics of Generative Transformers](https://mlsanity.com/paper/2402.13512/) (2024)
- [Implicit Bias and Fast Convergence Rates for Self-attention](https://mlsanity.com/paper/2402.05738/) (2024)
- [A phase transition between positional and semantic learning in a solvable model of dot-product attention](https://mlsanity.com/paper/2402.03902/) (2024)
- [Cognitive Memory in Large Language Models](https://mlsanity.com/paper/2504.02441/) (2025)
- [Linear Transformers are Versatile In-Context Learners](https://mlsanity.com/paper/2402.14180/) (2024)
- [Self-attention Networks Localize When QK-eigenspectrum Concentrates](https://mlsanity.com/paper/2402.02098/) (2024)
- [Implicit Regularization of Gradient Flow on One-Layer Softmax Attention](https://mlsanity.com/paper/2403.08699/) (2024)
- [Enhancing Graph Transformers with Hierarchical Distance Structural Encoding](https://mlsanity.com/paper/2308.11129/) (2023)
- [Selective Attention: Enhancing Transformer through Principled Context Control](https://mlsanity.com/paper/2411.12892/) (2024)
- [Trained Transformer Classifiers Generalize and Exhibit Benign Overfitting In-Context](https://mlsanity.com/paper/2410.01774/) (2024)
- [Transformers Learn Low Sensitivity Functions: Investigations and Implications](https://mlsanity.com/paper/2403.06925/) (2024)
- [What Improves the Generalization of Graph Transformers? A Theoretical Dive into the Self-attention and Positional Encoding](https://mlsanity.com/paper/2406.01977/) (2024)
- [Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis](https://mlsanity.com/paper/2410.02167/) (2024)
- [Non-asymptotic Convergence of Training Transformers for Next-token Prediction](https://mlsanity.com/paper/2409.17335/) (2024)
- [The Rich and the Simple: On the Implicit Bias of Adam and SGD](https://mlsanity.com/paper/2505.24022/) (2025)
- [Gating is Weighting: Understanding Gated Linear Attention through In-context Learning](https://mlsanity.com/paper/2504.04308/) (2025)
- [How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias](https://mlsanity.com/paper/2505.00926/) (2025)
- [Generalization or Hallucination? Understanding Out-of-Context Reasoning in Transformers](https://mlsanity.com/paper/2506.10887/) (2025)
- [Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformers](https://mlsanity.com/paper/2405.15926/) (2024)
- [Theoretical Insights into Fine-Tuning Attention Mechanism: Generalization and Optimization](https://mlsanity.com/paper/2410.02247/) (2024)
- [Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis](https://mlsanity.com/paper/2410.09605/) (2024)
- [Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection](https://mlsanity.com/paper/2410.14581/) (2024)

## 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.