（MoBA）大語言模型長文本處理新解法：塊注意力混合

在自然語言處理領(lǐng)域，高效處理長文本一直是個(gè)老大難問題。隨著大語言模型在閱讀、理解和生成文本方面的能力不斷提升，其處理輸入的核心——注意力機(jī)制，卻逐漸成了瓶頸。在典型的Transformer架構(gòu)里，這種機(jī)制要把每個(gè)詞元（token）和其他所有詞元進(jìn)行比較，這就導(dǎo)致計(jì)算成本會(huì)隨著序列長度的增加呈二次方增長。如今，我們把語言模型應(yīng)用到長篇文檔、多章節(jié)書籍、法律文書或是大型代碼庫這類需要處理海量文本信息的任務(wù)中時(shí)，這個(gè)問題就更加突出了。要是模型得處理幾萬甚至幾十萬個(gè)詞元，簡單粗暴地計(jì)算全注意力的成本實(shí)在高得離譜。

以往解決這個(gè)問題的方法，常常是設(shè)置固定結(jié)構(gòu)或者采用近似算法，但這些在某些場景下可能會(huì)影響效果。比如說，滑動(dòng)窗口機(jī)制只能讓詞元關(guān)注局部區(qū)域，這樣就容易忽略掉重要的全局關(guān)系。還有些方法，像直接用全新結(jié)構(gòu)替換softmax注意力機(jī)制，從根本上改變了基礎(chǔ)架構(gòu)，可這就需要從頭開始大規(guī)模重新訓(xùn)練模型，很難利用現(xiàn)有的預(yù)訓(xùn)練模型。所以，研究人員一直都在尋找一種新方法，既能保留原始Transformer設(shè)計(jì)的優(yōu)勢(shì)，也就是適應(yīng)性強(qiáng)、能捕捉廣泛依賴關(guān)系，又不會(huì)在處理超長序列時(shí)，產(chǎn)生傳統(tǒng)全注意力機(jī)制帶來的巨額計(jì)算開銷。

來自Moonshot AI、清華大學(xué)和浙江大學(xué)的研究人員，帶來了塊注意力混合（MoBA）這個(gè)創(chuàng)新方法，它把專家混合（MoE）的原理運(yùn)用到了注意力機(jī)制當(dāng)中。MoBA會(huì)把輸入劃分成一個(gè)個(gè)便于處理的“塊”，再通過一個(gè)可訓(xùn)練的門控系統(tǒng)，來確定每個(gè)查詢?cè)~元應(yīng)該關(guān)注哪些塊，這樣就解決了模型挨個(gè)比較詞元時(shí)效率低下的問題。和那些強(qiáng)制使用局部注意力或者窗口注意力的方法不同，MoBA能讓模型自己學(xué)習(xí)該重點(diǎn)關(guān)注哪里。這個(gè)設(shè)計(jì)遵循“l(fā)ess structure”原則，意思就是架構(gòu)不會(huì)提前規(guī)定哪些詞元必須相互作用，而是把這些決策交給訓(xùn)練好的門控網(wǎng)絡(luò)。

MoBA有個(gè)很關(guān)鍵的特點(diǎn)，就是能和現(xiàn)有的基于Transformer的模型完美配合。它沒有拋棄標(biāo)準(zhǔn)的自注意力接口，而是像個(gè)“插件”一樣，直接就能替換使用。MoBA的參數(shù)數(shù)量和原模型一樣，不會(huì)讓架構(gòu)變得臃腫，而且還保留了因果掩碼，保證自回歸生成的準(zhǔn)確性。在實(shí)際應(yīng)用中，MoBA可以在稀疏注意力和全注意力之間靈活切換。處理超長輸入時(shí)，用稀疏注意力能提高速度；在訓(xùn)練的某些層或者階段，如果有需要，還能切換回標(biāo)準(zhǔn)的全注意力模式。

可訓(xùn)練的區(qū)塊稀疏注意力

完整的上下文被劃分成多個(gè)區(qū)塊，每個(gè)查詢?cè)~元都能學(xué)會(huì)關(guān)注最相關(guān)的鍵值（KV）區(qū)塊，這樣就能高效處理長序列數(shù)據(jù)。

無參數(shù)門控機(jī)制

MoBA引入了一種很新穎的無參數(shù)top - k門控機(jī)制，專門為每個(gè)查詢?cè)~元挑選最相關(guān)的區(qū)塊，確保模型只聚焦在信息量最大的區(qū)塊上。

在完全注意力和稀疏注意力之間無縫過渡

MoBA設(shè)計(jì)得非常靈活，可以完美替代全注意力機(jī)制，在全注意力模式和稀疏注意力模式之間自由切換。

技術(shù)細(xì)節(jié)和優(yōu)勢(shì)

MoBA的核心操作，就是把上下文劃分成多個(gè)區(qū)塊，每個(gè)區(qū)塊都包含一連串連續(xù)的詞元。門控機(jī)制會(huì)計(jì)算查詢?cè)~元和每個(gè)區(qū)塊之間的“親和度”分?jǐn)?shù)，一般是通過把查詢和區(qū)塊里鍵的聚合表示進(jìn)行對(duì)比來實(shí)現(xiàn)。然后，它會(huì)選出得分最高的那些區(qū)塊。這樣一來，最終的注意力分配就只和最相關(guān)區(qū)塊里的詞元有關(guān)。而且，包含查詢?cè)~元本身的那個(gè)區(qū)塊一定會(huì)被納入，保證局部上下文信息隨時(shí)都能獲取到。同時(shí)，MoBA還采用了因果掩碼，讓詞元不會(huì)關(guān)注未來的位置，維持從左到右的自回歸特性。

經(jīng)過這樣的處理，MoBA的注意力矩陣比原始Transformer的要稀疏得多。但它依然很靈活，要是有需要，查詢也能獲取到遠(yuǎn)距離的信息。舉個(gè)例子，要是在文本末尾提出的問題，得參考文本開頭的細(xì)節(jié)才能回答，門控機(jī)制就能給開頭相關(guān)的區(qū)塊打高分。從技術(shù)層面來講，這種基于區(qū)塊的方法把詞元比較的數(shù)量減少到了亞二次方級(jí)別，隨著上下文長度增加到幾十萬甚至上百萬個(gè)詞元，效率提升就更加明顯了。

MoBA還有個(gè)很大的優(yōu)勢(shì)，就是和現(xiàn)代加速器以及專用內(nèi)核的兼容性很好。研究人員把MoBA和FlashAttention結(jié)合起來，F(xiàn)lashAttention是一個(gè)高性能庫，能實(shí)現(xiàn)快速、低內(nèi)存消耗的精確注意力計(jì)算。他們根據(jù)所選的區(qū)塊，仔細(xì)對(duì)查詢 - 鍵 - 值操作進(jìn)行分組，進(jìn)一步優(yōu)化了計(jì)算過程。研究人員表示，處理一百萬個(gè)詞元時(shí)，MoBA相比傳統(tǒng)的全注意力機(jī)制，速度能提升大約6倍，這在實(shí)際應(yīng)用場景里優(yōu)勢(shì)非常大。

結(jié)果和洞察

根據(jù)技術(shù)報(bào)告，MoBA在各種任務(wù)中的表現(xiàn)和全注意力機(jī)制不相上下，但處理長序列時(shí)，計(jì)算成本卻低得多。在語言建模數(shù)據(jù)測(cè)試?yán)铮?dāng)序列長度達(dá)到8192或者32768個(gè)詞元時(shí)，MoBA的困惑度和全注意力Transformer的很接近。更重要的是，當(dāng)研究人員把上下文長度逐漸增加到128000甚至更長時(shí)，MoBA對(duì)長上下文的理解能力依然很強(qiáng)。研究人員還做了“trailing token”評(píng)估，主要測(cè)試模型預(yù)測(cè)長提示末尾詞元的能力，這個(gè)測(cè)試通常能暴露那些過度依賴近似算法的方法的弱點(diǎn)。而MoBA在處理這些末尾位置的詞元時(shí)，預(yù)測(cè)質(zhì)量并沒有明顯下降。

研究人員還研究了MoBA對(duì)區(qū)塊大小和門控策略的敏感度。有些實(shí)驗(yàn)發(fā)現(xiàn)，細(xì)化粒度（也就是用更小的區(qū)塊，但選更多的區(qū)塊）能讓模型的效果更接近全注意力機(jī)制。就算MoBA忽略了大部分上下文，自適應(yīng)門控也能找到對(duì)查詢真正重要的區(qū)塊。此外，還有一種“混合”模式，采取了平衡策略：一部分層繼續(xù)用MoBA來提高速度，少數(shù)層則切換回全注意力模式。在進(jìn)行有監(jiān)督微調(diào)時(shí)，這種混合模式特別有用，因?yàn)橛?xùn)練目標(biāo)可能會(huì)屏蔽掉輸入里的某些位置。在少數(shù)上層保留全注意力，模型就能覆蓋更廣泛的上下文，對(duì)那些需要全局視角的任務(wù)很有幫助。

總的來說，這些研究結(jié)果表明，MoBA特別適合處理需要大量上下文的任務(wù)，像是長篇文檔閱讀理解、大規(guī)模代碼補(bǔ)全，還有需要參考完整對(duì)話歷史的多輪對(duì)話系統(tǒng)。MoBA能有效提升效率，而且對(duì)性能的影響很小，是讓大語言模型在大規(guī)模應(yīng)用中更高效的理想選擇。

結(jié)論

塊注意力混合（MoBA）為大語言模型處理長文本提供了一種更高效的方法，而且不用大幅修改Transformer架構(gòu)，也不會(huì)降低性能。通過在注意力模塊里融入專家混合的理念，MoBA用一種可學(xué)習(xí)的稀疏方式，讓模型聚焦在超長輸入的關(guān)鍵部分。它的設(shè)計(jì)非常靈活，尤其是能在稀疏注意力和全注意力之間無縫切換，這對(duì)正在進(jìn)行的和未來的訓(xùn)練流程都很有吸引力。研究人員可以靈活調(diào)整修剪注意力模式的程度，要是任務(wù)需要全面覆蓋信息，也可以有選擇地使用全注意力模式。

雖然目前對(duì)MoBA的研究主要集中在文本領(lǐng)域，但它的底層機(jī)制在其他數(shù)據(jù)模態(tài)上可能也有很大潛力。只要序列長度長到會(huì)引發(fā)計(jì)算或者內(nèi)存方面的問題，用MoBA這種把查詢分配給不同區(qū)塊“專家”的思路，就能在不影響處理全局依賴關(guān)系的前提下，有效緩解瓶頸。隨著語言應(yīng)用中的序列長度不斷增加，像MoBA這樣的方法，在提升神經(jīng)語言模型的可擴(kuò)展性和性價(jià)比方面，可能會(huì)發(fā)揮至關(guān)重要的作用。

鏈接

github鏈接：https://github.com/MoonshotAI/MoBA論文：  https://arxiv.org/html/2406.14909v1

本文轉(zhuǎn)載自 ??柏企科技圈??，作者：柏企