Transformer模型雖然強大，但它在解碼時的效率問題一直是讓人頭疼的事兒。不過，韓國科學技術研究院、LG和DeepMind的研究人員們，這次可是給我們帶來了一個驚喜——他們提出了一種新的Transformer架構，叫做Block Transformer，直接讓解碼速度提升了10到20倍!

這到底是怎麼做到的呢?原來，他們把Transformer的注意力機制給“切塊”了。這樣一來，原始Transformer那種每生成一個Token就要訪問一次全局KV緩存的低效方式，就被徹底顛覆了。

image.png

研究人員們分析了原始Transformer的短板:GPU的有效利用率竟然不到1%，剩下的99%都用在了內存訪問上。這顯然是不合理的，於是他們提出了Block Transformer，這個新架構通過塊級注意力和塊內注意力的分解，讓模型的推理吞吐量直接爆表。

具體來說，Block Transformer的工作流程是這樣的:先把序列進行切塊，然後用Embedder把每個塊轉換成嵌入向量。Block Decoder負責處理塊嵌入向量，捕捉塊與塊之間的全局依賴;Token Decoder則負責處理Token之間的局部依賴，生成Token序列。

image.png

這種方法不僅提升了推理速度，還大幅度降低了內存開銷。有網友表示，自己之前也有過類似的思路，但結果模型的性能不足，現在這個方法看上去確實有效削減了KV緩存。

而且，Block Transformer在多個零樣本任務上的準確率與同等大小的原始Transformer相當甚至略高，這證明了它在提高效率的同時，並沒有犧牲質量。

這項研究的意義可不止於此。它還降低了模型的訓練成本，全局注意力的二次內存訪問開銷減少了16倍，GPU利用率也從1%提升到了44%。

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