Kimi背后的長(zhǎng)上下文處理機(jī)制曝光了！

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

概括而言，這項(xiàng)耗時(shí)一年半的工作主要看點(diǎn)在：

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

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

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

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

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

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

與此同時(shí)，再次令人感到戲劇的是，Kimi的這項(xiàng)研究又又又和DeepSeek撞車了？？

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

經(jīng)過Kimi 1M長(zhǎng)上下文驗(yàn)證

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

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

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

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

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

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

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

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

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

OMT：又和DeepSeek撞車了

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

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

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

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