SHARPNESS-AWARE MINIMIZATION AND MUON: ROBUSTNESS UNDER THE SPECTRAL NORM

This paper investigates whether the geometric assumptions underlying Sharpness-Aware Minimization (SAM) can be improved by aligning them with the matrix structure of neural network weights. While standard SAM uses Euclidean geometry to define 'small' perturbations, the authors hypothesize that spectral-norm-based perturbations—which respect the linear operator structure of weight matrices—provide a more effective way to encourage flat, well-generalizing minima when paired with matrix-aware optimizers like Muon.

The authors introduce a generalized SAM framework that allows for flexible inner perturbation geometries and outer update rules. They specifically implement 'SpecSAM,' which uses a layerwise spectral norm to generate adversarial perturbations for matrix-valued parameters. They then test this against standard Euclidean SAM across a grid of outer optimizers, including AdamW, SGDW, and Muon. The experiments focus on ImageNet-1K classification using ViT-Small/16 and ResNet-50 architectures, evaluating validation accuracy, ReaL accuracy, and out-of-distribution performance on ImageNet-R.

The study reveals a significant interaction between the inner perturbation geometry and the outer optimizer. While Euclidean SAM provides consistent improvements regardless of the outer optimizer, spectral-norm-based perturbations are highly sensitive to the choice of the outer step. Specifically, SpecSAM performs poorly when paired with AdamW but achieves state-of-the-art results when paired with Muon. This suggests that the benefits of spectral-norm-based sharpness are realized only when the optimization process is consistently 'matrix-aware' in both the perturbation and the update phases.