Research map / Architectures and efficiency / RNNs and normalization
$\boldsymbol{\mu}\mathbf{P^2}$: Effective Sharpness Aware Minimization Requires Layerwise Perturbation Scaling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Derives an infinite-width (Tensor Programs) theory of SAM showing standard SAM only effectively perturbs the last layer, and proposes µP2, a layerwise perturbation scaling that makes all layers effectively perturbed and enables joint transfer of learning rate and perturbation radius.
Problem
Standard SAM's scaling behaviour with network width is not understood; under global perturbation scaling it is unclear which layers are effectively perturbed and whether optimal learning rate and perturbation radius transfer across widths.
Approach
Introduces bcd-parameterizations that add layerwise perturbation-radius scalings to abc-parameterizations, characterizes perturbation regimes (unstable, vanishing, nontrivial, effective), and derives the unique layerwise perturbation scaling µP2 that effectively perturbs all layers under µP.
Key results
- Standard SAM under global perturbation scaling is equivalent to last-layer-only perturbation in the infinite-width limit, confirmed empirically in a 3-layer MLP on CIFAR10.
- µP2 achieves joint transfer of optimal learning rate and perturbation radius across MLP widths (Figure 1) and ρ-transfer in ViT-S/16 on ImageNet1K.
- On ResNet-18 (width multiplier 4, CIFAR10), SAM-ON in µP2 reaches 97.29±0.06 test accuracy (+1.26 over SGD in SP), versus 97.00 for SAM with global scaling in SP.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace global perturbation scaling with layerwise µP2 scaling | SAM in µP · CIFAR10, ResNet-18 width multiplier 4 | test accuracy: 97.19±0.05 -> 97.23±0.08 | Helped |
| replace global perturbation scaling with layerwise µP2 scaling | SAM in SP · CIFAR10, ResNet-18 width multiplier 4 | test accuracy: 97.00±0.03 -> 97.00±0.03 | No clear effect |
| use SAM-ON instead of standard SAM with µP2 scaling | SAM µP2 · CIFAR10, ResNet-18 width multiplier 4, SP | test accuracy: 97.00±0.03 -> 97.29±0.06 | Helped |
| use elementwise ASAM instead of standard SAM with µP2 scaling | SAM µP2 · CIFAR10, ResNet-18 width multiplier 4, SP | test accuracy: 97.00±0.03 -> 97.15±0.01 | Helped |
| use SAM-ON instead of standard SAM with µP2 scaling | SAM µP2 · CIFAR10, ResNet-18 width multiplier 4, µP | test accuracy: 97.23±0.08 -> 97.34±0.08 | Helped |
| use elementwise ASAM instead of standard SAM with µP2 scaling | SAM µP2 · CIFAR10, ResNet-18 width multiplier 4, µP | test accuracy: 97.23±0.08 -> 97.32±0.05 | Helped |
| apply global perturbation scaling ρ · n−1/2 instead of µP2 layerwise scaling (ablation) | SAM in µP · CIFAR10, 3-layer MLP across widths | test accuracy | Hurt |
| use standard SAM perturbations across layers instead of perturbing only the last layer (ablation) | SAM in µP with global scaling ρ · n−1/2 · CIFAR10, 3-layer MLP across widths | activation updates | No clear effect |
| use µP2 layerwise perturbation scaling instead of global perturbation scaling (ablation) | SAM in µP · CIFAR10, 3-layer MLP across widths | test accuracy | Helped |
| use µP2 instead of µP-global perturbation scaling (ablation) | SAM · ImageNet1K, ViT-S/16 across widths | Top 1 accuracy | Helped |
| use µP2 instead of global perturbation scaling (ablation) | SAM variants · CIFAR10, ResNet-18 across widths | test accuracy | Helped |
| use layerwise gradient normalization instead of the paper's perturbation normalization (ablation) | SAM µP2 · Supplemental experiments | generalization performance | Hurt |
Surprising findings
- Standard SAM with optimal hyperparameters effectively perturbs only the last layer in wide networks, so inner backward passes beyond the last layer could be dropped with nearly SGD cost.
- SP networks sometimes show hyperparameter transfer in η and ρ after multi-epoch training, contradicting the infinite-width prediction of output blowup under large learning rates.
Limitations
- Theory assumes constant batch size and training time as width goes to infinity and does not cover generalization.
- Experiments use fine-tuned hyperparameters and moderate scales; gains over SP and µP-global may be larger at larger model sizes but were not tested there.
- Achieving µP2 via weight multipliers is incompatible with unit scaling for low precision training.
Experimental setup
Datasets: CIFAR10 ImageNet1K
Benchmarks: CIFAR10 ImageNet1K
Metrics: test accuracy top-1 accuracy
Baselines: SGD in µP SAM with global perturbation scaling (µP-global) SP SAM-ON Elementwise ASAM Layerwise ASAM AdamW in µP
Builds on: Tensor Programs maximal update parameterization (µP) Sharpness Aware Minimization Adaptive SAM SAM-ON
Closest papers by meaning
- Sharpness-Aware Minimization Efficiently Selects Flatter Minima Late In Training (ICLR 2025 · similarity 0.91)
- Sharpness-Aware Minimization: General Analysis and Improved Rates (ICLR 2025 · similarity 0.91)
- A Universal Class of Sharpness-Aware Minimization Algorithms (ICML 2024 · similarity 0.90)
- Unveiling m-Sharpness Through the Structure of Stochastic Gradient Noise (NeurIPS 2025 · similarity 0.90)
- Friendly Sharpness-Aware Minimization (CVPR 2024 · similarity 0.90)
- Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos (ICLR 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.