Research map / Architectures and efficiency / Attention and transformers
Scale-invariant attention
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 scale-invariant attention, a position-dependent affine transform of attention logits derived under a Gaussian assumption, combined with p-RoPE, to improve zero-shot length generalisation of GPT-2-style models.
Problem
Attention mechanisms trained on short contexts degrade when run on much longer contexts, as attention spreads out and local context attention vanishes.
Approach
Defines scale-invariant total attention and scale-invariant attention sparsity as desiderata, and shows that under Gaussian logits a position-dependent scaling a_t and bias m_t (a_t^2 = 2[log(t/tau+1) - log alpha + beta/alpha], m_t = -a_t^2 + beta/alpha) satisfies them, yielding a logit transform that keeps local attention while sparsifying global attention.
Key results
- Scale-invariant p-RoPE achieves 3.247 validation loss at 64k when trained at 4k (162M model), versus 5.260 for RoPE and 3.270 for ALiBi.
- On needle-in-a-haystack after fine-tuning, scale-invariant p-RoPE reaches 0.969 accuracy at 64k, versus 0.939 for LogN+p-RoPE and 0.000 for RoPE.
- On Llama 2 7B continual pretraining, scale-invariant p-RoPE gives 2.252 loss at 64k, far below RoPE+NTK (7.523) and LogN+p-RoPE (6.224), though it is worse at 4k (2.163 vs 1.968 for RoPE).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add scale-invariant attention to p-RoPE | 162M GPT-2-style model · FineWeb validation, Train @4k / Val @64k | validation loss: 3.317 ± 0.005 -> 3.247 ± 0.001 | Helped |
| Add scale-invariant attention to p-RoPE | 162M GPT-2-style model · FineWeb validation, Train @4k / Val @16k | validation loss: 3.262 ± 0.002 -> 3.235 ± 0.001 | Helped |
| Add scale-invariant attention to p-RoPE | 162M GPT-2-style model · Needle-in-a-haystack, Val @64k | accuracy: 0.939 ± 0.009 -> 0.969 ± 0.005 | Helped |
| Replace p-RoPE with RoPE while retaining scale-invariant attention (ablation) | Scale-invariant p-RoPE · FineWeb validation, Train @4k / Val @4k | validation loss | No clear effect |
| Replace p-RoPE with RoPE while retaining scale-invariant attention (ablation) | Scale-invariant p-RoPE · FineWeb validation, Train @4k / Val @64k | validation loss | Hurt |
| Replace p-RoPE with NoPE while retaining scale-invariant attention (ablation) | Scale-invariant p-RoPE · FineWeb validation, Train @4k / Val @4k | validation loss | Hurt |
| Replace p-RoPE with NoPE while retaining scale-invariant attention (ablation) | Scale-invariant p-RoPE · FineWeb validation, Train @4k / Val @64k | validation loss | Helped |
| Vary the lengthscale τ (ablation) | Scale-invariant attention · FineWeb validation at training context length (4k) | validation loss | No clear effect |
| Set the lengthscale τ to 10 (ablation) | Scale-invariant attention · FineWeb validation at out-of-distribution context lengths (16k and 64k) | validation loss | Helped |
| Add scale-invariant attention to p-RoPE | Llama-2 7B · Validation at 64k | validation loss: 6.224 -> 2.252 | Helped |
Surprising findings
- Scale-invariant RoPE (without p-RoPE) did not generalise well to long contexts, suggesting low-frequency RoPE components interfere with the transform.
- Scale-invariant attention trades a slightly worse in-distribution loss at 4k (e.g. 7B model 2.163 vs 1.968 for RoPE) for much better long-context behaviour.
Limitations
- Pretraining from scratch only at 162M and 304M parameters; 7B results are from continual pretraining on ~50M tokens.
- Only dense attention is studied; extension to sparse or other settings left open.
- Scale-invariance proven only in expectation under a Gaussian logit assumption; sparsity is weak (sqrt(log t) empirically), not strong.
Experimental setup
Datasets: FineWeb C4
Benchmarks: FineWeb validation loss (4k/16k/64k) needle-in-a-haystack (4k/16k/64k)
Metrics: validation loss retrieval accuracy
Baselines: RoPE p-RoPE NoPE RoPE+NTK YaRN LogN+RoPE LogN+p-RoPE ALiBi Infini-attention
Builds on: RoPE LogN scaling / Scalable-Softmax p-RoPE FlexAttention ALiBi
Closest papers by meaning
- Unveiling and Harnessing Hidden Attention Sinks: Enhancing Large Language Models without Training through Attention Calibration (ICML 2024 · similarity 0.89)
- SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking (NeurIPS 2026 · similarity 0.88)
- LLM Maybe LongLM: SelfExtend LLM Context Window Without Tuning (ICML 2024 · similarity 0.88)
- Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps (NeurIPS 2026 · similarity 0.88)
- Selective Attention: Enhancing Transformer through Principled Context Control (NeurIPS 2024 · similarity 0.87)
- Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study (ICLR 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.