Sharpness-Aware Minimization and Muon: Robustness under the Spectral Norm

2026-07-28

AI Summary

This work addresses the lack of principled criteria in existing Sharpness-Aware Minimization (SAM) methods for selecting perturbation geometry, which limits generalization performance. The authors propose a novel inter-layer spectral-norm-aware perturbation strategy that explicitly models the geometric structure of weight matrices within the SAM framework: the inner-loop optimization applies directional perturbations scaled by the spectral norm of each layer’s weights, while the outer-loop update leverages either Muon or AdamW/SGDW optimizers. This approach is the first to integrate spectral-norm-guided matrix geometry into SAM’s bilevel optimization process. Empirical results demonstrate substantial improvements in generalization on ImageNet-1K, achieving state-of-the-art validation accuracy with both ViT-Small/16 and ResNet-50 architectures.

Abstract

Sharpness-Aware Minimization (SAM) aims to improve generalization by encouraging insensitivity to small, worst-case parameter perturbations. However, the notion of a "small" perturbation is inherently geometry-dependent: while existing SAM variants have explored a wide range of choices, a clear perspective on which geometries are most effective in practice remains elusive. Recent work on matrix-aware optimization, particularly the Muon optimizer, suggests that respecting the matrix structure of hidden-layer weights can lead to strong empirical performance. Motivated by this, we study matrix-aware geometry in both stages of SAM: we introduce a layerwise spectral inner perturbation for matrix-valued hidden-layer parameters and combine it with either AdamW/SGDW or Muon in the outer update. Across ImageNet-1K experiments on ViT-Small/16 and ResNet-50, we find that the combination of a spectral inner step with a Muon outer step performs consistently strongly, achieving the best validation accuracy on both models among the evaluated methods.

Problem

• Sharpness-Aware Minimization
• matrix-aware geometry
• spectral norm
• parameter perturbation
• generalization

Innovation

• Sharpness-Aware Minimization
• spectral norm
• matrix-aware optimization
• Muon optimizer
• layerwise perturbation