随着大语言模型规模的不断扩大,我们似乎已经习惯了一个发展模式:更大的模型、更多的参数、更强的性能。然而,这种"暴力美学"正在遭遇瓶颈。最近,Google DeepMind联合KAIST AI和Mila发布的一篇论文为我们展示了另一种可能性。
这篇题为"Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation"的论文,并没有沿着主流的注意力机制优化路线,而是选择重新审视Transformer架构的一个基本假设:所有Token是否真的需要经过相同深度的计算?
现有架构的根本问题
当前的Transformer架构存在两个核心问题:
计算资源分配不均:模型对简单词汇(如"the"、"and")和复杂概念(如"epistemology"、"superconductivity")投入相同的计算资源。模型消耗大量能量处理这些简单词汇,使用与复杂术语相同的计算深度,这显然不合理。
KV缓存瓶颈:这是真正的性能瓶颈。推理过程中,模型必须为每层的每个Token保存键值对的"短期记忆",这个KV缓存呈二次增长并消耗大量GPU内存。这是运行大型模型成本高昂的最主要原因,也是上下文窗口面临巨大工程挑战的根本所在。
多年来,业界尝试用量化、剪枝和专门缓存等技巧来解决这些问题,但这些都只是在根本有缺陷的设计上打补丁。
MoR:自适应计算的新范式
Mixture-of-Recursions(MoR)架构通过两个核心设计解决了这些根本问题:
递归计算机制
MoR用单个高度优化的"递归块(Recursion Block)"替代了传统的多层结构。当任务(Token)需要处理时,它被送到这一层。如果任务复杂需要更多处理,它不会被送到新的层,而是被送回同一个精英团队进行另一轮处理。这就是递归的精髓。
这种设计立即大幅减少了独特参数的数量,使模型本身变得更小、更节省内存。但真正的创新在于它如何决定哪些Token需要更多处理。
智能路由系统
MoR引入了小巧轻量的"路由器(Router)",充当智能调度员。它的任务是为每个输入的Token动态决定"递归深度":
- 简单功能词如"and":路由器决定一轮处理即可
- 复杂高价值名词如"photosynthesis":路由器分配三轮深度处理
模型现在能够在每个Token的基础上分配其最宝贵的资源------计算力,智能地将算力只分配给真正需要的地方。这是直接融入架构DNA的自适应计算(Adaptive Computation)。
多路性设计(Multiway)
这是MoR最创新的设计。作者认为在Transformer的一个Block里,Q(查询)、K(键)、V(值)和R(残差输入)虽然都来自上一层,但它们的使命完全不同。MoR为它们设立了独立的专属通道。
MoR将下一层Transformer块的输入解耦为独立的Q、K、V、R四个流,并为每个流设计了独立模块:
- 生成下一层的Query输入时,使用专门的DA_Q模块独立聚合前层信息
- 生成Value输入时,使用DA_V模块进行不同的聚合
这是一种深度方向的多头注意力机制,通过允许信息在不同层之间进行更丰富的Q/K/V交互,极大增加了层间的通信带宽。
技术实现细节
两种路由策略
论文探索了两种不同的路由实现方法:
专家选择(Expert-Choice) :每个递归深度作为"专家"。"深度3专家"可以挑选它最喜欢处理的Token(例如"最复杂的前10%Token")。这种方法保证了每个批次的计算负载固定,但在训练过程中有个副作用------它可以"预知"序列的未来内容,研究人员用巧妙的辅助损失来缓解这个问题。
Token选择(Token-Choice) :每个Token可以选择自己的处理路径(比如"我是Token'递归',选择进行3层深度处理")。这更灵活,但可能造成"负载不平衡"------如果句子中每个Token都很复杂,都想进行3层深度处理,就可能造成计算瓶颈。
研究人员发现,配备简单线性路由器和辅助损失的专家选择方法表现最佳。
缓存优化策略
MoR为KV缓存问题提供了两种智能策略:
递归级缓存(Recursion-wise Caching) :只有当前在递归循环中活跃的Token才能看到该循环的彼此内存。非活跃(已退出)的Token会被忽略。这大幅减少了内存使用和注意力机制的FLOPs。
递归共享(Recursive Sharing) :这是极致的内存节省模式。所有递归循环都重用第一轮循环的KV缓存。这在内存效率上达到最优,对预填充(prefill)很有利,但如果路由非常精确,性能可能会稍有下降。
这是经典的工程权衡:最高性能vs最高效率。MoR能够提供选择本身就是一个强大的特性。
实验结果与性能分析
论文通过isoFLOP分析验证了MoR的有效性。
核心性能指标
参数效率显著提升:在每个模型规模(从135M到1.7B参数),MoR架构都持续超越了标准递归模型。167M参数的MoR模型实现了比315M参数Transformer模型更好的准确性,同时速度更快。
推理速度翻倍:更小模型、更智能计算和减少的KV缓存相结合,带来了2.06倍的吞吐量提升。它可以用相同硬件以两倍速度生成文本。
打破扩展定律:MoR创造了新的帕累托前沿(Pareto frontier),彻底打破了速度、成本和准确性之间的传统权衡。在更大规模下,它不仅匹配甚至经常超越Transformer的性能,特别是在较低计算预算下。
详细实验数据
实验显示,MoR从预训练开始,损失就显著低于所有基线模型。在计算预算相同的情况下:
- 要达到MUDFormer的损失水平,普通Transformer需要约1.89倍的算力
- 在深层配置下,差距进一步扩大到2.08倍算力
传统Transformer越深收益越低的"边际递减"问题,在MoR这里得到了有效缓解。MoR即使在更深的配置下,依然能保持强劲的性能增长。
多任务表现
MoR在多个下游任务上都能以小搏大,匹敌甚至超越更大模型:
- 在需要长距离上下文关联的5-shot场景下,它甚至能与约4.2倍计算量的更大模型正面对决
- 应用到Vision Transformer(ViT)进行图像分类,效果同样显著
与MoE架构结合
MoR与混合专家(MoE)架构结合产生了1+1>2的效果。虽然两者都利用动态权重,但作用机制不同(跨层聚合vs层内专家选择),属于正交且互补的技术。实验结果显示,MoR连接对MoE模型同样有效,带来了额外的性能增益。
深层意义与影响
计算范式的转变
MoR代表了从"暴力美学"到"精确艺术"的转变:
- 从规模驱动到效率驱动:不再是谁的模型大谁厉害,而是谁的架构更聪明谁胜出
- 从静态计算到动态推理:递归循环就像模型内部的"思维链(chain of thought)",让模型能够在继续处理之前对复杂Token进行多步"思考"
潜在推理能力
论文还暗示了一个重要概念:潜在推理(latent reasoning)。递归循环充当模型内部的思维链,允许它在继续之前对复杂Token进行多步"思考"。这是朝着不仅能预测,更能推理的模型迈出的重要一步。
资源消耗优化
仅增加约0.23%的参数和0.4%的计算量,却在多种规模(405M--2.8B)和多种架构(Decoder-only/ViT)上稳定超越原Transformer,需要的额外资源几乎可以忽略不计。
局限性与展望
论文对其局限性很诚实。这些实验是在最大1.7B参数的模型上进行的,真正的考验将是看当扩展到当今主流的70B、180B甚至更大的模型时,这些令人印象深刻的收益是否依然存在。
不过,论文中的趋势线非常有希望。MoR和Transformer之间的性能差距随着模型变大而扩大,这表明该架构在根本上更具可扩展性。
结语
MoR为我们展示了一条超越"大力出奇迹"的发展路径。它提醒我们,真正的突破往往来自对基本假设的重新思考,而不是简单的规模扩展。
虽然标准Transformer不会立即退出历史舞台,但MoR已经种下了变革的种子。它提供了一条更智能、更便宜、更快速的前进道路,这很难反驳。这将是我们未来多年都会讨论的论文之一。
在当前AI发展的关键节点,这种从根本架构层面提升效率的思路显得尤为重要。它不仅为技术发展提供了新的可能性,也为整个行业的可持续发展指明了方向。