Transformer模型的解碼速度一直是限制其應用的關鍵瓶頸。近日，韓國科學技術研究院、LG和DeepMind的研究人員聯合攻克了這個難題，他們提出了一個名為Block Transformer的新型架構，將解碼速度提升了驚人的10到20倍！這項突破性進展的核心在於巧妙地「切塊」了Transformer的注意力機制，有效解決了傳統Transformer中GPU利用率低下的問題，並大幅降低了記憶體開銷。
Transformer模型雖然強大，但它在解碼時的效率問題一直是個令人頭痛的事兒。不過，韓國科學技術研究院、LG和DeepMind的研究人員們，這次可是為我們帶來了一個驚喜——他們提出了一種新的Transformer架構，叫做Block Transformer，直接讓解碼速度提升了10到20倍!
這到底是怎麼做到的呢?原來，他們把Transformer的注意力機制給「切塊」了。這樣一來，原始Transformer那種每生成一個Token就要訪問一次全域KV快取的低效方式，就被徹底顛覆了。
研究人員分析了原始Transformer的短板:GPU的有效利用率竟然不到1%，剩下的99%都用在了記憶體存取。這顯然是不合理的，於是他們提出了Block Transformer，這個新架構透過區塊級注意力和區塊內注意力的分解，讓模型的推理吞吐量直接爆表。
具體來說，Block Transformer的工作流程是這樣的:先把序列進行切塊，然後用Embedder把每個塊轉換成嵌入向量。 Block Decoder負責處理區塊嵌入向量，捕捉區塊與區塊之間的全域依賴;Token Decoder則負責處理Token之間的局部依賴，產生Token序列。
這種方法不僅提升了推理速度，也大幅降低了記憶體開銷。有網友表示，自己之前也有類似的思路，但結果模型的效能不足，現在這個方法看起來確實有效削減了KV快取。
而且，Block Transformer在多個零樣本任務上的準確率與同等大小的原始Transformer相當甚至略高，這證明了它在提高效率的同時，並沒有犧牲品質。
這項研究的意義不止於此。它也降低了模型的訓練成本，全域注意力的二次記憶體存取開銷減少了16倍，GPU利用率也從1%提升到了44%。
論文網址：https://arxiv.org/abs/2406.02657
Block Transformer的出現，為Transformer模型的應用拓展了新的可能性，也為未來大型語言模型的效率最佳化提供了新的方向。其在速度和效率上的顯著提升，可望推動AI技術的進一步發展和應用。