μP$^2$: Effective Sharpness Aware Minimization Requires Layerwise Perturbation Scaling
This addresses a scaling problem in SAM for neural network training, offering a more robust optimization method for large models, though it is incremental as it builds on existing SAM frameworks.
The paper tackles the issue that Sharpness Aware Minimization (SAM) effectively reduces to applying perturbations only in the last layer in wide neural networks, limiting its effectiveness. It introduces μP², a parameterization with layerwise perturbation scaling, which ensures balanced perturbations across all layers and achieves hyperparameter transfer across model scales, as demonstrated in experiments with MLPs, ResNets, and Vision Transformers.
Sharpness Aware Minimization (SAM) enhances performance across various neural architectures and datasets. As models are continually scaled up to improve performance, a rigorous understanding of SAM's scaling behaviour is paramount. To this end, we study the infinite-width limit of neural networks trained with SAM, using the Tensor Programs framework. Our findings reveal that the dynamics of standard SAM effectively reduce to applying SAM solely in the last layer in wide neural networks, even with optimal hyperparameters. In contrast, we identify a stable parameterization with layerwise perturbation scaling, which we call $\textit{Maximal Update and Perturbation Parameterization}$ ($μ$P$^2$), that ensures all layers are both feature learning and effectively perturbed in the limit. Through experiments with MLPs, ResNets and Vision Transformers, we empirically demonstrate that $μ$P$^2$ achieves hyperparameter transfer of the joint optimum of learning rate and perturbation radius across model scales. Moreover, we provide an intuitive condition to derive $μ$P$^2$ for other perturbation rules like Adaptive SAM and SAM-ON, also ensuring balanced perturbation effects across all layers.