Transformer终结者?Google DeepMind新架构实现2倍推理速度和一半内存占用

随着大语言模型规模的不断扩大,我们似乎已经习惯了一个发展模式:更大的模型、更多的参数、更强的性能。然而,这种"暴力美学"正在遭遇瓶颈。最近,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发展的关键节点,这种从根本架构层面提升效率的思路显得尤为重要。它不仅为技术发展提供了新的可能性,也为整个行业的可持续发展指明了方向。

相关推荐
乌恩大侠8 分钟前
自动驾驶的未来:多模态传感器钻机
人工智能·机器学习·自动驾驶
光锥智能1 小时前
AI办公的效率革命,金山办公从未被颠覆
人工智能
GetcharZp1 小时前
爆肝整理!带你快速上手LangChain,轻松集成DeepSeek,打造自己的AI应用
人工智能·llm·deepseek
猫头虎2 小时前
新手小白如何快速检测IP 的好坏?
网络·人工智能·网络协议·tcp/ip·开源·github·php
GeeJoe2 小时前
凡人炼丹传之 · 我让 AI 帮我训练了一个 AI
人工智能·机器学习·llm
小和尚同志3 小时前
Dify29. 为你的 Dify API 穿层衣服吧
人工智能·aigc
不会学习的小白O^O3 小时前
神经网络----卷积层(Conv2D)
人工智能·深度学习·神经网络
努力一点9483 小时前
ubuntu22.04系统入门 linux入门(二) 简单命令 多实践以及相关文件管理命令
linux·运维·服务器·人工智能·gpu算力