Research map / Architectures and efficiency / Attention and transformers
Transformers Learn Faster with Semantic Focus
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 empirically and theoretically (via softmax input-stability and Lipschitz constants) that input-dependent top-k sparse attention speeds up convergence and improves generalization of transformers on ListOps and NNCH tasks, while input-agnostic sparse attention gives no consistent benefit.
Problem
Sparse attention in transformers is usually viewed only as a speed-up of full attention; it is unclear whether sparsity can improve learning convergence and in-distribution generalization compared with full attention.
Approach
Link the Lipschitz constant of the training loss (and hence SGD convergence and algorithmic-stability generalization) to the input-stability of softmax, and show that input-dependent heavy-hitter (top-k) attention reduces the softmax semantic dispersion and increases separation, improving these constants, whereas input-agnostic regular sparsity does not.
Key results
- Top-k attention (k=5) reaches 95% training accuracy in 1.37x fewer epochs than full attention on ListOps (121 vs 167) and up to 8.83x faster on Even Pairs (6 vs 53 epochs) in the main text.
- On ListOps, final held-out accuracy of top-k(5) is 31.1±0.7% versus 28.9±1.4% for full attention, while input-agnostic banded and block-local attention reach at most ~29-31% and often fail to reach 95% training accuracy.
- Empirically, the estimated Lipschitz constant of top-k attention is smaller than full attention at points distant from the optimum, consistent with the theory.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace full attention with top-k heavy-hitter attention | Transformer · ListOps | iterations to 95% training accuracy: 167\sstd{18} -> 122\sstd{ 8} | Helped |
| Replace full attention with top-k heavy-hitter attention | Transformer · MissingDups | iterations to 95% training accuracy: 96\sstd{21} -> 48\sstd{10} | Helped |
| Replace full attention with top-k heavy-hitter attention | Transformer · EvenPairs | iterations to 95% training accuracy: 53\sstd{15} -> 6\sstd{3} | Helped |
| Replace full attention with top-k heavy-hitter attention | Transformer · ListOps, best holdout accuracy | accuracy: 35.09\sstd{0.60} -> 36.25\sstd{0.29} | Helped |
| Replace full attention with top-k heavy-hitter attention | Transformer · ListOps, final holdout accuracy | accuracy: 28.92\sstd{1.40} -> 31.33\sstd{0.85} | Helped |
| Replace full attention with top-k heavy-hitter attention | Transformer · MissingDups and EvenPairs | best and final holdout accuracy: 100.00\sstd{0.00} -> 100.00\sstd{0.00}; 99.98\sstd{0.02} -> 99.95\sstd{0.06}; 100.00\sstd{ 0.00} -> 100.00\sstd{ 0.00} | No clear effect |
| Use input-agnostic banded or block-local attention instead of full attention (ablation) | Transformer · ListOps, MissingDups, and EvenPairs | training convergence and holdout accuracy | Hurt |
| Add global tokens to input-agnostic banded or block-local attention (ablation) | Transformer · NNCH tasks | training convergence | Helped |
| Vary the top-k attention mask size (ablation) | Transformer · ListOps, MissingDups, and EvenPairs | holdout accuracy | No clear effect |
| Replace ReLU with GELU or Mish in the MLP (ablation) | Full-attention and top-k Transformer · ListOps, EvenPairs, and MissingDups | training convergence and generalization | No clear effect |
| Vary the number of transformer blocks (ablation) | Full-attention and top-k Transformer · ListOps | training convergence | No clear effect |
| Increase the number of attention heads (ablation) | Full-attention and top-k Transformer · ListOps | training convergence | Helped |
Surprising findings
- Input-agnostic sparse attention (banded, block-local) converges slower than full attention and often fails to reach 95% training accuracy, despite the same parameter count.
- Full attention generalization on ListOps degrades from 35.1% best to 28.9% final, whereas top-k attention does not overfit as much.
Limitations
- Experiments are limited to small controlled synthetic benchmarks trained from scratch on encoder-only models, not large-scale LLMs.
- Theoretical bounds are worst-case upper bounds with constants that may differ from practical values.
- Only a few sparse attention variants (banded, block-local, top-k) are evaluated; in-distribution generalization only, no length generalization.
Experimental setup
Datasets: ListOps (LRA) NNCH tasks (Parity, Even Pairs, Missing Duplicates, Cycle Navigation, Stack Manipulation, Modular Arithmetic with Bracke
Benchmarks: ListOps Parity Even Pairs Missing Duplicates Stack Manipulation Modular Arithmetic Solve Equation Cycle Navigation
Metrics: training cross-entropy training accuracy held-out accuracy iterations to 95% training accuracy estimated Lipschitz constant
Baselines: full softmax attention banded attention block-local attention global tokens Adam
Builds on: top-k attention algorithmic stability generalization bounds Neural Networks and Chomsky Hierarchy benchmark Long-range Arena
Closest papers by meaning
- Rethinking Attention: Exploring Shallow Feed-Forward Neural Networks as an Alternative to Attention Layers in Transformers (Student Abstract) (AAAI 2024 · similarity 0.90)
- Theory, Analysis, and Best Practices for Sigmoid Self-Attention (ICLR 2025 · similarity 0.90)
- On the Role of Attention Masks and LayerNorm in Transformers (NeurIPS 2024 · similarity 0.90)
- Sparse Attention as Compact Kernel Regression (NeurIPS 2026 · similarity 0.90)
- Learning Linear Attention in Polynomial Time (NeurIPS 2025 · similarity 0.89)
- Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape (ICML 2024 · similarity 0.89)
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.