PaperlayerView exact revision on arXiv

Source-linked AI summary

Sharpness-Aware Muon: Matrix-Aware Optimization with Momentum-Guided Perturbations

Zhiwei Zhai, Wenjing Yan, Ying-Jun Angela Zhang

arXiv:2610.37225v1cs.LG

TL;DR

Muon structures matrix-valued update directions but does not explicitly control local worst-case loss, motivating a sharpness-aware formulation. The paper introduces MuSAM and then MoMuSAM, which uses momentum for sharpness probing; both achieve the same asymptotic stationarity order, while MoMuSAM has a more stable stochastic linearization and sharpness awareness improves Muon in benchmark experiments.

  • Problem

    Muon shapes update geometry but does not explicitly control neighborhood worst-case loss or favor flat regions, and sharpness-aware training has not been systematically developed for matrix-aware orthogonalized optimization.

    Source: §1 · p. 2

  • Method

    MuSAM integrates SAM-style perturbations with Muon’s orthogonalized update, while MoMuSAM uses normalized momentum as a temporally aggregated perturbation direction and evaluates the gradient once at the perturbed point.

    Source: §1 · p. 3

  • Results

    Both methods attain the same asymptotic first-order stationarity order, MoMuSAM has a sharper certified directional-stability bound, and experiments on standard benchmarks show that sharpness awareness consistently strengthens Muon.

    Source: §1 · p. 3, §6 · p. 13

  • Takeaways & Limitations

    Muon’s internal momentum can serve not only as an update variable but also as a history-aware geometric probe of local sharpness.

    Source: §1 · pp. 2–3, §6 · p. 13

    MuSAM’s perturbation requires an extra gradient and uses a single-minibatch direction that can be batch-local and noisy, potentially failing to reflect broader local geometry.

    Source: §3.1 · p. 7

Abstract

from arXiv · show

Muon is a matrix-aware optimizer that combines momentum with orthogonalization to produce structured update directions with strong empirical behavior. However, it only shapes the geometry of the update itself and does not explicitly control neighborhood worst-case loss or bias training toward flat regions, which are often linked to stronger generalization. To address this limitation, we introduce the sharpness-aware formulation of Muon. We first develop \textsc{MuSAM}, which incorporates Sharpness-Aware Minimization (SAM)-style perturbation into Muon's orthogonalized update and establishes a principled bridge between sharpness-aware training and matrix-aware optimization. We then show that this formulation faces an efficiency--stability tension: constructing the perturbation requires an extra gradient evaluation, while the perturbation itself is generated from a single minibatch gradient and is therefore noisy. To resolve this issue, we propose \textsc{MoMuSAM}, which turns Muon's momentum from an update variable into a sharpness-probing signal. By using normalized momentum as the perturbation direction, \textsc{MoMuSAM} removes the extra clean-gradient computation and replaces the single-batch perturbation proxy with a temporally aggregated one. We prove that both \textsc{MuSAM} and \textsc{MoMuSAM} attain the same asymptotic first-order stationarity order, and further show that \textsc{MoMuSAM} yields a more stable stochastic linearization than \textsc{MuSAM}. Experiments on standard benchmarks demonstrate that sharpness awareness consistently improves Muon.

1 Introduction

Muon shapes matrix-aware update directions but does not explicitly control local worst-case loss, leaving update geometry and flatness-oriented training uncoordinated. The paper introduces MuSAM and then MoMuSAM to integrate sharpness probing with Muon, with theoretical and benchmark results supporting the resulting methods.

1 Introduction

  • Muon controls update geometry but leaves local worst-case loss—and therefore explicit preference for flat regions—unmodeled.Sharpness-aware methods target neighborhood worst-case loss, while structured optimizers shape the update used to explore the landscape.Source: §1 · p. 2
  • MuSAM couples SAM-style local worst-case perturbation with Muon’s orthogonalized momentum update, making explicit sharpness control part of the matrix-aware framework.Its perturbation requires an additional gradient and is built from a single minibatch gradient, creating an efficiency–stability tension.Source: §1 · p. 2
  • MoMuSAM uses normalized momentum as a temporally aggregated sharpness probe, removing the extra clean-gradient evaluation and replacing the noisy single-batch perturbation proxy.This turns momentum from a pure update variable into a history-aware geometric probe of local sharpness.Source: §1 · p. 3
  • Both methods attain O(T^-1/4) average first-order stationarity, while MoMuSAM improves the certified average directional-linearization stability bound from O(T^-3/4) to O(T^-1).The guarantees cover both clean iterates and corresponding perturbed query points, and MoMuSAM retains one gradient evaluation per iteration.Source: §1 · p. 3
  • Sharpness awareness consistently strengthens Muon on standard benchmarks, and MoMuSAM delivers a markedly better generalization–efficiency tradeoff than MuSAM.The paper presents these experiments as evidence that the sharpness-aware methods improve Muon while differentiating MoMuSAM from the direct formulation.Source: §1 · p. 3

2 Preliminaries

The preliminaries establish the stochastic matrix-optimization setting and the assumptions used in the convergence analysis. They then describe Muon’s momentum orthogonalization and its distinction between update geometry and explicit sharpness control.

2.1 Problem Formulation

  • The setup treats each matrix block with stochastic loss and minibatch gradients, applying the same construction blockwise to matrix-valued hidden-layer parameters.The paper also specifies matrix norms, inner product, polar notation, learning rate, momentum coefficient, and perturbation radius.Source: §2.1 · p. 4
  • The convergence analysis adopts standard stochastic nonconvex assumptions, including a lower-bounded objective, sample-wise smoothness, and conditionally unbiased minibatch gradients.The minibatch estimator’s variance decreases proportionally to batch size under the stated noise assumption.Source: §2.1 · pp. 4–5

2.2 Momentum Orthogonalization

  • Muon forms a matrix-valued momentum term and orthogonalizes it to produce a structured search direction that preserves dominant directional subspaces while removing singular-value magnitudes.The generic update uses momentum matrix M_t and orthogonalized direction O_t.Source: §2.2 · p. 5
  • In practice, Muon approximates the polar factor with a few Newton–Schulz iterations, but this update-direction normalization remains sharpness-agnostic and leaves local worst-case loss uncontrolled.Thus, Muon does not explicitly encourage convergence to flat regions.Source: §2.2 · p. 5

3 Algorithm Design

MuSAM adds a stochastic sharpness probe before Muon’s orthogonalized update, exposing extra gradient cost and a batch-local perturbation. MoMuSAM reuses momentum to probe sharpness, reducing that cost and improving the certified directional-stability bound while retaining the same asymptotic stationarity order.

3.1 A Sharpness-Aware Extension of Muon: MuSAM

  • MuSAM perturbs parameters using a stochastic first-order approximation of local worst-case loss, then accumulates the perturbed gradient and orthogonalizes momentum for the Muon update.The perturbation searches within a small Frobenius ball to measure local loss sensitivity.Source: §3.1 · pp. 6–7

3.2 A Momentum-Guided Refinement: MoMuSAM

  • MuSAM’s single-minibatch perturbation is noise-sensitive and batch-local, while MoMuSAM’s momentum aggregates gradients over time to reduce sensitivity to any one minibatch.The momentum combines perturbed stochastic gradients from multiple iterations with exponentially decaying weights.Source: §3.1 · p. 7, §3.2 · p. 8
  • MoMuSAM uses Muon’s momentum state as a sharpness-probing signal, reusing temporally aggregated gradient information rather than a fresh clean gradient.The momentum state is available before the current backward pass and aggregates information across previous iterations.Source: §3.2 · p. 7
  • Unlike MuSAM, MoMuSAM removes the clean-gradient evaluation, reducing per-iteration gradient cost from two evaluations to one and recovering vanilla Muon’s computational profile.It evaluates the gradient at the perturbed point before proceeding with the standard Muon update.Source: §3.2 · p. 8
  • MoMuSAM improves the certified average achieved-linearization gap from O(T^-3/4) for MuSAM to O(T^-1), controlling stochastic fluctuation and momentum-tracking error.The comparison holds under the same horizon-dependent hyperparameters.Source: §3.2 · pp. 8–9
  • MoMuSAM reuses the optimizer’s momentum for neighborhood probing without an auxiliary estimator or optimizer state, while preserving Muon’s orthogonalized update.The same momentum is orthogonalized for the current update and normalized for the next perturbation.Source: §3.2 · pp. 8–9

4 Theoretical Analysis

MuSAM and MoMuSAM attain the same asymptotic first-order stationarity order, while MoMuSAM improves stochastic-linearization stability and avoids MuSAM’s extra clean-gradient evaluation. Its certified stability bound improves from O(T^-3/4) to O(T^-1) under the stated hyperparameter choices.

Stability of stochastic linearization

  • The analysis evaluates stochastic-linearization error by comparing a direction induced by a slope proxy with the population-gradient direction.The gap measures the loss in population linearization from using the proxy direction.Source: §4 · pp. 10–11
  • MoMuSAM improves the certified average stochastic-linearization stability bound from O(T^-3/4) for MuSAM to O(T^-1).The improvement follows from averaging minibatch noise across iterations, while momentum-tracking error remains controlled at O(T^-1).Source: §4 · p. 11

Convergence

  • Both MuSAM and MoMuSAM preserve the O(T^-1/4) first-order stationarity order at clean and perturbed iterates.The comparison holds under the stated assumptions and horizon-dependent hyperparameters.Source: §4 · pp. 10–11

Where MoMuSAM improves

  • MoMuSAM uses one gradient evaluation per iteration, removing the extra clean-gradient evaluation MuSAM needs to construct its perturbation.Its temporally aggregated slope proxy is also more stable than MuSAM’s raw-gradient perturbation.Source: §4 · pp. 10–11

5 Experiments

Across image-classification and language-modeling benchmarks, sharpness-aware Muon improves final performance, with MoMuSAM delivering the strongest overall results. Sensitivity tests also show that the proposed methods remain comparatively robust across perturbation radii and learning rates.

Experimental setup

  • The evaluation compares sharpness-aware baselines SAM, MSAM, and VASSO with matrix-aware Muon across image classification and language modeling.Image tasks use FMNIST, CIFAR-10, and CIFAR-100; language modeling uses a 124M-parameter GPT-2 model trained on OpenWebText.Source: §5 · pp. 11–12

Performance comparison

  • Across FMNIST, CIFAR-10, CIFAR-100, and OpenWebText, sharpness-aware Muon improves final performance, with MoMuSAM delivering the strongest overall results.Figure 1 compares test accuracy on the image datasets and validation loss on OpenWebText.Source: Figure 1 · p. 12, §5 · p. 12

Sensitivity analysis on ρ

  • MuSAM and MoMuSAM maintain strong accuracy across a broader range of ρ than SAM, MSAM, and VASSO, with MoMuSAM best or near-best.The SAM-style baselines deteriorate rapidly when ρ becomes too large.Source: §5 · pp. 12–13, Figure 2 · p. 12

Sensitivity analysis on η

  • As η decreases, MuSAM and MoMuSAM remain comparatively stable and retain stronger final accuracy than SAM, MSAM, and VASSO.The reported sensitivity results indicate the proposed methods are easier to tune across a broad range of stepsizes.Source: Figure 3 · p. 12, §5 · p. 13

6 Conclusion

The paper concludes that sharpness awareness extends Muon’s matrix-aware updates with explicit local-geometry control, and that MoMuSAM improves the generalization–efficiency tradeoff. The related-work discussion identifies the absence of explicit local-sharpness control in existing Muon work and the cost and noise challenges of SAM-style perturbations.

6 Conclusion

  • MoMuSAM turns momentum into a sharpness-probing signal, removes the extra clean-gradient computation, and yields a more stable perturbation while retaining MuSAM’s stationarity order.The paper reports that sharpness awareness consistently strengthens Muon and that MoMuSAM offers a markedly better generalization–efficiency tradeoff.Source: §6 · p. 13

A Related Works

  • Existing Muon work reshapes update directions but does not establish an explicit mechanism for controlling local sharpness or steering toward flat regions.The related-work discussion distinguishes update geometry from local loss geometry.Source: A Related Works · p. 18
  • SAM-style perturbations require an additional forward–backward pass and use a noisy, batch-local proxy for sharpness.The paper connects these limitations to instability and the challenge of integrating sharpness awareness with Muon’s structured update geometry.Source: A Related Works · p. 18

B Proof

The proof excerpts outline descent and expectation arguments under the stated assumptions, but do not provide the full lemma conclusions. They identify a lower bound on the objective as part of the telescoping argument.

B.1 Useful lemmas

  • Lemma B.1 invokes the descent lemma for the Muon+SAM update and derives its guarantee using the reverse triangle inequality and substitutions across intermediate equations.The supplied proof steps refer to Eqs. (17)–(22).Source: §B.1 · pp. 19–20
  • The proof’s stated assumptions for both lemmas are Assumptions 2.1 to 2.3.The excerpt does not include the full guarantee statements.Source: §B.1 · pp. 19–20
  • Lemma B.2 is obtained by taking expectations, summing over iterations, and applying a telescoping identity.The argument uses the lower bound f⋆ on f.Source: §B.1 · p. 20

B.2 Theoretical Analysis for Algorithm 1

Theorem B.4 analyzes estimator errors and gives an averaged stationarity bound for Algorithm 1, including a case under a prescribed parameter schedule. The section also presents the MoMuSAM update procedure.

B.2 Theoretical Analysis for Algorithm 1

  • The analysis decomposes the same-minibatch estimator and tracks momentum-recursion error, including a pointwise mismatch bound.The recursion separates the momentum term from stochastic-gradient and perturbation terms.Source: §B.2 · pp. 20–22
  • Under Assumptions 2.1–2.3, Theorem B.4 gives an averaged stationarity bound and a specialization for a prescribed parameter schedule.The schedule sets η = ρ = T^-3/4 and 1 − μ = T^-1/2.Source: §B.2 · pp. 22–23
  • MoMuSAM updates momentum from the perturbed stochastic gradient, orthogonalizes it for the parameter update, and uses normalized momentum for the next perturbation.The perturbation is δt+1 = ρ Mt / ||Mt||F.Source: §B.2 · p. 27

B.3 Theoretical Analysis for Algorithm 2

The analysis tracks the error between momentum and the population gradient at the perturbed point, using a recursion and stochastic-gradient assumptions. It derives bounds for this estimator error.

B.3 Theoretical Analysis for Algorithm 2

  • The analysis studies the perturbed-point mismatch E_t = M_t − ∇f(W_{t−1}+δ_t) recursively via the momentum update; conditioned on F_{t−1}, the prior mismatch and gradient difference are deterministic, making the cross term vanish.Source: §B.3 · p. 27
  • The proof treats μ = 0 separately and shows the same recursion remains valid in that case.At μ = 0, the estimator error is the stochastic-gradient error at the perturbed point.Source: §B.3 · pp. 27–28
  • The analysis combines Eq. (26) and Eq. (72); according to Eq. (57), it further sets η = ρ = 1/2.Source: §B.3 · pp. 29–30

B.4 Proof of Theorem 4.6

Theorem 4.6 compares a slope proxy with the population-gradient perturbation and analyzes perturbation stability for MuSAM and MoMuSAM. The MoMuSAM analysis decomposes momentum into tracking and martingale-noise components.

B.4 Proof of Theorem 4.6

  • Theorem 4.6 establishes a direction-sensitive comparison between a slope proxy and the population-gradient perturbation.The proof notes that the slope proxy is nonnegative for the indicated comparison.Source: §B.4 · p. 31

B.4.1 Stability of the MuSAM perturbation

  • The MuSAM perturbation analysis proves the stated stability result for a direction constructed from the raw minibatch gradient.The proof applies Lemma B.5 to that stochastic-gradient direction.Source: §B.4.1 · pp. 31–32

B.4.2 Stability of the MoMuSAM perturbation

  • The proof bounds the tracking component using the exact-polar update and treats the martingale-noise component with conditional mean-zero noise.The two components are identified separately in the stability argument.Source: §B.4.2 · pp. 32–33
  • The MoMuSAM stability proof establishes a finite-time bound by decomposing momentum into tracking and martingale-noise components.It analyzes bias-corrected momentum as a weighted average of past perturbed stochastic gradients; positive rescaling preserves its normalized direction.Source: §B.4.2 · pp. 32, 34

B.4.3 Average stability under the prescribed schedule

Under the prescribed schedule, the proof combines averaged bounds to establish its stated result. The included visuals identify a comparison of algorithms' loss landscapes and a test-accuracy table across datasets.

B.4.3 Average stability under the prescribed schedule

  • Choosing η = ρ = T^-3/4 and 1 − μ = T^-1/2, Eqs. (102) and (103) prove Eq. (16), with the last equality using ρ = T^-3/4.Source: §B.4.3 · p. 35
  • Figure 4 presents 3D loss landscapes for different algorithms.The caption identifies the visual's subject but does not specify axes or a result.Source: Figure 4 · p. 36
  • Table 2 reports test accuracy under different datasets.The supplied caption does not provide cell values or identify a comparison outcome.Source: Table 2 · p. 36

C Additional Experiments

MoMuSAM improves Muon’s accuracy and local flatness while retaining nearly the same computational cost. Geometry diagnostics and fixed-time comparisons favor MoMuSAM over MuSAM, whose extra clean-gradient evaluation makes its advantage less pronounced under a fixed time budget.

Experimental setup

  • For MuSAM and MoMuSAM, experiment hyperparameters come from theoretical prescriptions and depend only on the total iteration count T, without additional task-specific search.Baseline optimizer hyperparameters are instead tuned by grid search, with the best-performing configuration reported for each method.Source: C Additional Experiments · p. 36

Geometry analysis

  • Both MuSAM and MoMuSAM consistently improve upon Muon on three image-classification benchmarks, with MoMuSAM’s largest margin on the harder CIFAR-100 dataset.The reported geometry diagnostics also show that MoMuSAM reduces top Hessian eigenvalues further than MuSAM, while its loss landscape has the broadest, smoothest basin among the three methods.Source: C Additional Experiments · p. 37, Table 3 · p. 37
  • MoMuSAM’s momentum-based perturbation is temporally aggregated and less noisy than MuSAM’s single-minibatch-gradient direction, and the comparison associates it with flatter solutions and stronger final generalization.The authors describe this signal as more informed by Muon’s accumulated optimization history.Source: C Additional Experiments · p. 37

Wall-clock time comparison

  • MoMuSAM consistently matches or exceeds Muon across almost the entire wall-clock range and achieves the best final CIFAR-10 accuracy overall, without additional per-iteration overhead.It reuses Muon’s momentum state to construct perturbations instead of computing an additional gradient.Source: C Additional Experiments · p. 37, Table 4 · p. 37
  • MuSAM’s extra clean-gradient evaluation makes its accuracy improvement less pronounced under a fixed time budget, whereas MoMuSAM improves accuracy at nearly the same computational cost as Muon.The paper identifies this as a better practical accuracy–efficiency tradeoff.Source: C Additional Experiments · pp. 37–38
Loading 2610.37225v1…