（MoBA）大語言模型長文本處理新解法：塊注意力混合

在自然語言處理領域，高效處理長文本一直是個老大難問題。隨著大語言模型在閱讀、理解和生成文本方面的能力不斷提升，其處理輸入的核心——注意力機制，卻逐漸成了瓶頸。在典型的Transformer架構里，這種機制要把每個詞元（token）和其他所有詞元進行比較，這就導致計算成本會隨著序列長度的增加呈二次方增長。如今，我們把語言模型應用到長篇文檔、多章節書籍、法律文書或是大型代碼庫這類需要處理海量文本信息的任務中時，這個問題就更加突出了。要是模型得處理幾萬甚至幾十萬個詞元，簡單粗暴地計算全注意力的成本實在高得離譜。

以往解決這個問題的方法，常常是設置固定結構或者采用近似算法，但這些在某些場景下可能會影響效果。比如說，滑動窗口機制只能讓詞元關注局部區域，這樣就容易忽略掉重要的全局關系。還有些方法，像直接用全新結構替換softmax注意力機制，從根本上改變了基礎架構，可這就需要從頭開始大規模重新訓練模型，很難利用現有的預訓練模型。所以，研究人員一直都在尋找一種新方法，既能保留原始Transformer設計的優勢，也就是適應性強、能捕捉廣泛依賴關系，又不會在處理超長序列時，產生傳統全注意力機制帶來的巨額計算開銷。

來自Moonshot AI、清華大學和浙江大學的研究人員，帶來了塊注意力混合（MoBA）這個創新方法，它把專家混合（MoE）的原理運用到了注意力機制當中。MoBA會把輸入劃分成一個個便于處理的“塊”，再通過一個可訓練的門控系統，來確定每個查詢詞元應該關注哪些塊，這樣就解決了模型挨個比較詞元時效率低下的問題。和那些強制使用局部注意力或者窗口注意力的方法不同，MoBA能讓模型自己學習該重點關注哪里。這個設計遵循“less structure”原則，意思就是架構不會提前規定哪些詞元必須相互作用，而是把這些決策交給訓練好的門控網絡。

MoBA有個很關鍵的特點，就是能和現有的基于Transformer的模型完美配合。它沒有拋棄標準的自注意力接口，而是像個“插件”一樣，直接就能替換使用。MoBA的參數數量和原模型一樣，不會讓架構變得臃腫，而且還保留了因果掩碼，保證自回歸生成的準確性。在實際應用中，MoBA可以在稀疏注意力和全注意力之間靈活切換。處理超長輸入時，用稀疏注意力能提高速度；在訓練的某些層或者階段，如果有需要，還能切換回標準的全注意力模式。

可訓練的區塊稀疏注意力

完整的上下文被劃分成多個區塊，每個查詢詞元都能學會關注最相關的鍵值（KV）區塊，這樣就能高效處理長序列數據。

無參數門控機制

MoBA引入了一種很新穎的無參數top - k門控機制，專門為每個查詢詞元挑選最相關的區塊，確保模型只聚焦在信息量最大的區塊上。

在完全注意力和稀疏注意力之間無縫過渡

MoBA設計得非常靈活，可以完美替代全注意力機制，在全注意力模式和稀疏注意力模式之間自由切換。

技術細節和優勢

MoBA的核心操作，就是把上下文劃分成多個區塊，每個區塊都包含一連串連續的詞元。門控機制會計算查詢詞元和每個區塊之間的“親和度”分數，一般是通過把查詢和區塊里鍵的聚合表示進行對比來實現。然后，它會選出得分最高的那些區塊。這樣一來，最終的注意力分配就只和最相關區塊里的詞元有關。而且，包含查詢詞元本身的那個區塊一定會被納入，保證局部上下文信息隨時都能獲取到。同時，MoBA還采用了因果掩碼，讓詞元不會關注未來的位置，維持從左到右的自回歸特性。

經過這樣的處理，MoBA的注意力矩陣比原始Transformer的要稀疏得多。但它依然很靈活，要是有需要，查詢也能獲取到遠距離的信息。舉個例子，要是在文本末尾提出的問題，得參考文本開頭的細節才能回答，門控機制就能給開頭相關的區塊打高分。從技術層面來講，這種基于區塊的方法把詞元比較的數量減少到了亞二次方級別，隨著上下文長度增加到幾十萬甚至上百萬個詞元，效率提升就更加明顯了。

MoBA還有個很大的優勢，就是和現代加速器以及專用內核的兼容性很好。研究人員把MoBA和FlashAttention結合起來，FlashAttention是一個高性能庫，能實現快速、低內存消耗的精確注意力計算。他們根據所選的區塊，仔細對查詢 - 鍵 - 值操作進行分組，進一步優化了計算過程。研究人員表示，處理一百萬個詞元時，MoBA相比傳統的全注意力機制，速度能提升大約6倍，這在實際應用場景里優勢非常大。

結果和洞察

根據技術報告，MoBA在各種任務中的表現和全注意力機制不相上下，但處理長序列時，計算成本卻低得多。在語言建模數據測試里，當序列長度達到8192或者32768個詞元時，MoBA的困惑度和全注意力Transformer的很接近。更重要的是，當研究人員把上下文長度逐漸增加到128000甚至更長時，MoBA對長上下文的理解能力依然很強。研究人員還做了“trailing token”評估，主要測試模型預測長提示末尾詞元的能力，這個測試通常能暴露那些過度依賴近似算法的方法的弱點。而MoBA在處理這些末尾位置的詞元時，預測質量并沒有明顯下降。

研究人員還研究了MoBA對區塊大小和門控策略的敏感度。有些實驗發現，細化粒度（也就是用更小的區塊，但選更多的區塊）能讓模型的效果更接近全注意力機制。就算MoBA忽略了大部分上下文，自適應門控也能找到對查詢真正重要的區塊。此外，還有一種“混合”模式，采取了平衡策略：一部分層繼續用MoBA來提高速度，少數層則切換回全注意力模式。在進行有監督微調時，這種混合模式特別有用，因為訓練目標可能會屏蔽掉輸入里的某些位置。在少數上層保留全注意力，模型就能覆蓋更廣泛的上下文，對那些需要全局視角的任務很有幫助。

總的來說，這些研究結果表明，MoBA特別適合處理需要大量上下文的任務，像是長篇文檔閱讀理解、大規模代碼補全，還有需要參考完整對話歷史的多輪對話系統。MoBA能有效提升效率，而且對性能的影響很小，是讓大語言模型在大規模應用中更高效的理想選擇。

結論

塊注意力混合（MoBA）為大語言模型處理長文本提供了一種更高效的方法，而且不用大幅修改Transformer架構，也不會降低性能。通過在注意力模塊里融入專家混合的理念，MoBA用一種可學習的稀疏方式，讓模型聚焦在超長輸入的關鍵部分。它的設計非常靈活，尤其是能在稀疏注意力和全注意力之間無縫切換，這對正在進行的和未來的訓練流程都很有吸引力。研究人員可以靈活調整修剪注意力模式的程度，要是任務需要全面覆蓋信息，也可以有選擇地使用全注意力模式。

雖然目前對MoBA的研究主要集中在文本領域，但它的底層機制在其他數據模態上可能也有很大潛力。只要序列長度長到會引發計算或者內存方面的問題，用MoBA這種把查詢分配給不同區塊“專家”的思路，就能在不影響處理全局依賴關系的前提下，有效緩解瓶頸。隨著語言應用中的序列長度不斷增加，像MoBA這樣的方法，在提升神經語言模型的可擴展性和性價比方面，可能會發揮至關重要的作用。

鏈接

github鏈接：https://github.com/MoonshotAI/MoBA論文：  https://arxiv.org/html/2406.14909v1

本文轉載自 ??柏企科技圈??，作者：柏企