PaperlayerView exact revision on arXiv

Source-linked AI summary

Fantastic Pretraining Optimizers and Where to Find Them II: Hyperball Optimization

Kaiyue Wen, Xingyu Dang, Kaifeng Lyu, Tengyu Ma, Percy Liang

arXiv:2606.16899v1cs.LG

TL;DR

Matrix-based optimizer gains over AdamW have been observed to shrink as model and data scale grow, while weight decay’s role is puzzling for scale-invariant weights. Hyperball fixes matrix weight and update norms, replacing implicit norm control with an explicit constraint. On 1.2B-parameter Qwen3-style models, MuonH achieves 20–30% token-equivalent speedup over its weight-decay counterpart, and Hyperball improves learning-rate transfer across widths and depths.

  • Problem

    Matrix-based optimizer speedups over AdamW shrink from roughly 30% to about 10% as model size and data scale grow, while weight decay is puzzling as classical ℓ2 regularization for scale-invariant weights.

    Source: §1 · pp. 1–2

  • Method

    Hyperball wraps a base optimizer by explicitly constraining matrix weight and update norms, replacing weight decay’s implicit norm control.

    Source: §1 · pp. 1–2

  • Results

    20–30% token-equivalent speedup over its weight-decay counterpart is achieved by MuonH on 1.2B-parameter Qwen3-style models, versus about 10% for MuonW.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Hyperball improves learning-rate transfer across model widths and depths compared with decoupled weight decay.

    Source: §6 · p. 17

    Frobenius constraints approximate spectral constraints when the singular-value spectrum is not dominated by one direction; concentrated spectra may call for a spectral Hyperball variant.

    Source: §2 · p. 4

Abstract

from arXiv · show

Matrix based optimizers such as Muon can substantially speed up language model pretraining, but their gains over AdamW are observed to shrink as model size and data scale grow when using standard constant decoupled weight decay. We propose Hyperball, a simple optimizer wrapper that addresses this issue. Given a base optimizer such as Adam or Muon, Hyperball sets the Frobenius norms of weight matrices and their corresponding optimizer updates to fixed constants. On Qwen3 style models up to 1.2B parameters, Muon Hyperball achieves 20--30% token equivalent speedup over weight decay baselines. Hyperball also improves learning rate transfer across widths and depths compared to decoupled weight decay. This method is motivated by prior theory showing that training with weight decay leads to an equilibrium weight norm that only depends on the training hyperparameters. Through this mechanism, the weight decay then decides the angular learning rate, i.e. how fast the direction of the weight matrix changes.

1 Introduction

Hyperball addresses shrinking matrix-optimizer gains at larger scales by replacing weight decay’s implicit norm control with fixed weight and update norms. On 1.2B-parameter Qwen3-style models, MuonH delivers 20–30% token-equivalent speedup over its weight-decay counterpart.

1 Introduction

  • Matrix-based optimizer speedups over AdamW shrink from roughly 30% to about 10% as model size and data scale grow.For scale-invariant matrices, weight decay is puzzling as classical ℓ2 regularization because the loss is unchanged by rescaling the weights.Source: §1 · pp. 1–2
  • Hyperball replaces weight decay’s soft norm control with explicit constraints on weight and update norms.The weights remain on a sphere of radius R, and each step uses a normalized update followed by projection back onto the sphere.Source: §1 · p. 2, Figure 1 · p. 2
  • 20–30% token-equivalent speedup over its weight-decay counterpart is achieved by MuonH on 1.2B-parameter Qwen3-style models, while MuonW gains about 10%.Across depth and width sweeps, the maximal learning-rate-window drift is about 1.4× for AdamH and MuonH, versus 2–4× for AdamW and MuonW.Source: §1 · p. 2
  • Under the cited theory, weight decay’s equilibrium radius shapes angular step size, which Hyperball controls directly by fixing radius and update norm.The analysis relates angular step size to learning rate and weight decay mainly through their product ηλ.Source: §1 · p. 2

2 Method

Hyperball applies a Frobenius-norm constraint to matrix updates and projects weights back to a fixed-radius sphere. Its Frobenius version is computationally cheaper than exact spectral projection, while its approximation to spectral control depends on the weight matrix’s singular-value spectrum.

2 Method

  • 20–30% speedup over the AdamW scaling law is sustained by MuonH across Chinchilla ratios 1×–8× at 1.2B parameters, versus about 10% for MuonW.The gap grows with training duration in both the setup without QK-Norm and the QK-Norm setup.Source: Figure 2 · p. 3
  • Hyperball takes an unconstrained optimizer step with norm ηtR, then radially renormalizes the trial point to radius R.The base optimizer update can be supplied by any optimizer.Source: §2 · pp. 2–3
  • Hyperball constrains attention and MLP weight matrices, while embeddings, normalization gains, and other scalar parameters use a standard optimizer because their norms can carry semantic information.Adam is used for those unconstrained parameters in the experiments.Source: §2 · p. 3
  • Frobenius-sphere projection costs O(N^2) per matrix, whereas exact spectral projection generally requires an SVD and costs O(N^3).Frobenius control behaves similarly to spectral control when the singular-value spectrum is not dominated by a single direction; concentrated spectra motivate a spectral variant.Source: §2 · p. 4

3 Hyperball Experiments

The experiments compare Adam and Muon with decoupled weight decay against their Hyperball variants. They report end-to-end speedups, improved learning-rate transfer across model scales, and lower best C4 loss in an overtrained setting.

3.1 Setup

  • Experiments compare Adam and Muon weight-decay and Hyperball variants, measuring token-equivalent speedup and best-learning-rate drift across scales.The speedup metric compares each method’s final loss with the AdamW scaling law; drift tracks the range of optimal learning rates across scales.Source: §3.1 · p. 4

3.2 End-to-end speedup

  • In the weight-decay baseline setting, 1.2B Qwen3-style models trained over 1×–8× Chinchilla ratios show MuonH sustaining 20–30% token-equivalent speedup over the AdamW scaling law, versus MuonW’s ≈10%.At 8B parameters, MuonH yields a further 0.04 loss improvement over the AdamW baseline, with both runs using manually chosen hyperparameters.Source: §3.2 · p. 4, Figure 3 · p. 5
  • 3.2780 average validation loss in 3125 steps: KL-SOAP-H is the fastest loss-at-step result in the Track 3 comparison, and Hyperball variants improve matched weight-decay baselines.KL-SOAP-H combines KL-SOAP with Hyperball.Source: Figure 4 · p. 5, §3.2 · p. 5

3.3 Hyperparameter transfer

  • ≈1.4× optimal learning-rate drift: Hyperball variants show this across depths L ∈ {4, …, 512} and widths d ∈ {128, …, 2048}, versus ≈3× for AdamW and ≈4× for MuonW at depth.Transfer sweeps use 10B tokens per run.Source: §3.3 · p. 6, Figure 5 · p. 6, Figure 6 · p. 6

3.4 Overtrained Setting

  • Across 1B–128B tokens, MuonH attains lower best C4 validation loss than MuonW; fitted asymptotes are 3.065 versus 3.079.The fits use L(N) = L∞ + AN^-α.Source: §3.4 · p. 6, Figure 7 · p. 7

4 Theory

This section explains how scale constraints preserve expressivity in normalized networks and how weight decay shapes equilibrium norms and angular learning speed. It examines the resulting norm, gradient, loss, and transfer patterns, including how Hyperball makes radius and normalized update length explicit.

4.1 Expressivity

  • In normalized networks, fixing a weight matrix’s Frobenius norm need not reduce represented functions because a trainable normalization gain can absorb its scale.The joint rescaling of the weight and gain leaves the represented function unchanged.Source: §4.1 · p. 7

4.2 Optimization

  • Under the stated assumptions, weight decay drives the matrix to an equilibrium radius and angular step determined by update norm U and autocorrelation sum Cα.For scale-invariant blocks, a larger radius means smaller angular movement for a fixed update norm; AdamW and SVD Muon/Moonlight use different autocorrelation quantities.Source: §4.2 · pp. 8, 13
  • For scale-invariant losses, gradient norms scale inversely with the weight radius when the weight direction is fixed.The theory predicts ∥∇W L(Wt)∥F ∝ 1/Rt under that condition.Source: §4.2 · pp. 13–14
  • Hyperball makes weight decay’s coupled effects explicit by fixing the weight norm and normalized update length directly.Weight decay otherwise sets the radial scale, which in turn determines angular learning speed through ∥ut∥F /R⋆.Source: §4.2 · p. 14

4.3 Empirical Validation

  • In a 1.2B AdamW cosine-decay run, Q/K/V norms rose during warmup then fell during decay, while their gradient norms increased late in training.The norm pattern agrees with the radius prediction, and the gradient pattern with the inverse-radius prediction for scale-invariant blocks.Source: §4.3 · p. 14
  • With ηλ fixed, AdamW runs at (η, λ) = (0.002, 0.2) and (0.004, 0.1) had nearly overlapping train-loss curves, while equilibrium Q/K norms roughly doubled in the larger-learning-rate run.The theory predicts nearly fixed angular step size and equilibrium radius proportional to η under this rescaling.Source: §4.3 · p. 14
  • Weight decay’s larger effective step amplifies early hill-direction oscillations but accelerates progress along the river, so decay later suppresses oscillations and reveals that progress [Wen et al., 2024].The weight-decay run maintains a larger effective step size throughout training and ultimately reaches lower loss.Source: §4.3 · pp. 15–16
  • With the same learning-rate schedule, weight decay starts at higher loss but reaches lower validation loss, while maintaining a larger late angular proxy.Without weight decay, the weight norm grows and the angular proxy decays.Source: §4.3 · pp. 14, 16, Figure 11 · p. 16
  • Transfer is less sensitive to initialization scale than weight-decay scaling; Hyperball fixes radius and normalized update length, improving learning-rate transfer across widths and depths [Kosson et al., 2025, Blake et al., 2024, Fan et al., 2025, Wang and Aitchison, 2024, Qiu et al., 2025].The radial recursion forgets the initial radius and converges to a norm set by learning rate, weight decay, and optimizer-dependent update norm; changing weight-decay scaling changes equilibrium norm and angular step size.Source: §4.3 · p. 16

5 Related work

Prior work links weight decay in normalized networks to norm dynamics and effective angular learning rates, rather than mainly to classical capacity penalties. Hyperball is distinguished from earlier normalization methods by directly fixing both matrix and update Frobenius norms.

Weight decay in normalized networks

Norm constraints on weights and updates

Fixed norms in LLM pretraining and manifold optimization

  • Recent language-model methods use columnwise, per-channel, rowwise, or target-variance constraints [Loshchilov et al., 2024; Fu et al., 2025; Franke et al., 2025; Owen et al., 2025].Other approaches optimize on Stiefel or spectral-sphere manifolds [Thinking Machines, 2025; Su, 2025a; Bernstein, 2025, Cesista, 2025; Newhouse et al., 2025]; Hyperball uses a Frobenius sphere at O(N^2) cost per matrix versus O(N^3) for spectral projections.Source: §5 · p. 17

6 Conclusion

Hyperball replaces implicit weight-norm control from weight decay with explicit constraints on matrix and update norms. Experiments report improved scaling behavior for matrix-based optimizers and more reliable learning-rate transfer, while leaving open which constraint best fits different models and optimizers.

6 Conclusion

  • Across experiments, explicit norm constraints improve matrix-based optimizer scaling and learning-rate transfer across widths, depths, and training budgets.The constraints replace weight decay’s implicit norm control on matrix and update norms.Source: §6 · p. 17
  • Hyperball uses Frobenius constraints, but spectral, rowwise, columnwise, hybrid, or architecture-dependent constraints may better match some models and optimizers.The conclusion also identifies sharper theory for weight-normalized training and explicit norm constraints as an open direction.Source: §6 · p. 17
  • The paper leaves open how removing or fixing the radial degree of freedom shapes the training trajectory.Source: §6 · p. 17

Related papers

Loading 2606.16899v1…