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
Primary area: foundation or frontier models, including LLMs
Keywords: Linear Attention, Linearizing Transformers, Low-rank Adaptation, Large Language Models, Architecture Distillation

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.

Decision: accept (poster)
Reviewer scores: 8, 6, 8, 5

Reviews and discussion:
<thread>
creation date:
1730666527774

thread_type:
official_review

<thread_object>
creator: reviewer_povv

<content>
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.

strengths:
- The Introduction reads well (but unfortunately not the case for the rest)

- An interesting problem to study.

- 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. 

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

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

questions:
1. Can authors point out which line of appended code correspond to a. Eq 7?

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

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

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

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

5. Is there a study of different window size?

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

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

8. I don't really see what's the help from the linearization to parameter-efficient finetuning. I still don't understand why PEFT can't work on other methods. In addition, full FT should still outweigh PEFT. I believe authors should also show their method with full parameter tuning to see what's the difference.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730698421571

thread_type:
official_review

<thread_object>
creator: reviewer_dj9p

<content>
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.

strengths:
- The method is shown to scale to unprecedentedly large models (70B and 405B parameters), which was not possible with previous linearization techniques.
- The method is applied to various models, showing its broad applicability and potential impact on the field of natural language processing.

weaknesses:
- There is a lack of an overall summary description of the LOLCATS method. An algorithm description or pseudo code can be added.
- There are some writing errors, such as Line 294: APP.???

questions:
- Why was the proposed method not validated on a smaller model, such as llama1B?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730705635108

thread_type:
official_review

<thread_object>
creator: reviewer_m1ap

<content>
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.

strengths:
1. The LOLCATS approach significantly reduces computational and memory costs through a two-step process involving "attention transfer" and low-rank adaptation (LoRA), effectively lowering the complexity of large models. 

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

3. LOLCATS is the first to successfully apply linearization to large models with 70B and 405B parameters, expanding the applicability of linearization techniques.

weaknesses:
1. The paper mentions the differences in errors across layers and their impact on model performance, but it does not sufficiently discuss the underlying reasons, such as why lower soft-attention entropy leads to higher errors.

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

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

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

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

questions:
1. I would like the author to supplement and polish the article based on the weaknesses.

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

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

4. Does this approach extend to other variants of Linear Attention, such as Mamba 1/2 and Hgrn 1/2?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730718038174

thread_type:
official_review

<thread_object>
creator: reviewer_nsvb

<content>
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.

strengths:
- I believe this is a good work. The paper presents a novel method for linearizing LLMs that addresses key limitations of existing techniques. By focusing on approximating softmax attention outputs and using low-rank adjustments, the authors offer a fresh perspective on reducing computational complexity without sacrificing model quality.
- The LOLCATS method reduces the amount of training required, both in terms of model parameters updated and training data used. This efficiency makes the method practical for widespread use, especially in environments with limited computational resources.
- Demonstrating the ability to linearize large models up to 405B parameters is a notable achievement. The scalability of the method suggests it can be applied to future, even larger models.  The authors provide extensive experimental results, including comparisons with existing methods on multiple benchmarks. 
- The paper delves into the reasons why previous linear attentions struggled to approximate softmax attention effectively. By identifying issues like attention entropy and layer-wise errors, the authors provide valuable insights that inform their improved architecture.

weaknesses:
- The introduction and analysis of the preliminaries are too lengthy, with the core improvements in LOLCATs not appearing until the seventh page, which makes the reading experience somewhat disjointed.
- The experimental results are promising given the training budgets. However, I notice that even though previous works perform significantly worse for more challenging benchmarks (like MMLU in the setting), LOLCAT still considerably underperforms the original models. What will happen when the benchmarks become more challenging? (e.g., complex reasoning).
- While the authors claim improvements in inference efficiency, quantitative metrics such as actual speedup factors, memory utilization during inference, or comparisons of throughput are not extensively reported.

questions:
- How were the hyperparameters for the attention transfer and low-rank adaptation chosen?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732209838812

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response (1/1)

comment:
Many thanks to all reviewers for their helpful feedback and thoughtful comments, and to the ACs for chairing. Below we recap our paper + reviewer comments, and go over updates in our revision.

---
### **Recap**
To recap, towards obtaining LLMs with subquadratic efficiency, we study how to convert or “linearize” modern Transformer-based LLMs into linear attention variants with state-of-the-art quality, training and model parameter efficiency, and scalability.

As highlights, our method (called LoLCATs):  
* Gets state-of-the-art quality on popular LM Eval tasks, outperforming prior linearizing methods and subquadratic LLMs trained from scratch  
* Uses only 40M training tokens (prior linearizing methods use 2500x the tokens [1]), while only training 0.2% of LLM parameter counts via LoRA  
* Scales up linearizing to 70B and 405B LLMs for the first time

We appreciate that reviewers consistently acknowledged our method's **effectiveness and potential impact**, noting that our work:  
* Offers a “**fresh perspective on reducing computational complexity without sacrificing model quality**”, is “practical for widespread use, especially in environments with limited computational resources”, and provides “valuable insights that inform their improved architecture” (NsVB)  
* “**Opens new avenues for scaling LLMs further**”, “expanding the applicability of linearization techniques” (M1ap) 
 
* “Scales to unprecedentedly large models (70B and 405B parameters), which was not possible with previous linearization techniques,” and “is applied to various models, showing its broad applicability and **potential impact on the field of natural language processing**” (Dj9p)  
* Provides an “interesting problem to study” (povv)

---
### **Revision**  
Thanks to reviewer feedback, we uploaded a revision with updates highlighted in blue. We group these updates into two themes by reviewers:

**Writing clarity + presentation**  
* We update the methods section to better present our overall method & first contribution (NsVB, povv):  
  * We study + show for the first time that we can use LoRA to convert Transformer LLMs into viable linear attention variants. All prior methods use full model training, making it unclear before our work if our proposed LoRA linearizing is feasible [1, 2, 3] (povv)   
* We add discussion on prior linear attentions' limitations in this low-rank setup (L294-308) (M1ap)  
* We clarify how these results motivate our additional technical contributions to improve quality (Section 3.3) (povv)    
* We explicitly define terms like “attention transfer” (L199) (M1ap)  
* We summarize with algorithms 1 and 2 (L397 - 409) (DJjp), and update pseudocode to walk thru all components (App. C.1) (DJjp, povv)

**Expanded experimental analysis**  
* *Additional study + evaluation*. We:  
  * Find LoLCATs also gets **state-of-the-art quality when linearizing 1B LLMs** (Llama 3.2 1B, Table 12; Phi 1.5 1.3B, Table 13), outperforming other linearizing methods by 0.2-1.7 points, while only training 0.22% of their parameters (Dj9p)   
  * Provide standard deviations and means over three seeds (Table 11) (M1ap)  
  * Study how to improve quality on challenging tasks like MMLU* (App. B.42, L1275) (NsVB)  
 
  * Include needle-in-a-haystack / passkey retrieval evals: we recover softmax attention-like recall by linearizing with retrieval data (App. B.42, L1323) (M1ap)  
  * Conduct additional layer-wise analysis, tracking MSE error (App. B.6.1) and LoRA weight training dynamics (A and B low-rank matrices, App. B.6.2) (M1ap)

* *Expanded ablations*. We add ablations on:  
  * LoRA rank (including full finetuning) (Table 15) (M1ap, povv)   
  * Linearizing attention sliding window size (Table 17) (povv)  
  * Attention transfer, sliding window, and feature map results on all zero-shot LM Eval tasks (Table 10) (povv)
We supplement these reviewer-requested ablations with more experiments on LoRA projection layer (Table 16), linearizing data (App. B.4.1), and training token budgets (App. B.5)   

*In late-breaking results, by leaving some layers as softmax attention (50% like in prior work [3]), we substantially close the MMLU gap for Llama 3 8B:

|  | Mamba2-Llama (50% softmax attn) [3]   | LoLCATs (0% softmax attn)   | LoLCATs (50% softmax attn) | Llama 3 8B |
|---|:---:|:---:|:---:|:---:|
| 5-shot MMLU % | 55.7 | 52.8 | 65.8 | 66.6 |
---

We thank all reviewers again for their constructive comments, which we believe have strengthened the paper’s presentation. They also led to many more experiments and findings to improve overall content.

Please find our responses for individual reviewer comments below. We are happy to follow up with any questions.

**References**  
[1]  Linearizing Large Language Models, Mercat et al., 2024  
[2]  Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models, Bick and Li et al., 2024  
[3]  The Mamba in the Llama: Distilling and Accelerating Hybrid Models, Wang and Paliotta et al., 2024
</content>
</thread_object>

</thread>


Metareview:
metareview:
## Summary

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

## Decision

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