如何扩展专家混合：从 muP 到最大规模稳定参数化

最近的前沿 大语言模型 主要依赖于专家混合 (MoE) 架构。尽管取得了实证进展，但对于超参数应如何随着网络宽度 $N$、专家宽度 $N_e$、专家数量 $M$、稀疏度 $K$ 和深度 $L$ 进行扩展以确保稳定性和规模上的最佳性能，仍然没有原则性的理解。我们通过分析三种不同的缩放机制，为解决这一差距迈出了原则性的一步：(I) 共同缩放 $N\asymp N_e$，(II) 共同缩放 $N\asymp M\asymp K$，以及 (III) $N、N_e、M$ 和 $K$ 的完全比例缩放。对于每个机制，我们开发了一种新颖的动态平均场理论（DMFT）来描述 MoE 的限制训练动态，为我们的分析提供了正式的基础。在此框架内，我们推导出满足所有最大更新 ($μ$) 需求的 SGD 和 Adam 的独特参数化。然后我们表明，由此产生的 $μ$P 处方并不能可靠地通过规模或稳健的学习率迁移引起单调改进。我们将这些病理现象追溯到聚合动力学中与尺度相关的可观察量，这激发了一系列完善的需求，我们将其称为最大尺度稳定性。在这一原则的指导下，我们在所有三种尺度范围内推导了 SGD 和 Adam 的最大尺度稳定参数化 (MSSP)，并通过单独的 DMFT 分析表征了相应的极限动态 - 与 $μ$P 极限在质量上不同。实验验证了 MSSP 能够稳健地恢复学习率转移和跨机制规模的单调改进。结合现有的深度缩放理论，这些结果为 MoE 架构提供了完整的缩放处方，作为宽度、深度、专家宽度和专家数量的函数。