# Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection arXiv: https://arxiv.org/abs/2410.14581 · HTML: https://mlsanity.com/paper/2410.14581/ ## Key facts - Cited by 9 papers in the map (data through 14 Oct 2026). - Builds on 57 of its 65 references that are in the map. - First posted on arXiv 18 Oct 2024. ## Builds on (most cited first, 30 of 57) - [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) - [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) - [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021) - [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) - [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) - [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017) - [Convex Optimization: Algorithms and Complexity](https://mlsanity.com/paper/1405.4980/) (2014) - [Implicit Regularization in Deep Matrix Factorization](https://mlsanity.com/paper/1905.13655/) (2019) - [Gradient Descent Maximizes the Margin of Homogeneous Neural Networks](https://mlsanity.com/paper/1906.05890/) (2019) - [Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss](https://mlsanity.com/paper/2002.04486/) (2020) - [Characterizing Implicit Bias in Terms of Optimization Geometry](https://mlsanity.com/paper/1802.08246/) (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) - [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) - [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) - [On the Implicit Bias in Deep-Learning Algorithms](https://mlsanity.com/paper/2208.12591/) (2022) - [Risk and parameter convergence of logistic regression](https://mlsanity.com/paper/1803.07300/) (2018) - [A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity](https://mlsanity.com/paper/2302.06015/) (2023) - [Transformers as Support Vector Machines](https://mlsanity.com/paper/2308.16898/) (2023) - [Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality](https://mlsanity.com/paper/2402.19442/) (2024) - [Max-Margin Token Selection in Attention Mechanism](https://mlsanity.com/paper/2306.13596/) (2023) - [JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention](https://mlsanity.com/paper/2310.00535/) (2023) - [On the Role of Attention in Prompt-tuning](https://mlsanity.com/paper/2306.03435/) (2023) ## Built on by (most cited first, 9 of 9) - [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) - [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) - [Implicit Bias of Mirror Flow in Homogeneous Neural Networks: Sparse and Dense Feature Learning](https://mlsanity.com/paper/2605.19458/) (2026) - [Dynamics of Spontaneous Topic Changes in Next Token Prediction with Self-Attention](https://mlsanity.com/paper/2501.06382/) (2025) - [Diagonalizing the Softmax: Hadamard Initialization for Tractable Cross-Entropy Dynamics](https://mlsanity.com/paper/2512.04006/) (2025) - [Attention with Trained Embeddings Provably Selects Important Tokens](https://mlsanity.com/paper/2505.17282/) (2025) - [The Implicit Bias of Depth: From Neural Collapse to Softmax Codes](https://mlsanity.com/paper/2605.23087/) (2026) ## 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.