鷺羽 發自 凹非寺
量子位 | 公眾號 QbitAI
超越Transformer,谷歌推出全新底層架構——
Mixture-of-Recursions(MoR),注意不是MoE,它能推理速度提高2倍,而KV內存直接減半!
而且All in One,首次在單一框架中實現,用同一組參數處理不同任務的同時,進行動態分配計算資源。
就像給LLM開了個雙層增強buff,模型性能和效率全都要。
谷歌DeepMind聯合KAIST AI、Mila人團隊通過統一參數共享、自適應遞歸深度和高效KV緩存,在保持大模型性能的同時降低計算和內存成本,形成新的效率最優解。
不少網友甚至將它形容為Transformer Killer。
更有甚者表示,該架構的出現或許能代表,潛在空間推理也許將會成為下一個LLM突破所在。
具體MoR創新在哪些方面?下面一一詳解。
MoR:首次統一參數共享與自適應計算
Transformer的出現雖然帶來了優秀的少樣本泛化和推理能力,但隨之而來龐大的計算和內存需求還是讓訓練和部署成為難題。
目前相關優化方法主要是參數共享和自適應計算,但往往只能二選一,無法同時兼顧。
于是研究人員提出了遞歸混合模型MoR,可以在單一遞歸Transformer中同時融合兩個效率維度。
首先采用的遞歸Transformer,相比通過多個獨特層構建token的標準Transformer,其直接將模型劃分為遞歸塊,復用一組共享參數池。
主要包含三種參數共享策略:
- Cycle:循環復用層。
- Sequence:連續復用同一層。
- Middle變體:保留首尾層獨特參數,僅共享中間層。
利用參數共享,可以減少獨特參數數量,提升分布式訓練效率,并通過連續深度批處理消除計算 “氣泡”,提高推理吞吐量。
然后MoR采用動態路由機制,通過輕量級路由器為每個token分配不同遞歸深度,并集中計算在復雜token上,可以分為兩種:
- Expert-choice路由:將每個遞歸步驟視作 “專家”,基于隱藏狀態計算分數,通過閾值選擇合適的token繼續計算,采用層級過濾,優先為復雜token分配計算。
- Token-choice路由:初始階段為每個token分配固定遞歸深度,通過softmax/sigmoid確定專家,然后token按分配深度依次完成遞歸。
另外,MoR本身還借助KV緩存策略,管理鍵值的存儲與使用,同時確保內存效率的提升:
- Recursion-wise緩存:僅緩存當前遞歸步驟中活躍token的KV對,將注意力計算限制在本地緩存,降低內存和IO需求。
- Recursive KV共享:復用首次遞歸的KV對供后續步驟使用,確保所有token可訪問歷史上下文,減少預填充操作,而此時注意力計算量下降幅度較小。
在三種策略共同作用下,MoR通過在解碼每個token時直接進行潛在思考,路由機制讓模型能進行自適應推理,突破了先前固定思考深度的限制,實現了參數效率與自適應計算的統一。
性能超越Transformer
研究人員在135M到1.7B不同參數規模的模型上,就原始Transformer、遞歸基線模型和MoR進行對比實驗。
實驗表明,在相同的16.5e18 FLOPs的訓練預算下,MoR使用了將近50%的更少參數,但實現了更低的驗證損失和更高的平均少樣本準確率43.1%。
而vanilla模型的少樣本準確率此時是42.3%,說明MoR擁有更高的計算效率,可以在相同FLOPs預算下處理更多訓練token。
在訓練固定的20B token時,MoR也減少了25%的訓練FLOPs,訓練時間還縮短了19%,峰值內存減少25%。
另外,通過分析路由策略,發現Expert-choice路由的性能會在一定程度上優于Token-choice路由,說明路由粒度會對性能產生重要影響。
研究人員還對MoR進行了IsoFLOP分析,發現在135M、360M、730M和1.7B參數規模,以及2e18、5e18、16.5e18的FLOPs預算下,MoR始終優于遞歸基線模型。
雖然受遞歸容量瓶頸限制,在135M時略遜于vanilla模型,但在360M及規模的進一步擴大,MoR性能接近甚至超過Vanilla模型,且參數僅為后者的1/3,驗證了MoR的可擴展性。
在推理吞吐量評估上,360M規模的MoR模型在固定批大小和最大批大小設置下,均優于vanilla。
遞歸深度的增加讓更多token提前退出,KV緩存占用減少,吞吐量顯著提升,驗證了深度批處理與早期退出結合對部署效率的提升。
谷歌對底層架構的再思考
這已經不是谷歌第一次對底層架構進行重新思考,甚至可以說,谷歌始終希望利用架構創新重構計算范式,重新尋找AI新的平衡。
例如混合專家模型(MoE)正是這一理念的集中體現。
最早是在2017年,谷歌首次將MoE引入LSTM層,通過稀疏門控機制,僅激活部分專家網絡從事輸入處理,但仍能讓參數量高達137B的模型保持高效訓練。
后面推出的GShard將MoE與Transformer結合,可實現動態負載均衡,2021年的Switch Transformer又進一步地簡化了路由機制。
而Gemini 1.5 Pro就是采用的分層MoE架構,將專家網絡與多模態處理深度結合,可以處理更為復雜的多模態任務,訓練和服務效率也得到顯著提升。
MoE的底層邏輯設計突破了傳統全連接模型的計算缺陷,現在已經成為了許多超大規模模型的優先選擇,為應對算力瓶頸提供了新范式。
另外還有像TokenFormer等可擴展架構,將模型參數視為可學習的token,通過增量訓練無縫擴展模型規模,為未來千億級模型的低成本迭代提供了可能。
所以有網友認為,關于谷歌如今推出的MoR,會在未來徹底改變AI世界規則,會超越Transformer嗎?一切交給時間驗證。
參考鏈接:
[1]https://x.com/deedydas/status/1945313404958466519
[2]https://www.alphaxiv.org/abs/2507.10524
[3]https://x.com/reza_byt/status/1945498424536862841
[4]https://arxiv.org/abs/1701.06538
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.