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

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

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

• 把完整上下文劃分成塊，讓每個查詢token自動去關(guān)注最相關(guān)的KV塊，這樣就能高效處理長序列數(shù)據(jù)；
• 提出一種新的參數(shù)無關(guān)的top-k門控機制，它能給每個查詢token挑選出最相關(guān)的塊，保證模型只聚焦在最有用信息的塊上；
• 支持在全注意力和稀疏注意力模式之間輕松切換；

一言以蔽之，MoBA將MoE（專家混合）應(yīng)用于注意力機制，通過遵循一種“l(fā)ess structure” 原則，允許模型自主決定關(guān)注哪些區(qū)域或位置。

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

而且按照論文得出的結(jié)論：

MoBA開箱即用，它能輕松融入現(xiàn)有模型，無需高昂訓(xùn)練成本。

另外，在深扒論文細(xì)節(jié)的同時，我們也在作者名單中發(fā)現(xiàn)了楊植麟本人的身影。

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

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

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

將MoE原理應(yīng)用于注意力機制

還是先來看論文細(xì)節(jié)。

首先，團(tuán)隊回答了為什么要開展這項研究？

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

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

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

基于以上種種因素，一種新型注意力機制——MoBA架構(gòu)應(yīng)運而生。

顯而易見，不像傳統(tǒng)那樣關(guān)注全部鍵值（key），MoBA通過僅關(guān)注部分鍵值來提升效率。

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

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

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

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

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

此外，為了保證基于前文預(yù)測下一個token的準(zhǔn)確性，MoBA還采用了兩項關(guān)鍵設(shè)計：

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

與此同時，論文還公開了其他關(guān)鍵設(shè)計選擇。

比如更細(xì)粒度的塊分割。研究發(fā)現(xiàn)，把上下文劃分成更細(xì)的塊，能讓模型更好地捕捉信息。

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

到了具體實現(xiàn)上，MoBA還結(jié)合了Flash Attention（能讓注意力計算更高效）和MoE的優(yōu)化手段。 

完整過程小結(jié)如下：

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

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

第三步：計算每個KV塊的注意力輸出，用Flash Attention技術(shù)，讓模型去 “理解” 對應(yīng)塊里的信息，得出相關(guān)結(jié)果；

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

經(jīng)過Kimi 1M長上下文驗證

實驗階段，研究還得出了幾個值得關(guān)注的發(fā)現(xiàn)。

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

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

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

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

在監(jiān)督微調(diào)（SFT）中，層混合策略（部分層用全注意力，其余用MoBA）可顯著降低SFT損失。

以Llama 3.1 8B模型為基礎(chǔ)，對MoBA在多種長上下文基準(zhǔn)測試中評估，結(jié)果顯示其性能與全注意力模型相當(dāng)，在RULER基準(zhǔn)測試中二者得分接近，在1M上下文長度的 “大海撈針” 基準(zhǔn)測試中也表現(xiàn)良好。

總之，MoBA的計算復(fù)雜度隨著上下文長度增加而優(yōu)勢明顯。

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

OMT：又和DeepSeek撞車了

回到一開頭提到的，事實上，Kimi這篇論文一發(fā)，就有網(wǎng)友在底下惋惜：

而且還有人當(dāng)面提起了“傷心事”：

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

好家伙，連續(xù)兩次“撞車”，妥妥的宿命感這不就來了！（doge）