PaperlayerView exact revision on arXiv

Source-linked AI summary

The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry

Michael Zhang, Kush Bhatia, Hermann Kumbong, Christopher Ré

arXiv:2402.04347v1cs.LGcs.CL

TL;DR

Existing linear attention mechanisms often fail to match softmax attention in modeling quality, achieving 4–6 worse perplexity on WikiText-103 when trained from scratch. Hedgehog trains single-layer MLP feature maps to mimic softmax attention’s spiky, monotonic weights while maintaining O(nd^2) complexity. Hedgehog is competitive with softmax-based attention across scratch training and model-conversion regimes, recovering >99% of original model performance on average in finetuned-conversion.

  • Problem

    Existing linear attention mechanisms often fail to match softmax attention in modeling quality, achieving 4–6 worse perplexity on WikiText-103 when trained from scratch.

    Source: §1 · p. 2

  • Method

    Hedgehog trains single-layer MLP feature maps to mimic softmax attention’s spiky, monotonic weights while maintaining O(nd^2) complexity.

    Source: §1 · p. 2

  • Results

    Hedgehog is competitive with softmax-based attention across scratch training and model-conversion regimes, recovering >99% of original model performance on average in finetuned-conversion.

    Source: §1 · p. 3, §6 · p. 10

  • Takeaways & Limitations

    Hedgehog supports training linear attention models from scratch and converting existing Transformers into linear attention variants.

    Source: §6 · p. 10

    The degree-2 Taylor approximation is inefficient, with d′ = 1 + d + d^2 and O(nd^3) attention complexity.

    Source: §4.1 · p. 6

Abstract

from arXiv · show

Linear attentions have shown potential for improving Transformer efficiency, reducing attention's quadratic complexity to linear in sequence length. This holds exciting promise for (1) training linear Transformers from scratch, (2) "finetuned-conversion" of task-specific Transformers into linear versions that recover task performance, and (3) "pretrained-conversion" of Transformers such as large language models into linear versions finetunable on downstream tasks. However, linear attentions often underperform standard softmax attention in quality. To close this performance gap, we find prior linear attentions lack key properties of softmax attention tied to good performance: low-entropy (or "spiky") weights and dot-product monotonicity. We further observe surprisingly simple feature maps that retain these properties and match softmax performance, but are inefficient to compute in linear attention. We thus propose Hedgehog, a learnable linear attention that retains the spiky and monotonic properties of softmax attention while maintaining linear complexity. Hedgehog uses simple trainable MLPs to produce attention weights mimicking softmax attention. Experiments show Hedgehog recovers over 99% of standard Transformer quality in train-from-scratch and finetuned-conversion settings, outperforming prior linear attentions up to 6 perplexity points on WikiText-103 with causal GPTs, and up to 8.7 GLUE score points on finetuned bidirectional BERTs. Hedgehog also enables pretrained-conversion. Converting a pretrained GPT-2 into a linear attention variant achieves state-of-the-art 16.7 perplexity on WikiText-103 for 125M subquadratic decoder models. We finally turn a pretrained Llama-2 7B into a viable linear attention Llama. With low-rank adaptation, Hedgehog-Llama2 7B achieves 28.1 higher ROUGE-1 points over the base standard attention model, where prior linear attentions lead to 16.5 point drops.

1 Introduction

Linear attention promises more efficient Transformers across training-from-scratch, finetuned-conversion, and pretrained-conversion settings, but often falls short of softmax attention in quality. Hedgehog learns feature maps to mimic softmax attention’s spikiness and monotonicity while retaining linear-attention complexity, and the paper reports strong results across all three settings.

1 Introduction

  • Existing linear attentions often underperform softmax: scratch-trained models show 4–6 worse perplexity on WikiText-103, while converted models may need quadratic modules to close the gap.The passage also notes theoretical results suggesting high-quality truly subquadratic softmax approximations may be impossible at large sequence lengths.Source: §1 · p. 2
  • The paper identifies low-entropy spikiness and dot-product monotonicity as softmax properties that prior linear attentions lack and hypothesizes that their absence relates to the quality gap.Spiky weights selectively emphasize relevant tokens; monotonic weights rise as query-key dot products increase.Source: §1 · p. 2
  • Taylor-polynomial feature maps can recover softmax-like spikiness and monotonicity when query and key dot products are bounded, but degree p ≥2 makes them inefficient to compute.Their time and space cost is O(nd^(p+1)).Source: §1 · p. 2
  • Hedgehog replaces fixed kernel functions with single-layer MLP feature maps trained against softmax attention weights, maintaining O(nd^2) complexity while producing low-entropy, monotonic weights.The mappings are from R^d to R^d and are trained using softmax attention weights as soft-label cross-entropy targets.Source: §1 · p. 2
  • Hedgehog recovers >99% of original model performance on average in finetuned-conversion, and closes the linear-attention gap by 68.6% on WikiText-103 training-from-scratch.The conversion result covers bidirectional 110M BERT-base models finetuned on GLUE and causal 125M GPT models finetuned on WikiText-103.Source: §1 · p. 3
  • A 125M Hedgehog-GPT-2 reaches 16.7 ppl on WikiText-103, while Hedgehog-Llama2 7B gains up to 28.1 ROUGE-1 points over the base standard attention model.Prior linear attentions lead to 16.5 ROUGE-1 point drops for the Llama comparison.Source: §1 · p. 3

2 Preliminaries and Related Work

Standard attention computes query-key similarities across the sequence, but its quadratic time and memory costs make long sequences inefficient. Kernel feature maps reformulate attention to scale linearly in sequence length; prior maps include positive activations and randomized or low-rank kernel approximations.

2 Preliminaries and Related Work

  • Computing standard attention for all outputs requires O(n^2d) time and memory, motivating alternatives that improve efficiency without sacrificing expressivity.For causal attention, similarities are computed only for keys at positions j ≤ i.Source: §2 · p. 3
  • Replacing the exponential kernel with K(x, x′) = ϕ(x)^Tϕ(x′) permits linear time and space over sequence length n.The feature map ϕ maps vectors from R^d to R^d′.Source: §2 · p. 3
  • Prior feature maps include 1+ELU or ReLU for positive weights, randomized features for softmax or Gaussian kernel approximations, and low-rank approximations.The figure captions compare softmax, prior linear attentions, and Hedgehog on weight spikiness and monotonicity; they report Hedgehog recovering both properties.Source: §2 · p. 3, Figure 2 · p. 4, Figure 3 · p. 4

3 Improving Linear Attention via Spiky and Monotonic Weights

The section argues that strong linear attention requires both selective, low-entropy weights and monotonicity in query-key dot products. Its experiments connect these properties to associative-recall performance and to recovering a finetuned BERT’s original score after conversion.

3.1 Properties for Expressive Attention Maps

  • The proposed properties are low-entropy spikiness, which selectively upweights relevant tokens, and monotonicity, which makes weights rise with query-key dot products.The section hypothesizes that lacking spikiness corresponds to worse train-from-scratch performance, while lacking both properties corresponds to failure in finetuned conversion.Source: §3 · p. 4, §3.1 · pp. 4–5, §3.2 · p. 5
  • Existing linear attention maps can produce much higher-entropy, more uniform weights than softmax, reducing modeling capacity and performance quality; their non-monotonic query-key behavior may cause conflicting gradients when swapped into trained models.Source: §3.1 · p. 5

3.2 Explaining the Linear Attention Performance Gap

  • Softmax solves associative recall perfectly, while prior linear attentions achieve under 20% accuracy alongside much higher attention-weight entropy.Figure 4 reports a strong correspondence between lower entropy and higher associative-recall accuracy.Source: §3.2 · p. 5, Figure 4 · p. 5
  • Prior linear attentions do not fully recover the finetuned BERT’s CoLA Matthew’s correlation of 58.8, including the spiky ϕ2 map.The passage links this shortfall to the methods’ failure to capture monotonicity over the trained model’s query-key dot products.Source: §3.2 · p. 5, Table 1 · p. 5

4 Hedgehog: Expressive Linear Attention via Softmax Mimicry

The section motivates Hedgehog by showing that Taylor feature maps can recover softmax-like behavior but are costly, then introduces trainable feature maps to mimic softmax attention. Hedgehog uses MLPs trained against softmax attention weights while retaining linear-attention scaling.

4.1 Simple Polynomial Approximations to Softmax Attention

  • The second-degree Taylor approximation retains spikiness and monotonicity and near-matches softmax performance when BERT query-key dot products lie in the regime where monotonicity holds.Figure 5 presents the property recovery; the cited passage reports near-matching performance in the tested settings.Source: §4.1 · p. 6, Figure 5 · p. 6
  • The second-degree Taylor map has feature dimension d′ = 1 + d + d^2, raising attention complexity to O(nd^3) and creating an efficiency-effectiveness trade-off.More generally, degree-p polynomial feature maps take O(nd^p) time and space and track the exponential with low error only in bounded regimes.Source: §4.1 · p. 6

4 Hedgehog: Expressive Linear Attention via Softmax Mimicry

  • Hedgehog learns feature maps with single-layer MLPs, trained to mimic softmax attention while retaining prior linear attentions’ O(nd^2) complexity.The approach replaces a fixed functional form with trainable maps intended to recover spiky and monotonic attention properties.Source: §4 · p. 6, §4.2 · p. 7

4.2 Learnable Linear Attentions for Mimicking Softmax

  • Hedgehog applies trainable MLP feature maps after query and key projections, using an element-wise exponential activation to induce spikiness.The maps are trained by minimizing cross-entropy between linear-attention weights and the weights produced by softmax attention.Source: §4.2 · p. 7

5 Experiments

The experiments assess Hedgehog’s expressivity and efficiency, then test it in training-from-scratch, finetuned-conversion, and pretrained-conversion settings. Across these settings, the results show strong softmax-attention matching and competitive task performance.

5 Experiments; 5.1 Benchmarking Hedgehog for Expressivity and Efficiency

  • The experiments evaluate softmax-property recovery, efficiency, weight matching, and transfer to new data and longer contexts before downstream task benchmarks.For train-from-scratch associative recall, Hedgehog is trained end-to-end without the distillation loss; the reported results associate its modeling and complexity with spiky and monotonic properties.Source: §5 · p. 7, §5.1 · pp. 7–8, Table 3 · p. 7

5.1 Benchmarking Hedgehog for Expressivity and Efficiency

  • Hedgehog achieves near 6x faster inference and similar memory to FlashAttention, while Taylor features are slower and use more memory.The benchmark uses one 12-head attention layer with head dimension 64 and sequences up to 32K tokens.Source: Figure 6 · p. 7, §5.1 · p. 8
  • Distilled Hedgehog attention weights match softmax with higher fidelity than prior linear attentions, and distillation is necessary for recovering weights while the spiky MLP also helps.The evaluation visualizes weights in BERT-FT CoLA; the ablations compare distillation with ReLU features and untrained MLPs.Source: §5.1 · p. 8, Figure 7 · p. 8, Figure 8 · p. 8
  • Feature maps distilled on CoLA or WikiText-103 better match softmax weights on new GLUE data, and their KL divergence stays consistent from 256 to 4096 tokens.The reported context lengths reach 8x the default 512-token context.Source: Table 4 · p. 8, Table 5 · p. 8, Figure 9 · p. 8

5.2 Learning Sequence Modeling From Scratch; 5.3 Finetuned Conversion of Quadratic to Linear Transformers

  • On WikiText-103, Hedgehog closes the language-modeling gap by up to 6 PPL; it recovers 99.3% of original finetuned BERT GLUE performance and achieves 99% ViT accuracy.For training from scratch, it achieves the best average accuracy for both LRA and WikiText-103 tasks among linear attentions.Source: §5.2 · p. 9, Table 8 · p. 10, Table 9 · p. 10

5.3 Finetuned Conversion of Quadratic to Linear Transformers

  • For finetuned conversion, the authors distill attention weights before finetuning converted BERT models on their original tasks.The setup also evaluates ViT-B/16 models trained on ImageNet-1K.Source: §5.3 · p. 10, Table 8 · p. 10, Table 9 · p. 10

5.4 Pretrained Conversion for Subquadratic Task Transfer

  • Pretrained conversion is evaluated with a 125M GPT-2 on WikiText-103 and a Llama-2 7B converted before LoRA finetuning on SAMSum.The reported comparison includes T2R; Hedgehog is also applied to Llama-2 to enable linear-attention training with LoRA.Source: §5.4 · p. 10
  • Hedgehog pretrained-converted GPT-2 outperforms T2R and modern subquadratic models on WikiText-103, while finishing 1 PPL behind fully quadratic finetuned GPT-2.The table caption reports Hedgehog ahead of other subquadratic models by 1.8 PPL.Source: §5.4 · p. 10, Table 10 · p. 10

6 Conclusion

The conclusion presents Hedgehog as a learnable linear attention method for mimicking softmax attention. Its implementation uses head- and layer-specific feature maps, with a practical negation mapping that preserves the stated linear-attention complexity.

6 Conclusion

  • The Llama-2 conversion results are identified as ROUGE comparisons, but the supplied caption gives no numerical scores.The passage provides the Table 11 caption only.Source: Table 11 · p. 10
  • Hedgehog combines low-entropy spiky attention maps and monotonicity in query-key dot products to mimic softmax attention across three conversion and training regimes.The conclusion reports competitive performance in training-from-scratch, finetuned-conversion, and pretrained-conversion settings.Source: §6 · p. 10

A.1 Mechanics for Hedgehog feature map

  • A negation mapping adds a fixed factor of 2 to feature dimension while retaining O(nd2) time and space complexity.The additional mapping is intended to better account for negative dimensions that prior feature maps such as ReLU ignore.Source: §A.1 · p. 15

A.2 Hedgehog feature map and model architecture

  • Hedgehog feature maps are applied separately at each head and layer as single linear layers matching the Transformer head dimension.The head-specific MLPs follow each query and key projection, akin to adapters.Source: §A.2 · p. 15

A.3 Hedgehog Distillation and Finetuning Implementation Details

Hedgehog’s implementation trains feature maps end-to-end for models trained from scratch and uses a two-stage procedure for converted models. Conversion first distills softmax attention weights, then finetunes model parameters for the target task.

Training-from-scratch

  • For training from scratch, Hedgehog MLPs are inserted after query and key projections and trained jointly with the model using the task objective.For language modeling, the example objective is next-token-prediction cross-entropy.Source: §A.3 · p. 16

Finetuned / pretrained conversion

  • For finetuned or pretrained conversion, Hedgehog first inserts MLPs and follows a two-stage training process.The first stage freezes original Transformer weights and trains feature maps to match softmax attention weights using soft cross-entropy or KL divergence.Source: §A.3 · p. 16

Attention distillation

  • Attention distillation jointly optimizes Hedgehog feature maps across heads and layers using task data and a summed attention-matching loss.The original Transformer weights remain frozen during this stage.Source: §A.3 · p. 16

Original parameter finetuning

  • After distillation, the implementation unfreezes model weights for task-specific finetuning, while allowing selected layers to remain frozen or LoRA to be used.LoRA is explored with Llama-2 models.Source: §A.3 · p. 16

B.1 Associative recall analysis (Section 3.2)

The appendix gives the associative-recall setup and training details for conversion, scratch training, and pretrained experiments. It also reports an extended LRA comparison in which Hedgehog has the best average accuracy among the compared attention-based methods.

B.1 Associative recall analysis (Section 3.2)

  • Associative recall uses 40 possible tokens and 128-token sequences, with 10,000 training samples and 2,000 newly generated test samples.The task tests recalling pairings that occur three times on average in-context.Source: §B.1 · p. 18, Table 12 · p. 18

B.2 BERT-base finetuned on CoLA conversion (Section 3.2)

  • The BERT-base CoLA conversion experiment replaces attention in a finetuned checkpoint and trains for up to 10 epochs with early stopping.It uses batch size 8, learning rate 1e-5, zero weight decay, and AdamW.Source: §B.2 · p. 19

B.3 Hedgehog training from scratch (Section 5.2)

  • WikiText-103 scratch training uses a 125M GPT-2-style, 12-layer decoder with 12 heads and head dimension 64.The described training uses learning rate 6e-4, weight decay 0.01, and AdamW.Source: §B.3 · p. 19

B.4 Hedgehog finetuned conversion (Section 5.3)

  • For GLUE conversion, BERT attention layers are distilled for up to five epochs before each converted model is trained on its task.The ImageNet-1K ViT setup distills frozen attention for two epochs, then trains all parameters for up to 10 epochs.Source: §B.4 · p. 19

B.5 Hedgehog pretrained conversion (Section 5.4)

  • The Llama-2 7B SAMSum setup distills feature maps for two epochs, then applies LoRA to query, key, value, and output projections.Inputs are concatenated into 1024-token sequences, and generated summaries are evaluated with ROUGE metrics.Source: §B.5 · p. 20

C.1 Extended comparison to attention models on LRA

  • Hedgehog obtains the best average accuracy on LRA among the compared attention-based methods.The extended comparison draws alternatives from the official benchmark leaderboard and recent subquadratic-attention works.Source: §C.1 · p. 20

C.2 Hedgehog feature map generalization to new data

Hedgehog feature maps trained on one dataset often generalize to new datasets, producing attention weights closer to softmax and supporting improved performance after cross-task BERT conversion.

C.2 Hedgehog feature map generalization to new data

  • Hedgehog feature maps almost always generalize better than prior linear attention maps, and cross-task BERT finetuning still improves over prior linear attentions.The section evaluates fidelity and downstream performance when attention maps are trained on CoLA or WikiText-103 and used on other GLUE tasks.Source: §C.2 · pp. 20–21
  • Hedgehog achieves the best average performance across Transformers and subquadratic variants in training-from-scratch LRA experiments.Source: Table 13 · p. 21
  • The setup distills Hedgehog attentions from pretrained BERT-base-uncased on CoLA or 512-token WikiText-103 chunks before transfer to non-CoLA GLUE tasks.Comparisons include T2R-HH, Hedgehog without attention distillation, Performer, and cosFormer.Source: §C.2 · p. 21

C.2.1 Qualitative evidence of Hedgehog data generalization

Visualizations show that Hedgehog maps trained on one dataset can produce attention weights resembling softmax on new samples; ablations indicate both the feature-map form and distillation matter for fidelity.

C.2.1 Qualitative evidence of Hedgehog data generalization

  • Replacing the Hedgehog feature-map form with T2R-HH or leaving maps untrained lowers fidelity, while attention distillation is critical for retaining reasonably similar weights.Source: §C.2.1 · p. 21
  • Hedgehog maps often produce attention weights similar to softmax on new MRPC and QNLI samples, despite training on CoLA or WikiText-103.Source: §C.2.1 · p. 21, Figure 10 · p. 22, Figure 11 · p. 22
  • The visual comparison samples the third head in BERT-base-uncased’s first, sixth, and twelfth layers on MRPC and QNLI.Source: Figure 10 · p. 22, Figure 11 · p. 22

C.2.2 Quantitative analysis of Hedgehog data generalization

The authors quantify cross-dataset attention fidelity with KL divergence and report that Hedgehog maps trained on CoLA or WikiText-103 remain strong matches to softmax on other GLUE tasks.

C.2.2 Quantitative analysis of Hedgehog data generalization

  • Hedgehog produces better-matching attention weights to softmax on different GLUE tasks, measured by significantly smaller KL divergences.The comparison uses BERT-base-uncased softmax weights as the reference.Source: §C.2.2 · p. 22, Table 14 · p. 23
  • Hedgehog maps trained on CoLA or WikiText-103 best match softmax attention weights on different GLUE tasks under the lower-is-better KL-divergence measure.Source: Table 14 · p. 23
  • BERT models with Hedgehog attentions trained on CoLA or WikiText-103 achieve best GLUE performance when finetuned on different GLUE tasks.The table caption connects this downstream result with improved attention-weight fidelity.Source: Table 15 · p. 23

C.2.3 Hedgehog data generalization via GLUE task transfer

Cross-task finetuning of BERT models with Hedgehog attentions trained on other data achieves best or second-best performance. Sample Llama-2 generations and additional attention visualizations extend the appendix’s qualitative evidence.

C.2.3 Hedgehog data generalization via GLUE task transfer

  • Hedgehog-BERT models achieve best or second-best performance when finetuned on GLUE tasks different from the data used to train their attention maps.The reported transfer results correspond with the preceding attention-fidelity observations.Source: §C.2.3 · p. 23

C.3 Llama-2 SAMSum Generations

  • Four SAMSum test samples compare standard-attention Llama-2 with Hedgehog- and T2R-converted linear models, all LoRA-finetuned; the authors report that Hedgehog makes parameter-efficient pretrained conversion feasible on larger models.The setup uses one A6000 GPU and unquantized bfloat16 weights.Source: §C.3 · p. 24
  • Further work is proposed on generation quality, parameter-training efficiency, and transferring Hedgehog’s linear-attention efficiency to longer contexts.Source: §C.3 · p. 24

C.4 Additional attention weight visualizations

  • Additional visualizations compare softmax and linear attention weights across GLUE tasks, layers, and heads, with rows representing the first, middle, and last layers.Source: §C.4 · p. 27
Loading 2402.04347v1…