超越Transformer，谷歌推出全新底層架構(gòu)——

Mixture-of-Recursions（MoR），注意不是MoE，它能推理速度提高2倍，而KV內(nèi)存直接減半！

而且All in One，首次在單一框架中實(shí)現(xiàn)，用同一組參數(shù)處理不同任務(wù)的同時(shí)，進(jìn)行動(dòng)態(tài)分配計(jì)算資源。

就像給LLM開了個(gè)雙層增強(qiáng)buff，模型性能和效率全都要。

谷歌DeepMind聯(lián)合KAIST AI、Mila人團(tuán)隊(duì)通過統(tǒng)一參數(shù)共享、自適應(yīng)遞歸深度和高效KV緩存，在保持大模型性能的同時(shí)降低計(jì)算和內(nèi)存成本，形成新的效率最優(yōu)解。

不少網(wǎng)友甚至將它形容為Transformer Killer。

更有甚者表示，該架構(gòu)的出現(xiàn)或許能代表，潛在空間推理也許將會(huì)成為下一個(gè)LLM突破所在。

具體MoR創(chuàng)新在哪些方面？下面一一詳解。

MoR：首次統(tǒng)一參數(shù)共享與自適應(yīng)計(jì)算

Transformer的出現(xiàn)雖然帶來了優(yōu)秀的少樣本泛化和推理能力，但隨之而來龐大的計(jì)算和內(nèi)存需求還是讓訓(xùn)練和部署成為難題。

目前相關(guān)優(yōu)化方法主要是參數(shù)共享和自適應(yīng)計(jì)算，但往往只能二選一，無法同時(shí)兼顧。

于是研究人員提出了遞歸混合模型MoR，可以在單一遞歸Transformer中同時(shí)融合兩個(gè)效率維度。

首先采用的遞歸Transformer，相比通過多個(gè)獨(dú)特層構(gòu)建token的標(biāo)準(zhǔn)Transformer，其直接將模型劃分為遞歸塊，復(fù)用一組共享參數(shù)池。

主要包含三種參數(shù)共享策略：

1. Cycle：循環(huán)復(fù)用層。
2. Sequence：連續(xù)復(fù)用同一層。
3. Middle變體：保留首尾層獨(dú)特參數(shù)，僅共享中間層。

利用參數(shù)共享，可以減少獨(dú)特參數(shù)數(shù)量，提升分布式訓(xùn)練效率，并通過連續(xù)深度批處理消除計(jì)算 “氣泡”，提高推理吞吐量。

然后MoR采用動(dòng)態(tài)路由機(jī)制，通過輕量級(jí)路由器為每個(gè)token分配不同遞歸深度，并集中計(jì)算在復(fù)雜token上，可以分為兩種：

• Expert-choice路由：將每個(gè)遞歸步驟視作 “專家”，基于隱藏狀態(tài)計(jì)算分?jǐn)?shù)，通過閾值選擇合適的token繼續(xù)計(jì)算，采用層級(jí)過濾，優(yōu)先為復(fù)雜token分配計(jì)算。
• Token-choice路由：初始階段為每個(gè)token分配固定遞歸深度，通過softmax/sigmoid確定專家，然后token按分配深度依次完成遞歸。

另外，MoR本身還借助KV緩存策略，管理鍵值的存儲(chǔ)與使用，同時(shí)確保內(nèi)存效率的提升：

• Recursion-wise緩存：僅緩存當(dāng)前遞歸步驟中活躍token的KV對(duì)，將注意力計(jì)算限制在本地緩存，降低內(nèi)存和IO需求。
• Recursive KV共享：復(fù)用首次遞歸的KV對(duì)供后續(xù)步驟使用，確保所有token可訪問歷史上下文，減少預(yù)填充操作，而此時(shí)注意力計(jì)算量下降幅度較小。

在三種策略共同作用下，MoR通過在解碼每個(gè)token時(shí)直接進(jìn)行潛在思考，路由機(jī)制讓模型能進(jìn)行自適應(yīng)推理，突破了先前固定思考深度的限制，實(shí)現(xiàn)了參數(shù)效率與自適應(yīng)計(jì)算的統(tǒng)一。

性能超越Transformer

研究人員在135M到1.7B不同參數(shù)規(guī)模的模型上，就原始Transformer、遞歸基線模型和MoR進(jìn)行對(duì)比實(shí)驗(yàn)。

實(shí)驗(yàn)表明，在相同的16.5e18 FLOPs的訓(xùn)練預(yù)算下，MoR使用了將近50%的更少參數(shù)，但實(shí)現(xiàn)了更低的驗(yàn)證損失和更高的平均少樣本準(zhǔn)確率43.1%。

而vanilla模型的少樣本準(zhǔn)確率此時(shí)是42.3%，說明MoR擁有更高的計(jì)算效率，可以在相同F(xiàn)LOPs預(yù)算下處理更多訓(xùn)練token。

在訓(xùn)練固定的20B token時(shí)，MoR也減少了25%的訓(xùn)練FLOPs，訓(xùn)練時(shí)間還縮短了19%，峰值內(nèi)存減少25%。

另外，通過分析路由策略，發(fā)現(xiàn)Expert-choice路由的性能會(huì)在一定程度上優(yōu)于Token-choice路由，說明路由粒度會(huì)對(duì)性能產(chǎn)生重要影響。

研究人員還對(duì)MoR進(jìn)行了IsoFLOP分析，發(fā)現(xiàn)在135M、360M、730M和1.7B參數(shù)規(guī)模，以及2e18、5e18、16.5e18的FLOPs預(yù)算下，MoR始終優(yōu)于遞歸基線模型。

雖然受遞歸容量瓶頸限制，在135M時(shí)略遜于vanilla模型，但在360M及規(guī)模的進(jìn)一步擴(kuò)大，MoR性能接近甚至超過Vanilla模型，且參數(shù)僅為后者的1/3，驗(yàn)證了MoR的可擴(kuò)展性。

在推理吞吐量評(píng)估上，360M規(guī)模的MoR模型在固定批大小和最大批大小設(shè)置下，均優(yōu)于vanilla。

遞歸深度的增加讓更多token提前退出，KV緩存占用減少，吞吐量顯著提升，驗(yàn)證了深度批處理與早期退出結(jié)合對(duì)部署效率的提升。

谷歌對(duì)底層架構(gòu)的再思考

這已經(jīng)不是谷歌第一次對(duì)底層架構(gòu)進(jìn)行重新思考，甚至可以說，谷歌始終希望利用架構(gòu)創(chuàng)新重構(gòu)計(jì)算范式，重新尋找AI新的平衡。

例如混合專家模型（MoE）正是這一理念的集中體現(xiàn)。

最早是在2017年，谷歌首次將MoE引入LSTM層，通過稀疏門控機(jī)制，僅激活部分專家網(wǎng)絡(luò)從事輸入處理，但仍能讓參數(shù)量高達(dá)137B的模型保持高效訓(xùn)練。

后面推出的GShard將MoE與Transformer結(jié)合，可實(shí)現(xiàn)動(dòng)態(tài)負(fù)載均衡，2021年的Switch Transformer又進(jìn)一步地簡化了路由機(jī)制。

而Gemini 1.5 Pro就是采用的分層MoE架構(gòu)，將專家網(wǎng)絡(luò)與多模態(tài)處理深度結(jié)合，可以處理更為復(fù)雜的多模態(tài)任務(wù)，訓(xùn)練和服務(wù)效率也得到顯著提升。

MoE的底層邏輯設(shè)計(jì)突破了傳統(tǒng)全連接模型的計(jì)算缺陷，現(xiàn)在已經(jīng)成為了許多超大規(guī)模模型的優(yōu)先選擇，為應(yīng)對(duì)算力瓶頸提供了新范式。

另外還有像TokenFormer等可擴(kuò)展架構(gòu)，將模型參數(shù)視為可學(xué)習(xí)的token，通過增量訓(xùn)練無縫擴(kuò)展模型規(guī)模，為未來千億級(jí)模型的低成本迭代提供了可能。

所以有網(wǎng)友認(rèn)為，關(guān)于谷歌如今推出的MoR，會(huì)在未來徹底改變AI世界規(guī)則，會(huì)超越Transformer嗎？一切交給時(shí)間驗(yàn)證。