SAM Meets Muon: Why Spectral Geometry Unlocks Better Flat Minima

Paper: 2607.26001 Authors: Wenzhi Zhong, Edward Milsom, Michael Murray Categories: cs.LG, stat.ML

The Gap

SAM’s core promise is simple: find parameters that stay good even when slightly perturbed. But “slightly” is doing a lot of work here — it depends entirely on which geometry you use to measure perturbations. Prior work has explored ℓ₂ balls, ℓ∞ balls, per-layer normalizations, and various adaptive geometries (ASAM, LookSAM, etc.). Yet no principled framework tells us which geometry actually respects the structure of neural network weights.

Separately, the Muon optimizer has shown strong empirical results by treating hidden-layer weight matrices as matrices rather than flattened vectors — applying orthogonal updates via Newton-Schulz iterations. This matrix-aware perspective hasn’t been connected to SAM’s perturbation geometry.

This paper bridges that gap: what happens when you apply spectral-norm-aware geometry to SAM’s inner perturbation step and pair it with a matrix-aware outer optimizer?

SAM variants use ad-hoc        Muon optimizer respects
perturbation geometries         matrix structure of weights
(no consensus on best)          (strong empirically)
        |                               |
        +---------> GAP <--------------+
                   |
        What if SAM's perturbation
        geometry also respected
        matrix structure?
                   |
          Spectral inner perturbation
          + Muon outer update
                   |
          Consistent wins on ViT-Small,
          ResNet-50 across ImageNet-1K

The Increment

One sentence: Before this paper, SAM treated weight perturbation geometry as a design choice disconnected from the optimizer’s own structural awareness; after this paper, we know that aligning SAM’s inner-perturbation geometry (spectral norm) with the outer optimizer’s matrix-awareness (Muon) yields the strongest and most consistent generalization gains.

Core Mechanism

SAM has two stages. The inner step computes an adversarial perturbation: find the parameter direction that maximizes loss within some neighborhood. The outer step then updates parameters to minimize loss at that worst-case point. The geometry — how you measure “small” — matters in both stages.

This paper’s contribution focuses on the inner perturbation. Standard SAM uses Euclidean (ℓ₂) perturbation uniformly across all parameters. The authors instead compute perturbations layerwise for matrix-valued hidden-layer weights using the spectral norm (the largest singular value). Concretely: they compute the spectral gradient of the loss with respect to each weight matrix, scale the perturbation by the spectral norm, and use this to define the adversarial direction. For the outer update, they plug in either AdamW, SGDW, or Muon.

Input data
    |
    v
[Forward pass] --> compute loss L(theta)
    |
    v
[Inner step: spectral perturbation]
    |  For each weight matrix W_i:
    |    compute grad_i = dL/dW_i
    |    project onto spectral direction
    |    epsilon_i = rho * grad_i / ||grad_i||_spectral
    |    theta_adv = theta + epsilon
    |
    v
[Outer step: update rule]
    |  Compute loss L(theta_adv)
    |  Compute grad at theta_adv
    |  Apply Muon (or AdamW/SGDW) update
    |
    v
[Updated theta]

Here’s a structural metaphor. Imagine you’re a military strategist fortifying a castle (the model). SAM says: “Before you decide where to reinforce walls, first send scouts to find the weakest point an attacker could exploit, then shore up against that.” The geometry is the terrain your scouts patrol. Standard SAM sends scouts on flat ground (Euclidean ball). This paper sends scouts along the natural joints and hinges of the castle — the directions where the structure is most vulnerable *given how it’s actually built. The spectral norm identifies exactly those joints: the directions of maximum stretch in each weight matrix. And the Muon outer update is like the master builder who understands how each wall segment (matrix) connects, so reinforcements are applied structurally, not haphazardly.

Key Concepts

• Sharpness-Aware Minimization (SAM): Think of your loss landscape as a mountain range. Most optimizers find the lowest valley they can. SAM finds the valley that’s *also flat — because if the valley is narrow and steep, a tiny shift in data will catapult you uphill. Concretely, SAM perturbs parameters by a small amount in the direction that increases loss the most, then updates to reduce loss at that worst-case point. The result: solutions that generalize better because they aren’t memorizing narrow valleys.

• Spectral Norm: Every matrix has a “maximum stretch factor” — the largest amount it can stretch a vector when you multiply by it. This is the spectral norm (the largest singular value). For a weight matrix in a neural layer, the spectral norm tells you the most extreme direction that layer can amplify a signal. Using spectral geometry for SAM’s perturbation means you’re probing along exactly that most-amplifying direction, which is precisely where the model is most sensitive.

• Muon Optimizer: Standard optimizers like Adam treat all parameters as a flat list of numbers. Muon instead recognizes that hidden-layer weights are *matrices with structure. It computes updates that respect this matrix geometry — specifically, it uses Newton-Schulz iterations to approximate an orthogonal update (one that preserves the matrix’s singular vectors). Think of it as updating each weight matrix in a way that rotates it toward better solutions rather than just nudging individual entries.

Framework Shift

Before (mainstream approach):        After (this paper):

SAM with uniform perturbation        SAM with spectral perturbation
across all params                    per weight matrix

  [param vec]                          [matrix W_1]
       |                                    |
  [flat L2 ball]                       [spectral ball]
       |                                (singular direction)
  [Adam/SGD outer]                         |
                                      [Muon outer update]
                                      (matrix-aware)

From treating SAM’s perturbation geometry as a flat, uniform design choice to aligning it with the intrinsic matrix structure of neural network layers, the core shift is that the geometry of the perturbation and the geometry of the update should speak the same language.

Expert Assessment

Problem choice: This is a real gap. The question of which geometry SAM should use has been lingering since the original SAM paper, and the Muon optimizer’s empirical success has created a natural bridge to explore. The paper is well-positioned at the intersection of two active threads. Not a paradigm shift, but a legitimate and timely investigation.

Method maturity: The spectral inner perturbation is a clean, principled idea — not brute force. However, the paper is essentially a systematic combination of two known ingredients rather than a fundamentally new algorithm. The Newton-Schulz iteration cost for Muon and the spectral decomposition for the inner step add computational overhead that isn’t deeply analyzed. Simpler per-layer normalizations (like in ASAM) might close some of the gap at lower cost.

Experimental integrity: The baselines are reasonable — they compare SAM, ASAM, and other variants with different outer optimizers. ImageNet-1K with ViT-Small/16 and ResNet-50 is standard and credible. The gains are consistent but modest in absolute terms (we’re talking about validation accuracy differences on the order of fractions of a percent). The paper would benefit from ablation studies separating the spectral perturbation contribution from the Muon contribution more cleanly, and from testing on tasks beyond image classification to assess generality.

Writing quality: The paper reads clearly but cuts corners on the theoretical motivation — why *should spectral geometry be the right choice beyond empirical results? The related work section could be tighter. The experimental section would benefit from error bars or confidence intervals across multiple seeds; without them, it’s hard to judge whether the differences are statistically meaningful.

Verdict: weak accept — Clean combination of two strong ideas with consistent empirical gains, but the contribution is incremental rather than foundational, and the experiments need more rigor to fully convince.

Takeaways

Three concrete things a practitioner can steal:

1. Spectral perturbation for SAM is easy to implement and worth trying. If you’re already using SAM, replace the uniform ℓ₂ perturbation with a layerwise spectral perturbation on weight matrices. Compute the top singular vector (one power iteration is usually enough), scale by the spectral norm, and use that as your adversarial direction. This is a low-cost upgrade.

2. Match your perturbation geometry to your optimizer geometry. The paper’s key insight is structural: if your optimizer is matrix-aware (Muon), your SAM perturbation should be too (spectral). More generally, don’t treat SAM’s geometry as independent of the rest of your training pipeline.

3. Muon is worth a look as a drop-in outer optimizer for SAM. Even if you don’t adopt the spectral perturbation, swapping AdamW for Muon in the outer step of SAM may yield free gains on vision tasks with weight matrices. The Muon library is increasingly accessible.