Research map / Architectures and efficiency / Attention and transformers
KV Shifting Attention Enhances Language Modeling
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 KV shifting attention, which mixes each token's key and value with its neighbor's via learnable shifts, so a single attention layer can implement induction heads, and reports gains in toy and pretraining (2.9B/19B) experiments.
Problem
Induction heads normally require at least two attention layers and sufficient width, which may limit in-context learning and language modeling efficiency in decoder-only transformers.
Approach
Decouple keys and values by shifting them with the previous token using four learnable per-head coefficients (alpha1, alpha2, beta1, beta2), so a one-layer attention can represent induction heads; the extra cost is O(ND) and only 4 parameters per head.
Key results
- On the 19B model at 200B tokens, KV shifting reaches 38.83 average vs 38.06 for vanilla across the benchmark suite.
- On the 2.9B model at 500B tokens, KV shifting averages 38.57 vs 36.45 for vanilla; MMLU 5-shot 40.88 vs 37.26.
- On iGSM (Tr12-Te15), KV shifting achieves 0.8909 accuracy vs 0.8154 for vanilla.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use KV shifting attention instead of vanilla attention | 2.9B language model · Benchmarks at 500B tokens | average score: 36.45 -> 38.57 | Helped |
| Use KV shifting attention instead of vanilla attention | 19B language model · Benchmarks at 200B tokens | average score: 38.06 -> 38.83 | Helped |
| Use KV shifting attention instead of vanilla attention | 19B language model · Winogrande, 200B tokens | accuracy: 49.01 -> 48.38 | Hurt |
| Use KV shifting attention instead of vanilla attention | 19B language model · CMMLU, 200B tokens | accuracy: 42.92 -> 42.40 | Hurt |
| Use KV shifting attention instead of vanilla attention | language model · iGSM Tr12-Te15 | accuracy: 0.8154 -> 0.8909 | Helped |
| Use KV shifting attention instead of vanilla attention | language model · iGSM Tr21-Te24 | accuracy: 0.8711 -> 0.9062 | Helped |
| Shift keys and values rather than use vanilla attention | toy induction task · Induction accuracy | accuracy | Helped |
| Use KV shifting rather than vanilla attention | toy 3-gram task · 3-gram learning | accuracy | No clear effect |
| Remove key shifting (ablation) | KV shifting attention · Language modeling | training loss | Hurt |
| Remove value shifting (ablation) | KV shifting attention · Language modeling | training loss | Hurt |
| Extend the shift window from [i-1,i] to [i-2,i] or [i-3,i] (ablation) | KV shifting attention · Language modeling | training loss | No clear effect |
| Constrain the KV shifting parameters with a gate or to the range 0 to 1 (ablation) | KV shifting attention · 1.5B model, 10B tokens | training loss | Hurt |
Surprising findings
- Vanilla diverges at learning rate 1e-2 while KV shifting does not, suggesting a flatter optimization landscape.
- Learned alpha/beta coefficients become negative during training, and constraining them to [0,1] does not improve performance.
Limitations
- Most experiments run once due to compute limits, so variance is not well characterized.
- Pretraining data is private and not released, limiting reproducibility.
- Theory is for simplified conditions (no residuals, MLP, normalization, position embeddings) and large-d limits.
Experimental setup
Datasets: private pretraining corpus similar to FineWeb-edu WebText (validation) iGSM synthetic induction data
Benchmarks: lambada winogrande hellaswag ARC-E ARC-C CMMLU MMLU MATH iGSM
Metrics: accuracy training loss validation loss error rate on hop tasks
Baselines: Vanilla transformer (2-layer / standard attention)
Builds on: induction heads virtual attention heads Alibi Llama grouped query attention
Closest papers by meaning
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.91)
- Which Attention Heads Matter for In-Context Learning? (ICML 2025 · similarity 0.88)
- The Closeness of In-Context Learning and Weight Shifting for Softmax Regression (NeurIPS 2024 · similarity 0.87)
- KVSmooth: Mitigating Hallucination in Multi-modal Large Language Models through Key-Value Smoothing (CVPR 2026 · similarity 0.87)
- Selective Attention Improves Transformer (ICLR 2025 · similarity 0.87)
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · 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.