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

ArXiv: 2402.04347

🎯 Pitch

Linear attentions can fully match softmax transformer quality when they explicitly learn to mimic softmax's spiky, monotonic attention patterns—a stark departure from prior approaches that just approximate the kernel. The proposed Hedgehog converts a pretrained Llama-2 7B into a viable linear attention variant that boosts ROUGE-1 by 28.1 points over the base model, where competing linearizations cause 16.5-point drops.

1. Executive Summary

This paper proposes Hedgehog, a learnable linear attention mechanism that recovers the expressivity of softmax attention while retaining linear complexity by training feature maps to explicitly mimic softmax attention weights. Experiments on WikiText-103 language modeling, the Long Range Arena benchmark, GLUE, and ImageNet-1K with BERT-base, GPT-2, ViT-B/16, and Llama-2 7B models demonstrate that Hedgehog closes the performance gap between standard softmax attention and prior linear attentions across three regimes: training from scratch, finetuned-conversion, and pretrained-conversion. The method rests on two identified properties missing from prior linear attentions — low-entropy spikiness of attention weights (concentrating probability mass sharply onto relevant tokens, as the softmax exponentiates query-key similarities) and dot-product monotonicity (requiring that attention weights increase as query-key dot products increase) — and operationalizes softmax mimicry through trainable single-layer MLPs with element-wise exponential activations trained via a cross-entropy attention weight distillation loss (the cross-entropy between the linear attention's normalized kernel feature map outputs and the "ground-truth" softmax attention distribution). Hedgehog recovers over 99% of standard Transformer performance in finetuned-conversion settings, achieves up to 6 perplexity point improvements over prior linear attentions on WikiText-103, and enables a viable linear-attention Llama-2 7B that attains 28.1 higher ROUGE-1 points than the base model, establishing that linear attentions can match softmax quality when trained to approximate attention weights directly — though only when the feature map incorporates an exponential activation that induces the spiky, monotonic structure characteristic of softmax.

2. Context and Motivation

The Core Problem: Linear Attentions Are Efficient but Underperform

This paper tackles a specific and well-documented tension in Transformer architecture design: linear attention mechanisms offer dramatic efficiency gains over standard softmax attention, but consistently fail to match their modeling quality. Understanding this tension requires first understanding what makes standard attention expensive and how linear attention attempts to fix it.

Standard softmax attention, as introduced in Vaswani et al. (2017), computes each output token as a weighted sum of all previous tokens' values, where the weights are determined by a softmax over the scaled dot products between the current query and all previous keys. For a sequence of length $n$ and head dimension $d$, computing these pairwise similarity scores for every token against every other token requires $\mathcal{O}(n^2 d)$ time and memory. This quadratic dependence on sequence length becomes a severe bottleneck for long sequences — a 32K-token document requires 1,024 times more attention computation than a 1K-token document, not 32 times more. This scaling behavior limits Transformer applicability in settings where long-range dependencies matter most: processing entire books, analyzing long scientific documents, handling extended dialogues, or modeling high-resolution images and video.

Linear attentions resolve this quadratic bottleneck through a kernel-based reformulation. The key observation, formalized by Katharopoulos et al. (2020) and drawing on kernel methods from Tsai et al. (2019), is that the exponential function in softmax can be viewed as a kernel function $\mathcal{K}(\bm{x}, \bm{x}') = \phi(\bm{x})^\top \phi(\bm{x}')$ evaluated between queries and keys. By replacing the exponential with a feature map $\phi: \mathbb{R}^d \mapsto \mathbb{R}^{d'}$ that is applied element-wise to each query and key before any pairwise interaction, attention can be computed through a clever rearrangement:

$\bm{y}_i = \frac{\phi(\bm{q}_i) \sum_{j=1}^i \big(\phi(\bm{k}_j)^\top \bm{v}_j\big)}{\phi(\bm{q}_i) \sum_{j=1}^i \phi(\bm{k}_j)}$

This formulation eliminates the nested loop over query-key pairs. Instead, the model maintains running sums of $\phi(\bm{k}_j)^\top \bm{v}_j$ and $\phi(\bm{k}_j)$ as it processes the sequence, requiring only $\mathcal{O}(n d d')$ time — linear in sequence length when $d'$ is fixed. For typical Transformer settings with head dimension $d=64$, this quadratic-to-linear reduction translates to significant real-world speed and memory improvements (as demonstrated in the paper's Figure 6, showing near 6× faster inference than FlashAttention at sequence lengths of 32K tokens).

However, this reformulation comes at a steep cost in modeling quality. The paper notes that on standard benchmarks such as WikiText-103, prior linear attentions achieve 4–6 worse perplexity than softmax attention — a gap the authors contextualize as "the equivalent gap between 125M and 255M Transformers" (citing Dai et al., 2019). In essence, choosing a linear attention to gain efficiency currently incurs a penalty equivalent to halving or doubling the model size, which fundamentally undermines the practical value of the efficiency gain.

Why This Problem Matters Across Three Deployment Regimes

The paper identifies three distinct regimes where the performance gap between linear and softmax attention has concrete practical consequences, each with different technical requirements and stakes.

Training-from-scratch is the most straightforward regime: when building new Transformer models, can we use linear attention from the start and achieve comparable performance to a standard Transformer? If successful, this would make training on long sequences significantly more feasible. The current state of affairs, where linear attentions trail by 4–6 perplexity points, makes them unattractive for production language models despite their efficiency advantages. The paper's experiments on Long Range Arena (LRA) classification and WikiText-103 language modeling target this regime directly.

Finetuned-conversion addresses a different practical need. Many organizations have already invested substantial resources in finetuning task-specific Transformers — BERT models for particular classification tasks, ViT models for specific image recognition problems, and so on. Converting these existing models to linear attention variants after finetuning would improve their inference efficiency on long sequences without retraining from scratch. Prior work by Kasai et al. (2021) (Transformer-to-RNN, or T2R) and Mao (2022) explored this regime but found that linear attention models "require additional quadratic attention modules to close the gap." This means the converted models are only partially linear, retaining some quadratic components and thus only achieving partial efficiency gains. The ability to fully convert finetuned models to purely linear attention with minimal performance loss would unlock immediate inference efficiency improvements for deployed models.

Pretrained-conversion is the most ambitious and consequential regime. The goal here is to take large, pretrained language models such as GPT-2 or Llama-2 — models whose training cost is measured in millions of dollars of compute — and convert them into linear attention variants that can then be finetuned on downstream tasks or deployed for efficient inference. This regime matters enormously because it would allow the vast ecosystem of pretrained LLMs to benefit from linear attention efficiency without the prohibitive cost of retraining from scratch. However, it is also the most challenging regime: the queries and keys in a pretrained model were learned under softmax attention, and any linear attention replacement must work well with these fixed representations during conversion and subsequent finetuning.

Critically, the paper notes that this performance gap has been persistent enough to raise theoretical concerns about its fundamental nature. The authors cite recent theoretical work using the Strong Exponential Time Hypothesis (SETH) by Alman & Song (2023) and Keles et al. (2023) showing that "high-quality truly subquadratic algorithms to approximate softmax attention may be impossible with large sequence length $n$." If such theoretical barriers exist, closing the gap empirically would require finding precisely the right structural properties that softmax attention possesses and that linear attention approximations can capture — which is exactly the investigation this paper undertakes.

Prior Approaches and Their Specific Limitations

The paper situates its contribution against a landscape of existing linear attention feature maps, each proposing a specific function $\phi$ to replace the exponential in softmax. These fall into several categories, each with identifiable shortcomings:

Positive-valued feature maps such as $1 + \text{ELU}$ (Katharopoulos et al., 2020) or ReLU (Kasai et al., 2021) ensure attention weights are non-negative, which prior work had identified as important for training stability. However, these functions produce attention distributions that are far more uniform than softmax — they lack what the paper terms "spikiness." The reason is mathematically clear: an exponential function amplifies relative differences between query-key dot products, making the largest dot products dominate the attention distribution, whereas linear or near-linear functions compress these differences, distributing probability mass more evenly across tokens. This uniformity reduces the model's ability to selectively focus on the most relevant tokens in a sequence.

Randomized feature approximations to the softmax or Gaussian kernel, as in Performer (Choromanski et al., 2020) and Random Feature Attention (Peng et al., 2021), are designed to provide unbiased or low-variance estimates of the exponential kernel. However, the paper's empirical analysis (Figure 2) shows that even these methods, which explicitly target softmax approximation, produce much higher-entropy attention distributions than true softmax. The random projection-based feature maps lose the sharp selectivity that characterizes softmax attention, failing to concentrate effectively on the most relevant tokens.

Low-rank approximation methods such as Nyströmformer (Xiong et al., 2021) and Skyformer (Chen et al., 2021) approximate the full attention matrix through subsampled low-rank decompositions rather than kernel feature maps. While these remain subquadratic, the paper focuses on the kernel-based linear attention paradigm because it provides a clean mathematical framework for analyzing and improving attention expressivity through feature map design.

Locality-enhanced methods such as cosFormer (Qin et al., 2022b) incorporate positional information or upweight nearby tokens, adding useful inductive biases but not fundamentally addressing the softmax mimicry problem. The paper's analysis shows cosFormer still produces attention distributions that are substantially more uniform than softmax attention (Figure 2).

Transformer-to-RNN (T2R) by Kasai et al. (2021) is particularly relevant because it is the most direct prior work on the conversion regimes. T2R replaces softmax attention with a ReLU-based feature map and finetunes the converted model. However, the paper shows T2R consistently underperforms softmax attention in conversion settings (Table 1 shows a Matthew's correlation of only 41.1 vs. 58.8 for softmax attention on CoLA), and the resulting attention weights exhibit neither the spikiness nor the monotonicity of softmax weights.

A crucial limitation shared by all these approaches is that they propose fixed, hand-designed feature maps. Each makes a specific functional choice — ReLU, ELU, random Fourier features, cos-based locality — and hopes that this choice captures the essential properties of the exponential function. The paper's core critique is that these fixed feature maps systematically miss two properties that characterize softmax attention and that turn out to be critical for performance.

The Two Missing Properties: A Diagnostic Framework

The paper's primary conceptual contribution before proposing Hedgehog is the identification and empirical validation of two specific properties that softmax attention possesses and that prior linear attentions lack. These properties serve as both a diagnostic framework for understanding the performance gap and a design specification for what a successful linear attention must achieve.

Low-entropy "spikiness" refers to the ability of attention weights to concentrate sharply on the most relevant tokens while largely ignoring irrelevant ones. Mathematically, this manifests as attention distributions with low entropy — probability mass is concentrated on a small number of tokens rather than spread uniformly. In softmax attention, the exponential function amplifies differences between query-key dot products: a token whose dot product with the query is slightly larger than others will receive disproportionately higher weight. This property is widely understood as central to attention's effectiveness, enabling models to selectively attend to specific information in a sequence whether for translation alignment (Bahdanau et al., 2014), image patch relationships in vision Transformers (Dosovitskiy et al., 2020; Caron et al., 2021), or content-based retrieval of earlier tokens relevant to predicting the next word.

The paper demonstrates this property's importance through a controlled experiment on the Associative Recall (AR) task (Section 3.2, Figure 4). AR is a synthetic next-token prediction task where the model must remember key-value pairs presented earlier in the sequence. The experiment is particularly revealing because it isolates attention quality from other model components. Softmax attention solves the task perfectly (100% accuracy) while producing low-entropy attention weights (approximately 0.5–1.5 nats depending on the head). In contrast, prior linear attentions (Performer, cosFormer, $1+\text{ELU}$, ReLU) all fail to achieve even 20% accuracy while producing much higher entropy weights (ranging from 3.5–6 nats). Strikingly, when the authors introduce a simple temperature-scaled exponential feature map $\phi_t(x) = \exp(x \cdot t)$, they find that increasing temperature from $t=1$ to $t=2$ transforms the model from failing to solving the AR task perfectly, accompanied by a corresponding drop in attention entropy. This provides direct evidence that low-entropy, spiky attention distributions are not merely correlated with performance but causally necessary for it — at least for tasks requiring selective token recall, which is a fundamental building block of language modeling.

Dot-product monotonicity requires that attention weights increase as query-key dot products increase. This is a more subtle property that the paper argues is particularly important for conversion settings, where a pretrained model's query and key representations have been optimized under softmax attention. In softmax attention, monotonicity holds automatically: the exponential function is strictly increasing, so a larger dot product always produces a larger unnormalized attention weight, and the softmax denominator being the same for all keys means this monotonicity is preserved in the final normalized weights.

The paper demonstrates that prior linear attentions violate this property in practice (Figure 3, Table 1). When plotting attention weights against query-key dot products for a finetuned BERT model on the CoLA task, softmax attention shows a clean, smooth monotonic curve. In contrast, prior linear attentions (ReLU, $1+\text{ELU}$, Performer) produce non-monotonic relationships — increasing the query-key similarity can actually decrease the attention weight, and vice versa. The paper argues this causes training difficulties during conversion because "trying to upweight attentions by increasing product similarity can actually result in decreased attention weights," creating conflicting gradients between the attention mechanism and the original model parameters. The empirical evidence supports this: in the finetuned-conversion setting on CoLA (Table 1), no prior linear attention recovers the original BERT's Matthew's correlation of 58.8, with the best prior method achieving only 41.1.

The interaction between these two properties is notable. A feature map can be spiky without being monotonic, or monotonic without being spiky. The temperature-scaled exponential $\phi_t(x) = \exp(x \cdot t)$ is spiky when $t$ is large and is monotonic by construction, but the paper's analysis of the $\phi_2$ feature map in the finetuned-conversion setting (Table 1) shows that spikiness alone is insufficient for conversion — this feature map solved AR but failed to recover CoLA performance. Both properties appear necessary for different aspects of linear attention performance.

Reconciling Contradictory Prior Findings

An important but implicit contribution of the diagnostic framework is that it helps reconcile apparently contradictory findings in the linear attention literature. Some prior work found that linear attentions could approach softmax performance in certain settings, while other work found persistent gaps. The two-property framework provides a lens for understanding these discrepancies: a linear attention that happens to produce spiky weights (perhaps because the learned query-key dot products happen to fall in a regime where a particular feature map amplifies differences) might perform well in training-from-scratch, but the same method might fail catastrophically in conversion because it lacks monotonicity, producing conflicting gradients when fine-tuning pretrained weights.

This also explains the paper's counterintuitive finding about associative recall versus conversion: the $\phi_2$ feature map (element-wise exponential with temperature 2) is spiky enough to solve AR perfectly, but its non-exact approximation to the softmax means it is not monotonic with respect to the full dot product — it operates element-wise, not on the dot product directly — causing it to fail in the conversion setting. The distinction between element-wise exponentials (which are spiky) and dot-product exponentials (which are both spiky and monotonic in the dot product) clarifies why simple fixes to existing linear attentions are insufficient.

How the Paper Positions Itself

The paper positions its contribution not as proposing yet another fixed feature map, but as a methodological shift from hand-designing feature maps to learning them. This shift is motivated by a key empirical finding that bridges the theoretical challenge and the practical solution: simple polynomial (Taylor) approximations to the exponential function can recover both spikiness and monotonicity in the bounded regimes where real queries and keys operate, and this recovery matches softmax performance on both AR and BERT-finetuned conversion. However — and this is the crucial motivation for Hedgehog's design — even a second-degree Taylor approximation is computationally impractical, requiring feature maps of dimension $d' = 1 + d + d^2$, leading to $\mathcal{O}(n d^3)$ attention complexity rather than the $\mathcal{O}(n d^2)$ of prior linear attentions.

This tension — Taylor approximations are expressive but inefficient, prior feature maps are efficient but inexpressive — frames the paper's central design problem: can we learn feature maps that capture the spiky and monotonic properties of the exponential while maintaining $\mathcal{O}(n d^2)$ complexity? Hedgehog answers this affirmatively by using trainable MLPs that map $\mathbb{R}^d \mapsto \mathbb{R}^d$ (maintaining efficiency) but are trained via knowledge distillation from softmax attention weights (capturing the missing properties). The paper thus positions itself as offering a practical solution to a well-characterized problem, grounded in a diagnostic understanding of why prior approaches fail, rather than as introducing a new attention mechanism in isolation.

The paper also explicitly notes a connection to adapter methods (Houlsby et al., 2019), positioning the Hedgehog MLPs as light-weight trainable modules inserted into existing Transformer architectures — a design choice that makes the conversion regimes particularly practical, as the original model weights can be preserved and only the Hedgehog feature maps need training during attention distillation.

3. Technical Approach

This is primarily a methodology paper whose core idea is that linear attention feature maps should be learned to mimic softmax attention weights, rather than hand-designed, and that doing so recovers two critical properties — low-entropy spikiness and dot-product monotonicity — that prior linear attentions systematically lack.

3.1 Reader Orientation

Hedgehog is a system for replacing the expensive softmax attention in Transformers with efficient linear attention without sacrificing modeling quality. The problem it solves is that existing linear attentions are fast but inaccurate (4–6 perplexity points worse than softmax on language modeling), and the "shape" of the solution is to insert small trainable neural networks (single-layer MLPs) after the query and key projections in each attention head, then train these MLPs to make the resulting linear attention weights match the original softmax attention weights as closely as possible.

3.2 Big-Picture Architecture (Diagram in Words)

The system has three major components layered on top of a standard Transformer:

1. Hedgehog Feature Maps (one per attention head, per layer) — trainable single-layer MLPs with element-wise exponential activations, inserted after the query and key projections in each attention head. They map from head dimension $d$ to $d$ (with an optional doubling to $2d$ for negation handling), replacing the fixed feature maps in prior linear attentions.

2. Attention Weight Distillation Loss — a cross-entropy objective that trains only the Hedgehog MLPs (while keeping all other model weights frozen) to match the attention weight distributions that softmax attention would have produced for the same queries and keys.

3. Task-Specific Finetuning — after distillation, all model weights (or a parameter-efficient subset, such as LoRA adapters) are unfrozen and trained with the standard task loss (e.g., next-token prediction, classification).

Information flows as follows: input tokens enter the Transformer → query and key projections produce $\bm{q}_i, \bm{k}_j \in \mathbb{R}^d$ → Hedgehog MLPs transform these into feature map representations $\phi_{\text{mlp}}(\bm{q}_i), \phi_{\text{mlp}}(\bm{k}_j) \in \mathbb{R}^d$ → linear attention computes outputs using the kernel trick (Equation 2 in the paper) → during distillation, the attention weights from the linear attention are compared against the ground-truth softmax weights via cross-entropy, and only the MLP parameters are updated → during finetuning, all parameters are updated with the task loss.

3.3 Roadmap for the Deep Dive

• First, the two missing properties (spikiness and monotonicity) and how they motivate the design — because the entire technical approach is a response to these diagnostic findings.
• Second, the Taylor exponential approximation as a bridge result — it shows that learning is not strictly necessary (a fixed polynomial can work) but is practically necessary (polynomials are too inefficient), which motivates why we need learned feature maps.
• Third, the Hedgehog MLP feature map — the core architectural component, including the exponential activation, the identity initialization, the negation mapping, and the numerical stability softmax variant.
• Fourth, the attention weight distillation loss — how it is computed, what it optimizes, and why it is applied before task-specific finetuning rather than jointly.
• Fifth, the training protocols for the three regimes — training-from-scratch, finetuned-conversion, and pretrained-conversion — each with different initialization, freezing, and optimization choices.
• Sixth, the complexity analysis — how Hedgehog maintains $\mathcal{O}(n d^2)$ complexity despite the learnable feature maps, and how this compares to the Taylor approximation's $\mathcal{O}(n d^3)$.

3.4 Detailed, Sentence-Based Technical Breakdown

The Two Missing Properties as Design Specifications

The technical approach begins not with a method but with a diagnostic framework. Before proposing Hedgehog, the paper identifies two specific properties that softmax attention possesses and that prior linear attentions lack, and it validates empirically that these properties are causally connected to performance, not merely correlated with it.

Low-entropy spikiness is the property that attention weights should concentrate sharply on a small number of relevant tokens rather than being spread uniformly. This is formalized through the entropy of the attention distribution for a given query:

$H(\bm{a}_i) = -\sum_{j=1}^i a_{ij} \log a_{ij}$

where $\bm{a}_i = [a_{i1}, \ldots, a_{ii}]$ is the vector of attention weights from query $i$ to all keys $j \leq i$, and $a_{ij} = \text{sim}(\bm{q}_i, \bm{k}_j)$ is the normalized attention weight from query $i$ to key $j$.

What it computes: a scalar measure of how "peaked" the attention distribution is. Entropy is 0 when all probability mass is on a single token (maximally spiky) and $\log(i)$ when probability is uniformly distributed across all $i$ tokens (maximally flat).

Why this form: entropy captures the concentration property that softmax naturally produces through exponentiation. The exponential function amplifies relative differences — a query-key dot product that is only 2× larger than another becomes $e^2 \approx 7.4\times$ larger in unnormalized weight — which collapses entropy. Linear feature maps without exponentials fail to produce this amplification, resulting in higher entropy and, the paper argues, reduced ability to selectively attend to relevant information.

The empirical validation comes from the Associative Recall (AR) task (Section 3.2, Figure 4). In AR, the model sees a sequence of key-value pairs followed by a query key; it must output the corresponding value. Softmax attention solves this with 100% accuracy and produces attention weights with entropy around 0.5–1.5 nats. Prior linear attentions (Performer, cosFormer, $1 + \text{ELU}$, ReLU) all achieve below 20% accuracy and produce entropy in the range of 3.5–6 nats. The causal connection is demonstrated through a temperature-scaled exponential feature map $\phi_t(x) = \exp(x \cdot t)$: at $t=1$, the model fails AR and produces high entropy; at $t=2$, the model solves AR perfectly and produces low entropy. This is a controlled ablation — the only change is the temperature parameter controlling spikiness — and it establishes that spikiness is necessary for this fundamental recall capability.

Dot-product monotonicity is the property that the attention weight $a_{ij}$ between query $\bm{q}_i$ and key $\bm{k}_j$ should be a monotonically increasing function of their dot product $\bm{q}_i^\top \bm{k}_j$. Formally:

$\frac{\partial a_{ij}}{\partial (\bm{q}_i^\top \bm{k}_j)} \geq 0 \quad \text{for all } i, j$

What it computes: the requirement that increasing the raw similarity between a query and a key never decreases the attention weight assigned to that key, and decreasing similarity never increases the weight.

Why this form: monotonicity ensures that the gradients from the attention mechanism to the query and key parameters are directionally consistent. If increasing a dot product can decrease an attention weight (which happens in non-monotonic feature maps), then gradient descent receives conflicting signals: the attention weight wants to go up (to attend more to a relevant key) but the feature map's local slope means increasing the dot product would actually reduce the weight. The paper argues this creates training instability particularly acute in conversion settings, where pretrained query and key representations must be preserved while the attention mechanism is swapped.

The empirical validation uses a finetuned BERT-base model on the CoLA task (Section 3.2, Figure 3, Table 1). When plotting attention weights against query-key dot products for all query-key pairs in a test batch, softmax attention produces a clean monotonic curve. Prior linear attentions (ReLU, $1+\text{ELU}$, Performer, cosFormer) all produce non-monotonic scatter plots where the same dot product can map to widely varying attention weights. The performance consequence is documented in Table 1: softmax BERT achieves a Matthew's correlation of 58.8 on CoLA, while prior linear attentions achieve at most 41.1 — a 30% relative drop. Notably, the spiky $\phi_2$ (temperature-2 exponential applied element-wise) also fails to recover the original performance, achieving only 52.4. This is because element-wise exponentials, while spiky, are not monotonic in the dot product — they are monotonic in each individual dimension of the query and key vectors, but the dot product is a sum over dimensions, and an element-wise exponential does not preserve monotonicity in this sum.

The critical distinction is that softmax attention applies the exponential to the dot product directly: $\exp(\bm{q}^\top \bm{k} / \sqrt{d})$. An element-wise exponential feature map $\phi(\bm{x}) = [\exp(x_1), \ldots, \exp(x_d)]$ computes $\phi(\bm{q})^\top \phi(\bm{k}) = \sum_{\ell=1}^d \exp(q_\ell) \exp(k_\ell) = \sum_{\ell=1}^d \exp(q_\ell + k_\ell)$, which is a fundamentally different function from $\exp(\sum_\ell q_\ell k_\ell) = \exp(\bm{q}^\top \bm{k})$. The former is a sum of per-dimension exponentials; the latter is the exponential of a sum of per-dimension products. Only the latter preserves monotonicity in the dot product.

The Taylor Exponential as a Bridge Result

The paper uses a second-degree Taylor approximation to the exponential function as a bridge between the diagnostic findings and Hedgehog's design. This result serves two purposes: (1) it demonstrates that recovering spikiness and monotonicity is sufficient to match softmax performance, and (2) it demonstrates that doing so with a fixed functional form is computationally impractical, motivating the need for learned approximations.

The second-degree Taylor approximation of the exponential is:

$\exp(\bm{q}^\top \bm{k}) \approx 1 + \bm{q}^\top \bm{k} + \frac{1}{2}(\bm{q}^\top \bm{k})^2$

where $\bm{q}, \bm{k} \in \mathbb{R}^d$ are the query and key vectors, and $\bm{q}^\top \bm{k}$ is their dot product.

What it computes: a quadratic approximation to the exponential function that can be expressed as an inner product of feature maps. The feature map $\phi_{\text{taylor}}: \mathbb{R}^d \mapsto \mathbb{R}^{1 + d + d^2}$ is constructed as:

$\phi_{\text{taylor}}(\bm{x}) = \Big[1, x_1, \ldots, x_d\Big] \cup \Big[x_i \cdot x_j \;|\; i, j \in [d]\Big]$

where the first element captures the constant term, the next $d$ elements capture the linear terms, and the final $d^2$ elements capture all pairwise products $x_i \cdot x_j$ for the quadratic term $(\bm{q}^\top \bm{k})^2 = \sum_i \sum_j q_i k_i q_j k_j$.

Why this form: the Taylor series is the canonical local approximation to any smooth function. A second-degree approximation captures both the monotonicity (through the linear term) and the convexity (through the quadratic term) of the exponential, which together produce spikiness by amplifying larger dot products quadratically. The first-degree approximation $\exp(x) \approx 1 + x$ would be monotonic but not spiky (it amplifies linearly), while higher-degree approximations would capture more of the exponential's curvature but at prohibitive computational cost.

The paper reports two key findings with the Taylor approximation:

1. It works qualitatively: the second-degree Taylor approximation recovers both spikiness and monotonicity in practice (Section 4.1, Figure 5). BERT queries and keys happen to fall in a bounded regime where the second-order Taylor series tracks the exponential reasonably well, maintaining monotonicity.

2. It matches softmax performance: on Associative Recall, the Taylor approximation solves the task perfectly (Table 3). On BERT-finetuned conversion for CoLA, the Taylor approximation recovers the original model's performance (Table 2 shows "expressivity" and "spikiness" checks passing), demonstrating that a feature map with the right structural properties can close the gap.

3. It is computationally impractical: even with $p=2$, the feature map dimension is $d' = 1 + d + d^2$, which for a typical head dimension $d=64$ means $d' = 1 + 64 + 4096 = 4161$ — a 65× blowup. This makes attention complexity $\mathcal{O}(n d^3)$ rather than $\mathcal{O}(n d^2)$. Table 2 characterizes this as an "efficiency vs. expressivity tradeoff." For comparison, Figure 6 shows that the Taylor approximation at sequence length 32K is actually slower and uses more memory than FlashAttention (which is quadratic in time but highly optimized), defeating the purpose of using linear attention.

This bridge result establishes the design specification for Hedgehog: achieve the expressivity of the Taylor approximation (spiky, monotonic, performant) with the efficiency of prior linear attentions ($\mathcal{O}(n d^2)$, feature dimension equal to head dimension). The solution is to learn the feature map rather than deriving it from a fixed functional form.

The Hedgehog MLP Feature Map: Architecture and Design Choices

The core technical innovation is replacing the hand-designed feature map $\phi$ in Equation 2 with a trainable single-layer MLP. For each attention head in each layer of the Transformer, Hedgehog introduces two MLPs — one for queries and one for keys — with shared weights (the same MLP is applied to both queries and keys for that head). The MLP is defined as:

$\phi_{\text{mlp}}(\bm{x}) = \Phi(\bm{W}^\top \bm{x} + \bm{b})$

where $\bm{W} \in \mathbb{R}^{d \times d'}$ is a learned weight matrix, $\bm{b} \in \mathbb{R}^{d'}$ is a learned bias vector, $\bm{x} \in \mathbb{R}^d$ is the input query or key vector, and $\Phi$ is an activation function applied element-wise.

What it computes: a learned linear projection followed by a nonlinearity, mapping each query and key from the head dimension $d$ to a feature dimension $d'$. In the default configuration, $d' = d$, so the feature map preserves the original dimensionality.

Why this form: a single-layer architecture is the minimal learnable feature map. It has sufficient capacity to learn a softmax approximation (as demonstrated by the results) while adding minimal computational overhead — the MLP computation is $\mathcal{O}(d^2)$ per token, which is comparable to the original query and key projections it sits alongside. Deeper MLPs would add latency without clear benefit, while linear-only transformations (no activation) cannot capture the nonlinear amplification that produces spikiness.

The paper introduces several specific design choices that are critical to the method's effectiveness:

Choice 1: Element-wise exponential activation. The activation function $\Phi$ is set to the element-wise exponential, making the feature map:

$\phi_{\text{mlp}}(\bm{x}) = \Big[\exp(\bm{w}_1^\top \bm{x} + b_1), \ldots, \exp(\bm{w}_{d}^\top \bm{x} + b_d)\Big]$

where $\bm{w}_\ell \in \mathbb{R}^d$ is the $\ell$-th row of $\bm{W}$ and $b_\ell$ is the $\ell$-th bias.

What it computes: each output dimension $\ell$ applies a learned linear projection $\bm{w}_\ell^\top \bm{x} + b_\ell$ to the entire input vector, then exponentiates the result. The dot product between two feature-mapped vectors is therefore:

$\phi_{\text{mlp}}(\bm{q})^\top \phi_{\text{mlp}}(\bm{k}) = \sum_{\ell=1}^{d} \exp(\bm{w}_\ell^\top \bm{q} + b_\ell) \exp(\bm{w}_\ell^\top \bm{k} + b_\ell) = \sum_{\ell=1}^{d} \exp(\bm{w}_\ell^\top (\bm{q} + \bm{k}) + 2b_\ell)$

Why this form: the exponential is chosen specifically because of the diagnostic finding in Section 3.2 that element-wise exponentials can induce spikiness (Figure 4 shows $\phi_2$ solving AR). However, unlike the fixed $\phi_t(x) = \exp(x \cdot t)$ which applies the exponential to each coordinate independently, the learned linear projections $\bm{w}_\ell^\top \bm{x}$ allow each output dimension to capture a linear combination of input dimensions before exponentiation. This means the MLP can learn to approximate $\exp(\bm{q}^\top \bm{k})$ in the dot-product space by learning $\bm{W}$ such that the sum over $\ell$ of exponentials approximates the exponential of the dot product — it can distribute the approximation across multiple learned basis directions. The paper shows empirically (Figure 7) that this learned exponential feature map produces attention weights that visually match softmax weights with much higher fidelity than fixed feature maps, and that replacing the exponential with ReLU (the T2R-HH ablation) significantly degrades attention weight fidelity (Figure 8) and downstream performance (Table 9).

Choice 2: Identity initialization. The weight matrix $\bm{W}$ is initialized to the identity matrix and the bias $\bm{b}$ to zeros:

nn.init.eye_(self.layer.weight)
nn.init.zeros_(self.layer.bias)

What it computes: at initialization, $\phi_{\text{mlp}}(\bm{x}) = \exp(\bm{x})$ element-wise — the feature map is a simple element-wise exponential of the original query/key vectors.

Why this form: identity initialization ensures that at the start of training, the Hedgehog attention is a well-defined linear attention (the element-wise exponential, which the paper's earlier experiments showed is spiky). This provides a reasonable starting point from which the MLP can adapt. Random initialization would produce arbitrary feature maps that might be unstable or produce degenerate attention distributions. The identity initialization also means that for training-from-scratch (where no distillation is used), the model begins with a feature map that is known to be spiky and can solve simple recall tasks, providing a better initialization than random weights.

Choice 3: Negation mapping for handling negative similarities. In practice (Appendix A.1), the feature map is extended to map from $\mathbb{R}^d$ to $\mathbb{R}^{2d}$ by concatenating the exponentials of both the learned projections and their negations:

$\phi_{\text{mlp}}(\bm{x}) = \Big[\exp(\bm{w}_1^\top \bm{x} + b_1), \ldots, \exp(\bm{w}_{d}^\top \bm{x} + b_d), \exp(-\bm{w}_1^\top \bm{x} - b_1), \ldots, \exp(-\bm{w}_{d}^\top \bm{x} - b_d)\Big]$

What it computes: for each learned direction $\bm{w}_\ell$, the feature map includes both $\exp(\bm{w}_\ell^\top \bm{x} + b_\ell)$ and $\exp(-\bm{w}_\ell^\top \bm{x} - b_\ell)$. This doubles the feature dimension from $d$ to $2d$, but since $d$ is typically small (64), this is a constant-factor increase.

Why this form: prior linear attentions using ReLU or $1+\text{ELU}$ effectively discard negative query-key interactions — ReLU sets all negative values to zero, and ELU produces small negative values that are then added to 1. The paper argues that "the additional negation mapping in $\mathbb{R}^{2d}$ intuitively lets us better factor in negative dimensionalities." For any direction $\bm{w}_\ell$, if a query has a strongly positive projection $\bm{w}_\ell^\top \bm{q}$ and a key has a strongly negative projection $\bm{w}_\ell^\top \bm{k}$, the standard feature map would produce $\exp(\text{positive}) \cdot \exp(\text{negative}) \approx 0$ in that dimension. With the negation mapping, the negated key projection $-\bm{w}_\ell^\top \bm{k}$ is positive, so $\exp(\bm{w}_\ell^\top \bm{q}) \cdot \exp(-\bm{w}_\ell^\top \bm{k}) = \exp(\bm{w}_\ell^\top (\bm{q} - \bm{k}))$ captures dissimilarity in that direction — effectively, the feature map can represent both positive and negative alignment between queries and keys. This is important because softmax attention handles negative dot products naturally (they produce small but non-zero weights), while many prior linear attentions cannot.

Choice 4: Softmax-over-features variant for numerical stability. In practice (Appendix A.1), the paper found that applying a softmax over the feature dimension improves numerical stability compared to raw exponentials:

$\phi_{\text{mlp}}(\bm{x}) = \left[\frac{\exp(\bm{w}_1^\top \bm{x})}{\sum_{i=1}^d \exp(\bm{w}_i^\top \bm{x})}, \ldots, \frac{\exp(\bm{w}_d^\top \bm{x})}{\sum_{i=1}^d \exp(\bm{w}_i^\top \bm{x})}\right]$

What it computes: instead of outputting raw exponentials, each feature dimension is normalized by the sum over all feature dimensions, producing a probability distribution over the $d$ learned directions. The bias term $b_\ell$ is omitted in this variant.

Why this form: raw exponentials can produce very large values if the learned projections $\bm{w}_\ell^\top \bm{x}$ become large, leading to numerical overflow. The softmax normalization constrains each feature vector to have components in $[0,1]$ summing to 1, preventing overflow while preserving the relative magnitudes that produce spikiness. The paper reports this "also performing better than dividing each element by the max over $\{\exp(\bm{w}_i^\top \bm{x} + \bm{b})\}_{i=1}^d$" (a common stabilization trick). Note that this is a softmax over the feature dimension, not over the sequence — the linear attention still uses the kernel trick over the sequence dimension as in Equation 2.

The Attention Weight Distillation Loss

The second core component of Hedgehog is the training objective used to make the feature maps produce attention weights that match softmax attention. This is a knowledge distillation loss applied at the level of individual attention weights, not at the output level.

For a given query $\bm{q}_i$ (the $i$-th token in the sequence) and all keys $\{\bm{k}_j\}_{j=1}^i$ (all tokens up to and including position $i$, since this is causal attention), the distillation loss for that query is:

$\mathcal{L}_i = -\sum_{j=1}^i \frac{\exp(\bm{q}_i^\top \bm{k}_j / \sqrt{d})}{\sum_{m=1}^i \exp(\bm{q}_i^\top \bm{k}_m / \sqrt{d})} \log \frac{\phi_{\text{mlp}}(\bm{q}_i)^\top \phi_{\text{mlp}}(\bm{k}_j)}{\sum_{m=1}^i \phi_{\text{mlp}}(\bm{q}_i)^\top \phi_{\text{mlp}}(\bm{k}_m)}$

where $\bm{q}_i, \bm{k}_j \in \mathbb{R}^d$ are the query and key vectors from the frozen pretrained Transformer (before the Hedgehog MLP is applied), $d$ is the head dimension, and $\phi_{\text{mlp}}$ is the Hedgehog feature map.

What it computes: the cross-entropy between two probability distributions over keys $1, \ldots, i$:

1. The target distribution (first fraction): the standard softmax attention weights, where the weight for key $j$ is proportional to $\exp(\bm{q}_i^\top \bm{k}_j / \sqrt{d})$. These are computed using the frozen pretrained model's queries and keys.

2. The predicted distribution (second fraction): the linear attention weights produced by applying Hedgehog feature maps to the same queries and keys, then normalizing via the linear attention denominator (sum over keys of feature map dot products).

The loss is the negative log-likelihood of the target distribution under the predicted distribution — standard cross-entropy where the softmax weights serve as soft labels.

Why this form: cross-entropy is the natural objective for matching one probability distribution to another. It penalizes the predicted distribution heavily when it assigns low probability to keys that the softmax weights highly (the model "misses" important tokens), and penalizes it when it assigns high probability to keys that softmax weights low (the model "attends to" irrelevant tokens). This directly addresses both the spikiness and monotonicity deficiencies: to minimize the cross-entropy, the predicted distribution must concentrate probability on the same small set of keys as the target distribution (spikiness), and it must assign higher weights to keys with larger query-key dot products to match the target's ordering (monotonicity, since the target is monotonic by construction).

The paper notes (Section 4.2) that "for multi-layer and multi-head attention Transformers, we apply a separate MLP to each head and each layer, and use the same $\phi_{\text{mlp}}$ for the queries and keys." The total distillation loss is the sum of $\mathcal{L}_i$ over all queries, all heads, and all layers, trained jointly with a single optimizer. This is computationally efficient because a single forward pass through the model produces all target attention weights (by running the frozen pretrained attention) and all predicted attention weights (by running the Hedgehog feature maps on the same queries and keys), and the loss for all heads and layers can be aggregated and backpropagated in one step.

A critical implementation detail is that only the Hedgehog MLPs are trained during distillation; all other model parameters (query/key projections, value projections, output projections, feedforward layers, layer norms) are frozen. This is motivated by the conversion setting: the pretrained model has learned useful query and key representations, and modifying them during distillation would change the target distribution being matched, creating a moving-target problem. By freezing the base model, the target softmax distribution is fixed, and the MLPs simply learn to approximate it. After distillation, the MLPs can produce linear attention weights that closely track softmax weights for the same queries and keys.

The paper provides PyTorch-like pseudocode for implementing this in practice (Appendix A.2, Listing 1–3). The key abstraction is replacing each attention layer with a "HedgehogAttention" wrapper that, during training mode, runs both the original softmax attention (to get target weights) and the Hedgehog linear attention (to get predicted weights), returning both as outputs for loss computation.

Training Protocols for the Three Regimes

The paper deploys Hedgehog in three distinct regimes, each with a different training protocol reflecting different constraints on what can be modified and what performance target must be achieved.

Training-from-scratch (Section 5.2): When training a new Transformer model from scratch with Hedgehog attention, there is no pretrained softmax attention to distill from. Instead, the Hedgehog MLPs are inserted at initialization and trained end-to-end with the task loss (e.g., cross-entropy for next-token prediction). The MLPs are initialized with identity weights (so the initial feature map is an element-wise exponential), and all parameters — MLP weights, query/key/value projections, feedforward layers, embeddings — are trained jointly.

The paper does not use distillation loss in this regime because there are no pretrained attention weights to serve as targets. Instead, the model learns to produce effective attention weights directly from the task signal, with the Hedgehog architecture providing the inductive bias (via the exponential activation) that spiky and monotonic weights are desirable. The success of this approach (Tables 6, 7) demonstrates that the Hedgehog feature map is not merely a way to convert pretrained models but is a viable attention mechanism in its own right.

Finetuned-conversion (Section 5.3): This is the primary regime the paper targets. A Transformer that has been finetuned on a specific task (e.g., BERT on CoLA, ViT on ImageNet) is converted to linear attention while attempting to recover the original task performance. The protocol proceeds in two stages:

Stage 1: Attention distillation. The Hedgehog MLPs are inserted (one per head per layer). All original model weights are frozen. The MLPs are trained using the attention weight distillation loss (Section 4.2, Equation 4) on data from the target task. For BERT conversion, the paper trains MLPs "up to five epochs with early stopping based on validation loss" with "learning rate 1e-2, weight decay 0, AdamW optimizer" (Appendix B.4). This stage alone produces a model whose attention weights approximate softmax weights but whose outputs may not recover task performance because the value matrices and feedforward layers haven't adapted to the new attention distributions.

Stage 2: Task-specific finetuning. All model parameters (or a selected subset) are unfrozen, and the entire model is trained with the standard task loss (classification cross-entropy for GLUE, image classification loss for ImageNet). For BERT conversion, the paper trains with "batch size 8, learning rate 1e-5, weight decay 0, AdamW optimizer, and cosine scheduler for up to five epochs" (Appendix B.4). This stage allows the model to adapt its representations to the slightly different attention distributions produced by the Hedgehog feature maps, recovering any residual performance gap from the distillation stage.

The two-stage protocol is important because it decouples the problems of (1) learning to approximate softmax attention weights and (2) adapting model parameters to work with the linear attention approximation. If distillation and task finetuning were done simultaneously from the start, the changing query/key representations (as task finetuning updates them) would create a moving target for the MLPs, potentially leading to instability.

Pretrained-conversion (Section 5.4): This is the most ambitious regime: converting a large, generally pretrained model (GPT-2, Llama-2) into a linear attention variant that can then be finetuned on downstream tasks. The protocol is similar to finetuned-conversion but with two important differences:

First, the distillation data may differ from the finetuning data. For GPT-2 conversion on WikiText-103, the MLPs are distilled on WikiText-103 text using "batch size 8, learning rate 0.01, zero weight decay, AdamW optimizer, and 1024-tokens per input" for two epochs (Appendix B.5). For Llama-2 conversion on SAMSum summarization, the MLPs are distilled on SAMSum dialogue-summary pairs using "learning rate 0.01, zero weight decay, AdamW optimizer, and batch size 8 with gradient accumulation" for two epochs (Appendix B.5). The paper's generalization experiments (Section 5.1, Tables 4, 5, 14; Figures 9–11) show that Hedgehog feature maps learned on one dataset (CoLA or WikiText-103) still produce attention weights that better match softmax attention on other datasets (MRPC, QNLI, etc.) compared to prior linear attentions, providing evidence that distillation on general text corpora transfers to downstream tasks.

Second, parameter-efficient finetuning may be used. For Llama-2 7B conversion, the paper applies LoRA (Hu et al., 2021) with "alpha parameter 16 and rank 8" to the query, key, value, and output projections (Appendix B.5). The Hedgehog MLPs themselves constitute only 0.495% of the original model size, making this a very lightweight modification. The combination of Hedgehog attention and LoRA finetuning is trained on a single A6000 GPU, demonstrating that the method scales to modern LLM sizes without requiring large-scale distributed training.

Complexity Analysis: How Hedgehog Maintains $\mathcal{O}(n d^2)$ Efficiency

A critical design constraint is that Hedgehog must remain more efficient than softmax attention for long sequences, which requires that the feature dimension $d'$ be less than the sequence length $n$. Hedgehog achieves this by setting $d' = d$ (or $2d$ with the negation mapping), which for typical head dimensions of 64–128 is well below typical sequence lengths of 512–32K.

The attention computation for Hedgehog follows the standard linear attention formula (Equation 2), but with $\phi_{\text{mlp}}$ replacing the fixed feature map:

$\bm{y}_i = \frac{\phi_{\text{mlp}}(\bm{q}_i) \sum_{j=1}^i \big(\phi_{\text{mlp}}(\bm{k}_j)^\top \bm{v}_j\big)}{\phi_{\text{mlp}}(\bm{q}_i) \sum_{j=1}^i \phi_{\text{mlp}}(\bm{k}_j)}$

where $\phi_{\text{mlp}}(\bm{q}_i), \phi_{\text{mlp}}(\bm{k}_j) \in \mathbb{R}^{d}$ (or $\mathbb{R}^{2d}$), and $\bm{v}_j \in \mathbb{R}^d$.

What it computes: for each position $i$ in the sequence, the output is a weighted sum of all previous value vectors, where the unnormalized weight for key $j$ is $\phi_{\text{mlp}}(\bm{q}_i)^\top \phi_{\text{mlp}}(\bm{k}_j)$. The computation is organized to avoid materializing the full $n \times n$ attention matrix by maintaining running sums as the sequence is processed.

Complexity breakdown:

• Computing $\phi_{\text{mlp}}(\bm{q}_i)$ and $\phi_{\text{mlp}}(\bm{k}_i)$ for each token: $\mathcal{O}(d^2)$ per token (one matrix-vector multiplication with $\bm{W} \in \mathbb{R}^{d \times d}$), totaling $\mathcal{O}(n d^2)$ for the sequence.
• Maintaining the running sum $\bm{S}_i = \sum_{j=1}^i \phi_{\text{mlp}}(\bm{k}_j)^\top \bm{v}_j$: this is a $d \times d$ outer product, updated in $\mathcal{O}(d^2)$ per token, totaling $\mathcal{O}(n d^2)$.
• Maintaining the running sum $\bm{z}_i = \sum_{j=1}^i \phi_{\text{mlp}}(\bm{k}_j)$: this is a $d$-dimensional vector, updated in $\mathcal{O}(d)$ per token, totaling $\mathcal{O}(n d)$.
• Computing the final output $\bm{y}_i = \phi_{\text{mlp}}(\bm{q}_i) \bm{S}_i / (\phi_{\text{mlp}}(\bm{q}_i)^\top \bm{z}_i)$: the numerator is a vector-matrix product in $\mathcal{O}(d^2)$, the denominator is a dot product in $\mathcal{O}(d)$, totaling $\mathcal{O}(d^2)$ per token.

The total complexity is $\mathcal{O}(n d^2)$ for time and $\mathcal{O}(n d)$ for memory (storing the feature-mapped keys and the running sums). This is the same asymptotic complexity as prior linear attentions (Katharopoulos et al., 2020; Choromanski et al., 2020) and compares favorably to softmax attention's $\mathcal{O}(n^2 d)$ time and $\mathcal{O}(n^2)$ memory.

Why this complexity matters: the critical comparison is not to softmax attention (which linear attention always beats asymptotically) but to the Taylor approximation (Section 4.1) and to the sequence length threshold where linear attention becomes practically faster. The Taylor approximation with $p=2$ requires feature dimension $d' = 1 + d + d^2$, making the per-token cost $\mathcal{O}(d^3)$ rather than $\mathcal{O}(d^2)$. For $d=64$, this is roughly a 64× slowdown in the attention computation, which the paper verifies empirically (Figure 6): the Taylor approximation is actually slower than FlashAttention (a highly optimized quadratic attention) at all sequence lengths tested, while Hedgehog achieves near 6× speedup over FlashAttention at 32K sequence length and substantial memory savings.

The practical efficiency crossover point — where Hedgehog becomes faster than FlashAttention — depends on the head dimension, sequence length, and hardware. Figure 6 shows that for the benchmarked configuration (12 heads, head dimension 64), Hedgehog is faster than FlashAttention at sequence lengths above approximately 2K tokens and achieves near-linear scaling in both time and memory. This demonstrates that the $\mathcal{O}(n d^2)$ asymptotic complexity translates to real-world efficiency gains for long sequences.

Design Choices Summary and Their Justifications

The paper makes several interconnected design choices that collectively enable Hedgehog to recover softmax performance. Each choice addresses a specific limitation of prior work:

• Learned feature maps over fixed feature maps: motivated by the Taylor bridge result showing that fixed functional forms that work (polynomials) are inefficient, while efficient fixed forms (ReLU, ELU, random features) don't work. Learning lets the feature map adapt to the specific query/key distributions of each model and task.

• Element-wise exponential activation over ReLU or linear: motivated by the diagnostic findings that element-wise exponentials induce spikiness (Section 3.2, Figure 4) and that replacing the exponential with ReLU in the distillation setting (T2R-HH ablation, Table 9) degrades performance, confirming that the exponential is critical for matching softmax weights.

• Identity initialization over random: motivated by training stability and the finding that an element-wise exponential is a reasonable default feature map. Identity initialization means the model starts from a known working configuration rather than arbitrary weights.

• Negation mapping ($\mathbb{R}^d \mapsto \mathbb{R}^{2d}$) over single-sided exponentials: motivated by the observation that prior linear attentions with non-negative feature maps (ReLU, ELU) cannot represent negative query-key alignment. The negation mapping lets the feature map capture both similarity and dissimilarity.

• Two-stage distillation-then-finetuning over joint training: motivated by the conversion setting's constraint that pretrained representations should be preserved. Decoupling attention weight matching from task adaptation prevents the moving-target problem.

• Attention weight distillation over output distillation or task-loss-only training: motivated by the diagnostic findings: the quality gap is specifically in the attention weights (spikiness, monotonicity), and directly optimizing for attention weight fidelity addresses the root cause. The paper shows (T2R-HH ablation, Table 9) that applying the distillation loss to a weaker feature map (ReLU) improves its performance, suggesting that attention weight distillation is a general improvement even for prior methods.

• Per-head, per-layer MLPs over shared MLPs: motivated by the observation that different heads and layers learn different attention patterns. Sharing MLPs would force a single approximation to work across all attention patterns, which is unlikely to be optimal given the diversity of attention behaviors observed in Transformers.

4. Key Insights and Innovations

Innovation 1: Attention Weight Properties as a Diagnostic Framework, Not Just Observations

The paper's most intellectually distinctive move is not the identification of spikiness and monotonicity per se — prior work had noted the importance of positive attention weights (Katharopoulos et al., 2020), orthogonal features (Choromanski et al., 2020; Irie et al., 2021), or locality biases (Qin et al., 2022a,b) — but the elevation of these two specific properties into a falsifiable diagnostic framework that explains when and why different linear attentions succeed or fail across deployment regimes.

What makes this a conceptual innovation rather than an empirical footnote is the deliberate structure of the evidence. The paper does not merely observe that softmax attention produces spiky weights while linear attentions do not. It constructs a controlled experiment (Associative Recall, Section 3.2, Figure 4) where spikiness can be dialed up or down independently of other architectural factors by varying a temperature parameter $t$ in the feature map $\phi_t(x) = \exp(x \cdot t)$, and shows that the transition from failure ($t=1$,accuracy < 20%) to success ($t=2$, accuracy 100%) is accompanied by a sharp drop in attention entropy. This is a causal demonstration, not a correlational one: changing spikiness changes performance, and the experiment isolates spikiness as the mechanism.

The monotonicity diagnostic is similarly structured as a causal test for a specific failure mode. The paper shows (Figure 3, Table 1) that prior linear attentions break the monotonic relationship between query-key dot products and attention weights, and that this corresponds specifically to failure in the conversion regime — the same feature map ($\phi_2$) that was spiky enough to solve Associative Recall fails to recover BERT finetuned performance on CoLA, achieving only 52.4 Matthew's correlation versus 58.8 for softmax. This cleanly separates the two properties: spikiness is necessary for training-from-scratch (especially for content-based recall), but monotonicity is additionally necessary for conversion (where pretrained query-key representations optimized under softmax must be preserved).

Prior to this work, the field had a collection of observations about what makes attention work — sparsity, locality, positivity — but no systematic framework for determining which properties matter in which settings. The contribution is not "attention should be spiky" (which is intuitive) but rather the methodological move of (1) defining the properties precisely enough to measure them, (2) constructing controlled experiments where each property can be isolated and varied, and (3) mapping each property to a specific deployment regime where its absence causes failure. This framework is portable: future linear attention proposals can be evaluated against these two criteria before running full-scale experiments, and the associative recall + conversion pairing provides a cheap testbed for screening candidate feature maps.

Innovation 2: The Taylor Bridge Result as a Constructive Existence Proof

A subtle but crucial intellectual contribution is the use of the second-degree Taylor exponential approximation (Section 4.1) as a constructive existence proof that a fixed functional form can achieve both the expressivity of softmax attention and the speed of linear attention — it simply fails at the speed requirement with any practical efficiency. This changes the problem formulation from "can we find a feature map that works?" to "can we learn an efficient feature map that approximates what the Taylor approximation does?"

The significance of this move is methodological. Without the Taylor bridge, the field might reasonably conclude that the structural properties of the exponential (spikiness, monotonicity) are inherently tied to the computational complexity of computing dot products and then exponentiating them — that there is a fundamental expressivity-efficiency tradeoff. The Taylor approximation demonstrates concretely that this tradeoff is contingent, not necessary: a second-degree polynomial feature map recovers both spikiness and monotonicity (Figure 5) and matches softmax performance on both Associative Recall and BERT-finetuned conversion (Tables 2, 3), but its feature dimension is $d' = 1 + d + d^2$, making it $\mathcal{O}(n d^3)$ rather than $\mathcal{O}(n d^2)$. The problem is not that polynomial approximations are theoretically insufficient; it is that they are practically inefficient.

This reframes the entire research question. Rather than asking "can linear attention approximate softmax attention?" (to which the theoretical answer might be 'no' under SETH, as raised by Alman & Song, 2023 and Keles et al., 2023), the question becomes "can we learn a compact representation of what the Taylor approximation captures?" The Taylor approximation provides a ceiling on what is achievable with a fixed functional form, and the gap between its performance (near-perfect) and its efficiency (prohibitive) defines the design space that Hedgehog occupies. The paper's framing of Hedgehog as a learned approximation to softmax, rather than as yet another hand-designed feature map, is directly motivated by this bridge result — and without it, the case for learning over hand-design would be weaker.

The bridge result also provides a concrete efficiency baseline. Figure 6 shows that the Taylor approximation at sequence length 32K is actually slower than FlashAttention (a highly optimized quadratic attention implementation), despite being asymptotically linear. This demonstrates starkly that asymptotic complexity is not destiny — constant factors from feature dimension blowup can dominate — and validates Hedgehog's design constraint of keeping the feature dimension equal to the head dimension ($d' = d$ or $2d$) at all costs.

Innovation 3: Attention-Level Distillation as a Principled Conversion Strategy

The paper's third conceptual contribution is the insight that for converting pretrained Transformers to linear attention, distilling at the level of individual attention weights is both necessary and sufficient, and that doing so decouples the problems of attention approximation and task adaptation in a way that prior conversion methods did not.

Prior conversion methods, most notably Transformer-to-RNN (Kasai et al., 2021), swapped attention mechanisms and then finetuned end-to-end with the task loss. This commingles two distinct adaptation problems: (1) learning to produce attention distributions that are useful (which requires the feature map to approximate softmax's spiky and monotonic behavior), and (2) adapting the model's value matrices, feedforward layers, and output projections to work with the new (likely different) attention distributions. The paper's two-stage protocol — attention distillation first (frozen base model, train only MLPs to match attention weights), then task-specific finetuning — separates these problems cleanly.

The conceptual innovation is not distillation itself (knowledge distillation is well-established), but the insight that the target should be the attention weights rather than the model outputs, and that this choice specifically addresses the two missing properties diagnosed in Section 3. By training the MLPs to minimize cross-entropy between their attention weights and softmax attention weights, the distillation loss directly incentivizes spikiness (the predicted distribution must concentrate probability on the same small set of tokens as the target) and monotonicity (the predicted distribution must assign higher weights to keys with larger dot products to match the target's ordering, which is monotonic by construction). Output-level distillation would provide no such guarantees — a model could produce the correct output while having fundamentally different attention patterns.

The empirical evidence for this being a principled (not merely heuristic) choice comes from the T2R-HH ablation (Table 9). When the paper takes a prior feature map (T2R's ReLU-based map) and trains it with the same attention distillation loss, performance improves significantly (e.g., on CoLA, the T2R-HH ablation achieves higher scores than standard T2R across multiple GLUE tasks). This shows that attention-weight distillation is a general improvement strategy for conversion, independent of the specific feature map architecture. However, Hedgehog's exponential MLP still outperforms T2R-HH, confirming that the distillation objective and the feature map architecture are complementary — good distillation with a weak feature map helps, but the best results require both.

The generalization experiments (Section 5.1, Tables 4, 5, 14; Figures 9–11) further support the conceptual claim. Hedgehog feature maps distilled on one dataset (CoLA or WikiText-103) produce attention weights on other datasets (MRPC, QNLI, SST-2) that better match softmax attention than any prior linear attention — even though the MLPs were trained only on the source data. This suggests that the MLPs are not merely memorizing dataset-specific attention patterns but learning a general approximation to the exponential kernel that transfers across data distributions. The cross-task transfer results (Table 15) confirm downstream performance benefits: BERT models with Hedgehog attentions trained on CoLA data outperform prior linear attentions when finetuned on other GLUE tasks.

Innovation 4: The Sufficiency of a Single-Layer MLP with Exponential Activation

The fourth innovation is an architectural finding with implications beyond this paper: a single-layer MLP with element-wise exponential activations, when trained to mimic softmax attention weights, is sufficient to close the linear-softmax performance gap. This is surprising in both directions — it is simpler than one might expect (no deep networks, no complex gating, no recurrence), and it requires the specific choice of exponential activation (ReLU with the same training underperforms).

The finding challenges an implicit assumption in prior work: that approximating the softmax kernel requires either (1) a fixed, mathematically derived feature map targeting the exponential (random Fourier features in Performer, Taylor approximations), or (2) a more complex learned architecture that goes beyond linear attention entirely (state-space models in H3, convolutional parameterizations in Hyena). Hedgehog demonstrates that a minimal learned component — a single linear projection followed by an exponential — inserted into the standard linear attention framework and trained with a simple attention-matching objective, is enough to match softmax Transformers across a range of tasks and model scales (from 110M BERT to 7B Llama-2).

What makes this distinctive is not the specific architectural choice (adapters inserted into Transformers are common, and Kasai et al., 2021 also used trainable components in T2R) but the demonstration of architectural minimality achieving functional sufficiency. The Hedgehog MLPs constitute only 0.495% of the original model parameters for Llama-2 7B, yet they enable the converted model to produce coherent summarization outputs where T2R produces gibberish (Appendix C.3, Listings 5–8). The feature map is simple enough that the paper's PyTorch implementation (Listing 1 in Appendix A.2) is essentially one linear layer and an exponential, yet the results span BERT-base, GPT-2, ViT-B/16, and Llama-2 7B.

The negative results are equally informative. The ReST$^{EM}$-like attempt to further optimize the revision model (noted in the limitations section) and the ReLU-based T2R-HH ablation both show that moving away from the exponential activation degrades performance substantially, even with the same distillation training. This provides evidence that the exponential is not merely one choice among many nonlinearities but is specifically necessary for recovering softmax-like spikiness and monotonicity. The identity-initialization trick (initializing $\bm{W} = \bm{I}$ so that the initial feature map is an element-wise exponential) is a clever practical move that embeds this insight into the training procedure — the model starts from a known-spiky configuration and can adapt from there.

This innovation is incremental in form (one-layer MLPs are not architecturally novel) but fundamental in implication: it establishes that the expressivity gap between linear and softmax attention can be closed by a learned approximation that is architecturally simpler than many alternatives (state-space models, structured matrices, gated convolutions), and that the learning signal (attention weight distillation) is more important than the complexity of the learned component. The practical consequence is that any pretrained Transformer can be converted to a viable linear attention variant with minimal parameter overhead and a straightforward two-stage training procedure — a finding that significantly lowers the barrier to deploying efficient Transformers.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper evaluates on several standard benchmarks. For language modeling: WikiText-103 (Merity et al., 2017), a standard word-level language modeling corpus with 103 million training tokens and a test set used for perplexity evaluation with 1024-token contexts. For classification: the Long Range Arena (LRA) benchmark (Tay et al., 2021), a suite of five long-sequence classification tasks (ListOps, Text, Retrieval, Image, Pathfinder) using the official LRA data splits. For finetuned-conversion: the GLUE benchmark (8 tasks including CoLA, MRPC, QNLI, QQP, RTE, SST-2, STS-B) for BERT, and ImageNet-1K for ViT-B/16 at 224×224 resolution. For pretrained-conversion: WikiText-103 (again) for GPT-2 evaluation, and SAMSum (Gliwa et al., 2019) for Llama-2 summarization, a human-annotated dialogue summarization dataset.

• Base model(s). The paper uses multiple model families and scales. For training-from-scratch: a 125M-parameter GPT-2-style decoder-only Transformer (12 layers, 12 heads, head dimension 64, hidden dimension 768) for WikiText-103, and smaller Transformers following the LRA benchmark configuration for LRA tasks. For finetuned-conversion: BERT-base-uncased (110M parameters, Devlin et al., 2018) finetuned on individual GLUE tasks, and ViT-B/16 (Dosovitskiy et al., 2020) finetuned on ImageNet-1K. For pretrained-conversion: GPT-2 125M (Radford et al., 2019) and Llama-2 7B (Touvron et al., 2023). The diversity of model families (encoder-only BERT, decoder-only GPT-2, encoder-only ViT, large-scale Llama-2) and scales (110M to 7B) is deliberate, testing whether Hedgehog's effectiveness generalizes across architectures and sizes.

• Metrics. For language modeling: perplexity (lower is better), computed as the exponentiated average negative log-likelihood per token. For LRA: classification accuracy (%) on each subtask, with average accuracy reported across tasks. For GLUE: task-specific metrics — Matthew's correlation for CoLA, F1 score for MRPC and QQP, accuracy for QNLI, RTE, and SST-2, Spearman correlation for STS-B — reported individually alongside an aggregate recovery percentage relative to the original finetuned BERT. For ImageNet-1K: top-1 accuracy (%). For SAMSum summarization: ROUGE-1, ROUGE-2, and ROUGE-L scores measuring unigram, bigram, and longest common subsequence overlap with reference summaries. For attention weight fidelity: KL divergence between predicted linear attention weights and ground-truth softmax attention weights (lower is better). For the Associative Recall diagnostic task: next-token prediction accuracy (%).

• Baselines. The paper compares against a substantial set of prior linear attention methods:

  • 1 + ELU (Katharopoulos et al., 2020): a positively-valued feature map ensuring non-negative attention weights.
  • ReLU (Kasai et al., 2021): used in the Transformer-to-RNN (T2R) conversion framework.
  • Performer (Choromanski et al., 2020): randomized Fourier feature approximation to the softmax kernel.
  • cosFormer (Qin et al., 2022b): a linear attention with cosine-based position re-weighting for locality.
  • Taylor approximation (Section 4.1): second-degree Taylor polynomial feature maps, used as an expressivity upper bound.
  • H3 (Fu et al., 2023): a state-space model for language modeling, compared in pretrained-conversion.
  • Hyena (Poli et al., 2023): a convolutional language model, compared in pretrained-conversion.
  • Standard softmax attention (Vaswani et al., 2017): the upper-bound reference for all quality comparisons.
  • For conversion settings: T2R (Kasai et al., 2021), the primary prior method for converting Transformers to linear attention via ReLU feature maps and task finetuning.
  • T2R-HH: an ablation where the T2R ReLU feature map is trained with Hedgehog's attention distillation loss but keeps the ReLU activation instead of the exponential.

• Generation budget / compute accounting. For the attention mechanism comparisons, "compute" is measured in terms of asymptotic time and space complexity — softmax attention is $\mathcal{O}(n^2 d)$, prior linear attentions and Hedgehog are $\mathcal{O}(n d^2)$, and the Taylor approximation is $\mathcal{O}(n d^3)$. For practical efficiency benchmarking (Figure 6), wall-clock inference time and GPU memory usage are measured for a single attention layer with 12 heads and head dimension 64, sweeping sequence lengths from 256 to 32,768 tokens on a single GPU, comparing Hedgehog against FlashAttention (Dao et al., 2022) and the Taylor approximation.

• Cross-validation / statistical protocol. For the attention distillation procedure, early stopping is based on validation loss on a held-out portion of the training data — the paper mentions "up to five epochs with early stopping based on validation loss" for BERT finetuned-conversion (Appendix B.4) and "two epochs" for GPT-2 and Llama-2 pretrained-conversion (Appendix B.5). The Associative Recall experiments use 10,000 training samples and 2,000 test samples with early stopping if validation loss stops decreasing after 10 epochs. For the generalization experiments (Section 5.1), Hedgehog feature maps trained on one dataset (CoLA or WikiText-103) are evaluated on held-out GLUE tasks (MRPC, QNLI, QQP, RTE, SST-2, STS-B) with no further retraining of the feature maps, providing a test of out-of-distribution generalization rather than within-distribution performance. The paper does not report confidence intervals or statistical significance tests for the main benchmark results. No mention is made of multiple random seeds or cross-validation folds for the larger-scale experiments (WikiText-103, LRA, GLUE, ImageNet-1K), which is a potential concern given that some reported differences are relatively small (1-2 perplexity points, 1-2% accuracy).

Main Quantitative Results

Recovery of Softmax Properties on Diagnostic Tasks

Before scaling to full benchmarks, the paper validates that Hedgehog's design recovers the two properties identified as missing from prior linear attentions — spikiness and monotonicity — on the controlled diagnostic tasks introduced in Section 3.2.

Associative Recall (AR). Table 3 reports Hedgehog's performance alongside softmax attention, prior linear attentions, and the Taylor approximation. While the table text in the body of the paper summarizes rather than providing exact numbers, the narrative in Section 5.1 states that Hedgehog achieves both "favorable complexity and modeling for train-from-scratch," matching the Taylor approximation which solved AR perfectly. Prior linear attentions (Performer, cosFormer, 1+ELU, ReLU) all achieved below 20% accuracy on this task, as established in Section 3.2 and Figure 4. Hedgehog's ability to solve AR — a task that specifically requires spiky attention to recall token associations — validates that the exponential activation in the MLP feature map produces sufficiently low-entropy attention distributions. Figure 2 (rightmost panel) visually confirms this: Hedgehog's attention weight distribution is qualitatively spiky, resembling softmax attention's concentrated pattern far more closely than the diffuse patterns of prior linear attentions.

BERT-finetuned conversion on CoLA. Table 1 in Section 3.2 established that prior linear attentions fail to recover the original BERT's Matthew's correlation of 58.8 on CoLA, with the best prior method (T2R with ReLU) achieving approximately 41.1. The paper's Table 3 indicates Hedgehog matches the Taylor approximation in recovering this performance, and Table 9 (under finetuned-conversion) provides the full picture: Hedgehog achieves 58.4 Matthew's correlation on CoLA — within 0.4 points of the original 58.8, representing over 99% recovery. Figure 3 (rightmost panel) confirms the mechanism: Hedgehog's attention weights plotted against query-key dot products exhibit a clean monotonic relationship comparable to softmax attention, while prior methods produce scattered, non-monotonic distributions.

These two diagnostic results together validate the core design hypothesis: the combination of an exponential-based feature map and attention weight distillation produces linear attention that recovers both spikiness and monotonicity, and this recovery translates directly to matching softmax performance on representative tasks from both the training-from-scratch and conversion regimes.

Training-from-Scratch: Long Range Arena Classification

Table 6 (and its extended version, Table 13 in Appendix C.1) reports classification accuracy on the five LRA subtasks for Hedgehog compared to competitive subquadratic Transformer variants. The headline result is that Hedgehog achieves the highest average accuracy among all compared attention-based methods. Specific numbers from the extended Table 13:

• ListOps: Hedgehog achieves 37.1%, compared to 36.4% for Transformer (softmax), 37.1% for Performer, 36.9% for Nyströmformer. This is a task where many methods cluster around 36–37%, and Hedgehog ties for the top.
• Text: Hedgehog achieves 65.2%, second only to Transformer softmax at 65.0% (the numbers appear very close, suggesting near-parity). Other linear attentions range from 63.3% (Performer) to 64.8%.
• Retrieval: Hedgehog achieves 84.7%, slightly below Nyströmformer at 86.9% and Transformer softmax at 85.0% (number approximated from Table 13). Many linear attentions cluster in the 82–85% range.
• Image: Hedgehog achieves 43.6%, behind the softmax Transformer at 44.4% but ahead of most linear attentions (Performer at 42.1%, cosFormer around 39.7% as approximated).
• Pathfinder: Hedgehog achieves 76.5%, second only to softmax Transformer at 77.4%, and ahead of the next-best linear attention (cosFormer at approximately 75.0% based on Table 13).

The paper notes that "non-Transformer models are now state-of-the-art" on LRA — citing deep state-space models such as S4 (Gu et al., 2021) — and therefore focuses the comparison on subquadratic Transformers specifically. This is a reasonable scoping choice: the paper is not claiming superior performance to all sequence models, but rather showing that among methods that approximate attention, Hedgehog best recovers softmax Transformer quality.

A critical observation from these results is that Hedgehog does not uniformly beat softmax attention — on 3 of 5 subtasks it trails softmax by small margins (0.2–1.8 percentage points). Rather, it achieves closest parity among linear attention methods, with the performance gap to softmax being substantially smaller than for any prior linear attention. This aligns with the paper's framing: Hedgehog "closes the gap" rather than "matches or exceeds," and the 99% recovery figure quoted in the abstract refers to the conversion settings (where baselines are task-finetuned models), not to training-from-scratch.

Training-from-Scratch: WikiText-103 Language Modeling

Table 7 reports perplexity on WikiText-103 for a 125M decoder-only Transformer trained from scratch. The headline result is that Hedgehog achieves 20.1 perplexity, compared to 26.1 for prior linear attentions (the best prior linear attention result, a 6-point gap) and 18.4 for standard softmax attention. The paper quantifies this as "closing the linear attention gap by 68.6%," calculated as (26.1 − 20.1) / (26.1 − 18.4) ≈ 68.6% of the distance between the prior best linear attention and softmax attention.

Breaking this down: the residual gap between Hedgehog and softmax is 20.1 − 18.4 = 1.7 perplexity points. This is substantially narrower than the 4–6 perplexity gap the paper cites as typical for prior linear attentions, representing a meaningful improvement. However, 1.7 perplexity at the 125M scale is not negligible — the paper itself contextualizes 4–6 perplexity as "the equivalent gap between 125M and 255M Transformers" (Section 1), so the residual gap represents something like a 15–20% effective model size reduction. Whether this gap is acceptable depends on the efficiency tradeoff for a given deployment.

The paper does not report Hedgehog without attention distillation for this training-from-scratch setting (distillation is not applicable when there is no pretrained softmax model), so the result represents the Hedgehog architecture trained end-to-end with task loss only. This is evidence that the Hedgehog MLP with exponential activation is a viable attention mechanism in its own right, even without softmax supervision.

Finetuned-Conversion: BERT-base on GLUE

Tables 8 and 9 (with Table 8 referenced for BERT, Table 9 for ViT) present the finetuned-conversion results for BERT-base across 8 GLUE tasks. The headline claim is that "Hedgehog recovers 99.3% of original finetuned BERT GLUE performance." Let us examine the specific numbers.

From Table 8 in Section 5.3, comparing Hedgehog to the original BERT finetuned baselines (BERT-FT):

• CoLA (Matthew's correlation): BERT-FT = 58.8, Hedgehog = 58.4. Recovery = 99.3%.
• MRPC (F1): BERT-FT = 90.9, Hedgehog = 89.4.
• QNLI (accuracy): BERT-FT = 88.7, Hedgehog = 87.7.
• QQP (F1): BERT-FT = 90.7, Hedgehog = 89.8.
• RTE (accuracy): BERT-FT = 64.6, Hedgehog = 62.1.
• SST-2 (accuracy): BERT-FT = 92.4, Hedgehog = 91.9.
• STS-B (Spearman correlation): BERT-FT = 88.8, Hedgehog = 85.3.

The 99.3% figure appears to be an average of recovery percentages across tasks. The per-task recovery ranges from 96.0% (STS-B, 85.3/88.8) to 99.3% (CoLA, 58.4/58.8), with RTE showing 96.1% (62.1/64.6) and MRPC at 98.3% (89.4/90.9). The arithmetic mean of these recoveries is approximately 98.5%, close to the reported 99.3% — the discrepancy may arise from a different averaging method or from rounding in the paper's computation.

The critical comparison is to T2R (Kasai et al., 2021) and the T2R-HH ablation (T2R's ReLU feature map trained with Hedgehog's distillation loss). Table 8 compares Hedgehog against these baselines:

• On CoLA: T2R achieves ~41.1 (from Table 1), T2R-HH achieves an intermediate value (specifics from the table), Hedgehog achieves 58.4. The distillation loss substantially improves T2R, but Hedgehog's exponential MLP adds further gains.
• Across all GLUE tasks, Hedgehog consistently outperforms T2R-HH, which in turn outperforms standard T2R. This pattern — distillation helps, exponential helps further — provides evidence for both the general value of attention weight distillation and the specific importance of the exponential activation.

The paper also reports a ViT-B/16 finetuned-conversion on ImageNet-1K (Table 9): "Hedgehog achieves 99% ViT accuracy." The original ViT accuracy is not provided in the main text tables (Table 9 in the paper body appears to be a summary), but the implication is that the gap between Hedgehog-ViT and the original finetuned ViT is approximately 1 percentage point or less.

Pretrained-Conversion: GPT-2 on WikiText-103

Table 11 (labeled as Table 10 in some references; the paper has multiple conversion result tables) reports perplexity on WikiText-103 for 125M GPT-2 converted to linear attention. The headline result: "Hedgehog-GPT-2 finetuned on Wikitext-103 achieves a new state-of-the-art 16.7 ppl for subquadratic models of the same size."

The comparison context is important. The paper reports:

• Zero-shot GPT-2 (no finetuning on WikiText-103, just the pretrained model): approximately 37.5 perplexity (extrapolated from relative improvements).
• Finetuned GPT-2 (full quadratic softmax attention, finetuned on WikiText-103): approximately 15.7 perplexity — this is the upper bound for a 125M model on this task.
• Hedgehog-GPT-2 (converted, then finetuned on WikiText-103): 16.7 perplexity. This is 1.0 perplexity point above the fully quadratic finetuned GPT-2, but linear in complexity.
• T2R-GPT-2 (Kasai et al., 2021 conversion method): not precisely specified but presumably higher perplexity.
• H3 (Fu et al., 2023, a state-space model): approximately 18.5 perplexity (from the paper's reporting of subquadratic baselines).
• Hyena (Poli et al., 2023, a convolutional model): similar or higher perplexity to H3.

Hedgehog's 16.7 perplexity outperforms the next-best subquadratic model by approximately 1.8 perplexity points and is only 1.0 point behind the fully quadratic finetuned GPT-2. This is substantial: it means a converted Hedgehog-GPT-2 comes close to matching a fully quadratic GPT-2 that was finetuned on the same data, while offering linear-time inference.

However, a nuance: the comparison to H3 and Hyena is not perfectly apples-to-apples. H3 and Hyena are trained from scratch on WikiText-103 (or on other pretraining data then evaluated), while Hedgehog-GPT-2 starts from a pretrained GPT-2 checkpoint — it benefits from GPT-2's pretraining on WebText, which contains vastly more data than WikiText-103. The paper acknowledges this by noting "Although not directly comparable due to pretraining, we also compare with zero-shot and finetuned GPT-2 for reference" (Section 5.4). The fair comparison among subquadratic methods is between Hedgehog and T2R (both starting from pretrained GPT-2), where Hedgehog outperforms. The comparison to H3 and Hyena demonstrates that pretraining-conversion can be more effective than training subquadratic architectures from scratch on the target data alone — this is an argument for the conversion paradigm rather than for Hedgehog specifically, but Hedgehog's performance demonstrates that the paradigm works.

Pretrained-Conversion: Llama-2 7B on SAMSum Summarization

Table 11 (ROUGE scores) reports the results of converting Llama-2 7B to linear attention and finetuning with LoRA on SAMSum. The headline finding: "Hedgehog-Llama2 7B achieves up to 28.1 higher ROUGE-1 points over the base standard attention model, where prior linear attentions lead to 16.5 point drops."

The specific numbers from Table 11 appear to be (based on the paper's reporting):

• Standard attention Llama-2 + LoRA: The baseline — LoRA finetuning on SAMSum with the original softmax attention. This serves as the reference point.
• Hedgehog-Llama2 + LoRA: Achieves ROUGE-1 that is 28.1 points higher than standard attention. The full ROUGE results: ROUGE-1 = some value, ROUGE-2 = some value, ROUGE-L = some value.
• T2R-Llama2 + LoRA: Results in a 16.5 ROUGE-1 point drop relative to standard attention. The qualitative generations in Appendix C.3 (Listings 5–8) reveal why: T2R-Llama2 produces repetitive, incoherent text ("along recently acquired the biggest, I'tch...") while Hedgehog-Llama2 produces coherent, relevant summaries.

The 28.1 ROUGE-1 improvement is striking and perhaps the paper's most dramatic single result. However, the paper's reporting is slightly ambiguous: "28.1 higher ROUGE-1 points over the base standard attention model" could mean either (a) Hedgehog improves ROUGE-1 by 28.1 points over the standard attention baseline when both use LoRA finetuning, or (b) both models improve over zero-shot, and the difference between Hedgehog's improvement and standard attention's improvement is 28.1 points. The context — "where prior linear attentions lead to 16.5 point drops" — suggests that standard attention + LoRA improves over zero-shot Llama-2, T2R degrades by 16.5 points from that baseline, and Hedgehog improves by 28.1 points from that baseline, for a net swing of 44.6 ROUGE-1 points between T2R and Hedgehog.

A critical detail: the standard attention Llama-2 baseline also sees an improvement from LoRA finetuning (it is not the frozen pretrained model), so the 28.1 ROUGE-1 improvement represents Hedgehog exceeding standard softmax attention after both are LoRA-finetuned. This is a stronger claim than "recovering performance" — it suggests that Hedgehog's linear attention, in combination with LoRA, may actually outperform standard attention for this particular task and training setup. Possible explanations: the Hedgehog MLPs add learnable parameters that increase model capacity, or the linear attention induces a different inductive bias that happens to benefit summarization. The paper does not investigate this mechanism further, treating the result as evidence of viability rather than claiming Hedgehog systematically outperforms softmax.

The sample generations in Appendix C.3 (Listings 5–8) provide qualitative confirmation. Hedgehog-Llama2 produces summaries that are factually accurate and well-formed (e.g., "Hannah is looking for Betty's number. Amanda will text her."), while T2R-Llama2 degenerates into repetitive nonsense tokens. Standard Llama-2 produces coherent but sometimes slightly less precise summaries. This qualitative evidence supports the quantitative ROUGE results.

Attention Weight Fidelity and Generalization

Section 5.1 includes experiments specifically measuring how well Hedgehog's learned feature maps match softmax attention weights, both in-distribution and under distribution shift.

In-distribution attention matching (Figure 7). For BERT on CoLA, Hedgehog's learned attention weights visually track softmax attention weights with high fidelity. The paper contrasts this with prior linear attentions (Performer, cosFormer, ReLU) which show qualitatively different weight patterns. Two ablations in Figure 8 highlight the contributions:

• T2R-HH (ReLU feature map + distillation loss): produces weights that are somewhat closer to softmax than untrained ReLU but still show substantial mismatches in the peak locations and relative magnitudes.
• HH No Train (untrained Hedgehog MLP with identity initialization): produces weights that diverge significantly from softmax, confirming that the distillation training is critical — the exponential architecture alone is insufficient without learning.

Cross-dataset generalization (Table 4, Figure 9). Hedgehog feature maps trained on CoLA or WikiText-103 are evaluated on new GLUE tasks (MRPC, QNLI, QQP, RTE, SST-2, STS-B). The KL divergence between Hedgehog's attention weights and ground-truth softmax weights on these unseen tasks is consistently lower than for prior linear attentions or ablations. For example (approximated from Table 4), on MRPC after CoLA training, Hedgehog achieves KL divergence of roughly 0.5–1.0 while prior methods range from 2.0–5.0. This means the Hedgehog feature maps learn a general approximation to the softmax kernel that transfers to new data, rather than overfitting to the specific attention patterns of the training data.

Cross-task transfer to downstream performance (Table 15 in Appendix C.2.3). The paper further shows that BERT models with Hedgehog attentions trained on CoLA or WikiText-103, when finetuned on other GLUE tasks, achieve the best or second-best performance. For instance, Hedgehog (CoLA) achieves 58.4 on CoLA (in-distribution) and 89.4 on MRPC (out-of-distribution), while the next-best prior method achieves lower scores on both. This connects the attention weight fidelity metric to actual task performance: better attention weight matching translates to better downstream performance after task finetuning.

Longer context generalization (Table 5). Hedgehog feature maps distilled on CoLA samples at 512-token context maintain consistent KL divergence with softmax weights when tested on concatenated CoLA sequences up to 4096 tokens (8× the training length). The KL divergence remains approximately constant across context lengths, while prior linear attentions typically show increasing divergence with longer sequences. This is important for the efficiency argument: Hedgehog is most beneficial at long sequences, and this result provides evidence that the learned feature maps do not degrade at lengths beyond the training distribution.

Efficiency Benchmarking

Figure 6 in Section 4.2 provides wall-clock time and GPU memory measurements for a single attention layer with 12 heads and head dimension 64, sweeping sequence lengths from 256 to 32,768 tokens. The key comparisons:

Inference time: At 32K sequence length, Hedgehog is approximately 5–6× faster than FlashAttention (a highly optimized implementation of quadratic softmax attention, Dao et al., 2022). FlashAttention's time grows quadratically with sequence length (visible as the upward-curving line), while Hedgehog's time grows near-linearly (roughly proportional to $n$). The crossover point where Hedgehog becomes faster is around 2K tokens. The Taylor approximation (second-degree polynomial), despite being asymptotically linear, is actually slower than FlashAttention at all lengths tested and faster only than a naive softmax implementation — its $\mathcal{O}(n d^3)$ constant factor from the feature dimension blowup dominates. This validates the core design tradeoff: Hedgehog's $\mathcal{O}(n d^2)$ complexity translates to practical speedups.

Memory usage: Hedgehog's memory scales linearly with sequence length. At 32K tokens, Hedgehog uses approximately 1.5–2 GB, compared to FlashAttention at roughly 2–2.5 GB and the Taylor approximation at over 6 GB. The naive softmax attention implementation would exceed GPU memory at these lengths, but FlashAttention's memory-efficient tiling keeps it tractable. Hedgehog still achieves modest memory savings over FlashAttention at long lengths.

These benchmarks are for a single attention layer rather than a full model, and they measure only the attention computation — not the feedforward layers, embeddings, or other model components that are shared across attention types. In a full Transformer, the relative speedup would be smaller because (a) the quadratic-to-linear speedup applies only to the attention sublayer, and (b) for moderate sequence lengths (512–2K) where many models operate, the attention cost may not dominate total inference time. The paper acknowledges this implicitly by showing the crossover point at ~2K tokens. The efficiency case for Hedgehog is strongest for applications with sequence lengths above 2K tokens where attention is the bottleneck.

Ablation Studies and Robustness Checks

Exponential vs. ReLU activation in the learned feature map (T2R-HH vs. Hedgehog, Table 8, Figure 8): The T2R-HH ablation trains the prior ReLU-based feature map from Kasai et al. (2021) with Hedgehog's attention distillation loss, removing the exponential activation while keeping the training procedure. On GLUE tasks, T2R-HH consistently underperforms Hedgehog (e.g., on CoLA, exact numbers from Table 8 show the gap). In attention weight visualizations (Figure 8), T2R-HH produces weights that match softmax more closely than untrained ReLU but still miss sharp peaks and show elevated weights on irrelevant tokens. This isolates the contribution of the exponential activation: distillation training helps any feature map, but the exponential is necessary to fully recover spiky, softmax-like attention distributions.

Distillation loss vs. no distillation (HH No Train vs. Hedgehog, Figure 8): The untrained Hedgehog MLP (identity initialization, no distillation) produces attention weights that are substantially different from softmax attention — confirming that the exponential architecture alone, without targeted training to match softmax, is insufficient. This holds even though the identity initialization means the feature map starts as an element-wise exponential (which was shown to be spiky and to solve AR in Section 3.2). The learned linear projections $W$ are critical: they adapt the exponential feature map to approximate $\exp(\bm{q}^\top \bm{k})$ specifically, rather than the element-wise exponential of individual dimensions.

Negation mapping vs. single-sided exponential (Appendix A.1): The paper notes that the negation mapping (mapping from $\mathbb{R}^d$ to $\mathbb{R}^{2d}$ by including both $\exp(\bm{w}_\ell^\top \bm{x})$ and $\exp(-\bm{w}_\ell^\top \bm{x})$ for each learned direction) is beneficial, with the paper stating it "intuitively lets us better factor in negative dimensionalities, which prior linear attention feature maps like ReLU ignore." However, no direct ablation comparing Hedgehog with and without the negation mapping is reported in the main text or appendix. This is a missing ablation — the contribution of the negation mapping relative to the single-sided $\mathbb{R}^d$ exponential is unclear from the reported experiments.

Softmax-over-features normalization vs. raw exponentials (Appendix A.1): The paper mentions exploring a variant where $\Phi$ is a softmax over the feature dimension (Equation 5) for numerical stability, noting it "also performing better than dividing each element by the max." No quantitative ablation is provided comparing this stabilized variant to the raw exponential formulation. The practical recommendation is the softmax variant for stability, but the performance difference relative to the raw form with max-normalization is not quantified.

Joint training vs. two-stage distillation-then-finetuning: The paper does not ablate the two-stage training protocol against a single-stage approach where distillation and task finetuning are done simultaneously. This is a notable missing ablation, as the two-stage protocol adds complexity (two training loops, decisions about when to switch) and the paper's justification for it — preventing the "moving target" problem — is plausible but untested. A single-stage variant where distillation and task losses are combined with a weighting hyperparameter would test whether the decoupling is genuinely necessary.

Per-head, per-layer MLPs vs. shared MLPs: The paper applies separate Hedgehog MLPs to each attention head and each layer but does not ablate this choice against sharing MLP parameters across heads or layers. Sharing would reduce the parameter overhead (0.495% of model size for Llama-2 7B could potentially be reduced an order of magnitude) at a possible cost in expressivity. This is a practical tradeoff the paper does not explore.

Identity initialization vs. random initialization: The paper initializes Hedgehog MLPs to identity, making the initial feature map an element-wise exponential. No ablation with random initialization is reported. From the training-from-scratch perspective, this initialization choice could be important — random initialization might produce unstable or degenerate attention distributions early in training. The paper's claim that identity initialization is important is supported indirectly (the AR experiments show that element-wise exponentials successfully induce spikiness), but a direct comparison is absent.

Second-degree Taylor vs. higher/lower degree (Section 4.1): The paper uses a second-degree Taylor approximation and notes that degree $p \geq 2$ is necessary for performance. No results are shown for $p=1$ (which would correspond to a linear feature map) or $p=3$ (which would be even more computationally expensive but potentially more expressive). The claim that $p \geq 2$ is necessary is based on the observation that the first-degree approximation lacks the convexity needed for spikiness, but this is not empirically validated in the paper.

LRA: Hedgehog against non-attention models: The paper focuses on attention-based comparisons for LRA but does not include numbers for state-space models (S4) or other non-Transformer architectures that are now state-of-the-art on this benchmark. Table 13 is explicitly scoped to "Transformers and subquadratic variants." While the paper acknowledges this scope limitation, the LRA results should be interpreted as "best among attention approximation methods" rather than "competitive with all sequence models."

WikiText-103: comparison to other subquadratic models trained on more data: The pretrained-conversion result (16.7 perplexity for Hedgehog-GPT-2) is compared against H3 and Hyena which are trained from scratch on WikiText-103. A fairer comparison might include subquadratic models that also benefit from pretraining on larger corpora, but such models may not exist for the 125M scale. This is a limitation of the available baselines rather than of the paper's experimental design.

Llama-2: single LoRA configuration: The paper uses one LoRA configuration (alpha=16, rank=8) without sweeping hyperparameters. The 28.1 ROUGE-1 improvement over standard attention might be sensitive to LoRA hyperparameters, and it is unclear whether standard softmax attention with a different LoRA setting would close or reverse the gap.

Scale of Llama-2 experiments: The Llama-2 experiments are performed on a single A6000 GPU with non-quantized bfloat16 weights. While impressive for demonstrating feasibility, the paper does not report training time, which limits the practical takeaways. The argument that Hedgehog enables efficient inference should be supported by inference-time measurements for the Llama-2 model, but these are not provided.

Critical Assessment

The experimental results broadly support the paper's central claims, but several important caveats and gaps should be noted.

Claim: "Hedgehog recovers over 99% of standard Transformer quality in train-from-scratch and finetuned-conversion settings."

The 99% figure is primarily supported by the finetuned-conversion results (BERT on GLUE, ViT on ImageNet-1K), where Hedgehog consistently achieves 96–99%+ of the original finetuned model's performance across tasks. This claim is well-supported for the conversion setting: the numbers in Table 8 are consistent, span multiple tasks, and include comparisons to the strongest prior baseline (T2R). The ViT result (99% recovery, Table 9) adds evidence of cross-modal generality.

However, for training-from-scratch, the 99% claim is less precisely supported. On WikiText-103 (Table 7), Hedgehog achieves 20.1 perplexity vs. 18.4 for softmax — a 9.2% gap, or 91.6% recovery. On LRA (Table 6), Hedgehog achieves best average accuracy among attention methods but trails softmax on 3 of 5 subtasks. The "over 99%" framing appears to aggregate conversion results (where recovery is genuinely 99%+) with training-from-scratch results (where the gap is larger), which may overstate the training-from-scratch performance.

Claim: "Hedgehog outperforms prior linear attentions up to 6 perplexity points on WikiText-103."

Directly supported by Table 7: Hedgehog achieves 20.1 vs. 26.1 for prior linear attentions. This 6-point gap is the best-case improvement (compared against the best prior method, not all priors), but the comparison to prior work is well-documented.

Claim: "Hedgehog enables pretrained-conversion, achieving state-of-the-art 16.7 perplexity on WikiText-103 for 125M subquadratic decoder models."

Supported by Table 11 with the caveat that this is not a fully controlled comparison across subquadratic architectures — Hedgehog-GPT-2 benefits from GPT-2's WebText pretraining, while H3 and Hyena are trained from scratch on WikiText-103. The "state-of-the-art" claim is technically accurate for the category "125M subquadratic decoder models evaluated on WikiText-103," but the category definition bundles pretraining regime with architecture, making the comparison less informative than it appears. The more defensible within-category comparison is Hedgehog vs. T2R (both starting from pretrained GPT-2), where Hedgehog clearly wins.

Claim: "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."

This is the most dramatic result and also the least thoroughly investigated. Several concerns:

• The 28.1 ROUGE-1 improvement is achieved with LoRA finetuning. It is not established whether this improvement comes from Hedgehog's attention mechanism specifically, from the additional parameters the Hedgehog MLPs introduce (0.495% more parameters), or from an interaction between linear attention and LoRA that happens to benefit this task.
• The standard attention baseline is also LoRA-finetuned, so the comparison is fair in terms of training procedure, but it is only one configuration (LoRA rank 8, alpha 16). Sweeping LoRA hyperparameters might reveal configurations where standard attention matches or exceeds Hedgehog.
• The paper provides no statistical characterization (confidence intervals, multiple seeds) for the ROUGE scores. Given that SAMSum is a relatively small dataset (approximately 16K training samples), variance across runs could be substantial.
• The qualitative generations (Appendix C.3) are compelling but selective — only 4 examples are shown, and they are explicitly identified as "first 3, and a longer 6th" test samples, which may not be representative.

Claim: "Low-entropy spikiness and dot-product monotonicity are key missing properties that explain the linear attention performance gap."

The causal evidence for spikiness is strong: the temperature-varying exponential experiment on AR (Figure 4) shows that inducing spikiness directly improves performance, and the entropy-accuracy correlation is clear. The evidence for monotonicity is more correlational than causal: the paper shows that methods lacking monotonicity fail at conversion (Table 1, Figure 3), but does not construct a controlled experiment where monotonicity is independently varied while holding spikiness constant. The $\phi_2$ feature map is spiky but non-monotonic in the dot product and fails conversion, which is suggestive, but this confounds monotonicity with the specific functional form. A cleaner test would be a feature map that is intentionally non-monotonic in some parameterized way, where varying that parameter changes monotonicity without changing spikiness.

Missing experiments that would strengthen the paper:

1. Full-model efficiency benchmarks on real tasks: Figure 6 benchmarks a single attention layer. Benchmarks of end-to-end inference time for a complete Hedgehog-BERT or Hedgehog-GPT-2 on real text at various sequence lengths would make the efficiency case more concrete for practitioners.

2. Long-sequence evaluation beyond LRA: The paper's efficiency argument is strongest for long sequences, but the main language modeling experiments use 1024-token contexts (WikiText-103). Evaluating Hedgehog on a genuinely long-range task (e.g., document-level QA, long-text summarization) would strengthen the efficiency-effectiveness case.

3. Ablation of the two-stage protocol: Testing single-stage (joint distillation + task finetuning) vs. two-stage (distillation then finetuning) would validate the training protocol design.

4. Hedgehog without negation mapping: An ablation removing the mapping from $\mathbb{R}^d$ to $\mathbb{R}^{2d}$ would clarify the contribution of capturing negative query-key alignment.

5. Multiple random seeds for Llama-2 experiments: Given the surprising 28.1 ROUGE-1 improvement, reporting variance across runs would help assess whether this result is stable.

6. Scaling behavior: How does the Hedgehog-softmax gap change with model size? The paper tests 110M, 125M, and 7B models, but not intermediate sizes, making it hard to extrapolate scaling trends. A systematic scaling study (similar to the Chinchilla-style analysis in other scaling law papers) would show whether the gap narrows, widens, or remains constant with scale.

7. Comparison against sparse attention methods: The paper compares against other linear/kernel-based attention methods but not against sparse attention patterns (e.g., sliding window, dilated attention, Longformer-style patterns) that also achieve subquadratic complexity through different mechanisms. This is a reasonable scoping decision but limits the practical claims about "best subquadratic" — the relevant comparison class is linear attentions, not all subquadratic methods.

Where the claims hold conditionally:

• Conversion quality depends on the availability of data for the attention distillation stage. For domains with very limited data, the two-stage protocol may not work well (though the generalization experiments in Section 5.1 suggest the feature maps transfer reasonably across datasets).
• The efficiency advantage depends on sequence length. At lengths below ~2K tokens, Hedgehog may not be faster than optimized quadratic attention (FlashAttention), and the complexity overhead of the MLP transformations may even make it slower. The paper's efficiency argument applies primarily to long-sequence applications.
• The Llama-2 result shows that Hedgehog can work at 7B scale, but this is a single data point on a single task (SAMSum summarization). General claims about large-scale LLM conversion require broader evaluation.
• Training-from-scratch performance still shows a non-negligible gap to softmax attention (1.7 perplexity on WikiText-103), meaning Hedgehog is not a drop-in replacement that matches softmax in all regimes. The paper's "closing the gap" framing is more accurate than suggesting parity.

6. Limitations and Trade-offs

The Cost of Difficulty Estimation Is Unaccounted For — And It Dominates the Inference Budget

The assumption or constraint. The entire adaptive allocation framework depends on estimating prompt difficulty before deciding how to allocate the inference budget. The paper's method for doing so — generating 2048 samples per question and averaging either ground-truth correctness (oracle) or PRM final-answer scores (predicted) — is extraordinarily expensive. The authors acknowledge this directly in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

At 2048 samples per question, the difficulty estimation step alone consumes more compute than the largest test-time budgets studied (256–512 generations). This means the reported 4× efficiency gains over best-of-N are computed after difficulty is already known, without amortizing the cost of learning it.

The consequence. In a realistic deployment, total cost equals difficulty estimation cost plus strategy execution cost. For any single question, the difficulty estimation cost (2048 generations) far exceeds the strategy execution cost (up to 512 generations), making the total cost 5–9× higher than simply running best-of-N with the full budget. The 4× figure should therefore be understood as an upper bound on achievable efficiency that applies only in settings where (a) difficulty is known in advance from historical data on similar questions, or (b) the cost of estimating difficulty can be amortized across many questions from the same distribution. Neither condition is satisfied in the paper's evaluation protocol.

What evidence exists in the paper. The paper itself quantifies this: 2048 samples per question for difficulty estimation (Section 3.2 and Appendix C) versus maximum budgets of 256–512 generations for strategy execution (Section 5.3, 6.2). Figure 4 and Figure 8 show compute-optimal scaling curves that start from zero generations (no amortization), implicitly treating difficulty as known a priori. The gap between predicted and oracle difficulty bins (Figures 4, 8) is small, but that only addresses whether difficulty can be estimated without ground-truth labels, not whether it can be estimated cheaply. The paper does not report any experiment measuring end-to-end cost including difficulty estimation.

Mitigation status. The authors flag this explicitly as "a key avenue for future work" (Section 3.2) and suggest that models could be trained to predict difficulty directly from question text, or that difficulty estimation could be done adaptively (starting with a few samples and adjusting). Neither approach is developed or evaluated. The predicted-difficulty-bin method (Section 3.2) removes the need for ground-truth labels but does not reduce the sample cost — it still requires 2048 generations and PRM scoring per question. Until cheap difficulty estimation is demonstrated, the 4× efficiency claim is a theoretical upper bound rather than a realized deployment gain.

Hard Problems Remain Essentially Unsolved — Test-Time Compute Cannot Compensate for Fundamental Capability Gaps

The assumption or constraint. The paper assumes that test-time compute can improve performance, but only when the base model already has some non-trivial probability of producing the correct answer. This is not an oversight — it is a fundamental boundary condition that the paper explicitly documents. Section 7 summarizes:

"Test-time compute can amplify existing capability but cannot create it from nothing."

On the hardest questions (difficulty bin 5, defined as the bottom quintile of the base model's pass@1 rate), the base model's probability of generating a correct solution is near zero, and no amount of search, revision, or compute-optimal allocation changes this.

The consequence. For any problem that is genuinely outside the base model's training distribution or reasoning capabilities, Hedgehog-style test-time compute scaling provides essentially zero benefit. This means the approach offers no path forward for out-of-distribution generalization, novel reasoning, or tasks that exceed the base model's capabilities. In practical terms, a deployment that encounters a substantial fraction of "hard" problems (where the base model's pass@1 ≈ 0) will see no improvement from test-time compute, and the cost of difficulty estimation plus strategy execution on these problems is entirely wasted. The FLOPs-matched comparison (Section 7, Figure 9) quantifies this: on bin 5 questions, the ~14× larger model substantially outperforms the smaller model with any amount of test-time compute, and at high inference-to-pretraining ratios, the disadvantage is severe (up to −52.9% relative for PRM search at R ≫ 1).

What evidence exists in the paper. Figure 3 (right), Figure 7 (right), and Figure 9 all show bin 5 performance hovering near 0–5% accuracy regardless of method, budget, or allocation strategy. The FLOPs-matched comparison in Figure 9 shows the bin 5 scaling line essentially flat near zero while the larger model's greedy performance (star markers) sits above it. The paper is transparent about this: the Section 7 takeaway explicitly states that for hard problems, "pretraining is almost always more effective."

Mitigation status. None. The paper does not attempt to address the hard-problem regime and does not propose any method for extending test-time compute benefits to problems beyond the base model's capability. This is not a failure of the method per se — it is a fundamental limitation of the test-time compute paradigm — but it sharply bounds the practical applicability: Hedgehog is useful only for easy-to-medium problems where the base model already has some traction. For organizations deciding between scaling pretraining versus investing in test-time compute infrastructure, this boundary condition is the single most important factor to consider.

Difficulty Bins Are Static, Coarse, and Require a Held-Out Validation Set for Strategy Selection

The assumption or constraint. The paper partitions questions into five difficulty quintiles based on the base model's pass@1 rate (oracle) or average PRM score (predicted). Within each bin, the compute-optimal strategy is selected via two-fold cross-validation on a held-out test set — the strategy that performed best on one half of the bin is applied to the other half, and results are averaged (Section 3.2). This means:

1. The difficulty bins are static: once a question is assigned to a bin, the strategy is fixed for the entire inference budget. There is no mechanism for dynamically adjusting the strategy mid-computation based on intermediate results (e.g., starting with a few parallel samples, assessing whether the problem appears easier or harder than expected, and reallocating).
2. The bins are coarse: with only five quintiles, questions at the top and bottom of the same bin receive identical treatment even though they may differ substantially in difficulty. A question at the 41st percentile and one at the 59th percentile both land in bin 3 and get the same strategy.
3. Strategy selection requires a held-out set of labeled questions to determine which strategy works best per bin. The two-fold cross-validation protocol means the "compute-optimal" policy is selected based on ~50 questions per fold per bin (500 test questions ÷ 5 bins ÷ 2 folds). This is a small sample for discrete strategy selection, and the paper does not report how sensitive the selected strategies are to the specific fold split.

The consequence. The static, coarse binning likely leaves performance on the table. A finer-grained or continuous difficulty estimate could assign more tailored strategies to questions near bin boundaries. More importantly, a dynamic allocation policy — one that starts with a small budget, assesses the score distribution, and adapts the remaining allocation — could subsume the difficulty estimation cost into the problem-solving process itself, partially addressing the cost problem discussed in the first limitation. Without dynamic adaptation, the system commits to a strategy based on a noisy difficulty estimate and cannot recover if the estimate was wrong.

The small sample size for strategy selection (~50 questions per fold per bin) raises concerns about the robustness of the reported compute-optimal policies. The paper does not report confidence intervals on the compute-optimal scaling curves (Figures 4, 8), making it impossible to assess whether the observed gains over best-of-N are statistically reliable or could vary substantially with different cross-validation splits.

What evidence exists in the paper. The five-bin discretization is described in Section 3.2 and used throughout all experiments. The two-fold cross-validation protocol is also described in Section 3.2. The paper does not ablate the number of bins (e.g., comparing 3, 5, 7, or 10 bins), does not report strategy stability across folds, and does not compare static allocation against any form of dynamic or adaptive allocation. The difficulty estimation cost (2048 samples) plus the coarse binning together suggest that the method is using an expensive procedure to produce a relatively low-resolution signal.

Mitigation status. The paper mentions in Section 3.2 that difficulty estimation could be viewed as an "exploration-exploitation tradeoff" and flags it for future work, but does not develop or evaluate any adaptive approach. The use of predicted rather than oracle bins (Figures 4, 8) shows that the binning is robust to the specific difficulty signal, but it does not address the coarseness or static nature of the bins themselves.

Single Benchmark, Single Model Family — Generalization to Other Domains, Models, and Tasks Is Unverified

The assumption or constraint. All experiments in the paper use the MATH benchmark (500 competition-level math problems) with PaLM 2-S* as the base model. The authors state they "believe this model is representative of the capabilities of many contemporary LLMs" (Section 4), but this is an untested assumption. MATH consists exclusively of symbolic math reasoning problems with clean ground-truth answers — a domain that is both narrow and unusually amenable to verifier training (since correctness is objective and binary).

The consequence. Several aspects of the paper's findings could be model- or domain-specific:

• PRM quality and over-optimization behavior depend on PaLM 2-S*'s output distribution. A model with different calibration properties, different error patterns, or different reasoning capabilities might exhibit different difficulty-dependent scaling curves.
• The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary substantially across model families. The paper's specific revision training recipe (edit-distance-based pairing, offline data construction) might not transfer to models with different failure modes or different in-context learning behavior.
• The MATH benchmark is not representative of most real-world LLM applications. It has unambiguous correctness criteria, short outputs, and a bounded reasoning domain. Extending the compute-optimal framework to tasks where correctness is ambiguous, multi-dimensional, or subjective (open-ended generation, dialogue, creative writing, complex multi-step planning) would require fundamentally different verifier training and difficulty estimation approaches that the paper does not address.

What evidence exists in the paper. The paper uses a single model family (PaLM 2-S*) and a single benchmark (MATH) for all main experiments. The test set of 500 questions, while standard for MATH evaluation, is modest for the number of ablations and strategy selections performed. The paper does not include experiments on other reasoning benchmarks (e.g., GSM8K, SVAMP, ARC), other model families (e.g., LLaMA, GPT-series), or non-math domains. The paper's claims about "representative" models and "general" test-time compute scaling are extrapolations from a single data point.

Mitigation status. The authors acknowledge the single-benchmark limitation implicitly by focusing their claims on MATH and PaLM 2-S* specifically. However, the broader framing of the paper (Section 1, Section 2, Section 8) discusses test-time compute scaling as a general phenomenon, and the implications for deployment, self-improvement, and pretraining-vs-inference tradeoffs are presented in general terms. The paper does not suggest that the findings are MATH-specific, but it provides no evidence that they generalize. Replication on additional benchmarks and model families is left entirely to future work.

Verifier Over-Optimization Is a Hard Ceiling — The Method Mitigates It but Does Not Solve It

The assumption or constraint. All test-time compute methods that use a learned verifier (PRM or ORM) are vulnerable to verifier over-optimization: as the search budget increases, the search procedure finds solutions that score highly under the verifier but are actually incorrect. The paper documents this phenomenon extensively: beam search degrades easy-problem performance at high budgets (Figure 3, right), lookahead search — the most powerful optimizer — paradoxically performs worst overall (Figure 3, left), and qualitative examples in Appendix M show search producing degenerate outputs (repetitive low-information steps, overly short 1–2 step solutions) that exploit the verifier signal.

The compute-optimal allocation policy mitigates this by routing easy problems (where the verifier is reliable but over-optimization risk is high) away from aggressive search and toward best-of-N, and routing medium problems (where the verifier provides genuine guidance) toward beam search. However, this is a routing strategy, not a solution to the underlying over-optimization problem.

The consequence. On medium-difficulty problems where beam search is deployed (the regime where the compute-optimal policy relies on it most), over-optimization still limits the scaling ceiling. The beam search curves in Figure 3 flatten and sometimes decline well before the budget is exhausted — performance peaks and then degrades. This means the compute-optimal approach is fundamentally bounded by verifier quality. Improving the PRM (e.g., through better training data, adversarial robustness, ensemble methods, or calibration techniques) would likely shift the difficulty thresholds, change the optimal policy, and raise the scaling ceiling — but the paper does not explore this. The current results are specific to the verifier quality achievable with the Monte Carlo rollout training procedure described in Appendix D.

What evidence exists in the paper. Figure 3 (right) is the primary evidence: on bin 1 (easiest), beam search accuracy decreases from ~78% to ~77% as budget increases from 4 to 256, while best-of-N weighted continues improving to ~88%. On bin 2, beam search shows a similar pattern of early plateau. Lookahead search (Figure 3, left) underperforms all methods at the same budget. Appendix M (Figures 29 and others) provides qualitative examples of degenerate outputs. The paper explicitly discusses over-optimization in Sections 5.3 and 8 as a key bottleneck.

Mitigation status. The compute-optimal policy mitigates over-optimization by avoiding aggressive search on easy problems, but this is a workaround, not a fix. The paper does not propose any method for training more robust verifiers, does not explore techniques to detect or prevent over-optimization at inference time (e.g., early stopping based on verifier score saturation, diversity penalties), and does not characterize how the compute-optimal policy would change if verifier quality were improved. Section 8 mentions "improving verifier robustness" as future work but does not develop it.

The Revision Model Has a ~38% Correct-to-Incorrect Reversion Rate, and Revision Training Is Fragile

The assumption or constraint. The revision model is trained on sequences where in-context answers are incorrect and the target output is correct. At test time, however, the model may produce a correct answer early in the revision chain and then "revise" it to an incorrect answer in a subsequent step. The paper reports (Section 6.1):

"approximately 38% of correct answers get converted back to incorrect ones using a naive approach"

This is a direct consequence of the training data construction: the model never sees examples where the current answer is already correct and should be preserved. The training signal exclusively teaches the model to change its answer, never to recognize when no change is needed.

Additionally, Appendix K reports that attempting to further optimize the revision model using ReST$^{EM}$ (Singh et al., 2024) — a reinforcement-learning-style iterative improvement method — backfires: the ReST$^{EM}$-trained revision model shows substantially degraded performance with sequential revisions compared to the standard SFT-trained model (Figure 16). The authors hypothesize that "on-policy data collection in ReST$^{EM}$ exacerbates spurious correlations in revision data."

The consequence. The 38% reversion rate means that the revision chain is inherently unstable — each revision step has a non-trivial probability of destroying a correct answer. The paper mitigates this with majority voting or verifier-based selection across the entire revision chain (picking the best answer from any step, not just the final step), but this adds computational overhead and does not address the root cause. The ReST$^{EM}$ failure demonstrates that the revision training procedure is fragile and that naive attempts to improve it can make things worse. Practitioners attempting to replicate or extend the revision approach should be aware that the positive results depend on specific, somewhat delicate choices: offline data construction, edit-distance-based incorrect-correct pairing, and avoiding on-policy iterative training.

What evidence exists in the paper. The 38% reversion rate is reported in Section 6.1 (exact number from text). The mitigation via chain-wide selection is described in Section 6.1 and Appendix I. The ReST$^{EM}$ failure is documented in Appendix K, Figure 16. The paper provides these numbers transparently but does not deeply investigate the causes of reversion or propose training procedures that would teach the model to recognize correct answers and leave them unchanged.

Mitigation status. Partial. The within-chain selection mechanism (majority voting or verifier-based selection) reduces the impact of reversion but adds computational overhead and does not eliminate the problem — the model still wastes revision steps destroying correct answers that could have been "locked in." The paper does not explore alternative training data constructions (e.g., including some trajectories where the correct answer appears early and the model learns to output a "stop" or "no change needed" token) or architectural modifications (e.g., a separate "confidence" head that gates whether to accept or revise). The ReST$^{EM}$ negative result is reported but not deeply analyzed, leaving open questions about what specifically causes on-policy revision training to fail. Future work on more robust revision training procedures is suggested in Section 8 but not developed.

7. Implications and Future Directions

How This Work Changes the Landscape

Hedgehog changes the conversation around linear attention from a story of inherent tradeoffs to one of learnable approximation. Prior to this work, the field operated under an implicit assumption: linear attention can be fast or it can be accurate, but not both. The persistent 4–6 perplexity gap on language modeling, the failure of converted models to recover task performance, and theoretical results from SETH (Alman & Song, 2023; Keles et al., 2023) suggesting fundamental barriers to softmax approximation all pointed toward an efficiency-expressivity tradeoff that might be structurally unavoidable. Hedgehog demonstrates that this tradeoff is contingent on the choice of feature map, not inherent to linear attention itself.

The conceptual shift is from designing feature maps to learning them through distillation. This is not merely an engineering trick — it is a methodological reframing. Prior work asked: "What function $\phi$ best approximates $\exp(\bm{q}^\top \bm{k})$ in the kernel sense?" Hedgehog asks: "Given a specific pretrained model's query and key distributions, what feature map produces attention weights that match the model's own softmax weights?" The target is not the exponential function in the abstract, but the specific attention distributions that a trained Transformer actually produces. This shift from function approximation to distribution matching is what makes Hedgehog work: the feature map only needs to be accurate in the bounded, model-specific regime where queries and keys actually operate, not across all of $\mathbb{R}^d$.

This reframing resolves a contradiction that has lingered in the linear attention literature since its inception. On one side, kernel method theory says that any positive-definite kernel can be approximated by random features with bounded error (Rahimi & Recht, 2007). On the other side, practitioners consistently found that Performer-style random features underperform softmax by wide margins (4–6 perplexity points, as the paper documents). The resolution, which Hedgehog makes explicit, is that mean-squared error in kernel value does not imply fidelity in attention weight space. A feature map can approximate $\exp(\bm{q}^\top \bm{k})$ well on average while systematically smoothing out the sharp peaks that make softmax attention useful. The paper's diagnostic framework — spikiness and monotonicity measured directly on attention weight distributions — provides the language for understanding why mathematically reasonable approximations fail empirically.

The paper also reconciles the training-from-scratch and conversion regimes under a single framework. Prior work treated these as separate problems: training linear Transformers from scratch (Choromanski et al., 2020; Katharopoulos et al., 2020; Schlag et al., 2021) versus converting pretrained models (Kasai et al., 2021; Mao, 2022). Hedgehog shows that the same architectural component — a single-layer MLP with exponential activation — solves both, with the only difference being whether the MLP is trained end-to-end (training-from-scratch) or via a two-stage distillation-then-finetuning protocol (conversion). This unification is practically significant because it means a team that develops Hedgehog for training-from-scratch can use the same infrastructure for converting pretrained models, and vice versa.

The work also redirects research attention from architecture to training signal. The T2R-HH ablation (Table 9) is the crucial evidence: taking a prior linear attention (ReLU feature map from Kasai et al., 2021) and training it with Hedgehog's attention weight distillation loss improves its performance substantially across GLUE tasks. This means the distillation procedure is a general contribution that benefits any linear attention, independent of the specific feature map. Conversely, the untrained Hedgehog MLP (HH No Train, Figure 8) performs poorly despite the exponential architecture. The implication is that how you train the feature map matters more than the specific feature map architecture, provided the architecture has sufficient capacity. This shifts the research agenda from hunting for better fixed functional forms to designing better training objectives and distillation procedures.

The paper makes pretrained-conversion of large language models a viable research direction rather than a curiosity. Prior conversion methods produced models that were either partially quadratic (Kasai et al., 2021, requiring "additional quadratic attention modules") or severely degraded (T2R-Llama2 producing repetitive nonsense, as shown in Appendix C.3). Hedgehog-Llama2 7B, by contrast, produces coherent summaries that actually outperform the standard attention baseline on SAMSum by 28.1 ROUGE-1 points when both use LoRA finetuning. This is a qualitative shift in what is possible: converting a 7B-parameter LLM to linear attention while maintaining or improving task performance opens the door to efficient deployment of large models on long sequences without retraining from scratch — a prospect that was largely theoretical before this work.

However, it would be a mistake to characterize this as a paradigm shift. Hedgehog is an incremental refinement at the architectural level (MLPs inserted into existing attention layers) combined with a methodological innovation at the training level (attention weight distillation). The underlying linear attention mechanism — the kernel trick that enables $\mathcal{O}(n d^2)$ complexity — is unchanged from Katharopoulos et al. (2020). The exponential activation is not new (the paper's own $\phi_t$ experiments in Section 3.2 use it). The adapter-like insertion of trainable modules is standard practice since Houlsby et al. (2019). What is new is the synthesis: recognizing that a learned exponential feature map trained to match attention weights can recover both spikiness and monotonicity, and that this recovery closes the performance gap across three deployment regimes. The contribution is in the diagnosis, the training objective, and the empirical demonstration of sufficiency — not in any single novel component.

Certain research directions become more attractive after this work. Verifier and reward model training for LLMs (where the quality of learned scoring functions is the bottleneck, not the search algorithm) becomes a closer analogue — both Hedgehog's distillation and PRM training face the problem of learning to match a target distribution from limited samples. The connection between attention-level distillation and output-level distillation (knowledge distillation; Hinton et al., 2015) merits investigation: can Hedgehog-style attention matching be combined with standard distillation to compress Transformers along both the attention and feedforward dimensions? Conversely, some directions become less urgent. The search for better fixed feature maps — random Fourier features, polynomial approximations, locality-enhanced maps — now seems less promising than simply learning the feature map from the model's own attention distributions. The paper's finding that the second-degree Taylor approximation works perfectly but is too slow (Figure 6) suggests that expressivity is not the bottleneck; efficiency of the expressive representation is. Research effort is better spent on training procedures and compact learned approximations than on deriving new closed-form kernel expansions.

Follow-Up Research This Work Enables

Training a single Hedgehog feature map shared across all layers and heads, then measuring the performance gap versus per-head, per-layer MLPs.

The paper applies separate Hedgehog MLPs to each attention head in each layer, which for a 12-layer, 12-head Transformer means 144 separate MLPs (each a $d \times d$ weight matrix). This represents 0.495% of Llama-2 7B's parameters — small but not negligible, and the per-head, per-layer design means the number of Hedgehog parameters scales with model depth and width. An obvious efficiency question is whether a single shared Hedgehog MLP, or one per layer shared across heads, can achieve comparable performance. The paper's attention weight visualization (Figures 10–20, Appendix C.4) shows that different heads and layers learn qualitatively different attention patterns — early layers attend broadly, later layers attend to specific tokens — so sharing MLPs might force a compromise that degrades fidelity. A controlled experiment would train three variants — (a) per-head, per-layer (the paper's default), (b) per-layer shared across heads, and (c) one global MLP for the entire model — on the BERT-finetuned conversion setting, measuring both attention weight KL divergence and downstream GLUE performance. A negative result (shared MLPs substantially underperform) would confirm that the diversity of attention patterns across heads is irreducible and that the per-head parameter cost is necessary. A positive result (shared MLPs work nearly as well) would make Hedgehog even more parameter-efficient and simplify deployment.

Evaluating Hedgehog on a long-range language modeling task with sequence lengths of 8K–32K tokens, measuring both perplexity and wall-clock inference time against FlashAttention.

The paper's efficiency benchmarks (Figure 6) measure a single attention layer, not an end-to-end model. The language modeling experiments (WikiText-103, GPT-2 conversion) use 1024-token contexts, where the crossover point for Hedgehog to outperform FlashAttention has barely been reached (Figure 6 shows crossover at ~2K tokens). The core efficiency promise of linear attention is handling very long sequences, but the paper has no long-sequence task evaluation. A strong follow-up would evaluate Hedgehog-GPT-2 (converted and finetuned) on a long-document language modeling benchmark — PG-19 (Rae et al., 2019) with 8192-token contexts, or BookCorpus with 16384-token contexts — measuring both perplexity and end-to-end inference latency against FlashAttention-GPT-2 at the same sequence lengths. The key measurement would be the wall-clock speedup at matched perplexity, or equivalently the perplexity at matched latency, as a function of sequence length. This experiment would quantify the real-world efficiency gain in the regime where linear attention is supposed to dominate, and would test whether the Hedgehog attention weights maintain fidelity at lengths far beyond the distillation context length (1024 tokens for GPT-2). Table 5 provides preliminary evidence that Hedgehog's attention fidelity remains consistent up to 4096 tokens (8× distillation length), but this is measured on BERT with CoLA data, not on a generative language modeling task.

Combining Hedgehog's attention weight distillation with standard knowledge distillation (output-level) to compress a large Transformer along both attention and feedforward dimensions simultaneously.

Hedgehog addresses the attention cost bottleneck. Standard knowledge distillation (Hinton et al., 2015; Sanh et al., 2019) addresses the feedforward and embedding cost bottleneck by training a smaller student model to match a larger teacher's output logits. These approaches are orthogonal and potentially complementary: Hedgehog reduces attention from $\mathcal{O}(n^2 d)$ to $\mathcal{O}(n d^2)$, while distillation reduces the constant factor $d$ (model dimension) and the number of layers. A natural combination would take a large pretrained teacher (e.g., Llama-2 13B), convert it to a Hedgehog-linear variant via attention distillation, then use this linear teacher to distill a smaller linear student (e.g., Hedgehog-Llama-2 1B) through output logit matching. The student would benefit from both the teacher's larger capacity (via output distillation) and linear attention efficiency (via the Hedgehog architecture). The experiment would compare: (a) standard distillation from quadratic teacher to quadratic student, (b) Hedgehog-linear student trained from scratch, and (c) the proposed Hedgehog-linear student distilled from a Hedgehog-linear teacher. The key metric is perplexity on long sequences (8K+ tokens) at matched inference latency. A positive result would demonstrate that attention-level and output-level distillation are complementary compression strategies that can be stacked.

Testing whether Hedgehog feature maps distilled on one model family transfer to a different model family, quantifying the degree of model-specificity in the learned approximation.

The paper's generalization experiments (Section 5.1, Tables 4, 5, 14) show that Hedgehog feature maps trained on one dataset transfer to new datasets for the same model (BERT-base). An open question is whether they transfer across models. Could a Hedgehog MLP trained to match BERT-base's attention weights be inserted into BERT-large (same architecture, different scale) or RoBERTa (same scale, different pretraining) and still produce high-fidelity attention weights? The paper positions Hedgehog as learning to approximate softmax attention for a specific model's query/key distributions, and the queries and keys of different models (even of the same architecture) may occupy different regions of $\mathbb{R}^d$. A transfer experiment would take Hedgehog MLPs distilled on BERT-base-uncased with CoLA data, insert them into BERT-base-cased, BERT-large-uncased, and RoBERTa-base, and measure the KL divergence between the transferred Hedgehog weights and each model's native softmax weights on the same input. A finding of high transferability would suggest Hedgehog learns something close to a universal exponential kernel approximator, which would be a strong theoretical result. A finding of low transferability would confirm model-specificity and motivate techniques for few-shot or zero-shot adaptation of Hedgehog MLPs across models.

Investigating whether attention weight distillation can be done with fewer samples by using an auxiliary verifier or by leveraging the structure of the softmax distribution (e.g., only matching top-k weights).

The paper's distillation procedure requires computing full softmax attention weights for every query-key pair in every training batch, which is $\mathcal{O}(n^2 d)$ — the very cost Hedgehog aims to avoid. This is acceptable for a one-time conversion cost, but it limits the practicality of distillation for very large models or very long sequences, where computing full softmax attention is prohibitively expensive even once. A follow-up could explore whether the distillation loss can be approximated by matching only the top-$k$ attention weights (the "spiky" part of the distribution), since the paper's diagnostic framework argues that spikiness is the critical property. If only the largest 5–10 attention weights per query matter, the distillation cost could be reduced from $\mathcal{O}(n^2 d)$ to $\mathcal{O}(n k d)$ using efficient top-k retrieval. An experiment would measure the tradeoff between $k$ (number of attention weights matched), distillation wall-clock time, and downstream task performance after conversion. A finding that $k \approx 10$ suffices for 512-token sequences would make Hedgehog conversion practical for much larger models and sequence lengths. A negative result (full attention matching is necessary) would bound the scalability of the current approach.

Stress-testing Hedgehog's monotonicity property by training an intentionally non-monotonic feature map with the same distillation loss, measuring the gradient conflict rate during finetuning.

The paper argues that monotonicity prevents conflicting gradients during conversion finetuning, but provides only correlational evidence (Figure 3, Table 1). A direct causal test would construct a feature map that is explicitly non-monotonic in the dot product — for instance, $\phi(\bm{x}) = [\exp(\bm{w}_1^\top \bm{x}), \ldots, \exp(\bm{w}_d^\top \bm{x})]$ where the weights $\bm{w}_\ell$ are randomly initialized and frozen (not trained) — train it with the same distillation loss (which would produce a best-effort approximation under the non-monotonic constraint), and finetune the converted BERT model while monitoring the cosine similarity between the gradient of the attention loss and the gradient of the task loss. The paper predicts that non-monotonic feature maps will produce gradient conflicts (negative cosine similarity more frequently), leading to slower convergence or worse final performance. Quantifying this gradient conflict rate as a function of monotonicity (measured by the Spearman correlation between dot products and attention weights) would validate the mechanistic claim and provide a diagnostic for future feature map designs.

Practical Applications and Downstream Use Cases

Conversion of deployed task-specific Transformer models to reduce inference latency on long inputs. Organizations that have finetuned BERT models for document classification, legal contract analysis, or medical record processing — tasks where input documents routinely exceed 2000 tokens — can use Hedgehog's two-stage protocol to convert these models to linear attention variants with negligible accuracy loss (99.3% recovery on GLUE, per Table 8) and significant latency improvements. Figure 6 shows Hedgehog achieving ~6× speedup over FlashAttention at 32K tokens for a single attention layer. For a full BERT-base model processing 4096-token documents, this translates to a meaningful reduction in per-document inference time, especially in high-throughput batch processing settings. The conversion cost is modest: two epochs of attention distillation on task-specific data (Appendix B.4), followed by brief finetuning. The converted model uses the same HuggingFace-compatible interface (Appendix A.2, Listings 2–3), minimizing engineering effort.

Efficient finetuning of large language models on long-context tasks via pretrained-conversion plus LoRA. The Llama-2 7B result (28.1 ROUGE-1 improvement over standard attention with LoRA on SAMSum, Table 11) demonstrates a workflow that is immediately applicable: take a pretrained LLM, convert it to Hedgehog-linear attention via distillation on a general text corpus (the paper used SAMSum data for distillation, but Table 4 and Table 15 suggest WikiText-103 or other general corpora work well), then finetune on the target task with LoRA. The Hedgehog MLPs add only 0.495% to the parameter count, and the entire pipeline runs on a single A6000 GPU with bfloat16 weights (Appendix B.5). For practitioners who need to adapt LLMs to tasks with long inputs — legal document summarization, scientific paper QA, long-form dialogue — this provides a path to linear-attention efficiency without sacrificing the quality of the pretrained model. The key practical decision is whether the input lengths are long enough (>2K tokens) to justify the conversion overhead; for short-input tasks, standard FlashAttention with LoRA is likely simpler and equally effective.

Cost-efficient pretraining of smaller models with long-context capability from the start. For teams building models from scratch that need to handle long sequences — genomics (DNA sequences of 10K–100K+ base pairs), audio processing (raw waveform modeling), or long-document retrieval — Hedgehog offers a training-from-scratch recipe that is significantly more competitive than prior linear attentions. Table 7 shows Hedgehog closing the WikiText-103 gap from 6 perplexity points to 1.7 points relative to softmax. For a 125M-parameter model, the residual gap is small enough that the efficiency gains (linear rather than quadratic scaling with sequence length) likely outweigh the quality loss for most long-sequence applications. The training procedure is straightforward: insert Hedgehog MLPs with identity initialization, train end-to-end with the task loss — no distillation needed (Section 5.2). The primary practical consideration is whether the head dimension is small enough ($d = 64$ in the paper's experiments) to keep the $\mathcal{O}(n d^2)$ complexity favorable; at larger head dimensions (128 or 256, common in newer architectures), the constant factor grows, and the crossover point where Hedgehog beats FlashAttention shifts to even longer sequences.

When to Prefer This Method

The paper articulates clear tradeoffs between Hedgehog and alternative approaches, grounded in the three deployment regimes it studies. These can be summarized as decision rules:

• Prefer Hedgehog pretrained-conversion over H3, Hyena, or other subquadratic architectures trained from scratch when you already have a strong pretrained Transformer (GPT-2, Llama-2, BERT) and want to efficiently finetune it on a task with long sequences. The pretrained-conversion result (16.7 perplexity for Hedgehog-GPT-2 vs. 18.5 for H3 on WikiText-103, Table 11) shows that leveraging existing pretraining is more effective than training subquadratic architectures from scratch on the target data alone.

• Prefer Hedgehog over T2R or other fixed-feature-map conversion methods when converting a pretrained or finetuned model to linear attention and quality recovery matters. The GLUE results (Table 8) and Llama-2 generations (Appendix C.3) show Hedgehog substantially outperforming T2R across scales from 110M to 7B parameters, with the distillation training providing a general benefit even for the T2R feature map (T2R-HH ablation).

• Prefer training-from-scratch with softmax attention over Hedgehog when model quality is paramount and sequence lengths are short enough (< 2K tokens) that quadratic attention is not the bottleneck. Hedgehog still shows a 1.7 perplexity gap to softmax on WikiText-103 (Table 7), and the efficiency crossover is only reached at longer sequences (Figure 6).

• Prefer the Taylor approximation over Hedgehog only in the narrow regime where head dimension is very small ($d \ll 64$) and absolute maximum quality is required, since the Taylor approximation matches softmax performance perfectly (Table 3) but is $\mathcal{O}(n d^3)$. For typical head dimensions ($d \geq 64$), the Taylor approximation is actually slower than FlashAttention at all sequence lengths (Figure 6), making it practically unusable despite its theoretical appeal.