Research map / Architectures and efficiency / Attention and transformers
Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that a one-layer transformer with stochastic positional encoding trained by gradient descent provably learns the q-sparse token selection task with width O(d+q log T), while any FCN approximating it needs first-layer width Omega(Td), and the learned model length-generalizes.
Problem
Whether the expressivity separation between transformers and fully-connected networks on sparse token selection translates into learnability under gradient descent, and whether such trained transformers generalize to longer sequences.
Approach
Uses a stochastic (random near-orthogonal, RIP-based) positional encoding with a query-subset encoding, reparameterizes Q and K into one matrix W, and analyzes GD on the population loss to prove global convergence with near-optimal width; also proves an average-case (population) width lower bound for FCNs.
Key results
- Theory: GD with stochastic PE converges to zero population loss with transformer width Theta(d+q log T), matching the near-optimal width.
- Any FCN with first-layer width at most Td-1 incurs expected squared loss at least (T-q)/(Tq(T-1)), giving an average-case separation.
- Experiments show fixed-PE transformers fail to reach zero OOD loss while stochastic-PE transformers extrapolate to unseen lengths.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use stochastic positional encoding instead of fixed positional encoding | one-layer transformer on sparse token selection · out-of-distribution sequence lengths | validation loss | Helped |
| use fixed positional encoding instead of stochastic positional encoding | one-layer transformer on sparse token selection · in-distribution training | training loss convergence | No clear effect |
Surprising findings
- Fixed positional encoding leaves a strictly nonzero OOD loss even when in-distribution loss converges to zero.
Limitations
- Analysis uses population loss rather than finite-sample ERM/SGD; no sample complexity guarantee.
- Simplified architecture merges Q and K into one matrix W and restricts query/value attention, with q treated as constant.
- Length generalization corollary is informal and only covers the stated sparse-token-selection setting.
Experimental setup
Datasets: synthetic sparse token selection data (Gaussian tokens, uniform q-subsets)
Benchmarks: synthetic sparse token selection (in-distribution and OOD length)
Metrics: population squared loss OOD validation loss cosine similarity to ground-truth weights
Baselines: fully-connected networks fixed absolute positional encoding
Builds on: q-sparse averaging (qSA) stochastic positional encoding restricted isometry property (RIP) encodings
Closest papers by meaning
- Random Sparse Lifts: Construction, Analysis and Convergence of finite sparse networks (ICLR 2024 · similarity 0.85)
- Transformer Learns Optimal Variable Selection in Group-Sparse Classification (ICLR 2025 · similarity 0.85)
- Unraveling the Gradient Descent Dynamics of Transformers (NeurIPS 2024 · similarity 0.84)
- Non-asymptotic Convergence of Training Transformers for Next-token Prediction (NeurIPS 2024 · similarity 0.84)
- Learning Linear Attention in Polynomial Time (NeurIPS 2025 · similarity 0.84)
- Separations in the Representational Capabilities of Transformers and Recurrent Architectures (NeurIPS 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.