引言

Transformer 架构已成为自然语言处理领域无可争议的王者。从最初的《Attention Is All You Need》论文到如今像 GPT-4 这样的大型语言模型 (LLM) ，其基本配方在很大程度上保持不变: 深层堆叠的相同层。数据从底部进入，逐层按顺序处理，直到从顶部输出。

这种设计依赖于严格的“深度有序惯例 (depth-ordered convention) ”。第 5 层必须等待第 4 层，第 4 层必须等待第 3 层。但这这种僵化的层级结构真的是必须的吗？

想想人类是如何处理信息的。当你阅读像“the”这样简单的单词时，你的大脑动用的认知深度与阅读“量子纠缠 (quantum entanglement) ”这种复杂概念时是不一样的。然而，在标准 Transformer 中，每一个 Token——无论其复杂程度如何——都被迫通过完全相同的计算路径，消耗相同数量的参数和算力。

在研究论文 “Mixture-of-Modules: Reinventing Transformers as Dynamic Assemblies of Modules” 中，研究人员对这一惯例提出了挑战。他们提出了一种看待神经网络构建的全新视角。不再把它想象成一座固定的摩天大楼，而是想象成一桶积木——独立的模块——可以根据输入 Token 的具体需求进行实时动态组装。

这种被称为 混合模块 (Mixture-of-Modules, MoM) 的新架构允许模型即时构建自己的计算图。它将专家混合 (MoE) 和层跳过 (Layer Skipping) 等现有技术统一到一个框架中，实现了比标准 Transformer 更优越的性能和效率。

“堆叠”的问题

要理解 MoM 的重要性，我们首先需要看看“原生 (Vanilla) ” Transformer 的局限性。

原生 Transformer 与混合模块对比。

如上图 Figure 1 (a) 所示，原生 Transformer 由 \(N\) 层组成。每一层包含两个主要的子层: 多头注意力 (MHA) 块和前馈网络 (FFN) 。处理过程是严格线性的。这产生了两个主要问题:

  1. 过度参数化 (Over-parameterization) : 众所周知，大语言模型的参数严重过量。许多神经元起着冗余的作用。在僵化的堆叠结构中，如果不永久性地剪枝整个模型，我们无法轻易绕过这些冗余部分。
  2. 低效 (Inefficiency) : 每个 Token 都要付出最大的计算成本。没有机制让“简单”的 Token 提前退出或跳过繁重的处理步骤。

先前解决这个问题的尝试，如 专家混合 (Mixture-of-Experts, MoE) , 在每层引入了多个 FFN (专家) ，并选择其中几个来激活。虽然有效，但 MoE 仍然在很大程度上遵循深度有序的范式；Token 仍然按照固定的顺序从浅层流向深层。

混合模块 (Mixture-of-Modules) 架构

MoM 的作者提议完全颠覆这一范式。他们的直觉简单而有力: 任何一层，无论其原始位置如何，只要具备所需的处理能力，都可以用来计算 Token。

在这个框架下，模型不再是一个固定的结构，而是一个 动态组装 (dynamic assembly) 的过程。

1. 模块集合 (\(\mathcal{M}\))

MoM 的基础是一个有限的可用模块池。不再有“第 1 层注意力”和“第 2 层 FFN”之分，我们只有一个注意力模块集合和一个前馈模块集合。

研究人员将模块集合 \(\mathcal{M}\) 定义如下:

定义包含注意力模块、FFN 模块和跳过模块的模块集合 M 的公式。

其中:

  • \(\{m_i^A\}\) 是多头注意力 (MHA) 模块。
  • \(\{m_i^F\}\) 是前馈网络 (FFN) 模块。
  • \(m^S\) 是一个特殊的 “SKIP” (跳过) 模块。这是一个恒等函数，允许模型选择在特定步骤 不 执行操作，实际上就是让数据原样通过。

2. 路由器 (\(\mathcal{R}\))

如果我们有一池无序的模块，我们怎么知道该用哪一个呢？这就是 路由器 (Router) 发挥作用的地方。MoM 使用两个不同的路由器: \(\mathcal{R}^A\) 用于选择注意力模块，\(\mathcal{R}^F\) 用于选择 FFN 模块。

对于每个 Token，路由器输出所有可用模块上的概率分布:

将输入 x 映射到分布 r 的路由器函数公式。

不同于做出一次性决策的标准 MoE 路由器，MoM 路由器需要更聪明。构建计算图是一个序列决策过程。第 2 步的模块选择完全取决于第 1 步发生了什么。

为了处理这种依赖关系，作者在路由器内部采用了 门控循环单元 (GRU) 。 这赋予了路由器“记忆”。它维护一个隐藏状态 \(s_h\)，该状态在组装过程的每一步都会更新。

显示路由器隐藏状态 GRU 更新的公式。

然后根据这个隐藏状态生成路由权重:

根据隐藏状态计算路由权重 W 的公式。

这种循环机制确保路由器知晓其迄今为止构建的路径，从而能够就接下来插入哪个模块做出连贯的决策。

3. 动态组装

MoM 模型的前向传播是一个迭代构建的过程。该过程持续 \(H\) 步。在每一步中，模型执行两个动作: 组装注意力模块，然后组装 FFN 模块。

如 Figure 1 (b) 所示，对于特定步骤 \(h\) 的特定 Token:

  1. 路由器 \(\mathcal{R}^A\) 从池中选择前 \(K\) 个最相关的注意力模块。
  2. 这 \(K\) 个模块被组装 (并行计算并聚合) 以更新 Token。
  3. 路由器 \(\mathcal{R}^F\) 选择前 \(K\) 个 FFN 模块。
  4. 这些 FFN 模块被组装以进一步更新 Token。

组装函数 \(\phi\) 形式化地描述了这种选择和聚合:

组装函数 phi 的公式。

单个组装步骤的总体流程涉及残差连接，类似于标准 Transformer:

显示通过组装模块更新隐藏状态 u 和 x 的公式。

组装注意力模块

当路由器选择了 \(K\) 个注意力模块时，它们是如何组合的？输出是各个模块输出的加权和，权重由路由器的概率分数决定。

详述注意力模块输出加权求和的公式。

如果 SKIP 模块 (\(m^S\)) 被选为前 \(K\) 个选择之一，它实际上会将自己从计算中移除，这意味着只使用了 \(K-1\) 个实际的注意力模块。如果 \(K=1\) 且选中了 SKIP 模块，则该 Token 在该步骤完全跳过注意力处理。

组装 FFN 模块

前馈网络的聚合方式类似。所选 FFN 的输出由路由器分数加权并求和:

详述 FFN 模块输出加权求和的公式。

这种机制允许模型混合搭配各种能力。在一个步骤中，它可能会结合来自原始浅层的“语法专注”FFN 和来自深层的“上下文专注”FFN。

MoM 作为一个统一框架

MoM 论文最优雅的地方之一在于它如何统一了各种 Transformer 变体。作者证明，许多“高级”架构实际上只是 MoM 在组装函数 (\(\phi\)) 和路由器上施加特定约束的特例。

让我们看看 Figure 2 中的比较。

原生、参数共享、层跳过、MoE 和 MoM 模型的前向计算路径可视化。

  • 原生 Transformer (Vanilla Transformer) : 一个僵化的堆叠，其中路由器被硬编码为按顺序选择模块 1，然后是模块 2，然后是模块 3。
  • 层跳过 (Layer-Skip) : 这是一个 MoM，其模块集中包含一个 SKIP 模块，路由器决定是选择标准层还是 SKIP 模块。

作为 MoM 特例的层跳过公式。

  • 参数共享 (Parameter Sharing) : 在像 ALBERT 这样的模型中，层之间共享权重。用 MoM 的术语来说，这仅仅是路由器在不同步骤重复选择同一个物理模块索引。

作为 MoM 特例的参数共享公式。

  • 专家混合 (MoE) : 在 MoE 中，注意力机制是标准的，但 FFN 层选择专家子集。在 MoM 中，这对应于注意力的固定路由，但 FFN 的动态路由，其中模块集 \(\mathcal{M}\) 包含所有专家。

作为 MoM 特例的 MoE 公式。

通过泛化这些架构，MoM 提供了一个灵活的沙盒来探索最佳的模型结构。

训练策略: 两阶段方法

你可能会问: 为什么不从头开始训练一个 MoM 模型？

作者发现，用随机权重初始化 MoM 模型会导致“退化问题 (degeneration issue) ”。因为路由器是随机开始的，它会将 Token 均匀地分配给所有模块。结果，所有模块倾向于学习同质化 (相似) 的功能。如果所有模块做的事情都一样，路由器就没有动力去学习专门的路径，模型也就无法超越标准 Transformer。

为了解决这个问题，研究人员引入了 两阶段训练方法 (Two-Phase Training Approach) :

  1. 第一阶段 (预训练) : 在数据集上训练一个标准的原生 Transformer。这迫使各层学习专门的特征 (例如，低层学习语法，高层学习语义) 。
  2. 第二阶段 (MoM 训练) : “分解”预训练的 Transformer。原生模型的层成为模块集 \(\mathcal{M}\) 的初始群体。路由器从头开始初始化。然后进一步训练模型，让路由器学习如何动态组装这些预训练的、专业化的模块。

这种“热启动 (warm-start) ”确保了模块从第二阶段一开始就具备独特的能力，从而给路由器提供了有意义的选择。

实验与关键结果

研究人员使用 OpenWebText 语料库进行训练，在语言建模任务上评估了 MoM，并在 GLUE (语言理解) 和 XSUM (文本生成) 上测试了基准性能。

他们测试了三种规模: Small (122M)、Medium (346M) 和 Large (774M)。具体配置如下:

显示 MoM 模型配置的表格。

性能 vs. 效率

结果总结在 Table 1 中，非常引人注目。

比较 MoM 性能与 GPT-2、MoD 和 MoE 基线的表格。

作者重点介绍了 MoM 的三种特定配置:

  • MoM-P (高性能版, Performant): 最大化深度 (\(H=6\)) 和选择 (\(K=2\))。该模型在 GLUE 和 XSUM 上始终优于原生 GPT-2 和 MoE 基线，尽管计算成本较高。
  • MoM-E (高效版, Efficient): 最小化深度 (\(H=1\)) 同时增加模块选择 (\(K=3\))。该模型实现了与 GPT-2 相当的性能，但 减少了 16-17% 的浮点运算 (FLOPs) ，并显著降低了内存使用 。
  • MoM-I (插值版, Interpolated): 两者之间的平衡 (\(K=3, H=2\))。它略微优于 GPT-2，同时仍能节省内存。

这展示了该架构的灵活性。你可以调整 \(K\) 和 \(H\) 超参数，以优先考虑原始能力或推理速度。

关于过度参数化的洞察

最有趣的发现之一来自于对模块冗余的分析。研究人员采用了一个预训练模型，并逐步从集合 \(\mathcal{M}\) 中移除模块，观察性能如何下降。

显示相对于注意力 (N_A) 和 FFN (N_F) 模块数量的验证损失热力图。

Figure 3 展示了随着注意力模块数量 (\(N_A\)) 和 FFN 模块数量 (\(N_F\)) 减少，验证损失的热力图。

  • 横轴 (\(N_A\)): 注意当你从右向左移动 (将注意力模块从 4 减少到 2) 时，颜色主要保持深蓝色 (低损失) 。这意味着 注意力模块是高度冗余的。 你可以移除一半而不会显著损害模型。
  • 纵轴 (\(N_F\)): 当你从下往上移动 (减少 FFN) 时，损失增加得更加明显。FFN 似乎存储了更多关键的、非冗余的知识。

这证实了 NLP 社区长期以来的怀疑: Transformer 拥有过多的注意力容量。MoM 允许我们动态地修剪这种冗余。

路由器学到了什么？

路由器是真的学会了跳转，还是仅仅重建了原始的顺序？

展示路由器中模块间转移概率的图表。

Figure 6 可视化了路由模式。箭头显示了从一个模块转移到另一个模块的概率。

  1. 非顺序性: 路由器 没有 严格遵循 \(1 \to 2 \to 3 \to 4\)。我们看到了跳跃和循环，表明模型正在利用重用模块或跳过模块的自由度。
  2. 不平衡: 某些模块的使用率远高于其他模块 (由自循环和转移的权重表示) 。作者指出，他们没有强制负载平衡 (不像典型的 MoE 模型) 。他们认为这种不平衡是自然的——语言处理中的某些子任务本身就更常见或更困难，需要特定的“主力”模块。

超参数搜索

研究人员还探索了选择模块的数量 (\(K\)) 和组装步骤数 (\(H\)) 如何影响学习。

显示不同 K 和 H 设置下验证损失的图表。

Figure 4 表明，增加计算图深度 (\(H\)) 通常能提高性能 (降低损失) 。然而，收益是递减的。像 \(K=3, H=2\) 这样的配置 (右图中的绿线) 在压平深度的同时增加了每步的宽度，达到了低损失和高效率的“最佳平衡点”。

结论

混合模块 (Mixture-of-Modules) 代表了我们对 Transformer 架构理解的重大进步。通过打破僵化的、按深度排序的惯例，作者表明:

  1. 灵活性致胜: 动态组装允许模型根据 Token 的难度调整其计算。
  2. 冗余是真实的: 注意力机制的很大一部分可以在不损失性能的情况下被移除或重用。
  3. 统一理论: MoM 提供了一个数学框架，将 MoE、层跳过和参数共享解释为同一主题的变体。

这对未来的启示令人兴奋。当我们努力构建更大、更强的模型时，简单地堆叠更多层可能不是答案。相反，AI 的未来可能更像是一盒动态的乐高积木，实时自我组装，构建出当前思维所需的精确结构。