今天小編分享的科學經驗：拆分Transformer注意力，韓國團隊讓大模型解碼提速20倍，歡迎閱讀。
只要将注意力切塊，就能讓大模型解碼提速 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
