SAM Meets Muon: Spectral Norm Perturbations Improve Robustness and Accuracy

Forum topic · 小凯 · 2026-07-30

Summary

This arXiv paper (2607.26001) by Wenzhi Zhong, Edward Milsom, and Michael Murray investigates how geometry affects Sharpness-Aware Minimization (SAM), the technique that improves generalization by encouraging insensitivity to small worst-case parameter perturbations. Since what counts as a 'small' perturbation depends on the chosen geometry, the authors study matrix-aware geometry in both stages of SAM. Inspired by the Muon optimizer, which respects the matrix structure of hidden-layer weights, they introduce a layerwise spectral inner perturbation for matrix-valued hidden-layer parameters. This inner perturbation is combined with an outer update based on AdamW/SGDW or Muon. Experiments on ImageNet-1K with ViT-Small/16 and ResNet-50 show that combining the spectral inner step with Muon's outer step performs consistently strongly, achieving the best validation accuracy among the evaluated methods on both architectures. The work suggests that respecting matrix structure in SAM's perturbation stage is a key factor for improved robustness and generalization.

Paper Overview

Field: Machine Learning Authors: Wenzhi Zhong, Edward Milsom, Michael Murray arXiv: 2607.26001

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, the authors study matrix-aware geometry in both stages of SAM: they 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.

Key Findings

The notion of a "small" perturbation in SAM is geometry-dependent, and the best geometry choice has been unclear.
A layerwise spectral inner perturbation is proposed for matrix-valued hidden-layer parameters.
The spectral inner step is combined with outer updates using AdamW/SGDW or Muon.
On ImageNet-1K with ViT-Small/16 and ResNet-50, the combination of spectral inner steps with Muon outer steps performs consistently strongly, achieving the best validation accuracy among evaluated methods on both models.