PaperlayerView exact revision on arXiv

Source-linked AI summary

Controlled LLM Training on Spectral Sphere

Tian Xie, Haoming Luo, Haoyu Tang, Yiwen Hu, Jason Klein Liu, Qingnan Ren, Yang Wang, Wayne Xin Zhao, Rui Yan, Bing Su, Chong Luo, Baining Guo

arXiv:2601.08393v3cs.LGcs.AI

TL;DR

Large-model training needs fast convergence without losing width-invariant stability, but Muon controls updates while allowing weights to drift. The paper introduces SSO, which constrains both on the spectral sphere, and reports consistent advantages over AdamW and Muon across diverse architectures, alongside practical stability benefits.

  • Problem

    µP provides width-invariant activation control, but Muon is only “half-aligned” because it constrains updates while allowing weights to drift.

  • Method

    SSO derives steepest descent on the spectral sphere while enforcing constraints on both weights and updates, with a Megatron implementation for large-scale training.

  • Results

    SSO consistently outperforms AdamW and Muon across Dense 1.7B, MoE 8B-A1B, and 200-layer DeepNet pretraining experiments.

  • Takeaways & Limitations

    SSO preserves stable µP learning-rate transfer while improving MoE router load balancing, suppressing deep-network outliers, and bounding activations.

  • Takeaways & Limitations

    The current bracket-and-bisect root solver introduces non-trivial latency and synchronization overhead, motivating GPU-native and faster solver improvements.

Abstract

from arXiv · show

Scaling large models requires optimization strategies that ensure rapid convergence grounded in stability. Maximal Update Parametrization ($\boldsymbolμ$P) provides a theoretical safeguard for width-invariant $Θ(1)$ activation control, whereas emerging optimizers like Muon are only ``half-aligned'' with these constraints: they control updates but allow weights to drift. To address this limitation, we introduce the \textbf{Spectral Sphere Optimizer (SSO)}, which enforces strict module-wise spectral constraints on both weights and their updates. By deriving the steepest descent direction on the spectral sphere, SSO realizes a fully $\boldsymbolμ$P-aligned optimization process. To enable large-scale training, we implement SSO as an efficient parallel algorithm within Megatron. Through extensive pretraining on diverse architectures, including Dense 1.7B, MoE 8B-A1B, and 200-layer DeepNet models, SSO consistently outperforms AdamW and Muon. Furthermore, we observe significant practical stability benefits, including improved MoE router load balancing, suppressed outliers, and strictly bounded activations.

1 Introduction

The paper frames stable, rapid convergence as requiring µP-aligned control of both weights and updates. SSO addresses this by deriving constrained steepest descent on the spectral sphere and demonstrating broad training benefits.

  • Motivation: µP stabilizes width scaling by controlling spectral norms of weights and updates, whereas unconstrained weight drift can destabilize effective step sizes.Soft regularization and initialization are described as insufficient over long training horizons.
  • Motivation: Muon is “half-aligned” because it constrains updates but leaves weights unconstrained, with unstable activations still observed during training.The paper specifically notes attention-logit explosions and reliance on architectural patches for stability.
  • Method: SSO derives a steepest descent direction within the spectral-sphere geometry while constraining both weights and updates.Its constrained tangent-space solution uses Lagrange multiplier search followed by retraction.
  • Implementation: SSO is implemented in Megatron with spectral scaling guidelines and distributed atomic-module sharding to support large-scale training.The implementation also addresses solver workload imbalance and matrix-operation efficiency.
  • Results: SSO consistently outperforms AdamW and Muon across Dense 1.7B, MoE 8B-A1B, and 200-layer DeepNet pretraining experiments.The reported benefits include improved MoE router balancing, suppressed deep-network outliers, and tunably bounded activations.

2 Preliminary

The preliminary framework connects µP scale invariance to spectral-norm constraints on activations, weights, and updates, and interprets optimizers through their steepest-descent geometries. Muon uses spectral-norm steepest descent for updates but does not constrain forward weights.

  • 2.1 Maximal Update Parametrization (µP): For y = W x, µP requires RMS activations to remain Θ(1), equivalent to an RMS-to-RMS operator-norm condition on W.The operator norm characterizes how activation norms change through a linear layer.
  • 2.1 Maximal Update Parametrization (µP): The corresponding spectral constraint scales the weight norm as Θ((dout/din)^1/2), with an analogous requirement for parameter updates.Together, these requirements are called the spectral µP condition.
  • 2.2 Steepest Descent under Different Norms: Different optimizers correspond to different steepest-descent geometries: SGD uses the Frobenius norm, AdamW the ℓ∞ norm, and Shampoo the spectral norm.This interpretation is presented without convexity assumptions after disabling exponential moving averages.
  • 2.2 Steepest Descent under Different Norms: Steepest descent minimizes a first-order loss approximation under a hard step-size constraint determined by a chosen norm and learning rate.The gradient is G := ∇W L(W).
  • 2.3 Muon Optimizer: Under the spectral norm, the unique steepest-descent direction is the matrix sign function, which orthogonalizes the gradient across active singular directions.Muon approximates this operation efficiently with Newton–Schulz iterations on GPUs.
  • 2.3 Muon Optimizer: Muon constrains only the update direction, leaving forward weights unconstrained and thereby permitting unstable µP feature learning in hidden-state RMS.This motivates constraining both weights and updates.
  • 2.3 Muon Optimizer: Across 70M–1.8B models, Spectral Sphere achieves stable learning-rate transfer and lower optimal loss than Muon.The comparison spans a 25× range of model sizes.

3 Method

SSO performs steepest descent on a spectral sphere while constraining both weights and updates. It enforces the weight constraint through tangent-space optimization, root solving, and retraction.

  • 3.1 Optimization Target Formulation: SSO targets a spectral sphere of radius R for hidden-layer weights and defines unit updates Φ under spectral-norm constraints.The update is parameterized as ΔW = ηRΦ with both ∥Φ∥2 = 1 and ∥W − ηRΦ∥2 = ∥W∥2 = R.
  • 3.2 First-Order Tangent Space Constraint: With a unique top singular value, Θ = u1v1^⊤ is the gradient of the spectral norm and defines the tangent constraint.The first-order invariance condition requires ⟨Θ, Φ⟩ = 0.
  • 3.2 First-Order Tangent Space Constraint: SSO solves the constrained steepest-descent direction by maximizing ⟨G + λΘ, Φ⟩ under ∥Φ∥2 = 1 using a Lagrange multiplier.For fixed λ, the direction is obtained analytically through the matrix-sign formulation, while λ is found numerically.
  • 3.2 First-Order Tangent Space Constraint: The root function h(λ) is monotonic, and λ⋆ lies within a finite interval, enabling bracketing and bisection.Empirical curves for random matrices show the root lies close to zero, while the theoretical interval is [−2∥G∥∗, 2∥G∥∗].
  • 3.3 Second-Order Manifold Constraint: Retraction restores ∥W∥2 = R after updates, preventing accumulated second-order drift from moving weights off the spectral sphere.The implementation reuses one power-iteration singular triplet for both retraction and the tangent projector, and bounded weights make hidden-weight decay redundant.
  • 3.4 Overall Algorithm & Interpretation: The algorithm combines momentum normalization, power iteration, spectral retraction, bisection for λ⋆, and a µP-scaled update.The update is Wt+1 ← Wt − ηRΦt, with Φt determined by the tangent-constrained matrix-sign direction.

4 Algorithm Details

Algorithm details calibrate spectral radius, learning-rate scaling, and module granularity for stable optimization. Ablations support moderate radius scaling, the Spectral µP scaler, and finer-grained module treatment.

  • 4.1 Spectral Radius Scale: The radius scale c controls branch output magnitude relative to the residual stream and therefore influences depth-wise signal propagation.The paper describes c as balancing Attention/FFN contributions against the skip connection.
  • 4.1 Spectral Radius Scale: A moderate radius scale, such as c = 2.0, achieves the best performance in the radius ablation.AbsMax follows radius scale monotonically, while RMS exhibits power-law scaling with c.
  • 4.2 Learning Rate Scaler: The evaluated scalers include Spectral µP, Align-Adam-RMS, and Spectral Kaiming, each targeting a different norm or initialization-based invariance.Spectral Kaiming targets the spectral norm under Kaiming initialization and is described as preventing vanishing pre-activations in bottleneck layers.
  • 4.2 Learning Rate Scaler: The Spectral µP scaler outperforms Align-Adam-RMS under the spectral µP condition.The scaler enforces RMS-to-RMS operator-norm invariance for weights and updates, whereas Align-Adam-RMS aligns update RMS with AdamW.
  • 4.3 Module Granularity: Fused tensors are decomposed into independent modules because their distinct functional roles make unified spectral constraints suboptimal.The default granularity splits attention QKV per head and FFN gate/up projections for separate initialization and optimization.
  • 4.3 Module Granularity: Splitting QKV per head produces the largest performance gain, while splitting FFN gate/up weights yields nearly the same loss as no splitting.The FFN split is retained to respect the projections’ distinct functional roles.

5 Infrastructure Design

SSO’s infrastructure addresses the root-solver overhead of spectral updates with sharding, balancing, adaptive kernels, parallelism, caching, and mixed precision, while several latency and numerical-efficiency limitations remain.

  • Bottleneck Analysis: The bracket-and-bisect solver introduces workload imbalance, extra matrix computations, and GPU–CPU synchronization overhead at every update.Variable solver depths affect device workloads, while each root-search evaluation adds computation and synchronization.
  • Optimization Pipeline: Atomic module sharding preserves full weight matrices required for spectral operations, unlike flattened-buffer sharding designed for element-wise optimizers.Parameters are partitioned as minimal independent matrices so spectral updates can be computed without flattening.
  • Optimization Pipeline: A ping-pong allocation strategy assigns size-sorted atomic modules across data-parallel ranks to balance variable solver workloads.The alternating zigzag assignment outperforms greedy size-descent sorting and default round-robin allocation in the reported experiments.
  • Optimization Pipeline: Adaptive kernel selection uses JIT PyTorch kernels for matrices below 512 and Triton SYRK kernels for matrices at least 512.The dispatcher accounts for matrix-dimension sensitivity and exploits symmetry in Newton–Schulz iterations for large matrices.
  • Optimization Pipeline: Multi-stream execution, BF16 power iteration, FP32 msign, and cached singular vectors reduce launch, precision, and convergence costs.Independent matrices use multiple CUDA streams, while cached singular vectors initialize subsequent power iterations.
  • Bottleneck Analysis: The current solver remains limited by 5–7 bisection steps, projection-accuracy requirements, residual-stream outliers, and unresolved low-precision-training questions.Future work targets GPU-native root solving, optimized kernels, fully manifold-constrained architectures, and low-bit training.

6 Scaling Experiments

Scaling experiments evaluate SSO across dense, MoE, and extreme-depth models under specified training protocols. Across these settings, SSO improves efficiency, validation loss, routing balance, and deep-network stability relative to AdamW and Muon.

  • Experimental Setup: The experiments use 100 billion training tokens, 1 billion validation tokens, BF16 training, and a 0.1 weight decay shared across optimizers.The dense-model protocol uses a 1.7B architecture, 500 warmup steps, and cosine decay to 10% of peak learning rate.
  • Dense 1.7B: 19% fewer steps let Spectral Sphere reach AdamW’s final validation loss of 2.588, while Muon required 12% fewer steps under AdamW’s learning-rate setting.The comparison uses 23k steps as AdamW’s reference and favors AdamW because its learning rate was selected rather than the higher optimum for spectral methods.
  • MoE 8B-A1B: The MoE 8B-A1B model uses 64 experts with top-4 routing plus one shared expert, following a DeepSeek-V3-like configuration.Its 27 layers comprise one dense FFN followed by 26 MoE layers.
  • MoE 8B-A1B: Spectral Sphere achieves the lowest validation loss and best expert load balance in MoE training, while AdamW shows larger MaxVio and frequent routing spikes.Constraining each expert on the spectral sphere further improves load balance compared with Muon.
  • DeepNet 200-Layer: In the 200-layer DeepNet stress test, Spectral Sphere outperforms baselines with lower training loss and higher stability.The model extends the baseline’s 28 layers to 200 layers to evaluate optimizer stability under extreme depth.

7 Discussion

The discussion presents SSO as a spectral-µP optimizer that constrains weights and updates while providing a Megatron implementation and practical guidance. It also distinguishes SSO’s flexible spectral bound from stricter Stiefel constraints and acknowledges solver latency.

  • Discussion: SSO derives steepest descent constrained within both weight and update manifolds to control activations at the desired µP scale.The paper identifies the spectral sphere as the geometry for stable feature learning and reports rapid convergence grounded in training stability.
  • Discussion: SSO constrains only the maximal singular value, allowing internal singular values to evolve below the bound unlike the exactly unit-singular-value Stiefel manifold.This relaxation avoids the Stiefel manifold’s overly rigid isotropy.
  • Discussion: The Megatron-LM recipe covers atomic granularity, ping-pong load balancing, learning-rate scaling, and spectral-radius selection for spectral optimizers.The paper presents these components as a systematic practical recipe and outlines pathways to reduce root-solver latency.
  • Discussion: For infrastructure-cost-sensitive scenarios, the paper recommends MuonSphere as a lower-overhead variant retaining equivalent activation control.This recommendation accompanies the acknowledgment that SSO’s current root solver introduces non-trivial latency.

A.1 Duality with Spectral Norm

The appendix establishes the dual relationship between the nuclear and spectral norms through singular-vector structure. This result supports the spectral steepest-descent derivation used by SSO.

  • Duality Result: Theorem A.1 states that the nuclear norm is the dual norm of the spectral norm.The proof starts from the singular value decomposition of G and uses unit-norm singular vectors.
  • Proof: For a feasible matrix T with spectral norm 1, singular-vector inner products provide the bound needed to attain the nuclear norm.The construction has nonzero singular values equal to 1, making T feasible and achieving the upper bound.

A.2 Proofs: Localization of the Root of h(λ)

The appendix proves that h(λ) is monotonic and has a root λ⋆, then localizes every such root to a bounded interval. Numerical experiments further support the theoretical behavior and the practicality of the root solver.

  • The appendix states that the computational overhead of the iterative λ-solver is negligible relative to the full training process.
  • h(λ) is monotonic non-decreasing in λ because the maximization defining Φ⋆(λ) preserves the ordering of ⟨Θ, Φ⋆(λ)⟩.
  • A root λ⋆ satisfying h(λ⋆)=0 exists, using the limiting values of h and its intermediate value property.
  • Every root λ⋆ lies in [−2∥G∥∗, 2∥G∥∗], because h(λ)>0 above 2∥G∥∗ and h(λ)<0 below −2∥G∥∗.
  • Random-matrix experiments show that h(λ) is monotonic and its root lies close to λ=0, consistent with the theoretical analysis.

A.3 Dynamic Spectral Weight Decay

Dynamic spectral retraction provides a learning-rate-scaled alternative to exact projection, while the appendix also describes spectral constraints and an MoE scaling factor intended to stabilize training.

  • Spectral retraction counters accumulated higher-order errors by mapping weights back toward the spectral sphere during training.The exact constraint is ∥W∥2=R; retraction compensates for drift introduced by Taylor-expansion remainder terms.
  • The hard variant enforces ∥W∥2=R at every step, whereas the dynamic variant applies a soft, learning-rate-scaled correction toward the target radius.
  • Without weight constraints, Muon training dynamics become unstable and would hurt performance, contrasting with the constrained Spectral Sphere setup.
  • Both retraction variants remain first-order equivalent to steepest descent on the spectral manifold, with timing differences of O(η^2).
  • For MoE routing with one shared expert and Top-4 sigmoid routing, numerical simulation finds M≈2.0 to balance routed and shared contributions.The routed experts otherwise have lower variance because sigmoid probabilities are typically small, causing the optimizer to neglect them.
  • The width-scaling experiment varies model sizes from 70M to 1.8B and trains them on 30B tokens using the Spectral µP LR Scaler.
  • AdamW shows worse validation loss with a drifting optimal learning rate in the reported µP LR grid search.
Loading 2601.08393v3…