PaperlayerView exact revision on arXiv

Source-linked AI summary

Sharpness-aware Orthogonalized Momentum Algorithms

Geetika, Shivani Singh Chauhan, Gauranshi Gupta, Anant Jain, Kaushik Kalyanaraman, Bapi Chatterjee

arXiv:2610.32465v1cs.LGmath.OC

TL;DR

Sharpness-aware optimization with Muon remains underexplored, and SAM’s behavior under heavy-tailed noise is not established despite that noise’s relevance to LLM training. SOMA reuses prior stochastic-gradient information to form low-cost spectral perturbations for Muon. The paper reports gains on controlled-tail synthetic problems and evaluates SOMA in language-model training, quantization, and fine-tuning settings.

  • Problem

    It is unknown whether SAM remains beneficial under heavy-tailed gradient noise, and existing MUON pairings with SAM lack the low cost and convergence guarantees sought for LLM training.

    Source: §1 · p. 2

  • Method

    SOMA combines MUON with spectral-friendly perturbations built from the previous perturbed gradient and scaled MUON momentum, avoiding new gradient evaluations for perturbation construction.

    Source: §1 · p. 2

  • Results

    On controlled-tail synthetic problems, spectral variants reduce MUON’s optimality gap by up to 40%; SOMA matches MUON’s quality in language-model training.

    Source: §7 · p. 12

  • Takeaways & Limitations

    The results indicate that spectral perturbations can improve MUON’s optimality gap under heavy-tailed noise, while SOMA uses fewer gradient evaluations than its fresh-perturbation counterpart.

    Source: §1 · p. 2

    The convergence analysis assumes nonconvex objectives, i.i.d. minibatches, operator-norm smoothness, and a bounded p-th central moment for stochastic gradients.

    Source: §5 · pp. 6–7

Abstract

from arXiv · show

Adversarial sharpness penalty per iteration of a base optimizer, known as sharpness-aware minimization (SAM), is an accepted technique to improve generalization in deep learning. Recently, Muon has emerged as a dominant optimization algorithm for training large language models (LLMs). It achieves lower sample complexity for heavy-tailed gradient updates - a phenomenon common in LLM training. However, sharpness penalties with Muon updates are yet largely unexplored in the literature. Moreover, the behavior of SAM under heavy-tailed noise is unknown both theoretically and empirically. A potential reason for this gap is the additional cost of extra gradient computation per iteration for penalty perturbations in SAM, which becomes prohibitive for LLMs. This paper introduces Sharpness-aware Orthogonalized Momentum Algorithm (SOMA) to combine a computationally efficient and theoretically grounded perturbation scheme with the Muon algorithm. The key idea is to use stale (purely) stochastic gradient noise for adversarial perturbations. Theoretically, we show that SOMA's sample complexity is invariant to the perturbation direction. We provide explorations and insights into SOMA's robustness to heavy-tailed gradient noise. We extensively explore SOMA's efficacy in (a) mitigating quality degradation due to post-training quantization, and (b) catastrophic forgetting during fine-tuning of LLMs.

1 Introduction

The introduction identifies a gap in combining sharpness-aware optimization with Muon under heavy-tailed noise, then presents SOMA as a low-cost approach with theoretical guarantees and empirical evaluations. Its reported experiments include improvements on controlled-tail synthetic problems and studies of quantization and fine-tuning.

1 Introduction

  • The paper motivates fine-tuning and post-training quantization as important LLM concerns beyond generalization, relating both to model training dynamics.Fine-tuning can involve catastrophic forgetting, while PTQ can degrade model performance.Source: §1 · p. 2
  • Prior SAM analyses assume deterministic gradients or bounded variance, leaving its theoretical and empirical behavior under heavy-tailed noise unstudied.The introduction notes that stochastic gradient noise is heavy-tailed in deep networks, particularly attention-based language models.Source: §1 · p. 2
  • Existing SAM variants largely target SGD, while the only noted MUON pairing uses vanilla SAM’s doubled gradient cost and lacks convergence guarantees.The introduction also identifies as open whether a sharpness-aware perturbation can be added to MUON at essentially no extra cost while retaining its heavy-tail guarantees.Source: §1 · p. 2
  • SOMA reuses quantities already held by MUON to build a perturbation, uses T + 1 gradient evaluations instead of 2T, and is proved to converge under heavy-tailed noise.The guarantee holds for any perturbation direction, recovers MUON’s convergence rate when ρ = 0, and has an irreducible error floor of order ρν.Source: §1 · p. 2
  • Under controlled heavy-tailed noise, spectral variants reduce MUON’s optimality gap by up to 40%, while SAM and F-SAM on SGD degrade sharply.Spectral perturbations also outperform Frobenius perturbations; the paper evaluates proposed methods on 4-bit PTQ and observes catastrophic forgetting in nanoGPT and Pythia-70M.Source: §1 · p. 3

2 Related Work

Related work covers methods that reduce SAM’s computational overhead or refine its perturbation and trajectory, alongside developments in Muon and convergence analysis under heavy-tailed noise. The authors position their contribution at the intersection of these lines, where sharpness-aware perturbations remain absent from the cited heavy-tail Muon analyses.

Sharpness-aware optimization

  • SAM research reduces its two-gradient overhead, refines sharpness measurement, or stabilizes optimization trajectories, with approaches including LookSAM, ESAM, ASAM, and EISAM.Some methods use periodic or selective evaluations, while others change the perturbation or optimization trajectory; the survey says its list is inexhaustive.Source: §2 · p. 3

MUON and its variants

  • Muon variants extend its use through LLM scaling, stability clipping, element-wise adaptivity, memory reductions, and layer-wise norm choices.Zhong et al. pair a layer-wise spectral SAM perturbation with a MUON outer update and report vision-classification gains.Source: §2 · p. 3

Convergence under heavy-tailed noise

  • Heavy-tailed-noise analyses cover SGD and MUON, but the cited MUON convergence work does not include a sharpness-aware perturbation.The cited literature includes MUON analyses under p-BCM noise, including nuclear-norm stationarity without clipping or the extra dimension dependence of Euclidean methods.Source: §2 · p. 3

3 Problem Formulation

This section frames sharpness-aware optimization alongside the spectral geometry of MUON. It contrasts existing SAM, friendly-SAM, momentum-SAM, and spectral-SAM perturbations to motivate a low-cost perturbation that respects MUON’s matrix geometry.

Sharpness-aware minimization

  • Standard SAM perturbs parameters using a gradient-based sharpness penalty and then evaluates a second gradient, doubling per-iteration cost.The perturbation is constrained to a neighborhood of radius ρ, and the cited formulation uses an ℓ2-norm perturbation.Source: §3 · pp. 3–4

Friendly SAM

  • F-SAM perturbs along stochastic noise after decomposing the minibatch gradient into full-gradient and noise components.The noise residual is formed using an exponential moving average of gradients.Source: §3 · p. 4

Momentum-SAM

  • MSAM reuses the base optimizer’s momentum buffer and perturbs in the negative momentum direction to avoid SAM’s second gradient evaluation.The stated rationale is that this direction increases the loss and can estimate sharpness.Source: §3 · p. 4

MUON

  • MUON orthogonalizes matrix-parameter momentum; its Ortho operator is steepest descent under an operator-norm constraint, with implementations approximating it using Newton-Schulz iterations.MUON preserves singular directions while setting nonzero singular values to 1.Source: §3 · p. 4

Spectral SAM MUON

  • SSM combines spectral perturbations with a MUON update, but uses the full minibatch gradient and a fresh gradient evaluation, motivating a low-cost spectral residual.The section contrasts this with F-SAM’s stochastic-noise perturbation and MSAM’s momentum reuse without matrix-geometry adaptation.Source: §3 · p. 5

4 Algorithm

The algorithm section develops spectral perturbations for MUON, first adapting F-SAM and then reusing prior-iteration quantities to reduce gradient cost. It also describes projection variants that address alignment with previous MUON steps.

FP-SOMA

  • FP-SOMA adapts F-SAM to spectral orthogonalization and a MUON outer update, using a fresh gradient and a separate EMA buffer at the cost of two gradient evaluations per iteration.The fresh gradient dominates the perturbation, while the subsequent perturbed gradient drives descent.Source: §4 · p. 5

SOMA

  • SOMA replaces FP-SOMA’s fresh gradient with the previous perturbed gradient and derives its EMA from MUON momentum, requiring T + 1 rather than 2T gradient evaluations.The existing momentum buffer supplies the EMA, so no separate EMA buffer is needed; it tracks gradients at past perturbed iterates.Source: §4 · p. 5

OP-SOMA

  • SOMA’s perturbation can align with the previous descent direction, prompting OP-SOMA to remove the aligned component from the stale friendly residual.The section describes projection either before or after the Newton-Schulz transformation; orthogonality to raw momentum does not guarantee orthogonality to the previous MUON step.Source: §4 · pp. 5–6

Perturbation rules

  • The convergence properties of SOMA remain invariant to the choice of perturbation when the perturbation is zero for the stated zero-direction case.For P_t = 0, the method takes E_t = 0.Source: §4 · p. 6

MSOMA

  • MSOMA perturbs along negative prior MUON momentum, requiring no additional gradient evaluation, but does not project out alignment with the previous descent direction.Its perturbation may therefore remain aligned with the previous movement.Source: §4 · p. 6

5 Convergence Analysis

The analysis establishes convergence guarantees for SOMA under heavy-tailed noise and examines how perturbation radius and stale-gradient reuse affect those guarantees. It also shows that reusing the previous perturbed gradient incurs little accuracy cost while saving a gradient evaluation.

5 Convergence Analysis

  • Under the stated smoothness and p-bounded central-moment assumptions, SOMA attains an ϵ-stationarity guarantee with an irreducible error floor O(χ_q L ρ^ν).The floor grows with the perturbation radius but is independent of the noise distribution.Source: §5 · pp. 6–7
  • Setting ρ = 0 removes the perturbation term and recovers MUON’s convergence bound under heavy-tailed noise.The first term then dominates over the stated range of ν and δ.Source: §5 · p. 7
  • With nonzero ρ, the SAM bound has a reducible term that vanishes as T grows and an irreducible floor O(χ_q L ρ^ν), independent of the noise.The Newton–Schulz approximation factor χ_q approaches 1 as q grows.Source: §5 · p. 7

5.1 SOMA Analysis

  • Reusing the previous iteration’s perturbed gradient adds bias L(ρ + η), leaves noise unchanged, and shifts curvature by O(ρ^2(ρ + η)).When η ≲ ρ, the bias matches the usual SAM error, so the stale-gradient reuse saves a gradient evaluation without significant error.Source: §5.1 · p. 8

6 Experiments

Across the experiments, MUON-based methods generally retain strong training quality, while sharpness-aware variants show task-dependent benefits. SOMA adds little pretraining overhead and obtains part of the quantization robustness benefit of its two-gradient counterpart.

6 Experiments

  • The experiments span controlled synthetic heavy-tailed noise, continued pretraining of two language models and a CNN, 4-bit quantization, and fine-tuning for forgetting.The optimizer comparison includes SGD-, ADAMW-, and MUON-based methods, including SOMA variants.Source: §6 · p. 8, §6.2 · p. 10

6.1 Synthetic Experiments under Heavy-tailed Noise

  • Under heavier synthetic noise, MUON-based methods reach about 38–40% accuracy at α = 1.1, versus about 36% for SGD-based and 29% for ADAMW-based methods.SOMA improves on MUON by 1.0–2.3 points and stays within ±0.8 points of FP-SOMA while using half the gradient evaluations.Source: Figure 1 · p. 9, §6.1 · pp. 9–10

6.2.2 Post-training Quantization

  • On Pythia-70M, FP-SOMA lowers MUON’s perplexity increase after quantization from 10.96% to 10.20%, while SOMA lowers it to 10.55%.Across the MUON family, nanoGPT INT4 perplexity is nearly unchanged at 23.85–23.87; Table 1 compares full-precision and INT4 held-out perplexity.Source: Table 1 · p. 11, §6.2.2 · pp. 11–12

6.2.1 Continued Pretraining

  • On language-model pretraining, MUON-based methods end with the lowest validation perplexity, and SOMA matches MUON while costing 3% more time per nanoGPT epoch.Two-gradient SSM and FP-SOMA have slightly higher nanoGPT perplexity than MUON (24.55 versus 24.41); on CIFAR-10, SSM records 87.63% versus MUON’s 87.27%.Source: §6.2.1 · pp. 10–11, Figure 2 · p. 10

6.2.2 Post-training Quantization

  • Sharpness-aware perturbations moderately improve MUON-trained models’ quantization robustness, and SOMA captures a substantial part of that benefit at almost the cost of MUON.On Pythia-70M, the reported gains differ by optimizer family: sharpness-aware training reduces degradation for MUON and ADAMW, but not SGD.Source: §6.2.2 · pp. 11–12

7 Conclusion

The conclusion presents SOMA as a sharpness-aware MUON variant that reuses existing optimizer quantities. It reports convergence under heavy-tailed noise and practical results on language-model training, quantization robustness, and forgetting.

7 Conclusion

  • SOMA reuses quantities already computed by MUON to perturb weights, and its convergence error floor grows with radius but not with noise.The paper reports that SOMA maintains MUON’s language-model quality at about 3% extra runtime and moderately improves quantization robustness, while reducing forgetting only slightly.Source: §7 · p. 12
  • For MUON-based methods, spectral perturbations consistently help more than the choice of perturbation direction.SOMA matches its two-gradient counterpart in accuracy despite using a stale gradient; its best radius shrinks as noise gets lighter.Source: §7 · p. 12
  • The paper identifies testing SOMA at larger scale and fine-tuning with the pretraining optimizer as natural next steps.Source: §7 · p. 12

AI use statement

The authors report using generative AI for research support and paper preparation, while excluding it from theorem proving and interpretation of results.

AI use statement

  • Generative AI assisted with hypothesis refinement, methodology feedback, experiments, literature work, and drafting or editing parts of the paper.The authors state that they reviewed AI-assisted work and take responsibility for the final content.Source: AI use statement · p. 12
  • The authors report that generative AI was not used to prove theorems, formulate mathematical claims, develop theoretical models, or interpret results.Source: AI use statement · p. 12
  • The authors state that they reviewed the AI-assisted work and take responsibility for the final content.Source: AI use statement · p. 12

Supplementary Material

The supplement lists SAM–Muon variants and outlines a convergence analysis under heavy-tailed gradient noise. The proof uses a descent lemma under (ν, L)-smoothness.

A Variants of SAM + Muon

  • The listed SAM–Muon variants include momentum updates using gradients evaluated at perturbed weights, followed by Newton–Schulz orthogonalization.The algorithm listings also include variants with different perturbation constructions, including OP-SOMA and freshly perturbed SOMA.Source: A Variants of SAM + Muon · pp. 18–19

B Convergence Analysis under Heavy-Tailed Gradient Noise

  • The convergence analysis states that SOMA converges under heavy-tailed noise using a proof strategy adapted from MUON’s convergence analysis.The authors generalize the descent lemma to (ν, L)-smoothness and use a bounded p-th central moment assumption.Source: B Convergence Analysis under… · p. 20
  • The supplement proves a descent lemma for functions satisfying (ν, L)-smoothness, bounding function change by a linear term and an order-(ν+1) remainder.The bound applies for ν ∈ (0, 1].Source: B Convergence Analysis under… · p. 20

B.1 Proof of Theorem 1

The proof establishes Theorem 1 by bounding the descent error from momentum estimation and approximate orthogonalization. It combines those bounds with a gradient-gap inequality to obtain an ε-stationarity result.

B.1 Proof of Theorem 1

  • Finite Newton–Schulz iterations introduce a polar-approximation error, which the proof tracks through the descent bound.The proof defines the error by comparing the exact polar factor with the step produced by q Newton–Schulz iterations.Source: §B.1 · p. 21
  • The descent bound includes an error term for the gap between the scaled momentum and the gradient at the perturbed iterate.The proof bounds this term before applying the gradient-gap inequality.Source: §B.1 · p. 22
  • The proof bounds momentum error at the perturbed iterate and combines it with a gradient-gap inequality to obtain ε-stationarity.It decomposes the momentum discrepancy into stochastic-noise, iterate-drift, and initialization terms before closing the argument.Source: §B.1 · pp. 22, 25

C.1 Proof of Corollary 1

The corollary proof specializes the convergence argument to bounded variance and smoothness with ν = 1. It bounds the momentum error terms and then derives a stationary-point complexity.

C.1 Proof of Corollary 1

  • The proof treats the bounded-variance case by setting p = 2 and ν = 1 in the convergence argument.It uses the mini-batch variance bound in bounding the stochastic term.Source: §C.1 · p. 26
  • The perturbation and momentum-drift contribution to Term B is bounded using smoothness and the step-size history.The resulting bound includes the perturbation radius ρ and a term proportional to βηL/(1−β).Source: §C.1 · p. 29
  • After combining the error bounds and closing the descent argument, the proof states a complexity for finding an ε-nuclear norm stationary point.The supplied passage states the complexity result without its expression.Source: §C.1 · p. 30

D SOMA Analysis

The SOMA analysis bounds the bias from reusing a stale gradient by comparing its evaluation point with the current iterate. The additional heavy-tailed experiments were run on CPU hardware.

D.1 Proof of Lemma 1

  • For t ≥ 2, the reused gradient is evaluated at a point within operator-norm distance ρ + η of the current iterate, bounding its smoothness bias by L(ρ + η).For t ≥ 3, the minibatch is fresh conditional on the past, supporting the stochastic-noise bound.Source: §D.1 · p. 31

E Additional Heavy-tailed Experiments

  • The additional synthetic heavy-tailed experiments used CPU-only runs on a single laptop, with eight seeds run in parallel.A single seed took 4–5 seconds; the eight-seed batch took 30–40 seconds.Source: E Additional Heavy-tailed Experiments · p. 31

E.1 Experiment 1. Synthetic experiment with teacher-student model

These experiments compare sharpness-aware MUON methods under heavy-tailed noise, varying noise structure, perturbation geometry, and direction freshness. Across the synthetic settings, spectral perturbations often improve MUON, while the best radius and the balance between accuracy and optimization depend on the setting.

E.1 Experiment 1. Synthetic experiment with teacher-student model

  • In the nonlinear teacher–student experiment, spectral perturbations improve on MUON, while stale perturbations are more robust to large radii than two-gradient methods as tails lighten.At ρ = 1 and α = 3, SSM falls from 47.0% to 29.0% and FP-SOMA from 46.7% to 41.7%, whereas SOMA, OPSPRE, and OPSPOST lose at most 1.1 points from their best.Source: §E.1 · pp. 31–32

E.2 Experiment 2. Linear two-matrix network under heavy-tailed noise

  • At α = 1.1, SOMA lowers MUON’s optimality gap by about 36% under anisotropic noise, while SGD-based methods fail to converge with gaps about 8–14× MUON’s.The optimal SOMA radius falls from ρ⋆ = 1 at α = 1.1 to ρ⋆ = 0.1 at α = 3; at α = 3 only two-gradient methods still noticeably improve on MUON.Source: Figure 3 · p. 33, §E.2 · p. 33, Table 4 · p. 33

E.3 Isotropic Noise for Heavy-tailed experiments

  • With isotropic noise, SOMA improves MUON’s best validation accuracy by 1.4–2.8 points at every tail index and reduces its optimality gap by about 38% at α = 1.1 and 12% at α = 3.At α = 1.1, ADAMW-based methods lead both in accuracy and optimality gap; for α ≥ 2, ADAMW falls behind MUON on the gap.Source: §E.3 · pp. 34–35, Figure 4 · p. 34, Table 5 · p. 35

E.4 Ablation: Frobenius versus Spectral Geometry

  • Spectral geometry is never worse than Frobenius geometry across the tested direction–tail-index pairs, with accuracy gains larger at lighter tails and optimality-gap gains larger at heavier tails.At α = 1.1, spectral geometry lowers the gap of SAM-MUON and SOMA by about a third; for accuracy, its benefit reaches 0.8–2.4 points at α = 3.Source: §E.4 · pp. 36–37, Table 8 · p. 37, Table 9 · p. 37

E.5 Ablation: Fresh versus Stale Perturbations

  • Fresh perturbations improve the optimality gap in every tested case, but stale and fresh directions stay within 1.3 validation-accuracy points, and stale directions are often more accurate.In the spectral ball, stale gaps are 4–24% higher than fresh ones, versus 1–11% in the Frobenius ball; under the heaviest tails, SOMA’s gap is only 4% above FP-SOMA’s.Source: §E.5 · p. 37, Table 10 · p. 38, Table 11 · p. 38

F Additional LLM Experiments

  • The additional LLM experiments report pre-training loss and validation perplexity for nanoGPT and Pythia-70M, plus CIFAR-10 validation accuracy.The passage identifies Table 12 and says Section 6.2.1 discusses these outcomes, but supplies no values or comparisons.Source: F Additional LLM Experiments · p. 38

F.1 Catastrophic forgetting.

Among comparable runs, MUON fine-tuning is associated with more forgetting than ADAMW fine-tuning, while sharpness-aware perturbations reduce that increase only moderately. Forgetting also varies by fine-tuning dataset.

F.1 Catastrophic forgetting.

  • The comparison starts from a shared ADAMW-pretrained model and uses the same fine-tuning schedule and step count, comparing runs at a common fine-tuning loss.Forgetting is measured as the increase in held-out FineWeb-Edu loss after fine-tuning on four new datasets.Source: §F.1 · p. 38
  • For almost every optimizer, forgetting is largest after code or mathematics fine-tuning and smallest after music fine-tuning; SGD forgets least but starts from a clearly higher loss.Source: §F.1 · p. 38
  • Compared with ADAMW, MUON-based fine-tuning is associated with about 1.6× as much forgetting on nanoGPT and 2.1× on Pythia-70M.The figures show the FineWeb-Edu loss increase after fine-tuning; SGD runs are excluded from this comparison because their starting models were clearly weaker.Source: §F.1 · p. 38, Figure 6 · p. 39, Figure 7 · p. 40
  • On nanoGPT, FP-SOMA lowers MUON’s average forgetting increase by 16% (4.20% to 3.50%), but sharpness-aware perturbations only moderately narrow the ADAMW gap; Pythia-70M differences are within 1%.On nanoGPT, the reduction is mainly from code fine-tuning; the reported code increase falls from 0.37 for MUON to 0.24 for SSM and FP-SOMA.Source: §F.1 · pp. 39–40
  • On nanoGPT, one-pass SOMA adds 3% to MUON’s epoch time (272 s), while two-pass SSM and FP-SOMA take 2.4–2.9× as long as MUON.The runtime comparison is based on wall-clock time per epoch of 100 steps.Source: §F.1 · p. 40, Table 15 · p. 41

F.2 Hyperparameter Details

  • Fine-tuning learning rates are calibrated on codeparrot and tulu3, then searched for each dataset, pretrained model, and seed.The initial 19-value log-spaced grid selects a center using mean held-out fine-tuning loss; the subsequent grid uses 11 values.Source: §F.2 · p. 42

Related papers

Loading 2610.32465v1…