LoLCATs: On Low-Rank Linearizing of Large Language Models
Authors: Michael Zhang
Simran Arora
Rahul Chalamala
Benjamin Frederick Spector
Alan Wu
Krithik Ramesh
Aaryan Singhal
Christopher Re
Venue: ICLR 2025 Poster (2025)
Keywords: Linear Attention
Linearizing Transformers
Low-rank Adaptation
Large Language Models
Architecture Distillation
Primary area: foundation or frontier models, including LLMs

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 (LoLCATs), 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). LoLCATs 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, LoLCATs does so with only 0.2% of past methods' model parameters and 0.04-0.2% of their training tokens. Finally, we apply LoLCATs to create the first linearized 70B and 405B LLMs (50$\times$ that of prior work). When compared with prior approaches under the same compute budgets, LoLCATs 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.

Published as a conference paper at ICLR 2025
LOLCATS: ON LOW-RANK LINEARIZING OF LARGE LANGUAGE MODELS
Michael Zhang"1, Benjamin Spector", Aaryan Singhal', Krithik Ramesh 188, and Christopher Réi Stanford University, *Together AI, Califoomia Institute of Technology, MIT mzhang@es stanford. edu
Simran Arora "
Rahul Chalamala Alan Wu
ABSTRACT
Recent works show we can linearize large language models swappigs the quadratic attentions of popular Transformer-basee LLMs with subquadratic analogs, such as linear attention attention-avoiding the expensive pretraining costs. How- ever, 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 (LoLCATS), simple two-step method that improves LLM linearizing quality with orders magnitudes less memory and compute. We base these steps on findings. First, we replace 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 ("altention transfer"). Then, this enables adjusting for approximation errors and recovering LLM quality sim- ply with low-rank adaptation (LoRA). LOLCATS significantly improves lineariz- ing quality, training efficiency, and scalability. We significantly reduce the lin- earizing quality gap and produce state-of-the- -he-t-rrttttttt subquadratic LLMs from Llama 3 8B and Mistral 7B v0.1, leading to 20+ points of improvement on 5-shot MMLU. Furthermore, LOLCATS does so with only 0.2% of past methods' model parameters and 0.04-0.2% of their training tokens. Finally, to create the first linearized 70B and 405B LLMs (50x that prior work). When compared with prior approaches under the same compute budgets, LOLCATS significantly improves lincarizing quality, closing the gap between lincarized and original Llama 3.1 70B and 405B LLMs by 77.8% and 78.1% on 5-shot MMLU.
apply LOLCATS
INTRODUCTION
"Linearizing" large language models (LLMs)-ad or converting existing Transformer based LLMs into attention-fin free or subquadratic alternatives has shown promise for scaling up efficient architectures. While many such architectures offer complexity-level efficiency gains, like linear-lime and constant- memory generation, they are often limited to smaller models pretrained
academic budgets (Gu & Dao, 2023; Peng et al., 2023a; Yang et al., 2023; Arora et al., 2024; Beck et al., 2024). In complementary direction, linearizing aims start with openly available 7B+ parameters pretrained on trillions of tokens (AI, 2024; Jiang et al., 2023)-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.
those with
However, better realize this promise and allow anyone convert LLMs into subquadratic models, we desire methods that are (1) quality-preserving e.g., to recover the zero-shot abilities of modern LLMss (2) parameter and token efficient, to linearize L.LMs on widely accessible compute; and (3) highly scalable, to support linearizing 70B+ LLMs available today (Touvron et al., 2023a;b).
Existing methods present opportunities to improve all three criteria. On quality, despite using moti- vated subquadratic analogs such as RetNet-inspired linear attentions (Sun et al., 2023; Mercat al., 2024) or state-space model (SSM)- based Mamba layers (Gu Dao, 2023; Bick et al., 2024; Wang et al., 2024), prior works significantly reduce performance popular LM Evaluation Harness tasks (LM Eval) (Gao et al., 2023) (up to 23.4-28.2 pts on 5-shot MMLU (Hendrycks et al., 2020)). On pa- rameter and token efficiency, to adjust for architectural differences, prior methods update all model

---

Published as a conference paper at ICLR 2025
Linearizing Model Setup
Step 1: Attention Transfer
Step 2: Low-rank Linearizing
Replace attentions
I Layer wise attention matching
End to and next token prediction with LoRA 'We Aneanzed the chungus among
Transfrmer
Linear Attention
Softmat Attention Pass
fn
Aes Jell
W
W
W
We
Wh
W
N
"The cot sot the rat."
Lowwank
"We lnnannzed the chungus among ut."
Figure 1: LOLCATS framework. We linearize LLMs by (1) training attention analogs to approximate softmax attentions (attention transfer). before swapping attentions and (2) minimally adjusting (with LoRA).
parameters in at least one stage of training (Mercat et al., 2024; Wang et al., 2024; Bick et al., 2024), and use 20- 100B tokens 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 L.LMs (Dubey et al., 2024).
In this work, we thus propose LOLCATS (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) (Hu et al., 2021) ("low-rank linearizing").
In evaluating this hypothesis, we make several contributions. First, to better understand lineariz- ing feasibility, we empirically study attention transfer and low-rank linearizing with existing linear attentions. While bntuitive-by swapping in perfect subquadratic softmax attention approximators, we could get subquadratic LLMs with no additional training-priin -pinr works suggest linear attentions struggle to match softmax expressivity (Keles et al.. 2023; Qin et al.. 2022) or need full-model up- dates to recover linearizing quality (Kasai et al., 2021; Mercat et al., 2024). In contrast, we find that while either attention transfer or LoRA alone is insufficient, we can rapidly recover quality by simply doing both (Figure 3, Table 2). At the same time, we do uncover quality issues related to attention-matchin earized 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 (Zhang et al., 2024), we also find that later layers result in 200 x the MSE of earlier ones (Figure 7). We later find this issue aggravated by larger LLMs; jointly training all Llama 3.1 405B's 's 126 attention layers fails to viably lincarize.
architecture and training. With prior linear attentions, the best low-rank lin-
Next, to resolve these issues and improve upon our original criteria, detail LOLCATS' method components. For quality, we generalize prior notions of learnable linear attentions to sliding win- dow + linear attention variants. These remain subquadratic to compute yet consistently yield better attention transfer via lower mean-squared eтror (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 & 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 LOLCATS improves on each of our desired criteria.
On quality, when linearizing popular LLMs such as Mistral-7B and Llama 3 8B, LOLCATS substantially 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, LoLCATS for the first time closes the zero-shot LM Eval gap between linearized and Transformer models (73.1 vs 74.2 pts), while supporting xx the throughput and 64x the batch sizes vs. popular F (Dao, 2023)
2

---

Published as a conference paper at ICLR 2025
Name
Architecture
Quality Preserving
Parameter Ellicient
Efficient
Token Validated
Scale
Mistral LoLGAT (Ourg)
Cama 3 68 LOLGKS (Durs)
Uama 3 88
Pretrained
Allention
//
xx
Zero-shot Score (Avg.)
Zembba
SUPRA
Lincar Attention
X
K
Grittn
Mohawk
Mamba (2)
X
x
Mamba in Llama
Mamba (2)
Tioc eWWKV ve 78
Soflmax Approx
LoLCATs
/
W
Linear Attention
Figure 2: Linearizing comparison. LoLCATs significantly improves LLM linearizing quality and training efficiency. No. of or indicate relatively better or support.
Unearoad Supqasdistc
2-Umma
Training Takens (Log Scale)
models (generating 4096 tokens on an 80GB H100). As an alternative to pretraining, LOLCATS outperforms strong pretrained 7B subquadratic LLMs and hybrids by 1.2-9.9 pts (LM Eval avg.).
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, LOLCATS enables these gains while updating only <0.2% of past linearizing methods' model parameters, making a single 40GB GPU sufficient for 7B LLM linearizing. This also only takes 40M tokens, i.e., 0.003% and 0.04% of prior pretraining and linearizing methods' token counts.
On scalability, with LOLCATS we scale up linearizing to support Llama 3.1 70B and 405B
LLMs (Dubey et al., 2024). LOLCATS presents the first viable approach to linearizing larger LLMs. We create node, and the first linearized 405B LLM with combination of 5 hours on 14 80GB H100 GPUs (attention transfer) + 16 hours on three 8X8GGB 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 80GB A100s) (Wang et al., 2024). Furthermore, under these computational constraints, LOLCATS significantly improves quality versus prior linearizing approaches without attention transfer. With Llama 3.1 70B and 405B, we close 77.8% and 78. % of the 5-shot MMLU gap between Transformers and linearized variants respectively.
first linearized 70B LLM, taking only 18 hours on one 8x 80GB H100
2 PRELIMINARIES
To motivate LOLCATS, we first go over Transformers, attention, and linear attention. We then briefly discuss related works on linearizing Transformers and Transformer-base LLMs.
Transformers and Attention. Popular LLMs such Llama 8B (AI@ Meta. 2024a) and Mistral 7B (Jiang et al., 2023) are decoder-only Transformers, with repeated blocks of multi-head softmax attention followed by MLPs (Vaswani et al., 2017). For one head, attention computes outputs y E Rlxd from inputs x E Rixd (where is sequence length, d is head dimension) with query, key, and value weights W q: WE, W. Rdddd, causal language modeling, we compute q = x Wg k wWE x W.. o. before getting attention weights and outputs y via
exp(q, kt/ Vd)
On,i E exp(q, k/
yn = Σ Dn, for n in [1,....]
(1)
Multi-head attention maintains inputs, outputs, and weights for each head, egg. x ghxttd or Wg E hxxxdx (h being number of heads), and computes Eq. 1 for each head. In both cases, we compute final outputs by concatenating yn across heads, before using output weights W, E Ridxxdd to compute yr W. Rixhd, While expressive, causal softmax attention requires all {k,, v; Jisn to compute yn- For long context or large batch settings, this growing KV cache can incur prohibitive memory costs even with state-of-the-an implementations such as FlashAttention (Dao, 2023).
Linear Attention. To get around this, Katharopoulos et al. (2020) show a similar attention op- eration, 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 's exponential is kernel function K(gn, k.), which in general can be expressed as the dot product of feature maps Ø R + gd Swapping exp(q,) k Vd) with o(q.) ((k:) in Eq. 1 gives linear attention weights and outputs:
(gqn) T (k:)e ((q.)) EE o(k.)o,
in = Σ
(2)
Di=1 p(gn)) T
((g.)) T ((ks)
3

---

Published as a conference paper at ICLR 2025
This lets compute both the numerator 8, Ei=1 ( and denominator 2, Σ o (k;) as recurrent "KV states". With So = 0, 20 = 0, we recurrently compute linear attention outputs as
In (qqn) 8n for 8n = 8n + ((kn),, and Zn = 1 + ((kn)
2n Zn
Eq. 2 lets us compute attention over an imput sequence of length n in O(ndd") time and space, while Eq. lets us compute n new tokens in (nrdd") time and (cd'") memory. Especially during generation, when softmax attention has to compute new tokens sequentially anyway, Eq. 3 enables time and memory savings if d' < (prompt length + prior generated tokens).
Linearizing Transformers. To combine efficiency with quality, various works propose different 6, (e.g.. 6zz) = ELU(x) as in Katharopoulos et al. (2020)). 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 (Mercat al., 2024), and those for smaller Transformers--e. 110M BERTs (Devlin al., (118))-reasnnan adaptable to modern LLMs (Kasai et al., 2021; Mao, 2022; Zhang et al., 2024).
3 METHOD: LINEARIZING LLMS WITH LOLCATS
We now study how build 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 Sec- tion 3.2, we find that while this attention transfer works surprisingly well for low-rank linearizing with existing linear attentions, on certain tasks, still results in sizable quality gaps compared to prior methods. We also find that attention-transfer quality strongly corresponds with the final lin- earized model's performance. In Section 3.3, we use our learned findings to overcome prior issues, improving attention transfer subsequently improve low-rank linearizing quality.
3.1 LOLCATS PART I: A FRAMEWORK FOR LOW-COST LINEARIZING
In this section, we present our initial LOLCATS 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 minimal amount of subsequent training-e,8, that supported by low-rank uo recover LLM quality in a cost-effective manner effectively. We thus proceed in two steps.
1. Prrameter-Efficiieen Attention Transfer. For each softmax attention in LLM, we aim to learn appsxiaippro linear attention, i.e., one that computes attention outputs ý ~ for all natural imputs x. We call this "attention transfer". as we aim to transfer the attention modeling of existing softmax attentions into target linear attentions. Due to architectural similarity, we can treat this as a feature map learning problem, learning é to approximate softmax. For each head and layer, let dy and Φx be query and key feature maps. Per head, compute:
exp(q, k./Vd)
dg (gx) (k;)
Yn
Vi E exp(a, k. V.
yN Σ Σn == (g(an) &u (k:)
U
(4)
Softmax Attention
Lincar Attention
for all n E [1] with input E igldd and train Og de to minimize sample mean squared error (MSE)
H hm MSE
d Σ Cyn Dn,²
(MSE
ph.m MSE
MH
d
M= 1=1
==1
i.e., jointly for each head h in layer m. Similar to past work (Kasai et al., 2021; Zhang et al.,
2024), rather than manually design o, we parameterize each Φ IRd + Rd" as learnable layer:
(g.))) := f(g, ( + b()), db(ki) == fk; W()) + b(())
Here W E daxd and b E Rd are trainable weights and optional biases, f(.) is a nonlinear activation, and d' is an arbitrary feature dimension (set to equal head dimension d in practice).
4

---

Published as a conference paper at ICLR 2025
2. Low-rank Adjusting. After training the linearizing layers, replace the full- parameter train- ing of prior work with low-rank adaptation (LoRA) (Hu et al., 2021). Like prior work, to ad- just for the modifying layers and recover language modeling quality,
now train the mod- ified LLM end-to-end over tokens to minimize a sample next-token prediction loss fsem = Σ log (e ui:)). Here Pe is the modified LLM,, θ is the set of LLLM parameters, and we aim to maximize the probability of true Uz+1 given past tokens U1st (Fig. right). However, rather than train all LLM parameters, we only train the swapped linear attention Wg. q. WK, W. W. with LoRA, updating W < W + AW with AW as the product of two low-rank matrices BA, B E RddT, A E Rexd. This enables parameter efficiency for < d.
Training footprint and efficiency. Both steps remain parfmeter-efficieet.. For Step I, optimizing Eq. 5 is similar to a layer-by-layer rross-architecture distillation. We compute layer-wise (x, y) as pretrained attention imputs and outputs, using LLM forward pass over natural language samples (Fig. 1 middle). However, to keep our training footprint low, we freeze the original pretrained atten- tion layer's parameters and simply insert new oy dg after W,, We in each softmax attention (Fig. left). We compute outputs y, D with the same attention weights in separate passes (choosing either Eq. or Eq. 2: "teacher-foreing" by only sending softmax outputs y to future layers and preventing error propagation, Fig. 1 middle). For Llama 3 8B or Mistral 7B, training ②y. Φ& with d' = 64 then only takes 32 layers x 32 heads x 2 feature maps x (128 x 64) weights Rs 16.8M trainable weights (0.2% of LLM sizes). Furthermore, Eq. 5 keeps attention transfer memory-efficient et al. (2024) also train layer-wise & to approximate softmax attention, they supervise by matching on the n² attention weights computed for n-token inputs (App. D, Eq. 12). This scales poorly for large n. Instead, by only needing attention outputs, we can reduce memory from Onn²) to O(n), comput- ing softmax attention y with FlashAttention (Dao al., 2022; Dao, 2023) and linear attention ý with Eq. (or 6) both in (n) memory by not materializing attention weights. Despite these savings, we show similar attention weight recovery in App. E.2. In Step 2, LoRA with 8 on all attention projections suffices for state-of-the-ar quality. This updates just <0.09% of 7B parameters.
While Zhang
3.2 BASELINE STUDY: ATTENTION TRANSFER AND LOW-RANK LINEARIZING
As a first step, we aim to understand if attention transfer and low-rank adjusting as proposed are sufficient for linearizing LLMs. While simple, it is unclear whether these steps can lead to high- quality linearizing, as all prior works default to more involved approaches (Mercat et al., 2024; Bick et al. 2024; Wang al., 2024). They use linearizing layers featuring GroupNorms (Wu & He, 2018) and decay factors (Sun et al., 2023), or alternate SSM-based architectures (Gu & Dao, 2023; Dao & Gu, 2024). They also all use full-LLM training after swapping in the subquadratic layers. In contrast, as a first contribution 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.
Experimental Sctup. We test the LOLCATS framework by linearizing two popular base LLMs, Llama 3 8B (AI, 2024) and Mistral 7B v0.1 (Jiang et al., 2023). For linearizing lay- ers, we study two feature maps used in prior work (Table 1). To support the rotary posi- tional embeddings (RoPE) (Su et al., 2024) in these LLMs, we apply the feature maps after RoPE," i.e, computing query features (g) wish to see if LOLCATS with small amount of data can recover general zero-shot and instruction- following LLM abilities. We use the 50K samples of a cleaned Alpaca dataset', due to its ability to improve general cestruction-olloovinn in 7B LLMSs despite its relatively small size (Taori et al., 2023). We train all feature maps jointly. Training code and implementation details are in App. C.
Feature Map
o(g) (same for k)
Weight Shapes
T2R
ReLUig +b)
W: (128, 128), b: (128.)
Hedgehog (SM.(gW) @ SMa(-gW)) W: (128, 64)
Table 1: Learnable feature maps. Transformer to RNN (T2R) from Kasai et al. (2021), Hedgehog from Zhang et al. (2024), both (concat) and SMa (softmax) apply over feature dimension.
[(RoPE(g)) W,, b))- For linearizing data, we
To study the effects of attention transfer and low-rank linearizing across LLMs and linear attention architectures, we evaluate their validation set perplexity (Table 2, Fig. 3) and downstream LM Eval
Unlike prior works that apply & before RoPE (Mercat et al., 2024; Su et al., 2024), our choice preserves the linear attention kernel connection, where we can hope to learn oy, 6x for exp(q. ks/vd) so dg(g.) dh (kx). httpps: / //wgggi ing aee co/ /at//attassetee aaamaa paca cl cannd
5

---

Published as a conference paper at ICLR 2025
Ulma 3 89, T2R
Liama 3 86, Hedgenog
Mistral 7B. T2R
Mistral 7B, Hedgshog
der Sunddr
12"
22
12
Thaniig Siopo
Tiaal Thaiiing Shipo
Trla Toxking Stoz
Figure 3: Attention transfer training efficiency. Even accounting for initial training steps, low-rank lineariz- ing with attention transfer still consistently achieves lower perplexity faster across feature maps and LLMs.
zero-shot quality (Table 4). We use the same data for both stages, early stopping, and either 2 epochs for attention transfer and LoRA adjusting or 4 epochs for either alone (ss 40M total training tokens). We use LoRA r = 8 by popular default (Hu et al., 2021), training 0.2% of LLM parameter counts.
Llama 3 8B
Mistral 7B
Attention
T2R
Hedgehog
T2R
Hedgchog
Transfer?
PPL@0
PPL@2/4
P
PPL@2/4
PPL@0
PPL@2/4
PPL
PPL@2/4
No x
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
Table 2: Alpaca validation set perplexity (PPL.) of linearized L.LMs, comparing attention transfer, no LoRA adjusting (PPL.@D) and PPL. after training (PPL. @2/; 2 with attention transfer, without, for equal total steps).
Attention Transfer + LoRA Enables Fast LLM Linearizing. In Table 2 and Fig. 3, 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 2), single low-rank update rapidly recovers performance by 15-75 PPL (Fig. 3), 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 (Mercat al., 2024; Wang et al., 2024) (within 5 accuracy points on 4 / 6 popular LM Eval tasks; Table 4), 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 orders-of- mf--enniude less training costs than previously shown.
LOL SAD: Limitations of Low-Rank Linearizing. At the same time, we note quality limitations with the present framework. While sometimes close, low-rank lincarized LLMs perform worse than full-parameter alternatives and original Transformers on 5 5 6 LM Eval tasks (up to 42.4 points 5- shot MMLU; Table 4). To understand the issue, we study if the attention transfer stage can produce close linear attention approximations of LLM softmax attentions. We note three observations:
1. Attention transfer quality (via output MSE) strongly ties to final linearized LLM quality (via PPL.) (Fig. 5), suggesting we can improve quality by reducing MSE with softmax attentions.
2. However, larger MSEs coincide with lower softmax attention weight entropies (Fig. 6a). Zhang et al. (2024) find linear attentions struggle to approximate such "spikier" distributions, suggesting we may need better attention- matching layers to reduce MSE and improve final linearized quality.
3. When training layers jointly like in prior work (Zhang et al., 2024), larger MSEs also heavily concentrate in later layers (Fig. 6b). To bring the MSE in these layers down, we may thus need more fine- fine-grained grained layer-wise supervision, rather than the objective over all layers in Eq. 5.
Liama
Mistral
85
5.4
Medel
Tokess RIQA
ARC-F ARC-C
MMLIU
ledggehog
Heddgahog
Liama 3 MB
72.9
H0.3
53.3
Mamba2
100
74.1
7
4.8
LoRA Holgsheg
700
Miseril
SLIRA
100
4.4
LoRA Hedgehog
76.4
44
Figure 4: Linearizing comparison LM Eval. Task names in Table 4. Acc. norm: ARC-C, HS. Acc. otherwise. 5-shot MMLU, 0-shot otherwise.
55
3.
325
Sage
Siage
Figure 5: Attention MSE vs. PPL. Across feature maps,
LLMs; lower MSE coincides with better linearized quality.

---

Published as a conference paper at ICLR 2025
Liam 8U, Linear Attention
0.0100
0.0075
0.025
0.0200
0.0050
0.015
0.0025
6 19
0.0000
[1.68. 1.90] 13.56. 2.22] 12.23, 2:51 (2.51.2.7) 12.79, 3.06
8
12
16 Layers
20
28
Softmax Attention Weight Entropy
(a) Attention Entropy
(b) Model Layer
Figure 6: Sources of Attention Transfer Error with Llama 8B. We find two potential sources of
attention transfer difficulty: (a) low softmax attention entropy and (b) attentions in later layers.
3.3 LoLCATS PART 2: COMPONENTS TO IMPROVE LANK LINEARIZING
Following our motivating hypotheses, framework, and observations, we now introduce two simple improvements to improve linearized LLM quality by reducing MSE: (I) better attention attention-matching architectures (Section 3.3.1), and (2) finer-grained layer-wise attention transfer (Section 3.3.2).
3.3.1 ARCHITECTURE GENERALIZING LEARNABLE LINEAR ATTENTIONS
described, we can apply our framework with any linear attentions with learnable é (e.g. T2R and Hedgehog, Figure 3). However, to improve attention- aatention-matching quality, we introduce a hybrid ø parameterization combining linear attention and sliding window attention. Motivated by prior works that show quality improvements when combining attention layers with linear attentions (Arora et al., 2024: Munkhdalai et al., 2024), we combine short sliding windows of softmax attention (Beltagy et al., 2020; Zhu et al., 2021) (size 64 in experiments) followed by linear attention in single layer. This allows attending to all prior tokens for each layer while keeping the entire L.LM subquadratic. For window size and token indices [1,..., W.,... - n). apply the softmax attention over the w most recent tokens, and compute attention outputs Yn as
Yn Eien-w+ eepppp (., k./ Vd cn)us + (g (E R- j=1 On (kz) E w+1 exp(q.) k./Vd c..) + (g(gn) (E
(6)
j=1 &z (k₂))
y is a learnable mixing term, and en is a stabilizing constant as in log-sum-exp calculations (Cn, = max, kg [n w + I.....,(}}) Like before, we can pick any learnable ø.
Subquadratic efficiency. The hybrid layer retains linear time and constant memory generation. For n-token prompts, we initially require O(w2d) and O(nn w)dd") time and space for window and linear attention respectively, attending over a w-sized KV-cache and computing KV and K-states (Eq. 3). For generation, we only need Oww2d + dd') time and space for every token. We evict the KV-cache's first k. v, compute (k)),, and add de (k)o and Φx (k) to KV and K-states respectively.
Hardware-aware implementation. To make Eq. 6 competitive with modern softmax attentions like FlashAttention-22 (Dao, 2023), we provide "hardware-award" Eq. 6 with the Hedgehog feature map. For space, we defer implementation details App. A, C.2. We evaluate this version by default.
0.0000
Liama 88, Linear Actention
Liama 88, Linear Attention Sliding Window
00693
0.0080
0080
6.0020
00681
16
20
24
28
Solbrax Attendon Weight Entropy
Layer
Figure 7: Improving Attention matching MSE. Linearizing with linear + sliding window attention better matches LLM softmax attentions (lower MSE) over attention entropy values and LLM layers.
7

---

Published as a conference paper at ICLR 2025
3.3.2 TRAINING: LAYER (OR BLOCK)- WWIS ATTENTION TRANSFER
We describe our training approach and provide a simplified model to show its cost-quality tradeoff'. Based on the layer-wise MSE differences in Sec. 3.2 from training all layers jointly (Eq. 5), we instead generalize to training over finer- imer-grained b-layer blocks, and train each block independently:
6+6
1 bH
pblock MSE
ch.m MSE
(for blocks starting at layers 0, b, 2b....)
M A=1
We choose to balance quality and cost in both memory and training time, which we find partic- ularly helpful for linearizing larger LLMs. Several block-wise training approaches exist, including joint training with separate optimizer groups per block, sequentially training separate blocks, or pre- computing hidden states and training blocks in parallel across GPUs. For space, report results on LLM quality trade-offs in App. B.7.3, where we find smaller improves LOLCATS Llama 3.1 405B PPL. by 1.02 points (Table 24). Below we discuss primary cost trade-offs:
Compute: While the joint training of Llama 3.1 405B in 16-bit precision uses multiple nodes (e.g.. NVIDIA H100 x 80GB nodes), an individual block of b = 9 or fewer layers can be trained on a single GPU (e.g.. H100 80GB GPU) at sequence length 1024.
Time: However, doing so on a single GPU multiples total training time by M/b. With multiple GPUs, we can distribute training different blocks on difference devices in parallel. To train blocks at layers b, 2b... we need the outputs from each prior block. As we teacher-force with the "true" softmax attention outputs, we can simply precompute these with the original Transformer.
Lastly, we need to save each precomputed block's outputs disk. The total disk space required is 2 x T x d x & for total training tokens T, model dimension d, number of layers L and 2-byte (16-bit) precision. For Llama 3.1 405B, saving states per-layer = 1) for just 50M tokens requires over 200TB of disk space. Larger b divides this storage, potentially at the cost of quality.
Memory:
For each target LLM, we thus aim to make low-rank linearizing feasible by first doing block wise attention transfer, adjusting parameters based on linearizing quality and what is feasible in compute, time, and memory. We summarize LOLCATS with Alg. 1, 2, providing pseudocode in App. C.
Algorithm LOLCATS Step 1: Attn. Transfer
Algorithm 2 LOLCATS Step 2: LoRA Adjust
Input: Pretrained Transformer with atrn layers; inpor tokens u Input: Lincar attn. feature params. {d2" 4t" y : m [M]], window size wi; layer biock-size b, learning rate a
Input: Lineorized Transformer with M attn. layers; input tokens Imput: Trained feature map params. { 6z .7" window size w:; learning rate a [M do
[M]]
1: Freeze all Transformer parameters
1: for atin. layer
5 (Add LoRA weights) Init. LoRA A F2th for each WT WT W W
2: Initiatize block-wise losses { f- 0 [M
4: for attn. layer
3: Compute initial attn, inpat ze 4- embed(u)
[ad] do
Compute attentions
3. for attn. layer m [Md] do
2: Compute initial atm. input a + embed(u) (u)
D Compute attentions
5: q. k. v = x=" wm "x WT", "."".""
4:
q x ae""(W (W."
13 ATD)
> (LARA forward pass)
6: v sof:max.attanq
w) (Eiq.)
b (No grad)
5:
m (W + BT"AA) Ag)
7:
5 = =-inear_attn(q l. v, 10, 20 "t" o.. >") (Eq. 6)
6:
w
(W
AT)
8:
fmgh fmyy + fuSE (""". """
(Fq. 7)
9Th winer..ttnnng....... k, w. 00 (")) (Eq.
9: m+t mlp" ("""
(Teacher-fore force next layer)
8:
a m=++ = ml (W" B" A!"))
10:
for atin. layer [Md] do
Update feature maps
9: Compute next-token pred á In. > Train LoRA
11: for weights / f of " do
10: Compute sample loss & froettl (Mgg -1, m....)
12:
Update 0 0 77 p8
11: Update 4- 8 06 for all all
4 EXPERIMENTS
Through experiments, we study: (1) if LOLCATS linearizes LLMS with higher quality than ex- isting subquadratic alternatives and linearizations, and higher generation efficiency than original Transformers (See. 4.1); (2) how ablations on attention transfer loss, subquadratic architecture, and parameter and token counts impact LLM quality (Sec. 4.2); and (3) how LOLCATS' quality and efficiency holds up to 70B and 405B LLMs by linearizing the complete Llama 3.1 family (Sec. 4.3).
4.1 MAIN RESULTS: LOLCATS EFFICIENTLY RECOVERS QUALITY IN LINEARIZED LLMS
In our main evaluation, we linearize the popular base Llama 3 8B (AI, 2024) and Mistral 7B (Jiang et al., 2023) LLMs. We first test if LOLCATS can efficiently create high-quality subquadratic LLMs from strong base Transformers, comparing to existing linearized L.MSs from prior methods. We also
8

---

Published as a conference paper at ICLR 2025
Training Tokens (IB
ARC-c (norm)
HellaSwag (norm)
Wino- grande
MMLU (5-shot)
Avg.
Model
PIQA
ARC-e
Avg.
(no MMLJU)
Mistral 7B
82.1
53.8
74.0
62.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 (Ouns)
0.04
81.5
81.7
54.9
80.7
74.0
51.4
70.7
74.5
Llama 8B
79.9
80.1
53.3
79.1
73.1
66.6
72.0
73.1
Mambu2-Llamu.
20
76.8
74.1
48.0
70.8
58.6
43.2
61.9
65.6
Mamha2 laama 3, 50% Attn.
20
81.5
78.8
58.2
79.5
71.5
56.7
71.0
73.9
Llama 8B Hedgehog
0.04
77.4
71.1
40.6
66.5
54.3
24.2
55.7
62.0
IJama 88 LoLCATS (Ours)
0.04
80.9
81.7
54,9
79.7
74.1
52.8
70.7
74.2
Table LoLCATS comparison among linearized 7B+ LLMss. Among linearized 7B+ LLMss, LOLCATS-
linearized Mistral 7B and Llama 8B consistently achieve best or 2nd-best performance on LM Eval tasks. LOLCATS 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 linearize.
Model
Tokens (B)
PIQA
ARCle
ARC- (acc. norn)
HellaSwag (acc. norm)
Winegrande
(5-shot)
Avg. (W MMLU)
Avg. (no MMLU)
Transformer
Gemma
81.1
53.2
23.7
72.3
24.1
Mistral 7B
8000*
82.1
80.9
53.8
81.0
74.0
62.4
72.4
74.4
Lloma
15000
53.3
Subqundratic
Mamina 7H
1200
81.0
77.5
46.7
77.9
21.8
33.3
64.7
71.0
RWKV-6 World v2.1
46.3
TransNormerL..MM 7B
1400
80.1
75.4
44.4
75.2
66.1
43.1
64.
68.2
Hawk 7B
300
80.0
74.1
45.9
77.6
69.9
35.0
63.8
69.6
Hyhrid Softmax
3tipedHyens-Nons-78
78.8
77.2
40.0
76.4
66.4
26.0
60.8
67.8
Zamba
Linearized
Mistral 7B LoLCATs (Ouns)
0.04
$1.5
81.7
54,9
80.7
74.0
51.4
70.7
74.5
(Ours)
Table LOLCATS comparison to pretrained subquadratic LLMs. SOLCCAT--lincariza Mistral 7B and
Llama 8B outperform pretrained Transformer alternatives by 1.2 to 9.9 points (Avg.), only training 0.2% of their parameter counts 0.013 to 0.003% of their training token counts. 'Reported in Mercat al. (2024).
test if LOLCATS can create subquadratic LLMs that outperform modern Transformer alternatives pretrained from scratch. For space, we defer linearizing training details to Appendix A.
In Table 4, we report results on six popular LM Evaluation Harness (LM Eval) tasks (Gao et al., 2023). Compared to recent linearizing methods, LOLCATS significantly improves quality and training efficiency across tasks and LLMs. On quality, LOLCATS closes 79.8% and 86.6% of the Transformer-lnearrian
gap for Mistral 7B and Llama 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 ef- ficiency, 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 2500x improvement in "tokens-to to-model" efficiency). Among 7B LLMs, --in-rirriee LLMs further outperform strong subquadratic Transformer alternatives, representing RNNs or lin- ear attentions (RWKV-v6 (Peng al., 2024), Hawk, Griffin (De al., 2024). TransNormer (Qin et al., 2023)), state-space models (Mamba (Gu & Dao, 2023)), and hybrid architectures with some full attention (StripedHyenn (Poli al., 2023b), Zamba (Glorioso et al., 2024)).
4.2 LOLCATS COMPONENT PROPERTIES AND ABLATIONS
We next validate that LoLCATs enables subquadratic efficiency, and study how each of LOLCATS' components contribute to these quality gains. We include additional ablations in Appendix B.
Subquadratic Generation Throughput and Memory. We measure generation throughput and memory of LoLCATS LLMs, validating that lincarizing LLMs significantly improve their generation efficiency. We use the popular Llama 8B HuggingFace model?, and compare LoL-
1301
LECPP 1112
Monsie onden
2009
Bashh 301
Figure Generation Efficiency, Llama 8B.
"https://hugg" co/ metta amm /Meta- amlam--3

---

Published as a conference paper at ICLR 2025
LM Eval Metric
Swap & Finetune
+Attention Transfer
+Sliding Window, Attention Transfer
+ Sliding Window, No Attention Transfer
Feature Map
Average
44.20
55.32
70.66
68.78
Hodgchog
MMLU (5-shol)
23.80
23.80
52.77
45.80
Average
38.84
54.83
68.28
39.52
T2R
MMLU (5-shot)
23.20
23.10
40.70
23.80
Table 5: LOLCATS component ablations, linearizing Llama 3 8B over 1024-token sequences. LoLCATS default shaded. Across Hedgehog and T2R feature maps. LOLCATS' attention transfer and sliding window increasingly improve linearized LLM quality. Full task results in App. B.1.1.
PIQA
ARC Easy
ARC Challenge
HellaSwag (ace norm)
WinoGrande acc
MMLU (5-shot)
acc
acc
(ace norm)
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 31.44
LoL.CATS (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
Lincarized, no attn. transfer
81.99
80.89
54.44
82.29
71.19
28.74
LOLCAT's (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
Lincarized, no attn. transfer
84.44
86.62
64.33
86.
79.87
33.86
LOLCAT's (Ours)
85.58
88.80
67.75
87.41
80,35
72.20
Table 6: Linearizing Llama 3.1 8B, 70B, and 405B. Among the first linearized 70B and 405B LLMs (via low-rank linearizing), LOLCAT'S significantly improves zero- and few-shot quality.
CATs implemented in HuggingFace Transformers with the supported FlashAttention--2 (FA2) (Dao, 2023). We benchmark LoLCATS with Hedgehog feature map and linear + sliding window attention in FP32 and BF16 on one 80GB H100. Fig. 8 reports scaling batch size throughput (left) and memory (right). We measure throughput as (generated tokens batch size / total time), with 128- token prompts and 4096-token generations. With larger batch size, LoLCATSS-linerizee LLMs achieve higher throughput than FA2. This corresponds with lower memory, where FA2 exceeds memory at batch size 64. Meanwhile, LOLCATS supports 3000 tokens / second at batch size 2048 (Fig. 8 left), with fixed "KV state" vs. softmax attention's 's growing KV cache (Fig. 8 right).
Ablations. We study how attention transfer and linear sliding window attention in LoLCATS contribute to downstream linearized Llama 3 8B performance (Table 5). We start with prior linear attentions (Hedgehog, Zhang et al. (2024); T2R, Kasai al. (2021)), using the prior linearizing approach that swaps attentions and finetunes the model to predict next tokens (Mercat et al., 2024). We then add (i) attention transfer, (ii) sliding window attentions, or (iii) both. On average LM Eval score and 5-shot MMLU accuracy, LOLCATS' default performs best across feature maps.
4.3 SCALING UP LINEARIZING TO 70B AND 405B LLMS
We finally use LoLCATS to scale up linearizing Llama 3.1 70B and 405B models. In Ta- ble 6, LOLCATS provides the first practical solution for linearizing larger LLMs, achieving sig- nificant quality improvements over prior linearizing approaches of swapping attentions and fine- tuning (Mercat et al., 2024). With the same linear + sliding window layer, LOLCATS gets +39.0 points 5-shot MMLU accuracy on Llama 3.1 70B, and +38.3 on Llama 3.1 405B. These results highlight LOLCATS' 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.
CONCLUSION
We propose LOLCATS, efficient LLM linearizing method that (1) trains linear attentions to match an LLM's self-attentions, before (2) swapping the attentions and only finetuning the replacing attentions with LoRA. This reduces linearizing learning good softmax attention approximations, and we study how to do so. On popular LM Eval tasks, LOLCATS enables LLMs with high quality and inference efficiency, outperforming prior Transformers alternatives while only updating 0.2% of model parameters and requiring 0.003% of LLM pretraining tokens. Our findings substantially improve linearizing quality and accessibility, enabling the first linearized 70B and 405B LLMss.
10

---

Published as a conference paper at ICLR 2025
ETHICS STATEMENT
Our work deals with improving the efficiency of open-weight models. While promising for ben- eficial applications, increasing their accessibility also raises concerns about potential misuse. Bad actors could leverage our technique develop LLMs capable of generating harmful content, spread- ing misinformation, or enabling other malicious activities. We focus primarily on base models, but acknowledge that linearizing could also be used on instruction-tune 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.
L.LMs; research on whether
REPRODUCIBILITY
We include experimental details in Appendix A, and further implementation details with sample code for linearizing architectures and training in Appendix C.I. Our code is also available at /ttps://githiee com/ Haz Hawyyeseaaacc caceaarrch/oolcatooeo 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 grate- fully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF2247015 (Wareware-Aware), CCF1763315 (Beyond Sparsity). CCF1563078 (Volume to Veloc- ity), and 1937301 (RTML); US DEVCOM ARL under Nos. 1911NF-33-2-0184 (Long-context) and
(Interactive Human-AI Teaming); ONR under Nos. 122014 (Deep Signal Processing); Stanford HAI under No. 247183; NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASE, Accenture, Ericsson, Qualcomm, Ana- log 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. Govermment 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, endorsements, either expressed implied, of NIH, ONR, or the U.S. Government.
REFERENCES
URL petttp::/kaiokknndeee wit..biii....................... ceaithub.oooooo context.
Mistral AI. Mixtral of experts. Mistral Al - Frontier Al in your hands, May 2024. URL https:
/ http///w///miiieee -ofrralllo-epperss///o
AI@Meta. Llama 3 model card. 2024a. URL https://github.co wtwp:/thww.ccmmmtttlleleaae
CARD md.
AI@Meta. Llama 3.2:
Revolutionizing edge ai and vision with open, cus- URL https://ai. /ttp:://aaaaaaaa...// com/blog/
tomizable models, Sep 2024b.
1lama- 3- www.econnnnnnnn-neeeeee ootp---ooecte00----------------o
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall- thrlu-throughputt tradeoff. arXiv preprint arXiv:2402.18666 2024.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp. Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint 2405::2405.055//1 2024.
Iz. Beltagy, Mathew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05155 2020.
Aviv Bick, Kevin Y Li, Eric P Xing, Zico Kolter, and Albert Gu. Transformers to ssms: Distilling quadratic knowledge subquadratic models. arXiv preprint arXiv:2408. 10189, 2024.
11

---

Published as a conference paper at ICLR 2025
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle OBrien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp- 2397-2430. PMLR, 2023.
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Uni- fying sparse and low-rank attention approximation. arXiv preprint arXiv:2110.15344 2021a.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15595 2023a.
Yifan Chen, Qi Zeng, Heng Ji, and Yun Yang. Skyformer: Remodel self-attention with gaus- sian kernel and nystr) "om method. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021b. URL https: / //openzeview wwt/tfruuumeeeeeeee wttee//ooormmm.dd211777i.ii/
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309. 12307, 2023b.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Together Computer. Redpajama: An open source recipe reproduce Ilama training dataset, 2023. URL https::/gittiie com/t www.thercopuueeee.....oooooo /RedPa, jjama-Daaa Data.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691, 2023.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.2 1060, 2024.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances Neural Information Processing Systems, 14116344-163599 2022.
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Murann Al- bert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, al. Griffin: Mix- ing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402.1 19427, 2024.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.0444 04805, 2018.
Abhimanyu Dubey, Abhinav Jaubri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama 3 herd of models. arXiv preprint arXiv:2407.21788 2024.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa, Jason Phang, Laria Reynolds. Hailey Schoelkopf, Aviya Skowron, Lin- tang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023. URL https: /ttp:::///enodooo org/ reco/de/ 10356836.
Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. SAMSum corpus: A human- annotated dialogue dataset for abstractive summarization. In Lu Wang. Jackie Chi Kit Cheung, Giuseppe Carenini, and Fei Liu (eds.), Proceedings the 2nd Workshop New Frontiers in Sum- marization, pp. 70-79, Hong Kong, China, November 2019. Association for Computational Lin- guisties. doi: 10. 18.53///vv 101/3//11D119...0// 1/D99-5409. URL htw.s:tp.... / /acl antho ogy. 9- 5409.
12

---

Published as a conference paper at ICLR 2025
Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge. Zamba: compact 7b ssm hybrid model. arXiv preprint arXiv:2405.16712, 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300,, 2020.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2/06.09688 066806600966..... 2021.
Albert Q Jiang, Alexandre Sablayrolles. Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.066 2023.
Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A. Smith. Finetuning pretrained transformers into RNNs. In Proceed- ings of the 202/ Conference on Empirical Methods in Natural Language Processing, pp. 10630- 10643, Online and Punta Cana, Dominican Republic, November 2021, Association for Computa- tional Linguisties. doi: 10. 10/66/////222..oe// wwwnlp-main..... main.800.... URL org/2021, emnlp maapppmannn..mm 830.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference ma- chine learning, pp. 5156-5165. PMLR, 2020.
Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. In International Conference on Algorithmic Learning Theory, pp. 597-619. PMLR, 2023.
Woosuk Kwon, Zhuohan Li. Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Prin- ciples, pp. 61 611-626, 2023.
Yuanzhi Li, Sébastien Bubeck, Ronen Eldan, Allie Del Giorno, Suriya Gunasekar, and Yin Tat Lee. Textbooks are all you need ii: phi-1.5 technical report. arXiv preprint #xXv:2309.004633 2023.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 17 2017.
Huanru Henry Mao. Fine-tuning pre-trained transformers into decaying fast weights. In Proceed- ings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 10236- 10242, Abu Dhabi, United Arab Emirates, December 2022, Association for Computational Lin- guistics. doi: 2022. www.mmmwna......mmm 697.
URL https: httt::::///claathhoooooooo oorg/
Jean Mercat, Igor Vasiljevic, Sedrick Keh, Kushal Arora, Achal Dave, Adrien Gaidon, and Thomas Kollar. Linearizing large language models. arXiv preprint arXiv:2405. 066440, 2024.
Stephen Merity, Caiming Xiong. James Bradbury, and Richard Socher. Pointer sentinel mix- ture models. In International Conference Learning Representations, 2017. URL https: / //openrevies ew. net./ for forum? d=By 72ddxe.
Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. arXiv preprint ar0iv:2305. /6300, 2023.
13

---

Published as a conference paper at ICLR 2025
Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with infini-attention. April 2024.
Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay. The refinedweb dataset for falcon Ilm: outperforming curated corpora with web data, and web data only. arXiv preprint 2rX1v:2306.01116. 2023.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, et al. Rwkv: Reinventing mns for the transformer era. arXiv preprint arXiv:2305.13048, 2023a.
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemystaw Kazienko, et al. Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence. arXiv preprint arXiv:2404 058922 2024.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.0000 00000,0000..... 2023b.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. arXiv preprini urXiv:2302. 2086, 2023a.
Michael Poli, Jue Wang, Stefano Massaroli, Jeffrey Quesnelle, Ryan Carlow, Eric Nguyen, and Armin Thomas. StripedHyena: Moving Beyond Transformers with Hybrid Signal Processing Models, 12 2023b. URL https:/ //tpt////gthbbb///// cwwc woom//togetheccomptte..ooo
stripedhyena.
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 7025-7041, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10.18663/w1/2222.mmm)) emalp-main. 473. URL https:tpt... /acl anthology, org/ / 2 222..... emm ma in.iii... 473.
Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Bao- hong Lv, Fei Yuan, Xiao Luo, et al. Scaling transnormer to 175 billion parameters. arXiv preprint arXiv:2307.14999 2023.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019. URL https://apt
wwwp::tsssccooaa.ooooooooca (tt.oooppuii::1025332333113 160025553.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140))1-67, 2020.
Sebastian Raschka. Finetuning llms with lora and qlora: Insights from hundreds of experiments, catemmuntttoorriiinniiii
2023. URL https::///// //igh/tiigggigg ai ai/ /pwpp///////aaaeaaaaaeeeeess
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers secretly fast weight programmers. In International Conference on Machine Learning, pp. 9355-9366. PMLR, 2021.
Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. SCROLLS: Standardized CompaRison over long language sequences. Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 12007-12021 Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. URL https: http//acaaaanttoooogoo wttp:/lloooorrgg222222 org//022. emr.ppme mino..iin... 823.
Xuyang Shen, Dong Li, Ruitao Leng, Zhen Qin, Weigao Sun, and Yiran Zhong. Scaling laws for linear complexity language models. arXiv preprint arXiv:2406. (60.20006.660, 2024.
Benjamin Spector, Aaryan Singhal, Simran Arora, and Christopher R. Gpus go brTT, May 2024. 000-0055-12hh
URL https://///........ hatpess//haaaeeeee //haayyesearrhhh wtw.arrccssaaaa.deeeee..... hwt.ordddddubbbd//eoee/e//e/ 200211/g222222222
14

---

Published as a conference paper at ICLR 2025
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuging Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor transformer for large language models. arXiv preprint rrXiv:2307.08662 2023.
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. CommonsenseQQA: A ques- tion answering challenge targeting commonsense knowledge. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 4149-4158, Minneapolis, Minnesota, June 2019. Association for Com- putational Linguistics. doi: 10.18653/1//11--4222 URL https: http://ccaaatooooooooooooo / N19- 9- 1421.
Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-following llama model. https:/ /ttp//gttubb..t// coomtatsuse sat btab/sstan stanford_aapaca 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint 17X2::2302.3377,,022
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov. Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.09288. 2023b.
Szymon Tworkowski, Konrad Staniszewski, Mikolaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Mitos. Focused transformer: Contrastive training for context scaling. Advances in Neural Information Processing Systems, 36, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, al. An empirical study of mamba- based language models. arXiv preprini 2xXi:22406.07887, 2024.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 004611 2018
Junxiong Wang, Daniele Paliotta, Avner May, Alexander M Rush, and Tri Dao. The mamba in the llama: Distilling and accelerating hybrid models. arXiv preprint 2rXiv:2408.1522 2024.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp. 3-19, 2018
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nystrônformer: nyström-bar based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 14138-141483 2021.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:23/2.06635, 2023.
Songlin Yang, Bailin Wang. Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- ers with the delta rule over sequence length. arXiv preprint arXiv: 244444444000444444,4444444444444 2024.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283-17297, 2020.
15

---

Published as a conference paper at ICLR 2025
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Re.
The hedgehog & the por- cupine: Expressive linear attentions with softmax mimiery. The Tweifth International Confer- ence Learning Representations, 2024. URL https:// http:s/://ppnneeee..ettttttt uor
#gC212N2N4
Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. Advances in neural information processing systems, 34:17723-17736, 2021.
16

---

Published as a conference paper at ICLR 2025
A EXPERIMENTAL DETAILS
A.1 MAIN RESULTS, LINEARIZING 7B AND 8B LLMS
Setup. We describe our setup for linearizing Mistral 7B (v0.1) (Jiang et al., 2023), Llama 8B (AI@ Meta, 2024a), and Llama 3.1 8B (Dubey et al., 2024).
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 (Zhang et al., 2024).
For the sliding window implementation, we considered two options: a standard sliding win- dow where is the same for all tokens, and a "terraced" window where w changes based on token index (Figure 9). While found both comparable in quality (Table 18), the lat- ter lets exploit the new ThunderKittens (TK) DSL's (Spector et al., 2024) primitives for im- plementing fast CUDA kernels. Here we pre- fer contiguous blocks of size w 64, which can quickly be computed in parallel on modern
Libear
Softmax
Liner
Softmax
Sliding window
(b) TK "terrace"
Figure 9: We apply softmax attention locally and attend to all past tokens with linear attention.
GPUs. We use this "terrace" implementation in our main results, and include further implemen- tation details in Appendix C.2.
For linearizing data, we use the Alpaca lincarizing data setup in Section 3.2 unless otherwise noted. We also tried a more typical pretraining corpus (a subset" of RedPajama (Computer, 2023)), but found comparable performance when controlling for number of token updates (Appendix B,5. To linearize, 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 le-4. By default, use LoRA rank 8, and scale LoRA updates by (a 16 in HuggingFace PEFT*), amounting to training <0.09% of all model parameters. For both stages, we train with early stopping, Adam W optimizer (Loshchilov Hutter, 2017), and packing into 1024-token sequences with batch size 8. We evaluate the best checkpoints based on validation set perplexity.
Hy plpeepaaaaetter. We list all model and training hyperparameters in Table 7. For learning rates, we did an initial sweep over {le-2, 1e-3, le-4), choosing the best based on final validation set perplexity during step 2: low-rank adjusting, and checkpointing with early stopping. We did not tune batch size or choice of optimizer, and used default values informed by prior work for other design parameters such sliding window size (Arora et al., 2024), LoRA rank, and LoRA projection layers (Hu et al., 2021). In Appendix B, we study the effect of sweeping various values such as window sizes, ranks, and LoRA modules as ablations.
Compute Resources. For each linearizing run we use one NVIDIA 40GB A100 GPU. With batch size I and gradient accumulation over 8 batches, attention transfer takes Rs 2 hours and post-swap finetuning takes RV 4.5 hours, iee,, 6.5 total GPU hours to linearize an 8B L.I.M
A.2 LINEARIZING LLAMA 3.1 70B
We provide experimental details corresponding the 70B parameter results reported in Table 6.
Setup. We compare the quality of two lincarization approaches to the quality of the original Llama 3.1 70B model, including (1) the baseline linearization without attention transfer, which is represen- tative of the approach used in prior work (Mercat et al., 2024; Yang et al., 2024; Wang et al., 2024) and (2) our approach, LOLCATS. For both the baseline and LOLCATS, we start with Llama 3.1
httpps: https::///hgggg/a//ceee http://gga.c.cccaaaaattttteeeeoe dat.ddattttttgeeeeeeee /ttt://ttsttt//gggeeeroprpppepee
RedPa jama Daaaa-Dataallll 11'--111-5 11-Smpp 'ttpps: hottps://hugee http:ss:////ggggaaaaaa http:acccc..occc/eeeee/e/e///e/e Index
17

---

Published as a conference paper at ICLR 2025
Hedgehog
LOLCATS
Model
Precision
16-bit (bfloat16)
Sequence length Linearizing attention
1024
(Linear)
(Linear Sliding Window) Hedgehog
Linear attn feature map Linear attn feature dimension Linear attn feature activation Sliding window implementation Sliding window attn size
64 (effectively 128, see Table
Softmax (across feature dim)
N/A
Terrace
N/A
64
Optimizer and LR Schedule Optimizer Global batch size Gradient accumulation
AdamW
8
8
Gradient clipping threshold Learning rate schedule
1.0
Reduce LR on Plateau
Step 1: Attention Transfer Number of epochs Tokens per epoch Learning rate
2
10M
0.01
Step 2: Low-rank Adjusting Number of epochs Tokens per epoch Learning rate
2 20M
le-4
LoRA rank and alpha
=8, a =16
LoRA dropout
0.0
LoRA projections
W WR, W., W.
Table 7: Hyperparameters for Mistral 7B, Llama 8B, and Llama 3.1 8B experiments.
70B and replace the softmax attentions with the linear attention architecture defined in Section 3.3.1, Equation (6). The training procedure involves:
Baseline: We introduce LoRA parameters to the attention Wg: WR. W., W. projection matrices. We train the linear attention feature maps, learnable mixing term 7, and the LoRA parameters during the fine-tuning adjustment stage of the linearization process.
LOLCATS: We first perform layer-wise attention transfer following Equation (7), with k 80 (i.e., we optimize over all layers together). We then introduce LoRA parameters to the attention W Wi. W., u, W, projection matrices. During fine-tuning we only train the LoRA parameters, freezing the linear attention map and weights.
We MSE loss for the layer-wiss 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 LOLCATS approaches in Table 8, following the same sweep as the 8B models. Since the baseline does not use attention transfer, we mark these values with "N/A". We linearize each model using the same randomly sam- pled 20M tokens of the RedPajama pre-training corpus (Computer, 2023). We pack the sequences fill full context length, and evaluate the best checkpoints based on validation set perplexity.
Compute Resources. We linearize using single NVIDIA 8 x 80GB H100 node. Attention trans- fer takes 4 hours and fine-tume munint takes 14 hours. We use Py Torch FSDP with activation checkpoint- ing for distributed training.
18

---

Published as a conference paper at ICLR 2025
Baseline
LOLCATS
Model
Precision
16-bit (bfloat16)
Sequence length
1024
Linearizing attention
Linear Sliding Window
Linear altn feature map Linear attn feature dimension Linear attn feature activation Sliding window implementation Sliding window attn size
Hedgehog
64 (effectively 128, see Table Softmax (across feature dim)
Terrace
64
Optimizer and LR Schedule Optimizer Global batch size
Adam 8
Gradient accumulation
8
Gradient clipping threshold
1.0
Learning rate schedule
Reduce LR on Plateau
Stage 1: Attention Transfer
Number epochs
N/A
Tokens per epoch Learning rate
N/A
20M
N/A
0.01
Stage 2: Low-rank Adjusting Number of epochs Tokens per epoch Learning rate LoRA rank and alpha
1
20M
20M
le44
8, =66 0.0
LoRA dropout
LoRA projections
W W.. W., W.
Table 8: Hyperparameters for Llama 3.1 70B experiments.
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 repre- sentative of the approach used in prior work (Mercat et al., 2024) and (2) our approach, LoLCATS. For both the baseline and LOLCATS, we start with Llama 3.1 405B and replace softmax at- tentions with the linear attention architecture defined in Section 3.3.1, Equation (6). The training procedure involves:
Baseline: We introduce LoRA parameters to the attention Wg, Wt, W., W. projection matrices. We train the linear attention feature maps, learnable mixing term 7, and the LoRA parameters during the fine-tuning adjustment stage of the linearization process.
LOLCATS: We first perform block- wise attention transfer following Equation (7), with as the block size. To perform attention transfer for block i, we save the hidden states outputted by block 1 to disk and then use this training data for block We then introduce LoRA parameters to the attention W q: WR. W. o, W, projection matrices. During fine-tuning we only train the LoRA parameters, freezing the linear attention map and y weights.
For the reported checkpoints, we train the layer-wise attention transfer stage using use a weighted combination of the MSE loss attention outputs plus cross-entropy loss between the softmax and linear attention maps. We use cross-entropy loss for Stage Low-rank Linearizing.
Hyperparameters. We include the hyperparameters. for the baseline and LOLCATs approaches Table 9, following the same sweep as the 8B models. Since the baseline does not use attention transfer, mark it with "N/A". We linearize each model using the same randomly sampled 20M
19

---

Published as a conference paper at ICLR 2025
tokens of the RedPajama pre-training corpus (Computer, 2023). We pack the sequences to fill the full context length, and evaluate the best checkpoints based on validation set perplexity.
Baseline
LoLCATS
Model
Precision
16-bit (FP16)
Sequence length Linearizing attention Linear attn feature map Linear attn feature dimension Linear attn feature activation Sliding window implementation Sliding window attn size
1024 Linear Sliding Window
Hedgehog
64 (effectively 128, see Table Softmax (across feature dim)
Terrace
64
Optimizer and LR Schedule Optimizer Global batch size
AdamW 8
Gradient accumulation
8
Gradient clipping threshold Learning rate schedule
1.0
Reduce LR on Plateau
Stage 1: Attention Transfer
Number of epochs
N/A
Tokens per epoch Learning rate
N/A
20M
N/A
0.01
MSE and X-ent weights
N/A
1000,
Stage 2: Low-rank Adjusting Number of epochs
1 le44
Tokens per epoch
20M
20M
Learning rate LoRA rank and alpha
T = 4, a ==
LoRA dropout
0.5
LoRA projections
W W. W., W.
Table 9: Hyperparameters for Ilama 3.1 405B experiments.
Compute Resources. We linearize the baseline using three NVIDIA x S0GB H100 nodes, eval- uating the best validation checkpoint after 19.5 hours. For LOLCATS. we perform attention transfer with 80GB H100 GPU for hours per block, and we finetune with NVIDIA 8 80GB H100 nodes for 16 hours. We use PyTorch FSDP with activation checkpointing for distributed training.
B ADDITIONAL EXPERIMENTS AND RESULTS
To better understand LOLCATS' properties and performance, we now report extended results on LoLCATS. We first report expanded results of our main paper, including task-specific ablation numbers, multiple seeds, and additional comparison at the 1B L.M scale. We then extend our ablations by studying how different amounts of parameter updates, different data sources, and dif- ferent amounts training data affect LOLCATS quality for various tasks such as zero-shot LM Eval, 5-shot MMLU, and passkey retrieval. We finally expand the layer-wise training dynam- ics LoLCATs: how layer-wise MSE between learned and softmax attentions changes during LoRA adjusting (App. B.7.1), how LoRA weights update over time depending attention transfer and softmax attention-main matching quality (App. B.7.2), and how block-wirs attention transfer both is motivated by layer-wise MSE and improves linearized L.M quality for larger LLMMs (App. B.7.3).
20

---

Published as a conference paper at ICLR 2025
B.1 EXPANDED RESULTS
B.1.1 TASK-SPECIFIC RESULTS FOR LOLCATS COMPONENT ABLATIONS
We report the task-specific LM Eval results when ablating the attention transfer and linear + sliding window attention in LoLCATS, expanding on Table Table 10. Across all but one task (ARC- easy), we validate that the LOLCATS proposed combination leads to best performance.
+ Attention Transfer
+Sliding Window
ARC-c (acc. norm)
HellaSwag (acc, norm)
MMLUU (5-shot)
Feature Map
PIQA
ARC-e
Winogrande
Average
67.8
58.1
28.9
50.8
23.8
76.5
72.6
53.3
23.8
Hodgchog
x
/
80.5
80.6
53.4
78.8
73.6
45.8
68.8
80.9
81.7
54,9
79.7
74.1
52.8
70.7
x
x
62.0
42.1
24.7
32.7
48.3
23.2
38.8
x
76.1
72.8
40.8
63,6
52.6
23.1
54.8
T2R
54.8
26.3
26.2
56.4
49.6
23.8
39.5
80.7
82.0
54.6
79.5
72.2
40.7
68.3
Table 10: LOLCATS component ablations on individual LM Eval tasks, linearizing Llama 8B, expanded view of Table 5. LoLCATS default shaded. Across Hedgehog and T2R feature maps, LOLCATS' attention transfer and sliding window increasingly improve linearized L.LM quality on popular LM Eval tasks both on average and specifically for 5-shot MMLU scores.
B.1.2 LLAMA 8B LOLCATS RESULTS ACROSS MULTIPLE RUNS
In Table 11, we report the LM Eval results for Llama 8B models linearized with two linear attention approaches, Hedgehog (Zhang et al., 2024) and LOLCATS (ours), after doing linearizing across three seeds (0, 1, 2). We report means and standard deviations (in parentheses). Across all tasks, variation across seeds is low (under point, other than Winogrande), with the difference in task performance across methods being much higher (greater than 10 points).
Llama 8B
PIQA
ARC
ARC-c (norm)
HellaSwag (borm)
Winogrande
(S-sbot)
MMLU
Average
(ao MMLU)
Average
Hedgehog LoLCATs (Our))
(0.67)
40.76 (0.69)
24.22 (0.62)
55.72
(0.35)
80.79 (0.11) 81.62 (0.41)
54.73 (9.41)
79.48 (0.07)
72.92 (1.02)
52.74 (0.64)
70.38 (0.33)
73.91 (0.29)
Table 11: LoLCATS comparison with prior linearizing methods, multiple runs; expanded view of linear attention Llama 8B results in Table 3. We report mean and (standard deviation) across three seeds, comparing LOLCATS and prior linear attention Hedgehog linearizing method.
B.2 CLOSING PERFORMANCE GAPS WITH HYBRIDIZING LOLCATS ACROSS LAYERS
We now study how LOLCATs models perform while keeping some layers full softmax attention. For Llama 3 8B, we keep the first 16 layers (half of all layers) as full softmax attention and only linearize the last half with the sliding window linear attention layer for Llama 3 8B. We then just do attention transfer over Alpaca samples, instead of the two-stage attention transfer and subsequent LoRA finetuning as default in Table (keeping Stage I details the same). In Table 12, find that keeping just 50% softmax attention, while only doing Stage transfer on 20 million tokens, substantially closes the MMLU performance gap by 94.2% (65.8% hybrid versus 52.8% default LoLCATS, compared to Transformer 66.6% accuracy). Furthermore, while just training on 20 million tokens (half that of the default protocol), we recover performance on all other tasks.
B.3 LOLCATS EVALUATION FOR LINEARIZING 1B+ LLMS
We now study LOLCATS' performance when linearizing IB parameter LLMs. We choose two popular models, Llama 3.2 IB Meta, 2024b) and Phi 1.5 1.3B (Li et al., 2023), and linearize with LoLCATS using the same hyperparameters, architectures, and training details (e.g.. 40M to- kens over Alpaca) as Llama 3 8B (Table 7) except for linear attention feature dimension, instead adjusting this to the IB model head dimensions (Lee, 0.5 head dimension 64 = 32). To evaluate
21

---

Published as a conference paper at ICLR 2025
Softmax Attn. Layers
Tokens (B)
PIQA
ARC-E
ARC-C
HellaSwag
WinoGrando
MMLU
All (Llama 88 baseline)
79.9
80.1
53.3
79.1
73.1
66.6
0-15 (LoLCATS 50%, Just Altn Transfer) None (ATs, Attn Transfer LoRA)
0.02
79.5
80.2
53.4
79.2
73.6
65.8
0.04
80.9
81.7
54.9
79.0
74.1
52.8
Hybrid LoLCATS, Llama 3 8B, LM Eval. With just 50% softmax attention, training on 20 million tokens (half that of the default protocol) results substantially closing the gap to full softmax attention Transformers on MMLU (+13 points absolute), and maintaining performance on other LM Eval tasks.
Table 12:
Model
PIQA
ARC-e
ARC-c (acc. norm)
HellaSwag (acc. norm)
Winogrande
MMLU (5-sbot)
Avg.
Avg. (No MMLU)
Avg. of Transformer
Llama 3.2 1B
74.4
65.5
35.8
63.7
60.5
31,9
55.3
60.0
T2R
29.9
84.7
Hedgehog
70.1
55,8
29.8
47.7
50.7
23.0
46.2
50.8
84.7
LOLCATS (Ours)
74.6
63.0
35.1
63.7
61.5
27.3
54.2
59.6
99.3
Table 13: LOLCATS comparison with lincarizing methods, Llama 3.2 1B, LM Eval. Following the two-step linearizing procedure LoI.CATs, we compare LoLCATS architecture against prior Transformer-toRRNN (T2R) and Hedgehog linear attentions. At the IB scale, LOLCATS substantially closes the performance gap with original Llama 3.2 IB Transformer.
LOLCATS quality, we compare the LM Eval Harness performance of LOLCAT' with other avail- able linearizing methods on each Transformer LLM (Table 13, Table 14), and against competitive IB subquadratic LLMs pretrained from scratch (Table 15), reporting results from Bick et al. (2024).
In all evaluation settings, LOLCAT' demonstrates competitive ity. When first controlling for the original pretrained Transformer, LOLCATS is able to outperform both prior linear attentions, well as pure and hybrid Phi-Mamba models created with MOHAWK, another method to distill Transformers into SSM-based Mamba architectures (Bick et al., 2024). No- tably, LOLCATS makes this possible while only using 1.33% of the training tokens in MOHAWK (Table 14), and again only using cerrameetr--ffccien updates. Furthermore, by linearizing Llama 3.2 IB and Phi 1.5 1.3B, LOLCATS efficiently creates subquadratic IB+ LLMs that outperform various LLMs pretrained from scratch (Table 15). These results suggest LOLCATS remains competitive option for linearizing LLMs at multiple parameter counts.
mate-offtheeaat linearizing qual-
B.4 STUDY ON PARAMETER SHAAMETER-EPPCIENTE TRAINING
In our main results, we found that simple default initializations (e.g., rank 8, applied to all atten- tion 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.4.1 EFFECT OF LORA RANK
We study the effect of LoRA rank in post-swap finetuning for zero-shot linearized LLM perfor- mance. Following standard implementations (Hu et al., 2021), we consider two factors. First: rank which determines the rank of the low-rank matrices A, B we decompose the weight deltas into. Second: alpha a, where a/r is scaling factor that controls the degree to which B affect the output (i.e,, LoRA output y Wrr BBA.).
Setup. We sweep over ranks {4. 16. 32, 64, 128, 256}. and adjust such that a/r default scaling factor (Raschka, 2023). For comparison, we also do full finetuning run, where after attention transfer we do the stage adjustment but training all Llama 8B parameters. For all runs, we linearize Llama 3 8B. For LoRA runs, we use the same default experimental setup our main results, while for the full finetuning we decreased learning rate to le-5 due to training instability, 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 all attention weight projections (freezing all other parameters such those in MLPs or GroupNorms), finetuning for two more epochs over Alpaca using the hyperparameters in Table 7. We evaluate with LM Eval tasks.
22

---

Published as a conference paper at ICLR 2025
Model
Traizing Tokens (8)
PIQA
ARC=
ARC- (ace, nom)
HellaSwag (soc. nomi)
Winogrande
MMLU (5-ahot)
Aeg. (No MMLU)
Avg. % of Transformet
Phi 1.5 1.38 (Our nun)
150
76.6
76.1
47.6
62.6
72.8
43,6
67.1
Phi-Mamha 1.5
3
75.5
14.0
44.1
60.2
71.7
65.1
96.8
Phi-Mamba 1.5
60.6
Phi 1.3B T2R
0.04
71.0
89.1
36.6
46.2
53.6
21.3
55.3
82.4
Phi 1.5 1.38 Holgehig
0.04
72.7
70.9
38.0
49.4
54.1
23.5
$7.0
85.0
Phi 1.3B LoLCATs (Ours)
0.04
76.9
77.0
46.9
62.3
72.7
39.2
67.2
100.1
Table 14: LOLCATS comparison with linearizing methods, 1.5 1.3B, LM Eval. We compare LoL.- CAT' with linearizing with prior lincar attentions, and available results from Bick et al. (2024), who distill Phi 1.5B into Mamba and hybrid Mamba-Transfore Transformer architectures with their MOHAWK method (Phi-Mamba 1.5, Hybrid Phi-Mamba 1.5). LoLCATs similarly outperforms prior linearizing methods, closing the gap the Transformer Phi 1.5 1.3B PiQA, ARC-Easy and average zero-shot LM Eval (no MMLU).
Model
Training Tokens (B)
PIQA
ARC-e
ARC-c (acc. norm)
HellaSwag (acc. norm)
Winogrande
MMLU (5-shot)
Avg. (No MMLU)
Transformer
Pythia 144
300
71.1
60.6
26.0
52.1
57.3
26.6
53.4
Llama 3.2
9000
74.4
65.5
35,8
63.7
60.5
31.9
60.0
Phi 1.5 IB
150
76.6
76.1
47.6
62.6
72.8
43.6
67.1
Suhquadratic
xLSTM 1.4B
300
74.6
64.3
32.6
60.9
60.6
58.6
Finch 1.6B (RWKV-v6)
1100
72.6
64.2
34.1
$1.3
59.4
57.5
DeltaNet 1.3B
100
71.2
57.2
28.3
50.2
$3.6
52.1
GLA 1.3B
100
71.8
57.2
26.6
49.8
$3.9
51.9
Mamba 1.4B
315
74.2.
65.5
32,8
59.1
61.5
58.6
Mamba 1,3B
315
73.2
64.3
33.3
59.9
60.9
58.3
Llama 3.2 LoLCATs (Ours)
74.6
63.0
35.1
63.7
61.5
27.3
59.6
Phi 1,5 1,3B LoLCAT's (Ours)
0.04
76,9
77.0
46,9
62.3
72.7
39.2
67.2
Table 15: LOLCATS comparison to pretrained 1B LLMs. LOLCATS- -.ncAA--inncrriee Llama 3.2 1B and Phi 1.5 1.3B consistently outperform strong subquadratic IB+ LLMs pretrained from scratch, achieving best or second-best accuracy all tasks other than ARC-easy, Subquadratic results reported from Bick et al. (2024).
Results. We report results in Table 16. We find that when applying LoRA adjust in linearizing after attention transfer, larger rank updates, e.g.. r r= 128 or 256, do not necessarily lead to improved zero few-shot downstream performance. Full finetuning also does improve performance. Somewhat surprisingly, using just = leads to overall best performance, while r = 128 and r = achieve best and second-best 5-shot MMLU accuracy. While we leave further exploration for future work, we hypothesize that low-rank updates may improve quality by preventing large and potentially harmful updates with the linearizing data to pretrained weights. During the second step of adjusting, if the linearizing data is not particularly diverse or large (e.g., for efficient linearizing), LoRA can then reduce the risk of overfitting to the linearizing data and losing LLM pretrained generalization.
B.4.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 Appendix B.4.1, but fixing r = 8, = 16, we now apply LoRA to different combinations of W,, WR, W., W. weights after swapping in intention-transferee linear attentions. We use the same combination for each layer.
We report results in Table 17. Interestingly, when isolating for projections updated, LoRA on projec- tions not involved in computing layer-wise attention weights (value W, and output W, projections) improves quality compared to query W, or key projections W.. Somewhat surprisingly, updating just W. W. achieves comparable performance
updating all projections (c.f., average zeros- shot accuracy of 74.19% when updating just (W.). ). 74.09% for {W.} versus 74.24% updating (W, We, W., W.)) Meanwhile, updating just W, or W. performs significantly worse (72.68%, 72.29%; versus 74.24%). This suggests much of the quality recovery in Stage-2 low-rank fine- tuning 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 combination of adapting either W,, or W; with value and output projections, we may be able to achieve even more parameter-efficient linearizing with comparable quality-by focusing subsets with the latter.
23

---

W Published as a conference paper at ICLR 2025
LoRA Rank
ARC-c (norm)
HellaSwag (norm)
Wino- grande
MMLU (5-shot)
Avg.
PIQA
ARC-c
Avg.
(no MMLU)
4
80.7
82.5
56.0
79.6
73.6
50.1
71.3
74.5
80.9
81.7
54.9
79.7
74.1
52.8
70.7
74.2
16
80.7
81.9
54.5
79.7
73.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
80.9
81.9
54.5
79.3
72.1
51.7
71.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.
69.2
72.7
Full Finetune
80.6
81.9
54.3
79.4
72.4
52.1
70.1
73.7
Table 16: LoRA rank comparison. Evaluation on LM Evaluation Harness tasks. When adapting all pro- jections Wa. Wx, W., W. after attention transfer, we find smaller ranks r 4,8 are surprisingly sufficient. Larger ranks with more parameter-heary updates do not necessarily improve downstream performance.
LoRA Projection
PiQA
ARC-e
ARC-c (norm)
HellaSwag (norm)
Winogrande
Average
79.49
81.06
51.45
79.18
72.22
72.68
W
79.82
79.80
50.68
78.54
72.61
72.29
W.
80.69
82.49
56.66
79.28
71.82
74.19
W.
80.09
81.65
55.63
79.35
73.72
74.09
W. WA
79.76
81.19
51.02
79.29
72.22
72.70
W W.
80.96
81.90
54.35
79.01
72.30
73.70
W W.
81.28
81.86
54.78
79.20
73.72
74.17
W.. W.
80.74
82,62
56.83
79.26
71.98
74.28
WR, W.
80.69
82.15
55.46
79.53
73.01
69.33
W. W.
80.69
82.79
55.89
79.57
71.59
74.10
W WK. W
80.90
82.20
56.48
79.13
73.95
74.53
W WE. W.
80.41
80.89
54.18
79.18
74.35
73.80
W.. W. W.
80.36
82.32
54.61
79.13
73.56
74.00
W Wt. W W.
80.85
81.73
54.86
79.65
74.11
74.24
Table 17: LoRA projection comparison. Evaluation on zero-shot LM Evaluation Harness tasks. We apply LoRA with the same rank = 8 to different combinations of the attention projections, shading scores by increasing quality (darker is better). When isolating for projections updated. LoRA projections involved layer-wise attention weight computations (value W. and output W. projections) achieves higher quality over query W, or key projections W,. This suggests W., W. may more important adapt after attention transfer, although hest results involve a combination across attention weight and output projections.
B.4.3 EFFECT OF WINDOW SIZE
We now compare model performance using different window sizes in the LOLCATs linear + sliding window attention layer. With the standard sliding window implementation (Fig 9), compare LM Eval performance after linearizing with window sizes w E {4, 16, 64, 256]. We also compare against the
"terraced" implementation used in main experiments with window size 64. In Table 18, we find that in each of these settings, having more softmax attention generally improves performance (c.f., w 16,64 versus =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).
B.5 STUDY ON LINEARIZING DATA
While most our work focuses on architecture and training procedure for improving LLM lineariz- ing quality, now study how data selection affects LoI.CATS performance.
24

---

Published as a conference paper at ICLR 2025
Window Size
ARC-c (norm)
HellaSwag (norm)
MMLU
Avg.
PiQA
ARC-c
Winogrande
Avg.
(5-shot)
(no MMLU)
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
Table 18: Window size comparison. We ablate the window size in LOLCATS linear + sliding window atten- tion. For cach window size the layer applies softmax attention the w-most recent positions, combined with Hedgehog linear attention applied for all prior positions (Eq. 6, Figure 9). We compare {4, 16, 64, 256} using the standard sliding window implementation with our default w = terraced window setup motivated by ThunderKittens (TK 64). Window size 64 performs best, where both implementations perform comparably.
Lincarizing Dataset
ARC-c norm)
HellaSwag (ace. norm)
Wino- grande
MMLU
Avg. (no MMLU)
Model
PIQA
ARC-e
(5-shot)
Avg.
Mistral 7B (v0.1)
82.1
80.9
$3.8
81.0
74.0
62.4
72.4
74.4
LoLCCAT
Alpaca Clean
81.5
81.7
54,9
80.7
74.0
51.4
70.7
74.5
LoL.CAS
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
Lol CAT's
Alpaca Clean
80.9
81.7
54.9
79.7
74.1
52.8
70.7
74,2
LoLCATS
RedPujama 78.9
79.0
52.0
78.1
72.6
55.2
69.3
72.1
Table 19: Linearizing data comparison. For linearizing Mistral 7B (v0.1) and Llama 3 8B, LoLCATs with Alpaca and RedPajama subsets perform comparably (c.f. prior methods, Table 3), though we find that Alpaca actually leads to higher accuracy for most tasks other than 5-shot MMLU,
B.5.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 50K samples of cleaned Alpaca dataset (Taori al., 2023) could lead to surprisingly high performance on popular zero-shot LM Eval tasks, prior linearizing works (Mercat et al., 2024) use more typical pretraining datasets to linearize such as RefinedWeb (Penedo et al., 2023). We thus also try linearizing with a random subset of RedPajama (Computer, 2023) to evaluate how LOLCATS works with pretraining data, albeit without any special curation. For both setups, 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 the setup as described in Appendix A. for all other hyperparameters.
In Table 19, we find that across Mistral 7B (v0.1) and Llama 8B, using the Alpaca cleaned dataset actually leads better downstream task quality for all tasks except for 5-shot MMLU, where lin- earizing with RedPajama consistently leads to ~ percentage point improvements. with both of these datasets leads to comparable or higher performance than prior methods trained on 2500x the data (c.f., Table 3; SUPRA trained on 100B tokens gets Avg. accuracy of 64.0%), suggesting that LOLCATS can robustly improve linearizing quality over different data sources.
LOLCATS
B.5.2 MATCHING LINEARIZING DATA TO DOWNSTREAM (RETRIEVAL) TASK
Among LM Eval tasks, note sizable gap between linearized and Transformer-based LLMs on MMLU. We hypothesize one source of this gap is due to MMLU's evaluation setup, which not only tests for knowledge recall in pretrained weights, but also retrieval over the input context (Hendrycks et al., 2020). In the default 5-shot multiple choice setup, models must be able to retrieve and produce the letter associated with the right answer choice context. However, prior works have shown that linear attentions and non-softmax attention models perform worse on retrieval, both explicitly MMLU (Waleife et al., 2024) and in retrieval tasks at large (Waleffe et al., 2024; Shen et al., 2024).
To counteract these effects, we study if linearizing LLMs with data that explicitly reflects the target downstream task can improve performance. We report two such settings next.
cleaned
25

---

Published as a conference paper at ICLR 2025
Improving MMLU. First, to test improving performance on MMLU, we linearize with additional data from CommonsenseQA (Talmor et al., 2019) (CQA), another multiple-choice dataset. We construct linearizing dataset by using the same 5-shot in-context template in MMLU, i.e.,
<bos><questil on 1><answer choi ces 1 1><answer 1>, <quest i on 2><answer <question 5><answer choices 5> 5>_.<e0s>, where the next token to predict --- corresponds to the correct answer among <answer choices 52, using the ~ 10k samples of the CQA training set. For both stages, we linearize Llama 8B over the combined set of 5-shot CQA samples and our default ~50k Alpaca samples, comparing against just Alpaca or CQA samples. We pack all samples to context length 1024.
choices 2><answer 2>, ,
In Table 20, we report the MMLU scores following each linearizing data choice. We find that adding a small amount of multiple-choiee samples can substantially improve MMLJ accuracy (~ 2 points). However, CQA alone performs ~ 10 points worse than Alpaca. This suggests quantity and diversity of samples may still be necessary for linearizing; we break this down further and study the amount of linearizing needed during both attention transfer and low-rank adjusting in Appendix B.6.
Model Linearizing Data
Llama 3 8B LoLCATs
Llama 3 8B N/A
Alpaca
Alpaca + CQA
CQA
MMLU (5-shot) Acc.
52.8
54.5
43.9
66.6
Table 20: MMLU comparison with task-specific lincarizing data. Adding multiple-choice samples (Com- monsenseQA: CQA) to linearizing training data modestly improves downstream MMLU performance. How- ever, CQA alone is insufficient to achieve competitive quality.
Evaluating and improving saalle-inna-haystack retrieval. Next, to further test linearizing per- formance on downstream tasks, we evaluate the L.L.CATS--liearizze Llama 3 8B LLM on "needle- in-a-haystack" tasks. We use the passkey-retrieval task setup introduced Mohtashami & Jaggi (2023) and evaluated in various prior works (Chen et al., 2023a;b; Tworkowski et al., 2024), where a model must retrieve a hidden passkey uniform randomly placed inside a text span. We use ran- domly generated prompts and 5-digit passkeys (Listing 1), and test correctness by whether the model outputs the passkey exactly texact-matcht We evaluate over various text lengths (2048 to 10240 to- kens), and report accuracy binned by which decile the passkey occurs in.
Like in our MMLU study, we evaluate LOLCATS with two linearizing data setups. First, we use the default Alpaca data, but this time test linearizing Llama 3 8B at its max context length by packing the 50K Alpaca samples into 8192-token chunks. Second, to see if linearizing explicitly with retrieval data helps, we generate 10K passkey retrieval samples with 8819 tokens, and use these samples for both stages of LOLCATS linearizing, For both, we use the same hyperparameters as Table 7.
In Table 21, we first report the results for Llama models evaluated on 8192-token prompts. We compare the LOLCATS Llama models against the base Llama 3 8B, the instruction-cunee version (Llama 3 8B Instruct), and Transformer Llama 8B A-nn-finetuedd on Alpaca data (Llama 3 8B (Alpaca); using the same LoRA parameters and second stage training procedure as in Table 7). We find that with LOLCATS models, linearizing data plays a particularly strong role in downstream performance. As a potential drawback of LOLCATS, when just using packed Alpaca samples, the linearized Llama 8B fails to get even a single passkey retrieval sample correct (LOLCATs Llama 8B (Alpaca)). Meanwhile, LoRA foRA-iinetuning the non-linearized Llama 3 8B maintains high passkey retrieval accuracy. However, linearizing with passkey samples (LOLCATS Llama 8B (Passkey)) recovers 100% accuracy. This suggests that with LOLCATS linearizing, the linear + sliding window attention is able to do passkey retrieval with similar performance to full softmax attention models.
Pusskey Placement
0-10%
10-20%
20-30%
30-40%
40--0%
50-60%
60--00%
70-80%
80-9US
90-1019%
T.lama 3 8B
100.00
100.00
100.00
100.00
92.86
100.00
94.44
94. 12.
92.31
100.00
Llama 8B Instruct
100.00
100.00
100.00
100.00
100,000
100.00
100.00
100.00
100.00
100.00
Llama 8B (Alpaca)
100.00
100.00
100.00
100,00
100,00
100.00
100.00
100,00
100.00
100,00
LoLCATss Llama 3 SB (Alpaca)
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
LoLCAT's Llama 8B (Passkey)
100.00
100.00
100.00
100,00
100.00
100.00
100.00
100,00
100.00
100.00
Table 21: 8K Context Passkey Retrieval. We report passkey retrieval accuracy for various Llama 3 SB models, using 8192-token samples and binning by passkey placement decile. Linearizing with just packed Alpaca samples fails, but using passkey samples as the linearizing data recovers 100% retrieval performance.
26

---

Published as a conference paper at ICLR 2025
There is an important piece of infe hidden inside a iet of irrelevant text. Find it and memorize it. will quiz you about the important information there.
The grass is green, The sky is blue. The sun is yellow. Here we go. There and back again. The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again. The grass is green. The sky is blue. The sun is yellow. Here go. There and back again...
<EVEN MORE FILLER>
The pass key is <PASSKEY>, Remember it. <PASSKEY> is the pass key.
10 The grass is green. The sky is bluc. The sun is ycllow. Here we go. There and back again. The grass is green. The sky is bluc. The sun is yellow.
12 Here we go, There and back again...
13 <MORE FILLER>
14 What is the pass key? The pass key is
Listing 1: Passkey Retrieval Prompt Template
Passkey Retrieval Lamm
Pasakey Retrivval 1:8 instouct
KCN
LON
Ineatt Contest Leogt
200
KON
204
Pacckey Macemeett
Pyskkey Macomeet
Parsdin Retrinon Liaria
Penckey Betrinual. Uama LOUCA'S (Liner Passkey
Fisetnne Alpaca)
Irpai Cunsect Lar
18
Passary Rlacement
Pazalay Placerent
Figure 10: Needle-in---haystack with passkey retrieval. LoLCATs- Llama 8B with passkey retrieval linearizing data results comparable retrieval to original Llama 8B and instruction tuned variants.
We further test the robustness of this retrieval across various context lengths (2048 to 10240 tokens) in Figure 10. We report accuracies across input context lengths and passkey placements, finding similar strong retrieval performance for samples under Llama 8B's 8192 context length. Inter- estingly, only the Transformer Llama 3 8B LoRA-finetuee Alpaca is able do retrieval over longer context samples (10240 tokens).
B.5.3 SAMPLE LENGTHS: EFFECT OF EFFECTIVE SEQUENCE LENGTHS
We further study the impact of sample sequence length for linearizing quality. By default, for lin- earizing data we pack original data samples into sequences of consistent length, egg., 1024 tokens. As done in prior work (Raffel et al., 2020), 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 sam- ples 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 approximate attentions over "long enough" sequence lengths.
27

---

Published as a conference paper at ICLR 2025
Effective sequence length. To study this data effect, we define "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 we define the ESL per query as
exp(g kj/Va)
ESL(9)) = Σ (1- Em explq, km/ kn/VV)
We compute sample's ESL per head as the sum over all query ESLs, i.e., Ei= ESL(q. for sample with n tokens. We average this over all heads and layers measure sample's overall ESL.
We hypothesize that if our linearizing data only has samples with shorter ESL than those encountered test time, then we would poorly model these test samples. Conversely, we may be able improve linearizing quality by specifically filtering for samples with longer ESL. We report two findings next.
Alpaca
300
(Random) (Sample Top ESL)
RP (Random)
Sample Count
250
Sample Count
250
200
200
1.50
150-
100
200
250
300
35.
400
275
300
325
15.0
375
400
425
Effective Sequence Length
Effective Sequence Length
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 the cost of other LM Eval tasks (Table 22).
Finding 1: ESL corresponds with RedPajama versus Alpaca performance. In Figure 11, we first plot the distribution of sample ESLs computed with Llama 8B on Alpaca and RedPajama linearizing data subsets. We find RedPajama samples on average display longer ESLs, which conin- cides with improved MMLU score (c.f., Table B.5.1).
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 ESL.s, 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 22). However, this comes at cost for all other LM Eval tasks, dropping quality compared to random RedPajama packing by 0.2 to 2.1 points.
Finding 2: Filtering for higher ESL improves MMLU. Furthermore,
B.6 STUDY ON LINEARIZING TOKEN BUDGET
We further study how varying the number of tokens used for both attention transfer and low-rank adaptation impacts LoLCATS linearizing quality.
Impact 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 validation samples both in-distribution (held-out
HellaSwag (uce. norm)
Winogrande (acc. norm)
MMLU (5-shot)
Avg. (no MMLU)
Lineurizing Data
PIQA
ARC-easy
ARC-chullenge
Alpaca
80.9
81.7
54.9
79.7
74.1
52.8
74.2
RedPujama
78.9
79.0
52.0
78.1 1
72.6
55.2
72.1
RedPajama (Sample Top ESI.)
78.4
77.0
49.8
78.0
71.4
56.5
70.9
Table 22: Effect ESL on lincarized LLM quality, Llama 8B. While linearizing with longer ESLsse..g. RedPajama samples or specifically filtering for top ESLs in RedPajama (Sample Top ESL) imppooes 5-shot MMLU accuracy to 3.7 points, reduces quality on all other evaluated LM Eval tasks by 1.7 to 5.0 points.
28

---

Published as a conference paper at ICLR 2025
Alpaca samples) and out-of-distributio (RedPajama validation samples) over different combina- tions of steps (Figure 13). Without attention transfer, low-rank adaptation converges significantly higher on in-distribution samples (Figure 12a), suggesting poorer quality linearizing. However, find similar held-out perplexities after relatively few attention transfer steps (c.f. 1000 - 1800 up- dates, the former amounting to just 8 million tokens for attention transfer), where all runs improve in-distribution PPL by ~.23 points after LoRA finetuning for 2000 steps.
Atm. Transter Step
Atm. Trancher Step
Atm: Thansder Step
273
230
300
133
930
Peppectty
2-245
1993
1200
1993
285
282
500
1i40
2000
2400 AAN)
8AM
103
1993
2000
2040) 1092 8600
922
1300
2207 1997 Traniin Stops
1200
1992
Toal Training Steps
Training Stops
(a) Alpaca PPL
(b) Alpaca Val. PPL (Zoomed)
(c) RedPajama Val. PPL (Zoomed)
Figure 12: Evaluation curves over number of training steps. IJama 3 8B lincarized with Alpaca. We report the impact of steps allotted to attention transfer versus LoRA linearizing, using validation set perplexity (PPL) over both in-distribution (held-out Alpaca samples) and out-of-distibutioon (RedPajama validation samples). We first run linearizing with 0 I 1800 attention transfer steps, before LoRA-finetuning for up to 2000 steps. Without any attention transfer (0 steps), linearized LLMs get much higher perplexity (Fig 12a). On the other hand, we observe similar convergence after attention transfer over only 1000 steps.
In Table 23, 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-0622 higher points than the 50% softmax attention variant, c.f., Table 3). Without any attention transfer, linearized LLMs perform drastically worse out-of-distribution samples (Table 23, RedPajama and LM Eval metrics).
Impact of more pretraining data. We next study how linearizing over larger amounts of pretrain- ing 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 8B with differ- ent 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 13a) and general zero-shot LM Evaluation Harness quality (Table 13b). Increasing both Stage attention transfer steps and Stage low-rank adjusting steps notably improves validation perplexity. How- ever, 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 9x more unique tokens does not seem monotonically improve downstream quality (Table 13b). Meanwhile, we do find that across various amounts of attention transfer steps, subsequent low-rank adaptation consistently improves average zero-shot LM score by > points.
B.7 LAYER-WISE TRAINING DYNAMICS AND ANALYSIS
Finally, we further study LOLCATS layer-wise training dynamics, such as the resulting linear at- tention approximation quality softmax attention at every layer (App. B.7.1), how this corresponds
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
Table 23: Effect of attention transfer steps. With Llama 8B lincarized on Alpaca data, we report the final evaluation perplexities after 2000 LoRA steps in Fig 13, as well downstream LM Eval performance averaged over shoo tasks. We again find competitive quality with relatively few attention transfer steps.
29

---

Published as a conference paper at ICLR 2025
204
Perplesty
3000
LoRA Stepp:
1R Fall
5950
0000
7000
09000
0020
0
2000
4000
6000
8000
10000
Training Steps
Transfer
(a) RedPajama Validation Perplexity
(b) LM Eval Score
Figure 13: Evaluation quality from ablating RedPajama linearizing updates.
0.010
LOLCATS
MSF Crange
0.000
LoLCATS Transfer LORA
Hedgehog Attn. Transfer
0.00150
0.00175
Hedgehog MSE Change
Hedgehog Attn. Transfer LORA
0.001:25
0.001000
0.00075
0.002
8.000
26
Layer
(a) Layer-wise MSE before and after LoRA adjusting
(b) Layer-wise MSE change after LoRA adjusting
Figure 14: Liyer-wise MSE and deltas. Layer-wise MSE and change in MSE (absolute) before and after
LoRA adjusting with the Hedgehog linear attention and LOLCATs linear + sliding window attention on Llama 3 8B. Both use the Hedgchog feature map. With larger initial MSEs after attention transfer (Fig. 14a), LoRA adjusting with Hedgehog results in larger MSE improvements in MSE across layers (Fig. 14b).
with LoRA updates after attention transfer (App. B.7.2), and how these layer-wise attention MSEs (a) relate to model size and (b) further motivate block wise training (App. B.7.3).
B.7.1 LAYER-WISE SOFTMAX ATTENTION RECOVERY WITH LORA ADJUSTING
We study how LoRA explicitly improve softmax attention approximation. Given learned feature maps, can we recover softmax attention better by adjusting the attention projections?
In Figure 14a, we report the layer-wise MSE between learned linear attentions and softmax attention using either the pure linear attention (Hedgehog) or linear sliding window attention (LOLCATs). We plot the mean MSE computed over all samples in our Alpaca validation set, averaging over all heads, token positions, and heads per layer. In addition to plotting this metric after attention transfer (as Figure 7), we also plot the MSE after an additional round of LoRA adjusting (+ LoRA). Like before, we freeze the linear attention weights and add LoRA weights to query, key, value, and output projections. However, rather than train these weights end-to-end for next-token prediction, we update LoRAs to explicitly minimize the MSE between our trainable linear attention and original softmax attention outputs like in Stage 1. We use the same hyperparameters as Table
We report both the absolute MSEs (Figure 14a) and the change in MSE (Figure 14b). LoRA duces MSEs with both linear attentions frequently reducing MSE more when starting with worse approximations after attention consfer-whiee may suggest occurs both across layers for each linear attention, where we see greater MSE deltas for the later layers (c.f., LOLCATS MSE change, Figure 14b), and between linear attentions (c.f., Hedgehog MSE change, Figure 14b), where LoRA generally improves MSE more with the Hedgehog linear attention versus LOLCATS linear + sliding window. Despite these greater improvements, LoRA alone does not close the MSE gap (Figure 14a), suggesting that linearizing architecture still plays an important role in learning match softmax attention.
type of compensatory role. This
B.7.2 LAYER-SPECIFIC LORA TRAINING DYNAMICS
Next, further study how LOLCATS layers behave during LoRA adjusting, and plot the cumula- tive weight updates to LoRA low-rank A and weight matrices while training LLMs end-to-end for next-token prediction. As reference points, compare against LoRA finetuning (1) the original
30

---

Published as a conference paper at ICLR 2025
Transformer LLMs with softmax attention, (2) linearized LLMs without attention transfer (using the LoLCATs linear + sliding window attention, Hedgehog feature map, but not trained to match softmax first), (3) linearized LLMs with attention transfer and the Hedgehog pure linear attention, and (4) linearized LLMs with Hedgehog pure linear attention without attention transfer. We finetune Llama 3 8B with epochs on the Alpaca dataset, following Step 2 hyperparameters in Table 7.
We plot these updates per LoRA weight and projection in Figure 15, with layer-specific plots in Figure 16 and 17. We specifically show cumulative sums of the update magnitudes over 2500 steps.
LoRA A, cuproj
LoRA A, k. k_proj
LoRA V. v_pro
LoRA A, p.pro)
Soft Alterrion
LOLEAK
Hodachog
Hedgahog (No Attin Transfer)
3088
1330 3009
100
010
LoRA B, c_pro
LoRA B, k.k.k.K k_proj
LoRA B, v_pro
LoRA D, b_pppr
Softmac Altertion
LOLEAR
Hodachog
Hedgahog (No Atto. Transfer)
100
306
1330 350 300
100
800
300
Figure 15: LoRA A, B weight updates by attention projection over training, averaged over all layers.
Across plots, we find several interesting findings. First, LOLCATS with attention transfer results in much more similar cumulative updates softmax attention for value and output projection LoRAs, where both result in noticeably larger updates than any other configuration. This may suggest the LOLCATS attention transfer linear attentions are sufficiently similar to softmax attention, such that when LoRA finetuning end-to-end, the weights also behave similarly. Furthermore, for query and key projections- i.e., those involved in the attention weight computation -tee situation is reversed. The untrained or no attention transfer linearized LLMs display greater cumulative updates earlier during training for A weights and consistently for B weights than either trained linear attention or softmax attention. As we train all layers jointly, we find that with the untrained linear attentions, relatively more "weight" is diverted to learning to compute attention weights via query and key projection updates. How this impacts the downstream model-iee, by updating the value and output projections less than a softmax attention Transformer-is an interesting question for future work.
When plotting these updates per layer (Figure 16, 17) find similar dynamics. Interestingly, despite LOLCATS with attention transfer resulting in the worst softmax attention approximations for the layer 31 (c.f., attention MSE, Figure 7, 14a), the resulting LoRA weight deltas do not seem to track the softmax attention LoRA deltas noticeably worse than other layers.
B.7.3 BLOCK- WISE ATTENTION TRANSFER
Finally, we study how attention transfer properties change with model scale, motivating LOLCATS's biock-wise training approach (Section 3.3.2). In particular, we note that attention output MSEs can vary quite bit across layers, and this can be aggravated by model size (c.f., Llama 3.1 70B, Table 25; and Llama 3.1 405B, Table 26).
Recall that attention transfer involves training the LOLCATS 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 MSE already varies across layers after atten- tion transfer for Llama 3 8B (Figure 6b), 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. Corre- spondingly, in Table 24, 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 sufficient and performs similarly block-wise training the smaller scales (8B, 70B).
31

---

Published as a conference paper at ICLR 2025
Layer O LoRA RLAPN
Layer LORA A. SLPMO)
Layer 15 LoRA (LAM
Layer 23 LoRA A. C_SnR
Layer 31 LORA CLPm
Altr. Turafori
Hobgenog
Hedgetog
Layer K,00ol
pro))
K.0r3
Layer k.pro)
Layer Krpr
Tunndor
196 Alto. Turrsoor
Layer 0 LoPA v.orol
Layer LoRA A. pro)
Layer 15 L-RARAAAA v.rrg
Layer 23 LoRA A. y.p
Layer LORA v.
MALEAS; Tursfrr Hobganog
Hodgeiig
Layer
pr))
Layer
Layer c.p.
Layer p.pro.
Unio
LoLCAS
Tranfor
196 Altn. Transferr
Figure 16: LoRA A weight updates by attention projection and layer training.
Uniro 5
LALLANS
Trnnforr
Layer 0 LORA proe
Layer 7 LoRA B. pro)
Layer 15 LORA S,K. oro)
Layer 23 LoNk pro)
Layor 21 LOKA S. pro
He Alts Turnfor Hedgenng
Hodgehin Transferr
Unmo
Laye: 0 L.RAA B. o.
Layer 7 LoBA 8, pro)
Layer LONA B, o..rr
Layor 25 LoBk pno]
Layer 31 LoKA 8, pr
ran Tradden Hwdgehog
Hodgetog frander
Figure 17: LoRA weight updates by attention projection and layer over training.
Tradeoffss 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 (A 126) training, where we load the full model once and compute the loss as the sum of siser-wise MSEs, and (2) block-wise (k = 9) training, where we break the model into blocks and train each independently. Per our discussion in Sec. 3.3.2, this increases storage costs
Method
x M//b
PPL
Joint
126x1
4.23
Block-wise
9x14
3.21
Table 24: Validation per- plexity after LoRA adjusting for Llama 3.1 405B, compar- ing attention transfer with M/ blocks of layers. M 126 number of total layers.
32

---

Published as a conference paper at ICLR 2025
precomputed hidden states. However, we hypothesize that optimizing block attention transfer loss can improve quality.
For quality, compare the evaluation perplexity of both the attention transfer approaches, after LoRA adjusting on the same subset of RedPajama data. The block-wise and joint approaches per- form similarly the 8B and 70B scales, however, shown in Table 24, the block-wiss 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 the 405B scale.
Layer-wise MSE. We report the attention MSEs across model scales, finding that variation in MSE magnitudes corresponds with model size. Tables 25 and 26, we report magnitude of the MSE loss as the depth of the layers in the block increases. MSE variation increases with model size: while the largest MSEs at the 70B scale are between 2.5-3x x the MSE the final block at the 405B scale is 48. 66 (over 13x that of the next largest MSE by block, 3.62).
Block (Layer range)
Eval MSE
0-4
8e 0.03
5-9
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 25: 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 Ilama 3.1 70B model. Each block is trained the exact 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
Table 26: Attention transfer block-wise MSE We report the eval MSE by 9-layers block for the 14 blocks in the 126 Transformer layer Llama 405B model, Each block trained on the exact same RedPajama data at sequence length 1024.
33

---

Published as a conference paper at ICLR 2025
C IMPLEMENTATION DETAILS
C.I PSEUDOCODE IMPLEMENTATION
Below provide further details on implementing LoLCATS 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 linear attention. We define such class below.
import copy
import torch. as from einops import rearrange
class IolcatsilamaAttentionn (nn. (nn.Module)
def i (self,
feature dim: int,
base_attn: nn mnMucelle, original Transformer attn.
) -> None:
10
super () init
12
Inherit pretrained weights
13
self.q_proj pase_attn.q_prro
14
self.k_proj kase_atta.k_prot selo.v_proj base_attn.v_proo
15
self.o_proj base_attn.copro:
18
Inherit other attention things
19
self self.rotary_emt =
self.base_ batttsseetttn attn = base attn # keep for attention transfer
self.num heads base attn.num. heads self.head dim base atse_attn.head_de
22
23
24
Initialize feature maps, see Bedgehog definition below sclf. feature_map_q HedgehoggoatureMap self.num_her self.head_dd feature_dim
27
28
self. Teature_map_ copy. deyp.opepppyy (self. feature_map_q)
def (srward(self, torch. Tenser)
torch.
31
**"
32
Compute linear attention (assume no 6QA)
33
(b: batch_size, h: num_heads, 1: seq len, d: head_dim)
ş self.q (x)
all (b, h,
36
k self.k_prop (x)
37
v = self.v_prof (x)
38
Apply rotary embeddings
40
q sclf.rotary_cmb (q)
41
k self. (k)
42
Apply feature maps
a self self..eature_mma_d (q)
(b, h, feature_dim)
45
k feltrfe map (k) (b, h, 1, feature_dimt
46
47
# Compute linear attention
kv torch.
k, v)
49
Y torch. ('bhif, g, kv)
50
y /- torch.einsen ('bhif,
51
q, cucsum (dim=2) [..., None]
# Apply output projection
return self.o_proy (rearrange (y, "b h
b d)
Listing 2: LoLCATS Linear Attention Class
34

---

Published as a conference paper at ICLR 2025
Linear Sliding Window Attention. We can augment this linear attention with the linear attention and sliding window formulation described Eq. 6. We first define standalone functions for linear attention and sliding window softmax, before defining such as hybrid class below.
def
(q:
forch.Tenee torch. Tensor, v: Lorch,Tensor. window_size: int, hindow_factor: float):
Compute sliding window softmax attention 0 (n) time and space by not materializing o (n (n"2) attention weights
d g.shape[-1]
11
Compute windows for keys and values, shifting by window size window kwargs l'dimension': 2, 'size': window_sine 'step's 1 k = F.pad (k, (0, 0, sindow_size - 1, 0), value=0) (nfold window_kwwrrg kwargs)
12
13
F. p.d (v, (0, 0, window_size 1, 0), value=0) .nnfold(* window_kwargs)
15
16
Compute wi windowed_soffe soteeeaxssoaaxoooeeooeoe (qk); causal in its construction a_sm -
bhldd, bhldw->bhlw's q, k) (d ** -0.5)
heurist for zeroing out padding above sm a sm
19
-torch. finto (q. (qq.dtppe) .max
20
# Compute softmax terms for combining attentions (attn and sum)
a_sm_max torch. amax (a_sm, dim=-1, keepdin-True) a_sm window_factor
23
terch. terch.exp(e (a_sm a_sm_max) sum_sm = a_sm. sum (dim=-1, keepdin=True)
24
25
return torch. einsum (' bhlw, bhldw->bhid', a_sm, v), sum_am
Listing 3: Linear Attention in Linear + Sliding Window Attention
def under_window_lne wttt:llltaaaattttttioon (f_q: torch. Tensor, f_k: torch.Tensor. v: Tensor, window_size: int, linear_factor float):
# phi (q)
# phi (k)
Compute hybrid window attention dot product with linear complexity in q len
=**
10
dtype f_q. dtype
11
12
# Shift keys and valuves for window
13
window size
14
f_k F.pad (0, 0, w, 0), value=0) [:, 4, :-w, :]
15
v F.pad (0, 0, w, 0), value=0) [i, i, :-w, :]
16
Compute linear terms for combining attentions kv forch.einsume (wthhf/.hhldddtteeeeeeee
18
k, v)
19
gkv linear_factor torch.einsum ('bh1f, bhh q, kv)
20
sum_f_l flatt cumsum (ommuu(diim-)) to (dtype-dtype)
sum ck = linear factor terch. einsum ("bhid, bttd--hhbbldd_bh1""
f_ sum fum_f_kl ...., None]
25
return gkv, sum gk
Listing 4: Linear Attention in Linear + Sliding Window Attention
35

---

Published as a conference paper at ICLR 2025
class wtw.esss..aggg...oo.....ottaaaoa (io.eenntiiiiiiitittttttaiaaa
def __i_iint_________ (seif, windox_size int = 64, **wargs: any,):
super () -ii-i_t___ (**kwargs) self. windo_sieds_eee window_siad I sliding window size self. self.window_factors torch. torch.ones (1, self.num heads, 1, 1)
an.Parameter gamma mixing term
def ectentionnself.
10
q: torch. Tenser, torch. Tensor, f_q: Forch.Tenscr. f_k: torch. Tensor, v: torch.Tenser wi windou_factors torch. Tensor, linear_factor: terch. Tensor window_sizer int 64,1:
11
12
15
16
(n) hybrid linear sliding window attention
19
window_ _avdrw_ {' l'miensii''' 2, 'size' window_size, 'step' 11 1. Sliding window (softmax attention) with corch.no_gradiee
20
21
qkv_sm, sum sum_gk_sm sliding http:/_sftaaaaaaaeeeee
wriddingwwindwwwcm
q, k, v, sindow_size Sindow_factor"
24
25
2. Under window (linear attention) elv_in, sum_gk_in under_window Liw.oo___iierrattetttitee
26
f_q, f_k, v, window_size, sinear_facter
29
# 3. Combine
30
y = (gkv sm + qkv_In) (sum qk sm + sum qk In)
return
33
def (srward x: torch. Tensor) torch. Tensor:
34
**"
35
Compute linear attention (assume no GQA)
(b: batch_size, h: num_hoads, 1: seq_len, d: head
**"
38
q self.q_prooj (x) $ assume all are (b, h, 1, d)
39
k self.k_proji (x)
40
self.v_proj (x)
42
Apply rotary empeddings
43
g self.rotary_emb
44
k relf.rotary_embl (k)
Apply feature maps
47
f_q k = self, feltr, # (b, h, 1, _amm
(q)
(b, h, 1, feature_ dim)
48
49
Compute attention
window_factors F.sigmiid wwndow iinear fnneorfacttor # Eq.
52
53
attentttention (q, k, f_q, f_k, Lindo_facctors, linear_factors
half.window_sixe
56
Apply output projection return self.o o proj (rearrange (y, "b h 1 -> b (h d)
57
Listing 5: LoLCATs Linear + Sliding Window Attention Class
36

---

Published as a conference paper at ICLR 2025
Hedgehog Feature Map. We implement the Hedgehog feature map following Zhang al. (2024).
import torch. nn as
class HeagehoggeatureMap (nn.Module))
def init init_
num heads 32: int, # defaults for 82 LIMs
head_dim 128: int,
feature_din 64: int,
-> None:
super () _i._iitt
self. held.nun_heads heads num_heads
helf.head_dee dim head head_dim dim
12
sclf.foature_diit
fcature_dim,
13
14
Initialize trainable feature map weights self. weights = nn. Paraaerete torch, (sellffnummnee num_heads, sel heead_dim, self. feature_dim)
17
18
def self aetieatiiiiionn (self: terch. Tensor) -> Tensor: """Softmax return torch. cI
feature dims activation""
22
torch.softmat (x, dim=-1), corch.softmin (-x, dim=-1)
23
1, dim=-1)
def forward x: torch. Tensor) torch.Tensor:
26
.*.
27
Assume x, shape is (b, h, 1, d) (b: batch_size, hi nea_heads, 1: seq_len, di head_dim) **"
28
corch.einss (1 hdf, bhid->th1f', return self. activation (x)
f. weweights, x)
31
Listing Hedgehog Feature Map
Linearizing LLM Setup. To initialize an LLM for linearizing, we simply replace each softmax attention in Transformer's layers with our LoLCATs linear attention class. We illustrate this with Huggingface Transformer's class below.
from transformers import AutcModee1oorauusal..
def convert_ (model: AutoMooieLForCausa..1., window_size: int 64, feature_dim: int 64,): "**Setup linearizing icnentiong"!" for layer model.model.1ap
if window size 0:
layer.self attn = Iolcatst.lamattenntion feature dim=feature dim, baee_attnllayer.slf...tt
12
13
clsc:
14
layer self_aself_att =
feature_ cattur_ddin-feature__did base aase_attn=layer.selee
17
18
return model
Listing 7: Linearizing LLM Setup
37

---

Published as a conference paper at ICLR 2025
Attention Transfer Training. We then train LoLCATS layers in simple end-to-end loop- Although doing this attention transfer is akin to a layer-by-laye cross-architecture distillation, due to architectural similarities implement linearizing with the same footprint finetuning single model. Furthermore, as we freeze all parameters except for the newly introduced feature map weights, this amounts to parameter-efficien finetuning, training <0.2% of 7B+ LLM's parameters.
Example attention transfer training loop for Llama 3.1 8B
import torch.nn as from transformers import AutoMoodelForCausalll.
Load Llama 3.1 8B
model_ model_contig = 'pretrai ned sed___odel_____ _ 'meta- -1 lama Mat.teettalllaaaaaa 1-88"
10
11 model AutoModelFoorCausall. wttp:usssaaaaal.orrreeereeee from_pretraiined -model smodel_config)
12
13 Freeze all pretrained weights
for p in madam..prameters
15 p.requires_grad False
16
17 Frepare LoLCATS lincarizing layers
model convert_model (model)
20
Setup MSE less criterion
21
mse_ loss - nn. MMMSSLoss ()
22 block_size 32 # default end-to-end for VB+ LIMs
blocks 1en (modell // size
24
25
Get some linearizing data train_loader = dota (+ (+*aaaa
26
Train LeLCATs layers via attenticn transfer
29
for ix, input_ide ids (trainn_ttt
30
losses [0] * range mnmm_blocks, # Attention transfer loss here
31
x model. eeoell.embd__toen (input_ide) # Input embeddings from tokens
33
Forward pass thru model
34
for 1ix, layer in enumerate (model, 1ayers))
35
# *** Start Attention ***
fayer.input_layernoe
Just llaa things
38
# Attention Transfer part
39
with torch.no ()
40
wtw.esssel__ttt.aaeeeeeee Lx) y_pred Aayer.self_attn_3.
y_true
42
_dx lix // block_size
layer or block-wise MSE Icsses idx] += mse_loss pred, y_true)
45
44
x y_true
Pass true attention outputs
thru to rest of model
_
47
End Attention
48
49
$ +** Start MIP ***
poot_atteetioo_layermee _ self.mlp self.mlp(xx)
layer.post
52
x _x
53
# *** End MLP ***
for loss in losses:
i End-to-end attention transfer
lcss. basswaackward()
Listing 8: End-to-end attention transfer pseudocode (Stage 1).
38

---

Published as a conference paper at ICLR 2025
Low-rank Adjusting. Finally, after attention transfer, train the model end-to-end with next- token prediction. This allows the linearized LLM to adjust to the learned linear attentions, which may still not perfect approximations of the softmax attentions. However, with LOLCATS hope to make these errors small enough such that we can adjust and recover pretrained LLM capabilities with parameter-efficient low-rank updates (e.g., LoRA finetuning).
class LoRLayy (torch. www.orrch.....iiiiiii....... Module) def init _eiiiinnneel____________
base_layer: nn.Mcdule, rank: int 8, alpha: float 16); super __1___nntttt
"""Init low-rank parameters""*
in_dim base_layer. wew.yer.weiigt.... shapaerr...ghhhthhpp[[[]]]]
out_dit dim base_layer.weight.oo
10
self.A Pn.Parameter (terch. randn sii_dim, rank)) )
11
self.B = nn.Parameter (ttrch.. (rank, out (iim)
seli.alpha alpha
13
self.base_layer base_layer
14
15
def los_rank_forwardceelll x; terch. Tensor) >> torch. Tensor:
"""Conpute LoRA pass***
torch.einsum(" ... ..., x, self.A)
18
x torch.einsum(" .,, rd-,, rd- x, self.3)
19
return self.alpha
20
def (sewf, torch. Tenser) -> torch.Tensor:
"" """Actual forward***
23
x self.base_layerde + self, lww_rank_forwarrce
24
return
Listing 9: Defining a LoRA layer
model Eettentio__tranfffi (model)
Do Step of LOLCATS
0. Prepare model for LORA 1ora kwargs rark'' 8, 'alpha': 16)
examples
# 1. Freeze all pretrained weights for in model. paramotors (): p.requiros_grad False
Add LoRA weights Q,K, v,0 projections
10 for layer in model,
11
for prej in ['q_proj', 'k_proj', 'v_proj', 'o_proj']:
llaer getattr flayer.self_attn, proj]
layer LoRALayer Layer, sslora_kwargs)
14
tetattr wew.ayyerrrel...ee.m...e....eeee attn, proj, Lora layer)
15
16 3. Finetune LIM with LORA
xent_loss nn. www.rrsstttoooyoosssssss.s
18
train_loader load_data (+ dataa
19
20 for ix, input_ids ids in enumerate (tadii_loadrrr)r
# Process input tokens
next_token_idd ids model (input_ids)
23
Train model predict next tcken
24
Y-pred next_ text_tooee_idds..... :1] y_true input_ids [ 1:]
25
loss xent_loss (y_pred, y_true) loss. backwardd
Listing 10: End-to-end low-rank adjusting pseudocode (Stage
39

---

Despite the theoretical efficiency of linear attention, existing implementations have long underper- formed well-optimized attention implementations (c.g.. FlashAttention) wall clock speed (Dao et al., 2022). To translate the benefits of LoLCATS to wall clock speedups, we develop custom hardware-awand aware algorithm for LOLCAT' prefill using the ThunderKittens CUDA framework." We Published as a conference paper at ICLR 2025
C.2 HARDWARE-AWAR IMPLEMENTATION OF LOLCATS SLIDING WINDOW
first briefly review the GPU execution model and then detail our algorithm.
C.2.1 GPU EXECUTION MODEL
GPUs workloads are exccuted 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 SM threadblock can access shared memory, and threads within threadblock have independent register memory. To reduce the I/O costs, locality is key kernels should perform many operations as possible 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 newer generations of hardware. Tensor cores-speciall
compute units for matrix-matrix multiplications- 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 idcally exploit the temsor 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, classic principle in systems is to pipeline computation among parallel workers.
C.2.2 THUNDERKITTEN CUDA KERNEL FOR PREFILL
We describe our overall approach below and provide pseudocode in Algorithm 3, designed around the three principles above: memory locality, tensor core utilization, and pipelined execution.
The kernel fuses the entire LOLCATS layer, taking as input the attention queries, keys, and values, for q. k, E Nxé with sequence length N and head dimension d and outputting the result of the E R N Following Llama 3 (AI@ Meta, 2024a), we let = 128 in the discussion below.
Pipeline execution overview. Each thread block handles single batch and head element of size N x The kernel loops over chunks of length 64 along the sequence dimension N, loading 64 x 128 tiles of q, k, v, which we'll refer to qt, k,, Ut in each iteration 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 x 128 sized tiles, recall that the output simply the window attention result. At iteration the warpgroup loads q: küs v, into thread registers, use the tensor cores to multiply queries and keys, apply 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 warpgroup operations (e.g., WGMMA) introduced in the H100 architecture to perform these operations. We refer to the terraced window output as terrace..
40

---

Published as a conference paper at ICLR 2025
Because the window attention is relatively cheap, warpgroup 1 also helps handle loads and stores between HBM and SRAM for the entire kernel. We use tensor memory acceleration (TMA), 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 formula- tion on the left shows quadratic view, wherein ((q,)) and o(k,) multiplied first, while the right formulation shows a linear view, wherein k; and v. are multiplied first.
ON. Σ (blqn) (kka)) ((qn) (T (kk)) E. To(k))
((g.) E (k:))
Since LOLCATS no linear attention the diagonal tiles, the linear attention contribution for tile f is as follows, where queries are multiplied by the cumulative KV state from the prior iterations up 1:
((q.) (T.))(( (E (k)))
Ú. ((
(10)
(kk;
At t = O, the KV state and K state initialized to O, maintained in warpgroup 2's registers.
At iteration t, warpgroup 2 loads in the learned feature maps into register and computes qfi. This result gets multiplied by the running KV state far up until t (again 0 at iteration 0), and the result, linear, gets written to SMEM.
The warps then update the KV state to prepare for the next iteration by featurizing k2 to kft using the learned feature map, and multiplying by vs with WGMMA operations. Note that because the KV state in linear attention somewhat large (d x d), leave the state in register throughout the kernel execution to avoid I/O costs.
Combining the results. Warpgroup loads the linear, contribution from SMEM to its registers, adds the terraced, component, normalizes the overall result, and stores it back to HBM using TMA asynchronous store operations. We provide pseudocode in Algorithm
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).
Algorithm LOL.ATS ThunderKittens prefill kernel
Input: Attention queries, keys, and values k, v Nxd for head dimension d and sequence length N
Output: LOLCATS attention output E Nxd Let localıv be the cumulative recurrent state ("KV-state") initialized to 0 in warpgroup 2's registers.
1: for 1 + to 64 do
Warpgroup 1, Terraced attention
2: Load q:, kt, v. E RRoxxd from HBM to SMEM.
3: Compute the terraced attention output tile terrace 64xd register using WGMMA operations.
Warpgroup 2, Linear attention
4:
Featurize qi by multiplying with the learned feature map to obtain qft 5: Compute the linear attention output tile linear.t Re1xd in register, using qft and
localgy.
6: Write linear,, from register to SMEM
7:
Featurize k, by multiplying with the learned feature map to obtain kf:
8: Update localky by multiplying ff, and v,, and adding the result to local in place, all
in register
Warpgroup 1, Combine results
9:
Load linear... from SMEM register
10: Add o, linear.: + terrace, in register
Write o, to HBM
41

---

Published as a conference paper at ICLR 2025
D EXTENDED RELATED WORK
D.1 LINEARIZING TRANSFORMERS
In this work, we build upon both approaches explicitly proposed to linearize LLMs (Mercat et al., 2024), as well as prior methods focusing on smaller Transformers reasonably adaptable to modern LLMs (Kasai et al., 2021; Mao, 2022; Zhang et al., 2024). We highlight two approaches most related to LOLCATS and their extant limitations next.
Scalable UPtraining for Recurrent Attention (SUPRA). Mercat et al. (2024) linearize LLMs by swapping softmax attentions with linear attentions similar to Retentive Network (RetNet) lay- ers (Sun al., 2023), before jointly training all model parameters on the RefinedWeb pretraining dataset (Penedo et al., 2023). In particular, they suggest that linearizing LLMs with the vanilla linear attention in Eq. 2 unstable, and swap attentions with
yn GroupNorm
7 m- (lgn) To(k)))))
(11)
GroupNorm (Wu & He, 2018) is used as the normalization in place of the Ei= (qn) T<())
(k:)
denominator in Eq. 2, 7 is a decay factor as in RetNet, and Φ is modified learnable feature map from Transformer-to-RNN (T2R) (Kasai et al., 2021) with rotary embeddings (Su et al., 2024). In other words, <xx)- = RoPE W +)) with W E RRadd and E R as trainable weights and biases. With this approach, they recover zero-shot capabilities in linearized Llama 2 7B (Touvron et al., 2023b) and Mistral 7B (Jiang et al., 2023) models on popular LM Evaluation Harness (Gao al., 2023) and SCROLLS (Shaham al., 2022) tasks.
train linear attentions to approximate softmax at- tentions, 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 (Vaswani et al., 2017)) are first copied from an existing softmax attention. They then specify learnable feature maps <(x) = [softmax W + b) @ softmax(=xW> @ and the softmax are applied over the feature dimension) for q and & in each head and layer, and train & such that linear attention weights à match Transformer's original softmax weights a. Given some sample data, they update & with a cross-entropy-basee distillation to minimize:
Hedgehog. Zhang et al. (2024) show
b)] (where t denotes concatenation, and both
22
cxp(a, k.
((q,.) TK(k))
Cn
log
(12)
1-1 exp k./ vd) Σ p(qn) [kk,
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 (Wang et al., 2018) or WikiText-1033 language mod- eling (Merity et al., 2017)), 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 (Devlin et al., 2018) and 125M GPT-2s (Radford al., 2019)). They further show LLM linearizing by linearizing Llama 7B for a specific SAMSum summarization task (Gliwa et al., 2019).
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 (Katharopoulos et al., 2020; Choromanski et al., 2020; Chen et al., 202 bb; Xiong et al., 2021: Zhang et al., 2024; Qin et al., 2022; Schlag et al., 2021; Yang et al., 2023) and sparse attentions (Beltagy et al., 2020; Chen al., 2021a; Zaheer et al., 2020), other efficient architectures include convolutional models (Poli et al.,
2023a; Gu et al., 2021), recurrent neural nets (RNNs) (not strictly following the attention inductive
bias) (Peng et al., 2023a; De et al., 2024; Beck et al., 2024), and state-space model (SSM)-inspired neural net layers (Gu al., 2021; Gu & Dao, 2023; Dao & Gu, 2024).
42

---

Published as a conference paper at ICLR 2025
Improving LLMs for Long Contexts. We share motivation various works focus improving context length quality (Chen et al., 2023a; kai; Peng et al., 2023b). They that the rotary embed- dings (Su et al., 2024) of popular LLMs (Biderman al., 2023; Touvron et al., 2023b; AI@ Meta, 2024a; Jiang et al., 2023) struggle to handle contexts longer than those used training. Chen et al. (2023b) 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 im- prove efficiency by optimizing the KV Cache from systems standpoint (Kwon et al., 2023). These may present complementary optimizations linear attention's 's algorithmic improvements.
E ADDITIONAL ARTIFACTS
E. SAMPLE GENERATIONS
We include sample generations for various prompts, using Llama 8B models linearized with dif- ferent LoL.CATS-related linearizing setups on the Alpaca dataset. For reference, we also include responses from a standard (softmax attention) Llama 8B finetuned with LoRA (rank 8, alpha 16, applied to all attention projections) on the Alpaca training set. LOLCATS as presented often leads to greatest coherence among linearized LLMs.
Prompt
Below is an instruction that describes task. Write a response that appropriately completes the request.
< <Ibegin of ftteett >>
### Instruction:
What the difference between H&M and Balenciaga?
### Response:
Hedgehog (Linear Only) (No LoRA)
Hindi: the difference is that in the first one little bit more complicated, but the other one is that is the same. < end_doo of text
Hedgehog (Linear Only) (Attention Transfer + LoRA)
H&M and Balenciaga are two of the most popular fashion brands the world, but there are some differences between them.
H&M is a fast fashion retailer that sells trendy and affordable clothing. It was founded in Sweden in 1947 and is now available in over 60 countries worldwide. The brand is known for its fast fashion and affordable prices, making popular choice for those on a budget.
Balenciaga, on the other hand, a luxury fashion house that was founded in 1977. known for its high-end ready-to-wear pieces, including leather jackets, dresses, and shoes. The brand is known for its use luxurious fabrics and intricate designs, as well as its collaborations with other fashion houses such as Versace and Versace.
In summary, H&M and Balenci two popular fashion brands that offer different styles and price points. While H&M is known for its fast fashion and affordable fashion, Balenciaga is known for its luxury and high-end ready-to-wear /ieces.</eedd.of.tee
LoLCATS (Linear + Sliding Window) (No LoRA)
The difference between H&M and Balenciaga is that H&M is high street brand while Ba- Ienciaga is a luxury brand. H&M is more affordable than Balenciaga, which is luxury brand.< end of text
43

---

Published as a conference paper at ICLR 2025
LOLCATS (Linear + Sliding Window) (No Attention Transfer)
H&M and Balenciaga are both clothing brands, but there are some key differences between the
H&M fast fashion brand known for its affordable and trendy clothing. The company was founded in Sweden and has stores in over 60 countries. It is known for its seasonal collections, which include clothing. accessories, and home decor. H&M has reputation for its ethical and sustainable practices, and the company works with designers and brands to create its collections.
On the other hand, Burberry is British luxury fashion house founded in 1856. It is known for its high-end fashion, including clothing, accessories, and fragrances. The company has strong heritage, and its designs inspired by British culture and history. Burberry is known for its classic trench coats, and its signature check pattern is instantly recognizable.
Overall, both brands are high-end fashion companies, but they have different target audiences and design styles. H&M is more accessible and affordable, while Burberry is more expensive and known for its classic, timeless designs. < <lendof.texx
LOLCATS (Linear + Sliding Window) (Attention Transfer LoRA)
H&M and Balenciaga are both fashion brands, they have significant differences.
H&M is a fast-fashion brand that offers a wide range clothing, accessories, and home decor at affordable prices. The brand is known for its trendy and stylish designs, and it often collab- orates with high-profile designers to create limited-edition collections. H&M is a mass-market retailer, and its products available in its own stores, well in department stores and online.
Balenciaga, the other hand, is luxury fashion house founded in 1917 by Spanish designer Cristobal Balenciaga. The brand is known for its avant-garde designs, high-quality materials, and attention to detail. Balenciaga's collections are often bold, experimental, and daring, and the brand has strong following among fashion enthusiasts and celebrities. H&M, on the other hand, a fast-fashion retailer that offers affordable clothing and accessories inspired by the latest trends. While H&M also produces its own clothing and accessories, the brand is known for its collaborations with high-end designers and its ability to produce fashion at a much lower price point than other luxury brands.
Softmax Attention
H&M and Balenciaga two fashion brands that quite different from each other.
H&M is fast fashion brand that offers affordable clothing and accessories for both and women. The brand is known for its trendy and on-trend designs, and its collections often feature collaborations with high-profile designers and celebrities. H&M mass-market brand that aims to make fashion accessible to everyone.
Balenciaga, on the other hand, is luxury fashion brand that is known for its avant-garde designs and high-end materials. The brand is associated with couture and haute couture, and its collections often feature bold and experimental designs. Balenciaga is a high-end brand that caters to a more exclusive clientele.
In summary, the main difference between H&M and Balenciaga is their target audience, price point, and design philosophy. H&M is fast fashion brand that offers affordable and trendy designs, while Balenciaga is luxury brand that focuses on high-end and avant-garde de- jegnss//end.of_ter.
44

---

Published as a conference paper at ICLR 2025
Prompt
< (lbegin.of.text >> Below is an instruction that describes task. Write a response that appropriately completes the request.
### Instruction:
Explain why the following fraction is equivalent to 1/4. 4/16
#### Response:
Hedgehog (Linear Only) (No LoRA)
The number 4/4 fraction. This is fraction that is equivalent to 4/4. The number 4/4 is the number 4/4/4/4/4,< < end.of.text >>
Hedgehog (Linear Only) (Attention Transfer + LoRA)
The fraction 1/4 is equivalent 4/16. To why, denominator by 4 to obtain the equivalent fraction 16/16. This because the multiplication of two equal numbers will not change the value of the fraction.
multiply both the numerator and
The number 4/16 be rewritten 16/4, which is equivalent to 4/4. Similarly, the numerator 4/4 can be rewritten as 4/4, Therefore, the fraction 4/4 is equivalent to 4/16, which is the same as 1/4, end.of.eext
LOLCATS (Linear + Sliding Window) (No LoRA)
The fraction 4/16 is equivalent to 1/4 because the numerator and denominator are both divisible by 2. This is because 4/16 is equal to 4/16. This is because the numerator and denominator are both divisible by 2. This means that 4/16 is equal to 1/4. Therefore, 1/4 equal to 1/4, which is 1/4. Therefore, 4/16 is equal 1/4, and 1/4 is equal 1/4.< end_of_text
LOLCATS (Linear + Sliding Window) (No Attention Transfer) 1/4 4/16
explain why the fraction 1/4 is equivalent to 4/16, we first need to understand what fraction is. A fraction is way to represent part of whole, where the whole is divided into equal parts. In this case, the whole is the whole number 4, which is divided into 4 equal parts, each of which is represented by the fraction 1/4. The symbol used to represent a fraction is fraction bar, which is a horizontal line that goes between the numerator and the denominator. So, the fraction 1/4 can be written as 1/4 4/4.4 end_of_text
LOLCATS (Linear- Sliding Window) (Attention Transfer LoRA)
1/4 because both fractions represent the same value. When simplifying 4/16, we can divide both the numerator and denominator by 4 to get 1/4. This is because divided by 4 is equal to I, and 16 divided by also equal 4. So, we can cancel out the 4s to get 1/4.< end of text
The fraction 4/16 is equivalent
45

---

Published as a conference paper at ICLR 2025
Softmax Attention
The fraction 4/16 is equivalent to 1/4 because both fractions represent the same value. To see why this is the case, we can simplify the fraction 4/16 by dividing both the numerator and the denominator by their greatest common factor.
The greatest common factor (GCF) of 4 and 16 is 4. We simplify the fraction by dividing both the numerator and the denominator by 4:
+ 1
16 + 44
Thus, the simplified fraction is 1/4. Since the simplified fraction is equivalent to the original fraction, 4/16 is equivalent to 1/4.< end.of.text
Prompt
< <lbegin of ttext >> Below is an instruction that describes task. Write a response that appropriately completes the request.
### Instruction:
Why did we linearize the chungus among us?
### Response:
Hedgehog (Linear Only) (No LoRA)
The reason that the request is to request the request is to the request to quest.< |nd of text |>
Hedgehog (Linear Only) (Attention Transfer + LoRA)
The reason linearized the chunus among us is make it casier for understand and com- pare their behavior. By converting the non-linear behavior of the chunus into linear equations, we can more easily analyze and interpret their behavior, and make predictions about their future behavior. Additionally, it allows us to compare the behavior of the chunusus among themselves, and compare their behavior over time < end ff_etttl |>
LoLCATS (Linear + Sliding Window) (No LoRA)
Because have other choice.< /end_o >>
LOLCATS (Linear + Sliding Window) (No Attention Transfer)
The word "chungus" is not in the English language, earized.< end textext |>
SO it cannot be
lin-
LOLCATS (Linear + Sliding Window) (Attention Transfer + LoRA)
Linearizing the chungus among us is process of organizing and arranging the information in logical and coherent manner. This is done to make easier for people to understand and interpret the data, well to make it more accessible and usable. By linearizing the data, we are able to extract meaningful insights and patterns from the information, which can then used to make informed decisions and take action.< < end of text
Softmax Attention
I'm sorry, but cannot respond to this question do not have the necessary information to provide a meaningful response. Can you please provide more context or details about the topic you'l re referring to?< end of text
46

---

Published as a conference paper at ICLR 2025
E.2 SAMPLE ATTENTION WEIGHTS
We finally visualize sample attention weights from LOLCATS 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 LOLCATS 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). LOLCATS 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 LOLCATS training and architectures. Newly initialized LoL.- CATs linear + sliding window attentions do not capture softmax attention weights (init.), supporting carlier results that attention transfer is necessary to reduce architectural differences and improve lin- carizing efficiency. Trained LOLCATS attentions also match attention weights better than trained Hedgehog linear attentions (same feature map, but no sliding window). These results suggest LoL- CATs attention transfer and linear + sliding window layers allow to learn better approximations of softmax attention weights, coinciding with improved lienarizing quality.
Softmax
Softmax
Softmax
Softmax Layer
Softmax Layer
Layer
Layer
Layer
388
905
Querios
905
535
20
1k6 i0ls
1k9 196 1915
PA 1088
Keys
Knys
Keys
Keys
Knys
LOLCATS
LoLCATS
LeLCATS Layer Head
LoLGATs (MSE) Layer
LoLCATS (MSE) Layer
Layer Head
Layer Head
Quartes
330
10
20
1206 1013
20
20
PV 1mm 1010
10
50 Keys
10
589 Keys
300 3213
20
PM 10m8 sea
Keys
Keys
Keys
Lol. CATs (XENT) Head 0
CAT (XENT) Head 0
LOLCATS (XENT) Layer 16 Head
CATs (XENT)
LoLCAT: (XENT)
Head
Head
Quertss
10
20
196
1in 10m 1039 Keys
10
sis 190 Dis Keys
399 1085 103 Keys
Keys
Keys
Hedigehog (MSE)
Hedgehog (MSF) Layer Head B
Hodgehog (MSF) Layer 16 Head
Hedinhag (MSE) Layor 24 Head
Hedgehag (MSF)
Layer Head D
Layor 31 Head
Querfes
2000
3005
10W
X09
303
103
3006 ids
10B 1039
199
196
106
Keys
Keys
Keys
Keys
Lol CATs (init.) Layer D Head D
Lot CAT (init.) Layer 8 Head B
LATS 6nit.) Layer 16 Head
CATS (init.) Layer 24 Head
CoL (init.)
Layer 31 Head
Quaries
3308
1259
Keys
Keys
Keys
Keys
Hedgeeog (init.) Layer Head D
Hedgehog linit.) Layer Head D
Hedgehog linit.) Layer Head
Hedgehop (init.) Layer Head
Hediehog (init.) Layer Head
Querioo
1530
1239
199
1819
1889
15
14
Keys
Keys
Keys
Keys
Keys
Figure 18: Llama 3 8B attention weights; head layers 0, 8, 16, 24, 31.
47

---

Published as a conference paper at ICLR 2025
Softmax Layer Head
Softmax Layer Head
Softmax Layer Head
Softmax Layer Head
Softaax
Layer
188
56
Quarioos
1559
1k6 iels
Lib L23
549 1i6 1955
1239
A9 11t
Keys
Keys
Keys
Keys
Keys
LoLGATS Layer
LoLGATS (MSB) Layer
Layer
Layer
Layer
00
20
1kt 1816
20
10t 1819
29
20
1nin 1969
16
669 190 1915
99 10t 109
Keys
Keys
Keys
Keys
Keys
Lol CALC
LaLCAT: Layer Head
LOLCAT (XENT) Head
LoL (XENT)
LoLCAT: (XENT)
Quarlos
1299
00
20
196
1its
190 10W 1039
199 D00
10
190 195
39 10m 108 Keys
Keys
Keys
Keys
Keys
Heddgehog (MSE)
Hedgehog (MSE)
Hedgehog (MSE) Layer 16 Head
Hedgehog (MSE) Layer 24 Head
Hedgehog (MSE)
Layor Head
Layar Head 1
Layer 31 Head
1208
15m 1019
11m
Keys
Keys
Keys
Keys
Keys
LoLCATs (init.) Layer 0 Head
LoLCA's (init.) Layer B Head
LoLCATs (nni.) Layer 16 Head
LoLCATS (init.) Layer 26 Head
LoLCA's (init.) Layer 31 Head
12:9
Keys
Keys
Keys
Hedgeeoo (init.) Layer 0 Head
Hedgehog linit.) Layer Head
Hedgehog linit.) Layer 16 Head
Hedgehog (init.) Layer 24 Head
Hediehhg (init.) Layer 31 Head
1038
300
1015
103
Keys
Keys
Keys
Keys
Keys
Figure 19: Llama 3 8B attention weights; head layers 0, 24, 31.
48

---

Published as a conference paper at ICLR 2025
Softmax Layer Head
Softmax Layer Head
Softmax Layer Head
Softmax Layer Head
Softmax
Layer
188
56
Quarioos
1259
1k6 idla
Libe L23
549 196 1955
A9 10m
Keys
Keys
Keys
Keys
Keys
LoLCATS Layer
LoLGATS (MSB) Layer
(MSE)
Layer
Layer
Layer
00
20
19t 1816
20
20
19
10t
1819
29
20
180n 1969
16
669 190
1915
99 10t 109
Keys
Keys
Keys
Keys
Keys
Lol LALC
LaLCAT: Layer Head
LOLCAT (XENT) Head
LOL (XENT) Layer
LoLCAT: (XENT)
Quarlos
1299
10
20
196
101s
999 10W 1039
199 D00
10
190 195
99 10m 108 Keys
Keys
Keys
Keys
Keys
Heddgehog (MSE)
Hedgehog (MSE)
Hedgehog (MSE) Layer 16 Head
Hedgehog (MSE) Layer 24 Head
Hedgehog (MSE)
Layor D Head 2
Layar Head 2
Layer 31 Head
Quaries
1208
15th 1029
180
11m8
Keys
Keys
Keys
Keys
Keys
LoLCATs (init.) Layer 0 Head 2
LoLCA's (init.) Layer B Head 2
LoLCATs (nnt.) Layer 16 Head
LoLCATS (init.) Layer 26 Head
LoLCA's (init.) Layer 31 Head
12:9
Keys
Keys
Keys
Hedggeho (init.) Layer 0 Head 2
Hedgehog linit.) Layer Head 2
Hedgehog linit.) Layer 16 Head
Hedgehog (init.) Layer 24 Head
Hediehhg (init.) Layer 31 Head
1038
300
1015
103
Keys
Keys
Keys
Keys
Keys
Figure 20: Llama 8B attention weights; head 2,; layers 8, 16, 24, 31.
49

---

Published as a conference paper at ICLR 2025
Softmax Layer Head
Softmax Layer Head
Softmax Layer Head
Softmax Layer Head
Softmax Layer Head
188
58
Quarroos
1259
126 i0ls
1019
Libe 123
549 1is 195
A9 10t
Keys
Keys
Keys
Keys
Keys
LoLGATS Layer
LoLGATS (MSB) Layer
(MSE)
Layer
Layer
Layer
00
20
19h
20
16
10t
1819
29
20
1n0n 1969
16
669 190
1915
99 10t 109
Keys
Keys
Keys
Keys
Keys
LaLCAT: Layer Head
LOLCAT (XENT) Head
Lol (XENT)
LoLCAT: (XENT)
Head
Layer
Quarlos
1299
00
20
196
1015
190 10W 1039
19
199 D00
10
190 195
39 10m 105
Keys
Keys
Keys
Keys
Keys
Heddgehog (MSE)
Hedgehog (MSE)
Hedgehog (MSE) Layer 16 Head
Hedgehog (MSE) Layer 24 Head
Hedgehog (MSE)
Layar D Head 3
Layar Head 3
Layer 31 Head
Quaries
1208
115
15mh 1019
180
326
11m
Keys
Keys
Keys
Keys
Keys
LoLCATs (init.) Layer 0 Head
LoLCA's (init.) Layer B Head 3
LoLCATs (jnt.) Layer 16 Head
LoLCATS (init.) Layer 26 Head
LoLCA's (init.) Layer 31 Head
12:9
Keys
Keys
Keys
Hedggeoo (init.) Layer 0 Head
Hedgehog linit.) Layer Head
Hedgehog linit.) Layer 16 Head
Hedgehog (init.) Layer 24 Head
Hediehhg (init.) Layer 31 Head
Querle
1038
300
1015
103
Keys
Keys
Keys
Keys
Keys
Figure 21: Llama 3 8B attention weights; head layers 0, 24, 31.
50

---

Published as a conference paper at ICLR 2025
Softmax
Softmax Layer Head
Softmax
Softmax Layer Head
Softmax Layer Head
Layer Head 0
Layer Head
Quarioos
1559
1k6 i0ls
Libe L23
15
549 196
AA 110
Keys
Keys
Keys
Keys
Keys
LoLGATS (MSB)
(MSE)
Layer
Layer
Layer
Layer
Layer
00
20
1kn 18l6
10th 1819
1n 1963
16
190 191
99 16t6 129
Keys
Keys
Keys
Keys
Keys
Lol CATS
LaLCAT: Layer Head
LOLCATs (XENT) Head
Lol CATs (XENT)
LoLCAT: (XENT)
Head
Quarlos
1299
00
20
196
10ts
199 10W 1039
39
1999
10
390 195
99 10m 103 Keys
Keys
Keys
Keys
Keys
Heddgehog (MSE)
Hedgehog (MSE)
Hedgehog (MSE) Layer 16 Head
Hedgehog (MSE) Layer 24 Head
Hedgehog (MSE)
Layar D Head D
Layor Head 0
Layer 31 Head
3308
15mh 1019
180
11m8
Keys
Keys
Keys
Keys
Keys
LoLCATs (init.) Layer 0 Head 0
LoLCA's (init.) Layer B Head g
LoLCATs (nni.)) Layer 16 Head
LoLCATs (init.) Layer 26 Head
LoLCA's (init.) Layer 31 Head 0
12:9
Keys
Keys
Keys
Hedggeho (init.) Layer 0 Head 0
Hedgehog linit.) Layer Head 0
Hedgehog linit.) Layer 16 Head
Hedgehog (init.) Layer 24 Head
Hediehhg (init.) Layer 31 Head
Querle
1038
300
1015
103
Keys
Keys
Keys
Keys
Keys
Figure 22: Mistral 7B v0.1 attention weights; head layers 0, 8, 16, 24, 31.
51

---

Published as a conference paper at ICLR 2025
Softmax Layer Head
Softmax Layer Head
Softmax Layer Head
Softmax Layer Head
Softmax
Layer
188
56
Quarioos
1559
1k6 i01
Lib L23
549 1h6 191
A9 10m
Keys
Keys
Keys
Keys
Keys
LoLGATS Layer
LoLGATS (MSB)
Layer
Layer
Layer
Layer
00
20
10t 1819
16
659 190 191
99 10t 109
Keys
Keys
Keys
Keys
Keys
Lol
LalCATs Layer Head
LOLCAT (XENT) Head
Lol CATs (XENT)
LoLCAT: (XENT)
Head
Quarlos
1299
00
20
196
10ts
190 10W 1039
39
1999
10
390 195
399 10m 108
Keys
Keys
Keys
Keys
Keys
Heddgehog (MSE)
Hedgehog (MSE) Layor 8 Head 1
Hedgehog (MSE) Layer 16 Head
Hedgehog (MSE) Layer 24 Head
Hedgehog (MSE)
Layar D Head
Layer 31 Head
3208
15m 1019
11:8
Keys
Keys
Keys
Keys
Keys
LoLCATs (init.) Layer 0 Head
LoLCA's (init.) Layer B Head
LoLCATs (nit.) Layer 16 Head
LoLCATs (init.) Layer 26 Head
LoLCA's (init.) Layer 31 Head
1239
Keys
Keys
Keys
Hedgeeoo (init.) Layer 0 Head
Hedgehog linit.) Layer 8 Head
Hedgehog linit.) Layer 16 Head
Hedgehog (init.) Layer 24 Head
Hediehhg (init.) Layer 31 Head
Querle
1038
300
1015
103
Keys
Keys
Keys
Keys
Keys
Figure 23: Mistral 7B v0.1 attention weights; head 1; layers 0, 8, 16, 24, 31.
52

---

Published as a conference paper at ICLR 2025
Softmax Layer Head
Softmax Layer Head
Softmax
Softmax Layer Head
Softmax Layer Head
Layer Head
188
56
Quarioos
99
1559
1k6 idss
13b L23
15
549 196
1219
A9 10m
Keys
Keys
Keys
Keys
Keys
LoLCATS
LoLGATS (MSB)
(MSE)
Layer
Layer
Layer
Layer
Layer
00
20
vin 18le
10t 1819
29
190n 1969
16
651 190
191
999 10t 129
Keys
Keys
Keys
Keys
Keys
Lol
LaLCAT: Layer Head
LOLCATS (XENT) Head
LOL (XENT)
LoLCAT: (XENT)
Quarlos
1299
00
20
196
10ts
190 10W 1039
19
1999 D03
10
390 1915
99 10m 103 Keys
Keys
Keys
Keys
Keys
Heddgehog (MSE)
Hedgehog (MSE)
Hedgehog (MSE) Layer 16 Head
Hedgehog (MSE) Layer 24 Head
Hedgehog (MSE)
Layar Head 2
Layor Head 2
Layer 31 Head
3208
15m 1019
180
11:8
Keys
Keys
Keys
Keys
Keys
LoLCATs (init.) Layer 0 Head 2
LoLCA's (init.) Layer B Head 2
LoLCATs (nit.) Layer 16 Head
LoLCATs (init.) Layer 26 Head
LoLCA's (init.) Layer 31 Head
12:9
Keys
Keys
Keys
Hedggeoo (init.) Layer 0 Head 2
Hedgehog linit.) Layer 8 Head 2
Hedgehog linit.) Layer 16 Head
Hedgehog (init.) Layer 24 Head
Hediehhg (init.) Layer 31 Head
Querle
1038
300
1015
103
Keys
Keys
Keys
Keys
Keys
Figure 24: Mistral 7B v0.1 attention weights; head layers 0, 8, 16, 24, 31.
53

---

Published as a conference paper at ICLR 2025
Softmax Layer Head
Softmax Layer Head
Softmax Layer Head
Softmax Layer Head
Softmax Layer Head
188
56
Quarioos
1559
1k6 idss
Lib L25
549 1h6 125
AA 10t
Keys
Keys
Keys
Keys
Keys
LoLCATS Layer
LoLGATS (MSB)
(MSE)
Layer
Layer
Layer
Layer
00
20
10t 1819
29
190n 196
16
651 190 191
Keys
Keys
Keys
Keys
Keys
Lol
LaLCAT: Layer Head
LOLCATs (XENT) Head
CATs (XENT)
LoLCAT: (XENT)
Quarlos
1299
00
20
196
1its
390 10W 1039
39
1999 D00
10
390 195
99 10m 108
Keys
Keys
Keys
Keys
Keys
Heddgehog (MSE)
Hedgehog (MSE)
Hedgehog (MSE) Layer 16 Head 3
Hedgehog (MSE) Layer 24 Head
Hedgehog (MSE)
Layor D Head 3
Layor Head 3
Layer 31 Head
3208
15mh 1029
11m
Keys
Keys
Keys
Keys
Keys
LoLCATs (init.) Layer 0 Head 3
LoLCA's (init.) Layer B Head 3
LoLCATs (nit.) Layer 16 Head
LoLCATs (init.) Layer 26 Head
LoLCA's (init.) Layer 31 Head
1239
Keys
Keys
Keys
Hedgeeoo (init.) Layer 0 Head
Hedgehog linit.) Layer Head
Hedgehog linit.) Layer 16 Head
Hedgehog (init.) Layer 24 Head
Hediehhg (init.) Layer 31 Head
1038
300
1015
103
Keys
Keys
Keys
Keys
Keys
Figure 25; Mistral 7B v0.1 attention weights; head 3; layers 0, 8, 16, 24, 31.
54


Reviews:
[{"review_id": "sHj78mJ3Ew", "reviewer": "Reviewer_NsVB", "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.", "questions": "- How were the hyperparameters for the attention transfer and low-rank adaptation chosen?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 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.\n- 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.\n- 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. \n- 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.", "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.\n- 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).\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NEcx71xRFl", "reviewer": "Reviewer_M1ap", "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.", "questions": "1. I would like the author to supplement and polish the article based on the weaknesses.\n\n2. 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.\n\n3. 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?\n\n4. Does this approach extend to other variants of Linear Attention, such as Mamba 1/2 and Hgrn 1/2?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 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. \n\n2. 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.\n\n3. LOLCATS is the first to successfully apply linearization to large models with 70B and 405B parameters, expanding the applicability of linearization techniques.", "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.\n\n2. 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.\n\n3. 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.\n\n4. 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.\n\n[1] Xuyang Shen, Dong Li, Ruitao Leng, Zhen Qin, Weigao Sun, & Yiran Zhong. (2024). Scaling Laws for Linear Complexity Language Models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "x0EG4lHDOC", "reviewer": "Reviewer_Dj9p", "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.", "questions": "- Why was the proposed method not validated on a smaller model, such as llama1B?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The method is shown to scale to unprecedentedly large models (70B and 405B parameters), which was not possible with previous linearization techniques.\n- The method is applied to various models, showing its broad applicability and potential impact on the field of natural language processing.", "weaknesses": "- There is a lack of an overall summary description of the LOLCATS method. An algorithm description or pseudo code can be added.\n- There are some writing errors, such as Line 294: APP.???", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XQrYBh145Q", "reviewer": "Reviewer_povv", "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.", "questions": "1. Can authors point out which line of appended code correspond to a. Eq 7?\n\n2. 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.\n\n3. 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? \n\n3. 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.\n\n4. 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? \n\n5. Is there a study of different window size?\n\n6. 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.\n\n7. Also don't know why ablation is only done on MMLU, I am interested in seeing RCV-e and PiQA too.\n\n8. 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.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The Introduction reads well (but unfortunately not the case for the rest)\n\n- An interesting problem to study.\n\n- Provide the code in the Appendix", "weaknesses": "- The way authour present their methods interleaves many previous works, and thus it's difficult to precisely pinpoint what's their contribution. \n\n- Some lacking experiments make me uncertain if the claimed effectiveness is true.\n\n- The above two factors combined make me uncertain what's the real efficacy of the method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)