超参数迁移使矩阵预条件化优化器在不同规模下均能实现一致提升

最近几种基于矩阵预条件化的深度学习优化器在相对小规模实验中相较于当前主导优化器 AdamW 显示出有前景的加速效果。但尝试验证和复制这些成功结果的努力报告了混合的结果。为更好地理解这些优化器在大规模下的有效性，本文通过超参数迁移来扩展矩阵预条件化优化器的规模化方法，基于 $\mu$P 等先前工作。我们研究在广泛的优化器中，学习率和权重衰减应如何随模型宽度和深度进行比例缩放，包括 Shampoo、SOAP 和 Muon 等优化器，并考虑常见技术如分块和移植的影响。我们发现，根据 $\mu$P 对学习率进行比例缩放可改善迁移效果，但仍可能因显著的有限宽度偏差导致最优学习率漂移，这些偏差可通过分块和显式谱归一化来缓解。对于计算最优缩放，我们发现将权重衰减独立缩放为 $1/\mathrm{width}$ 在各种优化器中几乎最优。应用这些缩放规则后，我们表明 Muon、SOAP 和 Shampoo 在训练规模从 $190$M 到 $1.4$B 的 Llama 架构语言模型时，始终实现约 $1.4\times$ 的加速，而在错误缩放下，加速随规模迅速消失。基于这些结果和进一步的消融实验，我们认为，在给定现实调参预算的情况下，研究最优超参数迁移是可靠比较优化器规模化效果的关键。