Sharpness-aware Orthogonalized Momentum Algorithms

What it claims

SOMA adds a sharpness-aware perturbation to the Muon optimizer while using T+1 gradient evaluations over T iterations, against 2T for SAM-style methods. It builds the perturbation from quantities Muon already holds: the previous iteration's perturbed gradient and the scaled Muon momentum. The authors prove convergence to an ϵ-stationary point in the nuclear norm under heavy-tailed noise with a bounded p-th central moment. The bound has an irreducible floor of order ρν from the nonzero perturbation radius, holds for any perturbation direction, and recovers Muon's rate at ρ = 0. On synthetic problems with a controlled tail index, the spectral variants cut Muon's optimality gap by up to 40% even under the heaviest tails. SAM and F-SAM on SGD degrade sharply in the same setting.

How it works

The paper defines a general framework (Algorithm 1). At each step it builds a perturbation E_t = ρP_t and evaluates the minibatch gradient at W_{t-1} + E_t. It then updates momentum M_t = βM_{t-1} + G_t, orthogonalizes with q Newton-Schulz steps, and takes the step. Three perturbation rules are given:

• FP-SOMA adapts F-SAM to spectral geometry. It computes a fresh gradient G_t, keeps a separate EMA buffer V_t, and sets E_t = ρ·NS_q(G_t − γV_t), with γ in (0,1). This costs two gradient evaluations per step plus an extra buffer.
• SOMA replaces the fresh gradient with the stale perturbed gradient G_{t-1}. It replaces the EMA buffer with (1−β)M_{t-1}, the rescaled Muon momentum. At t = 1 it uses FP-SOMA's perturbation. Lemma 1 argues the stale reuse costs little accuracy.
• OP-SOMA removes the component of the perturbation aligned with the previous Muon step.

The residual G − γV follows F-SAM's argument that the full-gradient component of the perturbation hurts generalization. Orthogonalizing it makes the perturbation spectral rather than Euclidean.

flowchart LR
  G["Stale perturbed grad G(t-1)"] --> R["Residual G - γ(1-β)M"]
  M["Muon momentum M(t-1)"] --> R
  R --> NS["Newton-Schulz → E(t)"]
  NS --> PG["Gradient at W + E(t)"]
  PG --> MU["Muon momentum + step"]
  MU -->|reused next step| G
  MU --> M

Evidence

On synthetic heavy-tailed problems, the spectral variants reduce Muon's optimality gap by up to 40%. Spectral perturbations outperform Frobenius ones. The paper also reports that i.i.d. Gaussian perturbations do not work in practice, which agrees with prior work. The LLM experiments cover 4-bit post-training quantization and catastrophic forgetting on nanoGPT and Pythia-70M. The fetched text ends in Section 4, so those numbers are not available here.

How to use it

Code is at https://github.com/dcll-iiitd/soma. SOMA needs no memory beyond Muon's own state plus the stored previous gradient. FP-SOMA needs an extra EMA buffer and twice the gradient compute. The hyperparameters are ρ (radius), γ (EMA subtraction strength), β (momentum) and q (Newton-Schulz steps). The authors argue for Muon fine-tuning of Muon-pretrained open-weight models because switching optimizers between pre-training and fine-tuning can erode acquired knowledge.

Limits

The convergence bound does not vanish: it keeps a ρν floor whenever ρ > 0. The LLM evaluations use small models (nanoGPT, Pythia-70M), and no larger-scale results appear in the fetched text. The closest prior work is SSM by Zhong et al., a concurrent, unpublished paper that pairs a spectral SAM perturbation with Muon at 2× gradient cost. The fetched text does not show a direct comparison with SSM or with MSAM, which also avoids the extra gradient via momentum.