Research map / Architectures and efficiency / Attention and transformers
The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry
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 Hedgehog, a learnable linear attention whose per-head single-layer exponential MLP feature maps are trained via cross-entropy distillation to mimic softmax attention weights, recovering most softmax quality across training-from-scratch, finetuned-conversion, and pretrained-conversion settings.
Problem
Linear attentions usually underperform softmax attention in quality; the paper asks which properties (low-entropy spikiness and dot-product monotonicity) are missing and how to get them while keeping linear complexity.
Approach
Replace fixed linear-attention feature maps with trainable single-layer MLPs followed by element-wise exp (with a negation mapping), trained by cross-entropy against softmax attention weights, so the linear attention mimics softmax spikiness and monotonicity at O(nd^2) cost.
Key results
- On WikiText-103 training from scratch with 125M models, Hedgehog reaches 20.8 perplexity vs 18.6 for softmax Transformer, closing 68.6% of the gap to linear baselines.
- Finetuned-conversion of BERT on GLUE: Hedgehog recovers 99.3% of BERT-FT performance, versus 88.9% for T2R.
- Converting pretrained GPT-2 125M gives 16.7 perplexity on WikiText-103, best among subquadratic 125M models; on Llama-2 7B with LoRA, Hedgehog gets 47.4 ROUGE-1 vs 51.1 for softmax, while T2R gets 2.8.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace standard attention with Hedgehog linear attention | 125M decoder-only Transformer · WikiText-103 | perplexity: 25.6 -> 20.8 | Helped |
| replace standard attention with Hedgehog linear attention | 125M GPT-2 · WikiText-103 | perplexity: 19.4 -> 16.7 | Helped |
| convert Llama-2 7B to Hedgehog attention and finetune with LoRA | Llama-2 7B · SAMSum | ROUGE-1: 51.1 -> 47.4; 2.8 -> 47.4 | Mixed |
| replace the Hedgehog feature map with the T2R feature map while retaining attention distillation (ablation) | BERT-base-uncased conversion · GLUE | performance recovery (%): 99.3 -> 93.5 | Hurt |
| replace the Hedgehog feature map with the T2R feature map while retaining attention distillation (ablation) | ViT-B/16 conversion · ImageNet-1K | top-1 accuracy (%): 79.5 -> 77.0 | Hurt |
| use the T2R feature map with attention distillation instead of Hedgehog's exponential MLP (ablation) | BERT-base-uncased · BERT-FT CoLA | Matthew's correlation: 59.2 -> 56.9 | Hurt |
| remove training from the Hedgehog feature maps (ablation) | BERT-base-uncased attention matching · GLUE attention weights | KL divergence: 0.345 -> 0.851 | Hurt |
| replace the Hedgehog feature map with the T2R feature map (ablation) | BERT-base-uncased attention matching · GLUE attention weights | KL divergence: 0.345 -> 0.613 | Hurt |
| use a degree-2 Taylor exponential feature map instead of Hedgehog (ablation) | BERT-FT CoLA · BERT-FT CoLA | Matthew's correlation: 59.2 -> 58.4 | Hurt |
| use a degree-2 Taylor exponential feature map instead of Hedgehog (ablation) | Transformer trained from scratch · Associative Recall | accuracy: 100.0 -> 100.0 | No clear effect |
| increase the temperature of the exponential feature map from t = 1 to t = 2 (ablation) | linear attention model trained from scratch · Associative Recall | accuracy | Helped |
| use the exponential map with t = 2 rather than prior linear attention maps (ablation) | linear attention model trained from scratch · Associative Recall | accuracy: 17.0 -> 100.0 | Helped |
Surprising findings
- The simple exp(t=2) feature map solves associative recall while exp(t=1) fails, showing spikiness alone matters for recall.
- Naive T2R feature map collapses on Llama-2 conversion (2.8 ROUGE-1) while Hedgehog stays near softmax.
Limitations
- Llama-2 conversion evaluated only on SAMSum summarization with LoRA and a single GPU, limited scale of evaluation.
- Taylor polynomial feature maps are inefficient (O(nd^3)), motivating the learned approach; learned maps require a distillation stage before finetuning.
- LRA results lag non-Transformer state-space models, which the paper does not compare against directly.
Experimental setup
Datasets: CoLA GLUE WikiText-103 Long Range Arena ImageNet-1K SAMSum associative recall (synthetic)
Benchmarks: Long Range Arena WikiText-103 GLUE ImageNet-1K SAMSum associative recall
Metrics: accuracy perplexity Matthew's correlation ROUGE KL divergence wall-clock time memory
Baselines: 1 + ELU ReLU Performer cosFormer T2R T2R-HH H3 Hyena AFT Reformer Linformer BigBird
Builds on: Transformer-to-RNN (T2R) linear attention kernel feature maps FlashAttention LoRA
Closest papers by meaning
- Theory, Analysis, and Best Practices for Sigmoid Self-Attention (ICLR 2025 · similarity 0.89)
- Degrees of Freedom for Linear Attention: Distilling Softmax Attention with Optimal Feature Efficiency (NeurIPS 2025 · similarity 0.89)
- From Attention to Activation: Unraveling the Enigmas of Large Language Models (ICLR 2025 · similarity 0.88)
- STILL: Selecting Tokens for Intra-Layer Hybrid Attention to Linearize LLMs (NeurIPS 2026 · similarity 0.88)
- SEA: Sparse Linear Attention with Estimated Attention Mask (ICLR 2024 · similarity 0.88)
- Gated Linear Attention Transformers with Hardware-Efficient Training (ICML 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.