今天小編分享的科學經驗：新注意力讓大模型上下文内存占用砍半！精度不減還能加速2倍，歡迎閱讀。
大模型同樣的上下文視窗，只需一半内存就能實現，而且精度無損？
前蘋果 ASIC 架構師 Nils Graef，和一名 UC 伯克利在讀大學生一起提出了新的注意力機制Slim Attention。
它以标準多頭注意力（MHA）為基準，對其中的 value 緩存處理過程進行了調整，實現了更少的内存占用。
具體來說，Slim Attention 既可以讓KV 緩存大小減半，也可以在 KV 緩存大小不變的情況下讓上下文翻倍，都不會帶來精度損失。
此外，在内存帶寬受限的場景下，它還可以将模型的推理過程加速1.5-2 倍。
網友評價，Slim Attention 雖然簡單，但卻是一個很酷的想法。
還有 AI 創業者評論說，這是一項重大突破，可能重塑對模型訓練和部署的看法。
K-Cache is All You Need
在标準的 MHA 機制當中，對于輸入 X 會通過線性變換，經由三個投影矩陣 W_Q、W_K、W_V 得到 Q ( query ) 、K ( key ) 和 V ( value ) 三個矩陣。
在推理階段，每個輸入 token 計算得到的 K 和 V 向量都需要緩存起來，形成 KV cache 供後續 token 計算時使用。
Slim Attention 的核心思路是，利用 MHA 中 W_K 和 W_V 通常都是方陣的性質，只存儲 K 而不直接存儲 V，然後實時利用 K 計算出 V。
△原始 MHA（左）與改進版（右）對比
在訓練階段，Slim Attention 與标準 MHA 一樣，會對輸入 X 計算 Q、K、V 三個矩陣，注意力計算和梯度回傳也與标準 MHA 完全一致。
在 W_K 可逆的前提下，Slim Attention 引入一個新的參數矩陣 W_KV：
W_KV = W_K^ ( -1 ) · W_V
據此，可以得到：
V = X · W_V = X · W_K · W_K^ ( -1 ) · W_V = K · W_KV
推理過程則主要分為兩個階段——提示階段（并行計算）和生成階段（自回歸）。
提示階段與标準 MHA 一樣，将輸入的所有 token 并行計算 Q、K 矩陣，但不同的是，這裡不直接計算 V，而是将中間結果 K 緩存供後續使用。
生成階段每個時間步生成一個新 token，首先計算該時間步的 Q 向量 q，然後基于 q 和之前時間步緩存的 K 矩陣，計算注意力得（即 softmax 的輸入）。
在 softmax 之前，Slim Attention 通過公式 V = K · W_KV 實時計算 V 矩陣。具體有兩種方式 :
直接計算 V，然後将 softmax 結果與 V 相乘（矩陣乘法）得到注意力輸出；
先将 softmax 結果與 K 相乘，然後再與 W_KV 相乘，當序列較長時這種方式更高效。
剩餘流程（殘差連接、前饋層等）與标準 MHA 一致，最後将當前步的 k 向量添加到 K 緩存中，供下一時間步使用。
總之，Slim Attention 是标準 MHA 的精确數學重寫，因此與近似方法不同，可确保準确率不會下降。
以此為前提，Slim Attention 實現了 KV 緩存減半或上下文翻倍的效果。
前蘋果架構師與 UC 伯克利大學生成果
Slim Attention 的作者是 AI 初創公司 OpenMachine 的創始人兼 CEO Nils Graef，以及 UC 伯克利在讀大學生Andrew Wasielewski。
Nils 的主業是機器學習加速器的架構和設計，曾發表兩篇 IEEE 期刊論文和 30 多項專利，引用次數超過 900 次。
創立 OpenMachine 前，Nils 在知名推理加速平台 Groq（注意不是馬斯克的 Grok）擔任芯片架構師。
更早的時候，他先後擔任過谷歌 ML 加速器架構 & 設計工程師和蘋果 ASIC 架構師。
Andrew Wasielewski 是 UC 伯克利在讀大學生，專業是物理和 EECs（電氣工程與計算機科學），預計将于明年畢業。
根據論文署名信息顯示，Slim Attention 的工作是 Andrew 在 OpenMachine 完成的。
去年 7 月，Nils 和 Andrew 還與其他人合作，發表了一篇名為Flash normalization的論文，提出了一種更快的 RNS 歸一化方式。
此外在 Slim Attention 的致謝中還提到，艾倫實驗室的 Dirk Groeneveld，以及 SGLang 三作謝志強，對其工作提供了有益讨論；Transformer 作者之一、Character.AI 創始人 Noam Shazeer 給出了積極反饋。
論文地址：
https://arxiv.org/abs/2503.05840
參考鏈接：
https://x.com/rohanpaul_ai/status/1901092052282339474
