Kimi背后的長上下文處理機(jī)制曝光了! 這項(xiàng)名為MoBA的新型注意力機(jī)制,能將處理1M長文本的速度一下子提升6.5倍,而且還是經(jīng)過Kimi平臺(tái)實(shí)際驗(yàn)證的那種。 概括而言,這項(xiàng)耗時(shí)一年半的工作主要看點(diǎn)在: 把完整上下文劃分成塊,讓每個(gè)查詢token自動(dòng)去關(guān)注最相關(guān)的KV塊,這樣就能高效處理長序列數(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長文本的速度分別提升了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ā)展,處理長文本的能力將更加關(guān)鍵。然而,傳統(tǒng)注意力機(jī)制的計(jì)算復(fù)雜度隨著序列長度的增加而呈平方級(jí)增長,這一特性嚴(yán)重阻礙了模型對(duì)長序列的高效處理。 而且,一些現(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)練長上下文模型時(shí),因仍需處理大量上下文信息,訓(xùn)練成本未實(shí)質(zhì)降低,阻礙LLMs擴(kuò)展到長上下文場景; 而以Mamba、RWKV和RetNet為代表的線性注意力模型,用線性近似降低長序列計(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è)塊。這就像處理一篇長文章時(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ù)測下一個(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í),可以用全注意力模式,讓模型全面了解信息;而在處理長文本,對(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長上下文驗(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)。 在長上下文實(shí)驗(yàn)中,盡管MoBA最后塊損失略高,但差距逐漸縮小,表明其長上下文可擴(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ǔn)測試中評(píng)估,結(jié)果顯示其性能與全注意力模型相當(dāng),在RULER基準(zhǔn)測試中二者得分接近,在1M上下文長度的 “大海撈針” 基準(zhǔn)測試中也表現(xiàn)良好。 總之,MoBA的計(jì)算復(fù)雜度隨著上下文長度增加而優(yōu)勢(shì)明顯。 在1M token的測試中,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) 參考鏈接: 本文來源:量子位 |
原創(chuàng)欄目
IT百科
網(wǎng)友評(píng)論
聚超值•精選