# μP2: Effective Sharpness Aware Minimization Requires Layerwise Perturbation Scaling

arXiv: https://arxiv.org/abs/2411.00075 · HTML: https://mlsanity.com/paper/2411.00075/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 31 Oct 2024.

## Built on by (most cited first, 10 of 10)

- [Practical Efficiency of Muon for Pretraining](https://mlsanity.com/paper/2505.02222/) (2025)
- [Don't be lazy: CompleteP enables compute-efficient deep transformers](https://mlsanity.com/paper/2505.01618/) (2025)
- [Robust Layerwise Scaling Rules by Proper Weight Decay Tuning](https://mlsanity.com/paper/2510.15262/) (2025)
- [Scaling Diffusion Transformers Efficiently via μP](https://mlsanity.com/paper/2505.15270/) (2025)
- [Function-Space Learning Rates](https://mlsanity.com/paper/2502.17405/) (2025)
- [Do Flat Minima Improve Sparse Novel View Synthesis?](https://mlsanity.com/paper/2511.17918/) (2025)
- [Towards a Principled Muon under μmathsfP: Ensuring Spectral Conditions throughout Training](https://mlsanity.com/paper/2601.01306/) (2026)
- [Unlocking Feature Learning in Gated Delta Networks at Scale](https://mlsanity.com/paper/2606.04048/) (2026)
- [Gradient-Energy Guided Block-Wise Perturbations for Sharpness-Aware Minimization](https://mlsanity.com/paper/2607.18306/) (2026)
- [Learning Rate Transfer for Hybrid Transformer-SSM Architectures](https://mlsanity.com/paper/2610.01172/) (2026)

## Abstract (arXiv)

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 textitMaximal Update and Perturbation Parameterization (μP2), 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 μP2 achieves hyperparameter transfer of the joint optimum of learning rate and perturbation radius across model scales. Moreover, we provide an intuitive condition to derive μP2 for other perturbation rules like Adaptive SAM and SAM-ON, also ensuring balanced perturbation effects across all layers.
