今天小編分享的科學經驗：Kimi新論文再次“撞車”DeepSeek，都談到了長文注意力機制，歡迎閱讀。
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
