月球計畫 AI 研究介紹混合區塊注意力 (MoBA)：一種將混合專家 (MoE) 原則應用於注意力機制的新 AI 方法

有效處理長文本的挑戰

在自然語言處理中，有效處理長文本一直是一個挑戰。隨著大型語言模型的能力不斷增強，它們能夠閱讀、理解和生成文本，但注意力機制——這是它們處理輸入的核心——可能成為瓶頸。在典型的變壓器架構中，這個機制需要將每個字元與其他所有字元進行比較，這會導致計算成本隨著序列長度的增加而成平方增長。這個問題在需要查閱大量文本信息的任務中變得更加迫切，例如長篇文件、多章書籍、法律文件或大型代碼庫。當模型必須處理數十萬甚至數十萬字元時，簡單計算完整的注意力會變得非常昂貴。

以往的解決方案

以往的解決方案通常依賴於固定的結構或近似方法，這可能在某些情況下影響質量。例如，滑動窗口機制將字元限制在局部範圍內，這可能會遮蔽重要的全局關係。另一方面，徹底改變基本架構的方法——例如用全新的結構替代softmax注意力——可能需要從頭開始進行大量重訓練，這使得利用現有的預訓練模型變得困難。研究人員一直在尋找一種方法，能夠保持原始變壓器設計的關鍵優勢——其適應性和捕捉廣泛依賴的能力——而不會產生與傳統完整注意力在極長序列上相關的巨大計算開銷。

Mixture of Block Attention (MoBA)

來自Moonshot AI、清華大學和浙江大學的研究人員介紹了Mixture of Block Attention (MoBA)，這是一種創新的方法，將專家混合（Mixture of Experts, MoE）的原則應用於注意力機制。MoBA通過將輸入劃分為可管理的“區塊”，並使用可訓練的門控系統來決定哪些區塊對每個查詢字元是相關的，從而解決了模型在比較每個字元與其他字元時所產生的低效率。與強制執行局部或窗口注意力的方法不同，MoBA允許模型學習集中注意力的地方。這種設計遵循“較少結構”的原則，意味著架構不會預先定義哪些字元應該互動，而是將這些決策委託給學習到的門控網絡。

MoBA的特點

MoBA的一個關鍵特點是它能夠與現有的基於變壓器的模型無縫運作。MoBA不是丟棄標準的自注意力接口，而是一種“插件”或替代品。它保持相同的參數數量，因此不會使架構膨脹，並保留因果掩碼以確保自回歸生成的正確性。在實際應用中，MoBA可以在稀疏和完整注意力之間切換，使模型在處理極長的輸入時受益於速度提升，同時在需要時保留回退到標準完整注意力的能力。

技術細節與優勢

MoBA的核心在於將上下文劃分為區塊，每個區塊涵蓋一系列連續的字元。門控機制計算查詢字元與每個區塊之間的“親和力”分數，通常通過將查詢與區塊鍵的聚合表示進行比較來實現。然後，它選擇得分最高的區塊。因此，只有最相關區塊中的字元會對最終的注意力分佈產生影響。包含查詢本身的區塊始終被納入，確保局部上下文保持可訪問。同時，強制執行因果掩碼，以便字元不會關注未來的位置，保持從左到右的自回歸特性。

由於這一過程，MoBA的注意力矩陣比原始變壓器的要稀疏得多。然而，它仍然足夠靈活，允許查詢在需要時關注遙遠的信息。例如，如果在文本末尾提出的問題只能通過參考開頭的細節來回答，門控機制可以學習為相關的早期區塊分配高分。從技術上講，這種基於區塊的方法將字元比較的次數減少到次平方級別，帶來的效率提升在上下文長度達到數十萬甚至數百萬字元時尤為明顯。

結果與見解

根據技術報告，MoBA在各種任務中的表現與完整注意力相當，同時在處理長序列時提供了顯著的計算節省。在語言建模數據上的測試顯示，MoBA的困惑度在8,192或32,768字元的序列長度下依然接近完整注意力變壓器的水平。關鍵是，隨著研究人員逐漸延長上下文長度到128,000及以上，MoBA仍然保持強大的長上下文理解能力。作者提出了“尾隨字元”評估，專注於模型在長提示結尾附近預測字元的能力——這通常突顯出依賴重近似方法的弱點。MoBA能夠有效管理這些尾隨位置，幾乎沒有預測質量的重大損失。

結論

總結來說，Mixture of Block Attention (MoBA)為大型語言模型提供了一條更高效的長上下文處理途徑，而無需對變壓器架構進行大規模改造或性能下降。通過在注意力模塊中採用專家混合的理念，MoBA提供了一種可學習但稀疏的方式來專注於非常長的輸入的相關部分。其設計的適應性——特別是稀疏和完整注意力之間的無縫切換——使其在持續或未來的訓練流程中尤其具有吸引力。研究人員可以微調如何積極地修剪注意力模式，或選擇性地在需要全面覆蓋的任務中使用完整注意力。

雖然MoBA的注意力主要集中在文本上下文上，但其底層機制也可能對其他數據類型具有潛力。無論序列長度多大，若引發計算或內存問題，將查詢分配給區塊專家的概念都可以緩解瓶頸，同時保持處理重要全局依賴的能力。隨著語言應用中的序列長度不斷增長，像MoBA這樣的方法可能在推進神經語言建模的可擴展性和成本效益方面發揮關鍵作用。