只要將注意力切塊，就能讓大模型解碼提速20倍。

來自韓國科學(xué)技術(shù)研究院、LG和DeepMind的研究人員，提出了一種新的Transformer架構(gòu)。

不僅獲得了更快的推理速度，內(nèi)存開銷也大幅度下降。

研究人員詳細分析了原始Transformer推理速度慢的原因——

原始Transformer每生成一個Token就要訪問一次全局KV緩存，消耗了大量資源。

實際上，這種方法的GPU有效利用率不到1%，其余的99%都用在了內(nèi)存訪問上。

針對這一問題，團隊對Transformer的注意力機制進行了切塊調(diào)整，提出了名為Block Transformer的新架構(gòu)。

結(jié)果在沒有明顯質(zhì)量損失的情況下，推理吞吐量提升了10-20倍。

有網(wǎng)友表示，自己之前也有過類似的思路，但結(jié)果模型的性能不足，現(xiàn)在這個方法看上去確實有效削減了KV緩存。

“切開”Transformer的注意力

原始Transformer當(dāng)中，對全局KV的頻繁訪問，導(dǎo)致計算復(fù)雜度高、內(nèi)存占用大，但推理吞吐量卻很低。

針對這一問題，作者的核心思路是將原始Transformer的全局注意力分解，分成塊級注意力和塊內(nèi)注意力。

相應(yīng)地，塊級注意力和塊內(nèi)注意力分別由Block Decoder和Token Decoder進行處理。

具體的切塊數(shù)量根據(jù)總Token數(shù)和預(yù)設(shè)的塊大小決定，而塊大小的選擇，是全局和局部建模之間的平衡——

• 較大的塊可以減少塊的數(shù)量，從而降低Block Decoder的計算復(fù)雜度，但每個塊包含更多的token，可能影響局部依賴的建模能力；
• 較小的塊包含的Token更少，可以提高局部依賴的建模能力，但Block Decoder需要處理更多的塊，可能增加計算復(fù)雜度。

△不同塊大小的性能比較

工作流程上，Block Transformer拿到需要處理的序列之后，直接先進行切塊，然后利用Embedder將每個塊都轉(zhuǎn)換成一個嵌入向量。

具體來說，Embedder可以是一個簡單的查找表，將塊內(nèi)的token映射為對應(yīng)的嵌入向量，然后將這些嵌入向量拼接或累加得到塊嵌入向量。

完成塊的向量化之后，Block Decoder接收Embedder生成的塊嵌入向量序列作為輸入。

在其每個自注意力層中，都會對塊嵌入向量序列進行自注意力計算，捕捉塊與塊之間的全局依賴關(guān)系。

經(jīng)過多個自注意力層的處理，塊嵌入向量融合了全局上下文信息，所以，Block Decoder的輸出是一個全局上下文感知的塊嵌入向量序列。

完成塊級處理之后，Block Decoder的輸出會與塊內(nèi)已生成的Token向量一起被Token Decoder接收。

在Token Decoder中，塊嵌入向量首先被轉(zhuǎn)換為與Token嵌入向量相同維度的向量，然后在Token Decoder的多個自注意力層中進行處理，捕捉Token之間的局部依賴關(guān)系。

經(jīng)過多個自注意力層的處理，Token嵌入向量融合了局部上下文信息和來自塊嵌入向量的全局信息。

最終，Token Decoder的輸出是一個包含了局部上下文感知的Token嵌入向量序列，用于生成當(dāng)前塊的Token，Token Decoder重復(fù)這個過程，直到生成當(dāng)前塊的所有token。

回到整體上，Block Transformer通過交替執(zhí)行塊級自回歸建模和塊內(nèi)自回歸解碼，迭代生成整個輸出序列。

比如在生成第i個塊時，Block Decoder會根據(jù)前i-1個塊的嵌入向量預(yù)測第i個塊的嵌入向量，然后Token Decoder根據(jù)第i個塊的嵌入向量和已生成的Token，生成第i個塊的Token序列。

這個過程重復(fù)進行，直到生成整個輸出序列。

推理吞吐量最高提升20倍

對注意力的切塊帶來的效果立竿見影，模型的推理吞吐量直接提升了10-20倍。

例如，在decode-heavy設(shè)置下，85M參數(shù)的Block Transformer吞吐量達到了每秒13.5萬Tokens，而同等大小的原始Transformer僅有約6千Tokens。

針對更長的提示詞，Block Transformer同樣具有吞吐量優(yōu)勢——在提示詞長度為8K的情況下，Block Transformer的吞吐量超過了提示詞長度為2K的原始Transformer。

吞吐量的提升并沒有讓質(zhì)量下降，在HellaSwag、PIQA和ARC-easy等多個零樣本任務(wù)上，Block Transformer的準(zhǔn)確率與同等大小的原始Transformer相當(dāng)甚至略高。

進一步探究結(jié)果表明，Block Transformer這種全局-局部建模方式能在提高推理效率的同時保持較低的訓(xùn)練損失（圖a）。

同時這種方法還能有效利用全局上下文，在PG19測試集上，取得了與原始Transformer相似的位置損失（圖b）。

另外，在相同的訓(xùn)練計算量和推理吞吐量預(yù)算下，Block Transformer能達到比原始Transformer更低的訓(xùn)練損失，展現(xiàn)出了優(yōu)異的訓(xùn)練效率（圖c）。

除了帶來性能提升之外，Block Transformer也降低了模型的訓(xùn)練成本。

使用其默認的4個Token的塊長度，全局注意力的二次內(nèi)存訪問開銷減少了16倍。

反復(fù)讀取KV緩存帶來的內(nèi)存開銷也幾乎消除，1%的GPU利用率提升到了44%。

論文地址：https://arxiv.org/abs/2406.02657