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

一水 發自 凹非寺

量子位 | 公眾號 QbitAI

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

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

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

• 把完整上下文劃分成塊，讓每個查詢token自動去關注最相關的KV塊，這樣就能高效處理長序列數據；

• 提出一種新的參數無關的top-k門控機制，它能給每個查詢token挑選出最相關的塊，保證模型只聚焦在最有用信息的塊上；

• 支持在全注意力和稀疏注意力模式之間輕鬆切換；

一言以蔽之，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合併，將不同來源的結果整合得到綜合結果。就像把不同 「答案段落」 的信息整合，最終得出一個結論。

經過Kimi 1M長上下文驗證

實驗階段，研究還得出了幾個值得關注的發現。

首先，對比全注意力（使用Flash Attention實現）和MoBA訓練的語言模型，發現二者縮放趨勢相似，MoBA在高達75%稀疏度下性能與全注意力相當。

在長上下文實驗中，儘管MoBA最後塊損失略高，但差距逐漸縮小，表明其長上下文可擴展性。

消融實驗表明，細粒度塊分割確實對MoBA性能提升明顯。

其次，如果將MoBA與全注意力混合訓練，其模型在位置LM損失上接近全注意力模型，證明該訓練方法能平衡效率和性能。

在監督微調（SFT）中，層混合策略（部分層用全注意力，其餘用MoBA）可顯著降低SFT損失。

以Llama 3.1 8B模型為基礎，對MoBA在多種長上下文基準測試中評估，結果顯示其性能與全注意力模型相當，在RULER基準測試中二者得分接近，在1M上下文長度的 「大海撈針」 基準測試中也表現良好。

總之，MoBA的計算複雜度隨著上下文長度增加而優勢明顯。

在1M token的測試中，MoBA比全注意力快了6.5倍；到10M token時，則提速16倍。

OMT：又和DeepSeek撞車了

回到一開頭提到的，事實上，Kimi這篇論文一發，就有網民在底下惋惜：

而且還有人當面提起了「傷心事」：

原來，就在上個月（1月20日），也是在DeepSeek發了號稱比肩OpenAI-o1正式版的DeepSeek-R1之後，前後僅相隔僅2小時，Kimi才發佈了追平OpenAI-o1滿血版的多模態思考模型k1.5。

好傢伙，連續兩次「撞車」，妥妥的宿命感這不就來了！（doge）

參考鏈接：

[1]https://github.com/MoonshotAI/MoBA?tab=readme-ov-file

[2]https://x.com/Kimi_Moonshot/status/1891825059599352259

[3]https://x.com/deepseek_ai/status/1891745487071609327