Let's Scale Step by Step: Compute-Efficient Hyperparameter Transfer for Large-Scale Mixture-of-Experts

📄 arXiv: 2608.20061v1 📥 PDF

作者: Nayeon Kim, Hojin Lee, Yunju Bak, Jaesun Park, Boseop Kim

分类: cs.LG, cs.AI, cs.CL

发布日期: 2026-08-20

备注: COLM 2026

💡 一句话要点

提出计算高效的超参数转移框架以优化大规模MoE模型

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家 超参数优化 学习率转移 线性回归 大规模训练

📋 核心要点

  1. 现有方法在极大规模的MoE模型训练中，超参数优化尤其是学习率的选择计算开销巨大，难以高效实现。
  2. 本文提出了一种两步超参数转移框架，通过在不同宽度的模型间转移学习率，并利用线性回归外推到大规模训练。
  3. 实验结果表明，使用小型代理模型进行训练即可准确预测大规模MoE模型的最佳学习率，且在训练稳定性上表现良好。

📝 摘要（中文）

混合专家（MoE）架构在不成比例增加计算成本的情况下显著扩展了模型容量。然而，在极大规模的模型和令牌预算下，优化其超参数（特别是学习率）通过遍历方法仍然计算开销巨大。本文提出了一种计算高效的两步超参数转移框架，通过在宽度扩展的模型间转移学习率，并进一步外推到万亿令牌的训练范围。首先，我们为MoE架构制定了最大更新参数化（$μ$P）适应方案，利用多头潜在注意力（MLA）和Muon优化器，证明了最佳学习率在宽度扩展模型间的一致转移性。其次，通过建立预测性缩放法则，我们成功地将小型代理模型在有限预算下获得的最佳值外推到大规模训练中，达到高保真度（$R^2=0.95$）。

🔬 方法详解

问题定义：本文旨在解决在极大规模的混合专家（MoE）模型训练中，超参数（特别是学习率）优化的计算开销问题。现有方法在模型和令牌预算极大时，遍历学习率的过程非常耗时和资源密集。

核心思路：提出一种计算高效的两步超参数转移框架，首先在不同宽度的MoE模型间转移学习率，然后通过建立预测性缩放法则，将小型模型的最佳学习率外推到大规模训练中。

技术框架：该框架包括两个主要阶段：第一阶段是利用最大更新参数化（$μ$P）适应方案和多头潜在注意力（MLA）进行学习率的转移；第二阶段是通过线性回归对小型代理模型的最佳学习率进行外推，以适应大规模训练需求。

关键创新：最重要的创新在于提出了最大更新参数化（$μ$P）适应方案，证明了学习率在不同宽度模型间的一致性转移性，并建立了有效的预测性缩放法则。

关键设计：在设计上，采用了Muon优化器和多头潜在注意力机制，确保了在不同模型宽度下学习率的有效转移，同时通过线性回归实现了对大规模训练的学习率外推，达到高保真度（$R^2=0.95$）。

🖼️ 关键图片

fig_0
img_1
img_2

📊 实验亮点

实验结果显示，使用小型代理模型进行训练能够准确预测大规模MoE模型的最佳学习率，且在训练过程中表现出高达$R^2=0.95$的外推精度，验证了方法的有效性和稳定性。

🎯 应用场景

该研究的潜在应用领域包括大规模自然语言处理、计算机视觉等需要高效模型训练的任务。通过优化超参数，能够显著降低训练成本，提高模型性能，具有重要的实际价值和广泛的应用前景。

📄 摘要（原文）

Mixture-of-Experts (MoE) architectures significantly expand model capacity without a proportional increase in computational cost. However, optimizing their hyperparameters---particularly the learning rate---at extreme scales of both model size and token budget via sweeping remains computationally prohibitive. In this paper, we propose a compute-efficient, two-step hyperparameter transfer framework that estimates optimal learning rates for training large MoE models by transferring them across scaling model widths, and subsequently extrapolating to trillion-token horizons. First, we formulate a Maximal Update Parameterization ($μ$P) adaptation for MoE architectures utilizing Multi-head Latent Attention (MLA) and the Muon optimizer, demonstrating that optimal learning rates transfer consistently across width-scaled models. Second, we extend this transferability along the token dimension by establishing a predictive scaling law. By applying linear regression to the optimal values derived from small proxy models on limited budgets, we successfully extrapolate the ideal learning rate to massive training horizons (e.g., 10 trillion tokens) with high fidelity ($R^2=0.95$). Consequently, this indicates that proxy training on small models is sufficient to determine the optimal learning rate for the extensive training of large-scale MoEs. We apply the proposed methodology to pretrain our foundation model (155B total, 17B active parameters) from scratch, and the stable training and evaluation results validate that optimal configurations for full-scale target models can be accurately predicted with minimal ablation costs.