Transformer終結者!Google DeepMind全新MoR架構問世,新一代魔王來了

新智元報導

編輯:Aeneas KingHZ

【新智元導讀】Transformer殺手來了?KAIST、Google DeepMind等機構剛剛發布的MoR架構,推論速度翻倍、記憶體減半,直接重塑了LLM的性能邊界,全面碾壓了傳統的Transformer。網友們直呼炸裂:又一個改變遊戲規則的炸彈來了。

就在剛剛,KAIST、Mila和Google DeepMind團隊等放出重磅炸彈——

一個名為Mixture-of-Recursions的全新LLM模型架構。

這個嶄新的架構,被業界認為有潛力成為Transformer殺手!

它的推論速度提升2倍,訓練FLOP減少,KV快取記憶體直接減半。

最終,在135M到1.7B的參數規模下,MoR直接劃出了一個新的柏拉圖前緣:相同的訓練FLOPs,但困惑度更低、小樣本準確率更高,並且吞吐量提升超過2倍。

全面碾壓傳統的Transformer!

MoR與Transformer性能對比圖

論文連結:https://arxiv.org/abs/2507.10524

其實,學界很早就發現,Transformer複雜度太高,算力需求驚人。

比如最近CMU大牛、Mamba架構作者Albert Gu就表示,Transformer模型能力的局限太大,所謂token就是胡扯。

Albert Gu關於Transformer的評論

而Google產品負責人Logan Kilpatrick公開指出了注意力機制的缺陷——不可能實現無限上下文,還強調必須要在核心架構層進行全面創新。

今天Google DeepMind的這項研究,和這些大牛的觀點不謀而合了。

對此,網友們紛紛表示實在炸裂。

有人預測,潛在空間推論可能會帶來下一個重大突破。

網友評論截圖

顯然,對於程式碼、數學、邏輯這類分層分解問題的任務,MoR都是一個改變遊戲規則的重磅炸彈。

MoR在不同任務中的表現

MoR與傳統模型對比

甚至還有人評論道:看起來像是Hinton的膠囊網路重生了。

網友評論:膠囊網路重生

分隔線

Google DeepMind放大招

遞迴魔法讓LLM瘦身還提速

LLM發展到如今,接下來該怎樣做?靠堆疊參數、增加層數,讓它更聰明嗎?

這項研究告訴我們:真正的高手,從來都不是靠堆料,而是靠設計的藝術。

這次他們做出的MoR全新架構,直譯出來是「遞迴混合體」,直接讓LLM推論速度噌噌翻倍!

所以,MoR究竟做了什麼?

簡而言之,它做了以下兩點。

1. 不對所有token一視同仁

LLM在處理文本時,會把句子拆成一個個token,不過,像「的」「是」「在」這種詞,並不需要多高深的推論,只需要一次前向傳播就夠了。而複雜的token,則需多次經過同一層堆疊。

MoR的聰明之處就在於,因token而異。

MoR的秘密武器是小型路由器,會為每個token的隱藏狀態打分,僅高分token的會繼續循環,其餘的則提前退出。

MoR的遞迴路由機制

2. 循環復用:一個模組搞定全部

傳統Transformer的思路就是不斷「堆疊層數」,堆得越高,處理能力越強。但這樣的代價,就是記憶體和算力:模型會越來越慢,越來越貴。

而MoR則反其道而行之,專門設計了共享區塊,每個token最多循環4次,只要路由器說「完成」,就提前跳出循環。

總之,如果說Transformer是一個龐大的工廠流水線,那MoR就更像一支高效的特種部隊。未來的AI,恐怕不會再比拼誰更重,而是誰更會分工調度、節省力氣。

而Google DeepMind,已經敏銳地把握到了這一點,給我們演示了這一趨勢的早期範本。

分隔線

真自適應計算

只靠Scaling law,把語言模型做大,確實能讓它能力暴漲,但訓練、部署所需的算力和成本也跟著暴漲。

現在常見的「瘦身」招數,要麼是把參數共享(省顯存),要麼是按需計算(省算力)。

但目前仍缺乏一種能將兩者有機融合的架構。

「遞迴混合」(Mixture-of-Recursions, MoR),充分發揮了遞迴Transformer的潛力(見圖1),成功融合了兩者。

圖1:遞迴混合(Mixture-of-Recursions,MoR)概覽

圖1:遞迴混合(Mixture-of-Recursions,MoR)概覽

(左)每個遞迴步驟包含固定層堆疊和決定token是否繼續遞迴的路由器(中間灰框區域)。

(中)完整模型結構,其中共享遞迴步驟根據路由決策對每個token最多應用𝑁𝑟次。

(右)展示token級遞迴深度的路由模式範例,顏色越深表示該token在遞迴區塊中的計算越活躍。底部數字以不同顏色標註各文本token的遞迴步數:1步、2步和3步 。

在統一的架構中,MoR同時實現了三種效率優化:

透過共享權重壓縮參數量;透過動態路由減少冗餘計算;透過智慧快取減少記憶體開銷。

分隔線

遞迴混合架構

在預訓練和推論過程中,MoR為每個token動態調整遞迴步驟,依賴兩大組件:

路由機制和KV快取策略。

標題裝飾

路由機制:專家選擇vs.token選擇

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

在這種模式下,遞迴深度可看作「專家」,每輪遞迴時這些專家會挑選它們認為最值得處理的top-k個token。

圖2a:專家選擇路由

為了讓遞迴更具一致性,團隊還引入了分層篩選機制(hierarchical filtering):只有在第r層中被選中的token,才有資格參與第r+1層的評估。

這一設計模擬了早退出機制,讓模型在訓練初期就能自動「篩選」出需要深度處理的token,把算力集中到最困難的token上。

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

具體來說,根據第1層的隱藏狀態,模型會計算每個專家的打分(如透過softmax或sigmoid)。

假設有𝑁𝑟個專家,每個專家對應一次遞迴,那麼模型將token分配給得分最高的專家。 token將被送入前i層遞迴,每層都按順序進行處理。

在這種方式中,token在進入網路時就確定好遞迴深度,同時避免了每一層的重新選擇,提升了推論效率。

表2左比較了兩種方法:

expert-choice路由的優點在於,它可以實現理想的計算負載均衡。然而,它容易資訊洩露。

相比之下,token-choice路由天然不會洩露資訊。但這種方式負載分配不均。

表2:路由策略與鍵值快取策略的比較

表2:路由策略與鍵值快取策略的比較。(左)兩種路由策略總結:專家選擇與令牌選擇;(右)快取策略相對於普通Transformer的相對成本效率

圖2:混合遞迴(MoR)的架構組件

圖2:混合遞迴(MoR)的架構組件。(a)專家選擇路由;(b)token自主選擇路由;(c)KV快取策略

標題裝飾

KV快取策略:按遞迴層快取vs.跨層共享

針對MoR模型,研究人員提出了兩種KV快取策略:

按遞迴層快取和跨遞迴共享

1.按遞迴層快取(見圖2c上)是「選擇性快取」:只有被路由到某一遞迴層的Token,才會在該層生成並儲存它的KV對。

注意力計算僅在當前遞迴層的快取內進行,這種設計有助於實現局部化計算,顯著提升了記憶體使用效率,並減少I/O負擔。

2.跨遞迴共享(見圖2c):只在第一個遞迴層生成並快取KV對,然後在之後所有層中重複使用。這種機制下,每一層參與注意力計算的Query數量可能會減少。

也就是說,所有Token無論在後續層是否繼續參與計算,都可以完整地訪問歷史上下文,無需重新計算。

表2右對比了兩種快取策略:

  • 按遞迴層快取:KV記憶體與I/O負擔,被壓縮為原來的一半左右。

  • 跨遞迴共享:只能線性壓縮注意力計算量,而且KV的讀寫次數較高,可能會成為性能瓶頸。

表3:MoR、遞迴Transformer、普通Transformer的比較

表3:在等計算量與等token數條件下,MoR、遞迴Transformer、普通Transformer的比較

分隔線

實驗

研究者從零開始預訓練模型,採用基於Llama的Transformer架構,參考了SmolLM開源模型的配置,在FineWeb-Edu的驗證集和六個few-shot基準測試集上進行了評估。

標題裝飾

主要結果

在相同訓練計算預算下,MoR以更少參數優於基準模型

在相同的訓練預算(16.5e18 FLOPs)下,研究者將MoR模型與標準Transformer和遞迴Transformer進行了對比。

不同計算預算下的驗證損失

在四種模型規模(135M、360M、730M和1.7B參數)下,不同計算預算對應的驗證損失對如圖

如表3所示,MoR模型採用專家選擇路由和兩次遞迴(Nr=2),不僅在驗證損失上更低,在few-shot平均準確率上也優於標準基準。

這得益於MoR更高的計算效率,使其在相同FLOPs預算下能處理更多的訓練token。

在相同資料量下,MoR用更少計算量仍優於基準模型

為了隔離架構差異的影響,研究者在固定訓練token數量(20B)的前提下進行分析。

結果證實,在少了25%訓練FLOPs的情況下,MoR模型(𝑁𝑟=2)仍然實現了更低的驗證損失和更高的準確率,超越了標準和遞迴基準。

與標準基準相比,MoR模型的訓練時間減少了19%,峰值記憶體使用量降低了25%。

這就要歸功於專門設計的分層過濾機制和按遞迴進行的注意力機制。

此外,MoR的性能也會受路由與快取策略的影響。

標題裝飾

IsoFLOP分析

評估一種新模型架構設計的核心標準之一,是其在模型規模和計算量增長時,性能是否能持續提升。

因此,研究團隊全面對比了MoR與標準Transformer(Vanilla)和遞迴Transformer。

實驗設定

實驗的模型規模有四種:135M、360M、730M 和1.7B 參數。

對於遞迴Transformer和MoR配置,遞迴次數統一設為3。

在三個不同的計算預算下,進行預訓練:2e18、5e18和16.5e18 FLOPs。

MoR架構:可擴展且參數高效

如圖3所示,在所有參數規模和算力預算下,MoR始終優於遞迴基準模型。

儘管在最小規模(135M)時,MoR表現略遜於標準Transformer,但隨著模型規模擴大,這一差距迅速縮小。

當參數規模超過360M時,MoR不僅能夠與標準Transformer持平,甚至在低計算量和中等計算預算下,表現更加優越。

總體而言,這些結果表明,MoR具備良好可擴展性和高參數效率,可替代舊架構。

標題裝飾

推論吞吐量評估

透過參數共享,MoR能利用連續深度批次處理技術,在推論階段顯著提升了吞吐量。

這種機制在解碼過程中,舊序列完成後立刻填入新tokens,持續保持了GPU的高利用率。

實驗設定

在360M參數規模下,在不同遞迴深度(2、3和4)下,團隊測試了MoR模型。

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

如圖4a所示,在兩種設定下,MoR變體的推論吞吐量都超過了普通Transformer。

遞迴深度越高,越多tokens會提早退出,從而減少KV快取的使用,進一步大幅提升了推論速度。例如,在最大批次設定(𝐵=Max)下,MoR-4速度可提升2.06倍

實驗表明,結合深度批次處理機制與提前退出策略,可大幅加速MoR模型在實際的推論速度。

圖4a:MoR推論吞吐量提升

消融實驗等更多內容和細節,請參閱原文。

參考資料:

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模型架構

次標籤:大型語言模型記憶體效率推論最佳化遞迴神經網路


上一篇:Meta團隊突破:大型模型「幻覺」暴跌至5%!一句話提問竟成關鍵?

下一篇:大模型再爆弱點!舊記憶忘不掉,新記憶分不清,準確度暴跌 | ICML'25

分享短網址