如何扩展混合专家模型：从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在训练中的极限动力学，为我们的分析提供形式化基础。在此框架下，我们导出满足所有最大更新 ($\mu$) 要求的唯一SGD和Adam参数化。我们随后表明， resulting $\mu$P 方案不能可靠地诱导随规模的单调改进或稳健的学习率传递。我们将这些病态现象归因于聚合动力学中与尺度相关的可观测量，这激发我们提出一套细化的要求，称为最大尺度稳定性。循着这一原则，我们在所有三种缩放范例中推导最大尺度稳定参数化 (MSSP)，用于SGD和Adam，并刻画相应的极限动力学——与 $\mu$P 极限在本质上截然不同——通过独立的DMFT分析。实验验证了MSSP能够在各种缩放范例中稳健恢复学习率传递和随规模的单调改进。结合现有的深度缩放理论，这些结果为MoE架构提供了完整的缩放 prescription，作为网络宽度、深度、专家宽度和专家数量的函数。