Research map / Architectures and efficiency / Attention and transformers
Selective Attention Improves Transformer
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.
Proposes Selective Attention, a parameter-free modification that lets tokens mask earlier tokens by subtracting accumulated masking from attention logits, improving language modeling and enabling KV-cache pruning.
Problem
Standard transformer attention keeps the entire context and cannot let later tokens suppress unneeded earlier tokens, which degrades quality and inflates inference memory and compute.
Approach
Reuse an existing attention head's scores as a selection matrix, constrain it (ReLU, no BOS masking, no self-masking), cumulatively sum it over future tokens, and subtract it from attention logits; masked-out tokens can then be evicted from the KV cache under per-layer budgets.
Key results
- Selective attention matches the perplexity of standard transformers with ~2X more heads and parameters in the attention module.
- With the L_mem loss, selective attention reduces attention memory by 16X, 25X, and 47X at context sizes 512, 1,024, and 2,048 for d=12 models while matching baseline perplexity.
- Selective attention improves accuracy over the baseline in nearly all reported cases on ARC, HellaSwag, PiQA, CommonSenseQA, and OpenBookQA across model sizes d=16 to d=28.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add selective attention | Transformer d=28 · ARC-Easy, ARC-Challenge, HellaSwag, CommonSenseQA, OpenBookQA, and PiQA | accuracy: 45.64% -> 46.95%; 26.76% -> 27.37%; 53.50% -> 53.76%; 27.86% -> 28.51%; 37.33% -> 37.54%; 72.32% -> 72.85% | Helped |
| Add selective attention | Transformer d=12 · C4 validation | log-perplexity: 2.6815 -> 2.6372 | Helped |
| Add selective attention to the T5 decoder | T5-small, T5-base, and T5-large · T5 span-corruption validation | loss per non-padding token: 1.962 -> 1.952; 1.693 -> 1.691; 1.528 -> 1.522 | Helped |
| Replace shared attention-head selection with a separate bilinear form (ablation) | Selective attention on C4 · C4 validation, d=8 and d=12 | log-perplexity: 2.90 vs 2.91; 2.63 vs 2.63 | Mixed |
| Allow tokens to affect their own attention operation by removing the shift (ablation) | Selective attention · C4 validation, d=12 | log-perplexity: 2.831 -> 2.832 | Hurt |
| Allow selective attention to mask the <BOS> token (ablation) | Selective attention · C4 validation | log-perplexity: 2.6409 vs 2.6373; 2.5486 vs 2.5483; 2.4750 vs 2.4741; 2.4153 vs 2.4156; 2.3674 vs 2.3673; 2.2909 vs 2.2865 | Mixed |
| Allow tokens to mask themselves (ablation) | Selective attention · C4 validation, d=12, d=14, and d=18 | log-perplexity: 2.7348 -> 2.7251; 2.6510 -> 2.6423; 2.5261 -> 2.5209 | Hurt |
| Remove the ReLU constraint so selection can be negative (ablation) | Selective attention · C4 training | training convergence | Hurt |
| Use local attention instead of dense attention | Transformer d=12 · C4 validation | log-perplexity: 2.6815 -> 2.7860 (All-local 32); 2.6815 -> 2.7046 (Local-global 32) | Hurt |
| Add selective attention | Transformer d=3 · Variable Assignment, out-of-distribution | accuracy: 70% -> 100% | Helped |
| Add selective attention | Transformer d=3 · Copy and Parity* | accuracy | No clear effect |
Surprising findings
- Removing the ReLU constraint so that the selection matrix can be negative causes training not to converge.
- Adding a separate bilinear selection form gives the same or only slightly better results despite extra parameters and compute.
Limitations
- Experiments are limited to models up to about 1B parameters and to decoder-only setups (plus a small T5 test).
- Not tested with multi-query or grouped-query attention, nor during fine-tuning of existing models.
- Context pruning is performed without further training after budget reduction, leaving possible gains unexplored.
Experimental setup
Datasets: C4 WikiText lm1b ARC HellaSwag PiQA CommonSenseQA OpenBookQA
Benchmarks: C4 validation perplexity ARC HellaSwag PiQA CommonSenseQA OpenBookQA Variable Assignment Copy Parity*
Metrics: validation log-perplexity accuracy memory requirement of KV cache
Baselines: standard attention transformer H2O TOVA sparse attention local attention Window+4
Builds on: Transformer T5 H2O TOVA
Closest papers by meaning
- Transformers Learn Faster with Semantic Focus (NeurIPS 2025 · similarity 0.88)
- From Attention to Activation: Unraveling the Enigmas of Large Language Models (ICLR 2025 · similarity 0.87)
- SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking (NeurIPS 2026 · similarity 0.87)
- Theory, Analysis, and Best Practices for Sigmoid Self-Attention (ICLR 2025 · similarity 0.87)
- Unveiling and Harnessing Hidden Attention Sinks: Enhancing Large Language Models without Training through Attention Calibration (ICML 2024 · similarity 0.87)
- Heads up! Large Language Models Can Perform Tasks Without Your Instruction via Selective Attention Head Masking (ICML 2025 · similarity 0.87)
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.