Kimi新論文再次“撞車”DeepSeek，都談到了長文注意力機制

作者：量子位 2025-02-19 15:30:00

MoBA將MoE（專家混合）應用于注意力機制，通過遵循一種“less structure” 原則，允許模型自主決定關注哪些區域或位置。

Kimi背后的長上下文處理機制曝光了！

這項名為MoBA的新型注意力機制，能將處理1M長文本的速度一下子提升6.5倍，而且還是經過Kimi平臺實際驗證的那種。

概括而言，這項耗時一年半的工作主要看點在：

一言以蔽之，MoBA將MoE（專家混合）應用于注意力機制，通過遵循一種“less structure” 原則，允許模型自主決定關注哪些區域或位置。

最終結果是，在保持模型性能不變的前提下，將處理1M和10M長文本的速度分別提升了6.5倍和16倍。

而且按照論文得出的結論：

MoBA開箱即用，它能輕松融入現有模型，無需高昂訓練成本。

另外，在深扒論文細節的同時，我們也在作者名單中發現了楊植麟本人的身影。

與此同時，再次令人感到戲劇的是，Kimi的這項研究又又又和DeepSeek撞車了？？

差不多只比Kimi早5個小時，DeepSeek也公開了自己的注意力機制NSA（量子位今日另一篇文章有詳細介紹）。

好家伙，誰能想到這熟悉的一幕也才剛過去不久！（先賣個關子）

將MoE原理應用于注意力機制

還是先來看論文細節。

首先，團隊回答了為什么要開展這項研究？

隨著大語言模型（LLMs）向AGI發展，處理長文本的能力將更加關鍵。然而，傳統注意力機制的計算復雜度隨著序列長度的增加而呈平方級增長，這一特性嚴重阻礙了模型對長序列的高效處理。

而且，一些現有方法都存在或多或少的局限性：

像滑動窗口注意力機制，通過限制模型關注窗口內局部信息來減少計算量，但窗口設定依賴特定任務，缺乏通用性，嚴重限制模型泛化能力；
再如Quest、Minference和RetrievalAttention等動態稀疏注意力機制，推理時可減少計算量，但訓練長上下文模型時，因仍需處理大量上下文信息，訓練成本未實質降低，阻礙LLMs擴展到長上下文場景；
而以Mamba、RWKV和RetNet為代表的線性注意力模型，用線性近似降低長序列計算開銷，但與傳統Transformer模型差異大，轉換成本高昂，甚至需重新訓練模型，且在復雜推理任務中的性能缺乏足夠驗證，限制其實際應用。

基于以上種種因素，一種新型注意力機制——MoBA架構應運而生。

顯而易見，不像傳統那樣關注全部鍵值（key），MoBA通過僅關注部分鍵值來提升效率。

它會把上下文按照一定規則分成n個塊。這就像處理一篇長文章時，把文章按段落分成不同部分。

然后每個塊都包含一部分鍵值信息。

之后，利用MoE的top-k門控機制來為每個查詢token挑選相關塊。

這就好比從所有段落中找出與當前問題最相關的幾個段落。它會計算查詢token和每個塊之間的相關性分數，然后選擇分數最高的k個塊，這樣查詢token就能聚焦在最有用的信息上。

上述過程用公式表示如下：

此外，為了保證基于前文預測下一個token的準確性，MoBA還采用了兩項關鍵設計：

不關注未來塊：在語言生成任務中，為防止當前token依賴未來token信息，MoBA規定查詢token只能關注之前或當前位置的塊，不關注未來塊，避免信息提前泄露。
當前塊因果掩碼處理：查詢token所在當前塊計算時可能涉及未來token信息，MoBA用因果掩碼 “遮擋” 這些信息，計算注意力時只關注已出現的信息。

與此同時，論文還公開了其他關鍵設計選擇。

比如更細粒度的塊分割。研究發現，把上下文劃分成更細的塊，能讓模型更好地捕捉信息。

再比如MoBA與全注意力混合。即讓MoBA可以和全注意力模式相互切換，在剛開始訓練或者處理復雜問題時，可以用全注意力模式，讓模型全面了解信息；而在處理長文本，對效率要求高時，就切換到MoBA模式，節省計算資源。

到了具體實現上，MoBA還結合了Flash Attention（能讓注意力計算更高效）和MoE的優化手段。

完整過程小結如下：

第一步：確定查詢token到KV塊的分配，就像是給每個問題分配對應的 “答案段落”；

第二步：排序查詢token，比如把問相同主題問題的查詢token放在一起，方便統一處理；

第三步：計算每個KV塊的注意力輸出，用Flash Attention技術，讓模型去 “理解” 對應塊里的信息，得出相關結果；

第四步：重排注意力輸出并合并結果，把計算出的注意力輸出按原來順序重排，再用在線Softmax合并，將不同來源的結果整合得到綜合結果。就像把不同 “答案段落” 的信息整合，最終得出一個結論。