Technical Summary: Hyperball Optimization

Problem Statement

Matrix-based optimizers, such as Muon, have demonstrated substantial speedups over AdamW during language model pretraining. However, empirical observations indicate that these gains diminish as model size and data scale increase, particularly when using standard constant decoupled weight decay. In large-scale settings (e.g., 1.2B parameter models), the speedup of Muon over AdamW shrinks from approximately 30% to roughly 10%. This degradation motivates the search for an optimization strategy that can maintain high efficiency at higher compute scales.

The core issue lies in the mechanism of decoupled weight decay. In scale-invariant layers (common in modern Transformers with normalization), weight decay acts not merely as a regularizer but as an indirect controller of the "angular learning rate"—the speed at which the direction of the weight matrix changes. Theoretical analysis suggests that weight decay drives the weight norm to an equilibrium state determined by hyperparameters ($\eta$ and $\lambda$). This equilibrium norm, in turn, dictates the effective angular step size. However, this indirect calibration can become suboptimal as scales change, leading to the observed shrinkage in optimizer performance.

Methodology: Hyperball

The authors propose Hyperball, a simple optimizer wrapper designed to address the scaling limitations of weight decay. Hyperball replaces the implicit, soft control of weight norms via weight decay with an explicit constraint on the Frobenius norms of weight matrices and their updates.

Core Mechanism

Given a base optimizer (e.g., Adam or Muon) that produces an update $u_t$ for a parameter matrix $W_t$, Hyperball enforces two conditions:

1. The Frobenius norm of the weight matrix $\|W_t\|_F$ is fixed to a constant radius $R$ (initialized as the norm of the initial weights).
2. The update direction is normalized, and the step size is scaled explicitly.

The update rule is defined as:
$$W_{t+1} = R \cdot \text{Normalize}\left( W_t - \eta_t R \cdot \text{Normalize}(u_t) \right)$$
where $\text{Normalize}(X) = X / \|X\|_F$.

Geometrically, this constrains the optimization trajectory to the surface of a hypersphere with radius $R$. Each step moves along the normalized update direction by a distance proportional to $\eta_t R$ and is immediately projected back onto the sphere. This decouples the magnitude of the weights from the direction of the update, allowing the optimizer to navigate primarily through weight directions.

Implementation Details

• Scope: Hyperball is applied to attention and MLP weight matrices in prenorm Transformers.
• Exclusions: Embeddings, normalization gains, and scalar parameters are updated using standard optimizers (e.g., Adam) because their norms carry semantic information.
• Complexity: Projection onto the Frobenius sphere is $O(N^2)$ per matrix, which is computationally cheaper than spectral norm projection ($O(N^3)$) that requires SVD. The authors note that in high stable-rank regimes (typical in LLMs), Frobenius constraints behave similarly to spectral constraints.

Key Contributions

1. Hyperball Wrapper: A general-purpose wrapper that transforms any base optimizer into a variant with fixed weight and update norms.
2. Theoretical Insight: The paper provides a theoretical framework showing that for scale-invariant losses, decoupled weight decay implicitly sets an equilibrium radius and an angular step size. Hyperball makes this mechanism explicit by fixing the radius and the normalized update length directly, rather than relying on the equilibrium dynamics of weight decay.
3. Improved Scaling: The method is shown to maintain optimizer speedups at larger scales where standard weight decay baselines fail to do so.
4. Learning Rate Transfer: Hyperball significantly improves the transferability of learning rates across different model widths and depths compared to weight decay baselines.

Experimental Results

The authors evaluated Hyperball on Qwen3-style models (up to 1.2B parameters) and larger scales (up to 8B parameters) using various benchmarks.

• Token Equivalent Speedup: On 1.2B parameter models, Muon Hyperball (MuonH) achieved a 20–30% token equivalent speedup over AdamW baselines. In contrast, standard Muon with weight decay (MuonW) only achieved about 10% speedup at this scale. The gap between MuonH and MuonW widened with longer training durations.
• Benchmark Performance:
  • On the Marin speedrun benchmark, AdamH and MuonH matched weight decay baselines that were approximately 10% larger.
  • On the public modded-nanogpt Track 3 optimization benchmark, Hyperball variants (including KL-SOAP-H) reached lower validation losses faster than their weight decay counterparts. For instance, KL-SOAP-H reached a validation loss of 3.2780 in 3125 steps, outperforming the tuned MuonW baseline.
• Hyperparameter Transfer:
  • Depth Scaling: When sweeping depths from 4 to 512, the optimal learning rate drift for Hyperball variants was approximately 1.4×, compared to 2–4× for AdamW and MuonW.
  • Width Scaling: Similarly, across widths from 128 to 2048, Hyperball maintained a stable learning rate window with only 1.4× drift.
• Overtrained Setting: In data scaling experiments (130M parameter model, 1B–128B tokens), MuonH achieved a lower asymptotic validation loss ($L_\infty = 3.065$) compared to MuonW ($L_\infty = 3.079$).

Significance and Claims

The paper claims that Hyperball offers a more robust and efficient alternative to decoupled weight decay for training large language models, particularly when using matrix-based optimizers.

• Explicit Control: By explicitly fixing the weight norm and update norm, Hyperball removes the reliance on weight decay to indirectly regulate the angular learning rate. This leads to more predictable optimization dynamics.
• Scalability: The method successfully preserves the efficiency gains of matrix-based optimizers (like Muon) as model and data scales increase, addressing a critical bottleneck in current pretraining practices.
• Design Choice: The authors frame Hyperball as a design choice that turns the dependence on weight decay scaling (which affects transferability) into an explicit parameter (the radius $R$), thereby improving learning rate transfer across different model architectures.
• Proof of Concept: The work serves as a proof of concept that it is possible to remove weight decay entirely by designing optimizers that explicitly control weight norms, suggesting a new direction for optimizer design in normalized networks.

The authors remain modest regarding the specific constraint choice, noting that while the Frobenius norm is computationally efficient and theoretically motivated, other constraints (spectral, row-wise, column-wise) might be better suited for specific models or optimizers. They also highlight the need for sharper theories on weight-normalized training trajectories.