新型注意力機制TPA，姚期智院士團隊打造。

TPA對每個token做動態的張量分解，不存儲完整的靜態KV，而是保留分解的版本，內存占用節省90%（或者更多），而不會犧牲性能。

論文中還證明了流行的MHA、MQA、GQA都是TPA的特殊情況，用一個框架統一了現代注意力設計。

用此方法訓練的新模型T6，代碼已在GitHub開源。

論文發布后，有創業者表示，終于不用付那么多錢給云廠商了。

也有研究者認為，論文中的實驗看起來很有希望，不過實驗中的模型規模有點小，希望看到更多結果。

動態張量分解，無縫集成RoPE

盡管現有的注意力機制在眾多任務中取得了不錯的效果，但它還是有計算和內存開銷大的缺陷。

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

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

新方法在注意力計算過程中對QKV做分解。

與LoRA系列低秩分解方法相比，TPA將QKV分別構造為與上下文相關的分解張量，實現動態適應。

通過只緩存分解后的秩，設置合適的參數可使內存占用降低90%或以上。

對于流行的RoPE位置編碼，TPA可以與之無縫集成，實現以較低的成本旋轉分解KV，無需進行復雜的調整。

在實驗中，使用FineWeb-Edu 100B數據集訓練模型，TPA與其他注意力設計相比始終保持較低的困惑度。

在ARC、BoolQ、HellaSwag和MMLU等基準測試中測試了零樣本和少樣本性能。TPA和TPA-KVonly在大多數任務中都優于或匹配所有基線。

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

此外還有來自心動網絡Taptap的Qin Zhen。

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