生產(chǎn)級(jí)大模型應(yīng)用線性注意力的方法，來了。

線性Attention（包括RNN系列），再也不用困在幾B參數(shù)的范圍內(nèi)娛樂了。

一套方法，即可線性化現(xiàn)有各種量級(jí)的Transformer模型，上至Llama 3.1 405B，也只需要十來張顯卡在兩天內(nèi)搞定！

這就是斯坦福、MIT等科研機(jī)構(gòu)推出的低秩線性轉(zhuǎn)換LoLCATs（Low-rank Linear Conversion with Attention Transfer）。

論文與代碼：https://github.com/HazyResearch/lolcats

應(yīng)用LoLCATs，可以實(shí)現(xiàn)傳統(tǒng)注意力（softmax）到線性注意力的無縫轉(zhuǎn)移，

且轉(zhuǎn)換后僅需開銷很低的微調(diào)（LoRA），0.2%的參數(shù)更新即可恢復(fù)精度，對(duì)比同類的線性注意力模型或方法， 5-shot MMLU直接提高了20分左右！

也就是說，在幾乎不損失Transformer大模型語言能力的基礎(chǔ)上，將LLM的計(jì)算復(fù)雜度從二次方降到了線性。

線性Attention一事，前人之述備矣，然則，能夠真正做大做強(qiáng)，還是第一次。

尤其具有實(shí)用價(jià)值的是，LoLCATs實(shí)現(xiàn)了極小的開銷和接近原始模型的性能。

LoLCATs的線性化轉(zhuǎn)換只需兩個(gè)步驟：

首先使用線性Attention的形式替換原始Attention部分，并利用簡(jiǎn)單的MSE損失訓(xùn)練新增的參數(shù)，以近似softmax注意力；

然后通過低成本的微調(diào)（LoRA）來進(jìn)一步提高模型的精度。

為了實(shí)現(xiàn)可擴(kuò)展性，作者采用更精細(xì)的「block by block」訓(xùn)練，將LLM的每k層看成一個(gè)block，盡在塊內(nèi)聯(lián)合訓(xùn)練注意力，以提高分層注意力匹配。

就如上圖所表示的那樣，一個(gè)羊駝（Llama）可以看成多個(gè)小刺猬疊在一起，每個(gè)小刺猬擁有獨(dú)特的用于線性化的參數(shù)，并且相互之間可以獨(dú)立訓(xùn)練。

LoLCATS 加速 LLM

為了避免昂貴的訓(xùn)練成本，研究者們一直在不斷探索兩個(gè)方面：

make models fast 與 create fast models

諸如Mamba、RWKV、TransNormer、Hawk、 Griffin和 StripedHyena等高效的subquadratic models不斷出現(xiàn)，

而關(guān)于將流行的LLM線性化的工作也讓我們眼前一亮。

但是線性化LLM往往伴隨著模型質(zhì)量的顯著降低，你甚至能通過MMLU的測(cè)試分?jǐn)?shù)猜出一個(gè)模型是不是傳統(tǒng)的Attention架構(gòu)，或者傳統(tǒng)Attention塊在模型中的占比。

另外，從實(shí)用的角度講，只有拿下了生產(chǎn)級(jí)別的大模型，線性化的道路才能真正與傳統(tǒng)Transformer平分秋色。

預(yù)備知識(shí)

先打基礎(chǔ)：為什么要線性化？

正常的softmax注意力可以表示為下圖上面的公式：

由于softmax的緣故，只能先算Q乘K，導(dǎo)致中間緩存和計(jì)算量隨序列長(zhǎng)度的平方增長(zhǎng)；

線性化就是設(shè)計(jì)倆函數(shù)來近似softmax，從而把公式轉(zhuǎn)化成下面的形式。

此時(shí)Q和K不需要綁在一起了，就可以先算K乘V，這個(gè)順序的改變導(dǎo)致中間緩存和計(jì)算量隨向量長(zhǎng)度的平方增長(zhǎng)，而相對(duì)于序列長(zhǎng)度是線性關(guān)系。

這就是線性化的意思，這樣的Attention也就不懼怕長(zhǎng)序列帶來的壓力了。

開始線性化

本文中，作者的主要想法是向線性化Transformer中添加三個(gè)簡(jiǎn)單的概念：

1. Learnable (Linear) Attentions：可學(xué)習(xí)的（線性）注意力

2. Low-rank Adaptation：低秩適配

3. Layer-wise Optimization：分層優(yōu)化

Learnable Attentions

首先訓(xùn)練線性注意力來模擬和替換softmax注意力。這種「注意力轉(zhuǎn)移」的靈感來自作者之前的一篇工作：Hedgehog。

論文地址：https://arxiv.org/pdf/2402.04347

如何設(shè)計(jì)設(shè)計(jì)精妙復(fù)雜的函數(shù)來近似softmax注意力？

作者表示：與其讓人類煞費(fèi)苦心，不如交給AI自己去學(xué)！

相比于Hedgehog中只使用可學(xué)習(xí)的線性注意力，作者在LoLCATs中，將其推廣為可學(xué)習(xí)的線性注意力和 + 滑動(dòng)窗口。

研究人員將線性和softmax注意力統(tǒng)一在一個(gè)層中，訓(xùn)練一些新增的參數(shù)以從整體上近似softmax注意力。

對(duì)于N個(gè)token的序列，前W個(gè)token用于計(jì)算softmax注意力，后N-W個(gè)token用于計(jì)算線性注意力，然后將這些值組合。

在Hedgehog中，作者通過KL散度來訓(xùn)練特征圖以匹配注意力權(quán)重，而本文改為在注意力層的輸出上使用MSE 失。

這繞過了Hedgehog的一個(gè)限制：需要將所有注意力權(quán)重實(shí)例化為監(jiān)督目標(biāo)。

相反，LoLCATs可以使用FlashAttention來計(jì)算softmax注意力輸出，并將線性化注意力的內(nèi)存消耗保持在O(N)。

只需將這些特征圖插入到每個(gè)現(xiàn)有的注意力中，即可創(chuàng)建線性化的 LLM。凍結(jié)所有其他權(quán)重，只訓(xùn)練這些特征圖，對(duì)于7B的LLM來說，只需要調(diào)整0.2%的參數(shù)。

Low-rank Adaptation

之前的線性化工作，通常需要一個(gè)比較昂貴的端到端訓(xùn)練階段。

但在LoLCATs這里，可以通過簡(jiǎn)單地將低秩適應(yīng)（LoRA）應(yīng)用于注意力的QKVO權(quán)重來恢復(fù)模型的性能。

凍結(jié)所有其他內(nèi)容，只訓(xùn)練LoRA權(quán)重，在某些自然語言數(shù)據(jù)上，最大限度地減少LLM輸出的next-token預(yù)測(cè)損失。

Layer-wise Optimization

大多數(shù)情況下，只需要以上兩步就搞定了。但對(duì)于像Llama 3.1 405B這種規(guī)模的模型來說，還需要努力一下。

通過簡(jiǎn)單地聯(lián)合優(yōu)化所有層，可以成功地線性化7B到70B參數(shù)范圍的LLM，但整體訓(xùn)練時(shí)，后面層的MSE會(huì)比前面的層更大。

當(dāng)模型變得更大更深時(shí)，MSE升級(jí)為了微調(diào)Llama 3.1 405B的真正問題。

為此，研究人員使用了更精細(xì)的逐塊訓(xùn)練，將Llama 3.1 405B分成多個(gè)k層塊，并僅在每個(gè)塊內(nèi)聯(lián)合訓(xùn)練注意力。

當(dāng)使用一些線性化數(shù)據(jù)并行訓(xùn)練所有模塊時(shí)，只需為每個(gè)塊預(yù)先計(jì)算LLM的隱藏狀態(tài)。

可以調(diào)節(jié)k來平衡并行訓(xùn)練的速度與預(yù)計(jì)算的內(nèi)存，并將隱藏狀態(tài)保存到磁盤。不需要花哨的成本模型，對(duì)于50M token的線性化來說：

k = 1時(shí)，需要2字節(jié) × 126層 × 50M token × 16384（hidden size）= 200TB的磁盤空間來存儲(chǔ)隱藏狀態(tài)。

而k = 9時(shí)，磁盤空間的需求將減少為22TB，這時(shí)仍然能在單個(gè)GPU上并行訓(xùn)練每個(gè)塊（9層）。

——后者顯然更友好一點(diǎn)，所以作者將Llama 3.1 405B的126層拆分為14個(gè)9層塊，在14個(gè)GPU上并行進(jìn)行注意力的線性化，過程僅需5個(gè)小時(shí)。然后用LoRA將它們?nèi)科唇釉谝黄?，就得到了最終模型。

實(shí)驗(yàn)結(jié)果

質(zhì)量恢復(fù)

下表給出了6個(gè)流行的LLM評(píng)估任務(wù)的結(jié)果。

與最近的一些線性化方法相比，LoLCATs顯著提高了不同任務(wù)和不同LLM的質(zhì)量和訓(xùn)練效率。

盡管只訓(xùn)練了0.2% 的模型參數(shù)（40M token），LoLCATs將線性化與原始模型的性能差距平均縮小了80%以上，token to model的效率提高了500~2500倍。

在7B這個(gè)量級(jí)上，LoLCATs優(yōu)于所有的線性注意力（包括RNN系列）模型：Mamba、RWKV、TransNormer、Hawk、 Griffin和 StripedHyena。

挑戰(zhàn)405B大模型

最后，作者使用LoLCATs將線性化擴(kuò)展到Llama 3.1 70B和更大的405B模型。

與之前的線性化方法相比，首先是質(zhì)量上的顯著改進(jìn)。通過控制相同的線性 + 滑動(dòng)窗口層，對(duì)于Llama 3.1 70B，在5-shot MMLU上的精度實(shí)現(xiàn)了39點(diǎn)的提升，對(duì)于Llama 3.1 405B，同樣實(shí)現(xiàn)了38.3分的改進(jìn)。

其次是訓(xùn)練效率的提高，在單個(gè)8x80GB H100上線性化Llama 3.1 70B僅需18個(gè)小時(shí)，而線性化Llama 3.1 405B所花費(fèi)的時(shí)間比之前用于8B模型的方法還要少。

參考資料：

https://x.com/simran_s_arora/status/1845909074774475125