# Implicit Bias and Fast Convergence Rates for Self-attention

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

## Key facts

- Cited by 31 papers in the map (data through 14 Oct 2026).
- Builds on 66 of its 81 references that are in the map.
- First posted on arXiv 8 Feb 2024.

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

- [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)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [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)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (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)
- [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)
- [Implicit Regularization in Deep Matrix Factorization](https://mlsanity.com/paper/1905.13655/) (2019)
- [Nuanced Metrics for Measuring Unintended Bias with Real Data for Text Classification](https://mlsanity.com/paper/1903.04561/) (2019)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [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)
- [Loss landscapes and optimization in over-parameterized non-linear systems and neural networks](https://mlsanity.com/paper/2003.00307/) (2020)
- [Generalization Bounds of Stochastic Gradient Descent for Wide and Deep Neural Networks](https://mlsanity.com/paper/1905.13210/) (2019)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [Directional convergence and alignment in deep learning](https://mlsanity.com/paper/2006.06657/) (2020)
- [Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models](https://mlsanity.com/paper/2208.06677/) (2022)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [Towards Resolving the Implicit Bias of Gradient Descent for Matrix Factorization: Greedy Low-Rank Learning](https://mlsanity.com/paper/2012.09839/) (2020)
- [Stochastic Polyak Step-size for SGD: An Adaptive Learning Rate for Fast Convergence](https://mlsanity.com/paper/2002.10542/) (2020)
- [Representational Strengths and Limitations of Transformers](https://mlsanity.com/paper/2306.02896/) (2023)
- [Vision Transformer for Small-Size Datasets](https://mlsanity.com/paper/2112.13492/) (2021)

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

- [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)
- [Implicit Regularization of Gradient Flow on One-Layer Softmax Attention](https://mlsanity.com/paper/2403.08699/) (2024)
- [Trained Transformer Classifiers Generalize and Exhibit Benign Overfitting In-Context](https://mlsanity.com/paper/2410.01774/) (2024)
- [Implicit Optimization Bias of Next-Token Prediction in Linear Models](https://mlsanity.com/paper/2402.18551/) (2024)
- [Transformers Learn Low Sensitivity Functions: Investigations and Implications](https://mlsanity.com/paper/2403.06925/) (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)
- [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)
- [How Transformers Get Rich: Approximation and Dynamics Analysis](https://mlsanity.com/paper/2410.11474/) (2024)
- [Transformative or Conservative? Conservation laws for ResNets and Transformers](https://mlsanity.com/paper/2506.06194/) (2025)
- [On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recovery](https://mlsanity.com/paper/2410.13981/) (2024)
- [When Do Transformers Outperform Feedforward and Recurrent Networks? A Statistical Perspective](https://mlsanity.com/paper/2503.11272/) (2025)
- [Diagonalizing the Softmax: Hadamard Initialization for Tractable Cross-Entropy Dynamics](https://mlsanity.com/paper/2512.04006/) (2025)
- [Unlabeled Data Can Provably Enhance In-Context Learning of Transformers](https://mlsanity.com/paper/2601.10058/) (2026)
- [Breaking the Reversal Curse in Autoregressive Language Models via Identity Bridge](https://mlsanity.com/paper/2602.02470/) (2026)
- [Revisiting Transformers through the Lens of Low Entropy and Dynamic Sparsity](https://mlsanity.com/paper/2504.18929/) (2025)
- [How Transformers Learn In-Context Recall Tasks? Optimality, Training Dynamics and Generalization](https://mlsanity.com/paper/2505.15009/) (2025)
- [Attention with Trained Embeddings Provably Selects Important Tokens](https://mlsanity.com/paper/2505.17282/) (2025)
- [Learning Linear Regression with Low-Rank Tasks in-Context](https://mlsanity.com/paper/2510.04548/) (2025)
- [Transformers are almost optimal metalearners for linear classification](https://mlsanity.com/paper/2510.19797/) (2025)
- [Finite-Time Analysis of Gradient Descent for Shallow Transformers](https://mlsanity.com/paper/2601.16514/) (2026)
- [Variational Deep Learning via Implicit Regularization](https://mlsanity.com/paper/2505.20235/) (2025)
- [Gaussian Match-and-Copy: A Minimalist Benchmark for Studying Transformer Induction](https://mlsanity.com/paper/2602.07562/) (2026)
- [Transformer as an Euler Discretization of Score-based Variational Flow](https://mlsanity.com/paper/2604.23740/) (2026)
- [The Implicit Bias of Depth: From Neural Collapse to Softmax Codes](https://mlsanity.com/paper/2605.23087/) (2026)

## Abstract (arXiv)

We study the fundamental optimization principles of self-attention, the defining mechanism of transformers, by analyzing the implicit bias of gradient-based optimizers in training a self-attention layer with a linear decoder in binary classification. Building on prior studies in linear logistic regression, recent findings demonstrate that the key-query matrix Wt from gradient-descent (GD) converges in direction towards Wmm, which maximizes the margin between optimal and non-optimal tokens across sequences. However, this convergence is local, dependent on initial conditions, only holds asymptotically as the number of iterations increases, and leaves questions about the potential benefits of adaptive step-size rules unaddressed. To bridge this gap, we first establish scenarios for which convergence is provably global. We then analyze two adaptive step-size strategies: normalized GD and Polyak step-size, demonstrating finite-time convergence rates for Wt to Wmm, and quantifying the sparsification rate of the attention map. These findings not only show that these strategies can accelerate parameter convergence over standard GD in a non-convex setting but also deepen the understanding of the implicit bias in self-attention, linking it more closely to the phenomena observed in linear logistic regression despite its intricate non-convex nature.
