機器之心報道

編輯：冷貓

大型語言模型已展現出卓越的能力，但其部署仍面臨巨大的計算與內存開銷所帶來的挑戰。隨著模型參數規模擴大至數千億級別，訓練和推理的成本變得高昂，阻礙了其在許多實際應用中的推廣與落地。

這促使研究者們圍繞兩個主要方向積極探索高效化技術：一是通過權重共享提升參數效率，二是根據輸入復雜度動態分配計算資源，實現自適應計算

而一切的基礎則是 Transformer 架構。這次谷歌又一次坐不住了，在遞歸 Transformer 的基礎上再次進化，發布了名為 Mixture-of-Recursions （MoR）的新 LLM 模型架構，有學者認為這個新架構「有潛力成為 Transformer 的殺手」

MoR 框架引入了一種統一方法，同時實現參數共享和自適應計算。與以往孤立處理參數減少或自適應計算的方法不同，MoR 將動態 token 級路由集成到參數高效的遞歸 Transformer 中，創建了一種協同架構，實現了「無需承擔大模型成本的大模型質量」。

• 論文標題：Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation
• 論文鏈接：alphaxiv.org/abs/2507.10524

在本研究中，作者提出了Mixture-of-Recursions（MoR），一個統一框架，旨在充分發揮遞歸 Transformer 的潛力（如圖所示）。MoR 通過端到端訓練輕量級的路由，為每個 token 分配專屬的遞歸深度：即根據每個 token 所需的「思考」深度，決定共享參數模塊對其重復應用的次數，從而將計算資源精準地分配到最需要的部分。

這種基于 token 的動態遞歸機制，天然支持遞歸級的鍵值（KV）緩存，即根據每個 token 所分配的遞歸深度，有選擇地存儲和檢索對應的鍵值對。該定向緩存策略顯著降低了內存帶寬壓力，從而無需后處理操作就能提升推理吞吐量。

因此，MoR 在一個統一架構中同時實現了三項優化：

1.參數共享：通過權重綁定顯著減少參數量；

2.計算路由：通過 token 路由避免冗余計算（減少無效的 FLOPs）；

3.遞歸級緩存：按需緩存鍵值對，降低內存訪問開銷，提升吞吐效率。

MoR 以此兼顧了性能、效率和計算資源利用率，是遞歸 Transformer 架構的重要演進

MoR 總覽

左圖顯示每個遞歸步驟由一組固定堆疊的層與一個路由（router）組成。路由會為每個 token 決定是繼續通過遞歸塊（進行下一輪遞歸），還是提前退出。圖中灰色方框表示的就是該遞歸模塊。

中圖展示了完整的模型結構。對于每個 token，模型會根據路由的決策，對共享的遞歸模塊最多重復應用 N_r 次，從而實現動態遞歸深度。

右圖 是一個 token 級的路由示例，展示了每個 token 的實際遞歸次數。顏色越深的單元格表示該 token 在當前遞歸步驟中被激活參與了計算。底部顯示了每個文本 token 實際經歷的遞歸次數，使用不同顏色編碼標注（如：紫色 1 次、藍色 2 次、紅色 3 次）。

核心架構與方法

遞歸 Transformer 模型基礎

MoR 基于遞歸 Transformer，通過跨多個層共享權重來實現參數效率，而不是為每一層維護唯一參數。

多種參數共享模式

循環：遞歸塊循環復用（例如，重復層 [0,1,2]）

序列：每個塊連續地重用相同的層（例如，[0,0,0]，[1,1,1]）

中周期 / 中序列：保留首層和尾層的獨特參數，同時共享中間層權重的變體

中周期策略在經驗上被證明最為有效，在參數效率和模型表達能力之間提供了最佳平衡。

動態路由機制

MoR 的核心創新在于其動態路由系統，該系統根據單個 token 的計算需求分配不同的遞歸深度。實現了兩種主要的路由策略：

1、在每一個遞歸步驟 r，路由器 G 會對每個 token 當前的隱藏狀態進行評估，并選擇 top-k 的 token 繼續進入下一層遞歸。

這會形成一個逐步收窄的計算漏斗：復雜的 token 會接受更深層的處理，而簡單的 token 會更早退出計算。然而，這種方法在訓練過程中會面臨因果性挑戰，因此需要采用一些專門的技術：

• 輔助路由（Auxiliary Router）：一個單獨的輕量網絡，用于在推理階段預測 top-k 的選擇；
• 輔助損失（Auxiliary Loss）：采用二元交叉熵損失，鼓勵路由器輸出接近二值化，從而實現更清晰的閾值判斷。

2、Token 選擇路由（Token-Choice Routing）：每個 token 在一開始就做出一次路由決策，以確定其完整的遞歸深度。這種方式天然滿足因果性要求，但需要引入負載均衡機制：

• 均衡損失（Balancing Loss）：正則項，用于鼓勵遞歸深度之間的均勻分布；
• 路由偏置（Router Bias）：在訓練過程中動態調整各專家的偏置項。

KV緩存優化

MoR 引入了兩種針對動態深度模型的專用 KV 緩存策略：

遞歸式 KV 緩存：在每個遞歸深度選擇性地緩存 KV 對，注意力僅限于本地緩存的 token：

遞歸 KV 共享：所有 token 都會遍歷第一個遞歸塊，該塊的 KV 對會被緩存并在后續遞歸中重用：

實驗結果

MoR 在相同計算預算下以更少參數優于基線模型

在等效的訓練預算（16.5 × 101? FLOPs）下，作者將 MoR 模型與 原始 Transformer 及遞歸 Transformer 進行了比較。如表 3 所示，MoR 模型采用「專家選擇路由」（expert-choice routing）和兩層遞歸（N_r = 2），實現了更低的驗證損失，并在少樣本學習平均準確率上超越基線（43.1% 對比 42.3%）。值得注意的是，盡管參數量減少近 50%，MoR 依然取得了更優性能。這一優勢歸因于其更高的計算效率，使其在相同 FLOPs 預算下處理了更多訓練 token。

在固定 FLOPs（16.5e18）與固定 token 數（20B）兩種設定下，對 MoR、遞歸 Transformer（Recursive）和標準 Transformer（Vanilla）進行了對比。

MoR 在相同數據下以更少計算超越基線模型

為了單獨分析架構差異，作者在固定訓練 token 數量（20B）下進行評估。結果顯示，MoR 模型（N_r = 2）在使用比基線模型少 25% 的 FLOPs情況下，仍實現了更低的驗證損失和更高的準確率。

這種理論上的效率優勢在實踐中帶來了顯著提升：相比于基線模型，MoR 的訓練時間減少 19%，峰值內存使用量下降 25%。這些改進來源于 MoR 的層級過濾機制與遞歸注意力策略，有效縮短了序列長度，實現了更優的計算 - 準確率平衡，甚至在預訓練階段亦是如此。

MoR 性能受路由與緩存策略影響

作者還評估了 MoR 架構下的一些設計變體，特別是在輕量級的三層遞歸設置（N_r = 3）下仍具可比性的 Vanilla 模型。結果顯示，采用「token 選擇路由」時準確率較低（40.0%），而「專家選擇路由」表現更好（42.6%），表明路由粒度對模型性能具有關鍵影響。

此外，采用 KV 緩存共享略微降低了性能，但顯著提升了內存效率。在對內存占用有要求的部署場景中，這種性能與資源之間的權衡是值得接受的。

不同計算預算下的驗證損失，在四種模型規模（135M、360M、730M 和 1.7B 參數量）下進行對比。

MoR 是一種可擴展且參數高效的架構

如圖所示，MoR 在所有模型規模和計算預算下都穩定優于遞歸基線模型。盡管在最小規模（135M）下表現略低于原始 Transformer（可能由于遞歸容量瓶頸所致），但這種差距隨著模型規模的擴大迅速縮小。在參數量超過 360M 時，MoR 不僅能夠追平，而且在低至中等預算下常常超越原始 Transformer。

結果表明：MoR 是標準 Transformer 的一種可擴展且高效的替代方案

它在使用顯著更少參數的情況下，取得了優異的驗證性能，因而非常適合用于預訓練以及大規模部署。

MoR 通過連續深度批處理提升推理吞吐量

如圖 4a 所示，在上述兩種批處理設置下，所有 MoR 變體的推理吞吐量均高于 Transformer 基線模型。

隨著遞歸深度的增加，更多 token 可在早期階段退出，從而減少 KV 緩存的使用。這一機制顯著提升了推理速度：

例如，MoR-4 在最大批量配置下的吞吐量提升可達 2.06 倍

雖然這伴隨著輕微的性能下降，但考慮到吞吐量的大幅提升，這種權衡是值得的。

實驗結果表明：將深度方向批處理范式與提前退出機制相結合，能夠顯著提升 MoR 在實際部署場景中的推理吞吐能力。

圖中 (a) 在固定批量和最大批量設置下，MoR 與原始 Transformer 在推理吞吐量與對數似然（log-likelihood）之間的帕累托前沿（Pareto frontier）表現對比。

(b) 參數共享策略對遞歸 Transformer 的影響：展示了遞歸次數為 = 3 的遞歸 Transformers 在四種不同參數共享策略下的負對數似然（NLL）表現。這些模型在 100 億 token 上預訓練。圖中紅色虛線表示完整規模的原始 Transformer，黑色虛線表示參數量約為其三分之一的基線模型。

(c) 四種不同架構在啟用 KV 緩存共享時的 NLL 表現對比。其中：

• MoR 的綠色表示禁用遞歸級 KV 緩存（recursion-wise KV caching），
• 藍色表示啟用遞歸級 KV 緩存共享（recursive KV sharing）；
• MoR-E 表示采用專家選擇路由（expert-choice routing）的 MoR，
• MoR-T 表示采用 token 選擇路由（token-choice routing）的 MoR。

另有消融實驗，請參照原論文。

總結與討論

MoR 通過證明參數效率和自適應計算可以協同結合，為高效的 LLM 架構建立了一種新的范例。該框架實現大型模型質量的同時顯著降低了計算需求的能力，解決了語言建模中的基本可擴展性挑戰。

最重要的是，MoR 為開發更具認知啟發性的 AI 系統提供了基礎。其在生成過程中能夠按每個 token 自適應地分配「思考深度」的能力，與語言模型中新興的潛在推理和內部思考研究相一致。這表明 MoR 可以作為一個平臺，用于研究模型如何學會在保持常規任務效率的同時，更深入地思考復雜問題。

該方法的確提高了推理速度，降低了 KV 緩存，但其是否能被稱為「Transformer 殺手」，仍然存在很多質疑。

更多細節內容，請參考原論文。