新智元報(bào)道

　　編輯：KingHZ 犀牛

　　【新智元導(dǎo)讀】注意力機(jī)制的「平方枷鎖」，再次被撬開！一招Fenwick樹分段，用掩碼矩陣，讓注意力煥發(fā)對數(shù)級效率。更厲害的是，它無縫對接線性注意力家族，Mamba-2、DeltaNet 全員提速，跑分全面開花。長序列處理邁入log時(shí)代！

　　LLM苦算力太久了！

　　為緩解長序列建模中的算力瓶頸，研究界持續(xù)探索高效替代方案。

　　這次Mamba作者Tri Dao、華人AI領(lǐng)域大牛Eric P. Xing等聯(lián)手MIT、普林斯頓、CMU等機(jī)構(gòu)的研究人員，提出了全新的注意力機(jī)制：對數(shù)線性注意力（Log-Linear Attention）。

　　它具有以下特點(diǎn)：

　　- 訓(xùn)練效率：對數(shù)線性時(shí)間

　　- 推理性能：對數(shù)級別的時(shí)間和空間復(fù)雜度 - 硬件執(zhí)行：利用Triton內(nèi)核實(shí)現(xiàn)的高效執(zhí)行

　　論文鏈接：https://arxiv.org/abs/2506.04761

　　代碼鏈接：https://github.com/HanGuo97/log-linear-attention

　　此外，研究人員引入了新理論框架，統(tǒng)一了不同高效注意力機(jī)制的分析視角。

　　另外值得一提的是，兩位第一作者都是華人，均麻省理工學(xué)院計(jì)算機(jī)科學(xué)與人工智能實(shí)驗(yàn)室就讀。

　　結(jié)構(gòu)矩陣，一統(tǒng)注意力變體

　　2017 年，谷歌的八位研究人員提出了Transformer架構(gòu)，自此注意力機(jī)制（attention mechanism）開始主導(dǎo)LLM的發(fā)展。

　　然而，注意力機(jī)制存在「先天頑疾」：

　　它的計(jì)算復(fù)雜度與輸入序列長度N是平方關(guān)系，也就是O（N2）。

　　近年來，涌現(xiàn)了大量致力于實(shí)現(xiàn)次二次方計(jì)算復(fù)雜度（sub-quadratic compute）和次線性內(nèi)存消耗（sub-linear memory）的高效替代方案。

　　他們主要包括：線性注意力（linear attention）、狀態(tài)空間模型（state-space models）以及長卷積模型（long convolution models）。

　　盡管這些方法各有不同，但它們大多可以用以下方程統(tǒng)一表示：

　　其中A表示一個(gè)類Attention的交互矩陣，例如在線性注意力中，矩陣A就是Q和K的轉(zhuǎn)置矩陣的乘積矩陣；

　　而M是下三角形的因果掩碼矩陣，如線性注意力中的M的元素只能取值0和1。

　　從結(jié)構(gòu)矩陣視角，這種表示形式把交互項(xiàng)A與掩碼矩陣M拆分開，揭示了大量不同模型之間的結(jié)構(gòu)共性，如表1所示。

　　通常矩陣M，用于模擬不同時(shí)間步之間的「衰減關(guān)系」。

　　對掩碼矩陣M引入不同的結(jié)構(gòu)形式，還可以進(jìn)一步促進(jìn)訓(xùn)練和推理的高效實(shí)現(xiàn)。

　　掩碼矩陣M的結(jié)構(gòu)，決定了對高效算法的實(shí)現(xiàn)。

　　即便不使用softmax，如果采用無結(jié)構(gòu)的M（例如隨機(jī)下三角矩陣），注意力機(jī)制的計(jì)算和內(nèi)存復(fù)雜度，仍為與softmax注意力機(jī)制相當(dāng)。

　　這表明：提升效率的關(guān)鍵不只是去除softmax，而在于M本身是否具備合適的結(jié)構(gòu)。

　　在標(biāo)準(zhǔn)的線性注意力中，M是由1構(gòu)成的下三角矩陣。

　　這種結(jié)構(gòu)能對輸出O進(jìn)行分塊處理，從而將算法整體復(fù)雜度降至O(T)。

　　然而，在傳統(tǒng)注意力和這些線性時(shí)間變體之間，是否還存在其他可能性？

　　此方法還可以推廣到更復(fù)雜的門控機(jī)制中，此時(shí)的M擁有一種稱為「1-半可分結(jié)構(gòu)」（1-semiseparable structure）的特殊形式。

　　在狀態(tài)空間對偶建模框架中，這一方法已經(jīng)有所體現(xiàn)。

　　論文鏈接：https://arxiv.org/abs/2405.21060

　　另外，在長卷積模型（long convolution models）中，可以通過使用快速傅里葉變換（FFT）進(jìn)一步將復(fù)雜度降為O(TlogT)，相較于原始的O(T2)計(jì)算量，實(shí)現(xiàn)了顯著的效率提升。

　　對數(shù)線性注意力

　　在上一節(jié)中，已經(jīng)知道：注意力的計(jì)算效率和內(nèi)存消耗，取決于公式O=(A⊙M)V中掩碼矩陣M的結(jié)構(gòu)。

　　對數(shù)線性注意力機(jī)制（log-linear attention）就是在矩陣M引入特定結(jié)構(gòu)，讓計(jì)算復(fù)雜度在序列長度T上達(dá)到O(TlogT)，內(nèi)存復(fù)雜度降低到O(logT)。

　　該機(jī)制僅修改掩碼矩陣M，可無縫應(yīng)用于各種線性注意力模型。

　　作為應(yīng)用示例，研究人員展示了如何基于該框架構(gòu)建Mamba-2和Gated DeltaNet的對數(shù)線性版本。

　　圖1：標(biāo)準(zhǔn)線性注意力機(jī)制（上）與對數(shù)線性注意力機(jī)制（下）對比示意圖

　　特殊結(jié)構(gòu)：Fenwick樹劃分

　　在掩碼矩陣M上，對數(shù)線性注意力機(jī)制引入了一種特殊結(jié)構(gòu)，讓計(jì)算復(fù)雜度達(dá)到對數(shù)線性級別，內(nèi)存開銷則為對數(shù)級別。

　　為了實(shí)現(xiàn)這種多時(shí)間尺度的結(jié)構(gòu)化劃分，關(guān)鍵在于如何將前綴區(qū)間[0,t]分配給第t步的查詢向量。

　　根據(jù)Token的絕對位置s，可以簡單地把它劃入層級?=?log?s?。

　　但在自回歸解碼中，這種做法會導(dǎo)致對最近輸入的劃分粒度過大，進(jìn)而影響模型在關(guān)鍵位置上的預(yù)測精度。直覺上，越靠近當(dāng)前時(shí)間點(diǎn)的上下文信息越重要，應(yīng)該以更高分辨率來建模。

　　為了解決這一問題，研究者采用了另一種的分段策略。

　　從原理上看，這種結(jié)構(gòu)類似于Fenwick樹（也稱為樹狀數(shù)組）所使用的分層方式，將輸入序列按2的冪大小劃分為一系列區(qū)段。

　　Fenwick樹是一種支持單點(diǎn)修改和區(qū)間查詢的，代碼量小的數(shù)據(jù)結(jié)構(gòu)

　　在這種設(shè)計(jì)下，每個(gè)位置都會匯總一個(gè)以自身為終點(diǎn)的時(shí)間片段。

　　這能讓查詢操作只需關(guān)注少量（數(shù)量隨序列長度對數(shù)增長）的隱藏狀態(tài)，這些狀態(tài)能以不同時(shí)間粒度捕捉歷史上下文信息。

　　這種層次結(jié)構(gòu)使模型能夠以更精細(xì)的方式關(guān)注最近的token，同時(shí)在解碼過程中實(shí)現(xiàn)對數(shù)級別的時(shí)間和內(nèi)存效率。

　　圖2展示了這種劃分的可視化示意：每個(gè)Token被分配到若干層級桶中，最近的時(shí)間步被細(xì)致劃分，而越早的時(shí)間片則歸為更大的區(qū)段，從而實(shí)現(xiàn)了對時(shí)間上下文的層級壓縮建模。

　　為了生成最終的輸出向量，新方法會分別計(jì)算每個(gè)桶中的歷史記憶，并通過數(shù)據(jù)驅(qū)動(dòng)的標(biāo)量進(jìn)行加權(quán)。

　　該權(quán)重是輸入經(jīng)過線性變換后的結(jié)果，使得模型可以自適應(yīng)不同的時(shí)間尺度。

　　具體來說，輸出向量表達(dá)為：

　　如果所有標(biāo)量權(quán)重都相同或與層數(shù)?無關(guān)，則退化為線性注意力。

　　正是這些可區(qū)分的權(quán)重，賦予了模型捕捉多尺度時(shí)間結(jié)構(gòu)的能力。

　　為了更高效地在硬件上實(shí)現(xiàn)上述計(jì)算，可以將公式重構(gòu)為矩陣乘形式，方便批量并行：

　　其中，M^{H}根據(jù)s屬于t的哪一層?(t,s)來賦值。

　　在Fenwick分段下，這個(gè)矩陣呈現(xiàn)結(jié)構(gòu)化低秩模式，并能支持O(TlogT)的高效訓(xùn)練算法。

　　高效訓(xùn)練算法

　　線性注意力的分塊并行算法會將輸入序列劃分為若干長度為C的子塊，并對所有子塊進(jìn)行并行計(jì)算；當(dāng)需要跨塊傳遞信息時(shí)再進(jìn)行交互。

　　這種策略在「全并行計(jì)算」與「完全遞歸處理」之間找到平衡點(diǎn)，既減少了全局注意力的高計(jì)算成本，也提升了序列級別的并行效率。

　　同樣，分塊計(jì)算機(jī)制可以擴(kuò)展應(yīng)用于對數(shù)線性注意力機(jī)制。

　　首先注意到掩碼矩陣M^{H}的非對角區(qū)域具有低秩結(jié)構(gòu)，因此可將其分解為：

　　其中，D表示僅在塊內(nèi)部有效的對角矩陣，包含T?C個(gè)塊，每個(gè)塊記錄子塊內(nèi)的交互信息。

　　而M^{?}則表示第?層的跨塊依賴關(guān)系，

　　它通過一種類似樹狀結(jié)構(gòu)的方式，將較遠(yuǎn)位置之間的關(guān)聯(lián)壓縮成一個(gè)低秩表示（即對稱或重復(fù)性高的結(jié)構(gòu)），如圖3（左）所示。

　　基于這種結(jié)構(gòu)，研究者提出了分塊計(jì)算算法（見算法1和圖3右）。

　　這種方法在原有線性注意力的基礎(chǔ)上，僅引入了對數(shù)級別的額外開銷。

　　整個(gè)算法可分為兩個(gè)階段：

　　塊內(nèi)計(jì)算（?=0）：在每個(gè)子塊中，系統(tǒng)視其為無結(jié)構(gòu)數(shù)據(jù)，并使用標(biāo)準(zhǔn)的O(C2)計(jì)算完成塊內(nèi)交互。總共有T?C個(gè)子塊，因此整體塊內(nèi)計(jì)算成本為O(TC)。

　　塊間計(jì)算（?>0）：對于不同子塊之間的依賴，模型通過若干層次結(jié)構(gòu)表示進(jìn)行處理。這些結(jié)構(gòu)構(gòu)成了一個(gè)「分層可分矩陣」（SSS），允許在每層僅用少量操作完成跨塊傳遞。只要能調(diào)用諸如Mamba-2或GatedDeltaNet中那類高效的狀態(tài)傳遞模塊，每層的跨塊傳遞只需O(logT?C)次函數(shù)調(diào)用，每次耗費(fèi)O(T)的時(shí)間和內(nèi)存，因此總體跨塊成本為O(TlogT)。

　　該方法在原本線性注意力的計(jì)算程上，僅增加了對數(shù)級別的額外開銷，從而在保持高效性的同時(shí)提升了表達(dá)能力。

　　在圖3中，左圖展示了矩陣M的分解方式，右圖則是對應(yīng)的分塊計(jì)算算法（算法1）。

　　在Level 0，模型對每個(gè)小塊內(nèi)部進(jìn)行計(jì)算，采用的是相對于塊大小為二次復(fù)雜度的算法。由于每個(gè)塊本身較小，因此這一階段計(jì)算開銷低、效率高。

　　從Level 1開始，模型對不同塊之間進(jìn)行計(jì)算，方法是多次調(diào)用已有的跨塊計(jì)算算法組件。整體來看，該跨塊計(jì)算階段的復(fù)雜度相對于塊數(shù)是對數(shù)級別的，從而保證了整體計(jì)算過程的高效性。

　　這一方法實(shí)質(zhì)上是將經(jīng)典的scan掃描算法推廣到層級結(jié)構(gòu)中，研究者稱之為分塊并行掃描（chunkwise parallel scan）。

　　與傳統(tǒng)token級scan不同，它不再受限于內(nèi)存帶寬瓶頸，而是通過結(jié)構(gòu)優(yōu)化使?fàn)顟B(tài)以低成本在線上傳遞。

　　算法中每一層的系數(shù)，來自于掩碼矩陣的低秩項(xiàng)，可通過并行掃描算法（如Blelloch scan）進(jìn)行高效整合，從而提升整體訓(xùn)練效率和可擴(kuò)展性。

　　對Mamba-2和門控DeltaNet的對數(shù)線性推廣

　　這兩個(gè)模型的主要區(qū)別在于它們對轉(zhuǎn)換矩陣A的參數(shù)化方式不同。

　　研究團(tuán)隊(duì)的方法保留了每個(gè)模型中A的原始形式，同時(shí)將注意力掩碼與對數(shù)線性變體M進(jìn)行組合。

　　他們將得到的模型稱為對數(shù)線性Mamba-2和對數(shù)線性門控DeltaNet。

　　這一構(gòu)造體現(xiàn)了一個(gè)通用原則：任何具有結(jié)構(gòu)化記憶和高效分塊并行原語（chunkwise-parallel primitive）的線性注意力機(jī)制，都可以通過將其注意力掩碼與對數(shù)線性變體組合，擴(kuò)展為對數(shù)線性形式。

　　團(tuán)隊(duì)使用Triton實(shí)現(xiàn)了分塊并行掃描算法（chunkwise parallel scan algorithm）。

　　對數(shù)線性Mamba-2的定制內(nèi)核在序列長度超過8K時(shí)，性能超越了FlashAttention-2（前向+反向）。

　　在完整的訓(xùn)練設(shè)置中，吞吐量取決于模型架構(gòu)。值得注意的是，盡管對數(shù)線性Mamba-2（帶MLP）包含了Transformer中沒有的額外層（如深度卷積），但在序列長度達(dá)到32K時(shí)，其吞吐量依然超過了Transformer。

　　圖4：在不同序列長度下的訓(xùn)練吞吐量（左圖，數(shù)值越高越好）以及前向和反向傳播過程中內(nèi)核運(yùn)行時(shí)間（右圖，數(shù)值越低越好）。

　　圖4中，「Log-Linear Mamba-2 (naive)」表示簡單地重復(fù)使用現(xiàn)有的Mamba-2計(jì)算方法；

　　而「Log-Linear Mamba-2」」則采用了一種經(jīng)過優(yōu)化的自定義實(shí)現(xiàn)方式，其中包括層級融合（level fusion）等性能優(yōu)化手段。

　　當(dāng)序列長度達(dá)到131K時(shí)，訓(xùn)練吞吐量出現(xiàn)下降，這是由于引入了梯度檢查點(diǎn)（gradient checkpointing）以降低內(nèi)存使用所致。

　　所有實(shí)驗(yàn)均在H100 GPU上運(yùn)行，具體配置為：

　　batch size為2，注意力頭數(shù)為48，每個(gè)頭的維度為64，狀態(tài)維度為128，chunk size設(shè)置為64。

　　在（Log-Linear）Mamba-2中采用MVA，在FlashAttention-2中采用GQA。

　　實(shí)驗(yàn)結(jié)果

　　研究團(tuán)隊(duì)首先在多查詢關(guān)聯(lián)回憶（MQAR）上進(jìn)行實(shí)驗(yàn)，這是一個(gè)用于評估模型上下文回憶能力的標(biāo)準(zhǔn)測試基準(zhǔn)。

　　他們在一個(gè)包含1萬個(gè)樣本的數(shù)據(jù)集上訓(xùn)練了100個(gè)周期，并對學(xué)習(xí)率進(jìn)行了調(diào)整。

　　如圖5所示，隨著序列長度和鍵值對數(shù)量的增加，DeltaNet的性能顯著下降，而對數(shù)線性DeltaNet（Log-Linear DeltaNet）依然保持高準(zhǔn)確率。

　　需要注意的是，softmax注意力在所有設(shè)置下都能達(dá)到滿分準(zhǔn)確率。

　　語言建模

　　研究團(tuán)隊(duì)在Long-Data-Collections數(shù)據(jù)集上使用500億個(gè)token，從頭開始進(jìn)行學(xué)術(shù)規(guī)模的語言建模預(yù)訓(xùn)練，序列長度為16K。

　　所有模型都有21層，隱藏層大小為1536。

　　我們使用了以下模型：

　　帶16個(gè)注意力頭的Transformer，RoPE基數(shù)為50萬；

　　修改版的Mamba-2，包含48個(gè)頭和MLP層；

　　帶6個(gè)頭的門控DeltaNet。

　　這些模型的參數(shù)量分別是：Transformer（6.93億）、Mamba-2（8.02億）、門控DeltaNet（7.93億）。

　　標(biāo)準(zhǔn)基準(zhǔn)測試

　　團(tuán)隊(duì)在WikiText困惑度和幾個(gè)零樣本常識推理基準(zhǔn)上評估模型（表2）。這些都是短上下文任務(wù)，因此對模型狀態(tài)大小不太敏感。

　　對數(shù)線性Mamba-2在困惑度和一半的常識推理任務(wù)上優(yōu)于其線性版本。

　　對數(shù)線性門控DeltaNet表現(xiàn)更突出，在困惑度和除一項(xiàng)推理基準(zhǔn)外的所有任務(wù)上都超過了其線性版本。值得注意的是，它在所有指標(biāo)上都優(yōu)于層數(shù)匹配的Transformer，并且在一半指標(biāo)上優(yōu)于參數(shù)量匹配的Transformer。

　　逐位置損失

　　研究團(tuán)隊(duì)報(bào)告了模型在每個(gè)token位置的損失，以評估其處理長上下文的能力（圖6）。

　　如果隨著token位置增加，損失持續(xù)下降，說明模型能有效利用整個(gè)上下文。然而，如果損失在某一點(diǎn)后趨于平穩(wěn)，則表明模型難以利用序列中過于靠后的信息。在這項(xiàng)分析中，使用了來自Book-3的3900萬個(gè)token。

　　結(jié)果顯示，將Mamba-2和門控DeltaNet擴(kuò)展到它們的對數(shù)線性版本后，（平滑后的）損失在不同位置上均持續(xù)降低，表明長距離上下文利用能力有所提升。

　　對數(shù)線性門控DeltaNet的性能也與層數(shù)匹配的Transformer非常接近，盡管與參數(shù)量匹配的Transformer相比仍存在性能差距。

　　大海撈針

　　團(tuán)隊(duì)使用了RULER中的「大海撈針」（NIAH，圖7）基準(zhǔn)測試，在該測試中，模型需要根據(jù)隱藏在長上下文中的鍵來檢索一個(gè)值（針）。

　　在較簡單的單針任務(wù)中，對數(shù)線性Mamba-2在9個(gè)指標(biāo)中的8個(gè)上優(yōu)于其線性版本。

　　門控DeltaNet在多個(gè)情況下已達(dá)到完美準(zhǔn)確率，但在3個(gè)指標(biāo)上有所提升，另外3個(gè)保持不變。

　　在更具挑戰(zhàn)性的多針任務(wù)中，對數(shù)線性Mamba-2再次在9個(gè)指標(biāo)中的8個(gè)上有所改進(jìn)，而對數(shù)線性門控DeltaNet則在所有指標(biāo)上均取得進(jìn)步。

　　上下文檢索

　　團(tuán)隊(duì)在現(xiàn)實(shí)世界的、需要大量回憶的任務(wù)上評估模型（表3）。

　　由于這些基準(zhǔn)測試最初是為短序列（≤2K token）設(shè)計(jì)的，他們報(bào)告了序列長度為512、1024、2048以及（除NQ外）16K的結(jié)果。

　　結(jié)果發(fā)現(xiàn)，對數(shù)線性Mamba-2在大約一半任務(wù)（SQuAD、TriviaQA和NQ）上有所改進(jìn)。

　　相比之下，對數(shù)線性門控DeltaNet表現(xiàn)更為穩(wěn)定，在除DROP之外的所有任務(wù)上均匹配或優(yōu)于門控DeltaNet。

　　長上下文理解

　　最后，他們在LongBench（表4）上評估了模型的性能。

　　結(jié)果顯示，對數(shù)線性Mamba-2和門控DeltaNet在14個(gè)評估任務(wù)中的8個(gè)上均優(yōu)于基線Mamba-2和門控DeltaNet。

　　討論與局限性

　　雖然對數(shù)線性注意力在許多情況下優(yōu)于線性注意力，但仍有不少任務(wù)中它的表現(xiàn)未能超越線性注意力的基線。

　　由于計(jì)算資源限制，研究團(tuán)隊(duì)無法嘗試不同的λ項(xiàng)參數(shù)化（或超參數(shù)調(diào)整），而優(yōu)化λ的參數(shù)化可能會帶來更好的結(jié)果。

　　此外，與Transformer相比，所有基準(zhǔn)測試中仍存在顯著的性能差距。

　　對數(shù)線性注意力的工程復(fù)雜性較高。塊間計(jì)算在概念上類似于多次應(yīng)用線性注意力原語，但塊內(nèi)操作需要專門的實(shí)現(xiàn)。這些塊內(nèi)機(jī)制是導(dǎo)致速度差異的主要因素。

　　此外，反向傳播過程更為復(fù)雜，因?yàn)椴粌H需要（手動(dòng)）計(jì)算標(biāo)準(zhǔn)注意力組件的梯度，還需計(jì)算額外的λ項(xiàng)梯度。

　　最后，F(xiàn)enwick樹分區(qū)的使用引入了一種歸納偏差：近期token被分配更細(xì)粒度的內(nèi)存，而較遠(yuǎn)的token被更激進(jìn)地壓縮。

　　更多實(shí)驗(yàn)設(shè)置等細(xì)節(jié)，請參閱原文。

　　一作簡介

　　Han Guo，現(xiàn)任麻省理工學(xué)院計(jì)算機(jī)科學(xué)與人工智能實(shí)驗(yàn)室（MIT CSAIL）博士研究生，師從Yoon Kim教授與Eric P. Xing（邢波）教授。

　　此前，他曾在卡耐基梅隆大學(xué)語言技術(shù)研究所（CMU LTI）、北卡羅來納大學(xué)NLP研究組（UNC-NLP）， 與Mohit Bansal教授開展研究，度過數(shù)年寶貴學(xué)術(shù)時(shí)光。

　　他的研究方向聚焦可擴(kuò)展高效機(jī)器學(xué)習(xí)/自然語言處理的算法與系統(tǒng)設(shè)計(jì)，2022年榮獲微軟研究院博士生獎(jiǎng)學(xué)金（Microsoft Research PhD Fellowship）。

　　Songlin Yang，是麻省理工學(xué)院計(jì)算機(jī)科學(xué)與人工智能實(shí)驗(yàn)室（MIT CSAIL）的博士生，師從Yoon Kim教授。

　　她2020年獲得南方科技大學(xué)學(xué)士學(xué)位，2023年獲得上海科技大學(xué)碩士學(xué)位。

　　她聚焦機(jī)器學(xué)習(xí)系統(tǒng)與大型語言模型的交叉領(lǐng)域，特別關(guān)注：

　　? 面向硬件的高效序列建模算法設(shè)計(jì)

　　? 線性注意力模型（linear attention）的優(yōu)化與創(chuàng)新

　　參考資料：

　　https://x.com/HanGuo97/status/1930789829094297859

　　https://arxiv.org/abs/2506.04761