時空壓縮!劍橋大學提出注意力機制MTLA:推理加速5倍,顯存減至1/8
相鄰的 KV 緩存將合并為一個。在生成第一個字符時,KV 緩存長度為 1;生成第二個字符后,新生成的 KV 與前一個被合并,KV 緩存長度仍然保持為 1。這種動態合并機制有效壓縮了時間維度上的冗余信息。
然而,這也帶來了并行訓練上的挑戰:雖然兩個時間步的 KV 緩存長度相同,但它們所包含的信息不同,若不加以區分,容易導致訓練與推理行為不一致。
MTLA 通過一種優雅的方式解決了這一問題。正如下圖所示,在訓練階段,MTLA 保留了所有中間狀態的 KV 表達,并引入了步幅感知因果掩碼(stride-aware causal mask),確保每個 query 在訓練時訪問到與推理階段一致的 KV 區域,從而準確模擬增量推理中的注意力行為。
得益于這一設計,MTLA 能夠像標準注意力機制一樣通過矩陣乘法實現高效并行計算,在保持訓練效率的同時完成對時間維度的壓縮。
此外,MTLA 還引入了解耦的旋轉位置編碼(decoupled RoPE)來建模位置信息,并對其進行了時間維度上的壓縮,進一步提升了整體效率。
值得強調的是,MTLA 不僅是一種更高效的自注意力機制,它還具備極強的靈活性與可調性。例如,當將時間壓縮率 s 設置得足夠大時,MTLA 在推理過程中幾乎只保留一個 KV 緩存,這種形式本質上就退化為一種線性序列建模方法。換句話說,線性序列建模可以被視為 MTLA 的極端情況,MTLA 在注意力機制與線性模型之間架起了一座橋梁。
然而,在許多復雜任務中,傳統注意力機制所具備的二次計算復雜度雖然代價高昂,卻提供了更強的建模能力。因此,MTLA 所引入的 “可調時間壓縮率 s” 這一設計思路,恰恰為模型提供了一個在效率與性能之間靈活權衡的可能空間。
MTLA 的卓越性能
MTLA 在一系列任務中展現了出色的性能,包括語音翻譯,文本摘要生成,語音識別和口語理解。例如在語音翻譯中,MTLA 在保持與標準 MHA 相當的翻譯質量的同時,實現了超過 5 倍的推理速度提升,并將推理過程中的 GPU 顯存占用降低了超過 8 倍。
值得注意的是,僅當時間壓縮率 s=2 時,MTLA 對 KV 緩存的壓縮程度就已經與 MQA 相當,且在模型性能上更具優勢。而相比之下,MQA 所采用的減少 KV 頭數量的方法已達上限,而 MTLA 還有進一步的空間。
未來發展
MTLA 具備在大規模場景中部署的顯著潛力,尤其是在大語言模型參數規模不斷擴大、以及思維鏈等技術推動下生成序列日益增長的背景下,對 KV 緩存進行時空壓縮正是緩解推理開銷的關鍵手段。在這樣的趨勢下,MTLA 有望成為未來大語言模型中自注意力模塊的重要替代方案。
當然,與 DeepSeek 提出的 MLA 類似,MTLA 相較于 GQA 和 MQA,在工程落地方面的改動不再是簡單的一兩行代碼可以實現的優化。這也意味著要將其大規模應用到現有 LLM 框架中,還需要來自社區的持續推動與協同開發。
為促進這一過程,MTLA 的實現代碼已全面開源,希望能夠為研究者與工程實踐者提供便利,共同推動高效注意力機制在大模型時代的落地與普及。