LoLCATs: On Low-Rank Linearizing of Large Language Models Article Swipe

MAKE IMPACT

Michael Zhang , Simran Arora , Rahul Chalamala , Alan Wu , Benjamin Spector , Aaryan Singhal , Krithik Ramesh , Christopher Ré · YOU? · · 2024 · Open Access · · DOI: https://doi.org/10.48550/arxiv.2410.10254

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.4% of their training tokens. Finally, we apply LoLCATs to create the first linearized 70B and 405B LLMs (50x larger than 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.

Type

preprint

Language

en

Landing Page

http://arxiv.org/abs/2410.10254

PDF

https://arxiv.org/pdf/2410.10254

OA Status

green

Cited By

1

Related Works

10

OpenAlex ID

https://openalex.org/W4403570531

Raw OpenAlex JSON

OpenAlex ID

https://openalex.org/W4403570531

Canonical identifier for this work in OpenAlex

DOI

https://doi.org/10.48550/arxiv.2410.10254

Digital Object Identifier

Title

LoLCATs: On Low-Rank Linearizing of Large Language Models

Work title

Type

preprint

OpenAlex work type

Language

en

Primary language

Publication year

2024

Year of publication

Publication date

2024-10-14

Full publication date if available

Authors

Michael Zhang, Simran Arora, Rahul Chalamala, Alan Wu, Benjamin Spector, Aaryan Singhal, Krithik Ramesh, Christopher Ré

List of authors in order

Landing page

https://arxiv.org/abs/2410.10254

Publisher landing page

PDF URL

https://arxiv.org/pdf/2410.10254

Direct link to full text PDF

Open access

Yes

Whether a free full text is available

OA status

green

Open access status per OpenAlex

OA URL

https://arxiv.org/pdf/2410.10254

Direct OA link when available

Concepts

Rank (graph theory), Computer science, Mathematics, Combinatorics

Top concepts (fields/topics) attached by OpenAlex

Cited by

1

Total citation count in OpenAlex

Citations by year (recent)

2025: 1

Per-year citation counts (last 5 years)

Related works (count)

10

Other works algorithmically related by OpenAlex

Full payload