wpk打破地域界限,無論你身在何處,只需輕點鼠標,即可與各地的撲克高手同臺競技。感受不同文化的碰撞。
她2020年獲得南方科技大學(xué)學(xué)士學(xué)位,2023年獲得上海科技大學(xué)碩士學(xué)位。
對Mamba-2和門控DeltaNet的對數(shù)線性推廣
其中,M^{H}根據(jù)s屬于t的哪一層?(t,s)來賦值。
對數(shù)線性門控DeltaNet的性能也與層數(shù)匹配的Transformer非常接近,盡管與參數(shù)量匹配的Transformer相比仍存在性能差距。
對數(shù)線性Mamba-2在困惑度和一半的常識推理任務(wù)上優(yōu)于其線性版本。
他的研究方向聚焦可擴展高效機器學(xué)習(xí)/自然語言處理的算法與系統(tǒng)設(shè)計,2022年榮獲微軟研究院博士生獎學(xué)金(Microsoft Research PhD Fellowship)。
這些模型的參數(shù)量分別是:Transformer(6.93億)、Mamba-2(8.02億)、門控DeltaNet(7.93億)。
基于這種結(jié)構(gòu),研究者提出了分塊計算算法(見算法1和圖3右)。
對掩碼矩陣M引入不同的結(jié)構(gòu)形式,還可以進一步促進訓(xùn)練和推理的高效實現(xiàn)。
這種策略在「全并行計算」與「完全遞歸處理」之間找到平衡點,既減少了全局注意力的高計算成本,也提升了序列級別的并行效率。
在(Log-Linear)Mamba-2中采用MVA,在FlashAttention-2中采用GQA。
為緩解長序列建模中的算力瓶頸,研究界持續(xù)探索高效替代方案。
論文鏈接:https://arxiv.org/abs/2405.21060
結(jié)果發(fā)現(xiàn),對數(shù)線性Mamba-2在大約一半任務(wù)(SQuAD、TriviaQA和NQ)上有所改進。
門控DeltaNet在多個情況下已達到完美準確率,但在3個指標上有所提升,另外3個保持不變。
特殊結(jié)構(gòu):Fenwick樹劃分
此外,反向傳播過程更為復(fù)雜,因為不僅需要(手動)計算標準注意力組件的梯度,還需計算額外的λ項梯度。
對數(shù)線性注意力的工程復(fù)雜性較高。塊間計算在概念上類似于多次應(yīng)用線性注意力原語,但塊內(nèi)操作需要專門的實現(xiàn)。這些塊內(nèi)機制是導(dǎo)致速度差異的主要因素。
這一方法實質(zhì)上是將經(jīng)典的scan掃描算法推廣到層級結(jié)構(gòu)中,研究者稱之為分塊并行掃描(chunkwise parallel scan)。
對數(shù)線性門控DeltaNet表現(xiàn)更突出,在困惑度和除一項推理基準外的所有任務(wù)上都超過了其線性版本。值得注意的是,它在所有指標上都優(yōu)于層數(shù)匹配的Transformer,并且在一半指標上優(yōu)于參數(shù)量匹配的Transformer。
為了生成最終的輸出向量,新方法會分別計算每個桶中的歷史記憶,并通過數(shù)據(jù)驅(qū)動的標量進行加權(quán)。
Han Guo,現(xiàn)任麻省理工學(xué)院計算機科學(xué)與人工智能實驗室(MIT CSAIL)博士研究生,師從Yoon Kim教授與Eric P. Xing(邢波)教授。
整個算法可分為兩個階段:
在完整的訓(xùn)練設(shè)置中,吞吐量取決于模型架構(gòu)。值得注意的是,盡管對數(shù)線性Mamba-2(帶MLP)包含了Transformer中沒有的額外層(如深度卷積),但在序列長度達到32K時,其吞吐量依然超過了Transformer。
? 線性注意力模型(linear attention)的優(yōu)化與創(chuàng)新
團隊使用Triton實現(xiàn)了分塊并行掃描算法(chunkwise parallel scan algorithm)。
根據(jù)Token的絕對位置s,可以簡單地把它劃入層級?=?log?s?。
在掩碼矩陣M上,對數(shù)線性注意力機制引入了一種特殊結(jié)構(gòu),讓計算復(fù)雜度達到對數(shù)線性級別,內(nèi)存開銷則為對數(shù)級別。
線性注意力的分塊并行算法會將輸入序列劃分為若干長度為C的子塊,并對所有子塊進行并行計算;當需要跨塊傳遞信息時再進行交互。
該權(quán)重是輸入經(jīng)過線性變換后的結(jié)果,使得模型可以自適應(yīng)不同的時間尺度。
- 推理性能:對數(shù)級別的時間和空間復(fù)雜度 - 硬件執(zhí)行:利用Triton內(nèi)核實現(xiàn)的高效執(zhí)行
最后,F(xiàn)enwick樹分區(qū)的使用引入了一種歸納偏差:近期token被分配更細粒度的內(nèi)存,而較遠的token被更激進地壓縮。
該機制僅修改掩碼矩陣M,可無縫應(yīng)用于各種線性注意力模型。
團隊使用了RULER中的「大海撈針」(NIAH,圖7)基準測試,在該測試中,模型需要根據(jù)隱藏在長上下文中的鍵來檢索一個值(針)。
編輯:KingHZ 犀牛
它通過一種類似樹狀結(jié)構(gòu)的方式,將較遠位置之間的關(guān)聯(lián)壓縮成一個低秩表示(即對稱或重復(fù)性高的結(jié)構(gòu)),如圖3(左)所示。
在標準的線性注意力中,M是由1構(gòu)成的下三角矩陣。
圖4中,「Log-Linear Mamba-2 (naive)」表示簡單地重復(fù)使用現(xiàn)有的Mamba-2計算方法;
我們使用了以下模型:
在上一節(jié)中,已經(jīng)知道:注意力的計算效率和內(nèi)存消耗,取決于公式O=(A⊙M)V中掩碼矩陣M的結(jié)構(gòu)。
雖然對數(shù)線性注意力在許多情況下優(yōu)于線性注意力,但仍有不少任務(wù)中它的表現(xiàn)未能超越線性注意力的基線。
此外,研究人員引入了新理論框架,統(tǒng)一了不同高效注意力機制的分析視角。
這能讓查詢操作只需關(guān)注少量(數(shù)量隨序列長度對數(shù)增長)的隱藏狀態(tài),這些狀態(tài)能以不同時間粒度捕捉歷史上下文信息。
通常矩陣M,用于模擬不同時間步之間的「衰減關(guān)系」。
此外,與Transformer相比,所有基準測試中仍存在顯著的性能差距。
從原理上看,這種結(jié)構(gòu)類似于Fenwick樹(也稱為樹狀數(shù)組)所使用的分層方式,將輸入序列按2的冪大小劃分為一系列區(qū)段。
具體來說,輸出向量表達為:
研究團隊的方法保留了每個模型中A的原始形式,同時將注意力掩碼與對數(shù)線性變體M進行組合。
這次Mamba作者Tri Dao、華人AI領(lǐng)域大牛Eric P. Xing等聯(lián)手MIT、普林斯頓、CMU等機構(gòu)的研究人員,提出了全新的注意力機制:對數(shù)線性注意力(Log-Linear Attention)。
如果所有標量權(quán)重都相同或與層數(shù)?無關(guān),則退化為線性注意力。
然而,在傳統(tǒng)注意力和這些線性時間變體之間,是否還存在其他可能性?
【新智元導(dǎo)讀】注意力機制的「平方枷鎖」,再次被撬開!一招Fenwick樹分段,用掩碼矩陣,讓注意力煥發(fā)對數(shù)級效率。更厲害的是,它無縫對接線性注意力家族,Mamba-2、DeltaNet 全員提速,跑分全面開花。長序列處理邁入log時代!
所有實驗均在H100 GPU上運行,具體配置為:
此方法還可以推廣到更復(fù)雜的門控機制中,此時的M擁有一種稱為「1-半可分結(jié)構(gòu)」(1-semiseparable structure)的特殊形式。
而「Log-Linear Mamba-2」」則采用了一種經(jīng)過優(yōu)化的自定義實現(xiàn)方式,其中包括層級融合(level fusion)等性能優(yōu)化手段。
圖4:在不同序列長度下的訓(xùn)練吞吐量(左圖,數(shù)值越高越好)以及前向和反向傳播過程中內(nèi)核運行時間(右圖,數(shù)值越低越好)。
他們將得到的模型稱為對數(shù)線性Mamba-2和對數(shù)線性門控DeltaNet。
在這種設(shè)計下,每個位置都會匯總一個以自身為終點的時間片段。
它的計算復(fù)雜度與輸入序列長度N是平方關(guān)系,也就是O(N2)。
正是這些可區(qū)分的權(quán)重,賦予了模型捕捉多尺度時間結(jié)構(gòu)的能力。
該方法在原本線性注意力的計算程上,僅增加了對數(shù)級別的額外開銷,從而在保持高效性的同時提升了表達能力。
相比之下,對數(shù)線性門控DeltaNet表現(xiàn)更為穩(wěn)定,在除DROP之外的所有任務(wù)上均匹配或優(yōu)于門控DeltaNet。
從結(jié)構(gòu)矩陣視角,這種表示形式把交互項A與掩碼矩陣M拆分開,揭示了大量不同模型之間的結(jié)構(gòu)共性,如表1所示。
在Fenwick分段下,這個矩陣呈現(xiàn)結(jié)構(gòu)化低秩模式,并能支持O(TlogT)的高效訓(xùn)練算法。
此前,他曾在卡耐基梅隆大學(xué)語言技術(shù)研究所(CMU LTI)、北卡羅來納大學(xué)NLP研究組(UNC-NLP), 與Mohit Bansal教授開展研究,度過數(shù)年寶貴學(xué)術(shù)時光。
帶16個注意力頭的Transformer,RoPE基數(shù)為50萬;
而M^{?}則表示第?層的跨塊依賴關(guān)系,
為了更高效地在硬件上實現(xiàn)上述計算,可以將公式重構(gòu)為矩陣乘形式,方便批量并行:
算法中每一層的系數(shù),來自于掩碼矩陣的低秩項,可通過并行掃描算法(如Blelloch scan)進行高效整合,從而提升整體訓(xùn)練效率和可擴展性。
這一構(gòu)造體現(xiàn)了一個通用原則:任何具有結(jié)構(gòu)化記憶和高效分塊并行原語(chunkwise-parallel primitive)的線性注意力機制,都可以通過將其注意力掩碼與對數(shù)線性變體組合,擴展為對數(shù)線性形式。
對數(shù)線性Mamba-2的定制內(nèi)核在序列長度超過8K時,性能超越了FlashAttention-2(前向+反向)。
結(jié)果顯示,對數(shù)線性Mamba-2和門控DeltaNet在14個評估任務(wù)中的8個上均優(yōu)于基線Mamba-2和門控DeltaNet。
她聚焦機器學(xué)習(xí)系統(tǒng)與大型語言模型的交叉領(lǐng)域,特別關(guān)注:
所有模型都有21層,隱藏層大小為1536。
https://arxiv.org/abs/2506.04761
首先注意到掩碼矩陣M^{H}的非對角區(qū)域具有低秩結(jié)構(gòu),因此可將其分解為:
從Level 1開始,模型對不同塊之間進行計算,方法是多次調(diào)用已有的跨塊計算算法組件。整體來看,該跨塊計算階段的復(fù)雜度相對于塊數(shù)是對數(shù)級別的,從而保證了整體計算過程的高效性。
但在自回歸解碼中,這種做法會導(dǎo)致對最近輸入的劃分粒度過大,進而影響模型在關(guān)鍵位置上的預(yù)測精度。直覺上,越靠近當前時間點的上下文信息越重要,應(yīng)該以更高分辨率來建模。
這種方法在原有線性注意力的基礎(chǔ)上,僅引入了對數(shù)級別的額外開銷。
代碼鏈接:https://github.com/HanGuo97/log-linear-attention
塊間計算(?>0):對于不同子塊之間的依賴,模型通過若干層次結(jié)構(gòu)表示進行處理。這些結(jié)構(gòu)構(gòu)成了一個「分層可分矩陣」(SSS),允許在每層僅用少量操作完成跨塊傳遞。只要能調(diào)用諸如Mamba-2或GatedDeltaNet中那類高效的狀態(tài)傳遞模塊,每層的跨塊傳遞只需O(logT?C)次函數(shù)調(diào)用,每次耗費O(T)的時間和內(nèi)存,因此總體跨塊成本為O(TlogT)。
LLM苦算力太久了!
2017 年,谷歌的八位研究人員提出了Transformer架構(gòu),自此注意力機制(attention mechanism)開始主導(dǎo)LLM的發(fā)展。
圖2展示了這種劃分的可視化示意:每個Token被分配到若干層級桶中,最近的時間步被細致劃分,而越早的時間片則歸為更大的區(qū)段,從而實現(xiàn)了對時間上下文的層級壓縮建模。
由于計算資源限制,研究團隊無法嘗試不同的λ項參數(shù)化(或超參數(shù)調(diào)整),而優(yōu)化λ的參數(shù)化可能會帶來更好的結(jié)果。
掩碼矩陣M的結(jié)構(gòu),決定了對高效算法的實現(xiàn)。
更多實驗設(shè)置等細節(jié),請參閱原文。
其中A表示一個類Attention的交互矩陣,例如在線性注意力中,矩陣A就是Q和K的轉(zhuǎn)置矩陣的乘積矩陣;
而M是下三角形的因果掩碼矩陣,如線性注意力中的M的元素只能取值0和1。
帶6個頭的門控DeltaNet。
由于這些基準測試最初是為短序列(≤2K token)設(shè)計的,他們報告了序列長度為512、1024、2048以及(除NQ外)16K的結(jié)果。
塊內(nèi)計算(?=0):在每個子塊中,系統(tǒng)視其為無結(jié)構(gòu)數(shù)據(jù),并使用標準的O(C2)計算完成塊內(nèi)交互??偣灿蠺?C個子塊,因此整體塊內(nèi)計算成本為O(TC)。
近年來,涌現(xiàn)了大量致力于實現(xiàn)次二次方計算復(fù)雜度(sub-quadratic compute)和次線性內(nèi)存消耗(sub-linear memory)的高效替代方案。
這種結(jié)構(gòu)能對輸出O進行分塊處理,從而將算法整體復(fù)雜度降至O(T)。
其中,D表示僅在塊內(nèi)部有效的對角矩陣,包含T?C個塊,每個塊記錄子塊內(nèi)的交互信息。
如果隨著token位置增加,損失持續(xù)下降,說明模型能有效利用整個上下文。然而,如果損失在某一點后趨于平穩(wěn),則表明模型難以利用序列中過于靠后的信息。在這項分析中,使用了來自Book-3的3900萬個token。
然而,注意力機制存在「先天頑疾」:
修改版的Mamba-2,包含48個頭和MLP層;
作為應(yīng)用示例,研究人員展示了如何基于該框架構(gòu)建Mamba-2和Gated DeltaNet的對數(shù)線性版本。
團隊在現(xiàn)實世界的、需要大量回憶的任務(wù)上評估模型(表3)。
他們主要包括:線性注意力(linear attention)、狀態(tài)空間模型(state-space models)以及長卷積模型(long convolution models)。
對數(shù)線性注意力機制(log-linear attention)就是在矩陣M引入特定結(jié)構(gòu),讓計算復(fù)雜度在序列長度T上達到O(TlogT),內(nèi)存復(fù)雜度降低到O(logT)。
最后,他們在LongBench(表4)上評估了模型的性能。
在狀態(tài)空間對偶建??蚣苤?,這一方法已經(jīng)有所體現(xiàn)。
結(jié)構(gòu)矩陣,一統(tǒng)注意力變體
研究團隊報告了模型在每個token位置的損失,以評估其處理長上下文的能力(圖6)。
研究團隊首先在多查詢關(guān)聯(lián)回憶(MQAR)上進行實驗,這是一個用于評估模型上下文回憶能力的標準測試基準。
Fenwick樹是一種支持單點修改和區(qū)間查詢的,代碼量小的數(shù)據(jù)結(jié)構(gòu)
同樣,分塊計算機制可以擴展應(yīng)用于對數(shù)線性注意力機制。
在圖3中,左圖展示了矩陣M的分解方式,右圖則是對應(yīng)的分塊計算算法(算法1)。
這表明:提升效率的關(guān)鍵不只是去除softmax,而在于M本身是否具備合適的結(jié)構(gòu)。
在Level 0,模型對每個小塊內(nèi)部進行計算,采用的是相對于塊大小為二次復(fù)雜度的算法。由于每個塊本身較小,因此這一階段計算開銷低、效率高。
這兩個模型的主要區(qū)別在于它們對轉(zhuǎn)換矩陣A的參數(shù)化方式不同。
https://x.com/HanGuo97/status/1930789829094297859
結(jié)果顯示,將Mamba-2和門控DeltaNet擴展到它們的對數(shù)線性版本后,(平滑后的)損失在不同位置上均持續(xù)降低,表明長距離上下文利用能力有所提升。
在較簡單的單針任務(wù)中,對數(shù)線性Mamba-2在9個指標中的8個上優(yōu)于其線性版本。
為了實現(xiàn)這種多時間尺度的結(jié)構(gòu)化劃分,關(guān)鍵在于如何將前綴區(qū)間[0,t]分配給第t步的查詢向量。
在更具挑戰(zhàn)性的多針任務(wù)中,對數(shù)線性Mamba-2再次在9個指標中的8個上有所改進,而對數(shù)線性門控DeltaNet則在所有指標上均取得進步。
? 面向硬件的高效序列建模算法設(shè)計
研究團隊在Long-Data-Collections數(shù)據(jù)集上使用500億個token,從頭開始進行學(xué)術(shù)規(guī)模的語言建模預(yù)訓(xùn)練,序列長度為16K。
與傳統(tǒng)token級scan不同,它不再受限于內(nèi)存帶寬瓶頸,而是通過結(jié)構(gòu)優(yōu)化使狀態(tài)以低成本在線上傳遞。
圖1:標準線性注意力機制(上)與對數(shù)線性注意力機制(下)對比示意圖
如圖5所示,隨著序列長度和鍵值對數(shù)量的增加,DeltaNet的性能顯著下降,而對數(shù)線性DeltaNet(Log-Linear DeltaNet)依然保持高準確率。
Songlin Yang,是麻省理工學(xué)院計算機科學(xué)與人工智能實驗室(MIT CSAIL)的博士生,師從Yoon Kim教授。
為了解決這一問題,研究者采用了另一種的分段策略。
它具有以下特點:
另外值得一提的是,兩位第一作者都是華人,均麻省理工學(xué)院計算機科學(xué)與人工智能實驗室就讀。
- 訓(xùn)練效率:對數(shù)線性時間
batch size為2,注意力頭數(shù)為48,每個頭的維度為64,狀態(tài)維度為128,chunk size設(shè)置為64。
另外,在長卷積模型(long convolution models)中,可以通過使用快速傅里葉變換(FFT)進一步將復(fù)雜度降為O(TlogT),相較于原始的O(T2)計算量,實現(xiàn)了顯著的效率提升。
當序列長度達到131K時,訓(xùn)練吞吐量出現(xiàn)下降,這是由于引入了梯度檢查點(gradient checkpointing)以降低內(nèi)存使用所致。
這種層次結(jié)構(gòu)使模型能夠以更精細的方式關(guān)注最近的token,同時在解碼過程中實現(xiàn)對數(shù)級別的時間和內(nèi)存效率。
他們在一個包含1萬個樣本的數(shù)據(jù)集上訓(xùn)練了100個周期,并對學(xué)習(xí)率進行了調(diào)整。
盡管這些方法各有不同,但它們大多可以用以下方程統(tǒng)一表示:
論文鏈接:https://arxiv.org/abs/2506.04761
需要注意的是,softmax注意力在所有設(shè)置下都能達到滿分準確率。
團隊在WikiText困惑度和幾個零樣本常識推理基準上評估模型(表2)。這些都是短上下文任務(wù),因此對模型狀態(tài)大小不太敏感。
即便不使用softmax,如果采用無結(jié)構(gòu)的M(例如隨機下三角矩陣),注意力機制的計算和內(nèi)存復(fù)雜度,仍為與softmax注意力機制相當。