谷歌提出递归混合模型(MoR):通过参数共享和自适应计算提升Transformer效率

发布于 2025-8-8 08:28
浏览
0收藏

Transformer架构从根本上改变了人工智能的格局。从支撑大型语言模型(LLMs)的对话能力到实现实时语言翻译,Transformer已成为现代自然语言处理(NLP)应用的核心。然而,其成功背后伴随着巨大的代价:训练和部署往往需要超大规模数据中心的计算和内存资源。

这种计算需求给人工智能发展带来了巨大的经济压力。虽然科技巨头能够承担大规模GPU集群的成本,但小型组织和独立研究人员往往难以跟上步伐。因此,提高Transformer模型的效率已成为关键目标——不仅是为了提升AI能力,也是为了降低与训练和部署这些模型相关的财务和资源成本。

为应对这一挑战,已经出现了几种创新方法:

  • 递归Transformer:巧妙地在多个层中重用同一组权重的模型,大幅减少参数数量
  • 参数共享技术:如层绑定(layer tying)等在层间共享权重以缩小模型规模的方法
  • 提前退出机制:允许模型对简单输入跳过部分层,节省宝贵的计算资源

虽然这些方法各自解决了效率的特定方面,但它们通常要么专注于减少参数,要么专注于节省计算资源,很少能同时兼顾两者。而谷歌的递归混合模型(MoR) 正是在这一背景下应运而生。

MoR 通过将参数共享和自适应计算统一到单一框架中,实现了范式转变。它不仅仅是另一个增量改进——更是对如何让Transformer既高效又强大的根本性重新思考。

理解递归混合模型

谷歌提出递归混合模型(MoR):通过参数共享和自适应计算提升Transformer效率-AI.x社区

MoR示意图中不同颜色表示对每个token应用的递归次数不同。

MoR 的核心基于一个看似简单却意义深远的见解:不同的token需要不同程度的“思考”来有效处理。就像人类会在复杂概念上投入更多脑力,而对简单词汇则一带而过一样,MoR 允许Transformer按token自适应分配计算资源。

该架构通过三个相互关联的机制实现这一目标:

  1. 参数高效的递归块:有效重用权重
  2. 动态路由系统:基于token复杂度分配计算资源
  3. 优化的键值缓存:减少内存开销

参数共享策略

谷歌提出递归混合模型(MoR):通过参数共享和自适应计算提升Transformer效率-AI.x社区

基础:参数共享的工作原理

传统Transformer为每个层使用独特的参数,使得模型深度与参数数量呈线性关系。MoR 通过引入递归块打破了这一范式——递归块是可在多个步骤中重用的共享参数集。

这就像用一套单独的刀叉勺吃饭,而不是用一把勺子完成整个用餐过程。传统Transformer依赖不同的工具(层)执行每个操作,而递归块则在多个步骤中有效地重用同一工具。

四种策略:对比分析

谷歌提出递归混合模型(MoR):通过参数共享和自适应计算提升Transformer效率-AI.x社区

MoR 探索了四种不同的参数共享策略,每种策略都有其独特特点:

1. 循环策略(Cycle Strategy)

在循环方法中,权重在层间循环共享,形成重复模式。虽然简单直接,但这种方法可能会造成信息瓶颈,因为早期层需要同时处理初始处理和复杂推理任务。

这里的信息瓶颈源于在所有递归步骤中循环重用固定、相同的参数集。这迫使相同的共享层处理从初始特征提取到递归后期复杂迭代推理的广泛计算需求。这种僵化、无差别的重用限制了模型学习不同处理深度所需的多样化或专门化特征的能力。因此,共享参数块成为瓶颈,阻碍了信息在网络中高效且专门化的转换。这正是“中间循环”(Middle-Cycle)等策略被开发的原因——通过允许第一层和最后一层使用独特参数,在效率和表征灵活性之间取得平衡。

2. 序列策略(Sequence Strategy)

序列策略在进入下一层之前连续重用同一层。这允许在每种层类型中实现更深的专门化,但可能限制模型在不同处理类型之间切换的能力

3. 中间循环策略(Middle-Cycle Strategy)——最优选择

这正是MoR 的优势所在。中间循环策略为关键的输入处理层和输出生成层保留独特权重,而在中间层之间循环共享权重。这种策略通过在最重要的位置保持完整能力,同时在最有益的地方实现参数共享,平衡了效率和性能

4. 中间序列策略(Middle-Sequence Strategy)

与中间循环策略类似,但中间层采用序列共享方式。虽然可行,但其性能不如中间循环策略

为什么中间循环策略占优

实证结果一致表明中间循环策略是最优的。这反映了对Transformer信息处理方式的深刻理解:

  • 输入层需要独特参数来处理多样化的原始输入表征
  • 中间层可以共享参数,因为它们执行相似的抽象推理任务
  • 输出层需要独特参数将抽象表征映射回具体输出

这种架构选择减少了约三分之二的参数,同时保持了具有竞争力的性能——这是效率方面的显著成就。

路由机制:自适应token级计算

谷歌提出递归混合模型(MoR):通过参数共享和自适应计算提升Transformer效率-AI.x社区

固定计算的问题

传统Transformer对每个token施加相同的计算量,无论其复杂度如何。处理“the”和处理“量子纠缠”所获得的计算预算相同。这种统一方法本质上是浪费的。

谷歌提出递归混合模型(MoR):通过参数共享和自适应计算提升Transformer效率-AI.x社区

将全部处理能力用于像“.”、“The”或“—”这样的token可能被认为是不必要的——就像上面提到的类比一样。

MoR 通过复杂的路由机制解决了这一问题,该机制根据token的复杂度动态为单个token分配不同的递归深度。该系统包括两种主要路由策略:

专家选择路由(Expert-Choice Routing):层级token选择

专家选择路由将每个递归深度视为一个“专家”,由专家选择最具挑战性的token进行处理。其工作原理如下:

  1. token评分:在每个递归步骤r,轻量级路由器检查隐藏状态并计算每个token的复杂度分数
  2. Top-k选择:只有评分最高的k个token进入下一个递归深度
  3. 层级过滤:token必须在深度r被选中,才有资格进入深度r+1

这形成了一个自然的层级结构:

  • 所有token都经过第一次递归(基础处理)
  • 只有复杂token进入第二次递归(中间推理)
  • 最具挑战性的token进入最深层递归(高级推理)

这里,信息泄露是专家选择路由中的一个问题,即选择top-k专家的过程无意中使用了未来token的信息,违反了因果关系。这可能会影响模型在推理时的性能,解决方法是使用辅助机制或正则化来确保因果处理。

token选择路由(Token-Choice Routing):基于承诺的处理

token选择路由采用不同的方法,从一开始就为每个token分配完整的递归路径:

  1. 初始评估:在第一次递归步骤中,路由器分析每个token并将其分配给特定专家(递归深度)
  2. 完整路径执行:token随后沿着其分配的路径完成所有指定的递归步骤
  3. 无信息泄露:与专家选择路由不同,未来token不会影响早期路由决策,保持因果一致性

这种方法避免了专家选择路由中可能出现的因果违规,但需要仔细的负载均衡,以防止某些专家负担过重而其他专家闲置。

辅助组件:使路由工作

两种路由策略都需要复杂的辅助机制:

对于专家选择路由

  • 辅助损失:通过惩罚不良路由决策,确保路由器做出可靠选择
  • 线性路由器架构:轻量级、高效的路由器,无需大量计算开销即可确定token复杂度

对于token选择路由

  • 平衡损失:通过鼓励token在专家间均匀分布,防止负载不平衡
  • Z损失:通过防止路由权重变得过于极端来稳定训练

谷歌提出递归混合模型(MoR):通过参数共享和自适应计算提升Transformer效率-AI.x社区

内存优化:KV缓存策略

谷歌提出递归混合模型(MoR):通过参数共享和自适应计算提升Transformer效率-AI.x社区

KV缓存策略:矩阵中的每个方块表示一个token(行)是否关注另一个token的缓存键(列)。在“递归级KV缓存”(上)中,仅缓存每个递归步骤中当前选中(未被丢弃)的token的键(蓝色),且注意力仅限于这些条目。在“递归KV共享”(下)中,所有先前token的键在第一次递归步骤中被缓存(紫色),并在后续递归步骤中共享以进行注意力计算。

KV缓存的挑战

键值(KV)缓存对Transformer效率至关重要,它存储计算出的注意力键和值以避免重复计算。然而,传统缓存方法可能消耗大量内存,尤其是对于长序列。MoR 引入了两种创新缓存策略,与其递归架构无缝协作。

递归级缓存(Recursion-wise Caching):选择性内存管理

在递归级缓存中,MoR 为每个递归步骤维护单独的KV缓存,但巧妙之处在于:只有在每个深度实际被处理的token,其KV对才会在该级别被缓存。

这创建了一个自然的层级内存结构:

  • 深度1:所有token都有缓存的KV对(完整序列)
  • 深度2:只有复杂token有缓存的KV对(缩减序列)
  • 深度3:只有最复杂的token有缓存的KV对(最小序列)

内存节省非常显著。对于具有N_r次递归的模型,总KV内存使用量降至约为普通Transformer的​​(N_r + 1)/(2N_r)​​。对于3次递归模型,这意味着KV内存使用量减少67%。

递归KV共享(Recursive KV Sharing):最大化内存效率

递归KV共享通过在所有后续递归深度中重用第一次递归步骤的KV对,进一步优化了内存效率。这种方法利用了所有token都经过第一个递归块这一事实。

其工作原理如下:

  1. 第一次递归:所有token生成并缓存其KV对
  2. 后续递归:重用这些缓存的KV对,仅更新查询向量
  3. 注意力计算:后期递归可以使用共享的KV对关注完整上下文

这种策略提供了最大的内存效率,将KV内存使用量减少到原始需求的​​1/N_r​​。对于3次递归模型,这意味着内存减少83%。

其代价是性能略有下降,因为键和值表征不会随递归深度演变。然而,实证结果表明这种下降很小,使其成为内存受限环境的理想选择。

架构整合:各部分如何协同工作

统一框架

MoR 的成功不在于任何单一创新,而在于它如何将多种效率技术整合到一个连贯的架构中。让我们追踪token序列如何流经MoR 模型:

步骤1:初始处理

所有token通过独特的第一层进入,接受标准Transformer处理。这确保了强大的输入表征,不受参数共享的限制。

步骤2:路由决策

路由器分析每个token的隐藏状态并做出路由决策:

  • 专家选择:计算复杂度分数并选择top-k token进入下一次递归
  • token选择:为每个token分配特定的递归深度路径

步骤3:递归处理

token沿着其分配的递归路径流动:

  • 简单token(如“the”、“and”)可能在1-2次递归后完成处理
  • 复杂token(如技术术语、命名实体)接受3-4次递归
  • 每次递归应用相同的共享参数块,但处理的token集逐渐筛选

步骤4:内存管理

KV缓存策略启动:

  • 递归级:为每个活跃token集维护单独的缓存
  • 递归共享:在所有递归深度中重用初始KV对

步骤5:输出生成

处理后的token通过独特的最后一层进行输出生成,确保尽管中间层有参数共享,仍能获得高质量结果。

计算效率提升

这种整合方法在多个维度上实现了效率提升:

1. 参数效率:通过递归块减少约67%的独特参数

谷歌提出递归混合模型(MoR):通过参数共享和自适应计算提升Transformer效率-AI.x社区

与传统Transformer架构相比,MoR模型在不牺牲性能的情况下显著减少了参数数量。标准的Vanilla Transformer有3.15亿参数,而具有3次递归的MoR模型(Expert Cache M-Cyc 3)仅使用1.18亿独特参数——减少了67%。尽管如此,MoR在验证困惑度和少样本准确率方面始终达到或超过基线。例如,在相同训练计算量下,具有2次递归的MoR(M-Cyc 2)实现了更高的少样本准确率(43.1%),超过3.15亿参数的Vanilla基线(42.3%),同时使用的参数几乎少了50%。在更大规模(3.6亿+)下,MoR继续与Vanilla模型匹敌或超越,同时保持约三分之一的独特参数。

2. 计算效率:通过选择性token处理减少FLOPs

3. 内存效率:大幅降低KV缓存需求

谷歌提出递归混合模型(MoR):通过参数共享和自适应计算提升Transformer效率-AI.x社区

4. 吞吐量提升:连续的深度批处理消除了GPU空闲时间

谷歌提出递归混合模型(MoR):通过参数共享和自适应计算提升Transformer效率-AI.x社区

实际意义和现实应用

部署场景

MoR 的效率提升使其在多种部署场景中特别有价值:

边缘计算

凭借参数减少和内存效率,MoR 模型可以在资源受限的设备上运行,而传统大型模型在这些设备上是不可能运行的。

云成本优化

更小的内存占用和更高的吞吐量直接转化为AI服务提供商更低的云计算成本。

实时应用

吞吐量的提升使实时处理成为可能,如实时翻译、交互式聊天机器人和动态内容生成。

训练效率

MoR 的优势不仅限于推理,还延伸到训练:

  • 减少内存需求:训练期间降低峰值内存使用
  • 改进FSDP效率:参数共享与全分片数据并行(FSDP)训练配合极佳
  • 更快收敛:自适应计算可带来更高效的学习

技术挑战与解决方案

专家选择路由中的因果性

专家选择路由面临一个根本挑战:训练时,路由器在做选择决策时可以看到未来的token,但推理时却不能。这种“信息泄露”可能导致推理性能不佳。

MoR 尝试通过以下方式解决:

  1. 辅助损失:训练路由器仅基于可用信息做出良好决策
  2. 辅助路由器:专门为推理时行为训练的独立路由器

token选择路由中的负载均衡

token选择路由可能遭受负载不平衡问题,即某些专家(深度递归)接收大量token,而其他专家接收很少。这导致资源利用效率低下。

解决方案包括:

  1. 平衡损失:明确惩罚不均衡的token分布
  2. 无损失方法:使用路由器偏置鼓励均衡分配
  3. 动态容量:根据实际负载调整专家容量

KV缓存一致性

当token在递归处理中提前退出时,其KV对可能在后续token的注意力计算中缺失。MoR 通过以下方式解决这一问题:

  1. 结构化缓存:在递归深度间维持一致的缓存结构
  2. 选择性注意力:将注意力限制在可用的缓存对
  3. 共享上下文:确保所有token都能访问必要的上下文信息

展望未来:未来意义

缩放定律与MoR

MoR 表现出与传统Transformer不同的缩放行为。研究表明,MoR 从增大模型规模中获益比从增加训练数据中更多,这为递归架构提出了新的计算最优缩放策略。

这对我们进行模型开发的方式具有深远意义:

  • 资源分配:倾向于更大的模型和更短的训练周期
  • 架构设计:优先考虑参数效率而非原始参数数量
  • 训练策略:关注模型容量而非数据量

整合机会

MoR 的模块化设计使其与其他效率技术兼容:

  • 量化:递归块可以被量化以进一步节省内存
  • 剪枝:结构化剪枝可以进一步降低计算需求

在这种情况下,剪枝可能不是减少计算的有效方法。由于递归Transformer已经通过多次重用同一层来节省内存,引入剪枝可能会干扰这一机制。在某种意义上,剪枝和递归可能相互冲突。我非常好奇看到这方面的实证结果。

  • 蒸馏:知识蒸馏可以从更大的模型转移能力

多模态扩展

递归架构本质上与模态无关。未来的工作可以将MoR 扩展到:

  • 视觉Transformer:对不同图像区域的自适应处理
  • 音频处理:对不同音频片段的可变计算
  • 多模态模型:处理文本、图像和音频的统一架构

结论

递归混合模型(MoR) 提供的不仅仅是另一种效率机制——它代表了如何构建强大且经济可持续的AI系统的战略性重新思考。通过在单一框架中结合参数共享、自适应计算和内存优化,MoR 为在不牺牲性能的情况下降低计算成本开辟了道路。

其意义可能超出学术研究。MoR 的效率提升可能有助于降低部署大型语言模型的财务和基础设施要求。这使得小型组织、研究实验室和初创公司更有可能使用原本无法获得的先进AI能力。

随着AI发展的继续推进,MoR 表明创新并不总是需要扩大规模——有时,它来自优化资源的使用方式。从这个意义上说,MoR 支持向经济上可扩展的AI转变,在这种转变中,性能和资源效率是一致的。

参考文献

Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation

本文转载自​​AIGC深一度


收藏
回复
举报
回复
相关推荐