Transformer終結(jié)者!谷歌DeepMind全新MoR架構(gòu)問世,新一代魔王來了

0 評(píng)論 632 瀏覽 1 收藏 16 分鐘

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)化:

  1. 通過共享權(quán)重壓縮參數(shù)量;
  2. 通過動(dòng)態(tài)路由減少冗余計(jì)算;
  3. 通過智能緩存減少內(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é)議。

更多精彩內(nèi)容,請(qǐng)關(guān)注人人都是產(chǎn)品經(jīng)理微信公眾號(hào)或下載App
評(píng)論
評(píng)論請(qǐng)登錄
  1. 目前還沒評(píng)論,等你發(fā)揮!