Google發布MoR架構:兩倍推理速度,節省50%內存 原創 精華
Google DeepMind發布的遞歸混合(MoR):
(1)推理速度提升2倍,因為token可以提前退出共享循環
(2)訓練時減少了大量計算,每層的注意力機制FLOP減少約一半,相同預算可以訓練更多數據
論文展示了一種全新的方法,讓大語言模型在自己的推理循環中學會規劃步驟,而不是硬編碼單一的推理鏈。
其次,它證明了混合器理念的可擴展性。通過混合多個小型遞歸專家,讓模型選擇下一個調用哪個專家,團隊在數學和編程基準測試中提升了準確性,而無需增加參數數量。
遞歸混合(MoR)在內存中保持1個層堆棧,對困難的token進行循環處理,在準確性和速度上都超越了更大的普通模型。它通過讓一個小型路由器選擇每個token獲得多少次循環來實現這一點,然后僅為保持活躍的token保存緩存。更少的權重、更少的FLOP、更少的內存,但在135M-1.7B規模上獲得了更好的困惑度。
擴展Transformer通常意味著堆疊更多層,并在內存和計算上付出代價。MoR顛覆了這一習慣。它共享1個緊湊的塊,根據token難度運行最多4次,當路由器判斷"完成"時提前跳出循環。
圖1描繪了這種token級別的階梯結構,簡單單詞快速退出,復雜單詞繼續攀升。
通過共享層來減少內存:遞歸Transformer在深度上綁定權重,但過去的工作仍然讓每個token通過每個循環。MoR保持權重綁定的思想,但添加了"中間循環"共享,因此只有第一層和最后一層保持獨特,而中間的所有層在每個循環中重用一個小型三元組層。這種選擇保持梯度穩定,將獨特參數減少約3倍而不失表現力。由于相同權重重復使用,完全分片數據并行每步只收集一次權重,也減少了通信。
一個小型路由器為每個token評分,然后要么將其送回共享塊,要么讓它繼續前進。中間的示意圖顯示循環重復到設定上限,因此簡單token提前退出,復雜token繼續工作。右側的熱圖證明大多數token跳過后期循環,因此計算和內存隨任務難度而非序列長度縮放。這種選擇性深度帶來大約2倍的推理速度提升和約50%的鍵值緩存減少,同時保持準確性。
讓路由器決定深度:一個輕量級線性門在第一個循環后立即為每個token評分。得分靠前的token獲得通過共享塊的另一次機會,其余token以當前狀態繼續。訓練使用輔助損失,將選中的分數推向1,丟棄的分數推向0,這樣門在推理時表現相同并避免因果泄漏。這種專家選擇路由保持計算預算固定,仍然將深度聚焦在有幫助的地方。
隨token而非層增長的緩存:標準解碼為每層每個token存儲鍵值對,一旦循環就會爆炸式增長。MoR只緩存當前循環中仍然活躍的token,因此更深的循環觸及越來越小的內存片段。如果內存比計算更緊張,另一個變體為所有后續循環重用第一循環緩存,用困惑度的微小損失換取RAM和預填充延遲的大幅削減。
效果:在相同的16.5E FLOP預算下,具有3個循環的118M參數MoR在平均少樣本準確率(43.1% vs 42.3%)和驗證損失方面擊敗了315M的普通基線。當計算固定但數據變化時,MoR繼續獲勝,證明收益不僅僅是分詞器的運氣。token選擇路由比專家選擇落后約2個點,因此每循環的智能容量分配很重要。
生成時的速度提升:深度批處理將當前處于同一循環的token分組,因此GPU保持滿載而不是等待長序列完成。使用3個循環,批次32時吞吐量提升1.60倍,當批次增長到H100上的內存限制時提升2.06倍。更多循環意味著更多早期退出,因此MoR-4在速度上比普通模型快一倍,困惑度僅下降一小部分。
擴展性和挑戰 等FLOP圖顯示MoR在135M時落后于普通模型,因為單個塊變得擁擠,但在360M時匹配,在730M和1.7B時領先,同時僅持有三分之一的參數。計算最優擴展傾向于在預算攀升時擴大共享塊而非添加數據,因為每個額外參數在每個循環中都被重用。當路由器面對非常不均勻的文本時,MoR仍需要更好的負載均衡,其固定的top-k使動態改變延遲目標變得棘手。
論文標題:Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation
論文鏈接:???https://arxiv.org/abs/2507.10524v1??
本文轉載自???AI帝國???,作者:無影寺
