The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry
Authors: Michael Zhang
Kush Bhatia
Hermann Kumbong
Christopher Re
Venue: ICLR 2024 poster (2024)
Keywords: linear attention
transformers
Primary area: representation learning for computer vision, audio, language, and other modalities

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

Published as a conference paper at ICLR 2024
THE HEDGEHOG & THE PORCUPINE: EXPRESSIVE LINEAR ATTENTIONS WITH SOFTMAX MIMICRY
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré Department of Computer Science, Stanford University fmzhang, kushb, ch (erismreeeecs stanford, edu, kumbch@starffrd.edd,
ABSTRACT
Linear attentions have shown potential for improving Transformer efficiency, re- ducing attention's quadratic complexity to linear in sequence length. This holds exciting promise for (1) training linear Transformers from scratch, (2) "finetuned- conversion" of task-specific Transformers into linear versions that recover task performance, and (3) seetrained-conversionn of Transformers such as large lan- guage models into linear versions finetunable on downstream tasks. However, lin- ear attentions often underperform standard softmax attention in quality. To close this performance gap, we find prior linear attentions lack key properties of softmax attention tied to good performance: low-entropy (or "spiky") weights and dot- product monotonicity. We further observe surprisingly simple feature maps that retain these properties and match softmax performance, but inefficient to pute in linear attention. We thus propose Hedgehog, a learnable linear attention that retains the spiky and monotonic properties of softmax attention while main- taining linear complexity. Hedgehog uses simple trainable MLPs to produce atten- tion weights mimicking softmax attention. Experiments show Hedgehog recovers over 99% of standard Transformer quality in train-from- fro----rattttch and finetuned- conversion settings, outperforming prior linear attentions up to perplexity points on WikiTText- 103 with causal GPTss and to 8.7 GL. score points finetuned bidirectional BERTs. Hedgehog also enables netrained-cooveenee pretrained GPT-2 into a linear attention variant achieves state-of-the-art 16.7 per- plexity on WikiText- 103 for 125M subquadratic decoder models. We finally turn a pretrained Llama-2 7B into viable linear attention Llama. With low-rank adap- tation, Hedgehog-L -lama2 7B achieves 28.1 higher ROUGE-1 1 points over the base standard attention model, where prior linear attentions lead to 16.5 point drops.
Converting a
INTRODUCTION
Linear attentions are promising methods for improving Transformer efficiency. By replacing the softmax of attention's query and key dot products with kernel function feature maps, linear attentions reduce attention's 's time and space complexity from (n to O where is sequence length, is head dimension and d' the feature map dimension (Katharopoulos et al., 2020; Choromanski et al., 2020; Peng et al., 2021; Xiong et al., 2021; Schlag et al., 2021). ) For typical Transformer settings, e.g. with head dimension 64 and sequence lengths at 512 to 32K. this mundratic-t--linee scaling can result in significant speed and memory improvements (Fig. 6). As drop-in alternatives to popular softmax attention (Vaswani et al., 2017), linear attentions not only improve Transformer efficiency when training new models from scratch but can also improve inference efficiency by converting pretrained Transformers into corresponding linear variants (Kasai al., 2021; Mao, 2022). Linear attention enables efficient Transformers in variety of regimes:
Training-from-scratee training Transformer models with linear attention with the goal of match- ing standard Transformer performance, e.g., as tested on benchmarks such as Long Range Arena (LRA) classification (Tay al., 2021) and WikiText- 103 language modeling (Merity et al., 2017).
finetsunedcoonvrrsion: swapping the attentions of task-specific Transformers and finetuning them to convert existing models into linear versions, with the goal to recover original task per- formance with improved efficiency (Kasai al., 2021; Mao, 2022).

---

Published as a conference paper at ICLR 2024
(1) Trainuble Linear Attention Festure Map
(2) Supervised mimic
Leading Expressive Efficient Linear Attention
standard attention weights
(()) exp( (W' T x +
Concat
Concatt
Ma)TT(K
expl9 k))
Standard Baftmax Attention
Hedgehog Linear Attention
Inewr
Linnar
Lincar
Lincar
Lincar
Linear
Linear
Linco
Figure 1: Hedgehog learns a trainable linear attention feature map designed to mimic standard atten- tion, resulting in expressive yet efficient linear attentions for various Transformer training settings
Preirained-coovversion doing the same finetuned-conversion but for pretrained Transformers such as large language models (LLMs), e.g., to transfer to new tasks and longer contexts.
Unfortunately, existing linear attention mechanisms typically fail to match softmax attention in mod- eling quality. When training from scratch, linear attentions achieve 4-6 worse perplexity (ppl) than softmax attention on standard benchmarks such as WikiText-103 (Schlag et al., 2021; Irie et al., 2021; Fu et al., 2023), the equivalent gap between 125M and 255M Transformers (Dai et al., 2019).
When converting finetuned models, linear attention models require additional quadratic attention modules to close the gap (Kasai et al., 2021; Mao, 2022). One might worry that such gaps are fundamental; for example, recent theory using the Strong Exponential Time Hypothesis (SETH) showed that high-quality truly subquadratic algorithms to approximate softmax attention may be impossible with large sequence length (Alman & Song, 2023;
Keles et al., 2023).
We begin by empirically studying why this performance gap exists between standard softmax and proposed linear attentions. We identify two simple properties for softmax attention which prior linear attentions lack: 1) low-entropy "spikyness" and 2) dot-product monotonicity. We hypothesize that the quality gap in linear attentions corresponds with lacking these two properties:
Low-entropy "spikyness"" Intuitively, want attentions that attend to relevant tokens while ig- noring irrelevant ones via their query-key interactions. We observe these low-entropy or "spiky" aetention-weight distributions in standard Transformer attention but not prior linear attention maps where spikes enabled via the scaled dot-product softmax are lost via other feature maps (Fig. 2nd find this strongly corresponds to Transformer performance (Fig. 4).
Dot-product monotonicity: This property requires that attention weights increase as the dot prod- ucts of their corresponding queries and keys increase. Intuitively, the lack of this monotonicity can produce unstable gradients during training and finetuning, where increasing the query-key dot product can result in decreasing the attention weight the other way (and vice versa).
As first step to recover these properties, we explore simple feature maps-such as low-degree Tay- lor polynomial approximations to the exp() function-that satisfy the above two properties (albeit in restricted regimes of bounded query-key dot products). In practice, we find that queries and keys are often bounded, resulting linear attentions that recover softmax attention's spikiness, mono- tonicity, and subsequent performance. Unfortunately, while technically linear in sequence length, these polynomial feature maps remain inefficient to compute. They take O(d"++1) time and space, and we find degree 2 2 necessary for performance.
learnable linear attention trained to capture the spiky and monotonic softmax properties. Unlike prior works that propose a specific kernel func- tion (Katharopoulos et al., 2020; Choromanski et al., 2020; Qin et al., 2022b) and our polynomial feature maps, we learn these feature maps single-layer MLPs specifically trained match soft- max attention weights. By mapping from Ed + Rdd, we maintain prior linear attentions' (ndd)) complexity. However, training these mappings via softmax attention weights as cross-entropy soft- labels, we find Hedgehog can match softmax attention weights with much higher fidelity (Fig. 7), producing low-entropy and monotonic weights that match standard attention performance quality.
We thus propose Hedgehog, an
We validate experimentally that Hedgehog's 's improved expressivity translates to closing the softmax attention performance gap in the three regimes mentioned above:
2

---

Published as a conference paper at ICLR 2024
Trrining-fromm-cceee we find Hedgehog matches Transformers standard attention bench- marks such as Long Range Arena (LRA) (Tay et al., 2021) task, and closes the linear attention gap by 68.6% WikiText-109 language modeling (improving up to 6 ppl).
Finetuned-conversion: we find Hedgehog recovers >99% of original model performance on av- erage across bidirectional encoder-only 110M BERT-base models finetuned on GLUE and causal decoder-only 125M GPT models finetuned on Wikitext-103.
we find Hedgehog enables effective transfer to new tasks and efficient scaling longer contexts, while frequently outperforming modern subquadratic sequence archi- tectures by linearizing existing pretrained Transformers. A 125M Hedgehog-GPP -PPo-2 finetuned on Wikitext- 103 achieves new state-of-the-art 16.7 ppl for subquadratic models of the same size.
Pretrained
Finally, we demonstrate that Hedgehog can be scaled up to modern large language models; we convert pretrained Llama-2 7B into viable linear attention Llama. With low-rank adaptation, Hedgehog-Llama2 7B achieves up to 28.1 higher ROUGE-1 points over the base standard atten- tion model. In contrast, prior linear attentions result in models that struggle produce coherent text (with 16.5 ROUGE-1 point drops).
2 PRELIMINARIES AND RELATED WORK
We provide background on attention computation, describe kernel feature based linear attentions, and finally provide details existing linear attention mechanisms proposed in the literature.
Attention setup. Let {q: J==1- {k.)_ {w:}== denote the set of queries, keys, and values, with individual elements in Rd Let n denote sequence length and d denote head dimension. We compute attention outputs y/ E IRd by first computing similarities between each qi and every ky: for causal attention compute these similarities for i. The vanilla Transformer attention computes these similarities using the softmax dot products (Vaswani et al., 2017):
explq kg/vVd)
VI sim(q ky)u, where sim(g., k,) Em exp(q, km Som/vd
While very expressive, computing attention via Eq. for all {y, Hi requires O( time and memory, making this inefficient for long sequences. To improve efficiency without sacrificing qual- ity, we thus want alternative linear attention maps which maintain standard attention's 's expressivity.
Linear attention and kernel functions. Observe that the exp(-) in Eq. I can be viewed as a kernel function, which Tsai et al. (2019); Katharopoulos et al. (2020) show can be replaced in general with K(xx, x²) <(x) 6(xx). Here Φ Rd + Rdd is feature map applied to each vector. We can thus compute attention linear time and space over the sequence length n, seen by rewriting Eq. as:
y ((q.) E (o(k, v.)) σ(q.) Σド= ((k))
Prior feature maps. From the previous section, we observe that linear attentions are promising di- rections for improving Transformer efficiency at both training and inference time. Numerous prior works have proposed feature maps & aiming to remain more efficient (where linear attention is desir- able standard attention if d' n), while still being expressive and stable train. These range from & ensuring positive attention weights, egg., via + ELU (Katharopoulos et al., 2020) or ReLU (Ka- sai et al., 2021), softmax or Gaussian kernel approximations via randomized features (Rahimi & Recht, 2007; Choromanski et al., 2020; Peng et al., 2021; Choromanski et al., 2021; Zheng et al., 2023) or low-rank approximations (Xiong al., 2021; Chen et al., 2021).
IMPROVING LINEAR ATTENTION VIA SPIKY AND MONOTONIC WEIGHTS
We begin by identifying two key properties of attention weights which we hypothesize essential for good performance quality. The first, low-entropy spikyness, requires that the attention map is able to capture effectively capture sparse relevant tokens in sequence. The second, monotonicity over kee dot products, requires the attention map to increase with increasing dot products, and allows for smooth conversion of pretrained Transformers into linear variants.
3

---

Published as a conference paper at ICLR 2024
Softmax
Performer
EWU
cosformer
Hedgehog (Ours)
Likep Cife
Li
Nougle
Ovay
010
238
00
218
6.0
012
1/ 10. Wwle
000
600
Figure 2: Attention weight spikiness. (Plots 1 5): Softmax attention results in lower entropy and "spiky" selective weighting compared to prior linear attentions (training from scratch on associative recall (Sec. 3.2)). (Plot 6): By training mimic softmax attention, our proposed Hedgehog recovers this spikiness as a linear attention, corresponding with improved performance (Sec. 5).
Softmax
Parformar
BLU
BoLu
cosFormer
Hedychop
prodect
product
Figure 3: Attention weight monotonicity. (Plots 1- 5): In contrast to softmax attention, prior linear attentions are not smoothly monotonic over trained query-key dot products, resulting in poor performance when converting BERT models by replacing attentions (Table 1). (Plot 6): Hedgehog recovers this monotonicity, and thus recovers 99% of BERT performance after conversion (Table 8).
3.1 PROPERTIES FOR EXPRESSIVE ATTENTION MAPS
Here we describe the spiky and monotonic properties hypothesized for desirable linear attention. We note these add to past observations for more performant linear attentions, including positive attention weights (Katharopoulos et al., 2020), orthogonal features (Choromanski et al., 2020; Irie et al., 2021), or locality (upweighting nearby values) (Qin et al., 2022a;b). We validate these properties among past linear attentions in Sec. 3.2, and preview how our proposed Hedgehog linear attention recovers these properties in correspondence with improved performance (Sec. 5) in Fig. 2.
Low-entropy spikiness. Intuitively, one source of attention's 's effectiveness is its ability to selec- tively upweight relevant tokens in sequence. This is popular interpretation visualized in various Transformer architectures and settings ranging from encoder-decoder language translation (Bah- danau et al., 2014) to VIT image segmentation (Dosovitskiy et al., 2020; Caron et al., 2021). Me- chanically, the softmax over query-key dot products exponentiates relative similarities between query and each key, quantified via low-entropy or "spiky" attention weight distributions (Fig. 2).
Linear attention maps work by replacing the softmax with the normalized dot product of alternate feature maps (Eq. 2). With existing feature maps, we find the resulting attention weights can result in much higher entropy or more uniform distributions. This is true even for methods designed to approximate the softmax under mean-squared error bounds (Choromanski al., 2020) (Performer, Fig. 2) or imposed locality (Qin et al., 2022b) (cosFormer, Fig. 2). This uniformity in attention weights reduces the modeling capacity linear attentions leading to worse performance quality.
Monotonicity over query-key dot products. This property requires that the attention maps are monotonic over query-key dot products: when the dot product increases (decreases), the attention weight increases (decreases). In Fig. 3, we observe that while softmax attention exhibits this mono- tonicty (first subplot), the existing linear attentions do not. We believe this can training issues after swapping attentions due to conflicting gradients between attentions and original model param- eters. In Fig. 3, trying to upweight attentions by increasing product similarity actually result in
4

---

Published as a conference paper at ICLR 2024
BERT-FT
1 ELU
ReLU
Performer
cosFormer
exp(t= 1)
exp(t = 2)
Matthew's correlation
58.8
28.1
39.5
24.7
39.9
45.9
50.0
Table 1: Finetuned-converssoon performance of BERT finetuned on CoLA (BERT-FT), using prior linear attentions. With poor monotonicity (Fig. 3), prior methods fail to recover performance.
decreased attention weights. Later in See 3.2, we find this corresponds to failing to recover original performance when converting finetuned Transformers.
3.2 EXPLAINING THE LINEAR ATTENTION PERFORMANCE GAP
We validate the properties introduced above by showing that (1) lacking spikiness corresponds to significantly worse performance when training from scratch, and (2) lacking spikiness and mono- tonicity corresponds to failing to recover performance when converting finetuned models.
Training from scratch. We compare various Transformers' abilities to solve Associative Recall (AR) (Ba al., 2016), next-token prediction task previously studied a proxy for language modeling capability (Olsson et al., 2022). AR tests how well a model can recall specific content in an input sequence, structured as a list of key-value pairs which ends in key (Table 12).
100
As a control for evaluating our hypothesis, we also con- sider simple feature map designed to induce "spikiness" but not monotonicity: (x) = exp(x - t), which applies a temperature-4 scaled exponential element-wise.
Softmax
Exp (t=1)
Acc. (%)
Exp (t=2)
Performer
ELU
In Fig. 4, we observe a strong correspondence between low-entropy attention weights and AR accuracy. While softmax attention solves the AR task perfectly, prior lin- car attentions struggle to achieve even 20% accuracy, at the same time obtaining much larger attention weight en- tropies. As further support to our hypothesis, we see that while the exponential map Φ1 fails AR and produces sim- ilarly high entropy attention weights, increasing spikiness with 2 actually solves the task.
ReLU
cosFormer
20
Attention Weight Entropy
Figure 4: Associative recall perfor- mance strongly corresponds to lower at- tention entropy: present in softmax at- tention but not prior linear variants.
Finetuned-converssonn. We next compare how various linear attentions perform at recovering original softmax attention performance for finetuned- conversion. We adopt the procedure in Kasai et al. (2021), which takes a Transformer already finetuned on specific task, swaps the attention layers with linear attention variant, and further finetunes the entire model on the same task.
For this setting, we evaluate with a Corpus of Linguistic Acceptability (CoLA) task (Warstadt et al., 2019), where the goal is to classify whether sentence is grammatically correct. We compare the performance of the original (softmax attention) BERT model with the linear attention converted models. In Table 1, we find that just as no linear attention smoothly captures monotonicity over the trained model's query-key dot products, no linear attentions fully recover the original finetuned BERT's Matthew's 's correlation of 58.8. This includes the spiky Φ2 feature map which was sufficient in the maining-from-scratte regime.
model (Devlin et al., 2018) finetuned on the
4 HEDGEHOG: EXPRESSIVE LINEAR ATTENTION VIA SOFTMAX MIMICRY
We present Hedgehog, simple, efficient, and expressive feature map trained to mimic softmax at- tention. Hedgehog is predicated by (1) there existing linear attention approximations to the softmax that recover the spiky and monotonic properties of standard attention in practice, and (2) that we can efficiently compute similar approximations efficiently.
In Sec. 4.1, motivate Hedgehog and show that (1) is possible by revisiting low-degree Taylor polynomials. We find that for linear attention, the Taylor exponential works as a surprisingly simple feature map. recovering spikiness and monotonicity while matching standard Transformer perfor- mance. Unfortunately, we also find it introduces its own issues, where the feature map results in
5

---

Published as a conference paper at ICLR 2024
large query and key dimensions and becomes inefficient to compute. In Sec. 4.2, thus overcome these challenges, we propose and describe Hedgehog, a trainable linear attention trained to mimic softmax attention. In Sec. 5.1, show how this enables similar spiky and monotonic properties to the softmax and Taylor exponential attentions, while retaining past linear attentions' efficiency.
4.1 SIMPLE POLYNOMIAL APPROXIMATIONS SOFTMAX ATTENTION
From findings in Sec. 3, we seek an efficient linear alternative to the softmax which retains its spiky and monotonic properties. We first consider simple potential approach: approximating the exponential in softmax by a low-degree Taylor polynomial (Keles et al., 2023; Banerjee al., 2020).
While in general, a high-quality approximation to the softmax should retain its spiky, monotonic, and performant properties, we ground our investigation with two potential caveats for the Taylor polynomial. First, recall that feature maps for p-degree polynomial approximations can be computed ((nd)) time and space for every query and key vector. Thus, while this indeed subquadratic in sequence length, the question remains whether we can set p low enough to make the computation feasible while approximating exp reasonably. Second, as a general property of polynomials, the Taylor approximation only tracks its original function with low error in bounded regimes.
Setup. To test the Taylor approximation, we use the second-degree exp approximation, and evaluate on the prior fon--fro--scattch and feature map as exp(q k) x Oraylor Praylor(k) where @uayor (x) projects a d-dimensional query
settings (Sec. 3.2). We implement the
key to (dd) (dP)-dimensional features (mnge((a)) (z) 1,,1,...... T.d x: -z, i,j [d] Positive results. We find that the 2nd-degree Taylor approximation retains both the spikiness and monotonic properties (Fig. 5), and this corresponds to matehing softmax attention performance (Table 2). We also note that here, the BERT query-key dot products are bounded in regimes where the second- order Taylor series exp approximation maintains monotonicity (Fig. 5). This suggests we can enable expressive linear attentions for training from scratch and finetuned-conversion.
Unfortunately, the 2nd-degree Taylor approximation is not efficient. Even with p 2, the feature map dimension is now d' + d + dz, resulting in O(nd") attention complexity. As summarized in Table 2, this introduces an efficiency-effectieeees trade-off among functional attention approximations. Thus, the question remains whether we can recover the expressivity and modeling quality of softmax while achieving similar (ndd) scaling of past linear attentions.
Caveats.
4.2 LEARNABLE LINEAR ATTENTIONS FOR MIMICKING SOFTMAX
Our key insight is that rather than rely on fixed functional form that captures spiky and mono- tonic properties, we can learn linear attention feature maps that so. For each attention block, we propose feature maps trainable single-layer MLPs, which is similar to prior work (Kasai et al., 2021) and acts similarly to an adapter (Houlsby et al., 2019) inserted after the query and key projec- tions in Transformer attention layers (Fig. 1). However, unlike prior work, we explicitly train these feature maps such that the attention layers mimic the properties softmax attention. We describe these two core components below, and validate these design choices in Sec. 5.1.
Spiky MLP feature map. Recall the kernel based linear attention paradigm from Sec. 2, where feature map Rd Rd is applied to both queries and keys to compute causal self-attention outputs using equation However, unlike prior work that sticks to a pre-specified function feature map, we make the feature map a trainable MLP. particular, for the single-head attention
Method
Complexity
Spiky?
Mono- tonic?
Train-from- scratch (acc)
BERT-FT (MC)
Taylor Spikiness
Taylor Monotonidty
425
Softmax
(n)3d)
58.8
ELU
(ndd)
28.1
124
Performer
O (
x
17.0
24.7
CosFormer
(nd2)
17.0
39.9
Taylor Exp
(
100.0
58.4
3
10
Table 2: Summary of feature maps compared soft-
Figure 5: Taylor approximation recovers
max, exhibiting an efficiency vs. expressivity tradeoff.
spikiness and monotonicity

---

Published as a conference paper at ICLR 2024
setting, compute Pmip (q:) @mip(k)) with simple one-layer MLP @mlp(2) @ b) where the matrix W RRdxd' and the bias E Rd are learned, and + is an activation function. To induce spikiness, we set & as the element-wise exponential function studied in Sec. 3.2, resulting in
@mpp = [exp(1]] a + b), exp(ereeppii x + +b)
(3)
To learn a softmax approximation, we train @mlp to minimize the cross-entropy loss between the computed linear attention weights and those that would have been computed via softmax attention. For query q: and keys {k; Jr, we compute the sample losses as
Attention weight distillation loss.
explg. k2)
(usp(())) Qmlp (kg)
C mm=1 exp(q, km) Em=1 Quip (q.) Omip (k;)
log
(4)
For training Hedgehog attentions in multi-layer and multi-head attention Transformers, apply a separate MLP to each head and each layer, and use the same Dmip for the queries and keys. We include further implementation details and pseudocode in Appendix A.
5 EXPERIMENTS
In experiments, we evaluate whether Hedgehog recovers softmax attention expressivity while retain- ing linear attention efficiency (Sec. 5.1), and how this improves modeling quality in training-froe from- scratch (Sec. 5.2), finetuned-conversioon (Sec. 5.3), and pretrained-convvrsioo regimes (Sec. 5.4).
5.1 BENCHMARKING HEDGEHOG FOR EXPRESSIVITY AND EFFICIENCY
Before evaluating Hedgehog on downstream tasks, we aim to validate Hedgehog's design choices for efficiency and expressivity. We address: (1) Do Hedgehog's ' spiky feature map and distillation loss recover the spiky and monotonic properties of softmax attention on the prior associative recall and BERT CoLA tasks? (2) Does Hedgehog achieve improved efficiency
softmax attention? (3) For conversion, do the learned attention weights actually match those of "ggound--ruth"" softmax aftention? Once learned, does this transfer to longer contexts and different tasks?
Recovering softmax spiky and monotonic properties. We test Hedgehog in the same train-from- scratch associative recall (AR) and finetuned-conversion of BERT on CoLA settings Sec. 3.2. For traning-from-scratee on AR, we do not use the distillation loss, and train the model end-to-end with next-token -reeeccpreii after inserting the learnable MLPs. In Table. 3, we find that Hedgehog achieves both favorable complexity and modeling for train-from from--cratch and This corresponds respectively with the spiky (Fig. 2) and monotonic (Fig. 3) properties noted prior.
Recovering linear attention efficiency. We next find Hedgehog's s ((nd²) scaling in compute and memory can lead real-world efficiency gains. We benchmark inference in wall-clock time and memory usage for one attention layer with 12 heads and head dimension 64 on sequences up to n = 32K tokens long (Fig. 6). Hedgehog achieves near 6x faster inference and similar memory to FlashAttention (Dao et al., 2022) (linear in memory but quadratic in time). Meanwhile, the Taylor approximation, while (n). gets significantly larger memory and slower speed due to the extra d.
Recovering softmax attention weights. We next study the combination of Hedgehog's feature map and distillation loss for matching softmax attention weights. Beyond recovering the spiky
250
28.3
Method
Complexity
AR
BERT-FT
Taylor Exp Hodgghoog
FlasrAtiantion
13.
Teyler Tap
Elshhtttnt
Suntime (mn4
Softmax
(n²
100.0
58.8
Taylor Exp
(nd
100.0
58.4
Hedgehog
O (nd2)
100.0
59.2
Sequence Length
Sequence Length
Table Hedgehog matches performance on as- sociative recall (AR) and BERT-finetuned conver- sion (BERT-FT) with prior best approaches, while achieving better time and space complexity.
Figure 6: Hedgehog linear scaling in wall- clock time (left) and memory (right). Unlike the Taylor approx.. Hedgehog inference gets real-world gains over FlashAttention.
7

---

Published as a conference paper at ICLR 2024
Softmax
Hedgehog
Taylor Exp
+ ELU
ReLU
Performer
CosFormer
Exp
Figure 7: Compared to prior linear attentions, trained Hedgehog layers (2nd left) produce attention weights closely tracking softmax (left), with greater fidelity with both components (vs. Fig. 8).
HH (CoLA)
HH (WT-103)
T2R-HH (CoLA)
HH
+
T2R-HH
HH No Train
Dataset
(No Train)
ELU
Performer
CosFormer
CoLAA
0.172
0.352
0.191
0.694
1.223
1.293
1.196
MRPC
0.663
0.485
1.128
1,250
2.165
2.234
1.982
MNLI
0,345
0.382
0.613
0.851
1.51
1.581
1.338
QNLI
0.671
0.444
1.126
1.139
1.968
2.069
1.817
Table 4: We find Hedgehog feature maps trained via distillation on CoLA WikiText- 103 generalize new GLUE data, better matching softmax than prior linear attentions or ablations (reporting KL div.).
Figure 8: Hedgehog ab- lated attention weights.
Softmax
HIl (CoLA)
HG (WT-103)
Pos ELU
Performer
Seq. Len
256
1024
2048
4096
CoLA KL
0.182
0.187
0,190
0.181
Table 5: Hedgehog attention maintains fidelity with softmax attention over con- text lengths for BERT-FT CoLA.
Figure 9: Hedgehog trained on CoLA and WT-103 re- cover softmax attentions MRPC data.
and monotonic properties, learning exactly match the weights can be particularly effective for converting or "distilling" pretrained quadratic Transformers into linear variants. For evaluation, we visualize the attention weights for different linear attentions in our BERT-FT CoLA setting (Fig. 7). We find Hedgehog recovers linear attention weights that match softmax's 's with much higher fidelity.
To further understand the contribution of Hedgehog's (I) spiky MLP and (2) distillation loss in Sec. 4.2, we visualize ablated attention weights by (1) using the distillation loss with the ReLU feature map used in Transformer-to RNN (T2R-HH) (Kasai al. (2021)), and (2) using untrained MLPs, replacing the trainable weights with an identity function (HH No Train). We find that distil- lation training is necessary to recover attention weights, and that the spiky MLP also helpful for matching attentions (later supported by improved Transformer conversion in Sec. 5.3).
Generalization to new data and longer contexts. Finally, we investigate the generality of learned Hedgehog feature maps. We show Hedgehog attentions learned over specific data and context lengths can still better match softmax attention weights for new data and sequence lengths than prior linear attentions. We distill attentions for BERT models using CoLA WikiText- 103 (WT- 103) samples, and report attention weights compared to softmax attention on three other GLUE tasks: qualitatively (Fig. 9) and quantitatively via KI. divergence weights (Table 4). We include additional visualizations and comparisons Appendix 5.
the "ground-truth" softmax
In Table 9, further show that Hedgehog attention matching remains consistent across longer contexts. Post-distillation on CoLA samples, we concatenate CoLA samples into sequences 256 to 4096 tokens long (up 8x the default 512 context length). We then compute attention weights using softmax and learned Hedgehog feature maps, and find that their KL divergence remains consistent.
5.2 LEARNING SEQUENCE MODELING FROM SCRATCH
We evaluate Hedgehog Transformers trained from scratch on the popular LRA sequence classifi- cation and WikiText-103 language modeling benchmarks. For training from scratch, we initialize MLPs identity matrices for Hedgehog feature maps, and train the entire models end-to-end with the task-specific loss. We find Hedgehog achieves best average accuracy for both tasks among linear attentions (Table 6, 7). For LRA, while non- non-Transfer Transformer models now state-of-the-a (Gu al., 2021), our work focuses on approximating attention, so we compare with competitive subquadratic Transformers. We adopt the hyperparameter settings the official benchmark (Tay et al.,
8

---

Published as a conference paper at ICLR 2024
Model
ListOps
Text
Retrieval
Image
Pathfinder
Average
Transformer
36.37
64.27
57.46
42.44
71.40
54.39
Local
15.82
$2.98
53.39
41.46
66.63
46.06
Linear Trans.
16.13
65.90
53.09
42.34
75.30
50.55
Reformer
37.27
56.10
53.40
38.07
68.50
50.67
Sparse Trans.
17.07
63.58
59.59
44.24
71.71
51.24
Sinkhorn Trans.
33.67
61.20
53.83
41.23
67.45
51.29
Linformer
35.70
$3.94
$2.27
76.34
51.36
Performer
18.01
65.40
53.82
42.77
77.05
51.41
Synthesizer
36.99
61.68
54.67
41.61
69,45
52.88
Longformer
35.63
62.85
56.89
42.22
69.71
53.46
BigBird
36.05
64.02
$9.29
40.83
74.87
55.01
Nystrmformer'
37.15
65.52
79.56
41.58
70.94
58.95
cosFormer!
37,90
63.41
61.36
43.17
70.33
55.23
Skyformer! 39.25
64.70
82.06
40.77
70.73
59.50
Hedgehog 37.15
64.60
82.24
40.15
74.16
59.66
Table 6: Training- wrrnnng--foo--ee from-scratch LRA. Hedgehog achieves best avg, acc. (%) across most compet- itive Transformers (full results in Table 13, trends hold). ? indicates results reported from original works. All others reported from the official LRA benchmark (Tay et al., 2021). Best, 2nd-hest.
Model
Transformer
Performer
Reformer
AFT
ELU)
Hedgehog
Perplexity 18.6
26.8
25.6
28.2
25.6
20.8
Table 7: Training-from-sooe on WikiText-103. Among 125M decoder-only models, Hedgehog significantly closes the gap between standard Transformers and prior linear attention maps by 68.6%.
2021). On WikiText- we adopt the setting et al. (2023), evaluating a 125M decoder-only Transformer on perplexity over 1024 tokens. Hedgehog significantly closes the gap by up to 6 PPL.
5.3 FINETUNED CONVERSION QUADRATIC TO LINEAR TRANSFORMERS
For the finetuned Transformer conversion regime, we evaluate performance recovery for BERT-base models finetuned on GLUE, and ViT-B/16 models trained ImageNet- K For both settings, first swap attentions and train via our distillation loss (Sec. 4.2). We then finetune the converted BERT models on their original tasks in Transformer-to RNN (T2R) (Kasai et al., 2021).
For BERT, we compare Hedgehog to T2R in Table 8, and find that in contrast, Hedgehog conversion recovers near- 100% the original softmax attention performance. To further test Hedgehog's fea- ture map and attention distillation, we also compare against an ablation that trains the T2R feature map with our distillation loss (T2R-HH). We find that training to mimic softmax attentions boosts performance of T2R, suggesting that attention weight distillation may be general step to improving linear attention feature maps. However, Hedgehog's 's exponential still leads to superior performance. suggesting Hedgehog can also apply other modalities.
We find similar results for
5.4 PRETRAINED CONVERSION FOR SUBQUADRATIC TASK TRANSFER
We finally evaluate Hedgehog for converting pretrained Transformers into linear Transformers. We consider two settings: (1) To benchmark Hedgehog and the pretrained-convesiie
regime for sub- quadratic sequence modeling, we use the same WT-103 evaluation in Sec. 5.2 for converting 125M- parameter GPT-2. (2) As an early application for Hedgehog on larger models, we convert Llama-2 7B (Touvron al.. 2023) before finetuning with low-rank adapters (LoRA) (Hu et al., 2021) on SAMSum summarization (Gliwa et al., 2019). We include further training details in Appendix. B.5.
To most directly measure pretrained-comverte For GPT-2, we find Hedgehog both outperforms T2R, and further outperforms modern subquadratic sequence models such H3 (Fu et al., 2023) and Hyena (Poli et al., 2023) (Table 10). Although not directly comparable due to pretraining, we also compare with zero-shot and finetuned GPT-2 for reference. While Hedgehog is PPL. off the fully quadratic finetuned GPT-2, significantly
quality, for both settings we compare against T2R.

---

Published as a conference paper at ICLR 2024
improves over zero-shot while being linear to train. We finally apply Hedgehog for Llama-2 conver- sion, where Hedgehog enables linear attention Llamas that train via LoRA (see Appendix C.3 for sample generations).
Method
CoLA
SST2
MRPC
STS-B
QQP
MNLI
QNLI
RTE
(%) Recover
Top-1
Ace. %
BERT-FT
58.8
88.8
91.0
84.7
91.3
68.2
100.0
VIT-B/16
80.3
T2R
43.6
87.7
83.0
78.6
86.7
78.9
84.6
54.1
88.9
T2R-HH
56,9
90.9
89.1
77.7
90.0
77.4
84.5
56.3
93.5
T2R-HH
77.0
Hedgehog
59.2
92.6
90.1
87.4
91.0
82.6
89.6
69.3
99.3
Hedgehog
79.5
Table 8:
Finetuned-cooverssion evaluation.
Hedgehog recovers 99.3% of
Table 9: Hedgehog
original finetuned BERT (BERT-FT) GLUE performance.
achieve 99% ViT acc.
Llana 2
RI RI/RA RL
Method
GPT-2
GPT-2
Hybrid
Hyena
T2R-GPT-2
HH-GPT-2
Softmax (Zero-sbol)
/ 14.9
PPL
28.0
15.8
18.5
18.5
19.4
16.7
Softmax (LoRA)
51.1 27..2.4....
T2R (LoRA)
2.8/0.0/2.6
Table 10: Pretrained-converssoon for 125M GPT-2 on WT-103 lang. modeling. While finetuned GPT-2 gets lowest PPL., among sub- quadratic models Hedgehog significantly outperforms by 1.8 PPL.
Hedgehog (LoRA)
47.4/23A//
Table 11: Hedgehog Llama-2 conversion (ROUGE).
CONCLUSION
We present Hedgehog, a learnable linear attention mimic softmax attention. This enables training linear attention models from scratch and converting existing Transformers into linear attention vari- ants. To motivate Hedgehog we study why prior linear attentions underperform softmax attention, and identify two missing properties: (1) the ability to capture low entropy spiky attention maps and (2) to monotonic with respect to the underlying query-key dot products. We find training to match softmax attentions results in recovering many of its expressive properties, and that Hedgehog leads to competitive performance with softmax- based attention in training from scratch, finctuned- conversion, and pretrained conversion regimes.
10

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
We thank Armin Thomas, Gordon Downs, Krista Opsahl-Ong, Pun Waiwitlikhi, Schwinn Sacreesit- thipitak, Dan Fu, Simran Arora, Sabri Eyuboglu, Tri Dao, and anonymous reviewers for helpful discussions on linear attention and paper feedback. We also thank Dan Fu for prior versions of the pseudocode formatting in the appendix.
We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML)); US DEVCOM ARL under No. W911NF-21-2-025) (Interactive Human-Al Teaming); ONR un- der No. N000141712266 (Unifying Weak Supervision); ONR 500014-20-1-2488 and Applying Non-Euclidean Geometry in Machine Learning: NO001 120012225 (NEPTUNE); NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Cred- its for Research program, the Stanford Data Science Initiative (SDSI), and members the Stanford DAWN project: Meta, Google, and VMWare. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government.
Understanding
REFERENCES
Josh Alman and Zhao Song. Fast attention requires bounded entries. 1321:22302.12214. 2023.
arXiv preprint
Jimmy Ba, Geoffrey Hinton, Volodymyr Mnih, Joel Z Leibo, and Catalin Ionescu. Using fast weights to attend to the recent past. Advances in neural information processing systems, 29, 2016.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint urXiv:1409.0477 2014.
Kunal Banerjee, Vishak C., Rishi Raj Gupta, Kartik Vyas, Anushree H., and Biswajit Mishra. Exploring alternatives softmax function. ArXiv, abs/2011.11538, 2020. URL https:
//api.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle OBrien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 9650-9660, 2021.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15595, 2023.
Yifan Chen, Qi Zeng, Heng Ji, and Yun Yang. Skyformer: Remodel self-attention with gaus- sian kernel and nystr'" method. In A. Beygelzimer, Y. Dauphin, P. Liang, and Wort- man Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https: / //openzeview.....///...oe net /forfooria forumm wtt.mmmddddpptgggpppp
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint a0Xiv:2009. 14794, 2020.
Krzysztof Choromanski, Haoxian Chen, Han Lin, Yuanzhe Ma, Arijit Sehanobish. Deepali Jain, Michael S Ryoo, Jake Varley, Andy Zeng, Valerii Likhosherstov, et al. Hybrid random features. arXiv preprint 1rXiv:2110.04367. 2021.
11

---

Published as a conference paper at ICLR 2024
Zihang Dai, Zhilin Yang. Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer XL: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics. pp. -978-2988 Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 118633/11/119 1221. URL. https: / a/clclcl...o antho 1
P1 1285.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359 2022.
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized llms. arXiv preprint arXiv:2305. 14314, 2023.
Jacob Devlin, Ming-Weir Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint AXXv:20100.11999 2020.
Daniel Y Fu, Tri Dao, Khaled Kamal Saab, Armin W Thomas, Atri Rudra, and Christopher Re. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations, 2023. URL htttts:///ppeeeeetttt net/ pet/fozee forum? idd COZDyOWYGS
Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. SAMSum corpus: A human- annotated dialogue dataset for abstractive summarization. In Lu Wang. Jackic Chi Kit Cheung, Giuseppe Carenini, and Fei Liu (eds.), Proceedings of the 2nd Workshop on New Frontiers in Sum- marization, pp. 70-79, Hong Kong, China, November 2019. Association for Computational Lin- guistics. doi: 10. 18.13653/v I/D19-5409, URL httpt::///// htt///aaaaaanttooaoeeeeaea httphhoooooooooggooggooo/ 5409.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00330 2021.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, An- drea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for nlp. In International Conference on Machine Learning, pp. 2790-2799. PMLR, 2019.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint vrXiv:2/06.0966 2021.
Kazuki Irie, Imanol Schlag, Róbert Csordás, and Jürgen Schmidhuber. Going beyond linear trans- formers with recurrent fast weight programmers. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https:: wttp://eepeeeeeeeeee//eeee..eeee net / forum? weruumin==tt2RRii@/1
Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A. Smith. Finetuning pretrained transformers into RNNs. In Proceed- ings of the 202/ Conference on Empirical Methods in Natural Language Processing, pp. 10630- 10643, Online and Punta Cana, Dominican Republic, November 2021. Association for Computa- tional Linguisties. doi: 10. 10.18653/v1/2021.com www.pl-eemmmmmee.ee main.830. httLtttts////caatti org/202/ emnlp-main 830.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165. PMLR, 2020.
Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. In International Conference on Algorithmic Learning Theory, pp- 597-619. PMLR, 2023.
12

---

Published as a conference paper at ICLR 2024
Huanru Henry Mao. Fine-tuning pre-trained transformers into decaying fast weights. In Proceed- ings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 10236- 10242, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Lin- guistics. doi: 10.186534 20200222.mmpmmmimmi.... 2022. emn comnlp-maiin 697.
URL https: pettess:/aclll ogy. orrg/
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mix- ture models. International Conference on Learning Representations, 2017. URL https:
/ //openzeview...../oo.oe/oe.o.e/e net. wttt///ooommiddddeeee j72udxe.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209./1895, 2022.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.0211 2021.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. arXiv preprint arXivv2302. 2023.
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 7025-7041, Abu Dhabi, United Arab Emirates, De- cember 2022a. Association for Computational Linguisties. doi: 10. 18.53v 1/2022.mmmllmmaii. 473. URL https: /////aaaantttttt wwl.aannoooooyy....ooooooooooooa org/ 202gg2000 wwt.022emmmmmmaaam.... 473.
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Ling- peng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. arXiv preprint arXiv:2202.08791, 2022b.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models unsupervised multitask leamers. 2019. URL https://api.
1660022553..
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In J. Platt, D. Koller, Y. Singer, and S. Roweis (eds.), Advances in Neural In- formation Processing Systems, volume 20. Curran Associates, Inc., 2007. URL https:/ //proceed ngs neurips iss..pp... cc/paper_rpp..e hetlle/appee/////////// 013a006f03dbc533 92ef F1 (FFdaa755-a- Paperr.ppprr... pdf.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pp. 9355-9366. PMLR, 2021.
Yi Tay, Mostafa Dehghani. Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena A benchmark for efficient transformers. In International Conference on Learning Representations, 2021. URL https: / //openreview. net. worfmoorurm// wtw.:ffrrmmmiioeee wttp://orrmeeee////eeeee----geee
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov. Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint Stt.:22307.09288 2023.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis- Morency, and Ruslan Salakhutdinov. Transformer dissection: An unified understanding for transformer's attention via the lens of kernel. In Proceedings of the 2019 Conference on Empirical Methods in Natural Lan- guage Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNIPP, pp. 4344-4353, Hong Kong, China, November 2019. Association for Com- putational Linguistics. doi: 10. 10.18653/11/D199 1/1931444/ URL https: http://////aaaahoooooo.goo / D19-1443.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
13

---

Published as a conference paper at ICLR 2024
Alex Warstadt, Amanpreet Singh, and Samuel R Bowman. Neural network acceptability judgments. Transactions of the Association for Computational Lingusssics, 7:625-641, 2019.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nystrômformer: A nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 14138-14148, 2021.
Lin Zheng. Jianbo Yuan, Chong Wang, and Lingpeng Kong. Efficient attention via control variates. In The Eleventh International Conference on Learning Representations, 2023, URL https: /openeeview net/forum? id= uNfHKrj46
14

---

Published as a conference paper at ICLR 2024
A HEDGEHOG IMPLEMENTATION DETAILS
We provide further details on the Hedgehog feature map and attention weight distillation training.
A.1 MECHANICS FOR HEDGEHOG FEATURE MAP
To improve Hedgehog performance in practice, we explored variations along additional criteria for numerical stability and improved expressivity.
Numerical stability practice, we find that computing O as the softmax applied over the MLP output dimension also seems to work but with better stability. In this case, we expand Eq. 3 as
exp(w) x)
exp(w, x)
mmlp (a:)
(5)
E ii exp(w, x)
E exp( w
(also performing better than dividing each element by max over (exp(w, a + b))_
Negation mapping. To better compute dot products as a similarity measure between queries and keys, in practice we also + mapping from RRd R2dd e.g., via
(m))(a) exp x ex,,).,,,pp((.] x + b), exp(- x
b) (6)
where the additional negation mapping in Rdd intuitively lets better factor in negative dimen- sionalities, which prior linear attention feature maps like ReLU ignore. While this results larger feature dimension, it only scales fixed constant 2 such that the overall time and space complexity for Hedgehog linear attention is still ((nd). We further find that in practice, this still accomplishes favorable scaling and much faster inference with smaller memory than the Taylor exponential dis- cussed in Sec. 4.1 (see Fig. 6 for real-world wall-clock time and memory savings).
A.2 HEDGEHOG FEATURE MAP AND MODEL ARCHITECTURE
We apply Hedgehog feature maps for each head and layer individually in a standard Transformer architecture, where the addition of head-specific MLPs is akin inserting "adapters" (Houlsby et al., 2019) after every query and key projection. Each MLP is a single linear layer with input and output dimensions equal to the base Transformer's head dimension. Pytorch-like code is given below.
import torch import torch.nn as nn
class Hedgehoog (nn .nn.ule)e):
def
hem:_dim; int, activation: str 'exp') super () initt_ Trainable map self. nn Linea (head dim, head dim)
_
self.
11
def self. seif.initt
12
""*nitiialiee trainable map as identity* weight) mm.nnit.zero__ff layer.
13
15
16
def (self, x: torch. Tensor)
17
= = self. layer (x) $ shape b, h, 1, d
18
return torch. cat (Itorch. (x), torch. (-x) din=-1)
A.3 HEDGEHOG DISTILLATION AND FINETUNING IMPLEMENTATION DETAILS
We include additional details for training Hedgehog layers to obtain linear attention Transformers.
These fall under categories: (1) fram-snh-romm-craacch. and (2) finetuned pretrained conversion.
15

---

Published as a conference paper at ICLR 2024
1. Training-from--eeree hog MLP for each query and key projection of the randomly initialized Transformer (e.g., for each head of a multi-head attention layer, and for all such layers). We then train the Hedgehog MLPS jointly with all other model parameters end-to-end with a single objective function, egg. cross- entropy loss next-token prediction when training models for language modeling.
When training Hedgehog Transformers from scratch, insert Hedge-
2. Finetuned pretrained conversion. For both these regimes, we carry out training as a two stage process. Like mrining-from-scratchh we initially insert Hedgehog MLPs for query and key projections. Following this, we proceed in two stages:
Attention distillation. We first freeze the Transformer's original weights and specifically train the Hedgehog MLPs, such that the resulting linear attention weights match those produced via softmax attention over the same query and key tensors. For each head, we conceptually follow Listing below to compute soft cross-entropy KL-divergence between the "predicted" linear attention weights and "ground-truth" softmax attention weights. We compute these losses for each attention head and layer after forward pass of the entire model, using data samples from the target task. We find it sufficient to use one optimizer for joint training over all Hedgehog layers parallel, using the sum of each individual attention head distillation loss as the final criterion. This makes training simple and comparable to standard training loop; we further provide code Listing to do so with popular APIs such as HuggingFace Transformers'.
2. Original parameter finetuning. Following attention distillation, we simply unfreeze all model weights and train with a standard task-specific loss function. We find we can also keep certain layers frozen train with parameter-efffiient finetuning such low-rank adaptation (Hu et al., 2021); we explore this in Sec. 5.4 with Llama-2 models.
Hedgehog distillation loss for one attention head
def softmax_attn (q: torch. Tensor, k: "**Get softmax attention weights -> Assume k are both shape (b, h, 1, d) """
scale = q. shappee[11] ** 0.5
ok torch, torchh cinsum bhmd, bhnd->bhmn', q, k) scale return torch. softmax (qk, dim=-1)
def quaáratic interattcciinneaaatttat inner_attt (g: torch. Tensor, Norch.Tensory
10
=**
11
Get linear attention weights
12
-> Assume q, k are both shape (b, h, 1, d) and feature maps already applied
13
14
qk torch, einsum (' b'hmm bhnd->bhmm q, k)
15
return qk / qk. sum (dim=-1, keepdin-True)
def srapute_hedgehog.. loss (q: torch. Tensor,
18
k: torch.
19
hh_mlp_q: HedgehoggeastureMa.. hh_mlp_k: BedgehogFeaturrMMp))
20
***
22
Compute the attention distillation loss >> Assume (alternative) use Kı, divergence) Assume q and k are the queries and keys of pretrained Transformer,
23
is impl emented
24
27
e.g., via q self.q.proj (hidden_states) states)
28
29
true_attn = softmax_attn k)
prod_attn qud return soft waw.aaaooooo_.ccctrrrooooooodddo (pre_Ltttt, truc_atta)
(k)m
32
Listing 1: Hedgehog distillation loss for attention head
16

---

Published as a conference paper at ICLR 2024
Hedgehog Attention class for easy attention distillation
class HedgehogAttentiio (nn. Mon.M))le):
Sample code for Heedyhooattenttien, following HuggingFace API =**
def _ (self, base_atte training = True):
self base_attn base_atth attn # e.g., IlamaAttention
10
Trainable feature maps
11
self.mlp_q HedgchogFeaturoKkp (base_attn.head
12
self.mlp_k HedgehogFeaturekke (base_attn. aew.attttheeddii
Freeze original attention parameters for p self.base ptta.paraneters
15
16
p.requires grad = False
self.g_proo self @aas__attn.__pr0 pero. self haf.base__ttn.kk.rr
19
20
21
Whether we train attenticns or not self. taainiraiinng training
24
def forward
25
hidden_states: torch. Tensor, output_attentions: bool = True, +<base_kwargs: any):
26
29
if self, telf.training:
30
# Compute ground-truth attention weights outputs, true attns = base attn( hidden hidden_stateseeeeeeeeeeeeeeeeeee states=hidden states, output attenticns-True, **base_kwargs)
35
34
35
# Compute Redghog feature maps
q self. (self [self.q_proj (hidden_states) k self.mlp (self.k_pro) thidden_statess
38
39
40
attd_ quadratic Linear_ k)
42
it cutput_attentionn
Hook for attentions
43
return outputs, (pred_attns, true_attns)
44
End relevant
Listing 2: Hedgehog Attention class for easy attention distillation.
'In practice, to train all Hedgehog layers easily in a joint end-to-end fashion, we make use of pop- ular pretrained Transformer APIs such those in the HuggingFace transformers library. We im- plement Hedgehog equivalent of the base Transformers' attention class, which (1) abstracts away the Transformer-speeiff attention computation and (2) lets us hook attention weights calculated at True keyword args.
each layer to the model's final outputs. e.g., via output attentions We can subsequently substitute each attention layer with the "Hedgehog Attention equivalent, and train via simple loop over the data. We present Pytorch-like code in Listing 3.
17

---

Published as a conference paper at ICLR 2024
End-to-end joint attention distillation
from transformers import AutoModel
Load base model
base_moces model AutoModel, From fwtm from_pretrainedd...)
# Freeze original parameters
for p in basc_model; p prrequires_grad False
10
11 Convert attentions for all layers
12 for layer in base_model:
base_model..c attn - Hedgehogattention (base_model.
15 Define single optimizer for training all feature maps
16
optim = optimizer (base casse__odee..arr. parame..armmtters(
Train Hedgehog feature maps for data in dataloader:
19
20
21
Compute outputs and hook to attentions
outputs base sase__ddellleeeeeeeeee (**data, cutput_atten stttut__tttentissTrree outputs outputs. get ('attentiens')
24
25
0
26
for attns in enumerate (outputss: pred_attn, true_attn - attns
attentions for each layer
total_loss soft_labed ceff__labe__coss_etniee _rcossstttooooyyyoooooe (pred_att,, true_attn)
29
30
loss.backwardy Jointly optimize all feature maps
()tim)
Listing 3: End-to-end joint attention distillation.
B DEFERRED EXPERIMENTAL DETAILS
B.1 ASSOCIATIVE RECALL ANALYSTS (SECTION 3.2)
In Sec. 3.2, we compare various Transformers' abilities to solve Associative Recall (AR) (Ba et al., 2016), next-token prediction task previously studied a proxy for language modeling capabil- ity (Olsson et al., 2022). AR tests how well model can recall specific content in an input sequence, structured list of key-value pairs which ends in key Table 12.
Next
Vocab
Seq. Length
Input Sequence
Token
Size
9 99.33...ff flc
40
128
Table 12: Associative recall task. Example from Ba et al. (2016).
Dataset details. To understand the effects of more uniform attention weightings, we evaluate with 40 possible tokens and 128 sieen-long-sequencee such that models must recall pairings that only occur three times on average in-context. We generate 10,000 training samples following patterns described in Table 12, and evaluate 2000 newly-generated test samples (again using the same associative recall structure, but with different token associations).
Architecture details. For all experiements, we use a four layer Transformer with four heads-per layer, head dimension 64, and rotary embeddings. This is similar modern model families such
18

---

Published as a conference paper at ICLR 2024
as Pythia (Biderman et al., 2023) and LLaMA / Llama-2 (Touvron et al., 2023). We keep all parts consistent except for the multi-head attention, comparing popular linear attentions (c.f., Fig. 2).
Training details. For fair comparison to evaluate just the feature map / modeling architecture, train all models by sweeping learning rate {Ie-2, le-4), weight decay E {0, 5e-4), and batch size {8, 32} with AdamW optimizer. We train up to 100 epochs with early stopping (explicitly stopping training if validation loss stops decreasing after 10 epochs).
B.2 BERT-BASE FINETUNED ON COLA CONVERSION (SECTION 3.2)
Training details. For our fietuned-cooversion analysis, we replace the attentions of finetuned BERT-base-uncased model available on the HuggingFace model hub'. We train with batch size 8, learning rate 1e-5, zero weight decay, A optimizer, and up to 10 epochs with early stopping.
B.3 HEDGEHOG TRAINING FROM SCRATCH (SECTION 5.2)
LRA training and model details. On LRA, for fair comparison we implement Hedgehog in the existing PyTorch implementation provided by Xiong et al. (2021), deferring to the same model configurations and hyperparameters used in the original repository (Tay et al., 2021).
WikiText- 103 training and model details. For WikiText- 103, train 125M parameter GPT-2 style Transformer with learning rate 6e-4, weight decay 0.01, and AdamW optimizer. For close comparison, we follow the architectural details of GPT-2 125M, and use a 12 layer decoder-only network with 12 heads, head dimension = 64, hidden dimension 768, and MLP dimension 3072.
B.4 HEDGEHOG FINETUNED CONVERSION (SECTION 5.3)
Recovering finetuned BERT performance on GLUE tasks. For finetuned conversion, we first conduct Hedgehog attention distillation by training attention layers up to five epochs with early stop- ping based on validation loss. We train with learning rate Ie-2, weight decay Adam optimizer. We follow the same procedure for the Transformer-to-RNN (T2R) (Kasai et al., 2021) ablation. For regulard T2R and subsequently post attention distillation, we train each BERT model with batch size 8, learning rate le-5, weight decay 0, AdamW optimizer, and cosine scheduler for up to five epochs on the individual classification (all except STS-B) or regression tasks (STS-B) on the GLUE bench- mark. For all tasks, hosted at the HuggingFace model hub", and thank the original uploader for their contributions.
the corresponding available finetuned BERT- base-uncased checkpoints
Recovering finetuned Vision Transformer performance on ImageNet-1K. To demonstrate pnetuned--converiinn for the image domain, we use the vit-base pateeeppach----- 2224 checkpoint provided by Google HuggingFace' which is trained on ImageNet-21k before being finetuned on ImageNet- 1 K at resolution of 224 x 224 pixels (Dosovitskiy et al., 2020). For distillation, we freeze the original ViT weights, and train linear attention MLPS with batch size 32, learning rate 0.01, zero weight decay, and Adam W optimizer, and train for two epochs. We then train all parameters with learning rate 1e-3, zero weight decay and AdamW optimizer up 10 epochs with early stopping.
B.5 HEDGEHOG PRETRAINED CONVERSION (SECTION 5.4)
Linear GPT-2 125M conversion for WikiText-100% language modeling, We use the available GPT-2 125M pretrained checkpoint available on HuggingFace' from Radford et al. (2019). For Hedgehog, we first do attention distillation and train Hedgehog MLPs for two epochs over the WikiText-103 data, using batch size 8, learning rate 0.01, zero weight decay, AdamW optimizer, and 1024-tokens per input. For T2R-GPT-2 and the subsequent Bedgehog-GPT. model, we fine- tune all model parameters with learning rate 6e-4, weight decay 0.01, and Adam W optimizer and 1024 lokens-per-input
https:/huggingfacc www/eremiailZ/eettt------mmmm
"utps://huggginffac..c co/JeremiahZ/
bhtpss//huggingfacccco col hww.taaw...ooggggooeeeooo wwwggggll/vit-aaa-ppieee waaccc-pptchh--- 16-224
thtps:/hugginggfcc.com co/gpt2
19

---

Published as a conference paper at ICLR 2024
Linear Llama-2 7B conversion for SAMSum corpus summarization. We use the base Llama-2 7B model available via Meta and HuggingFace (1lama 1-lama22-7b- -2-7b-hf) from Touvron et al. (2023). For all experiments, and evaluations on a single A6000 GPU.
non-quantized model weights in bíloat16, and conduct all training runs
For dataset preparation, first convert individual document and summarization pairs into single next-token prediction samples, using the template in Listing 4. For both distillation and subsequent finetuning, we then chunk these samples into concatenated inputs 1024 tokens long.
For attention distillation, we freeze all original Llama-2 weights, and train Hedgehog MLPs for every head and layer (0.495% of the original model size). We then train for two epochs with learning rate 0.01, zero weight decay, Adam optimizer, and batch size 8 with gradient accumulation.
For finetuning and comparison to T2R and standard attention, we apply LoRA to query, key, value, and output projections of each layer. We use alpha parameter 16 and rank 8. We train with learning rate le-4, zero weight decay, Adam optimizer, and batch size with gradient accumulation.
For generation, we compute ROUGE metrics (R1, R2, RL; for overlap of unigrams, bigrams, and longest common subsequence) over model outputs. We generate sequences up 100 tokens long, and evaluate based on outputs up til the first 
 Llama stop token.
1lama-2 prompt template for SAMSUM
Summarize this dialog:
{input}
Summary: ) (cos_token)
Listing 4: Llama-2 prompt template for SAMSum corpus summarization
C ADDITIONAL RESULTS
C.I EXTENDED COMPARISON TO ATTENTION MODELS ON LRA
In Table 13, compare Hedgehog's performance LRA against fuller set of Transformer and subquadratic Transformer based alternatives sourced either from the official benchmark leader- board (Tay et al., 2021) or recent subquadratic attention works (where we display the most com- petitive alternatives in Table 6). We find Hedgehog on average obtains best accuracy. Although recently non-Transe -ro-Trasffrrmer models such deep state-space models have shown impressive results outperforming Transformers on LRA (Gu et al., 2021). as our work focuses on how to improve and recover the expressivity of standard softmax Transformers, we focus the comparison against other attention- based methods. We defer to Gu et al. (2021) and related works for their LRA results.
C.2 HEDGEHOG FEATURE MAP GENERALIZATION TO NEW DATA
We extend our analysis into how Hedgehog's s feature maps learned with one dataset generalize to attentions computed new dataset (c.f. Table 4 and Fig. in Sec. 5.1). As in the prior section, we find Hedgehog learned feature maps frequently generalize to new datasets. Despite training to match the softmax attentions on one model and dataset, we first find Hedgehog feature maps can produce attention weights that closely resemble softmax attention for the same model on another dataset (App. C.2.1). We next quantify this fidelity via KL divergence w.r.t. the softmax attentions (App. C.2.2). We find that Hedgehog learned feature maps almost always still generalize better than prior linear attention feature maps. We finally show that this attention matching generalization transfers to actual pretrained-conversin attentions with Hedgehog attentions trained on one task, and find finetuning with these converted models on another GLUE task still leads to improvements over prior linear attentions.
performance (App. C.2.3). We replace BERT-bass softmax
Setup. For all experiments, we begin by training Hedgehog attentions on "in-distribution" softmax model (Devlin al., 2018) as the Trans-
attention data. We use the pretrained
20

---

Published as a conference paper at ICLR 2024
Model
ListOps
Text
Retrieval
Image
Pathfinder
Average
Transformer
36.37
64.27
57.46
42.44
71.40
54.39
Local Att
15.82
52.98
53.39
41.46
66.63
46.06
Linear Trans.
16.13
65.90
53.09
42.34
75.30
50.55
Reformer
37.27
56.10
53.40
38.07
68.50
50.67
Sparse Trans.
17.07
63.58
59.59
44.24
71.71
51.24
Sinkhorn Trans.
33.67
61.20
53.83
41.23
67.45
51.29
Linformer
35.70
53,94
52.27
38.56
76.34
51.36
Performer
18.01
65.40
53.82
42.77
77.05
51.41
Synthesizer
36,99
61.68
54.67
41.61
69.45
52.88
Longformer
35.63
62.85
56.89
42.22
69.71
53.46
BigBird
36.05
64.02
59.29
40.83
74.87
55.01
Nystrmformer!
37.15
65.52
79.56
41.58
70.94
58.95
cosFormer!
37.90
63.41
61.36
43.17
70.33
55.23
Skyformer' 39.25
64.70
82.06
40.77
70.73
59.50
Hedgehog 37.15
64.60
82.24
40.15
74.16
59.66
Table 13: Training- wen-e-ffo---aate on LRA. Hedgehog achieves best average performance across Transformers and subquadratic variants. indicates method results reported from original works. other reported from the official LRA benchmark (Tay et al., 2021). Best, 2nd-best acc (%).
former we wish to convert, and distill two sets of Hedgehog attentions over (1) the GLUE CoLA task or (2) 512-token chunks of WikiText- 103 corpus. Thus, queries and keys computed with the FERT-baas--ncasse model over CoLA validation samples are "in-distribution" for the first set, and interested in seeing how attention weight fidelity downstream performance recovery affected when subsequently finetuning on non-CoLA GLUE data. We compare with various prior ablations and alternative feature maps, such as the Transformer-to-RNN. feature map (Kasai al., 2021) after attention distillation, Hedgehog without attention distillation, and prior representative linear attentions such as Performer (Choromanski al., 2020) and cosFormer (Qin al., 2022b).
C.2.1 QUALITATIVE EVIDENCE OF HEDGEHOG DATA
In Fig. 10 and Fig. 11, we visualize attention weights computed via various methods on heads in the 1st, 6th, and 12th layers of the BERT-base uncased model. We find Hedgehog can learn feature maps that lead to matching softmax attention weights, even when computed on new data samples. Interestingly, the Hedgehog feature maps result in significantly more similar attention weights versus alternative feature maps (quantified in the next section).
In addition, our comparisons to Hedgehog ablations suggest that the proposed Hedgehog feature map and distillation procedure are important for best generalization. Removing either the Hedgehog feature map form (via doing attention distillation using the prior Transformer-to-RNY feature map (T2R-HH) not training feature maps (HH (No Train)) leads to lower fidelity, where attention distillation seems critical for retaining weights reasonably similar to softmax attention.
C.2.2 QUANTITATIVE ANALYSIS OF HEDGEHOG DATA GENER ALIZATION
To quantify the above observations. compute the KL divergence between Hedgehog attention weights computed on various GLUE tasks and the "ground truth" softmax attention weights, using the pretrained the above visualizations, we find that Hedgehog feature maps do seem to produce better matching attention weights softmax attention via significantly smaller KI. divergences.
model. We report the KL divergence in Table 14. Similar to
C.2.3 HEDGEHOG DATA GENERALIZATION VIA GLUE TASK TRANSFER
We finally evaluate the Hedgehog attention generalization by finetuning the pretrained BERT mod- els with trained Hedgehog new GLUE tasks. We follow the same procedure described in Ap-
21

---

Published as a conference paper at ICLR 2024
Softmax
HH (CoLA)
HH (WT-103)
T2R-HH (CoLA) HH (No Train)
Pos ELU
Performer
CosFormer
Softmax
HH (CoLA)
HH (WT-103)
T2R-HH (CoLA) HH (No Train)
Pos ELU
Performer
CosFomee
Softmax
HH (CaLA)
HH (WT-103)
T2R-HH (CoLA) HH (No Train)
ELU
Performer
CosFomer
Qualitative generalization to MRPC. Attention weights for BERR-base-uneased queries and keys computed on MRPC samples. We compare attentions from the 3rd head in the 1st, 6th and I 12th layers (top, middle, bottom). Hedgehog feature maps trained on CoLA or WikiTText- 103 often still produce attention weights similar to those softmax attention on new data.
Figure 10:
Softmax
HH (CoLA)
HH (WT-103)
TZR-HH (CoLA) HH (No Train)
Pos ELU
Performer
Costomer
Softmax
HH (CoLA)
HH (WT-103)
T2R HH (CoLA) HH (No Train)
Pos ELU
Performer
Cosfomer
Softmax
HH (CoLA)
HH (WT-103) T2R-HH (CoLA) HH Train)
Pos ELU
Performer
CosFormer
Qualitative generalization QNLI. Attention weights for BERT-base- queries and keys computed on QNLI samples. We compare attentions from the 3rd head in the 1st, 6th and 12th layers (top, middle, bottom). Hedgehog feature maps trained on CoLA or WikiText-103 often still produce attention weights similar to those of softmax attention on new data.
Figure 11:
pendix B.4. In Table 15, we find that the above attention weight observations on Hedgehog general- ization also correspond with downstream task performance. BerdgehogBBRR' models achieve best or second-bee performance, despite using attention feature maps trained on different data. We leave further generalization studies, such how Hedgehog attentions trained on one model generalize an entirely different model for future work.
22

---

Published as a conference paper at ICLR 2024
Method
CoLA
MNLI
MRPC
QNLI
QQP
RTE
SST-2
STS-B
Hedgehog (CoLA)
0.173
0.340
0.652
0.673
0.374
0.932
0.211
0.275
Hedgehog (WT-103)
0.357
0.381
0.484
0.446
0.432
0.428
0.347
0.360
T2R-HH (CoLA)
0.191
0.603
1.124
1.141
0.658
1.357
0.254
0.444
Hedgehog Untrained
0.687
0.845
1.264
1.146
0.890
1.493
0.859
0.743
ELU Performer
1,231
1,500
2.150
1.947
1.509
2.491
1,505
1.285
1.293
1.592
2.239
2.059
1.588
2.594
1.558
1.352
CosFormer
1.191
1.341
1.987
1.840
1.330
2.398
1.443
1.142
Table 14: KL divergence attention weights generalizing new data. Hedgehog attentions trained on either CoLA (CoLA) or WikiText-105 (WT-103) data, still best match softmax attention weights computed on different GLUE tasks, despite being trained with task-specific data (measured via KL. divergence; lower is better).
Method
CoLA
MRPC
QNLI
QQP
RTE
SST-2
STS-B
Hedgehog (CoLA)
58.4
89.4
87.7
89.8
62.1
91.9
85.3
Hedgehog (WT-103)
47.2
89.4
89.2
90.4
62.5
91.4
86.7
HH (No Train)
50.3
83.3
85.9
86.5
55.6
89.5
79.3
1 ++ELU Performer CosFormer
26.8
81.9
78.5
89.1
55.6
85.9
41.8
24.7
81.4
75.8
86.5
55.6
85.1
39.8
41.1
82
82.6
89.3
54.9
88.4
76.6
Table 15: Attention generalization on downstream tasks. BERT models with Hedgehog atten- tions trained on either CoLA (CoLA) or WikiText-103 (WT-103) achieve best GLUE performance despite being finetuned on different GLUE tasks. This corresponds with prior observations gen- eralization via improved attention weight fidelity.
C.3 LLAMA-2 SAMSUM GENERATIONS
We include sample generations from the SAMSum corpus summarization task (Gliwa al., 2019), used to evaluate Hedgehog conversion of LLama-2 7B models in combination with low-rank adap- tation (LoRA). Via the generation quality, we find that contrast to prior conversion methods such as Transformer-to-RNN (T2R) (Kasai et al., 2021), Hedgehog makes pretrained-converiio with parameter-efficeen finetuning feasible on larger models.
We report generations for four test samples of the SAMSum test set (first 3, and longer 6th), com- paring standard attention Llama-2 models, linear attention Llama-2 models achieved via Hedgehog attention distillation, and linear attention Llama-2 models following T2R, all with LoRA finetuning on single A6000 GPU (unquantized b weights; LoRA on all query, key, value, and output projections).
For future work, we think further showing how Hedgehog and pretrained-coovvesionn can help apply larger models to tasks requiring longer contexts more efficient inference, promising avenues of exploration. Several considerations include additional explorations to improve generation quality, parameter training efficiency (e.g., combining Hedgehog with quantization (Dettmers al., 2023)), and further utilizing Hedgehog's linear attention efficiency via transfer to longer contexts (e.g., combining Hedgehog with positional interpolation (Chen et al., 2023)).
(Input) :

 Summarize this dialog:
Hannah: Hey, do you have Betty's number?
Amanda: Lemme check
Hannah: <file_gif>
Amanda: Scrry, can't find it.
23

---

Published as a conference paper at ICLR 2024
Amanda: Ask Larry
Amanda: Hе called her last time we were the park together
Hannah: I don't t know him well
Hannah: <file_gifx
Amanda: Don't be shy, he's very nice
Hannah: If you say so..
Hannah: I'd rather you texted him
Amanda: Just text him
Bannah: Urgh.. Alright
Eannah: Bye
Amanda: Bye bye
Summary:
(Llama-2):
Bannah looking for Betty's s number. Amanda can't find it. Amanda suggests Hannah to ask Larry.

(Hedgehog Llama-2):
Hannah is looking for Betty's s number. Amanda will text her.

(T2R lllma-2):
along recently acquired the biggest, I'tch. SCHPW182seud is the biggest, I't SCHPentua. SCHPentie't SCHPentua.
along ch.
along ch.
along ch.
along ch.
along ch. along ch.
Listing 5: SAMSum Llama-2 generations after LoRA finetuning.
24

---

Published as a conference paper at ICLR 2024
(Input):

 Summarize this dialog:
Bric: MACHINE!
Rob: That's gr8!
Eric: know! And shows how Americans see Russian ;)
Rob: And it's really funny!
Eric: I know! I especially like the train part!
Rob: Hahaha! No one talks to the machine like that!
Eric: this his only stand-up?
Rob: Iak. I'll check.
Eric: Surc.
Rob: Turns cut no! There are some of his stand-ups on youtube.
Eric: Gr8! I'11 watch them now!
Rob: Me too!
Eric: MACHINE!
Rob: MACHINE!
Eric: TTYL?
Rob: Sure
Summary:
(Llama-2) 2)
Eric and Rob watching stand-up by Russian comedian, 

(Hedgehog Llama-2):
Eric and Rob are watching stand-up conedy by a Rassian comedian. 

(T2R ilama-2)
along recently acqui red the biggest,
I'tch. SCHPentus. SCHPentua. SCHPentua. SCHPentua. SCHP
along ch.
along ch.
along ch. SCHP
along ch.
along ch.
along ch.
along for:
along for:
al
Listing SAMSum Llama-2 generations after LoRA finetuning.
25

---

Published as a conference paper at ICLR 2024
(Input):

 Summarize this dialog:
lenny: Babe, can you help me with semething?
Bob: Sure, what's up?
Lenny: Which one should I pick?
Bob: Send me photos
Lenny: <file_photo>
Lenny: <file_photo>
Lenny: <filc_photo>
Bob: I like the first ones best
Lenny: But already have purple trousers. Does it make sense to have two pairs?
Bob: I have four black pairs :D :D
lenny: yeah, but shouldn't pick different color?
what you' 11 give you the most outfit options Lenny: So I guess I'11 puy the first or the third pair then Bob: Pick the best quality then Lenny: right, thx Bob: no prob
Bob: what matters
Summary:
(Llama-2) Lenny locking new trousers. Bob suggests the first or the third pair.

(Heögehog Illmaa2)
kenny wants to buy pair of trousers. Hob suggesta the first two options.<//s
(T2R Llama-2):
along, I't a. SCHPentua. SCH? along ch. SCHPentua.
along ch. SCHP
along ch.
along ch.
along ch.
along for:
along ch.
along ch.
along for:
along ch.
along for:
Listing SAMSum Llama-2 generations after LoRA finetuning.
26

---

Published as a conference paper at ICLR 2024
(Input):

 Summarize this dialog:
Benjamin: Hey guys, what are we dcing with the keys today?
Hilary: I've got them. whoever wants them can meet at lunchtime after
Elliot: I'm ck, We' re meeting for the drinks the evening аnyway and I quess we'll be going back to the apartment together?
Hilary: Yeah, I guess
Daniel: I'm with Hilary atm and won't let go of her for the rest of the day, so any option you guys cheese is good for me
Benjamin: Hmm I might actually pass by at lunchtime, take the keys and go take a nap. I'm s00o tired after yesterday
Hilary: Scunds good. We'11 be having lunch with some French people (the ones who work on the history of food in colonial Mexico alreacy see you yawning your head cif)
Benjamin: YAAAANN
Where and where are you meeting?
Hilary: So I'm meeting them at the entrance to the conference hall at pm and then we'll head to this place called La Cantina. Italian cuisine, which quite funny, but that's what they' chosen Benjamin: Interesting
To be honest, Hilary, I almost feel like changing my mind. Wanting to take this nap might end up costing me to dear
Hilary: Oh come on Benjamin: All these terrible obstacles on mu way to bed might just prove to much to take Hilary: We' 11 try to avoid talking about their subject of research. 0h wait,
I'm actually meeting them because wanted to chat about
their research 1ol
Elliot:
Hilary: Do join us, we' re going to have fun. And then you' 11 take the
keys and take this most deserved of naps
Elliot: Sounds like plan
Hilary:
Elliot: See you at then xx
(Llama-2):
Bilary has the keys. Benjamin is tired after yesterday. Hilary is meeting seme French people at pm. Benjamin will join them, 

(Heagehog Llama-2):
Hilary and Elliot hall at pm. Hilary and Benjamin will meet there. Hilary and Benjamin will take the keys to the apartment. Hilary and Benjamin will meet Ellict at 2 pm. Hilary and Benjamin will take nap. 

going to meet the entrance the conference
(T2R Llama-2):
Most is the biggest, I's:
Most is the biggest, I's:
Most is the biggest, I's:
Most is the biggest, I's:
Most is the biggest, I's:
Most is the biggest, I's: M
Listing SAMSum Llama-2 generations after LoRA finetuning.
C.4 ADDITIONAL. ATTENTION WEIGHT VISUALIZZATIONS
We finally include additional visualizations of the attention weights computed via softmax attention in comparison to Hedgehog and alternate linear attention feature maps. We visualize attentions computed GLUE tasks (Sec. 5.4) from the 1st, 6th, and 12th (first, middle, last) layers of BERT models in top, middle, and bottom rows respectively, and for the 1st, 6th, and 12th heads.
27

---

Published as a conference paper at ICLR 2024
Softmax
Hedgehog
(WT-103)
Train)
Performer
CosFormer
Softmax
Hedgehog
(WT-103)
(No Train)
Performer
CosFormer
Softmax
Hedgehog
HH (WT-1003)
HH (No Train)
ReLU
Pos
Performer
CosFormer
Figure 12: BERT attention visualizations for CoLA. Head I; 1st, 6th, and 12th layers.
Softmax
Hedgehog
(WT-103)
HH Train)
ReLU
Pos ELU
Performer CosFormer
Softmax
Hedgehog
(WT-100)
Train)
ReLU
Pos FLU
Perfermer
CosFormer
Softmax
Hedgehog
(WT-103)
Train)
ReLU
Pos ELU
Performer
CosFormer
Figure 13: BERT attention visualizations for CoLA. Head 6; 1st, 6th, and 12th layers.
Softmax
Hedgehog
HH (WT-103)
HH (No Train)
ReLl
Pos ELU
Performer CosFormer
Softmax
Hedgehog
HH (WT-1003)
HH (No Train)
ReLU
Pos ELU
Performer
CosFormer
Softmax
Hedgehog
HH (WT-103)
HH (No Train)
ReLU
Pos ELU
Pertormer
CosFormer
Figure 14: BERT attention visualizations for CoLA. Head 12; 1st, 6th, and 12th layers.
28

---

Published as a conference paper at ICLR 2024
Softmax
Hedgehog
(WT-103)
Train)
Performer
CosFormer
Softmax
Hedgehog
(WT-1033)
(No Train)
ReLU
Performer
CosFormer
Softmax
Hedgehog
HH (WT-103)
HH (No Train)
ReLU
Pos
Performer
CosFormer
Figure 15: BERT attention visualizations for RTE. Head 0; 1st, 6th, and 12th layers.
Softmax
Hedgehog
(WT-103)
HH (No Train)
ReLU
Pos ELU
Performer CosFormer
Softmax
Hedgehog
(WT-1003
HH (No Train)
ReLU
Pos ELU
Perfermer
CosFormer
Softmax
Hedgehog
(WT-103)
Train)
ReLU
Pos ELU
Perfermer
CosFormer
Figure 16: BERT attention visualizations for RTE, Head 6; 1st, 6th, and 12th layers.
Softmax
Hedgehog
HH (WT-103)
HH (No Train)
ReLl
Pos ELU
Performer
CosFormer
Softmax
Hedgehog
HH (WT-103)
HH (No Train)
ReLU
Pos ELU
Performer
CosFormer
Softmax
Hedgehog
HH (WT-1033
HH (No Train)
ReLU
Pos ELU
Pertormer
CosFormer
Figure 17: BERT attention visualizations for RTE. Head 12; 1st, 6th, and 12th layers.
29

---

Published as a conference paper at ICLR 2024
Softmax
Hedgehog
(WT-1033
Train)
Performer
CosFormer
Softmax
Hedgehog
(WT-103)
(No Train)
ReLU
Performer
CosFormer
Softmax
Hedgehog
HH 103-103)
HH (No Train)
ReLU
Pos
Performer
CosFormer
Figure 18: BERT attention visualizations for SST2. Head 1; 1st, 6th, and 12th layers.
Softmax
Hedgehog
(WT-103)
HH Train)
ReLU
Pos ELU
Performer CosFormer
Softmax
Hedgehog
(WT-100)
Train)
ReLU
Pos ELU
Perfermer
CosFormer
Softmax
Hedgehog
(WT-103)
Train)
ReLU
Pos ELU
Perfermer
CosFormer
Figure 19: BERT attention visualizations for SST2, Head 6; 1st, 6th, and 12th layers.
Softmax
Hedgehog
HH (WT-103)
HH (No Train)
ReLU
Pos ELU
Performer CosFormer
Softmax
Hedgehog
HH (WT-103)
HH (No Train)
ReLU
Pos ELU
Performer
CosFormer
Softmax
Hedgehog
HH (WT-103)
HH (No Train)
ReLU
Pos ELU
Pertormer
CosFormer
Figure 20: BERT attention visualizations for SST2, Head 12; 1st, 6th, and 12th layers.
30


Reviews:
[{"review_id": "Txr5iZ6dGK", "reviewer": "Reviewer_ZtxG", "summary": "This paper address the quadratic scaling of attention mechanism in transformers. There are 2 main contributions:\n\n1. Defining characteristics that make an attention mechanism work well: (a) low-entropy \"spikyness\" and (b) dot-product monotonicity\n\n2. Two new attention mechanisms to evaluate their claims. Taylor and Hedgehog\n\n They evaluate attention mechanisms in several scenarios (1) Training-from-scratch, (2) Finetuned-conversion, and (3) Pretrained-conversion.", "questions": "There doesn't seem to be much downside to this method except for training from scratch? Should everyone be finetuning and converting their pretrained models to Hedgehog?\n\nWas there any investigation into the learned weights of Hedgehog to gain intuition on how it's approximating softmax?\n\nHow does Hedgehog's softmax approximation do with length generalization (sequence length longer than training data) or data that is out of distribution?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The exploration into \"spikyness\" and monotonicity is interesting and novel to me.\n\nThe regime for evaluating their attention mechanism is thorough and covers a wide variety of situations.\n\nThe ablations are well done and the experimental results on benchmarks and efficiency are strong.", "weaknesses": "I found the presentation a little confusing. In some figures and tables only Hedgehog is shown. In others only Taylor is shown. Should both be evaluated?\n\nIs the GPT-2 column in Table 6 unfair? It seems they should finetune with the full attention mechanism?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4M84HaQsnE", "reviewer": "Reviewer_y6Q1", "summary": "This paper proposes *Hedgehog*, a new linear attention mechanism for Transformers. \n\nThe authors start by noting that previous linear attention mechanisms still lag considerable behind full quadratic attention. The authors then investigate why this is the case, by studing the properties of the attention maps produced by full attention and linear approximations. They find that two properties in full attention that previous linear attention methods seem to lack:\n\n- low-entropy \"spiky\" attention weights\n- monotonicity over query-key dot products\n\nThe authors then attempt to measure the importance of these two properties by inducing these properties in some test tasks, finding that these properties seem to correlate with performance of the trained model.\n\nFinally, the authors propose Hedgehog: the core idea being using linear map from the query/key space into a lower-dimensional feature space, but learning this linear map to *explicitly* approximate the final attention maps produced by dense attention. \n\nExperimental results are quite convincing: Hedgehog outperforms most other linear attention methods on multiple tasks (including LM perplexity and popular downstream tasks), both when training models from scratch and when “linearizing” models with dense attention.", "questions": "- In the train-from-scratch setting, how is the Hedgehog attention initialize?. From my understanding, Hedgehog requires an already-trained dense attention model to finetune the feature map. Does this imply that in this scenario we still need to train a dense attention model anyway?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper provides a quite consistent story: they identified a problem (current linear-attention LLMs/transformers don’t work that well), studied potential reasons for their failure mode and verified their importance for the performance of models (the two properties of attention) and proposed and experimentally verify a fix to these problems that leads to linear-attention models that get closer to dense attention models)\n- The paper is well-written, and the problem studied very relevant: the quadratic complexity is one of the main bottlenecks preventing the scaling of transformers to longer contexts.\n- The identified properties of dense attention (and the subsequent analysis of their relation to downstream performance) are super interesting (and could be a contribution almost independently of the proposed method)\n- The proposed Hedgehog method seems to be effective for the model-size studied in this work, with Hedgehog leading to much less performance degradation than other linear-attention approaches (when compared to dense-attention model)", "weaknesses": "- My main criticism of the study is that they only studied relatively small transformer models:  they only studied 125M decoder only models. Previous work [1] has shown that linear attention approach are less amenable to scaling than dense attention (meaning as models get larger, the gap between dense and linear attention increases) and so it’s unclear if this approach would still lead to such small gaps with dense attentions for 1B+ models (which IMO, is the main use case for this technique). While I already really liked the paper, I think this paper could be much stronger impact on the community if results were validated on larger LLMs.\n\n[1] https://arxiv.org/pdf/2207.10551.pdf", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "a9VIdkeDR1", "reviewer": "Reviewer_PsQx", "summary": "The authors found that prior linear attentions miss two key properties of softmax attention that are crucial for good performance: having low-entropy (spiky) weights and maintaining dot-product monotonicity. A new learnable linear attention called Hedgehog was introduced. Hedgehog retains these crucial properties while keeping linear complexity. It uses trainable MLPs to produce attention weights similar to softmax attention. The results indicate that Hedgehog can achieve over 99% of the standard Transformer's performance, outdoing previous linear attentions in various settings.", "questions": "Since the function mimicking softmax is learned how scalable it is to other modalities?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Technically sound well-written paper dealing with an important research task.", "weaknesses": "Table 3 is confusing since there are many more higher performing solutions for LRA: https://paperswithcode.com/sota/long-range-modeling-on-lra\nTaylor expansion in transformers is well-known approach: eg. https://arxiv.org/pdf/2206.08898.pdf\nWhile Hedgehog aims to retain the properties of softmax attention in a linear attention framework, this could introduce its own complexities in terms of implementation and optimization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)