姚期智團(tuán)隊(duì)開(kāi)源新型注意力，節(jié)省90%內(nèi)存不降性能，一個(gè)框架統(tǒng)一MHA/MQA/GQA

夢(mèng)晨 2025-01-15 17:08:58 來(lái)源：量子位

夢(mèng)晨 發(fā)自 凹非寺
量子位 | 公眾號(hào) QbitAI

新型注意力機(jī)制TPA，姚期智院士團(tuán)隊(duì)打造。

TPA對(duì)每個(gè)token做動(dòng)態(tài)的張量分解，不存儲(chǔ)完整的靜態(tài)KV，而是保留分解的版本，內(nèi)存占用節(jié)省90%（或者更多），而不會(huì)犧牲性能。

論文中還證明了流行的MHA、MQA、GQA都是TPA的特殊情況，用一個(gè)框架統(tǒng)一了現(xiàn)代注意力設(shè)計(jì)。

用此方法訓(xùn)練的新模型T6，代碼已在GitHub開(kāi)源。

論文發(fā)布后，有創(chuàng)業(yè)者表示，終于不用付那么多錢給云廠商了。

也有研究者認(rèn)為，論文中的實(shí)驗(yàn)看起來(lái)很有希望，不過(guò)實(shí)驗(yàn)中的模型規(guī)模有點(diǎn)小，希望看到更多結(jié)果。

動(dòng)態(tài)張量分解，無(wú)縫集成RoPE

盡管現(xiàn)有的注意力機(jī)制在眾多任務(wù)中取得了不錯(cuò)的效果，但它還是有計(jì)算和內(nèi)存開(kāi)銷大的缺陷。

DeepSeek-v2中提出的MLA壓縮了KV緩存，但與RoPE位置編碼不兼容，每個(gè)注意力頭需要額外的位置編碼參數(shù)。

為了克服這些方法的局限性，團(tuán)隊(duì)提出張量積注意力（TPA，Tensor Product Attention）。

新方法在注意力計(jì)算過(guò)程中對(duì)QKV做分解。

與LoRA系列低秩分解方法相比，TPA將QKV分別構(gòu)造為與上下文相關(guān)的分解張量，實(shí)現(xiàn)動(dòng)態(tài)適應(yīng)。

通過(guò)只緩存分解后的秩，設(shè)置合適的參數(shù)可使內(nèi)存占用降低90%或以上。

對(duì)于流行的RoPE位置編碼，TPA可以與之無(wú)縫集成，實(shí)現(xiàn)以較低的成本旋轉(zhuǎn)分解KV，無(wú)需進(jìn)行復(fù)雜的調(diào)整。

在實(shí)驗(yàn)中，使用FineWeb-Edu 100B數(shù)據(jù)集訓(xùn)練模型，TPA與其他注意力設(shè)計(jì)相比始終保持較低的困惑度。

在ARC、BoolQ、HellaSwag和MMLU等基準(zhǔn)測(cè)試中測(cè)試了零樣本和少樣本性能。TPA和TPA-KVonly在大多數(shù)任務(wù)中都優(yōu)于或匹配所有基線。

論文由清華&上海期智研究員團(tuán)隊(duì)、UCLA顧全全團(tuán)隊(duì)合作，共同一作為清華博士生張伊凡與姚班校友、現(xiàn)UCLA博士生劉益楓。

此外還有來(lái)自心動(dòng)網(wǎng)絡(luò)Taptap的Qin Zhen。

論文地址：
https://arxiv.org/abs/2501.06425
開(kāi)源代碼：
https://github.com/tensorgi/T6

參考鏈接：
[1]https://x.com/yifan_zhang_/status/1879049477681741896