LoLCATs: On Low-Rank Linearizing of Large Language Models

## Abstract

Abstract Recent works show we can linearize large language models (LLMs)—swapping the quadratic attentions of popular Transformer-based LLMs with subquadratic analogs, such as linear attention—avoiding the expensive pretraining costs.
However, linearizing LLMs
often significantly degrades model quality, still requires training over billions of tokens, and remains limited to smaller 1.3B to 7B LLMs.
We thus propose Low-rank Linear Conversion via Attention Transfer ( \name ), a simple two-step method that improves LLM linearizing quality with orders of magnitudes less memory and compute.
We base these steps on two findings.
First, we can replace an LLM’s softmax attentions with closely-approximating linear attentions, simply by training the linear attentions to match their softmax counterparts with an output MSE loss (“ attention transfer ”).
Then, this enables adjusting for approximation errors and recovering LLM quality simply with low-rank adaptation (LoRA). \name significantly improves linearizing quality, training efficiency, and scalability. We significantly reduce the linearizing quality gap and produce state-of-the-art subquadratic LLMs from Llama 3 8B and Mistral 7B v0.1, leading to 20+ points of improvement on 5-shot MMLU. Furthermore, \name does so with only 0.2% of past methods’ model parameters and 0.4% of their training tokens.
Finally, we apply \name to create the first linearized 70B and 405B LLMs (50 × \times × larger than prior work).
When compared with prior approaches under the same compute budgets, \name significantly improves linearizing quality, closing the gap between linearized and original Llama 3.1 70B and 405B LLMs by 77.8% and 78.1% on 5-shot MMLU.

## 1 Introduction

“Linearizing” large language models (LLMs)—or converting existing Transformer-based LLMs into attention-free or subquadratic alternatives—has shown promise for scaling up efficient architectures. While many such architectures offer complexity-level efficiency gains, like linear-time and constant-memory generation, they are often limited to smaller models pretrained on academic budgets [, , , , ].
In a complementary direction,
linearizing aims to start with openly available LLMs—e.g., those with 7B+ parameters trained on trillions of tokens [, ]—and (i) swap their softmax attentions with subquadratic analogs, before (ii) further finetuning to recover quality. This holds exciting promise for quickly scaling up subquadratic capabilities.

However, to better realize this promise and allow anyone to convert LLMs into subquadratic models, we desire methods that are
(1) quality-preserving, e.g., to recover the zero-shot abilities of modern LLMs; (2) parameter and token efficient, to linearize LLMs on widely accessible compute; and (3) highly scalable, to support linearizing the various 70B+ LLMs available today [, ].

Existing methods present opportunities to improve all three criteria.
On quality, despite using motivated subquadratic analogs such
as RetNet-inspired
linear attentions [, ]
or state-space model (SSM)-based
Mamba layers [, , ],
prior works significantly reduce performance on popular LM Evaluation Harness tasks (LM Eval) [] (up to 23.4-28.2 pts on 5-shot MMLU []).
On parameter and token efficiency, to adjust for architectural differences, prior methods update all model parameters in at least one stage of training [, , ], and use 20 - 100B tokens to linearize 7B LLMs.
On scalability, these training costs make linearizing larger models on academic compute more difficult;
existing works only linearize up to 8B LLMs. This makes it unclear how to support linearizing 70B to 405B LLMs [].

Figure: Figure 1: \name framework. We linearize LLMs by (1) training attention analogs to approximate softmax attentions (attention transfer), before swapping attentions and (2) minimally adjusting (with LoRA).
Refer to caption: x1.png

In this work, we thus propose \name
(Low-rank Linear Conversion with Attention Transfer), a simple approach to improve the quality, efficiency, and scalability of linearizing LLMs.
As guiding motivation, we ask if we can linearize LLMs by simply reducing architectural differences, i.e.,

- 1.
Starting with simple softmax attention analogs such as linear attention (Eq. 2), and training their parameterizations explicitly to approximate softmax attention (“attention transfer”).
- 2.
Subsequently only training with low-cost finetuning to adjust for any approximation errors, e.g., with low-rank adaptation (LoRA) [] (“low-rank linearizing”).

In evaluating this hypothesis, we make several contributions.
First, to better understand linearizing feasibility, we empirically study attention transfer and low-rank linearizing with existing linear attentions. While intuitive—by swapping in perfect subquadratic softmax attention approximators, we could get subquadratic LLMs with no additional training—prior works suggest linear attentions struggle to match softmax expressivity [, ]
or need full-model updates to recover linearizing quality [, ].
In contrast, we find that while either attention transfer or LoRA alone is insufficient, we can rapidly recover quality by simply doing both (Figure 4, Table 1).
At the same time, we do uncover quality issues related to attention-matching architecture and training. With prior linear attentions, the best low-rank linearized LLMs still significantly degrade in quality
vs. original Transformers (up to 42.4 pts on 5-shot MMLU).
With prior approaches that train all attentions jointly [], we also find that later layers can result in 200$\times$ the MSE of earlier ones (Figure 8). We later find this issue aggravated by larger LLMs; jointly training all of Llama 3.1 405B’s 126 attention layers fails to viably linearize the LLM.

Next, to resolve these issues and improve upon our original criteria, we detail \name’ method components.
For quality, we generalize prior notions of learnable linear attentions to sliding window + linear attention variants. These remain subquadratic to compute yet consistently yield better attention transfer via lower mean-squared error (MSE) on attention outputs. For parameter and token efficiency, we maintain our simple 2-step framework of (1) training subquadratic attentions to match softmax attentions,
before (2) adjusting for any errors via only LoRA. For scalability, we use finer-grained “block-by-block”
training. We split LLMs into blocks of $k$ layers before jointly training attentions only within each block to improve layer-wise attention matching. We pick $k$ to balance the speed of training blocks in parallel with the memory of saving hidden state outputs of prior blocks (as inputs for later ones). We provide a simple cost model to navigate these tradeoffs.

Finally, in experiments, we validate that \name improves on
each of our desired criteria.

- •
On quality, when linearizing popular LLMs such as Mistral-7B and Llama 3 8B, \name significantly improves past linearizing methods (by
1.1$-$8.6
points (pts) on zero-shot LM Eval tasks;
+17.2 pts
on 5-shot MMLU)). With Llama 3 8B, \name for the first time closes the zero-shot LM Eval gap between linearized and Transformer models
(73.1 vs 74.2 pts),
while supporting 3$\times$ higher throughput and 64$\times$ larger batch sizes vs. popular FlashAttention-2 [] implementations (generating 4096 token samples on an 80GB H100). We further validate \name as a high-quality training method, outperforming strong 7B subquadratic LLMs (RWKV-v6 [], Mamba [], Griffin []) and hybrids (StripedHyena [], Zamba []) trained from scratch by
1.2 to 9.9 pts
on average over popular LM Eval tasks.
- •
On parameter and token-efficiency, by only training linear attention feature maps in Stage 1, while only using LoRA on linear attention projections in Stage 2, \name enables these gains while updating only $<$0.2% of past linearizing methods’ model parameters (doable on a single 40GB GPU). This also only takes 40M tokens, i.e., 0.003% and 0.04% of prior pretraining and linearizing methods’ token counts.
- •
On scalability, with \name we scale up linearizing to support Llama 3.1 70B and 405B LLMs []. \name presents the first viable approach to linearizing larger LLMs. We create the first linearized 70B LLM, taking only 18 hours on one 8$\times$80GB H100 node, and the first linearized 405B LLM with
a combination of 5 hours on 14 80GB H100 GPUs (attention transfer) +
16 hours on three 8$\times$80GB H100 nodes (LoRA finetuning) for Llama 3.1 405B.
For both models, this amount to under half the total GPU hours than prior methods reported to linearize 8B models (5 days on 8$\times$80GB A100s) []. Furthermore, under these computational constraints, \name significantly improves quality versus prior linearizing approaches without attention transfer. With Llama 3.1 70B and 405B, we close
77.8% and 78.1% of the 5-shot MMLU gap between Transformers and linearized variants respectively.

Our code is available at: [https://github.com/HazyResearch/lolcats](https://github.com/HazyResearch/lolcats).

Figure: Figure 2: Linearizing comparison. \name significantly improves LLM linearizing quality and training efficiency. Multiple ✓ or ✗ denoting relatively better or worse support.
[Uncaptioned image]: extracted/5954737/figures/lm_eval_all.png

## 2 Preliminaries

To motivate \name, we first go over Transformers, attention, and linear attention. We then briefly discuss related works on linearizing Transformers and Transformer-based LLMs.

##### Transformers and Attention.

Popular LLMs such as Llama 3 8B [] and Mistral 7B [] are decoder-only Transformers, with repeated blocks of multi-head softmax attention followed by MLPs [].
For one head, attention computes outputs $\bm{y}\in\mathbb{R}^{l\times d}$ from inputs $\bm{x}\in\mathbb{R}^{l\times d}$ (where
$l$ is sequence length,
$d$ is head dimension) with query, key, and value weights $\bm{W}_{q},\bm{W}_{k},\bm{W}_{v}\in\mathbb{R}^{d\times d}$. In causal language modeling, we compute $\bm{q}=\bm{x}\bm{W}_{q}$, $\bm{k}=\bm{x}\bm{W}_{k}$, $\bm{v}=\bm{x}\bm{W}_{v}$, before getting attention weights $\bm{a}$ and outputs $\bm{y}$ via

$$ $a_{n,i}=\frac{\exp(\bm{q}_{n}^{\top}\bm{k}_{i}/\sqrt{d})}{\sum_{i=1}^{n}\exp( \bm{q}_{n}^{\top}\bm{k}_{i}/\sqrt{d})},\;\;\;\bm{y}_{n}=\sum_{i=1}^{n}a_{n,i} \bm{v}_{i},\;\;\;\text{for $n$ in $[1,\ldots,l]$}$ (1) $$

Multi-head attention maintains inputs, outputs, and weights for each head, e.g., $\bm{x}\in\mathbb{R}^{h\times l\times d}$ or $\bm{W}_{q}\in\mathbb{R}^{h\times d\times d}$ ($h$ being number of heads), and computes Eq. 1 for each head. In both cases, we compute final outputs by
concatenating $\bm{y}_{n}$ across heads, before using output weights $\bm{W}_{o}\in\mathbb{R}^{hd\times hd}$ to compute $\bm{y}_{n}\bm{W}_{o}\in\mathbb{R}^{l\times hd}$.
While expressive, causal softmax attention requires all $\{\bm{k}_{i},\bm{v}_{i}\}_{i\leq n}$ to compute $\bm{y}_{n}$.
For long context or large batch settings,
this growing KV cache can incur prohibitive memory costs even with state-of-the-art implementations such as FlashAttention [].

##### Linear Attention.

To get around this, Katharopoulos et al. [] show a similar attention operation, but with linear time and constant memory over generation length (linear time and space when processing inputs).
To see how, note that
softmax attention’s exponential is a kernel function $\mathcal{K}(\bm{q}_{n},\bm{k}_{i})$, which in general can be expressed as the dot product of feature maps $\phi:\mathbb{R}^{d}\mapsto\mathbb{R}^{d^{\prime}}$. Swapping $\exp(\bm{q}_{n}^{\top}\bm{k}_{i}/\sqrt{d})$ with $\phi(\bm{q}_{n})^{\top}\phi(\bm{k}_{i})$ in Eq. 1 gives us *linear attention* weights
and outputs:

$$ $\hat{a}_{n,i}=\frac{\phi(\bm{q}_{n})^{\top}\phi(\bm{k}_{i})}{\sum_{i=1}^{n} \phi(\bm{q}_{n})^{\top}\phi(\bm{k}_{i})},\;\;\;\hat{\bm{y}}_{n}=\sum_{i=1}^{n} \hat{a}_{i,n}\bm{v}_{i}$ (2) $$

Rearranging terms via matrix product associativity, we get

$$ $\hat{\bm{y}}_{n}=\sum_{i=1}^{n}\frac{\phi(\bm{q}_{n})^{\top}\phi(\bm{k}_{i}) \bm{v}_{i}}{\sum_{i=1}^{n}\phi(\bm{q}_{n})^{\top}\phi(\bm{k}_{i})}=\frac{\phi( \bm{q}_{n})^{\top}\Big{(}\sum_{i=1}^{n}\phi(\bm{k}_{i})\bm{v}_{i}^{\top}\Big{) }}{\phi(\bm{q}_{n})^{\top}\sum_{i=1}^{n}\phi(\bm{k}_{i})}$ (3) $$

This lets us compute both the numerator $\bm{s}_{n}=\sum_{i=1}^{n}\phi(\bm{k}_{i})\bm{v}_{i}^{\top}$ and denominator $\bm{z}_{n}=\sum_{i=1}^{n}\phi(\bm{k}_{i})$ as recurrent “KV states”. With $\bm{s}_{0}=\bm{0},\bm{z}_{0}=\bm{0}$, we recurrently compute linear attention outputs as

$$ $\hat{\bm{y}}_{n}=\frac{\phi(\bm{q}_{n})^{\top}\bm{s}_{n}}{\phi(\bm{q}_{n})^{ \top}\bm{z}_{n}}\;\;\text{for}\;\;\bm{s}_{n}=\bm{s}_{n-1}+\phi(\bm{k}_{n})\bm{ v}_{n}^{\top}\;\;\text{and}\;\;\bm{z}_{n}=\bm{z}_{n-1}+\phi(\bm{k}_{n})\;\;\;$ (4) $$

Eq. 3 lets us compute attention over an input sequence of length $n$ in $\mathcal{O}(ndd^{\prime})$ time and space, while Eq. 4 lets us compute $n$ new tokens in $\mathcal{O}(ndd^{\prime})$ time and $\mathcal{O}(dd^{\prime})$ memory.
Especially during generation, when softmax attention has to compute new tokens sequentially anyway, Eq. 4 enables time and memory savings if $d^{\prime}<$ (prompt length + prior generated tokens).

##### Linearizing Transformers.

To combine efficiency with quality, various works propose different $\phi$, (e.g., $\phi(x)=1+\text{ELU}(x)$ as in  []). However, they typically train linear attention Transformers from scratch.
We build upon recent works that swap the softmax attentions of existing Transformers with linear attention before finetuning the modified models with next-token prediction to recover language modeling quality.
These include methods proposed for LLMs [], and those for smaller Transformers—e.g., 110M BERTs [])—reasonably adaptable to modern LLMs [, , ].

## 3 Method: Linearizing LLMs with \name

We now study how to build a high-quality and highly efficient linearizing method.
In Section 3.1, we present our motivating framework,
which aims to (1) learn good softmax attention approximators with linear attentions and (2) enable low-rank adaptation for recovering linearized quality.
In Section 3.2, we find that while this attention transfer works surprisingly well for low-rank linearizing with existing linear attentions, on certain tasks, it still results in sizable quality gaps compared to prior methods.
We also find that attention-transfer quality strongly corresponds with the final linearized model’s performance. In Section 3.3, we use our learned findings to overcome prior issues, improving attention transfer to subsequently improve low-rank linearizing quality.

### 3.1 A Framework for Low-cost Linearizing

In this section, we present our initial \name framework for linearizing LLMs in an effective yet efficient manner. Our main hypothesis is that by first learning linear attentions that approximate softmax, we can then swap these attentions in as drop-in subquadratic replacements. We would then only need a minimal amount of subsequent training—e.g., that is supported by low-rank updates—to recover LLM quality in a cost-effective manner effectively. We thus proceed in two steps.

- 1.
Parameter-Efficient Attention Transfer. For each softmax attention in an LLM, we aim to learn a closely-approximating linear attention, i.e., one that computes attention outputs $\hat{\bm{y}}\approx\bm{y}$ for all natural inputs $\bm{x}$.
We treat this as a feature map learning problem, learning $\phi$ to approximate softmax. For each head and layer, let $\phi_{q},\phi_{k}$ be query, key feature maps. Per head, we compute:
$\bm{y}_{n}=\underbrace{\sum_{i=1}^{n}\frac{\exp(\bm{q}_{n}^{\top}\bm{k}_{i}/
\sqrt{d})}{\sum_{i=1}^{n}\exp(\bm{q}_{n}^{\top}\bm{k}_{i}/\sqrt{d})}\bm{v}_{i}
}_{\text{Softmax Attention}},\;\;\;\;\hat{\bm{y}}_{n}=\underbrace{\sum_{i=1}^{
n}\frac{\phi_{q}(\bm{q}_{n})^{\top}\phi_{k}(\bm{k}_{i})}{\sum_{i=1}^{n}\phi_{q
}(\bm{q}_{n})^{\top}\phi_{k}(\bm{k}_{i})}\bm{v}_{i}}_{\text{Linear Attention}}$
(5)
for all $n\in[l]$ with input $\in\mathbb{R}^{l\times d}$,
and train $\phi_{q},\phi_{k}$ to minimize sample mean squared error (MSE)
$\ell_{\text{MSE}}=\frac{1}{MH}\sum_{m=1}^{M}\sum_{h=1}^{H}\ell_{\text{MSE}}^{h
,m}\;\;,\;\;\;\;\ell_{\text{MSE}}^{h,m}=\frac{1}{d}\sum_{n=1}^{d}(\bm{y}_{n}-
\hat{\bm{y}}_{n})^{2}$
(6)
i.e., jointly for each head $h$ in layer $m$.
Similar to past work [, ], rather than manually design $\phi$, we parameterize each $\phi:\mathbb{R}^{d}\mapsto\mathbb{R}^{d^{\prime}}$ as a learnable layer:
$\phi_{q}(\bm{q}_{n}):=f(\bm{q}_{n}\tilde{\bm{W}}_{(q)}+\tilde{\bm{b}}_{(q)})\;
\;,\;\;\;\phi_{k}(\bm{k}_{i}):=f(\bm{k}_{i}\tilde{\bm{W}}_{(k)}+\tilde{\bm{b}}
_{(k)})$
Here $\tilde{\bm{W}}\in\mathbb{R}^{d\times d^{\prime}}$ and $\tilde{\bm{b}}\in\mathbb{R}^{d^{\prime}}$ are trainable weights and optional biases,
$f(\cdot)$ is a nonlinear activation, and $d^{\prime}$ is an arbitrary feature dimension (set to equal head dimension $d$ in practice).
- 2.
Low-rank Adjusting. After training the linearizing layers, we replace the full-parameter training of prior work with low-rank adaptation (LoRA) [].
Like prior work, to adjust for the modifying layers and recover language modeling quality, we now train the modified LLM end-to-end over tokens to minimize a sample next-token prediction loss
$\ell_{\text{xent}}=-\sum\log P_{\Theta}(\bm{u}_{t+1}\mid\bm{u}_{1:t})$.
Here $P_{\Theta}$ is the modified LLM, $\Theta$ is the set of LLM parameters, and we aim to maximize the probability of true $\bm{u}_{t+1}$ given past tokens $\bm{u}_{1:t}$ (Fig. 1 right).
However, rather than train all LLM parameters, we only train the swapped linear attention $\bm{W}_{q},\bm{W}_{k},\bm{W}_{v},\bm{W}_{o}$ with LoRA.
Instead of full-rank updates, $\bm{W}^{\prime}\leftarrow\bm{W}+\Delta\bm{W}$, LoRA decomposes $\Delta\bm{W}$
as the product of two low-rank matrices $\bm{B}\bm{A}$, $\bm{B}\in\mathbb{R}^{d\times r}$, $\bm{A}\in\mathbb{R}^{r\times d}$.
For parameter efficiency,
we aim to pick small $r\ll d$.

Training footprint and efficiency. Both steps remain parameter-efficient.
For Step 1, optimizing Eq. 6 is similar to a layer-by-layer cross-architecture distillation. We compute layer-wise ($\bm{x},\bm{y}$) as pretrained attention inputs and outputs, using an LLM forward pass over natural language samples (Fig. 1 middle).
However, to keep our training footprint low, we freeze the original pretrained attention layer’s parameters and
simply insert new $\phi_{q},\phi_{k}$ after $\bm{W}_{q},\bm{W}_{k}$ in each softmax attention (Fig. 1 left). We compute outputs $\bm{y},\hat{\bm{y}}$ with the same attention weights in separate passes (choosing either Eq. 1 or Eq. 3; Fig. 1 middle).
For Llama 3 8B or Mistral 7B, training $\phi_{q},\phi_{k}$ with $d^{\prime}=64$ then only takes 32 layers $\times$ 32 heads $\times$ 2 feature maps $\times$ (128 $\times$ 64) weights $\approx$ 16.8M trainable weights (0.2% of LLM sizes). For Step 2, LoRA with $r=8$ on all attention projections suffices for state-of-the-art quality. This updates just $<$0.09% of 7B parameters.

### 3.2 Baseline Study: Attention Transfer and Low-rank Linearizing

We now aim to understand if attention transfer and low-rank adjusting are sufficient for linearizing LLMs.
It is unclear whether these simple steps can lead to high-quality LLMs, given that
prior works default to more involved approaches [, , ]. They use linearizing layers featuring GroupNorms [] and decay factors [], or alternate SSM-based architectures [, ]. They also all use full-LLM training after swapping in the subquadratic layers.
In contrast, we find that simple linear attentions can lead to viable linearizing, with attention transfer + LoRA obtaining competitive quality on 4 / 6 popular LM Eval tasks.

Figure: Figure 4: Attention transfer training efficiency. Even accounting for initial training steps, low-rank linearizing with attention transfer still consistently achieves lower perplexity faster across feature maps and LLMs.
Refer to caption: extracted/5954737/figures/train_efficiency_llama_t2r_logppl.png

**Table 1: Alpaca validation set perplexity (PPL) of linearized LLMs, comparing attention transfer, no LoRA adjusting (PPL@0) and PPL after training (PPL@2/4; 2 with attention transfer, 4 without, for equal total steps).**
|  | Llama 3 8B | Mistral 7B |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Attention | T2R | Hedgehog | T2R | Hedgehog |  |  |  |  |
| Transfer? | PPL@0 | PPL@2/4 | PPL@0 | PPL@2/4 | PPL@0 | PPL@2/4 | PPL@0 | PPL@2/4 |
| No  ✗ | 1539.39 | 16.05 | 2448.01 | 9.02 | 2497.13 | 8.85 | 561.47 | 4.87 |
| Yes  ✓ | 79.33 | 4.11 | 60.86 | 3.90 | 32.78 | 3.29 | 18.94 | 3.04 |

##### Experimental Setup.

We test the \name framework by linearizing two popular base LLMs, Llama 3 8B [] and Mistral 7B v0.1 []. For linearizing layers, we study two feature maps used in prior work (T2R [] and Hedgehog [], Table 2).
To support the rotary positional embeddings (RoPE) [] in these LLMs, we apply the feature maps $\phi$ after RoPE,(^2^22 Unlike prior works that apply $\phi$ before RoPE [, ], our choice preserves the linear attention kernel connection, where we can hope to learn $\phi_{q},\phi_{k}$ for $\exp(\bm{q}_{n}^{\top}\bm{k}_{i}/\sqrt{d})\approx\phi_{q}(\bm{q}_{n})^{\top} \phi_{k}(\bm{k}_{i})$.)
i.e., computing query features $\phi_{q}(\bm{q})=f(\text{RoPE}(\bm{q})\tilde{\bm{W}}_{q}+\tilde{\bm{b}}))$.
For linearizing data, we wish to see if \name with a small amount of data can recover general zero-shot and instruction-following LLM abilities. We use the 50K samples of a cleaned Alpaca dataset(^3^33[https://huggingface.co/datasets/yahma/alpaca-cleaned](https://huggingface.co/datasets/yahma/alpaca-cleaned)), due to its ability to improve general instruction-following in 7B LLMs despite its relatively small size []. Following [], we train all feature maps jointly.
We include training code and implementation details in App. C).

**Table 2: Learnable feature maps. Transformer to RNN (T2R) from [], Hedgehog from [], both $\oplus$ (concat) and $\text{SM}_{d}$ (softmax) apply over feature dimension.**
| Feature Map | $\phi(\bm{q})$ (same for $\bm{k}$) | Weight Shapes |
| --- | --- | --- |
| T2R | $\text{ReLU}(\bm{q}\tilde{\bm{W}}+\tilde{\bm{b}})$ | $\tilde{\bm{W}}$: (128, 128), $\tilde{\bm{b}}$: (128,) |
| Hedgehog | $[\text{SM}_{d}(\bm{q}\tilde{\bm{W})}\oplus\text{SM}_{d}(-\bm{q}\tilde{\bm{W})}]$ | $\tilde{\bm{W}}:(128,64)$ |

To study the effects of attention transfer and low-rank linearizing across LLMs and linear attention architectures, we evaluate their validation set perplexity (Table 1, Fig. 4) and downstream LM Eval zero-shot quality (Table 6).
We train both stages with the same data, evaluate with early stopping, and use either two epochs for both attention transfer and LoRA adjusting or four epochs with either alone ($\approx$ 40M total training tokens). For LoRA, we use $r=8$ as a popular default [], which results in training 0.2% of LLM parameter counts.

##### Attention Transfer + LoRA Enables Fast LLM Linearizing.

In Table 1 and Fig. 4, we report the validation PPL of linearized LLMs, ablating attention transfer and LoRA adjusting. We find that while attention transfer alone is often insufficient (c.f., PPL@0, Table 1), a single low-rank update rapidly recovers performance by 15–75 PPL (Fig. 4),
where training to approximate softmax leads to up to 11.9 lower PPL than no attention transfer.
Somewhat surprisingly, this translates to performing competitively with prior linearizing methods that train all model parameters [, ] (within 5 accuracy points on 4 / 6 popular LM Eval tasks; Table 6), while only training with 0.04$\%$ of their token counts and 0.2% of their parameter counts. The results suggest we can linearize 7B LLMs at orders-of-magnitude less training costs than previously shown.

Figure: Figure 5: Linearizing comparison on LM Eval. Task names in Table 5. Acc. norm: ARC-C, HS. Acc. otherwise. 5-shot MMLU. 0-shot otherwise.
Refer to caption: extracted/5954737/figures/mse_scatter_llama3_study.png

Figure: (a) Attention Entropy
Refer to caption: extracted/5954737/figures/method_entropy_vs_mse_linear_wide.png

##### LoL SAD : Limitations of Low-Rank Linearizing.

At the same time, we note quality limitations with the present framework.
While sometimes close, low-rank linearized LLMs perform worse than full-parameter alternatives and original Transformers on 5 / 6 LM Eval tasks (up to 42.4 points on 5-shot MMLU; Table 6).
To understand the issue, we study whether the attention transfer stage can produce high-fidelity linear approximations of softmax attention. We note four observations:

- 1.
Attention transfer quality (via output MSE) strongly correlates with the downstream low-rank linearization quality (Fig. 6). This suggests that MSE is a useful way to measure quality during attention transfer.
- 2.
Larger attention output MSEs coincide with lower softmax attention weight entropies (Fig. 7(a)). However, prior work shows that lower-entropy distributions are difficult to approximate with linear attentions [].
- 3.
Larger MSEs also heavily concentrate in attention input samples from later layers (Fig. 7(b)). Since the MSE metric is scale-sensitive, the later layers could impact our ability to train the earlier layers, using the simple loss summation across layers in Equation 6.
- 4.
The variation in MSE magnitudes increases with model scale. In Tables 17 and 18, we measure the magnitude of the MSE loss by layer depth at the 70B and 405B scales. The magnitudes range between 0.02 – 48.66 at the 405B scale versus 0.0008 – 2.91 at the 70B scale.

We thus hypothesize we can improve quality by reducing MSE in two ways.
First, to approximate samples with lower softmax attention weight entropies—i.e., “spikier” distributions more challenging to capture with linear attentions []—we may need better attention-matching architectures.
Second, to address the difficulty of learning certain attention layers, we may need more fine-grained layer-wise supervision instead of jointly training all layers, especially for larger models.

### 3.3 \name : Improved Low-rank Linearizing

We now introduce two simple improvements in architecture (Section 3.3.1) and linearizing procedure (Section 3.3.2) to improve low-rank linearizing quality.

#### 3.3.1 Architecture: Generalizing Learnable Linear Attentions

As described, we can apply our framework with any linear attentions with learnable $\phi$ (e.g., T2R and Hedgehog, Figure 4).
However, to improve attention-matching quality, we introduce a hybrid $\phi$ parameterization combining linear attention and sliding window attention. Motivated by prior works that show quality improvements when combining attention layers with linear attentions [, ],
we combine short sliding windows of softmax attention [, ] (size 64 in experiments) followed by linear attention in a single layer. This allows attending to all prior tokens for each layer while keeping the entire LLM subquadratic. For window size $w$ and token indices $[1,\ldots,n-w,\ldots,n]$, we apply the softmax attention over the $w$ most recent tokens,
and compute attention outputs $\hat{\bm{y}}_{n}$ as

$$ $\hat{\bm{y}}_{n}=\frac{{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{ rgb}{0,0,1}\sum_{i=n-w+1}^{n}\gamma\exp(\bm{q}_{n}^{\top}\bm{k}_{i}/\sqrt{d}- \bm{c}_{n})\bm{v}_{i}}+\phi_{q}(\bm{q}_{n})^{\top}\big{(}\sum_{j=1}^{n-w}\phi_ {k}(\bm{k}_{j})\bm{v}_{j}^{\top}\big{)}}{{\color[rgb]{0,0,1}\definecolor[named ]{pgfstrokecolor}{rgb}{0,0,1}\sum_{i={n-w+1}}^{n}\gamma\exp(\bm{q}_{n}^{\top} \bm{k}_{i}/\sqrt{d}-\bm{c}_{n})}\;+\phi_{q}(\bm{q}_{n})^{\top}\big{(}\sum_{j=1 }^{n-w}\phi_{k}(\bm{k}_{j})^{\top}\big{)}}$ (7) $$

$\gamma$ is a learnable mixing term, and $\bm{c}_{n}$ is a stabilizing constant as in log-sum-exp calculations ($\bm{c}_{n}=\max_{i}\big{\{}\bm{q}_{n}^{\top}\bm{k}_{i}/\sqrt{d}:i\in[n-w+1,
\ldots,n]\big{\}}$).
Like before, we can pick any learnable $\phi$.

Figure: Figure 8: Improving Attention matching MSE. Linearizing with linear + sliding window attention better matches LLM softmax attentions (lower MSE) over attention entropy values and LLM layers.
Refer to caption: extracted/5954737/figures/method_entropy_vs_mse_linear_fix.png

##### Subquadratic efficiency.

The hybrid layer retains linear time and constant memory generation. For $n$-token prompts, we initially require $\mathcal{O}(w^{2}d)$ and $\mathcal{O}((n-w)dd^{\prime})$ time and space for window and linear attention respectively, attending over a $w$-sized KV-cache and computing KV and K-states (Eq. 4).
For generation, we only need $\mathcal{O}(w^{2}d+dd^{\prime})$ time and space for every token. We evict the KV-cache’s first $\bm{k}$, $\bm{v}$, compute $\phi_{k}(\bm{k})$, and add $\phi_{k}(\bm{k})\bm{v}^{\top}$ and $\phi_{k}(\bm{k})$ to KV and K-states respectively.

#### 3.3.2 Training: Layer (or Block)-wise Attention Transfer

We describe the \name training approach and provide a simplified model to show its cost-quality tradeoffs.
Based on the limitations of low-rank linearizing identified in the prior section, we perform attention transfer at a block-by-block granularity for the large scale models. The MSE loss function is scale-sensitive and our study showed that the later layers consistently result in larger MSEs (Fig. 7(b)).
Thus, instead of computing the training loss (Eq. 6) over
all $m\in[M]$ for a model with $M$ layers, we compute the loss over $k$-layer blocks, and train each block independently:

$$ $\ell_{\text{MSE}}^{\text{block}}=\frac{1}{kH}\sum_{m=i}^{i+k}\sum_{h=1}^{H} \ell_{\text{MSE}}^{h,m}\;\;\;\;(\text{for blocks starting at layers }i=0,k,2k, \ldots),$ (8) $$

where the training data for the block starting at layer $i$ is the hidden states outputted by prior block starting at layer $i-k$, for each sequence in the training corpus. After attention transfer, we arrange the blocks back together in their original order. Like before, we then do end-to-end LoRA finetuning to recover quality.

**Table 3: Attention transfer for Llama 3.1 405B using $c$ blocks of $k$ layers.**
| Method | $k\times c$ | PPL |
| --- | --- | --- |
| Joint | 126$\times$1 | 4.23 |
| Block-wise | 9$\times$14 | 3.21 |

##### Tradeoffs between quality and efficiency.

We show that the block-wise approach improves linearization quality at large model scales and improves the user’s ability to flexibly balance compute and memory efficiency tradeoffs. We compare (1) joint ($k=126$) training, where we load the full model once and compute the loss as the sum of layer-wise MSEs, and (2) block-wise ($k=9$) training, where we break the model into blocks and train each independently.

For quality, we compare the evaluation perplexity of both the attention transfer approaches, after LoRA fine-tuning on the same subset of RedPajama data. The block-wise and joint approaches perform similarly at the 8B and 70B scales, however, as shown in Table 3, the block-wise approach performs 1.02 points better than joint attention transfer at the 405B scale. These results support our study in Section 3.2, which shows the variation in MSE magnitudes grows large at the 405B scale.

Next, for efficiency, we compare the compute and memory use for both methods:

- •
Compute:
While the joint training of Llama 3.1 405B in 16-bit precision uses multiple nodes (e.g., NVIDIA H100 8$\times$80GB nodes), an individual block of $k=9$ or fewer layers can be trained on a single GPU (e.g., H100 80GB GPU), at sequence length 1024. The block-wise approach is also amenable to using a different training duration (and training data) per block.
- •
Memory: To train block $i$, we need to save the hidden states from the prior block to disk. The total amount of disk space required is
$2\times T\times d\times\frac{L}{k}$
for total training tokens $T$, model dimension $d$, number of layers $L$ and 2-byte (16-bit) precision.
At the Llama 3.1 405B scale, saving states per-layer ($k=1$) for just $50$M tokens would require over 200TB of disk space, while $k=126$ (joint training) requires no additional disk space for the states.

## 4 Experiments

Through experiments, we study: (1) if \name linearizes LLMs with higher quality than existing subquadratic alternatives and linearizations, and higher generation efficiency than original Transformers (Section 4.1); (2) how ablations on attention transfer loss, subquadratic architecture, and parameter and token counts impact LLM downstream quality (Section 4.2); (3) how \name’ quality and efficiency holds up to 70B and 405B LLMs, where we linearize and compare model quality across the complete Llama 3.1 family (Section 4.3).

### 4.1 Main Results: \name Efficiently Recovers Quality in Linearized LLMs

In our main evaluation, we linearize the popular base Llama 3 8B [] and Mistral 7B [] LLMs. We first test if \name can efficiently create high-quality subquadratic LLMs from strong base Transformers, comparing to existing linearized LLMs from prior methods. We also test if \name can create subquadratic LLMs that outperform modern Transformer alternatives pretrained from scratch. For space, we defer linearizing training details to App. A.

**Table 4: \name comparison among linearized 7B+ LLMs. Among linearized 7B+ LLMs, \name-linearized Mistral 7B and Llama 3 8B consistently achieve best or 2nd-best performance on LM Eval tasks (only getting 2nd best to Mamba-Transformer hybrids). \name closes the Transformer quality gap by 79.8% (Mistral 7B) and 86.6% (Llama 3 8B) (average over all tasks; numbers except Hedgehog cited from original works), despite only using 40M tokens to linearize (a 2,500$\times$ improvement in tokens-to-model efficiency).**
| Model | Training<br>Tokens (B) | PiQA | ARC-e | ARC-c<br>(norm) | HellaSwag<br>(norm) | Wino-<br>grande | MMLU<br>(5-shot) | Avg. | Avg.<br>(no MMLU) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Mistral 7B | - | 82.1 | 80.9 | 53.8 | 81.0 | 74.0 | 62.4 | 72.4 | 74.4 |
| Mistral 7B SUPRA | 100 | 80.4 | 75.9 | 45.8 | 77.1 | 70.3 | 34.2 | 64.0 | 69.9 |
| Mistral 7B LoLCATs (Ours) | 0.04 | 81.5 | 81.7 | 54.9 | 80.7 | 74.0 | 51.4 | 70.7 | 74.5 |
| Llama 3 8B | - | 79.9 | 80.1 | 53.3 | 79.1 | 73.1 | 66.6 | 72.0 | 73.1 |
| Mamba2-Llama 3 | 20 | 76.8 | 74.1 | 48.0 | 70.8 | 58.6 | 43.2 | 61.9 | 65.6 |
| Mamba2-Llama 3, 50% Attn. | 20 | 81.5 | 78.8 | 58.2 | 79.5 | 71.5 | 56.7 | 71.0 | 73.9 |
| Llama 3 8B Hedgehog | 0.04 | 77.4 | 71.1 | 40.6 | 66.5 | 54.3 | 24.2 | 55.7 | 62.0 |
| Llama 3 8B LoLCATs (Ours) | 0.04 | 80.9 | 81.7 | 54.9 | 79.7 | 74.1 | 52.8 | 70.7 | 74.2 |

**Table 5: \name comparison to pretrained subquadratic LLMs. \name-linearized Mistral 7B and Llama 3 8B outperform pretrained Transformer alternatives by 1.2 to 9.9 points (Avg.), only training 0.2% of their parameter counts on 0.013 to 0.003% of their training token counts. ^∗Reported in Mercat et al. [].**
| Model | Tokens (B) | PiQA | ARC-easy | ARC-c<br>(acc. norm) | HellaSwag<br>(acc. norm) | Winogrande | MMLU<br>(5-shot) | Avg.<br>(w MMLU) | Avg.<br>(no MMLU) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Transformer |  |  |  |  |  |  |  |  |  |
| Gemma 7B | 6000 | 81.9 | 81.1 | 53.2 | 80.7 | 73.7 | 62.9 | 72.3 | 74.1 |
| Mistral 7B | 8000* | 82.1 | 80.9 | 53.8 | 81.0 | 74.0 | 62.4 | 72.4 | 74.4 |
| Llama 3 8B | 15000 | 79.9 | 80.1 | 53.3 | 79.1 | 73.1 | 66.6 | 72.0 | 73.1 |
| Subquadratic |  |  |  |  |  |  |  |  |  |
| Mamba 7B | 1200 | 81.0 | 77.5 | 46.7 | 77.9 | 71.8 | 33.3 | 64.7 | 71.0 |
| RWKV-6 World v2.1 7B | 1420 | 78.7 | 76.8 | 46.3 | 75.1 | 70.0 | - | 69.4 | 69.4 |
| TransNormerLLM 7B | 1400 | 80.1 | 75.4 | 44.4 | 75.2 | 66.1 | 43.1 | 64.1 | 68.2 |
| Hawk 7B | 300 | 80.0 | 74.4 | 45.9 | 77.6 | 69.9 | 35.0 | 63.8 | 69.6 |
| Griffin 7B | 300 | 81.0 | 75.4 | 47.9 | 78.6 | 72.6 | 39.3 | 65.8 | 71.1 |
| Hybrid Softmax |  |  |  |  |  |  |  |  |  |
| StripedHyena-Nous-7B | - | 78.8 | 77.2 | 40.0 | 76.4 | 66.4 | 26.0 | 60.8 | 67.8 |
| Zamba 7B | 1000 | 81.4 | 74.5 | 46.6 | 80.2 | 76.4 | 57.7 | 69.5 | 71.8 |
| Linearized |  |  |  |  |  |  |  |  |  |
| Mistral 7B LoLCATs (Ours) | 0.04 | 81.5 | 81.7 | 54.9 | 80.7 | 74.0 | 51.4 | 70.7 | 74.5 |
| Llama 3 8B LoLCATs (Ours) | 0.04 | 80.9 | 81.7 | 54.9 | 79.7 | 74.1 | 52.8 | 70.7 | 74.2 |

In Table 5, we report results on six popular LM Evaluation Harness (LM Eval) tasks [].
Compared to recent linearizing methods, \name significantly improves quality and training efficiency across tasks and LLMs. On quality, \name closes
79.8% and 86.6%
of the Transformer-linearizing gap for Mistral 7B and Llama 3 8B respectively, notably improving 5-shot MMLU by
60.9% and 40.9%
over next best fully subquadratic models (17.2 and 9.6 points). On efficiency, we achieve these results while only training $<$0.2% of model parameters via LoRA versus prior full-parameter training. We also only use 40M tokens versus the prior 20 – 100B (a 500 – 2500$\times$ improvement in “tokens-to-model” efficiency).
Among all 7B LLMs, \name-linearized LLMs further outperform strong subquadratic Transformer alternatives, representing RNNs or linear attentions (RWKV-v6 [], Hawk, Griffin [], TransNormer []), state-space models (Mamba []), and hybrid architectures with some full attention (StripedHyena [], Zamba []).

### 4.2 \name Component Properties and Ablations

We next validate that \name linearizing enable subquadratic efficiency, and study how each of \name’ components contribute to these linearizing quality gains.

Figure: (a) Batch size vs. Throughput
Refer to caption: extracted/5954737/figures/benchmark_thruput_bs.png

##### Subquadratic Generation Throughput and Memory.

We measure the generation throughput and memory of \name LLMs, validating that linearizing LLMs can significantly improve their generation efficiency. We use the popular Llama 3 8B HuggingFace checkpoint(^4^44[https://huggingface.co/meta-llama/Meta-Llama-3-8B](https://huggingface.co/meta-llama/Meta-Llama-3-8B)), and compare \name implemented in HuggingFace Transformers with the supported FlashAttention-2 (FA2) implementation []. We benchmark on a single 80GB H100 and benchmark two \name implementations with the Hedgehog feature map and (linear + sliding window) attention in FP32 and BF16.
In Fig. 9(a) and Fig. 9(b), we report the effect of scaling batch size on throughput and memory. We measure throughput as $($newly generated tokens $\times$ batch size $/$ total time$)$, using 128 token prompts and 4096 token generations. As batch size scales, \name-linearized LLMs achieve significantly higher throughput than FA2. We note this is primarily due to lower memory, where FA2 runs out of memory at batch size 64. Meanwhile, \name supports up to 3000 tokens / second with batch size 2048 (Fig. 9(a)), only maintaining a fixed “KV state” as opposed to the growing KV cache in all attention implementations (Fig. 9(b)).

##### Ablations.

We study how adding the attention transfer and linear + sliding window attention in \name contribute to downstream linearized LLM performance, linearizing Llama 3 8B over 1024-token long samples (Table 6). We start with standard linear attentions (Hedgehog, []; T2R, []), using the prior linearizing procedure of just swapping attentions and finetuning the model to predict next tokens []. We then add either (i) attention transfer, (ii) linear + sliding window attentions, or (iii) both, and report the average LM Eval score over the six popular zero-shot tasks in Table 5 and 5-shot MMLU accuracy. Across feature maps, we validate the \name combination leads to best performance.

**Table 6: \name component ablations, linearizing Llama 3 8B over 1024-token sequences. Default configuration highlighted. Across Hedgehog and T2R feature maps, \name’ attention transfer and sliding window increasingly improve linearized LLM quality on average LM Eval (Avg.) and 5-shot MMLU scores.**
| Feature Map | LM Eval<br>Metric | Swap &<br>Finetune | +Attention<br>Transfer | +Sliding Window,<br>+Attention Transfer | + Sliding Window,<br>No Attention Transfer |
| --- | --- | --- | --- | --- | --- |
|  | Avg. Zero-Shot | 44.20 | 55.32 | 70.66 | \ul68.78 |
| Hedgehog | MMLU (5-shot) | 23.80 | 23.80 | 52.77 | 45.80 |
|  | Avg. Zero-Shot | 38.84 | \ul54.83 | 68.28 | 39.52 |
| T2R | MMLU (5-shot) | 23.20 | 23.10 | 40.70 | 23.80 |

### 4.3 Scaling Up Linearizing to 70B and 405B LLMs

We finally use \name to scale up linearizing to Llama 3.1 70B and 405B models.
In Table 7, we find that \name provides the first practical solution for linearizing larger LLMs, achieving significant quality improvements over prior linearizing approaches of swapping in attentions and finetuning []. Controlling for the same linear + sliding window layer, for Llama 3.1 70B we achieve a
39.0
point improvement in 5-shot MMLU accuracy. For Llama 3.1 405B, \name similarly achieves a 38.3 point improvement. These results highlight \name’ ability to linearize large-scale models with greater efficiency and improved performance, showing for the first time that we can scale up linearizing to 70B+ LLMs.

**Table 7: Linearizing Llama 3.1 70B and 405B. Among the first linearized 70B and 405B LLMs (via low-rank linearizing), \name significantly improves zero- and few-shot quality.**
|  | PiQA | ARC Easy | ARC Challenge | HellaSwag | WinoGrande | MMLU (5-shot) |
| --- | --- | --- | --- | --- | --- | --- |
|  | acc | acc | (acc norm) | (acc norm) | acc | acc |
| Llama 3.1 8B | 79.87 | 81.52 | 53.58 | 79.01 | 73.48 | 66.11 |
| Linearized, no attn. transfer | 78.67 | 78.11 | 49.83 | 77.83 | 68.51 | 51.44 |
| \name (Ours) | 80.96 | 82.37 | 54.44 | 79.07 | 69.69 | 54.88 |
| Llama 3.1 70B | 83.10 | 87.30 | 60.60 | 85.00 | 79.60 | 78.80 |
| Linearized, no attn. transfer | 81.99 | 80.89 | 54.44 | 82.29 | 71.19 | 28.74 |
| \name (Ours) | 82.10 | 84.98 | 60.50 | 84.62 | 73.72 | 67.70 |
| Llama 3.1 405B | 85.58 | 87.58 | 66.21 | 87.13 | 79.40 | 82.98 |
| Linearized, no attn. transfer | 84.44 | 86.62 | 64.33 | 86.19 | 79.87 | 33.86 |
| \name (Ours) | 85.58 | 88.80 | 67.75 | 87.41 | 80.35 | 72.20 |

## 5 Conclusion

We propose \name, an efficient LLM linearizing method that (1) trains attention analogs—such as linear attentions and linear attention + sliding window hybrids—to approximate an LLM’s self-attentions, before (2) swapping the attentions and only finetuning the replacing attentions with LoRA.
We exploit the fidelity between these attention analogs and softmax attention, where we reduce the problem of linearizing LLMs to learning to approximate softmax attention in a subquadratic analog.
Furthermore, we demonstrate that via an MSE-based attention output-matching loss, we *are able* to train such attention analogs to approximate the “ground-truth” softmax attentions in practice.
On popular zero-shot LM Evaluation harness benchmarks and 5-shot MMLU, we find this enables producing high-quality, high-inference efficiency LLMs that outperform prior Transformer alternatives while only updating 0.2% of model parameters and requiring 0.003% of the training tokens to achieve similar quality with LLM pretraining. Our findings significantly improve linearizing quality and accessibility, allowing us to create the first linearized 70B and 405B LLMs.

#### Limitations and Future Work

While we focus on studying how to enable high quality yet highly efficient LLM linearizing with simple linear attentions, we note several areas for additional evaluation in both subquadratic capabilities and architectures. On subquadratic capabilities, by replacing each attention layer alternative, we eliminate the need to manage growing key-value (KV) caches and their associated memory overheads. However, it remains to be seen what kinds of capabilities we can enable with this cheaper inference, e.g., if linearized models can exploit quality-improving inference scaling laws suggested by recent works [, ]. Under a different motivation, while layers like linear attention achieve greater efficiency gains over softmax attention when processing longer contexts, we leave studying how low-rank linearizing applies to such long context scenarios as a motivated direction for future work. Finally, while we stick to “vanilla” linear + sliding window attentions in \name, many more recent subquadratic architectures improve linear attention quality with additional factors such as decay terms [] and additional gating []. Studying whether attention transfer and low-rank linearizing can help scale up these additional attention analogs is an interesting line of future work.

#### Ethics Statement

Our work deals with improving the efficiency of open-weight models. While promising for beneficial applications, increasing their accessibility also raises concerns about potential misuse. Bad actors could leverage our technique to develop LLMs capable of generating harmful content, spreading misinformation, or enabling other malicious activities. We focus primarily on base models, but acknowledge that linearizing could also be used on instruction-tuned LLMs; research on whether linearizing preserves guardrails is still an open question. We acknowledge the risks and believe in the responsible development and deployment of efficient and widely accessible models.

#### Reproducibility

We include experimental details in Appendix A, and further implementation details with sample code for linearizing architectures and training in Appendix C.1. Our code is also available at [https://github.com/HazyResearch/lolcats](https://github.com/HazyResearch/lolcats)

### Acknowledgements

We thank Mayee Chen, Ben Viggiano, Gautam Machiraju, Dan Fu, Sabri Eyuboglu, Tri Dao, and anonymous reviewers for helpful discussions on linear attention and paper feedback.
We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF2247015 (Hardware-Aware), CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); US DEVCOM ARL under Nos. W911NF-23-2-0184 (Long-context) and W911NF-21-2-0251 (Interactive Human-AI Teaming); ONR under Nos. N000142312633 (Deep Signal Processing); Stanford HAI under No. 247183; 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 Credits for Research program, the Stanford Data Science Initiative (SDSI), and members of 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 in 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.

## Appendix A Experimental Details

### A.1 Main Results, Linearizing 7B and 8B LLMs

##### Setup.

We describe our setup for linearizing Mistral 7B (v0.1) [], Llama 3 8B [], and Llama 3.1 8B [].

For linearizing layers, we replace softmax attentions with hybrid linear + sliding window analogs (Section 3.3.1), using Hedgehog’s feature map for its prior quality [].

Figure: (a) Sliding window
Refer to caption: extracted/5954737/figures/concept_sw_hybrid.png

For the sliding window implementation, we considered two options: a standard sliding window where $w$ is the same for all tokens, and a “terraced” window where $w$ changes based on token index (Figure 10). While we found both comparable in quality (Table 13), the latter lets us exploit the new ThunderKittens (TK) DSL’s [] primitives for implementing fast CUDA kernels. Here we prefer contiguous blocks of size $w=64$, which can quickly
be computed in parallel on modern GPUs. We use this “terrace” implementation in our main results, and include further implementation details in Appendix C.2.

For linearizing data, we use the Alpaca linearizing data setup in Section 3.2 unless otherwise noted. We also tried a more typical pretraining corpus (a subset(^5^55[https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T-Sample](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T-Sample)) of RedPajama []), but found comparable performance when controlling for number of token updates (Appendix B.2.1).
To linearize, we simply train all feature maps in parallel for two epochs with learning rate 1e-2, before applying LoRA on
the attention projection layers for two epochs with learning rate 1e-4. By default, we use LoRA rank $r=8$, and scale LoRA updates by $2$ ($\alpha$ = 16 in HuggingFace PEFT(^6^66[https://huggingface.co/docs/peft/en/index](https://huggingface.co/docs/peft/en/index))), amounting to training $<$0.09% of all model parameters.
For both
stages, we train with early stopping, AdamW optimizer [], and packing into 1024-token sequences with batch size 8. We evaluate the best checkpoints based on validation set perplexity.

##### Hyperparameters.

We list all model and training hyperparameters in Table 8.

**Table 8: Hyperparameters for Mistral 7B, Llama 3 8B, and Llama 3.1 8B experiments.**
|  | Hedgehog | \name |
| --- | --- | --- |
| Model |  |  |
| Precision | 16-bit (bfloat16) |  |
| Sequence length | 1024 |  |
| Linearizing attention | (Linear) | (Linear + Sliding Window) |
| Linear attn feature map | Hedgehog |  |
| Linear attn feature dimension | 64 (effectively 128, see Table 2) |  |
| Linear attn feature activation | Softmax (across feature dim) |  |
| Sliding window implementation | N/A | Terrace |
| Sliding window attn size | N/A | 64 |
| Optimizer and LR Schedule |  |  |
| Optimizer | AdamW |  |
| Global batch size | 8 |  |
| Gradient accumulation | 8 |  |
| Gradient clipping threshold | $1.0$ |  |
| Learning rate schedule | Reduce LR on Plateau |  |
| Step 1: Attention Transfer |  |  |
| Number of epochs | 2 |  |
| Tokens per epoch | 10M |  |
| Learning rate | 0.01 |  |
| Step 2: Low-rank Finetuning |  |  |
| Number of epochs | 2 |  |
| Tokens per epoch | 20M |  |
| Learning rate | 1e-4 |  |
| LoRA rank and alpha | $r=$8, $\alpha=$16 |  |
| LoRA dropout | $0.0$ |  |
| LoRA projections | $\bm{W}_{q},\bm{W}_{k},\bm{W}_{v},\bm{W}_{o}$ |  |

##### Compute Resources.

For each linearizing run we use one NVIDIA 40GB A100 GPU. With batch size 1 and gradient accumulation over 8 batches, attention transfer takes $\approx$ 2 hours and post-swap finetuning takes $\approx$ 4.5 hours, i.e., 6.5 total GPU hours to linearize an 8B LLM.

### A.2 Linearizing Llama 3.1 70B

We provide experimental details corresponding to the $70$B parameter results reported in Table 7.

##### Setup.

We compare the quality of two linearization approaches to the quality of the original Llama 3.1 70B model, including (1) the baseline linearization without attention transfer, which is representative of the approach used in prior work [, , ] and (2) our approach, \name. For both the baseline and \name, we start with Llama 3.1 70B and replace the softmax attentions with the linear attention architecture defined in Section 3.3.1, Equation 7. The training procedure involves:

- •
Baseline: We introduce LoRA parameters to the attention $\bm{W}_{q},\bm{W}_{k},\bm{W}_{v},\bm{W}_{o}$ projection matrices. We train the linear attention feature maps, learnable mixing term $\gamma$, and the LoRA parameters during the fine-tuning adjustment stage of the linearization process.
- •
\name
: We first perform layer-wise attention transfer following Equation 8, with $k=80$ (i.e., we optimize over all layers together). We then introduce LoRA parameters to the attention $\bm{W}_{q},\bm{W}_{k},\bm{W}_{v},\bm{W}_{o}$ projection matrices. During fine-tuning we only train the LoRA parameters, freezing the linear attention map and $\gamma$ weights.

We use an MSE loss for the layer-wise attention transfer stage and cross-entropy loss for the next-token prediction fine-tuning (Low-rank Finetuning) stage.

##### Hyperparameters.

We include the hyperparameters for the baseline and \name approaches in Table 9. Since the baseline does not use attention transfer, we mark these values with “N/A”.
We linearize each model using the same randomly sampled $20$M tokens of the RedPajama pre-training corpus []. We pack the sequences to fill full context length, and evaluate the best checkpoints based on validation set perplexity.

**Table 9: Hyperparameters for Llama 3.1 70B experiments.**
|  | Baseline | \name |
| --- | --- | --- |
| Model |  |  |
| Precision | 16-bit (bfloat16) |  |
| Sequence length | 1024 |  |
| Linearizing attention | Linear + Sliding Window |  |
| Linear attn feature map | Hedgehog |  |
| Linear attn feature dimension | 64 (effectively 128, see Table 2) |  |
| Linear attn feature activation | Softmax (across feature dim) |  |
| Sliding window implementation | Terrace |  |
| Sliding window attn size | 64 |  |
| Optimizer and LR Schedule |  |  |
| Optimizer | AdamW |  |
| Global batch size | 8 |  |
| Gradient accumulation | 8 |  |
| Gradient clipping threshold | $1.0$ |  |
| Learning rate schedule | Reduce LR on Plateau |  |
| Stage 1: Attention Transfer |  |  |
| Number of epochs | N/A | 1 |
| Tokens per epoch | N/A | 20M |
| Learning rate | N/A | 0.01 |
| Stage 2: Low-rank Finetuning |  |  |
| Number of epochs | 1 |  |
| Tokens per epoch | 20M | 20M |
| Learning rate | 1e-4 |  |
| LoRA rank and alpha | $r=$8, $\alpha=$16 |  |
| LoRA dropout | $0.0$ |  |
| LoRA projections | $\bm{W}_{q},\bm{W}_{k},\bm{W}_{v},\bm{W}_{o}$ |  |

##### Compute Resources.

We linearize using a single NVIDIA $8\times 80$GB H100 node. Attention transfer takes $4$ hours and fine-tuning takes $14$ hours. We use PyTorch FSDP with activation checkpointing for distributed training.

### A.3 Linearizing Llama 3.1 405B

##### Setup.

We compare the quality of two linearization approaches to the quality of the original Llama 3.1 405B model, including (1) the baseline linearization without attention transfer, which is representative of the approach used in prior work [] and (2) our approach, \name.
For both the baseline and \name, we start with Llama 3.1 405B and replace the softmax attentions with the linear attention architecture defined in Section 3.3.1, Equation 7.
The training procedure involves:

- •
Baseline: We introduce LoRA parameters to the attention $\bm{W}_{q},\bm{W}_{k},\bm{W}_{v},\bm{W}_{o}$ projection matrices.
We train the linear attention feature maps, learnable mixing term $\gamma$, and the LoRA parameters during the fine-tuning adjustment stage of the linearization process.
- •
\name
:
We first perform block-wise attention transfer following Equation 8, with $k=9$ as the block size. To perform attention transfer for block $i$, we save the hidden states outputted by block $i-1$ to disk and then use this as training data for block $i$.
We then introduce LoRA parameters to the attention $\bm{W}_{q},\bm{W}_{k},\bm{W}_{v},\bm{W}_{o}$ projection matrices.
During fine-tuning we only train the LoRA parameters, freezing the linear attention map and $\gamma$ weights.

For the reported checkpoints, we train the layer-wise attention transfer stage using use a weighted combination of the MSE loss on attention outputs plus a cross-entropy loss between the softmax and linear attention maps. We use cross-entropy loss for the next-token prediction fine-tuning (Low-rank Finetuning) stage.

##### Hyperparameters.

We include the hyperparameters for the baseline and \name approaches in Table 10. Since the baseline does not use attention transfer, we mark it with “N/A”.
We linearize each model using the same randomly sampled $20$M tokens of the RedPajama pre-training corpus []. We pack the sequences to fill the full context length, and evaluate the best checkpoints based on validation set perplexity.

**Table 10: Hyperparameters for Llama 3.1 405B experiments.**
|  | Baseline | \name |
| --- | --- | --- |
| Model |  |  |
| Precision | 16-bit (FP16) |  |
| Sequence length | 1024 |  |
| Linearizing attention | Linear + Sliding Window |  |
| Linear attn feature map | Hedgehog |  |
| Linear attn feature dimension | 64 (effectively 128, see Table 2) |  |
| Linear attn feature activation | Softmax (across feature dim) |  |
| Sliding window implementation | Terrace |  |
| Sliding window attn size | 64 |  |
| Optimizer and LR Schedule |  |  |
| Optimizer | AdamW |  |
| Global batch size | 8 |  |
| Gradient accumulation | 8 |  |
| Gradient clipping threshold | $1.0$ |  |
| Learning rate schedule | Reduce LR on Plateau |  |
| Stage 1: Attention Transfer |  |  |
| Number of epochs | N/A | 1 |
| Tokens per epoch | N/A | 20M |
| Learning rate | N/A | 0.01 |
| MSE and X-ent weights | N/A | 1000, 1 |
| Stage 2: Low-rank Finetuning |  |  |
| Number of epochs | 1 |  |
| Tokens per epoch | 20M | 20M |
| Learning rate | 1e-4 |  |
| LoRA rank and alpha | $r=4$, $\alpha=8$ |  |
| LoRA dropout | $0.5$ |  |
| LoRA projections | $\bm{W}_{q},\bm{W}_{k},\bm{W}_{v},\bm{W}_{o}$ |  |

##### Compute Resources.

We linearize the baseline using three NVIDIA $8\times 80$GB H100 nodes, evaluating the best validation checkpoint after 19.5 hours.
For \name, we perform attention transfer using a single $80$GB H100 GPU for 5 hours per block, and we finetune using three NVIDIA $8\times 80$GB H100 nodes for 16 hours. We use PyTorch FSDP with activation checkpointing for distributed training.

## Appendix B Additional Experiments

To better understand \name’ properties and performance for high quality yet parameter and training-efficient linearizing, we now study how different amounts of parameter training, different data sources, and different amounts of training data affect \name quality.

### B.1 Study on Parameter-Efficient Training

In our main results, we found that simple default initializations (e.g., rank $8$, applied to all attention projection layers) could recover high quality linearizing while only updating $<$0.2% of model parameters. In this section, we study how changing different aspects of low-rank adaptation and sliding window size impact linearizing performance.

#### B.1.1 Effect of LoRA Rank

We study the effect of LoRA rank in post-swap finetuning for zero-shot linearized LLM performance. Following standard implementations [], we consider two factors: rank $r$, which determines the rank of the low-rank matrices $\bm{A},\bm{B}$ we decompose the weight deltas into; and alpha $\alpha$, where $\alpha/r$ is a scaling factor that controls the degree to which $\bm{B}\bm{A}$ affect the output (i.e., LoRA output $y=\bm{W}x+\frac{\alpha}{r}\bm{B}\bm{A}x$).

##### Setup.

We sweep over ranks {4, 8, 16, 32, 64, 128, 256}, and adjust $\alpha$ such that $\alpha/r=2$ as a default scaling factor []. For all runs, we linearize Llama 3 8B, and use the same default experimental setup as our main results. We start with the linear + “terrace” window attentions, training the feature maps via attention transfer over 20M tokens over Alpaca (2 full epochs). We then freeze feature maps and apply LoRA with the above ranks on all attention weight projections (freezing all other parameters such as those in MLPs or GroupNorms), finetuning for two more epochs over Alpaca using the hyperparameters in Table 8. We evaluate with LM Eval tasks.

##### Results.

We report results in Table 11. We find that when applying LoRA to adjust in linearizing after attention transfer, larger rank updates, e.g., $r=$ 128 or 256, do not necessarily lead to improved zero or few-shot downstream performance. Somewhat surprisingly, using just $r=$ 4 leads to overall best performance, while $r=$ 128 and $r=$ 8 achieve best and second-best 5-shot MMLU accuracy.

**Table 11: LoRA rank $r$ comparison. Evaluation on LM Evaluation Harness tasks. When adapting all projections $\bm{W}_{q},\bm{W}_{k},\bm{W}_{v},\bm{W}_{o}$ after attention transfer, we find smaller ranks $r=4,8$ are surprisingly sufficient. Larger ranks with more parameter-heavy updates do not necessarily improve downstream performance.**
| LoRA<br>Rank | PiQA | ARC-e | ARC-c<br>(norm) | HellaSwag<br>(norm) | Wino-<br>grande | MMLU<br>(5-shot) | Avg. | Avg.<br>(no MMLU) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 4 | 80.7 | 82.5 | 56.0 | \ul79.6 | 73.6 | 50.1 | 71.3 | 74.5 |
| 8 | \ul80.9 | 81.7 | \ul54.9 | 79.7 | 74.1 | \ul52.8 | 70.7 | \ul74.2 |
| 16 | 80.7 | 81.9 | 54.5 | 79.7 | \ul73.8 | 48.9 | 69.9 | 74.1 |
| 32 | 81.1 | 81.5 | 54.5 | 79.7 | 72.8 | 51.0 | 70.1 | 73.9 |
| 64 | \ul80.9 | 81.9 | 54.5 | 79.3 | 72.1 | 51.7 | \ul71.1 | 73.8 |
| 128 | 80.5 | 80.9 | 52.8 | 78.4 | 72.2 | 53.4 | 69.7 | 73.0 |
| 256 | 80.7 | 80.2 | 52.1 | 78.8 | 71.4 | 52.1 | 69.2 | 72.7 |

#### B.1.2 Effect of LoRA Projection Layer

We next compare performance when applying LoRA to different weight matrices of the linear attention layers. With the same training and evaluation setup as Appendix B.1.1, but fixing $r=8,\alpha=16$, we now apply LoRA to different combinations of $\bm{W}_{q},\bm{W}_{k},\bm{W}_{v},\bm{W}_{o}$ weights after swapping in attention-transferred linear attentions. We use the same combination for each layer.

We report results in Table 12. Interestingly, when isolating for projections updated, LoRA on projections not involved in computing layer-wise attention weights (value $\bm{W}_{v}$ and output $\bm{W}_{o}$ projections) improves quality compared to query $\bm{W}_{q}$ or key projections $\bm{W}_{k}$. Somewhat surprisingly, updating just $\bm{W}_{v}$ or $\bm{W}_{o}$ achieves comparable performance to updating all projections (c.f., average zeros-shot accuracy of 74.19% when updating just $\{\bm{W}_{v}\}$, 74.09% for $\{\bm{W}_{o}\}$ versus 74.24% updating $\{\bm{W}_{q},\bm{W}_{k},\bm{W}_{v},\bm{W}_{o}\}$). Meanwhile, updating just $\bm{W}_{q}$ or $\bm{W}_{k}$ performs significantly worse (72.68%, 72.29%; versus 74.24%). This suggests much of the quality recovery in Stage-2 low-rank finetuning comes from adjusting values and outputs to the learned attention weights—as opposed to further refining attention weight computation. While best results do come from a combination of adapting either $\bm{W}_{q}$ or $\bm{W}_{k}$ with value and output projections, we may be able to achieve even more parameter-efficient linearizing—with comparable quality—by focusing on subsets with the latter.

**Table 12: LoRA projection comparison. Evaluation on zero-shot LM Evaluation Harness tasks. We apply LoRA with the same rank $r=8$ to different combinations of the attention projections, shading scores by increasing quality (darker is better). When isolating for projections updated, LoRA on projections not involved in layer-wise attention weight computations (value $\bm{W}_{v}$ and output $\bm{W}_{o}$ projections) achieves higher quality over query $\bm{W}_{q}$ or key projections $\bm{W}_{k}$. This suggests $\bm{W}_{v},\bm{W}_{o}$ may be more important to adapt after attention transfer, although best results involve a combination across attention weight and output projections.**
| LoRA Projection | PiQA | ARC-e | ARC-c (norm) | HellaSwag (norm) | Winogrande | Average |
| --- | --- | --- | --- | --- | --- | --- |
| $\bm{W}_{q}$ | 79.49 | 81.06 | 51.45 | 79.18 | 72.22 | 72.68 |
| $\bm{W}_{k}$ | 79.82 | 79.80 | 50.68 | 78.54 | 72.61 | 72.29 |
| $\bm{W}_{v}$ | 80.69 | 82.49 | 56.66 | 79.28 | 71.82 | 74.19 |
| $\bm{W}_{o}$ | 80.09 | 81.65 | 55.63 | 79.35 | 73.72 | 74.09 |
| $\bm{W}_{q},\bm{W}_{k}$ | 79.76 | 81.19 | 51.02 | 79.29 | 72.22 | 72.70 |
| $\bm{W}_{q},\bm{W}_{v}$ | 80.96 | 81.90 | 54.35 | 79.01 | 72.30 | 73.70 |
| $\bm{W}_{q},\bm{W}_{o}$ | 81.28 | 81.86 | 54.78 | 79.20 | 73.72 | 74.17 |
| $\bm{W}_{k},\bm{W}_{v}$ | 80.74 | 82.62 | 56.83 | 79.26 | 71.98 | 74.28 |
| $\bm{W}_{k},\bm{W}_{o}$ | 80.69 | 82.15 | 55.46 | 79.53 | 73.01 | 69.33 |
| $\bm{W}_{v},\bm{W}_{o}$ | 80.69 | 82.79 | 55.89 | 79.57 | 71.59 | 74.10 |
| $\bm{W}_{q},\bm{W}_{k},\bm{W}_{v}$ | 80.90 | 82.20 | 56.48 | 79.13 | 73.95 | 74.53 |
| $\bm{W}_{q},\bm{W}_{k},\bm{W}_{o}$ | 80.41 | 80.89 | 54.18 | 79.18 | 74.35 | 73.80 |
| $\bm{W}_{k},\bm{W}_{v},\bm{W}_{o}$ | 80.36 | 82.32 | 54.61 | 79.13 | 73.56 | 74.00 |
| $\bm{W}_{q},\bm{W}_{k},\bm{W}_{v},\bm{W}_{o}$ | 80.85 | 81.73 | 54.86 | 79.65 | 74.11 | 74.24 |

#### B.1.3 Effect of Window Size

We now compare model performance using different window sizes in the \name linear + sliding window attention layer. With the standard sliding window implementation (Fig 10), we compare LM Eval performance after linearizing with window sizes $w\in\{4,16,64,256\}$. We also compare against the ThunderKittens-motivated “terraced” implementation used in our main experiments with window size $64$.
In Table 13, we find that in each of these settings, having more softmax attention generally improves performance (c.f., $w=$ 16,64 versus $w=$4). However, more softmax attention does not always lead to better quality. Window size 256 results in up to an 8.8 point drop in 5-shot MMLU accuracy, suggesting we may not necessarily trade-off more softmax attention for higher quality. Comparing the standard sliding window and terracing implementations, we find similar performance (70.3 versus 70.7 average accuracy across tasks).

**Table 13: Window size comparison. We ablate the window size in \name linear + sliding window attention. For each window size $w$, the layer applies softmax attention to the $w$-most recent positions, combined with Hedgehog linear attention applied for all prior positions (Eq. 7, Figure 10). We compare $w\in\{4,16,64,256\}$ using the standard sliding window implementation with our default $w=64$ terraced window setup motivated by ThunderKittens (TK 64). Window size 64 performs best, where both implementations perform comparably.**
| Window<br>Size | PiQA | ARC-e | ARC-c<br>(norm) | HellaSwag<br>(norm) | Winogrande | MMLU<br>(5-shot) | Avg. | Avg.<br>(no MMLU) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 4 | 80.7 | 81.4 | 55.8 | 76.6 | 72.1 | 40.8 | 67.9 | 73.3 |
| 16 | 80.5 | 82.2 | 56.0 | 78.2 | 73.9 | 50.3 | 70.2 | 74.1 |
| 64 | 80.7 | 81.7 | 54.7 | 79.1 | 75.3 | 50.3 | 70.3 | 74.3 |
| 256 | 80.6 | 81.8 | 55.0 | 75.6 | 74.9 | 41.5 | 68.3 | 73.6 |
| TK 64 | 80.9 | 81.7 | 54.9 | 79.7 | 74.1 | 52.8 | 70.7 | 74.2 |

### B.2 Study on Linearizing Data

While most of our work focuses on architecture and training procedure for improving LLM linearizing quality, we now study how data selection affects \name performance.

#### B.2.1 Data Source: Alpaca versus RedPajama

We study the effect of linearizing data for downstream LLM performance. While we initially found that just using the $\sim$ 50K samples of a cleaned Alpaca dataset(^7^77[https://huggingface.co/datasets/yahma/alpaca-cleaned](https://huggingface.co/datasets/yahma/alpaca-cleaned)) [] could lead to surprisingly high performance on popular zero-shot LM Eval tasks, prior linearizing works [] use more typical pretraining datasets to linearize such as RefinedWeb []. We thus also try linearizing with a random subset of RedPajama [] to evaluate how \name works with pretraining data, albeit without any special curation. For both setups, we pack samples into 1024 token sequences and randomly subsample the RedPajama data so that we use the same number of training tokens (20M) for both attention transfer and finetune stages (40M tokens overall). We use the setup as described in Appendix A.1 for all other hyperparameters.

In Table 14, we find that across Mistral 7B (v0.1) and Llama 3 8B, using the Alpaca cleaned dataset actually leads to better downstream task quality for all tasks except for 5-shot MMLU, where linearizing with RedPajama consistently leads to $\sim$ 2 percentage point improvements. \name with both of these datasets leads to comparable or higher performance than prior methods trained on 2500$\times$ the data (c.f., Table 4; SUPRA trained on 100B tokens gets Avg. accuracy of 64.0%), suggesting that \name can robustly improve linearizing quality over different data sources.

**Table 14: Linearizing data comparison. For linearizing Mistral 7B (v0.1) and Llama 3 8B, \name with Alpaca and RedPajama subsets perform comparably (c.f. prior methods, Table 4), though we find that Alpaca actually leads to higher accuracy for most tasks other than 5-shot MMLU.**
| Model | Linearizing<br>Dataset | PiQA | ARC-e | ARC-c<br>(acc. norm) | HellaSwag<br>(acc. norm) | Wino-<br>grande | MMLU<br>(5-shot) | Avg. | Avg.<br>(no MMLU) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Mistral 7B (v0.1) | - | 82.1 | 80.9 | 53.8 | 81.0 | 74.0 | 62.4 | 72.4 | 74.4 |
| $\rightarrow$ LoLCATs (Ours) | Alpaca Clean | 81.5 | 81.7 | 54.9 | 80.7 | 74.0 | 51.4 | 70.7 | 74.5 |
| $\rightarrow$ LoLCATs (Ours) | RedPajama | 80.1 | 77.6 | 49.0 | 80.3 | 71.7 | 53.2 | 68.6 | 71.7 |
| Llama 3 8B | - | 79.9 | 80.1 | 53.3 | 79.1 | 73.1 | 66.6 | 72.0 | 73.1 |
| $\rightarrow$ LoLCATs (Ours) | Alpaca Clean | 80.9 | 81.7 | 54.9 | 79.7 | 74.1 | 52.8 | 70.7 | 74.2 |
| $\rightarrow$ LoLCATs (Ours) | RedPajama | 78.9 | 79.0 | 52.0 | 78.1 | 72.6 | 55.2 | 69.3 | 72.1 |

#### B.2.2 Sample Lengths: Effect of Effective Sequence Lengths

We further study the impact of sample sequence length for linearizing quality.
By default, for linearizing data we pack original data samples into sequences of consistent length, e.g., 1024 tokens. As done in prior work [], this allows us to pack multiple short data samples together into longer training sequences, improving training efficiency and removing any padding tokens. However, it may also introduce situations where our linearizing sequences only carry short-context dependencies, i.e., because we pack together many samples with few tokens, or split longer samples into multiple sequences. Especially with attention transfer, linearized LLMs may model longer samples less well (e.g., the 5-example in-context samples in 5-shot MMLU) because we never learn to approximate attentions over “long enough” sequence lengths.

##### Effective sequence length.

To study this data effect, we define an “effective sequence length” (ESL) metric. This roughly captures for each query how far back a layer needs to attend to capture all non-zero softmax attention weights. For query at position $i$, we define the ESL per query as

$$ $\text{ESL}({\bm{q}_{i}}):=\sum_{j=1}^{i}(i-j)\frac{\exp(\bm{q}_{i}^{\top}\bm{k }_{j}/\sqrt{d})}{\sum_{m\leq j}\exp(\bm{q}_{i}^{\top}\bm{k}_{m}/\sqrt{d})}$ (9) $$

We compute a sample’s ESL per head as the sum over all query ESLs, i.e., $\sum_{i=1}^{n}\text{ESL}(\bm{q}_{i})$ for a sample with $n$ tokens. We average this over all heads and layers to measure a sample’s overall ESL.

We hypothesize that if our linearizing data only has samples with shorter ESL than those encountered at test time, then we would poorly model these test samples. Conversely, we may be able to improve linearizing quality by specifically filtering for samples with longer ESL. We report two findings next.

Figure: Figure 11: Effective sequence length distributions. Although we pack data samples into 1024-token training sequences, different data sources can vary in effective sequence lengths (left). Furthermore, we can selectively filter for longer ESL samples (right). We find linearizing with longer ESLs coincides with improved MMLU scores, albeit at the cost of other LM Eval tasks (Table 15).
Refer to caption: extracted/5954737/figures/esl_alpaca_rp.png

**Table 15: Effect of ESL on linearized LLM quality, Llama 3 8B. While linearizing with longer ESLs—e.g., RedPajama samples or specifically filtering for top ESLs in RedPajama (Sample Top ESL)—improves 5-shot MMLU accuracy up to 3.7 points, it reduces quality on all other evaluated LM Eval tasks by 1.7 to 5.0 points.**
| Linearizing Data | PiQA | ARC-easy | ARC-challenge | HellaSwag<br>(acc. norm) | Winogrande<br>(acc. norm) | MMLU<br>(5-shot) | Avg.<br>(no MMLU) |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Alpaca | 80.9 | 81.7 | 54.9 | 79.7 | 74.1 | 52.8 | 74.2 |
| RedPajama | 78.9 | 79.0 | 52.0 | 78.1 | 72.6 | 55.2 | 72.1 |
| RedPajama (Sample Top ESL) | 78.4 | 77.0 | 49.8 | 78.0 | 71.4 | 56.5 | 70.9 |

##### Finding 1: ESL corresponds with RedPajama versus Alpaca performance.

In Figure 11, we first plot the distribution of sample ESLs computed with Llama 3 8B on Alpaca and RedPajama linearizing data subsets. We find RedPajama samples on average display longer ESLs, which conincides with improved MMLU score (c.f., Table B.2.1).

##### Finding 2: Filtering for higher ESL improves MMLU.

Furthermore, we can increase the ESLs in linearizing data by actively filtering for high ESL samples (Figure 11 right). Here we actively filter for the top 20,000 packed RedPajama samples with the highest ESLs, amounting to 20M tokens. When doing one epoch of attention transfer and low-rank linearizing with this subset, we further improve MMLU accuracy by 1.3 points (Table 15). However, this comes at a cost for all other LM Eval tasks, dropping quality compared to random RedPajama packing by 0.2 to 2.1 points.

### B.3 Study on Linearizing Token Budget

We further study how varying the number of tokens used for both attention transfer and low-rank adaptation impacts \name linearizing quality.

Impact of minimal tokens. To first test how efficient we can be with attention transfer, we linearize Llama 3 8B with varying numbers of attention transfer steps (0 - 1800), before low-rank adjusting for up to 2000 steps. We use the Alpaca dataset and the same packed random sampling as our main experiments, and measure evaluation perplexity on validation samples both in-distribution (held-out Alpaca samples) and out-of-distribution (RedPajama validation samples) over different combinations of steps (Figure 13).
Without attention transfer, low-rank adaptation converges significantly higher on in-distribution samples (Figure 12(a)), suggesting poorer quality linearizing. However, we find similar held-out perplexities after relatively few attention transfer steps (c.f., 1000 - 1800 updates, the former amounting to just 8 million tokens for attention transfer), where all runs improve in-distribution PPL by $\sim$0.23 points after LoRA finetuning for 2000 steps.

In Table 16, we report the numerical values for held-out perplexities at the end of linearizing (1800 attention transfer steps + 2000 low-rank adaptation steps), as well as the average LM Eval score over zero-shot tasks. We similarly find competitive generalized zero-shot quality with relatively few attention transfer steps (200 steps), all achieving 7.70–8.16 higher points than the next best Mamba-Llama model (0.16–0.62 higher points than the 50% softmax attention variant, c.f., Table 4). Without any attention transfer, linearized LLMs perform drastically worse on out-of-distribution samples (Table 16, RedPajama and LM Eval metrics).

Figure: (a) Alpaca PPL
Refer to caption: extracted/5954737/figures/ppl_train_token_alpaca.png

**Table 16: Effect of attention transfer steps. With Llama 3 8B linearized on Alpaca data, we report the final evaluation perplexities after 2000 LoRA steps in Fig 13, as well as downstream LM Eval performance averaged over zero-shot tasks. We again find competitive quality with relatively few attention transfer steps.**
| Attn. Transfer Steps | 0 | 200 | 400 | 600 | 800 | 1000 | 1200 | 1400 | 1600 | 1800 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Alpaca Eval PPL | 3.211 | 2.802 | 2.792 | 2.782 | 2.782 | 2.776 | 2.778 | 2.775 | 2.782 | 2.776 |
| RedPajama Eval PPL | 61.305 | 10.459 | 9.799 | 9.699 | 9.628 | 9.679 | 9.420 | 9.328 | 9.413 | 9.358 |
| Avg. Zero-shot LM Eval Acc. | 56.86 | 73.68 | 73.34 | 73.70 | 73.66 | 73.74 | 73.47 | 73.70 | 73.80 | 73.66 |

Impact of more pretraining data. We finally study how linearizing over larger amounts of pretraining data impacts quality. We randomly sample a larger set of unique RedPajama training sequences (1024-token packed; 72,000 such samples overall), allowing us to linearize Llama 3 8B with different combinations of up to 9,000 attention transfer updates and 2,000 low-rank linearizing updates. To test language modeling recovery, we report both held-out validation sample perplexity (Table 13(a)) and general zero-shot LM Evaluation Harness quality (Table 13(b)). Increasing both Stage 1 attention transfer steps and Stage 2 low-rank adjusting steps notably improves validation perplexity. However, we similarly find competitive zero-shot LM scores across all evaluated attention transfer steps. Across checkpoints at different numbers of low-rank updates, attention transfer with up to 9$\times$ more unique tokens does not seem to monotonically improve downstream quality (Table 13(b)). Meanwhile, we do find that across various amounts of attention transfer steps, subsequent low-rank adaptation consistently improves average zero-shot LM score by $>$1 points.

Figure: (a) RedPajama Validation Perplexity
Refer to caption: extracted/5954737/figures/lm_eval_long_ppl.png

### B.4 Study on Block-wise Attention Transfer

This section studies how the properties of attention transfer change with model scale, motivating \name’s block-wise training approach (Section 3.3.2).

Attention transfer involves training the \name linear attentions to match the outputs of softmax attention at each layer by minimizing the MSE between the softmax and linear attention outputs.
Since the MSE loss is scale-sensitive and the magnitude of the MSE loss varies across layers (Figure 7(b), we hypothesize that jointly training all $126$ Transformer layers in Llama 3.1 405B – by summing the MSE losses across all layers – may be difficult. Correspondingly, in Table 3, we find that block-wise attention transfer leads to lower language modeling perplexity for 405B linearized LLMs compared to joint training. However, we find that joint training is sufficient and performs similarly to block-wise training at the smaller scales (8B, 70B).

We study attention transfer across model scales and find that the variation in MSE magnitudes correlates with the benefit of block-wise training. In Tables 17 and 18, we report magnitude of the MSE loss as the depth of the layers in the block increases. We find that the variation in MSE magnitudes increases with the model size: while the largest MSEs at the 70B scale are between $2.5-3$, the MSE of the final block at the 405b scale is $48.66$ (over $13\times$ that of the next largest MSE by block, $3.62$).

**Table 17: Attention transfer block-wise MSE We report the eval MSE by $5$-layer block for each of the $16$ blocks in the $80$ Transformer layer Llama 3.1 70B model. Each block is trained on the exact same set of RedPajama data at sequence length 1024.**
| Block (Layer range) | Eval MSE |
| --- | --- |
| 0-4 | $8e-4$ |
| 5-9 | 0.03 |
| 10-14 | 0.06 |
| 15-19 | 0.10 |
| 20-24 | 0.28 |
| 25-29 | 0.73 |
| 30-34 | 0.28 |
| 35-39 | 0.28 |
| 40-44 | 0.25 |
| 45-49 | 1.08 |
| 50-54 | 0.26 |
| 55-59 | 0.18 |
| 60-64 | 0.45 |
| 65-69 | 0.50 |
| 70-74 | 2.91 |
| 75-79 | 2.56 |

**Table 18: Attention transfer block-wise MSE We report the eval MSE by $9$-layers block for the $14$ blocks in the $126$ Transformer layer Llama 3.1 405B model. Each block is trained on the exact same set of RedPajama data at sequence length 1024.**
| Block (Layer range) | Eval MSE |
| --- | --- |
| 0-8 | 0.03 |
| 9-17 | 0.02 |
| 18-26 | 0.09 |
| 27-35 | 0.04 |
| 36-44 | 0.36 |
| 45-53 | 0.42 |
| 54-62 | 1.59 |
| 63-71 | 2.75 |
| 72-80 | 3.62 |
| 81-89 | 2.49 |
| 90-98 | 0.29 |
| 99-107 | 1.11 |
| 108-116 | 3.42 |
| 117-126 | 48.66 |

## Appendix C Implementation Details

### C.1 Code Implementation

Below we provide further details on implementing \name with PyTorch-like code and example demonstrations from the HuggingFace Transformers library.

Learnable Linear Attention.
To start, we simply replace the softmax attentions in an LLM with a linear attention. We define such a class below.

Figure: Listing 1: \name Linear Attention Class

Hedgehog Feature Map. We implement the Hedgehog feature map following Zhang et al. [].

Figure: Listing 2: Hedgehog Feature Map

Linearizing LLM Setup. To initialize an LLM for linearizing, we simply replace each softmax attention in the Transformer’s layers with our LoLCATs linear attention class. We illustrate this with a Huggingface Transformer’s class below.

Figure: Listing 3: Linearizing LLM Setup

Attention Transfer Training. Finally, we can train \name layers in a simple end-to-end loop. Although doing this attention transfer is akin to a layer-by-layer cross-architecture distillation, due to architectural similarities we implement linearizing with the same footprint as finetuning a single model. Furthermore, as we freeze all parameters except for the newly introduced feature map weights, this amounts to parameter-efficient finetuning, training $<$0.2% of a 7B+ LLM’s parameters.

Figure: Listing 4: End-to-end attention transfer pseudocode.

### C.2 Hardware-aware implementation of \name sliding window

Despite the theoretical efficiency of linear attention, existing implementations have long underperformed well-optimized attention implementations (e.g., FlashAttention) in wall clock speed []. To translate the benefits of \name to wall clock speedups, we develop a custom hardware-aware algorithm for \name prefill using the ThunderKittens CUDA framework.(^8^88[https://github.com/HazyResearch/ThunderKittens](https://github.com/HazyResearch/ThunderKittens))
We first briefly review the GPU execution model and then detail our algorithm.

#### C.2.1 GPU execution model

GPUs workloads are executed by independent streaming multiprocessors (SMs), which contain warps, groups of $32$ threads, that operate in parallel.

##### Memory hierarchy.

ML workloads involve moving large tensors (weights, activations) in and out of memory to perform computation. GPUs have a memory hierarchy, which includes global memory (HBM), shared memory (SRAM), and registers. Reading from and writing data to memory, referred to as I/O operations, takes time. There is a large amount of HBM, which has high I/O costs, and a small amount of SRAM and registers have much costs. All SMs access global memory, warps within an SM threadblock can access shared memory, and threads within a threadblock have independent register memory. To reduce the I/O costs, locality is key – kernels should perform as many operations as possible on data that has already been loaded into fast memory (i.e., thread registers) before writing the results back to slower memory.

##### Compute units.

GPUs have increasingly heterogeneous compute units on newer generations of hardware. Tensor cores—specialized compute units for matrix-matrix multiplications—are the fastest units, operating at 1.0 PetaFLOPS on Nvidia H100 GPUs in contrast to 67 TeraFLOPS for the general non Tensor core units. ML workloads should thus ideally exploit the tensor cores.

##### Cost model.

Overall, workloads may either be compute or memory bound, depending on whether they are bottlenecked by the compute speed or I/O costs. To hide latencies from either expensive compute or I/O, a classic principle in systems is to pipeline computation among parallel workers.

#### C.2.2 ThunderKittens CUDA kernel for prefill

We describe our overall approach below and provide pseudocode in Algorithm 1, designed around the three principles above: memory locality, tensor core utilization, and pipelined execution.

The kernel fuses the entire \name layer, taking as input the attention queries, keys, and values, for $\bm{q},\bm{k},\bm{v}\in\mathbb{R}^{N\times d}$ with sequence length $N$ and head dimension $d$ and outputting the result of the $\bm{y}\in\mathbb{R}^{N\times d}$. Following Llama 3 [], we let $d=128$ in the discussion below.

##### Pipeline execution overview.

Each thread block handles a single batch and head element of size $N\times d$. The kernel loops over chunks of length $64$ along the sequence dimension $N$, loading $64\times 128$ tiles of $\bm{q},\bm{k},\bm{v}$, which we’ll refer to as $\bm{q}_{t},\bm{k}_{t},\bm{v}_{t}$, in each iteration $t$.
We use $8$ warps (workers) per thread block, splitting them into two groups of $4$ workers that pipline the computation. One “warpgroup” is in charge of launching memory loads and stores and computing the relatively cheap terraced window attention, while the other focuses on computing the more expensive linear attention computation and recurrent state updates.

##### Warpgroup 1 (Window attention).

For the diagonal $64\times 128$ sized tiles, recall that the output is simply the window attention result. At iteration $t$, the warpgroup loads $\bm{q}_{t},\bm{k}_{t},\bm{v}_{t}$ into thread registers, use the tensor cores to multiply queries and keys, apply a causal mask, apply the Softmax, and use the tensor cores to multiply the attention scores with the values. Note that we can use Nvidia’s new waprgroup operations (e.g., WGMMA) introduced in the H100 architecture to perform these operations. We refer to the terraced window output as $\mathbf{terrace_{o}}$.

Because the window attention is relatively cheap, waprgroup 1 also helps handle loads and stores between HBM and SRAM for the entire kernel. We use tensor memory acceleration (TMA), a new H100 capability for asynchronous memory movement, to perform these loads and stores.

##### Warpgroup 2 (Linear attention).

We briefly review the linear attention equation. The formulation on the left shows a quadratic view, wherein $\phi(\bm{q}_{n})^{\top}$ and $\phi(\bm{k}_{i})$ are multiplied first, while the right formulation shows a linear view, wherein $\bm{k}_{i}$ and $\bm{v}_{i}^{\top}$ are multiplied first.

$$ $\hat{\bm{y}}_{n}=\sum_{i=1}^{n}\frac{(\phi(\bm{q}_{n})^{\top}\phi(\bm{k}_{i})) \bm{v}_{i}}{\sum_{i=1}^{n}\phi(\bm{q}_{n})^{\top}\phi(\bm{k}_{i})}=\frac{\phi( \bm{q}_{n})^{\top}\Big{(}\sum_{i=1}^{n}\phi(\bm{k}_{i})\bm{v}_{i}^{\top}\Big{) }}{\phi(\bm{q}_{n})^{\top}\sum_{i=1}^{n}\phi(\bm{k}_{i})}$ (10) $$

Since \name uses no linear attention on the diagonal tiles, the linear attention contribution for tile $t$ is as follows, where queries are multiplied by the cumulative KV state from the prior iterations up to $t-1$:

$$ $\hat{\bm{y}}_{t}\frac{\phi(\bm{q}_{t})^{\top}\Big{(}\sum_{i=1}^{t-1}\phi(\bm{k }_{i})\bm{v}_{i}^{\top}\Big{)}}{\phi(\bm{q}_{n})^{\top}\sum_{i=1}^{t-1}\phi( \bm{k}_{i})}$ (11) $$

At $t=0$, the KV state and K state are initialized to 0 0, maintained in waprgroup 2’s registers.

At iteration $t$, warpgroup 2 loads in the learned feature maps into register and computes $\bm{qf}_{t}$. This result gets multiplied by the running KV state so far up until $t-1$ (again 0 0 at iteration 0 0), and the result, $\mathbf{linear_{o}}$ gets written to SMEM.

The warps then update the KV state to prepare for the next iteration by featurizing $\bm{k}_{t}$ to $\bm{kf}_{t}$ using the learned feature map, and multiplying by $\bm{v}_{t}$ with WGMMA operations. Note that because the KV state in linear attention is somewhat large ($d\times d$), we leave the state in register throughout the kernel execution to avoid I/O costs.

##### Combining the results.

Warpgroup 1 loads the $\mathbf{linear_{o}}$ contribution from SMEM to its registers, adds the $\mathbf{terraced_{o}}$ component, normalizes the overall result, and stores it back to HBM using TMA asynchronous store operations. We provide pseudocode in Algorithm 1.

To recap, our overall algorithm uses three classical systems ideas to run efficiently: (1) pipelining the different attention and I/O operations, (2) keeping the fastest compute—the tensor cores—occupied, and (3) keeping the recurrent (KV) state local in fast memory (thread registers).

Figure: Algorithm 1 \name ThunderKittens prefill kernel

## Appendix D Extended Related Work

### D.1 Linearizing Transformers

In this work, we build upon both approaches explicitly proposed to linearize LLMs [], as well as prior methods focusing on smaller Transformers
reasonably adaptable to modern LLMs [, , ]. We highlight two approaches most related to \name and their extant limitations next.

##### Scalable UPtraining for Recurrent Attention (SUPRA).

Mercat et al. [] linearize LLMs by swapping softmax attentions with linear attentions similar to Retentive Network (RetNet) layers [], before jointly training all model parameters on the RefinedWeb pretraining dataset []. In particular, they suggest that linearizing LLMs with the vanilla linear attention in Eq. 2 is unstable, and swap attentions with

$$ $\hat{\bm{y}}_{n}=\text{GroupNorm}\Big{(}\sum_{i=1}^{n}\gamma^{n-i}\phi(\bm{q}_ {n})^{\top}\phi(\bm{k}_{i})\bm{v}_{i}\Big{)}$ (12) $$

GroupNorm [] is used as the normalization in place of the $\sum_{i=1}^{n}\phi(\bm{q}_{n})^{\top}\phi(\bm{k}_{i})$ denominator in Eq. 2, $\gamma$ is a decay factor as in RetNet, and $\phi$ is a modified learnable feature map from Transformer-to-RNN (T2R) [] with rotary embeddings []. In other words, $\phi(\bm{x})=\text{RoPE}(\text{ReLU}(\bm{x}\bm{W}+\bm{b}))$ with $\bm{W}\in\mathbb{R}^{d\times d}$ and $\bm{b}\in\mathbb{R}^{d}$ as trainable weights and biases. With this approach, they recover zero-shot capabilities in linearized Llama 2 7B [] and Mistral 7B [] models on popular LM Evaluation Harness [] and SCROLLS [] tasks.

##### Hedgehog.

Zhang et al. [] show we can train linear attentions to approximate softmax attentions, improving linearized model quality by swapping in the linear attentions as learned drop-in replacements. They use the standard linear attention (Eq. 2), where query, key, value, and output projections (the latter combining outputs in multi-head attention []) are first copied from an existing softmax attention. They then specify learnable feature maps $\phi(\bm{x})=[\text{softmax}(\bm{x}\bm{W}+\bm{b})\oplus\text{softmax}(-\bm{x}
\bm{W}-\bm{b})]$ (where $\oplus$ denotes concatenation, and both $\oplus$ and the softmax are applied over the *feature dimension*) for $\bm{q}$ and $\bm{k}$ in each head and layer, and train $\phi$ such that linear attention weights $\hat{\bm{a}}$ match a Transformer’s original softmax weights $\bm{a}$. Given some sample data, they update $\phi$ with a cross-entropy-based distillation to minimize:

$$ $\mathcal{L}_{n}=-\sum_{i=1}^{n}\frac{\exp(\bm{q}_{n}^{\top}\bm{k}_{i}/\sqrt{d} )}{\sum_{i=1}^{n}\exp(\bm{q}_{n}^{\top}\bm{k}_{i}/\sqrt{d})}\log\frac{\phi(\bm {q}_{n})^{\top}\phi(\bm{k}_{i})}{\sum_{i=1}^{n}\phi(\bm{q}_{n})^{\top}\phi(\bm {k}_{i})}$ (13) $$

as the softmax and linear attention weights are both positive and sum to $1$. As they focus on task-specific linearization (e.g., GLUE classification [] or WikiText-103 language modeling []), for both attention and model training they use task-specific training data. By doing this “attention distillation”,
they show significant linearized quality improvements over T2R mainly on smaller Transformers (e.g., 110M parameter BERTs [] and 125M GPT-2s []). They further show LLM linearizing by linearizing Llama 2 7B for a specific SAMSum summarization task [].

### D.2 Efficient Architectures

##### Subquadratic Attention Alternatives.

Many prior works study more efficient sequence modeling modules compared to Transformer self-attention, commonly training proposed architectures from scratch. While our work is most similar to and compatible to linear attentions [, , , , , , , ] and sparse attentions [, , ], other efficient architectures include convolutional models [, ], recurrent neural nets (RNNs) (not strictly following the attention inductive bias) [, , ], and state-space model (SSM)-inspired neural net layers  [, , ].

##### Improving LLMs for Long Contexts.

We share motivation various works focus on improving context length quality [, , ]. They that the rotary embeddings [] of popular LLMs [, , , ] struggle to handle contexts longer than those used in training. [] notably study the training efficiency of this via LoRA. Our work is complementary, where we reduce the efficiency bottleneck by replacing self-attentions. Further work lies in how to improve long context quality of linearized LLMs, e.g., by integrating these methods. Other works also improve efficiency by optimizing the KV Cache from a systems standpoint []. These may present complementary optimizations to linear attention’s algorithmic improvements.

## Appendix E Additional Artifacts

### E.1 Sample Generations

We include sample generations for various prompts, using Llama 3 8B models linearized with different \name-related linearizing setups on the Alpaca dataset. For reference, we also include responses from a standard (softmax attention) Llama 3 8B finetuned with LoRA (rank = 8, alpha = 16, applied to all attention projections) on the Alpaca training set. \name as presented often leads to greatest coherence among linearized LLMs.

### E.2 Sample Attention Weights

We finally visualize sample attention weights from \name Mistral 7b and Llama 3 8B layers, computed over held-out 1024-token packed Alpaca samples. Despite only training to match attention outputs via an MSE loss, we find \name often recovers softmax attention weights qualitatively, with comparable quality to explicitly matching attention weights via a cross-entropy loss (XENT). We validate that this is not just due to upweighting softmax attention terms in our sliding window setup, plotting query-key interactions for 1024-long samples between the last 32 queries and the first and last 32 keys (spanning distances well over the 64-token window sizes). \name linear attentions frequently recover attention weights over these distances, suggesting that we learn both linear attention feature maps and weighting factors to match softmax attention weights.

Finally, we note two observations on \name training and architectures. Newly initialized \name linear + sliding window attentions do not capture softmax attention weights (init.), supporting earlier results that attention transfer is necessary to reduce architectural differences and improve linearizing efficiency. Trained \name attentions also match attention weights better than trained Hedgehog linear attentions (same feature map, but no sliding window). These results suggest \name attention transfer and linear + sliding window layers allow us to learn better approximations of softmax attention weights, coinciding with improved lienarizing quality.

Figure: Figure 14: Llama 3 8B attention weights; head 0; layers 0, 8, 16, 24, 31.
Refer to caption: x2.png

Figure: Figure 15: Llama 3 8B attention weights; head 1; layers 0, 8, 16, 24, 31.
Refer to caption: x3.png

Figure: Figure 16: Llama 3 8B attention weights; head 2,; layers 0, 8, 16, 24, 31.
Refer to caption: x4.png

Figure: Figure 17: Llama 3 8B attention weights; head 3; layers 0, 8, 16, 24, 31.
Refer to caption: x5.png

Figure: Figure 18: Mistral 7B v0.1 attention weights; head 0; layers 0, 8, 16, 24, 31.
Refer to caption: x6.png

Figure: Figure 19: Mistral 7B v0.1 attention weights; head 1; layers 0, 8, 16, 24, 31.
Refer to caption: x7.png

Figure: Figure 20: Mistral 7B v0.1 attention weights; head 2; layers 0, 8, 16, 24, 31.
Refer to caption: x8.png

Figure: Figure 21: Mistral 7B v0.1 attention weights; head 3; layers 0, 8, 16, 24, 31.
Refer to caption: x9.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: - How were the hyperparameters for the attention transfer and low-rank adaptation chosen?
soundness: 3
strengths: - I believe this is a good work. The paper presents a novel method for linearizing LLMs that addresses key limitations of existing techniques. By focusing on approximating softmax attention outputs and using low-rank adjustments, the authors offer a fresh perspective on reducing computational complexity without sacrificing model quality.
- The LOLCATS method reduces the amount of training required, both in terms of model parameters updated and training data used. This efficiency makes the method practical for widespread use, especially in environments with limited computational resources.
- Demonstrating the ability to linearize large models up to 405B parameters is a notable achievement. The scalability of the method suggests it can be applied to future, even larger models.  The authors provide extensive experimental results, including comparisons with existing methods on multiple benchmarks. 
- The paper delves into the reasons why previous linear attentions struggled to approximate softmax attention effectively. By identifying issues like attention entropy and layer-wise errors, the authors provide valuable insights that inform their improved architecture.
summary: This paper presents LoLCATs, a method for converting large language models (LLMs) with quadratic attention complexity into models with linear complexity while maintaining model quality. The key innovation is a two-step approach: (1) attention transfer - training linear attention layers to directly approximate the original softmax attention outputs, and (2) low-rank adaptation (LoRA) fine-tuning to adjust for approximation errors. The authors demonstrate that LoLCATs can effectively linearize models up to 405B parameters with significantly less compute and data compared to previous methods, while better preserving model capabilities.
weaknesses: - The introduction and analysis of the preliminaries are too lengthy, with the core improvements in LOLCATs not appearing until the seventh page, which makes the reading experience somewhat disjointed.
- The experimental results are promising given the training budgets. However, I notice that even though previous works perform significantly worse for more challenging benchmarks (like MMLU in the setting), LOLCAT still considerably underperforms the original models. What will happen when the benchmarks become more challenging? (e.g., complex reasoning).
- While the authors claim improvements in inference efficiency, quantitative metrics such as actual speedup factors, memory utilization during inference, or comparisons of throughput are not extensively reported.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: 1. I would like the author to supplement and polish the article based on the weaknesses.

2. The authors mention that LoRA effectively reduces approximation errors. Have they considered conducting a specific error propagation analysis? This is crucial for understanding how low-rank adaptation accumulates errors across different layers in deeper models, as it is vital for controlling cumulative errors.

3. In the LoRA method, the choice of the rank parameter is critical for the model's approximation effectiveness. How have the authors taken into account the impact of low-rank parameters on performance across different layers of the model?

4. Does this approach extend to other variants of Linear Attention, such as Mamba 1/2 and Hgrn 1/2?
soundness: 3
strengths: 1. The LOLCATS approach significantly reduces computational and memory costs through a two-step process involving "attention transfer" and low-rank adaptation (LoRA), effectively lowering the complexity of large models. 

2. LOLCATS effectively retains the performance of the original self-attention model. Experimental results show that this method can recover much of the original model's quality after linearization, using only a small portion of the parameters and training data.

3. LOLCATS is the first to successfully apply linearization to large models with 70B and 405B parameters, expanding the applicability of linearization techniques.
summary: The paper addresses the efficiency and scalability challenges of large language models (LLMs) caused by the quadratic complexity of traditional Transformer models. To overcome these limitations, it introduces LOLCATS (Low-rank Linear Conversion with Attention TranSfer), a novel method that linearizes attention mechanisms to reduce computational and memory demands. LOLCATS uses an “attention transfer” phase to approximate softmax attention efficiently and a low-rank adaptation (LoRA) to correct errors. This approach enables scalable training of LLMs with up to 405B parameters—50 times larger than previous models—while maintaining high performance. Experimental results show LOLCATS outperforms existing methods and opens new avenues for scaling LLMs further.
weaknesses: 1. The paper mentions the differences in errors across layers and their impact on model performance, but it does not sufficiently discuss the underlying reasons, such as why lower soft-attention entropy leads to higher errors.

2. There is insufficient explanation for new terms, such as "attention transfer," which may lead to misunderstandings regarding specific implementation details and processes. Clearly defining key concepts within the paper would improve overall clarity.

3. The paper does not provide averages, variances, or confidence intervals from multiple experiments, making it difficult for readers to assess whether performance differences in the model are statistically significant and robust.

4. As indicated in [1], there is a notable difference between Linear Attention and Softmax Attention in retrieval tasks, particularly in "needle-in-a-haystack" scenarios. Therefore, assessing the performance of models on such tasks before and after the application of LoLCAT would provide a more comprehensive validation of the method's effectiveness.

[1] Xuyang Shen, Dong Li, Ruitao Leng, Zhen Qin, Weigao Sun, & Yiran Zhong. (2024). Scaling Laws for Linear Complexity Language Models.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: - Why was the proposed method not validated on a smaller model, such as llama1B?
soundness: 3
strengths: - The method is shown to scale to unprecedentedly large models (70B and 405B parameters), which was not possible with previous linearization techniques.
- The method is applied to various models, showing its broad applicability and potential impact on the field of natural language processing.
summary: The paper presents a novel method for linearizing large language models (LLMs) to make them more efficient in terms of memory and compute resources. The authors propose a method called Low-rank Linear Conversion via Attention Transfer (LOLCATS), which aims to replace the quadratic attention mechanisms in popular Transformer-based LLMs with subquadratic alternatives, such as linear attention. This approach avoids the expensive pretraining costs associated with traditional LLMs.
weaknesses: - There is a lack of an overall summary description of the LOLCATS method. An algorithm description or pseudo code can be added.
- There are some writing errors, such as Line 294: APP.???

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. Can authors point out which line of appended code correspond to a. Eq 7?

2. I'd like to confirm the author is doing post-training approximation or training from scratch with this specialized architecture. (It's a bit unclear to me). And for all the baseline methods, are they doing training from scratch or post-training approximation.

3. My understanding is that actually there is no new technical thing proposed. Eq 5 is proposed by previous works, and sliding window attention was also proposed in the previous work. And it seems like combining this 2 idea is the only technical contribution. Is that correct? Or do you have other new proposed approximation that I overlooked? 

3. In that regard, adding a sliding window to any of the previous linearization techniques should improve the performance, and this has to be validated. If not, there is a need to investigate why synergy only happens here.

4. Apologize that I really don't get that's the purpose of Table 1, Figure 3 and Figure 5. What's the main point you want to convey here? 

5. Is there a study of different window size?

6. For ablation study I only see Table 5. Is there one with your proposed method? Or should I read it as Hedgehog == Eq 5 and your method is nothing butHedgehog + Sliding?  And somehow Hedgehog + sliding ==> 68.78 outperformed your proposed method a lot, which reads weird to me.

7. Also don't know why ablation is only done on MMLU, I am interested in seeing RCV-e and PiQA too.

8. I don't really see what's the help from the linearization to parameter-efficient finetuning. I still don't understand why PEFT can't work on other methods. In addition, full FT should still outweigh PEFT. I believe authors should also show their method with full parameter tuning to see what's the difference.
soundness: 2
strengths: - The Introduction reads well (but unfortunately not the case for the rest)

- An interesting problem to study.

- Provide the code in the Appendix
summary: Authors proposed a new method to approximate the quadratic attention operation to semi-linear ones (with a pre-specified quadratic window). The linearized approximation seems to work with PEFT so the computational burden is partially alleviated when accelerating the computation on large models.
weaknesses: - The way authour present their methods interleaves many previous works, and thus it's difficult to precisely pinpoint what's their contribution. 

- Some lacking experiments make me uncertain if the claimed effectiveness is true.

- The above two factors combined make me uncertain what's the real efficacy of the method.

Meta-review: ## Summary

The paper introduces LoLCATs, a method proposed for improving the efficiency and scalability of large language models (LLMs) by replacing quadratic softmax attention with subquadratic linear attention. LoLCATs use a two-step process: attention transfer, where linear attention approximates softmax attention with minimal error, and low-rank adaptation (LoRA) to refine the model’s quality. This method significantly reduces computational requirements while maintaining competitive performance. LoLCATs enable training larger models, including the first linearized 70B and 405B parameter LLMs, with reduced memory and token costs. Experiments show notable quality improvements over previous approaches, narrowing the performance gap between linearized and original LLMs.

## Decision

The proposed idea is novel and relevant to improving the efficiency of frontier models. The method is shown to scale to large llama3 models (70B and 405B parameters), which was not possible with the previous linearization techniques. The method is applied to various models, showing its broad applicability and potential impact on natural language processing. The results are convincing, and I believe that publishing this paper would be beneficial for the LLM community.

Decision: Accept (Poster)