Google DeepMind發(fā)布的遞歸混合（MoR）：

（1）推理速度提升2倍，因為token可以提前退出共享循環(huán)

（2）訓(xùn)練時減少了大量計算，每層的注意力機制FLOP減少約一半，相同預(yù)算可以訓(xùn)練更多數(shù)據(jù)

論文展示了一種全新的方法，讓大語言模型在自己的推理循環(huán)中學(xué)會規(guī)劃步驟，而不是硬編碼單一的推理鏈。

其次，它證明了混合器理念的可擴展性。通過混合多個小型遞歸專家，讓模型選擇下一個調(diào)用哪個專家，團隊在數(shù)學(xué)和編程基準(zhǔn)測試中提升了準(zhǔn)確性，而無需增加參數(shù)數(shù)量。

遞歸混合（MoR）在內(nèi)存中保持1個層堆棧，對困難的token進(jìn)行循環(huán)處理，在準(zhǔn)確性和速度上都超越了更大的普通模型。它通過讓一個小型路由器選擇每個token獲得多少次循環(huán)來實現(xiàn)這一點，然后僅為保持活躍的token保存緩存。更少的權(quán)重、更少的FLOP、更少的內(nèi)存，但在135M-1.7B規(guī)模上獲得了更好的困惑度。

擴展Transformer通常意味著堆疊更多層，并在內(nèi)存和計算上付出代價。MoR顛覆了這一習(xí)慣。它共享1個緊湊的塊，根據(jù)token難度運行最多4次，當(dāng)路由器判斷"完成"時提前跳出循環(huán)。

圖1描繪了這種token級別的階梯結(jié)構(gòu)，簡單單詞快速退出，復(fù)雜單詞繼續(xù)攀升。

通過共享層來減少內(nèi)存：遞歸Transformer在深度上綁定權(quán)重，但過去的工作仍然讓每個token通過每個循環(huán)。MoR保持權(quán)重綁定的思想，但添加了"中間循環(huán)"共享，因此只有第一層和最后一層保持獨特，而中間的所有層在每個循環(huán)中重用一個小型三元組層。這種選擇保持梯度穩(wěn)定，將獨特參數(shù)減少約3倍而不失表現(xiàn)力。由于相同權(quán)重重復(fù)使用，完全分片數(shù)據(jù)并行每步只收集一次權(quán)重，也減少了通信。

一個小型路由器為每個token評分，然后要么將其送回共享塊，要么讓它繼續(xù)前進(jìn)。中間的示意圖顯示循環(huán)重復(fù)到設(shè)定上限，因此簡單token提前退出，復(fù)雜token繼續(xù)工作。右側(cè)的熱圖證明大多數(shù)token跳過后期循環(huán)，因此計算和內(nèi)存隨任務(wù)難度而非序列長度縮放。這種選擇性深度帶來大約2倍的推理速度提升和約50%的鍵值緩存減少，同時保持準(zhǔn)確性。

讓路由器決定深度：一個輕量級線性門在第一個循環(huán)后立即為每個token評分。得分靠前的token獲得通過共享塊的另一次機會，其余token以當(dāng)前狀態(tài)繼續(xù)。訓(xùn)練使用輔助損失，將選中的分?jǐn)?shù)推向1，丟棄的分?jǐn)?shù)推向0，這樣門在推理時表現(xiàn)相同并避免因果泄漏。這種專家選擇路由保持計算預(yù)算固定，仍然將深度聚焦在有幫助的地方。

隨token而非層增長的緩存：標(biāo)準(zhǔn)解碼為每層每個token存儲鍵值對，一旦循環(huán)就會爆炸式增長。MoR只緩存當(dāng)前循環(huán)中仍然活躍的token，因此更深的循環(huán)觸及越來越小的內(nèi)存片段。如果內(nèi)存比計算更緊張，另一個變體為所有后續(xù)循環(huán)重用第一循環(huán)緩存，用困惑度的微小損失換取RAM和預(yù)填充延遲的大幅削減。

效果：在相同的16.5E FLOP預(yù)算下，具有3個循環(huán)的118M參數(shù)MoR在平均少樣本準(zhǔn)確率(43.1% vs 42.3%)和驗證損失方面擊敗了315M的普通基線。當(dāng)計算固定但數(shù)據(jù)變化時，MoR繼續(xù)獲勝，證明收益不僅僅是分詞器的運氣。token選擇路由比專家選擇落后約2個點，因此每循環(huán)的智能容量分配很重要。

生成時的速度提升：深度批處理將當(dāng)前處于同一循環(huán)的token分組，因此GPU保持滿載而不是等待長序列完成。使用3個循環(huán)，批次32時吞吐量提升1.60倍，當(dāng)批次增長到H100上的內(nèi)存限制時提升2.06倍。更多循環(huán)意味著更多早期退出，因此MoR-4在速度上比普通模型快一倍，困惑度僅下降一小部分。

擴展性和挑戰(zhàn) 等FLOP圖顯示MoR在135M時落后于普通模型，因為單個塊變得擁擠，但在360M時匹配，在730M和1.7B時領(lǐng)先，同時僅持有三分之一的參數(shù)。計算最優(yōu)擴展傾向于在預(yù)算攀升時擴大共享塊而非添加數(shù)據(jù)，因為每個額外參數(shù)在每個循環(huán)中都被重用。當(dāng)路由器面對非常不均勻的文本時，MoR仍需要更好的負(fù)載均衡，其固定的top-k使動態(tài)改變延遲目標(biāo)變得棘手。

論文標(biāo)題：Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation

論文鏈接：???https://arxiv.org/abs/2507.10524v1??

本文轉(zhuǎn)載自???AI帝國???，作者：無影寺