Transformer終結(jié)者!谷歌DeepMind全新MoR架構(gòu)問世,新一代魔王來了
Transformer殺手來了?KAIST、谷歌DeepMind等機(jī)構(gòu)剛剛發(fā)布的MoR架構(gòu),推理速度翻倍、內(nèi)存減半,直接重塑了LLM的性能邊界,全面碾壓了傳統(tǒng)的Transformer。網(wǎng)友們直呼炸裂:又一個(gè)改變游戲規(guī)則的炸彈來了。
就在剛剛,KAIST、Mila和谷歌DeepMind團(tuán)隊(duì)等放出重磅炸彈——
一個(gè)名為Mixture-of-Recursions的全新LLM模型架構(gòu)。
這個(gè)嶄新的架構(gòu),被業(yè)內(nèi)認(rèn)為有潛力成為Transformer殺手!
它的推理速度提升2倍,訓(xùn)練FLOP減少,KV緩存內(nèi)存直接減半。
最終,在135M到1.7B的參數(shù)規(guī)模下,MoR直接劃出了一個(gè)新的帕累托前沿:相同的訓(xùn)練FLOPs,但困惑度更低、小樣本準(zhǔn)確率更高,并且吞吐量提升超過2倍。
全面碾壓傳統(tǒng)的Transformer!
論文鏈接:https://arxiv.org/abs/2507.10524
其實(shí),學(xué)界很早就發(fā)現(xiàn),Transformer復(fù)雜度太高,算力需求驚人。
比如最近CMU大牛、Mamba架構(gòu)作者Albert Gu就表示,Transformer模型能力的局限太大,所謂token就是胡扯。
而谷歌產(chǎn)品負(fù)責(zé)人Logan Kilpatrick公開指出了注意力機(jī)制的缺陷——不可能實(shí)現(xiàn)無限上下文,還強(qiáng)調(diào)必須要在核心架構(gòu)層進(jìn)行全面創(chuàng)新。
今天谷歌DeepMind的這項(xiàng)研究,和這些大牛的觀點(diǎn)不謀而合了。
對(duì)此,網(wǎng)友們紛紛表示實(shí)在炸裂。
有人預(yù)測(cè),潛在空間推理可能會(huì)帶來下一個(gè)重大突破。
顯然,對(duì)于代碼、數(shù)學(xué)、邏輯這類分層分解問題的任務(wù),MoR都是一個(gè)改變游戲規(guī)則的重磅炸彈。
甚至還有人評(píng)論道:看起來像是Hinton的膠囊網(wǎng)絡(luò)重生了。
一、谷歌DeepMind放大招,遞歸魔法讓LLM瘦身還提速
LLM發(fā)展到如今,接下來該怎樣做?靠堆參數(shù)、加層數(shù),讓它更聰明嗎?
這項(xiàng)研究告訴我們:真正的高手,從來都不是靠堆料,而是靠設(shè)計(jì)的藝術(shù)。
這次他們做出的MoR全新架構(gòu),直譯出來是「遞歸混合體」,直接讓LLM推理速度噌噌翻倍!
所以,MoR究竟做了什么?
簡而言之,它做了以下兩點(diǎn)。
1. 不對(duì)所有token一視同仁
LLM在處理文本時(shí),會(huì)把句子拆成一個(gè)個(gè)token,不過,像「的」「是」「在」這種詞,并不需要多高深的推理,只需要一次前向傳播就夠了。而復(fù)雜的token,則需多次經(jīng)過同一層棧。
MoR的聰明之處就在于,因token而異。
MoR的秘密武器是小型路由器,會(huì)為每個(gè)token的隱藏狀態(tài)打分,僅高分token的會(huì)繼續(xù)循環(huán),其余的則提前退出。
2. 循環(huán)復(fù)用:一個(gè)模塊搞定全部
傳統(tǒng)Transformer的思路就是不斷「堆層」,堆得越高,處理能力越強(qiáng)。但這樣的代價(jià),就是內(nèi)存和算力:模型會(huì)越來越慢,越來越貴。
而MoR則反其道而行之,專門設(shè)計(jì)了共享塊,每個(gè)token最多循環(huán)4次,只要路由器說「完成」,就提前跳出循環(huán)。
總之,如果說Transformer是一個(gè)龐大的工廠流水線,那MoR就更像一支高效的特種部隊(duì)。未來的AI,恐怕不會(huì)再比拼誰更重,而是誰更會(huì)分工調(diào)度、節(jié)省力氣。
而谷歌DeepMind,已經(jīng)敏銳地把握到了這一點(diǎn),給我們演示了這一趨勢(shì)的早期范本。
二、真自適應(yīng)計(jì)算
只靠Scaling law,把語言模型做大,確實(shí)能讓它能力暴漲,但訓(xùn)練、部署所需的算力和成本也跟著暴漲。
現(xiàn)在常見的「瘦身」招數(shù),要么是把參數(shù)共享(省顯存),要么是按需計(jì)算(省算力)。
但目前仍缺乏一種能將兩者有機(jī)融合的架構(gòu)。
「遞歸混合」(Mixture-of-Recursions, MoR),充分發(fā)揮了遞歸Transformer的潛力(見圖1),成功融合了兩者。
圖1:遞歸混合(Mixture-of-Recursions,MoR)概覽
(左)每個(gè)遞歸步驟包含固定層堆棧和決定token是否繼續(xù)遞歸的路由器(中間灰框區(qū)域)。
(中)完整模型結(jié)構(gòu),其中共享遞歸步驟根據(jù)路由決策對(duì)每個(gè)token最多應(yīng)用????次。
(右)展示token級(jí)遞歸深度的路由模式示例,顏色越深表示該token在遞歸塊中的計(jì)算越活躍。底部數(shù)字以不同顏色標(biāo)注各文本token的遞歸步數(shù):1步、2步和3步 ?。
在統(tǒng)一的架構(gòu)中,MoR同時(shí)實(shí)現(xiàn)了三種效率優(yōu)化:
- 通過共享權(quán)重壓縮參數(shù)量;
- 通過動(dòng)態(tài)路由減少冗余計(jì)算;
- 通過智能緩存減少內(nèi)存開銷。
三、遞歸混合架構(gòu)
在預(yù)訓(xùn)練和推理過程中,MoR為每個(gè)token動(dòng)態(tài)調(diào)整遞歸步驟,依賴兩大組件:
路由機(jī)制和KV緩存策略。
1. 路由機(jī)制:專家選擇vs.token選擇
受top-k門控機(jī)制的啟發(fā),研究人員提出了專家選擇路由(Expert-choice routing)(見圖2a)。
在這種模式下,遞歸深度可看作「專家」,每輪遞歸時(shí)這些專家會(huì)挑選它們認(rèn)為最值得處理的top-k個(gè)token。
為了讓遞歸更具一致性,團(tuán)隊(duì)還引入了分層篩選機(jī)制(hierarchical filtering):只有在第r層中被選中的token,才有資格參與第r+1層的評(píng)估。
這一設(shè)計(jì)模擬了早退出機(jī)制,讓模型在訓(xùn)練初期就能自動(dòng)「篩選」出需要深度處理的token,把算力集中到最困難的token上。
與前者不同,token選擇路由(token-choice routing)(見圖2b),在一開始就決定了每個(gè)token要走多少次遞歸流程。
具體來說,根據(jù)第1層的隱藏狀態(tài),模型會(huì)計(jì)算每個(gè)專家的打分(如通過softmax或sigmoid)。
假設(shè)有????個(gè)專家,每個(gè)專家對(duì)應(yīng)一次遞歸,那么模型將token分配給得分最高的專家。 token將被送入前i層遞歸,每層都按順序進(jìn)行處理。
在這種方式中,token在進(jìn)入網(wǎng)絡(luò)時(shí)就確定好遞歸深度,同時(shí)避免了每一層的重新選擇,提升了推理效率。
表2左比較了兩種方法:
expert-choice路由的優(yōu)點(diǎn)在于,它可以實(shí)現(xiàn)理想的計(jì)算負(fù)載均衡。然而,它容易信息泄露。
相比之下,token-choice路由天然不會(huì)泄露信息。但這種方式負(fù)載分配不均。
表2:路由策略與鍵值緩存策略的比較。(左)兩種路由策略總結(jié):專家選擇與令牌選擇;(右)緩存策略相對(duì)于普通Transformer的相對(duì)成本效率
圖2:混合遞歸(MoR)的架構(gòu)組件。(a)專家選擇路由;(b)token自主選擇路由;(c)KV緩存策略
2. KV緩存策略:按遞歸層緩存vs.跨層共享
針對(duì)MoR模型,研究人員提出了兩種KV緩存策略:按遞歸層緩存和跨遞歸共享。
1.按遞歸層緩存(見圖2c上)是「選擇性緩存」:只有被路由到某一遞歸層的Token,才會(huì)在該層生成并存儲(chǔ)它的KV對(duì)。
注意力計(jì)算僅在當(dāng)前遞歸層的緩存內(nèi)進(jìn)行,這種設(shè)計(jì)有助于實(shí)現(xiàn)局部化計(jì)算,顯著提升了內(nèi)存使用效率,并減少I/O負(fù)擔(dān)。
2.跨遞歸共享(見圖2c):只在第一個(gè)遞歸層生成并緩存KV對(duì),然后在之后所有層中重復(fù)使用。這種機(jī)制下,每一層參與注意力計(jì)算的Query數(shù)量可能會(huì)減少。
也就是說,所有Token無論在后續(xù)層是否繼續(xù)參與計(jì)算,都可以完整地訪問歷史上下文,無需重新計(jì)算。
表2右對(duì)比了兩種緩存策略:
按遞歸層緩存:KV內(nèi)存與I/O負(fù)擔(dān),被壓縮為原來的一半左右。
跨遞歸共享:只能線性壓縮注意力計(jì)算量,而且KV的讀寫次數(shù)較高,可能會(huì)成為性能瓶頸。
表3:在等計(jì)算量與等token數(shù)條件下,MoR、遞歸Transformer、普通Transformer的比較
四、實(shí)驗(yàn)
研究者從零開始預(yù)訓(xùn)練模型,采用基于Llama的Transformer架構(gòu),參考了SmolLM開源模型的配置,在FineWeb-Edu的驗(yàn)證集和六個(gè)few-shot基準(zhǔn)測(cè)試集上進(jìn)行了評(píng)估。
1. 主要結(jié)果
在相同訓(xùn)練計(jì)算預(yù)算下,MoR以更少參數(shù)優(yōu)于基線模型
在相同的訓(xùn)練預(yù)算(16.5e18 FLOPs)下,研究者將MoR模型與標(biāo)準(zhǔn)Transformer和遞歸Transformer進(jìn)行了對(duì)比。
在四種模型規(guī)模(135M、360M、730M和1.7B參數(shù))下,不同計(jì)算預(yù)算對(duì)應(yīng)的驗(yàn)證損失對(duì)如圖
如表3所示,MoR模型采用專家選擇路由和兩次遞歸(Nr=2),不僅在驗(yàn)證損失上更低,在few-shot平均準(zhǔn)確率上也優(yōu)于標(biāo)準(zhǔn)基線。
這得益于MoR更高的計(jì)算效率,使其在相同F(xiàn)LOPs預(yù)算下能處理更多的訓(xùn)練token。在相同數(shù)據(jù)量下,MoR用更少計(jì)算量仍優(yōu)于基線模型
為了隔離架構(gòu)差異的影響,研究者在固定訓(xùn)練token數(shù)量(20B)的前提下進(jìn)行分析。
結(jié)果證實(shí),在少了25%訓(xùn)練FLOPs的情況下,MoR模型(????=2)仍然實(shí)現(xiàn)了更低的驗(yàn)證損失和更高的準(zhǔn)確率,超越了標(biāo)準(zhǔn)和遞歸基線。
與標(biāo)準(zhǔn)基線相比,MoR模型的訓(xùn)練時(shí)間減少了19%,峰值內(nèi)存使用量降低了25%。
這就要?dú)w功于專門設(shè)計(jì)的分層過濾機(jī)制和按遞歸進(jìn)行的注意力機(jī)制。
此外,MoR的性能也會(huì)受路由與緩存策略的影響。
2. IsoFLOP分析
評(píng)估一種新模型架構(gòu)設(shè)計(jì)的核心標(biāo)準(zhǔn)之一,是其在模型規(guī)模和計(jì)算量增長時(shí),性能是否能持續(xù)提升。
因此,研究團(tuán)隊(duì)全面對(duì)比了MoR與標(biāo)準(zhǔn)Transformer(Vanilla)和遞歸Transformer。
實(shí)驗(yàn)設(shè)置
實(shí)驗(yàn)的模型規(guī)模有四種:135M、360M、730M 和1.7B 參數(shù)。
對(duì)于遞歸Transformer和MoR配置,遞歸次數(shù)統(tǒng)一設(shè)為3。
在三個(gè)不同的計(jì)算預(yù)算下,進(jìn)行預(yù)訓(xùn)練:2e18、5e18和16.5e18 FLOPs。
MoR架構(gòu):可擴(kuò)展且參數(shù)高效
如圖3所示,在所有參數(shù)規(guī)模和算預(yù)算力下,MoR始終優(yōu)于遞歸基線模型。
盡管在最小規(guī)模(135M)時(shí),MoR表現(xiàn)略遜于標(biāo)準(zhǔn)Transformer,但隨著模型規(guī)模擴(kuò)大,這一差距迅速縮小。
當(dāng)參數(shù)規(guī)模超過360M時(shí),MoR不僅能夠與標(biāo)準(zhǔn)Transformer持平,甚至在低計(jì)算量和中等計(jì)算預(yù)算下,表現(xiàn)更加優(yōu)越。
總體而言,這些結(jié)果表明,MoR具備良好可擴(kuò)展性和高參數(shù)效率,可替代舊架構(gòu)。
3. 推理吞吐量評(píng)估
通過參數(shù)共享,MoR能利用連續(xù)深度批處理技術(shù),在推理階段顯著提升了吞吐量。
這種機(jī)制在解碼過程中,舊序列完成后立刻填入新tokens,持續(xù)保持了GPU的高利用率。
實(shí)驗(yàn)設(shè)置
在360M參數(shù)規(guī)模下,在不同遞歸深度(2、3和4)下,團(tuán)隊(duì)測(cè)試了MoR模型。利用深度批處理,MoR顯著提升推理吞吐量
如圖4a所示,在兩種設(shè)置下,MoR變體的推理吞吐量都超過了普通Transformer。
遞歸深度越高,越多tokens會(huì)提早退出,從而減少KV緩存的使用,進(jìn)一步大幅提升了推理速度。例如,在最大批設(shè)置(??=Max)下,MoR-4速度可提升2.06倍。
實(shí)驗(yàn)表明,結(jié)合深度批處理機(jī)制與提前退出策略,可大幅加速M(fèi)oR模型在實(shí)際的推理速度。
消融實(shí)驗(yàn)等更多內(nèi)容和細(xì)節(jié),請(qǐng)參閱原文。
參考資料:
https://arxiv.org/abs/2507.10524
https://x.com/rohanpaul_ai/status/1945342236310561091
https://www.rohan-paul.com/p/landmark-research-from-google-deepmind
新智元報(bào)道 編輯:Aeneas KingHZ
本文由人人都是產(chǎn)品經(jīng)理作者【新智元】,微信公眾號(hào):【新智元】,原創(chuàng)/授權(quán) 發(fā)布于人人都是產(chǎn)品經(jīng)理,未經(jīng)許可,禁止轉(zhuǎn)載。
題圖來自Unsplash,基于 CC0 協(xié)議。
- 目前還沒評(píng)論,等你發(fā)揮!