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