拆分Transformer注意力，韓國團隊讓大模型解碼提速20倍
克雷西 發自 凹非寺
量子位 | 公眾號 QbitAI
只要將注意力切塊，就能讓大模型解碼提速20倍。
來自韓國科學技術研究院、LG和DeepMind的研究人員，提出了一種新的Transformer架構。
不僅獲得了更快的推理速度，內存開銷也大幅度下降。
研究人員詳細分析了原始Transformer推理速度慢的原因——
原始Transformer每生成一個Token就要訪問一次全局KV緩存，消耗了大量資源。
實際上，這種方法的GPU有效利用率不到1%，其餘的99%都用在了內存訪問上。
針對這一問題，團隊對Transformer的注意力機制進行了切塊調整，提出了名為Block Transformer的新架構。
結果在沒有明顯質量損失的情況下，推理吞吐量提升了10-20倍。
有網友表示，自己之前也有過類似的思路，但結果模型的性能不足，現在這個方法看上去確實有效削減了KV緩存。
“切開”Transformer的注意力
原始Transformer當中，對全局KV的頻繁訪問，導致計算複雜度高、內存佔用大，但推理吞吐量卻很低。
針對這一問題，作者的核心思路是將原始Transformer的全局注意力分解，分成塊級注意力和塊內注意力。
相應地，塊級注意力和塊內注意力分別由Block Decoder和Token Decoder進行處理。
具體的切塊數量根據總Token數和預設的塊大小決定，而塊大小的選擇，是全局和局部建模之間的平衡——
較大的塊可以減少塊的數量，從而降低Block Decoder的計算複雜度，但每個塊包含更多的token，可能影響局部依賴的建模能力；
較小的塊包含的Token更少，可以提高局部依賴的建模能力，但Block Decoder需要處理更多的塊，可能增加計算複雜度。
△不同塊大小的性能比較
工作流程上，Block Transformer拿到需要處理的序列之後，直接先進行切塊，然後利用Embedder將每個塊都轉換成一個嵌入向量。
具體來說，Embedder可以是一個簡單的查找表，將塊內的token映射為對應的嵌入向量，然後將這些嵌入向量拼接或累加得到塊嵌入向量。
完成塊的向量化之後，Block Decoder接收Embedder生成的塊嵌入向量序列作為輸入。
在其每個自注意力層中，都會對塊嵌入向量序列進行自注意力計算，捕捉塊與塊之間的全局依賴關係。
經過多個自注意力層的處理，塊嵌入向量融合了全局上下文信息，所以，Block Decoder的輸出是一個全局上下文感知的塊嵌入向量序列。
完成塊級處理之後，Block Decoder的輸出會與塊內已生成的Token向量一起被Token Decoder接收。
在Token Decoder中，塊嵌入向量首先被轉換為與Token嵌入向量相同維度的向量，然後在Token Decoder的多個自注意力層中進行處理，捕捉Token之間的局部依賴關係。
經過多個自注意力層的處理，Token嵌入向量融合了局部上下文信息和來自塊嵌入向量的全局信息。
最終，Token Decoder的輸出是一個包含了局部上下文感知的Token嵌入向量序列，用於生成當前塊的Token，Token Decoder重複這個過程，直到生成當前塊的所有token。
回到整體上，Block Transformer通過交替執行塊級自迴歸建模和塊內自迴歸解碼，迭代生成整個輸出序列。
比如在生成第i個塊時，Block Decoder會根據前i-1個塊的嵌入向量預測第i個塊的嵌入向量，然後Token Decoder根據第i個塊的嵌入向量和已生成的Token，生成第i個塊的Token序列。
這個過程重複進行，直到生成整個輸出序列。
推理吞吐量最高提升20倍
對注意力的切塊帶來的效果立竿見影，模型的推理吞吐量直接提升了10-20倍。
例如，在decode-heavy設置下，85M參數的Block Transformer吞吐量達到了每秒13.5萬Tokens，而同等大小的原始Transformer僅有約6千Tokens。
針對更長的提示詞，Block Transformer同樣具有吞吐量優勢——在提示詞長度為8K的情況下，Block Transformer的吞吐量超過了提示詞長度為2K的原始Transformer。
吞吐量的提升並沒有讓質量下降，在HellaSwag、PIQA和ARC-easy等多個零樣本任務上，Block Transformer的準確率與同等大小的原始Transformer相當甚至略高。
進一步探究結果表明，Block Transformer這種全局-局部建模方式能在提高推理效率的同時保持較低的訓練損失（圖a）。
同時這種方法還能有效利用全局上下文，在PG19測試集上，取得了與原始Transformer相似的位置損失（圖b）。
另外，在相同的訓練計算量和推理吞吐量預算下，Block Transformer能達到比原始Transformer更低的訓練損失，展現出了優異的訓練效率（圖c）。
除了帶來性能提升之外，Block Transformer也降低了模型的訓練成本。
使用其默認的4個Token的塊長度，全局注意力的二次內存訪問開銷減少了16倍。
反覆讀取KV緩存帶來的內存開銷也幾乎消除，1%的GPU利用率提升到了44%。
論文地址：
https://arxiv.org/abs/2406.02657
— 完 —
量子位年度AI主題策劃正在徵集中！
歡迎投稿專題 一千零一個AI應用，365行AI落地方案
或與我們分享你在尋找的AI產品，或發現的AI新動向
點這裡👇關注我，記得標星哦～
