Transformer終結(jié)者！谷歌DeepMind全新MoR架構(gòu)問世，推理速度翻倍、內(nèi)存減半！

來源：新智元

【導(dǎo)讀】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)不謀而合了。

對此，網(wǎng)友們紛紛表示實(shí)在炸裂。

有人預(yù)測，潛在空間推理可能會帶來下一個(gè)重大突破。

顯然，對于代碼、數(shù)學(xué)、邏輯這類分層分解問題的任務(wù)，MoR都是一個(gè)改變游戲規(guī)則的重磅炸彈。

甚至還有人評論道：看起來像是Hinton的膠囊網(wǎng)絡(luò)重生了。

谷歌DeepMind放大招

遞歸魔法讓LLM瘦身還提速

LLM發(fā)展到如今，接下來該怎樣做？靠堆參數(shù)、加層數(shù)，讓它更聰明嗎？

這項(xiàng)研究告訴我們：真正的高手，從來都不是靠堆料，而是靠設(shè)計(jì)的藝術(shù)。

這次他們做出的MoR全新架構(gòu)，直譯出來是「遞歸混合體」，直接讓LLM推理速度噌噌翻倍！

所以，MoR究竟做了什么？

簡而言之，它做了以下兩點(diǎn)。

1. 不對所有token一視同仁

LLM在處理文本時(shí)，會把句子拆成一個(gè)個(gè)token，不過，像「的」「是」「在」這種詞，并不需要多高深的推理，只需要一次前向傳播就夠了。而復(fù)雜的token，則需多次經(jīng)過同一層棧。

MoR的聰明之處就在于，因token而異。

MoR的秘密武器是小型路由器，會為每個(gè)token的隱藏狀態(tài)打分，僅高分token的會繼續(xù)循環(huán)，其余的則提前退出。

2. 循環(huán)復(fù)用：一個(gè)模塊搞定全部

傳統(tǒng)Transformer的思路就是不斷「堆層」，堆得越高，處理能力越強(qiáng)。但這樣的代價(jià)，就是內(nèi)存和算力：模型會越來越慢，越來越貴。

而MoR則反其道而行之，專門設(shè)計(jì)了共享塊，每個(gè)token最多循環(huán)4次，只要路由器說「完成」，就提前跳出循環(huán)。

總之，如果說Transformer是一個(gè)龐大的工廠流水線，那MoR就更像一支高效的特種部隊(duì)。未來的AI，恐怕不會再比拼誰更重，而是誰更會分工調(diào)度、節(jié)省力氣。

而谷歌DeepMind，已經(jīng)敏銳地把握到了這一點(diǎn)，給我們演示了這一趨勢的早期范本。

真自適應(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ù)路由決策對每個(gè)token最多應(yīng)用次。

（右）展示token級遞歸深度的路由模式示例，顏色越深表示該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緩存策略。

路由機(jī)制：專家選擇vs.token選擇

受top-k門控機(jī)制的啟發(fā)，研究人員提出了專家選擇路由（Expert-choice routing）（見圖2a）。

在這種模式下，遞歸深度可看作「專家」，每輪遞歸時(shí)這些專家會挑選它們認(rèn)為最值得處理的top-k個(gè)token。

為了讓遞歸更具一致性，團(tuán)隊(duì)還引入了分層篩選機(jī)制（hierarchical filtering）：只有在第r層中被選中的token，才有資格參與第r+1層的評估。

這一設(shè)計(jì)模擬了早退出機(jī)制，讓模型在訓(xùn)練初期就能自動(dòng)「篩選」出需要深度處理的token，把算力集中到最困難的token上。

與前者不同，token選擇路由（token-choice routing）（見圖2b），在一開始就決定了每個(gè)token要走多少次遞歸流程。

具體來說，根據(jù)第1層的隱藏狀態(tài)，模型會計(jì)算每個(gè)專家的打分（如通過softmax或sigmoid）。

假設(shè)有個(gè)專家，每個(gè)專家對應(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路由天然不會泄露信息。但這種方式負(fù)載分配不均。

表2：路由策略與鍵值緩存策略的比較。(左)兩種路由策略總結(jié)：專家選擇與令牌選擇；(右)緩存策略相對于普通Transformer的相對成本效率

圖2：混合遞歸（MoR）的架構(gòu)組件。(a)專家選擇路由；(b)token自主選擇路由；(c)KV緩存策略

KV緩存策略：按遞歸層緩存vs.跨層共享

針對MoR模型，研究人員提出了兩種KV緩存策略：按遞歸層緩存和跨遞歸共享。

1.按遞歸層緩存（見圖2c上）是「選擇性緩存」：只有被路由到某一遞歸層的Token，才會在該層生成并存儲它的KV對。

注意力計(jì)算僅在當(dāng)前遞歸層的緩存內(nèi)進(jìn)行，這種設(shè)計(jì)有助于實(shí)現(xiàn)局部化計(jì)算，顯著提升了內(nèi)存使用效率，并減少I/O負(fù)擔(dān)。

2.跨遞歸共享（見圖2c）：只在第一個(gè)遞歸層生成并緩存KV對，然后在之后所有層中重復(fù)使用。這種機(jī)制下，每一層參與注意力計(jì)算的Query數(shù)量可能會減少。

也就是說，所有Token無論在后續(xù)層是否繼續(xù)參與計(jì)算，都可以完整地訪問歷史上下文，無需重新計(jì)算。

表2右對比了兩種緩存策略：

• 按遞歸層緩存：KV內(nèi)存與I/O負(fù)擔(dān)，被壓縮為原來的一半左右。

• 跨遞歸共享：只能線性壓縮注意力計(jì)算量，而且KV的讀寫次數(shù)較高，可能會成為性能瓶頸。

表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)測試集上進(jìn)行了評估。

主要結(jié)果

在相同訓(xùn)練計(jì)算預(yù)算下，MoR以更少參數(shù)優(yōu)于基線模型

在相同的訓(xùn)練預(yù)算（16.5e18 FLOPs）下，研究者將MoR模型與標(biāo)準(zhǔn)Transformer和遞歸Transformer進(jìn)行了對比。

在四種模型規(guī)模（135M、360M、730M和1.7B參數(shù)）下，不同計(jì)算預(yù)算對應(yīng)的驗(yàn)證損失對如圖

如表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的性能也會受路由與緩存策略的影響。

IsoFLOP分析

評估一種新模型架構(gòu)設(shè)計(jì)的核心標(biāo)準(zhǔn)之一，是其在模型規(guī)模和計(jì)算量增長時(shí)，性能是否能持續(xù)提升。

因此，研究團(tuán)隊(duì)全面對比了MoR與標(biāo)準(zhǔn)Transformer（Vanilla）和遞歸Transformer。

實(shí)驗(yàn)設(shè)置

實(shí)驗(yàn)的模型規(guī)模有四種：135M、360M、730M 和1.7B 參數(shù)。

對于遞歸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)。

推理吞吐量評估

通過參數(shù)共享，MoR能利用連續(xù)深度批處理技術(shù)，在推理階段顯著提升了吞吐量。

這種機(jī)制在解碼過程中，舊序列完成后立刻填入新tokens，持續(xù)保持了GPU的高利用率。

實(shí)驗(yàn)設(shè)置

在360M參數(shù)規(guī)模下，在不同遞歸深度（2、3和4）下，團(tuán)隊(duì)測試了MoR模型。

利用深度批處理，MoR顯著提升推理吞吐量

如圖4a所示，在兩種設(shè)置下，MoR變體的推理吞吐量都超過了普通Transformer。

遞歸深度越高，越多tokens會提早退出，從而減少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é)，請參閱原文。

參考資料：

https://arxiv.org/abs/2507.10524

https://x.com/rohanpaul_ai/status/1945342236310561091

https://www.rohan-paul.com/p/landmark-research-from-google-deepmind

為偉大思想而生！

AI+時(shí)代，互聯(lián)網(wǎng)思想（wanging0123)，

第一必讀自媒體