Research map / Architectures and efficiency / Attention and transformers
Selective Attention: Enhancing Transformer through Principled Context Control
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 Self-Attention (SSA), which applies learnable token- and position-dependent temperature scaling to query and value embeddings, improving language modeling with <0.5% extra parameters via weight sharing.
Problem
Standard softmax self-attention treats all queries identically, which limits its ability to control contextual sparsity (spikiness) of attention per query and per position, causing attention dilution in long contexts.
Approach
Augment softmax attention with data-dependent inverse-temperature scaling on queries (token-aware via tanh of a small MLP reusing attention weights, plus a position-aware log(n) term) and on values, decoupling semantic similarity from contextual sparsity; theory shows norm and expressivity benefits.
Key results
- SSA on Llama3-8B fine-tuning lowers Wikitext perplexity from 12.416 to 10.982 and raises average accuracy from 0.632 to 0.639.
- On passkey retrieval, fine-tuned Pythia-160m improves from 56.89% to 74.41% with SSA.
- SSA reaches comparable performance with 1.45x fewer pretraining steps when fine-tuning Llama on Wikitext.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add SSA query and value temperature scaling | Pythia-160m · Wikitext (fine-tune) | perplexity: 26.681 -> 26.514 | Helped |
| add SSA query and value temperature scaling | Llama3-8b · Wikitext (fine-tune) | perplexity: 12.416 -> 10.982 | Helped |
| add SSA attention | Pythia-160m · Passkey retrieval | retrieval accuracy: 56.89 -> 74.41 | Helped |
| add SSA attention | Llama · Passkey retrieval | retrieval accuracy: 77.62 -> 89.53 | Helped |
| apply value selectivity to suppress noisy tokens | one-layer attention denoising model · Synthetic denoising task | MSE risk: 1.390 -> 0.071 | Helped |
| use query-selective attention | one-layer attention model · Synthetic next-token prediction | cross-entropy loss: 0.0126 -> 0.009 | Helped |
| apply temperature scaling to queries only (ablation) | Pythia · Wikitext (fine-tune) | perplexity: 28.781 -> 27.416 | Helped |
| apply temperature scaling to keys only (ablation) | Pythia · Wikitext (fine-tune) | perplexity: 28.781 -> 28.715 | No clear effect |
| apply temperature scaling to values only (ablation) | Pythia · Wikitext (fine-tune) | perplexity: 28.781 -> 27.980 | Helped |
| apply temperature scaling to queries and values (ablation) | Pythia · Wikitext (fine-tune) | perplexity: 28.781 -> 26.514 | Helped |
| also apply temperature scaling to keys (ablation) | Pythia with query and value temperatures · Wikitext (fine-tune) | perplexity: 26.514 -> 26.603 | Hurt |
| use position-aware query temperature only instead of both position- and token-aware query scaling (ablation) | Pythia · Wikitext (fine-tune) | perplexity: 27.416 -> 27.995 | Hurt |
Surprising findings
- Scaling temperature on keys degrades or does not improve performance versus vanilla attention, unlike query and value scaling.
- Position-aware temperature gives much smaller benefit than token-aware temperature in the authors' setting.
Limitations
- Focuses on softmax attention, which retains quadratic computation cost; linear attention extension not explored.
- Feature-based temperature variant is sensitive to feature selection and varies across datasets.
- Gains are modest in several benchmarks (often small absolute differences).
Experimental setup
Datasets: SlimPajama Wikitext WebText The Pile
Benchmarks: Wikitext Lambada Piqa HellaSwag Winogrande ARC-E ARC-C passkey retrieval
Metrics: perplexity accuracy acc_norm passkey retrieval accuracy
Baselines: Yarn constant temperature scaling frequency-based temperature vanilla attention
Builds on: Yarn FlashAttention Mamba
Closest papers by meaning
- SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking (NeurIPS 2026 · similarity 0.90)
- A Training-Free Sub-quadratic Cost Transformer Model Serving Framework with Hierarchically Pruned Attention (ICLR 2025 · similarity 0.88)
- Theory, Analysis, and Best Practices for Sigmoid Self-Attention (ICLR 2025 · similarity 0.88)
- Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps (NeurIPS 2026 · similarity 0.88)
- Scaling Limits of Long-Context Transformers (NeurIPS 2026 · similarity 0.88)
- Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free (NeurIPS 2025 · similarity 0.88)
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.