只要將注意力切塊，就能讓大模型解碼提速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