μP$^2$：有效的锐度感知最小化需要逐层扰动缩放
机器学习
2025-02-12 v2 机器学习
摘要
锐度感知最小化（SAM）在各种神经架构和数据集上均能提升性能。随着模型规模不断扩大以提升性能，严谨理解SAM的缩放行为变得至关重要。为此，我们利用Tensor Programs框架，研究了使用SAM训练的神经网络在无限宽度极限下的行为。我们的发现揭示，即使在最优超参数下，标准SAM的动力学在宽神经网络中实际上退化为仅在最后一层应用SAM。相比之下，我们确定了一种具有逐层扰动缩放的稳定参数化方法，我们称之为\textit{最大更新与扰动参数化}（P），它确保在极限情况下所有层都既能进行特征学习又能被有效扰动。通过对MLP、ResNet和Vision Transformer的实验，我们实证证明了P能够在不同模型规模之间实现学习率和扰动半径联合最优值的超参数迁移。此外，我们提供了一个直观的条件，用于为其他扰动规则（如Adaptive SAM和SAM-ON）推导P，同样确保了所有层之间均衡的扰动效果。
引用
@article{arxiv.2411.00075,
title = {{\mu}P$^2$: Effective Sharpness Aware Minimization Requires Layerwise Perturbation Scaling},
author = {Moritz Haas and Jin Xu and Volkan Cevher and Leena Chennuru Vankadara},
journal= {arXiv preprint arXiv:2411.00075},
year = {2025}
}
备注
Final NeurIPS 2024 camera-ready version. Differences to v1: Cleaner Figure 1, added Appendix H.3.2 showing that even MLPs can transfer optimal HPs in some versions of SP on CIFAR-10, small improvements in writing