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

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

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

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

原始Transformer每生成一個(gè)Token就要訪(fǎng)問(wèn)一次全局KV緩存，消耗了大量資源。

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

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

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

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

“切開(kāi)”Transformer的注意力

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

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

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

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

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

△不同塊大小的性能比較

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

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

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

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

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

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

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

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

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

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

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

這個(gè)過(guò)程重復(fù)進(jìn)行，直到生成整個(gè)輸出序列。

推理吞吐量最高提升20倍

對(duì)注意力的切塊帶來(lái)的效果立竿見(jiàn)影，模型的推理吞吐量直接提升了10-20倍。

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

針對(duì)更長(zhǎng)的提示詞，Block Transformer同樣具有吞吐量?jī)?yōu)勢(shì)——在提示詞長(zhǎng)度為8K的情況下，Block Transformer的吞吐量超過(guò)了提示詞長(zhǎng)度為2K的原始Transformer。

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

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

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

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

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

使用其默認(rèn)的4個(gè)Token的塊長(zhǎng)度，全局注意力的二次內(nèi)存訪(fǎng)問(wèn)開(kāi)銷(xiāo)減少了16倍。

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

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