# Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot

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

## Key facts

- Cited by 22 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 67 references that are in the map.
- First posted on arXiv 11 Jun 2024.

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

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Show Your Work: Scratchpads for Intermediate Computation with Language Models](https://mlsanity.com/paper/2112.00114/) (2021)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [A Survey on In-context Learning](https://mlsanity.com/paper/2301.00234/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [The Impact of Positional Encoding on Length Generalization in Transformers](https://mlsanity.com/paper/2305.19466/) (2023)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Transformers Learn Shortcuts to Automata](https://mlsanity.com/paper/2210.10749/) (2022)
- [Looped Transformers as Programmable Computers](https://mlsanity.com/paper/2301.13196/) (2023)
- [Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel](https://mlsanity.com/paper/1908.11775/) (2019)
- [What Algorithms can Transformers Learn? A Study in Length Generalization](https://mlsanity.com/paper/2310.16028/) (2023)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability](https://mlsanity.com/paper/2209.15594/) (2022)
- [Transformer Language Models without Positional Encodings Still Learn Positional Information](https://mlsanity.com/paper/2203.16634/) (2022)
- [Representational Strengths and Limitations of Transformers](https://mlsanity.com/paper/2306.02896/) (2023)
- [On the Turing Completeness of Modern Neural Network Architectures](https://mlsanity.com/paper/1901.03429/) (2019)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [Randomized Positional Encodings Boost Length Generalization of Transformers](https://mlsanity.com/paper/2305.16843/) (2023)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Self-Attention Networks Can Process Bounded Hierarchical Languages](https://mlsanity.com/paper/2105.11115/) (2021)
- [Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer](https://mlsanity.com/paper/2305.16380/) (2023)
- [Transformers Can Achieve Length Generalization But Not Robustly](https://mlsanity.com/paper/2402.09371/) (2024)

## Built on by (most cited first, 22 of 22)

- [Training Dynamics of In-Context Learning in Linear Attention](https://mlsanity.com/paper/2501.16265/) (2025)
- [From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency](https://mlsanity.com/paper/2410.05459/) (2024)
- [Attention layers provably solve single-location regression](https://mlsanity.com/paper/2410.01537/) (2024)
- [On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures](https://mlsanity.com/paper/2410.11778/) (2024)
- [Gating is Weighting: Understanding Gated Linear Attention through In-context Learning](https://mlsanity.com/paper/2504.04308/) (2025)
- [On the Optimal Memorization Capacity of Transformers](https://mlsanity.com/paper/2409.17677/) (2024)
- [Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection](https://mlsanity.com/paper/2410.14581/) (2024)
- [Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networks](https://mlsanity.com/paper/2506.02651/) (2025)
- [How Transformers Get Rich: Approximation and Dynamics Analysis](https://mlsanity.com/paper/2410.11474/) (2024)
- [The Role of Sparsity for Length Generalization in Transformers](https://mlsanity.com/paper/2502.16792/) (2025)
- [Quantitative Bounds for Length Generalization in Transformers](https://mlsanity.com/paper/2510.27015/) (2025)
- [Transformers Can Overcome the Curse of Dimensionality: A Theoretical Study from an Approximation Perspective](https://mlsanity.com/paper/2504.13558/) (2025)
- [On the Robustness of Transformers against Context Hijacking for Linear Classification](https://mlsanity.com/paper/2502.15609/) (2025)
- [Neural Networks Generalize on Low Complexity Data](https://mlsanity.com/paper/2409.12446/) (2024)
- [When Do Transformers Outperform Feedforward and Recurrent Networks? A Statistical Perspective](https://mlsanity.com/paper/2503.11272/) (2025)
- [Towards Understanding Transformers in Learning Random Walks](https://mlsanity.com/paper/2511.23239/) (2025)
- [Incremental Learning of Sparse Attention Patterns in Transformers](https://mlsanity.com/paper/2602.19143/) (2026)
- [Transformers Provably Learn to Internalize Chain-of-Thought](https://mlsanity.com/paper/2605.28600/) (2026)
- [Attention with Trained Embeddings Provably Selects Important Tokens](https://mlsanity.com/paper/2505.17282/) (2025)
- [Provably Learning Attention with Queries](https://mlsanity.com/paper/2601.16873/) (2026)
- [Theory on Attention Dynamics for Out-of-Distribution In-Context Learning](https://mlsanity.com/paper/2609.36448/) (2026)
- [Training Dynamics of Transformers on Unbounded Inputs: Attention Learning Can Lag Output Layer Learning](https://mlsanity.com/paper/2610.20337/) (2026)

## Abstract (arXiv)

The transformer architecture has prevailed in various deep learning settings due to its exceptional capabilities to select and compose structural information. Motivated by these capabilities, Sanford et al. proposed the sparse token selection task, in which transformers excel while fully-connected networks (FCNs) fail in the worst case. Building upon that, we strengthen the FCN lower bound to an average-case setting and establish an algorithmic separation of transformers over FCNs. Specifically, a one-layer transformer trained with gradient descent provably learns the sparse token selection task and, surprisingly, exhibits strong out-of-distribution length generalization. We provide empirical simulations to justify our theoretical findings.
