循环MoE架构深度解析,稠密Transformer之外第三条大模型路线

有一种模型架构,能同时装下海量知识和超长思考,还不用堆算力!

看完这篇,你再也不会觉得大模型只能靠堆参数变聪明!

本文拆解循环专家混合(Looped MoE)这种新型神经网络架构的底层逻辑,讲清楚它跟稠密Transformer(Dense Transformer)和普通MoE的本质区别,为什么参数量和算力(FLOPs)的比值相同,跑起来却完全是两码事,涉及显存带宽、动态路由、片上缓存(SRAM)等硬件底层机制。


一个看似聪明的算术题,藏着大陷阱

先讲个背景。现在训练大语言模型(Large Language Model),工程师最头疼两件事:一是参数不够多,模型记不住那么多知识;二是计算不够深,模型想不了那么多步。

为了解决第一个问题,谷歌那帮人在2017年前后把一个老古董技术翻了出来,叫MoE,中文叫“专家混合”(Mixture of Experts)。简单说,就是把一个大网络切成很多小块,每个小块叫一个“专家”,每次处理一个词的时候,只叫醒其中几个专家干活,其他专家继续睡觉。这样参数总量可以塞得很大,但每次真正跑的计算量(也就是FLOPs,浮点运算次数)却很小。

为了解决第二个问题,另一批人搞出了循环Transformer(Looped Transformer),也叫通用Transformer(Universal Transformer)。做法反过来:参数不多,就那么一层或几层,但让这一层反复跑很多遍,跑一遍不够就跑十遍。参数少,但计算深度管够。

那问题就来了。一个是“参数多、算力少”,一个是“算力多、参数少”,把这两个拼起来,参数和算力的比值不就正好抵消,回到普通稠密Transformer那种“参数和算力一比一”的状态了吗?如果真是这样,绕这么一大圈干嘛,直接用稠密模型不就完了?

这就是那道看似成立、其实完全站不住脚的算术题。答案很干脆:算术上确实抵消了,架构行为上差了十万八千里!


稠密模型的路,是修好的高速公路

要看懂差别在哪,先得搞清楚稠密Transformer到底怎么跑。

稠密的意思就是“全部激活”。你输入一个词,比如“苹果”,这个词的向量表示会依次穿过第1层、第2层、一直到第96层(以GPT-3为例)。每一层里所有的参数都要参与计算,一个都不能少。每一层的权重矩阵是不一样的,第1层可能在识别基础的语法关系,第50层可能在处理抽象语义,第90层可能在整合上下文做预测。

这条路径是完全固定的。不管你输入的是“苹果”还是“量子力学”,走的路线一模一样,经过的层数一模一样,激活的参数一模一样。这就像一条修好的高速公路,所有车都必须从入口进、按顺序过每一个收费站、从出口出。

稠密Transformer的深度和宽度是强绑定的。你想让模型能处理更复杂的推理,就得加层数,加层数就得加参数,加参数就得加显存、加算力、加成本。四件事捆在一起,掰不开。


专家混合的路,是有分岔的立交桥

MoE架构改变了这个逻辑。

一个典型的MoE层长这样:它有一个“路由器”(Router),还有一堆“专家”(Experts),比如8个、64个甚至上千个。每个专家其实就是一个小型的前馈网络(FFN)。当一个词的向量走到这一层时,路由器会看一眼这个向量,然后决定把它送给哪几个专家处理,通常是选打分最高的前2个(叫Top-2路由)。

Mixtral 8x7B就是个典型例子。它每层有8个专家,每次只激活2个。所以虽然它总参数量是47B左右,但每次推理只用到大概13B的参数在跑。参数量看起来吓人,实际算力开销小得多。

但是,MoE的路径依然是“横向”的。它在同一层里做选择,选完之后还是要往下一层走。整个网络的深度并没有增加。所以MoE解决的是“怎么塞更多知识进去”,没解决“怎么想得更深”。

而且MoE有个隐藏的成本,就是显存带宽。虽然每次只激活2个专家,但你不知道会激活哪2个,所以8个专家的权重都得放在显存里随时待命。这就是为什么MoE模型对显存要求那么高。

MoE省的是每词元计算量,但它并没有省内存带宽。所有专家的参数都得放在显存里等着被调用,推理的时候系统得随时准备把任何一个专家从显存里加载出来。这叫驻留成本。你在手机上跑一个专家混合模型,专家参数体积太大,只能放在慢速存储里,每次加载专家就成了整个推理过程的关键路径,延迟和能耗全都飙上去。

所以MoE的核心矛盾在于,它解决了计算量的问题,却把压力转移到了内存上。这个矛盾会引出后面所有事情。

有研究专门量化了这件事,在推理阶段,MoE面临双重惩罚,大量驻留的专家池会挤占高带宽显存的可用空间,留给键值缓存(KV cache)的余量变小,长上下文场景下这个问题尤其突出。你在训练阶段省下来的计算量,到了推理阶段可能被内存带宽吃掉。


循环Transformer的路,是一个反复回头的圈

再看循环Transformer,它的做法是:把一个Transformer层(或者几层)当作一个模块,让输入的向量反复穿过这同一个模块很多次。第一次穿过去,得到一个中间结果;把这个中间结果再喂回去,穿第二次;再穿第三次,直到某个停止条件满足。

关键点:每次穿过去用的都是同一套权重。参数量没变,但等效的计算深度增加了。

这个思路在DeepMind 2018年的通用Transformer论文里就有了,最近几年又被重新捞出来,因为大家发现这种架构在做算法推理、状态追踪这类任务上表现特别好。比如让模型追踪一个变量的值在一连串操作后变成什么,稠密模型经常翻车,循环模型反而能算对。

原因也不难理解。像“把变量x加1,重复100次”这种任务,本质上就是一个循环。用循环架构去解循环问题,天然对味儿。用稠密架构去解,等于要把100步循环全部“展开”成100层网络,明显浪费。

这听上去很美。你只有几层参数,但跑十遍就等于有了十层的计算深度。但这里有个硬伤。循环的次数增加,每词元计算量跟着线性增长,因为每跑一遍都要重新算一遍注意力机制和前馈网络。参数量没变,但计算量翻了好几倍。

更要命的是,循环本身会撞上表达力天花板。你在同一组层里反复跑,每次用的参数完全一样,模型从输入里提取的信息模式就固定下来了。第一次跑提取的特征和第三次跑提取的特征没有本质区别,重复劳动,边际收益递减。有研究精确地刻画了这种递减,循环执行的贡献和独立块之间的等效关系呈现次线性指数,循环每次迭代带来的容量增益在衰减。

这就像你让同一个人反复审同一份合同,第一遍他能找出大部分问题,第二遍还能发现一两个遗漏,第三遍基本上就是浪费时间了。人没变,方法没变,视角没变,重复再多遍也发现不了新东西。

所以循环架构单独用的时候,会陷进两个泥潭。一方面计算量跟着循环次数涨,另一方面表达力的边际收益在衰减。有研究发现,稠密循环模型在规模化的时候,扩展性并不比标准Transformer好。你多花的计算量,并没有换来同比例的模型质量提升。

事情到这里还只是铺垫。真正有意思的问题是:把MoE和Loop拼起来,会发生什么?


拼起来之后,路变成了一张动态的地图

先看拼起来的字面意思。循环MoE就是:把MoE层当作被循环的模块,让向量反复穿过这个MoE层很多次。

从参数和算力的比值看,MoE把比值压下去,Loop把比值拉上来,两者互相抵消,看起来跟稠密模型没区别。

但真正跑起来,路径完全变了样。

假设有个token叫A,第1次穿过MoE层时,路由器把它送给了3号专家和7号专家;第2次穿过时,因为向量已经被上一轮改变了,路由器可能把它送给1号和5号;第3次可能送给2号和8号。每一步走的专家都不一样,而且这个选择是根据前一步的中间状态决定的。

同一时间,另一个token叫B,走的可能是完全另一套序列:先4号和6号,再4号和4号,再3号和1号。

这已经不是一条固定的路径了。这是一张根据内容动态生成的地图。每个token根据自己的推理进度,决定下一步该调用哪个专家。

这个特性,稠密Transformer没有,普通MoE没有,普通Loop也没有,只有循环MoE才有!


动态路径带来的第一个反转:知识和思考彻底分家

在稠密Transformer里,知识和思考是绑在一起的。你想让模型知道更多事,就得加参数;加了参数,计算量也跟着涨。

在循环MoE里,这两件事第一次被拆开了。

专家的数量决定“知识容量”。你可以有1000个专家,每个专家存一部分世界知识,需要哪块调哪块。这跟循环多少次没关系。

循环的次数决定“思考深度”。同一个MoE层跑一次是浅思考,跑十次是深思考。这跟专家有多少也没关系。

这就意味着你可以做出一个“知识超级多但思考很浅”的模型,用来快速回答事实性问题;也可以做出一个“知识一般但思考很深”的模型,用来解数学题。甚至可以让同一个模型根据任务难度自己决定循环几次,简单题少循环,难题多循环。

OpenAI的o1系列和DeepSeek的R1,这类被称为“推理模型”的产品,本质上都在试图让模型学会“想得更久”。循环MoE架构正是从底层结构上支持这件事,而不是靠外挂的“思考链”提示工程硬凑出来。

这就是第一次翻转:算术上看是抵消了,功能上看是彻底解耦了!


动态路径带来的第二个反转:硬件账反过来算了

再看第二个反转,也是更违反直觉的一个:硬件层面。

现在的GPU,比如英伟达H100,有两种“内存”。一种叫HBM,就是显卡上那80G显存,容量大但速度相对慢;另一种叫SRAM,是芯片内部的高速缓存,速度快得多但容量极小,一般只有几十MB到一百多MB。

大语言模型做推理的时候,绝大部分时间不是卡在算力上,而是卡在“把权重从HBM搬到SRAM”这个搬运过程上。这就是所谓的“显存带宽瓶颈”(Memory-Bound)。你算力再强,权重搬不过来也白搭。

稠密Transformer每处理一个token,得把整个模型的所有权重从HBM搬一遍。假设模型有70B参数,那就得搬70B。搬运本身就把时间吃光了。

普通MoE稍好一点。虽然显存里塞了很多专家,但每次只用几个,理论上只需要搬那几个专家的权重。不过实际操作中,因为你不知道下一个token会用哪几个专家,缓存策略很难做,收益打折扣。

循环MoE有个特殊的优化空间:如果被激活的那几个专家的权重加起来足够小,能塞进SRAM里,那么循环的时候可以让它们一直待在SRAM不出去。第一次循环搬运一次,后面九次循环就完全不用再搬了,白赚九次“免费计算”。

这就把一个原本是“显存带宽瓶颈”的任务,硬生生转成了“算力瓶颈”,而算力恰恰是现代GPU最富裕的资源。等于把GPU从堵车状态解放出来,让它真的能跑起来。

同样的参数量、同样的FLOPs,稠密模型跑起来慢得要命,循环MoE跑起来可以快好几倍。这个差别在纸面算术里根本看不出来。


表征空间的差异:走同一间屋子和走一栋楼

从更抽象的层面看,还有第三个反转。

稠密Transformer的每一层是不同的“房间”。向量从第1层的房间走到第2层的房间,房间的布置(也就是权重)完全不一样,向量在每个房间里被以不同方式改造。层数多,就是楼层多,向量一路从1楼爬到96楼。

循环Transformer的循环,是让向量反复走进同一间屋子。屋子的布置没变,但每次进来向量本身已经不一样了,所以每次改造的结果也不一样。数学上,这叫“迭代动力学”,跟解方程的迭代法、跟寻找不动点(Fixed Point)的算法,是同一个思路。

有一类问题特别适合用迭代动力学求解:搜索问题、约束满足问题、逻辑推理问题。它们的共同特点是,你不知道要迭代多少步才能收敛,但你知道每一步该怎么走。稠密模型对付这类问题很吃力,因为它必须提前决定用多少层,用少了不够用,用多了浪费。循环模型可以根据问题难度动态决定循环几次,恰到好处。

把MoE加进来之后,事情更精彩。每次循环调用的“专家组合”都可以不一样,等于向量在同一间屋子里走一圈之后,下一次进屋能召唤出不同的家具。这已经不是简单的迭代,这是一种可以根据当前状态动态调整“计算算子”的机制。

严格说,这已经不像传统神经网络,更像一台可编程的机器:专家是指令集,路由器是控制流,循环是程序循环。这个类比不是我瞎编的,2022年DeepMind在一篇叫《通用Transformer》的后续研究里就明确讨论过神经网络和图灵机(Turing Machine)的对应关系。


一个绕不开的麻烦:训练难度

讲了这么多好处,得说说坏处。循环MoE不是免费午餐。

训练这种架构,比训练稠密Transformer难得多。原因有几个。

其一,路由器很难训。路由器的输出是离散的(选专家3还是选专家7),离散决策没法直接求导,得用一些近似技巧比如Gumbel-Softmax或者辅助损失(Auxiliary Loss)来引导。稍不注意就会出现“专家坍缩”问题:所有token都被路由到同一个专家,其他专家饿死。

其二,循环次数难定。循环几次算够?训练时如果固定循环次数,模型学不到“动态停止”的能力;如果动态决定循环次数,又得引入停止机制(Halting Mechanism),停止机制本身也难训。

其三,权重共享带来的梯度问题。同一套权重在循环中被使用多次,反向传播时梯度会累积,容易出现梯度爆炸或消失。得用特殊的初始化和归一化技巧才能稳住。

这也是为什么这个架构在纸面上早就被提出、实际部署的产品却少之又少。工业界主流还是稠密Transformer加普通MoE的组合,比如GPT-4被广泛猜测是MoE架构,Claude和Gemini也走类似路线,但都没有大规模引入循环机制。

有意思的是,2024年之后,随着推理模型(Reasoning Model)越来越受关注,业界开始重新审视循环架构。因为“让模型多想一会儿”这个需求,用循环机制在底层实现,比在应用层用长思考链硬凑要优雅得多。


一个还没答案的问题

现在把镜头拉回到那道算术题。

参数量和算力的比值相同,为什么行为完全不同?答案已经很清楚了:一个模型的能力不取决于它有多少参数、能跑多少算力,而取决于这些参数和算力以什么样的方式被组织、被调用、被激活

稠密Transformer把所有东西固定死:固定的层数、固定的路径、固定的激活模式。它简单、稳定、好训练,但缺乏灵活性。

MoE引入了横向的动态性:不同的token走不同的专家,让参数总量可以爆炸增长,激活成本却不增加。

Loop引入了纵向的动态性:让计算深度可以根据任务需求延长,不用为此增加参数。

循环MoE把两个维度的动态性叠在一起,形成了一个既能装海量知识、又能想得很深、还能根据内容自己决定计算路径的架构。这跟稠密Transformer只是“看起来参数算力比值一样”而已,骨子里是两种完全不同的计算范式。

一个还没有公开答案的实验数据是这样的:如果把一个稠密70B模型的所有权重,改造成一个循环MoE架构,让参数量和平均激活算力保持完全一致,在GSM8K(一个数学推理测试集)上,循环MoE到底能领先多少?业界内部有传闻说差距能到15个百分点以上,但没有任何一家公司公布过严格对照的数据。这个数字要是真的,意味着过去两年所有稠密模型的推理能力评测,前提可能就错了。这事,还没完!