PaperlayerView exact revision on arXiv

Source-linked AI summary

Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining

Shuchen Zhu, Yuxin Fang, Mingze Wang, Kun Yuan

arXiv:2608.28442v1cs.LG

TL;DR

LLM pretraining faces noise-dominant gradients and ill-conditioned curvature that leave progress along flat directions slow. The paper introduces curvature-conditioned multiscale momentum with sphere constraints, and reports consistent Muon acceleration with lower terminal loss across diverse architectures and training setups. The study also identifies open questions about learning-rate schedules and slow-momentum preconditioners.

  • Problem

    Noise-dominant gradients and ill-conditioned curvature make progress along flat directions slow, although those directions primarily drive final loss reduction.

  • Method

    The method applies fast- and slow-decay multiscale momentum only in curvature-identified flat directions and uses sphere constraints with parallel transport to prevent norm inflation and effective learning-rate collapse.

  • Results

    MuonM consistently achieves the lowest terminal loss across dense models from 0.12B to 1.4B parameters, outperforming MuonS and Muon by approximately 0.02 and 0.03, respectively.

  • Takeaways & Limitations

    Sphere constraints provide a foundation for exploiting flat-direction multiscale momentum, and the method consistently accelerates Muon across diverse architectures and training setups.

  • Takeaways & Limitations

    The optimal learning-rate schedule for sphere-constrained methods remains open, and slow-momentum preconditioner design was not investigated.

Abstract

from arXiv · show

Pretraining accounts for a large fraction of the total computational cost in LLM training. However, noise-dominant gradients and the highly ill-conditioned loss landscape bring severe challenges. Although modern adaptive optimizers such as AdamW and Muon have achieved great success in large-scale pretraining, their reliance on gradient normalization offers limited mitigation of the ill-conditioned curvature. The progress along flat directions (eigen-directions of small eigenvalues), which dominates the final loss reduction, remains relatively slow. To enhance training dynamics along flat directions, we propose a curvature-conditioned multiscale momentum method with sphere constraints, delivering steady acceleration in LLM pretraining. This multiscale momentum, applied only along flat directions, pairs a slow-decay component for noise reduction with a fast-decay component for rapid curvature adaptation, harnessing their complementary strengths. Crucially, we employ a sphere constraint technique to prevent parameter inflation and excessively rapid effective learning rate decay that would otherwise arise from a naive combination. Extensive experiments show that the proposed method significantly accelerates Muon across diverse architectures (dense, MoE) and model sizes (0.12B--2.3B parameters). Theoretically, we verify the acceleration effect and provide insight into the design principles underlying the flat-direction multiscale momentum.

1 Introduction

LLM pretraining is slowed by noise-dominant gradients and ill-conditioned curvature, especially along flat directions that primarily drive final loss reduction. The paper proposes curvature-conditioned multiscale momentum with sphere constraints and evaluates it across diverse LLM settings.

  • Motivation: Noise dominates stochastic gradients because training tokens vastly exceed per-step batch size, while the loss landscape contains many flat and few sharp directions.The resulting curvature structure makes optimization highly nonconvex and ill-conditioned.
  • Motivation: Progress along flat directions primarily drives final loss reduction but remains slow under learning-rate constraints imposed by sharp directions and noise.Adaptive normalization in AdamW and Muon only partially compensates for ill-conditioned curvature.
  • Motivation: Amplifying flat-direction learning rates also amplifies stochastic noise, potentially destabilizing training.This motivates reducing gradient noise rather than only increasing the flat-direction learning rate.
  • Contributions: The proposed method combines fast- and slow-decay momentum in flat directions while deactivating the multiscale mechanism in sharp directions.The slow component targets noise reduction and the fast component supports adaptation and stability.
  • Contributions: Sphere constraints with parallel transport prevent norm inflation and preserve momentum inheritance across changing tangent spaces.The constraints address effective learning-rate collapse caused by a naive multiscale combination.
  • Contributions: MuonM consistently outperforms Muon with lower terminal loss across dense and MoE models from 0.12B to 2.3B parameters and under cosine-decay and WSD schedules.The loss gains persist under extended token budgets, suggesting scalability to longer training horizons.

2 Preliminaries and Related Works

The paper situates its method among work on ill-conditioned loss landscapes, adaptive preconditioning, multiscale momentum, and effective learning-rate control. Its focus is momentum design rather than improving preconditioner structure.

  • Landscape geometry: Deep-learning loss landscapes contain few sharp eigendirections and many near-zero or negative-eigenvalue flat directions, producing fast–slow training dynamics.Sharp directions exhibit oscillatory behavior, whereas flat directions support slower progress.
  • Adaptive optimizers: AdamW uses coordinate-wise preconditioning, while Muon applies matrix-level preconditioning and achieves favorable convergence speed and scaling behavior.Structured matrix preconditioners further seek to approximate curvature-related information.
  • Momentum: Heavy Ball and Nesterov momentum underpin modern adaptive optimizers, while multiscale methods combine decay rates to damp oscillations and accelerate dynamics.AggMo and AdEMAMix exemplify multiscale momentum extensions.
  • Sphere constraints: Sphere constraints are motivated by separating radial and angular updates because the angular update rate acts as an effective learning rate.Weight norm evolution therefore affects optimization speed through the effective learning rate.

3 Illustrative Examples Motivating Flat-Direction Multiscale Momentum

The motivating examples show that slow momentum reduces gradient-estimation variance but is safe mainly in flat directions. Flat-direction multiscale momentum combines slow variance reduction with fast adaptation while avoiding sharp-direction instability.

  • Momentum decay: Smaller α gives slower momentum decay and lower equilibrium variance when estimating a fixed true gradient.The heavy-ball estimator retains the same mean while its equilibrium covariance decreases with smaller α.
  • Momentum decay: Slow decay is safe mainly in flat directions because small Hessian eigenvalues make gradients change slowly; applying it in sharp directions can cause bias and instability.The simple quadratic example illustrates this curvature-dependent trade-off.
  • Anisotropic momentum: Increasing the flat-direction learning rate promotes progress but amplifies noise, producing oscillatory and relatively less stable acceleration.This motivates anisotropic momentum decay rather than learning-rate amplification alone.
  • Anisotropic momentum: Anisotropic momentum assigns faster decay to sharp directions and slower decay to flat directions through projections Pk and Qk.The decay coefficients satisfy 1 > αsharp ≫ αflat > 0.
  • Robustness: Projection errors and changing Hessian eigenspaces can inject slow-decay momentum into sharp directions or sharp momentum into flat directions, causing instability or persistent oscillations.Slow forgetting makes these effects especially problematic in flat regions.
  • Multiscale momentum: Flat-direction multiscale momentum interpolates fast-decay momentum for rapid adaptation with slow-decay momentum for variance reduction.The fast component is retained to help maintain stability.
  • Multiscale momentum: Under projection errors and varying curvature, the method yields less noisy flat-direction acceleration without amplifying sharp-direction oscillations.The stated mechanism uses a fast component together with the lower-noise slow component.
  • Local analysis: A second-order loss analysis chooses c < 1 so bias injection is negligible while the noise-variance term is suppressed, improving local loss reduction.The method excludes slow momentum from sharp directions, where its bias could be much larger than the gradient signal.

4 Obstacles with Naive Multiscale Momentum: Parameter Inflation and Effective Learning Rate Collapse

The flat-direction multiscale momentum extension estimates flat subspaces and adds slow momentum there, but naive use inflates parameter norms and accelerates effective learning-rate decay. On a 0.12B dense model, this produces only a marginal terminal-loss improvement despite an initially sharper loss drop.

  • Flat-direction momentum: The method estimates flat directions as the bottom 90% singular spaces of each momentum matrix and adds slow momentum there alongside Muon preconditioning.The projection is based on empirical alignment between top singular spaces of momentum and top Hessian eigenspaces.
  • Observed failure mode: The naive slow-momentum variant reduces loss initially, but its terminal-loss improvement remains below 0.01.Both the Muon baseline and the slow-momentum hyperparameters were tuned on the 0.12B dense model.
  • Observed failure mode: Slow momentum amplifies the Frobenius norm of the weights, which accelerates decay of the effective learning rate ∥∆Wk∥F/∥Wk∥F.The norm increase is illustrated for an up projection block in Figure 3.
  • Observed failure mode: The loss curve shows a sharp initial drop as the effective learning rate decays rapidly, followed by the Muon baseline catching up later.The slow-momentum method retains only a marginal edge at the end.
  • Implication: Increasing weight decay to λ = 0.2 restores the norm toward baseline but raises loss further, so hyperparameter tuning alone is insufficient.The default weight-decay coefficient is λ = 0.1.

5 The Flat-Direction Multiscale Momentum Method with Sphere Constraints

The proposed method combines curvature-conditioned fast- and slow-decay momentum in flat directions with sphere-constrained manifold updates. Fixed-norm directional updates, learnable radii, tangent-space operations, and parallel transport address norm inflation while preserving the multiscale design.

  • Sphere-constrained optimization: Sphere constraints replace weight decay to address the weight-norm inflation caused by flat-direction slow momentum.The method applies manifold optimization under a Frobenius sphere constraint.
  • Manifold operations: Updates are projected onto the sphere’s tangent space, retracted back to the sphere, and parallel-transported between successive tangent spaces.Parallel transport preserves momentum norms and prevents accumulated direction errors as the geometry changes.
  • Scope: The sphere-constrained algorithm extends the same construction to tensor parameters such as matrices.The preliminaries are formulated for vectors and then directly extended to other tensor types.
  • Directional parameterization: The method fixes each directional component’s Frobenius norm at R, making the learning rate act as angular velocity, while a learnable scalar radius adapts across blocks and layers.Embedding and output matrices use row-wise learnable radii to accommodate token-level heterogeneity.
  • Algorithm: MuonM combines fast- and slow-momentum updates under flat-direction conditioning and degenerates to MuonS when χ = 0.The complete transformer-block procedure uses Nesterov-type momentum and sphere-constrained Adam for vector parameters.
  • Implementation cost: Flat-direction projections are estimated online with power iterations, while the added computation is nearly negligible because it uses lightweight operations on smaller matrices.The main overhead also includes tangent projection, retraction, and parallel transport, without dense matrix–matrix products.

6 Experiments

Experiments show that MuonM improves terminal validation loss across dense and MoE models, with benefits from flat-direction momentum and sphere constraints extending to longer training horizons.

  • Algorithm Components: Replacing the flat-direction projection with a sharp-direction projection causes a large loss increase.The ablation indicates that applying multiscale momentum along sharp directions is detrimental.
  • Algorithm Components: Parallel transport is indispensable for slow momentum but has negligible impact for fast momentum.Removing transport for slow momentum significantly increases loss, whereas removing it for fast momentum produces little change.
  • Main Results: MuonM consistently achieves the lowest terminal loss across dense models from 0.12B to 1.4B parameters.It outperforms MuonS and Muon by approximately 0.02 and 0.03, respectively.
  • Scaling Behavior under Extended Training: MuonM maintains the lowest terminal loss across extended token budgets, with the performance gap widening as training budgets increase.The evaluation uses the 0.12B dense model over approximately 100–1000 TPP with a tuned learning-rate sweep at each budget.
  • Results on MoE Models: MuonM’s advantages extend from dense architectures to MoE models with 0.64B and 2.3B total parameters.MuonM also attains the lowest terminal loss under the WSD schedule on the 0.64B MoE model.

7 Theoretical Analysis

The theoretical analysis uses a power-law linear-regression proxy to study flat-direction multiscale momentum and its effect on excess-risk reduction. It shows that slow momentum can accelerate convergence, while the analysis is scoped to a finite intrinsic-time regime.

  • 7.1 Problem Formulation: The analysis approximates LLM pretraining dynamics with a quadratic linear-regression model whose Hessian spectrum follows a power-law structure.The model uses diagonal eigenvalues with a power-decayed spectrum and a target parameter satisfying a corresponding power-law assumption.
  • 7.1 Problem Formulation: Flat-direction multiscale momentum combines fast and slow decay components through a projection onto the flat eigenspace.The framework introduces Nesterov-type momentum and projects the slow-momentum mechanism onto directions indexed as flat.
  • 7.2 Theoretical Results: The asymptotic excess-risk analysis separates signal-learning and noise-accumulation contributions and derives an optimal learning rate, loss-reduction ratio, and iteration-complexity speedup.The corollary explicitly quantifies speedup from slow momentum, with the effect becoming more pronounced as χ/αslow increases.
  • 7.2 Theoretical Results: The bounds apply under the power-law assumptions and additional conditions, in the regime afullTk ≲ d^ν.The analysis notes that matching upper and lower bounds are unavailable when Tk ≫ d^ν/afull.
  • 7.2 Theoretical Results: For fixed fast-momentum and damping parameters, incorporating slow momentum with χ > 0 yields accelerated convergence after tuning the learning rate.The analysis attributes the improvement to choosing a proportionally smaller learning rate, which reduces both signal-learning error and noise accumulation.

8 Conclusion and Discussion

The paper proposes flat-direction multiscale momentum with sphere constraints to accelerate LLM pretraining across diverse architectures and training setups. It also identifies open questions about learning-rate schedules and slow-momentum preconditioning.

  • Conclusion: Sphere constraints resolve effective learning-rate collapse in flat-direction multiscale momentum and provide a stronger foundation than standard weight decay.The method is designed to prevent norm inflation while preserving the acceleration potential of multiscale momentum.
  • Conclusion: MuonM consistently accelerates Muon with substantial loss reductions across diverse model architectures and training setups.The conclusion presents this as the principal empirical finding.
  • Limitations: The study leaves the optimal learning-rate schedule for sphere-constrained methods open and does not investigate preconditioner design for slow momentum.The authors posit that improved slow-momentum preconditioning could further accelerate training.
  • Flat-direction estimation: Flat-direction projections estimate sharp subspaces from top singular spaces of stabilized momentum, motivated by their alignment with block Hessian eigenspaces.The method uses separate row-wise and column-wise projections and applies the complementary projection to identify flat directions.
  • Flat-direction estimation: The projection uses msign to estimate subspace projections, replacing computationally expensive QR decomposition.For a matrix U, the flat-direction projection is formed using identity-based complements and tangent-space projection.
  • Parameter-specific updates: MuonM applies to matrix parameters, while sphere-constrained Adam handles vector parameters and direct Adam handles scalar parameters.Embedding and output rows use row-wise sphere constraints with learnable radii.

B.2 Experimental Details of Section 6.2

The experiments compare optimizers across dense and MoE Transformer configurations using cosine and WSD schedules. MuonM uses curvature-conditioned momentum with fixed default coefficients and a warmup for slow momentum.

  • Baselines: Muon applies to 2D Transformer matrices, while AdamW handles RMSNorm, embeddings, and output parameters, using Nesterov momentum, weight decay, and gradient clipping.The Muon configuration uses θmuon = 0.95, λ = 0.1, and clipping threshold 1.0.
  • MuonM: MuonM defaults to χ = 0.2, αfast = 0.05, αslow = 0.001, and a top-10% sharp subspace for each block.The slow-momentum coefficient is warmed up using the learning-rate warmup iterations.
  • Models: The experiments cover dense and MoE decoder-only Transformer architectures with configurations summarized separately for each model class.Dense models use RoPE, SwiGLU, RMSNorm, Pre-Norm, and MHA; MoE models use MLA and partially activated experts.
  • Learning-rate schedules: Cosine schedules decay from peak learning rate to 0.05 × lr_max, whereas WSD holds lr_max before minus-sqrt decay over the final 30% of steps.Both schedules use linear warm-up with min{1000, T/50} steps.
  • Hyperparameters: On the 0.64B MoE model, searched peak learning rates are 5e-3 for Muon and 7e-3 for MuonS, while MuonM uses 4.5e-3.These values are reported under the WSD schedule.
  • Hyperparameters: For the 0.12B dense model, MuonM’s peak learning rate is set to approximately two-thirds of MuonS’s without an additional search.The cited table describes optimal Muon and MuonS peak learning rates across training steps.

C.1 Results with WSD Learning Rate Schedules

Under the WSD learning-rate schedule, MuonM achieves the lowest terminal validation loss among Muon, MuonS, and Muon on the 0.6B MoE model.

  • WSD results: MuonM achieves the lowest terminal loss among Muon, MuonS, and Muon on the 0.6B MoE model under WSD.The compared peak learning rates are 0.005, 0.007, and 0.0045, respectively, for Muon, MuonS, and MuonM.

D A Continuous-Time Perspective for Unifying Enhanced Momentum Methods

The paper places enhanced momentum methods in a shared continuous-time ODE framework and analyzes multiscale momentum through its variance and timescale structure. The Gaussian analysis shows that adding slow momentum can reduce stationary variance relative to fast momentum alone.

  • Unified momentum perspective: The continuous-time analysis derives common ODE forms for AdEMAMix, GPA, EMA-Nesterov, SODA, and MuonM after omitting preconditioning.The derivation assumes gradient-noise-free conditions and constant hyperparameters.
  • Related methods: The ODE derivations distinguish method-specific momentum states and parameter updates, including slow and fast moments in AdEMAMix and look-ahead states in EMA-Nesterov.GPA is instantiated with preconditioned Heavy Ball, while EMA-Nesterov wraps a base optimizer map.
  • Unified momentum perspective: For two-timescale momentum, the polynomial operator A has degree three, generalizing to degree n + 1 for n timescales.This provides a polynomial framework for extending the continuous-time representation.
  • Variance analysis: In the Gaussian analysis, slow momentum reduces stationary variance by a factor of up to αslow(2 −αfast) relative to fast momentum alone.The analysis models stochastic gradients as Gaussian and preserves the stationary mean of the update.
  • Variance analysis: Increasing the weight of slow momentum monotonically reduces stationary variance when 0 < αslow < αfast ≤ 1.The limiting cases are fast-momentum-only variance αfast/(2 −αfast) and slow-momentum-only variance αslow/(2 −αslow).

E.2 Weight Norm Inflation by Slow Momentum

Slow momentum can inflate parameter norms under noise-dominated, scale-invariant dynamics. Weight decay makes stationary norms finite, but stronger decay alone may worsen the loss, motivating weight-decay-free control methods.

  • For scale-invariant loss, Brownian-motion-like stochastic movement can drive parameter norms upward, while weight decay suppresses this growth.The analysis assumes i.i.d. Gaussian stochastic gradients and a nearly full-rank flat-direction projection.
  • The proposition establishes finite-iteration norm inflation for the slow-momentum process under ρ := 1 − ηλ > 0.
  • Slow momentum contributes to norm inflation through its positive coefficient χ > 0 and the asymptotic dominance of the slow component over the fast component.The asymptotic imbalance becomes stronger as the iteration index grows.

F Proof

The proof analyzes momentum dynamics through discrete coordinate recursions, transition-matrix spectra, spectral projections, and Hessian-aligned noise. It establishes root separation and uniform bounds under explicit assumptions and small-learning-rate conditions.

  • Noise structure: In later training, when E ≲ σ^2, gradient noise aligns well with the Hessian and stochastic gradients become noise-dominated.Earlier in training, the noise is instead dominated by the loss scale.
  • Discrete recursions: The discrete parameter-momentum recursion reduces each eigen-direction to a transition system with fast and slow decay factors ρfast and ρslow.The shared noise term appears across the three state entries, with the slow contribution selected by the sharp/flat indicator.
  • Spectrum analysis: For sufficiently small positive learning rates, the root issuing from 1 remains separated from the roots issuing from ρfast and ρslow.The characteristic polynomial governs the convergence rate, and the proof also establishes three real roots under the stated sign conditions.
  • Spectrum analysis: The largest-root displacement satisfies 0 ≤ δx ≤ 2ax ≤ αslow/4 and converges to zero as x approaches zero.This follows from a contraction-mapping argument on the interval [0, αslow/2].
  • Spectral projections: Gershgorin disks isolate one eigenvalue near 1 from the complementary two eigenvalues near the momentum decay factors.The resulting spectral projections are complementary, with P0 associated with the largest eigenvalue and P1 with the other two-dimensional subspace.
Loading 2608.28442v1…