Research map / Architectures and efficiency / Attention and transformers
Degrees of Freedom for Linear Attention: Distilling Softmax Attention with Optimal Feature Efficiency
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.
Selects per-layer feature dimensions for linear attention distilled from softmax attention using statistical degrees of freedom of the attention kernel, and trains PRF features layerwise.
Problem
Distilling pre-trained softmax attention into linear attention requires choosing a feature dimension that governs approximation quality, but existing methods use one fixed dimension across all layers.
Approach
Uses the statistical degrees of freedom tr(Σ(Σ+λI)^-1) of the attention kernel, estimated from queries and keys via a Gram matrix, to allocate a layer-specific feature dimension under a fixed average compute budget, and learns positive random features per layer with L2 or softmax losses instead of end-to-end training.
Key results
- On GPT-2, DoF with direct or softmax loss improves the seven-task average over Fix (about 0.419 vs 0.414 for direct and 0.396 for softmax), and both beat Performer (0.395) and DiJiang (0.341).
- On Pythia-1B with softmax loss, DoF reaches an average of 0.4118 vs 0.4014 for Fix and 0.4279 for the original model.
- Selected feature dimensions are larger in early and middle layers and smaller in later layers, and most heads have much smaller DoF than their layer maximum.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use DoF-selected dimensions instead of fixed dimensions with direct loss (ablation) | GPT-2 distilled linear attention · Downstream tasks (Average) | accuracy: 0.4141 -> 0.4182 | Helped |
| Use DoF-selected dimensions instead of fixed dimensions with softmax loss (ablation) | GPT-2 distilled linear attention · Downstream tasks (Average) | accuracy: 0.3958 -> 0.4190 | Helped |
| Use DoF-selected dimensions instead of fixed dimensions with L2 loss (ablation) | GPT-2 distilled linear attention · Downstream tasks (Average) | accuracy: 0.3980 -> 0.3965 | No clear effect |
| Cap DoF-selected feature dimensions at the head size, using direct loss (ablation) | GPT-2 distilled linear attention · Downstream tasks (Average) | accuracy: 0.4141 -> 0.4223 | Helped |
| Cap DoF-selected feature dimensions at the head size, using softmax loss (ablation) | GPT-2 distilled linear attention · Downstream tasks (Average) | accuracy: 0.3958 -> 0.4182 | Helped |
| Cap DoF-selected feature dimensions at the head size, using L2 loss (ablation) | GPT-2 distilled linear attention · Downstream tasks (Average) | accuracy: 0.3980 -> 0.4079 | Helped |
| Use layerwise softmax loss instead of direct end-to-end training (ablation) | GPT-2 distilled linear attention with DoF dimensions · Downstream tasks (Average) | accuracy: 0.4182 -> 0.4190 | No clear effect |
| Use layerwise L2 loss instead of direct end-to-end training (ablation) | GPT-2 distilled linear attention with DoF dimensions · Downstream tasks (Average) | accuracy: 0.4182 -> 0.3965 | Hurt |
| Use DoF-selected dimensions instead of fixed dimensions (ablation) | Pythia-1B distilled linear attention with softmax loss · Downstream tasks (Average) | accuracy: 0.4014 -> 0.4118 | Helped |
| Apply DoF + Clip instead of Fix (ablation) | GPT-2 inference · 1,000-token sequences | speed (sec/sequence): 15.589 -> 14.662 | Helped |
| Use DoF dimensions instead of fixed dimensions (ablation) | GPT-2 inference · 1,000-token sequences | speed (sec/sequence): 15.589 -> 15.579 | No clear effect |
| Use layerwise softmax loss instead of direct loss (ablation) | Feature training · Training samples | speed (ms/sample): 109.4 -> 90.4 | Helped |
Surprising findings
- The DoF varies strongly across heads within the same layer, though the method uses a per-layer maximum for implementation reasons.
- Layerwise softmax loss matches or beats end-to-end direct loss on downstream average for DoF selection, despite not optimizing next-token prediction directly.
Limitations
- A single feature dimension is shared across all heads within a layer, although DoF varies across heads.
- Experiments are limited to moderate-scale models (GPT-2 and Pythia-1B) and language tasks.
Experimental setup
Datasets: Wikipedia PiQA logiQA ARC-E ARC-C Winogrande MMLU WSC
Benchmarks: PiQA logiQA ARC-E ARC-C Winogrande MMLU WSC
Metrics: accuracy next-token prediction loss inference time per sequence
Baselines: Performer DiJiang Fix (uniform feature dimension) direct end-to-end loss
Builds on: Performer DiJiang positive random features Bochner's theorem
Closest papers by meaning
- Bridging the Divide: Reconsidering Softmax and Linear Attention (NeurIPS 2024 · similarity 0.90)
- STILL: Selecting Tokens for Intra-Layer Hybrid Attention to Linearize LLMs (NeurIPS 2026 · similarity 0.89)
- ZeroS: Zero‑Sum Linear Attention for Efficient Transformers (NeurIPS 2025 · similarity 0.89)
- In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention (ICML 2025 · similarity 0.89)
- The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry (ICLR 2024 · similarity 0.89)
- MetaLA: Unified Optimal Linear Approximation to Softmax Attention Map (NeurIPS 2024 · 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.