技术摘要：基于混合专家层（MoE）的超参数迁移

1. 问题陈述

混合专家（MoE）层已成为扩展现代神经网络的关键机制，它将前向传播中的可训练参数总数与活跃参数解耦。然而，训练稀疏 MoE 模型引入了超参数（HP）选择方面的显著复杂性。与稠密模型不同，MoE 架构引入了新的可训练参数（路由器权重）和新的架构维度（专家数量、专家大小），需要仔细调整。

在大规模上直接调整超参数（如学习率、初始化尺度和权重衰减）在计算上是不可行的。虽然针对稠密 Transformer 的超参数迁移技术已经存在——允许将小模型上发现的最优超参数外推到大模型——但这些方法尚未针对稀疏 MoE 层特定的扩展动态进行严格适配。核心挑战在于，如何在同时增加模型宽度、深度、专家数量和专家大小的情况下确定超参数的扩展方式，而无需在每个规模上重新调整，同时确保训练稳定性和性能。

2. 方法论

2.1. 提出的参数化方法

作者提出了一种针对带有 MoE 层的 Transformer 模型的新参数化方法，将此前为稠密 Transformer 开发的CompleteP参数化扩展，以包含 MoE 特定的扩展规则。其目标是定义一组规则，预测当模型维度扩展时，原始超参数值（初始化标准差 $\sigma$ 和学习率 $\eta$）应如何变化，以保持训练动态的一致性。

推导依赖于**最大更新参数化（$\mu$P）**原则，该原则要求网络组件（预激活值和残差更新）在初始化时保持 $O(1)$，并在每个训练步骤中获得 $\Theta(1)$ 的更新。作者将此扩展到 MoE，要求最大更新条件不仅适用于层输出，还适用于各个专家组件（混合系数和专家输出）。

为 MoE 模块推导的关键扩展规则包括：

• 路由器权重： 学习率按 $\eta \propto n_{embd}^{-1}$ 缩放。初始化按 $n_{embd}^{-\gamma}$ 缩放（其中 $\gamma \geq 0.5$）。
• 专家偏置： 初始化为零，学习率为常数 $\eta \propto 1$（独立于专家数量），前提是稀疏度固定。
• 专家 MLP 权重：
  • 上投影（$W_{up}$）：$\sigma_{init} \propto n_{embd}^{-1/2}$，$\eta \propto n_{embd}^{-1}$。
  • 下投影（$W_{down}$）：$\sigma_{init} \propto \alpha_{ffn}^{-1} n_{embd}^{-1/2}$，$\eta \propto \alpha_{ffn}^{-1} n_{embd}^{-1}$。
  • 其中，$\alpha_{ffn}$ 表示相对于嵌入维度的专家隐藏层大小乘数。

至关重要的是，作者在扩展专家总数（$n_{exp}$）和活跃专家数（$n_{act}$）时，将稀疏度比率 $\kappa = n_{act}/n_{exp}$（活跃专家的比例）固定为常数。这与固定活跃专家数量而增加总池大小的方法形成对比。

2.2. 理论依据：动态平均场理论（DMFT）

为了证明这些启发式扩展规则的合理性，作者采用了动态平均场理论（DMFT）。他们在无限宽度（$n_{embd}$）、深度（$L$）、专家大小（$n_{hid}$）和专家数量（$n_{exp}$）的同时极限下，分析了带有 MoE 层的残差网络的训练动态，同时保持恒定的激活稀疏度 $\kappa$。

分析揭示了一个新颖的三级平均场层级：

1. 残差流表示： 基于专家输出的平均场。
2. 专家输出： 基于单个专家神经元的平均场。
3. 单个神经元： 每个专家内部的平均场。

DMFT 分析表明，在提出的参数化下，极限训练动态是：

• 独立于 FFN 比率（$\alpha_{ffn}$）： 只要采取联合扩展极限，动态就不依赖于专家相对于嵌入维度的具体大小。
• 仅依赖于稀疏度 $\kappa$： 只要稀疏度比率固定，动态在所有扩展参数上都是一致的。
• 尺度不变性： 网络统计量（如逐层特征核）的演化在不同尺度上保持一致，从理论上保证了超参数迁移的可靠性。

2.3. 实验设置

作者在 FineWeb 和 C4 数据集上使用仅解码器的 Transformer 语言模型，对提出的参数化进行了实证验证。

• 基础模型： 在约 3800 万（38M）激活参数的小模型上进行调整。
• 扩展： 扩展至总参数高达 20 亿（2B）的模型，变化宽度、深度、专家数量和专家大小。
• Token 预算： 实验在固定的 10 亿（1B）Token 预算（2000 步）上进行，以隔离早期训练动态，同时也进行了更长的周期（高达 75 亿 Token）。
• 优化器： 标准 Adam 优化器。
• 负载均衡： 采用无辅助损失的策略，直接更新专家偏置以鼓励负载均衡，而无需在损失函数中添加正则化项。

3. 主要贡献

1. MoE 参数化： 本文将 CompleteP 参数化扩展到稀疏 MoE 模型，提供了针对路由器权重、专家偏置和专家 MLP 权重在宽度、深度、专家数量和专家大小上的明确扩展规则。
2. 基于 DMFT 的理论论证： 作者利用 DMFT 为其参数化提供了严格的理论基础。他们推导出了无限宽度/深度极限下训练动态的明确描述，证明了动态收敛到一个稳定的、尺度不变的系统，该系统仅依赖于稀疏度，而不依赖于专家维度的具体扩展方式。
3. 超参数迁移的实证验证： 该研究表明，在小基础模型（3800 万激活参数）上识别出的最优超参数（学习率和初始化尺度）可以可靠地迁移到各种架构维度下的大得多的模型（总参数高达 20 亿）。
4. 架构洞察： 作者通过实证验证，在其参数化下，增加专家数量（在保持总参数固定的情况下）比增加单个专家的大小能带来更好的性能。这一发现与最近的文献一致，但此处无需在每个规模上进行昂贵的超参数搜索即可实现。

4. 结果

• 可靠的超参数迁移： 在提出的扩展规则下，最优的基础学习率和初始化标准差在总参数从 5100 万到 20 亿的模型间有效迁移。扩展模型的损失曲线在早期迭代中与基础模型重合，随后发散（更大的模型实现更低的损失）。
• 稳定性： 该参数化确保了稳定的训练动态，包括均匀的专家负载均衡，即使在扩展专家数量时也是如此。作者指出，MoE 预训练对常数尺度超参数（被视为 $\Theta(1)$ 的乘数）特别敏感，调整这些参数对稳定性至关重要。
• 性能： 使用零样本超参数（从小模型迁移而来）训练的模型，在与激活参数数量匹配的情况下，其性能可与稠密基线（如 GPT-2 small/medium）竞争。
• 专家数量与大小： 实验证实，在固定参数数量的情况下，增加专家数量比增加专家大小更具参数效率。这一优势即使在更长的训练周期（50 亿 Token）中依然存在。
• 稀疏度敏感性： 研究强调，超参数迁移仅在稀疏度比率 $\kappa$ 固定时才有效。在固定活跃专家数量的情况下扩展专家数量（导致 $\kappa \to 0$）会破坏最优超参数的可迁移性。

5. 意义与主张

本文声称提供了一个实用且严谨的框架用于扩展 MoE 模型。通过将启发式参数化与 DMFT 分析相结合，作者提供了一种方法：

• 降低训练成本： 通过仅调整小型、廉价的基础模型，即可为大规模 MoE 模型选择最优超参数。
• 确保稳定性： 提供规则以防止训练不稳定（如专家坍塌或发散），这在稀疏 MoE 训练中很常见。
• 指导架构设计： 提供理论和实证证据，表明在固定稀疏度下，增加专家数量优于增加专家大小，有助于设计高效的大规模模型。

作者承认了局限性，指出当前工作侧重于学习率和初始化，其他超参数（批量大小、权重衰减、学习率调度）留待未来研究。他们还指出，虽然 DMFT 分析支持无限宽度极限，但小宽度迁移的行为需要进一步的理论研究。此外，本文并未声称解决 MoE 的“计算最优”扩展定律（如 Chinchilla 指数），因为由于稀疏性引起的硬件限制，MoE 中的 FLOP 与性能权衡与稠密模型显著不同。

总之，这项工作建立了一个基础参数化，允许基于新颖的三层训练动态平均场理论，将超参数从小的 MoE 模型可靠地外推到大的 MoE 模型。