正确迁移超参数可让矩阵预条件优化器在大模型上稳定提速1.4倍。
Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scales
- 通过分析学习率与权重衰减的缩放规律,实现跨规模优化器迁移。
- 在190M到1.4B参数的Llama模型上,新规则使优化器比AdamW快1.4倍。
- 适合关注大模型训练效率、优化器可扩展性的研究者和工程师。
近期引入的矩阵级预条件优化器在小规模实验中已展现出相比主流优化器AdamW的显著加速效果,但其在大规模场景下的表现仍存争议。本文研究如何通过超参数迁移有效放大这些优化器的性能优势,基于μP等前期工作,系统分析了多种优化器(包括Shampoo、SOAP、Muon)在模型宽度与深度变化时最优学习率和权重衰减的缩放规律,同时考虑了分块(blocking)和嫁接(grafting)等常用技术的影响。研究发现,按μP规则缩放学习率虽能改善迁移效果,但仍存在有限宽度偏差导致最优学习率漂移;通过分块与显式谱归一化可有效缓解此问题。对于计算最优缩放,发现权重衰减按1/宽度独立缩放几乎对所有优化器均接近最优。应用该缩放规则后,Muon、SOAP和Shampoo在190M至1.4B参数的Llama架构语言模型训练中均实现近1.4倍的加速,而错误缩放则使加速效果迅速消失。基于这些结果及进一步消融分析,作者认为在真实调参预算下,研究最优超参数迁移是可靠比较优化器大规模性能的关键。
原文摘要 · Abstract (English)
Several recently introduced deep learning optimizers utilizing matrix-level preconditioning have shown promising speedups relative to the current dominant optimizer AdamW, particularly in relatively small-scale experiments. However, efforts to validate and replicate their successes have reported mixed results. To better understand the effectiveness of these optimizers at scale, in this work we investigate how to scale preconditioned optimizers via hyperparameter transfer, building on prior works such as $μ$P. We study how the optimal learning rate and weight decay should scale with model width and depth for a wide range of optimizers, including Shampoo, SOAP, and Muon, accounting for the impact of commonly used techniques such as blocking and grafting. We find that scaling the learning rate according to $μ$P improves transfer, but can still suffer from significant finite-width deviations that cause drifting optimal learning rates, which we show can be mitigated by blocking and explicit spectral normalization. For compute-optimal scaling, we find scaling independent weight decay as $1/\mathrm{width}$ is nearly optimal across optimizers. Applying these scaling rules, we show Muon, SOAP and Shampoo consistently achieve near $1.4\times$ speedup over AdamW for training Llama-architecture language models of sizes ranging from $190$M to $1.4$B, whereas the speedup vanishes rapidly with scale under incorrect scaling. Based on these results and further ablations, we argue that studying optimal hyperparameter transfer is essential for reliably comparing optimizers at scale given a realistic tuning budget.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。