機器之心報道
編輯：冷貓
大型語言模型已展現出卓越的能力，但其部署仍面臨巨大的計算與內存開銷所帶來的挑戰(zhàn)。隨著模型參數規(guī)模擴大至數千億級別，訓練和推理的成本變得高昂，阻礙了其在許多實際應用中的推廣與落地。
這促使研究者們圍繞兩個主要方向積極探索高效化技術：一是通過權重共享提升參數效率，二是根據輸入復雜度動態(tài)分配計算資源，實現自適應計算。
而一切的基礎則是 Transformer 架構。這次谷歌又一次坐不住了，在遞歸 Transformer 的基礎上再次進化，發(fā)布了名為 Mixture-of-Recursions （MoR）的新 LLM 模型架構，有學者認為這個新架構「有潛力成為 Transformer 的殺手」。
MoR 框架引入了一種統(tǒng)一方法，同時實現參數共享和自適應計算。與以往孤立處理參數減少或自適應計算的方法不同，MoR 將動態(tài) token 級路由集成到參數高效的遞歸 Transformer 中，創(chuàng)建了一種協(xié)同架構，實現了「無需承擔大模型成本的大模型質量」。
在本研究中，作者提出了 Mixture-of-Recursions（MoR），一個統(tǒng)一框架，旨在充分發(fā)揮遞歸 Transformer 的潛力（如圖所示）。MoR 通過端到端訓練輕量級的路由，為每個 token 分配專屬的遞歸深度：即根據每個 token 所需的「思考」深度，決定共享參數模塊對其重復應用的次數，從而將計算資源精準地分配到最需要的部分。
這種基于 token 的動態(tài)遞歸機制，天然支持遞歸級的鍵值（KV）緩存，即根據每個 token 所分配的遞歸深度，有選擇地存儲和檢索對應的鍵值對。該定向緩存策略顯著降低了內存帶寬壓力，從而無需后處理操作就能提升推理吞吐量。
因此，MoR 在一個統(tǒng)一架構中同時實現了三項優(yōu)化：
1. 參數共享：通過權重綁定顯著減少參數量；
2. 計算路由：通過 token 路由避免冗余計算（減少無效的 FLOPs）；
3. 遞歸級緩存：按需緩存鍵值對，降低內存訪問開銷，提升吞吐效率。
MoR 以此兼顧了性能、效率和計算資源利用率，是遞歸 Transformer 架構的重要演進。
MoR 總覽
左圖顯示每個遞歸步驟由一組固定堆疊的層與一個路由（router）組成。路由會為每個 token 決定是繼續(xù)通過遞歸塊（進行下一輪遞歸），還是提前退出。圖中灰色方框表示的就是該遞歸模塊。
中圖展示了完整的模型結構。對于每個 token，模型會根據路由的決策，對共享的遞歸模塊最多重復應用 N_r 次，從而實現動態(tài)遞歸深度。
右圖 是一個 token 級的路由示例，展示了每個 token 的實際遞歸次數。顏色越深的單元格表示該 token 在當前遞歸步驟中被激活參與了計算。底部顯示了每個文本 token 實際經歷的遞歸次數，使用不同顏色編碼標注（如：紫色 1 次、藍色 2 次、紅色 3 次）。
核心架構與方法
遞歸 Transformer 模型基礎
MoR 基于遞歸 Transformer，通過跨多個層共享權重來實現參數效率，而不是為每一層維護唯一參數。
多種參數共享模式
循環(huán)：遞歸塊循環(huán)復用（例如，重復層 [0,1,2]）
序列：每個塊連續(xù)地重用相同的層（例如，[0,0,0]，[1,1,1]）
中周期 / 中序列：保留首層和尾層的獨特參數，同時共享中間層權重的變體
中周期策略在經驗上被證明最為有效，在參數效率和模型表達能力之間提供了最佳平衡。
動態(tài)路由機制
MoR 的核心創(chuàng)新在于其動態(tài)路由系統(tǒng)，該系統(tǒng)根據單個 token 的計算需求分配不同的遞歸深度。實現了兩種主要的路由策略：
1、在每一個遞歸步驟 r，路由器 G 會對每個 token 當前的隱藏狀態(tài)進行評估，并選擇 top-k 的 token 繼續(xù)進入下一層遞歸。
這會形成一個逐步收窄的計算漏斗：復雜的 token 會接受更深層的處理，而簡單的 token 會更早退出計算。然而，這種方法在訓練過程中會面臨因果性挑戰(zhàn)，因此需要采用一些專門的技術：
輔助路由（Auxiliary Router）：一個單獨的輕量網絡，用于在推理階段預測 top-k 的選擇；輔助損失（Auxiliary Loss）：采用二元交叉熵損失，鼓勵路由器輸出接近二值化，從而實現更清晰的閾值判斷。2、Token 選擇路由（Token-Choice Routing）：每個 token 在一開始就做出一次路由決策，以確定其完整的遞歸深度。這種方式天然滿足因果性要求，但需要引入負載均衡機制：
均衡損失（Balancing Loss）：正則項，用于鼓勵遞歸深度之間的均勻分布；路由偏置（Router Bias）：在訓練過程中動態(tài)調整各專家的偏置項。KV緩存優(yōu)化
MoR 引入了兩種針對動態(tài)深度模型的專用 KV 緩存策略：
遞歸式 KV 緩存：在每個遞歸深度選擇性地緩存 KV 對，注意力僅限于本地緩存的 token：
遞歸 KV 共享：所有 token 都會遍歷第一個遞歸塊，該塊的 KV 對會被緩存并在后續(xù)遞歸中重用：
實驗結果
MoR 在相同計算預算下以更少參數優(yōu)于基線模型
在等效的訓練預算（16.5 × 101? FLOPs）下，作者將 MoR 模型與 原始 Transformer 及遞歸 Transformer 進行了比較。如表 3 所示，MoR 模型采用「專家選擇路由」（expert-choice routing）和兩層遞歸（N_r = 2），實現了更低的驗證損失，并在少樣本學習平均準確率上超越基線（43.1% 對比 42.3%）。值得注意的是，盡管參數量減少近 50%，MoR 依然取得了更優(yōu)性能。這一優(yōu)勢歸因于其更高的計算效率，使其在相同 FLOPs 預算下處理了更多訓練 token。
在固定 FLOPs（16.5e18）與固定 token 數（20B）兩種設定下，對 MoR、遞歸 Transformer（Recursive）和標準 Transformer（Vanilla）進行了對比。
MoR 在相同數據下以更少計算超越基線模型
為了單獨分析架構差異，作者在固定訓練 token 數量（20B）下進行評估。結果顯示，MoR 模型（N_r = 2）在使用比基線模型少 25% 的 FLOPs 情況下，仍實現了更低的驗證損失和更高的準確率。
這種理論上的效率優(yōu)勢在實踐中帶來了顯著提升：相比于基線模型，MoR 的訓練時間減少 19%，峰值內存使用量下降 25%。這些改進來源于 MoR 的層級過濾機制與遞歸注意力策略，有效縮短了序列長度，實現了更優(yōu)的計算 - 準確率平衡，甚至在預訓練階段亦是如此。
MoR 性能受路由與緩存策略影響
作者還評估了 MoR 架構下的一些設計變體，特別是在輕量級的三層遞歸設置（N_r = 3）下仍具可比性的 Vanilla 模型。結果顯示，采用「token 選擇路由」時準確率較低（40.0%），而「專家選擇路由」表現更好（42.6%），表明路由粒度對模型性能具有關鍵影響。
此外，采用 KV 緩存共享略微降低了性能，但顯著提升了內存效率。在對內存占用有要求的部署場景中，這種性能與資源之間的權衡是值得接受的。
不同計算預算下的驗證損失，在四種模型規(guī)模（135M、360M、730M 和 1.7B 參數量）下進行對比。
MoR 是一種可擴展且參數高效的架構
如圖所示，MoR 在所有模型規(guī)模和計算預算下都穩(wěn)定優(yōu)于遞歸基線模型。盡管在最小規(guī)模（135M）下表現略低于原始 Transformer（可能由于遞歸容量瓶頸所致），但這種差距隨著模型規(guī)模的擴大迅速縮小。在參數量超過 360M 時，MoR 不僅能夠追平，而且在低至中等預算下常常超越原始 Transformer。
結果表明：MoR 是標準 Transformer 的一種可擴展且高效的替代方案。
它在使用顯著更少參數的情況下，取得了優(yōu)異的驗證性能，因而非常適合用于預訓練以及大規(guī)模部署。
MoR 通過連續(xù)深度批處理提升推理吞吐量
如圖 4a 所示，在上述兩種批處理設置下，所有 MoR 變體的推理吞吐量均高于 Transformer 基線模型。
隨著遞歸深度的增加，更多 token 可在早期階段退出，從而減少 KV 緩存的使用。這一機制顯著提升了推理速度：
例如，MoR-4 在最大批量配置下的吞吐量提升可達 2.06 倍。
雖然這伴隨著輕微的性能下降，但考慮到吞吐量的大幅提升，這種權衡是值得的。
實驗結果表明：將深度方向批處理范式與提前退出機制相結合，能夠顯著提升 MoR 在實際部署場景中的推理吞吐能力。
圖中 (a) 在固定批量和最大批量設置下，MoR 與原始 Transformer 在推理吞吐量與對數似然（log-likelihood）之間的帕累托前沿（Pareto frontier）表現對比。
(b) 參數共享策略對遞歸 Transformer 的影響：展示了遞歸次數為
京公網安備 11011402013531號