Summary (Overview)

• HyperP Framework: Introduces the first framework for transferring optimal learning rates across model width, depth, training tokens, and MoE granularity under Frobenius-sphere optimization with the MuonH optimizer.
• Key Theoretical Results: Proves that weight decay is a first-order no-op on the Frobenius sphere, demonstrates that Depth-µP remains necessary (refuting prior claims of inherent depth transferability), and discovers a "magic exponent" of 0.32 for data scaling that matches previous AdamW results.
• Compute Efficiency: A single base learning rate tuned at the smallest scale (208M parameters) transfers across all compute budgets, achieving 1.58× Compute Efficiency Leverage (CEL) over a strong Muon baseline at $6 \times 10^{21}$ FLOPs, with MoE models reaching 3.38× CEL over dense baselines.
• Transferable Stability: All six monitored instability indicators (Z-values, output RMS, activation outliers) remain bounded and non-increasing as training FLOPs scale from 913M to 13.3B total parameters.
• SqrtGate Mechanism: A novel square-root gating mechanism for MoE that preserves output RMS across granularities, reducing router Z-value peaks by 5× compared to standard gating.

Introduction and Theoretical Foundation

Background and Motivation

Neural scaling laws are central to compute-efficient LLM development, but identifying true scaling behavior requires each model to be trained with near-optimal hyperparameters. Existing hyperparameter transfer frameworks primarily study first-order optimizers and lack structural guarantees on training stability at scale.

Key challenges addressed:

• Training instability: Scaling up training FLOPs routinely triggers logit explosion, activation outliers, and loss spikes
• Hyperparameter complexity: Weight decay's optimal value depends intricately on learning rate, training duration, and model width
• MoE scaling: Classical gating mechanisms cause output RMS to shrink with granularity

Theoretical Foundation: Hypersphere Optimization

Hypersphere optimization constrains weight matrices to lie on a unit sphere under a chosen matrix norm:

$$W \leftarrow C \frac{W - \eta G}{\| W - \eta G\|}, \tag{1}$$

MuonH instantiates this with the Frobenius norm and Muon optimizer:

$$\widehat{G} = c_{G} \frac{G}{\|G\|_F}, \qquad W^{+} = c_{W} \frac{W - \eta \widehat{G}}{\|W - \eta \widehat{G}\|_F}, \tag{2}$$

where $c_W = \|W_0\|_F$ is the initial weight norm and $c_G = c_W$.

Key Theoretical Results

Theorem 1 (Weight decay is a first-order no-op): Under Frobenius renormalization, if $\Delta = -\eta G - \eta \lambda W$, then:

$$W^{+} - W = -\eta \Pi_T(G) + \mathcal{O}(\eta^2), \tag{4}$$

where $\Pi_T(\Delta) = \Delta - \frac{\langle \Delta, W\rangle_F}{\|W\|_F^2} W$ is the tangent-space projection. The weight decay term has no first-order effect.

Theorem 2 (Width transfer): With $\|W\|_F = C\sqrt{d_{out}}$, hypersphere optimization preserves width transfer without explicit $1/w$ learning rate scaling as in standard µP.

Theorem 3 (Depth scaling): For residual networks with normalized updates:

$$\eta_l = O\left(\frac{1}{L\alpha_L}\right) \tag{6}$$

With standard depth-stabilizing scaling $\alpha_L = L^{-1/2}$, this gives $\eta_l = O(L^{-1/2})$.

Methodology

Data Scaling Law

Through empirical studies varying training tokens from 10.4B to 166.4B with fine-grained LR sweeps, the optimal learning rate follows:

$$\eta^{*} = 24.27 \cdot T^{-0.320} \tag{8}$$

where $T$ is the total training tokens. Leave-one-out cross-validation gives a mean absolute prediction error of only 1.50% for optimal LR.

HyperP Parameterization Summary

The framework applies:

• Weight decay: Eliminated entirely (set to 0)
• Width scaling: Native transfer via Frobenius-sphere constraint
• Depth scaling: $\eta \propto 1/\sqrt{d}$ with residual multiplier $1/\sqrt{2d}$
• Data scaling: $\eta \propto T^{-0.32}$

SqrtGate for MoE

Proposition 5 (Classical gating is k-dependent): With classical softmax gating:

$$\|y_{route}\|_{rms} \approx \frac{r}{\sqrt{k}} \tag{13}$$

for near-uniform routing weights, showing output RMS shrinks with granularity.

Proposition 6 (SqrtGate is approximately k-invariant): By replacing $g_i$ with $\sqrt{g_i}$:

$$\|y'_{route}\|_{rms} \approx r \sqrt{\sum_{i=1}^{k}(\sqrt{g_i})^2} = r \tag{14}$$

The routed-branch RMS is approximately invariant to granularity $k$.

Training Setup

• Architecture: Transformer-Next with GQA (4 KV heads), head dimension 128, aspect ratio $\alpha = 128$ (width $w = 128d$), QK-Norm, headwise gated attention
• Data: SlimPajama dataset, context length 4K, batch size 2M tokens
• Schedule: Linear decay to 10% of peak without warm-up, momentum 0.95
• Scaling: Chinchilla law with Tokens Per Parameter (TPP) $= T/N$

Empirical Validation / Results

MuonH vs. Muon Optimality

MuonH achieves slightly better validation loss while entirely removing weight decay as a hyperparameter.

Depth Scaling Verification

Without Depth-µP, optimal LR decreases from $\eta^* = 0.016$ at d=8 to $\eta^* = 0.008$ at d=24. With Depth-µP, the optimal LR remains nearly constant at $\eta^* \approx 0.014$–$0.016$ from d=8 to d=24, validating the theoretical predictions.

Critical Batch Size

The optimal LR scales as $\eta^* = 4.66 \times 10^{-6} \cdot B^{0.558}$, with exponent ≈ 0.56 sitting between linear scaling (1.0) and square-root rule (0.5).

MoE Scaling Results

Auxiliary Balance Loss: The largest weight $\gamma = 10^{-1}$ achieves the best loss (2.332) with the lowest Mean MaxVio (0.086), contrasting with prior work suggesting auxiliary loss harms quality.

Sparsity Scaling: Optimal LR varies only mildly (0.012–0.016) across a 32× sparsity range, with increasing sparsity improving validation loss by 0.224 from S=1 to S=32.

Granularity Scaling: With SqrtGate, performance improves up to k=32 (best loss 2.310), whereas baseline saturates at k=16. SqrtGate consistently improves val loss at every k.

Compute Scaling Comparisons

At the largest budget ($5.96 \times 10^{21}$ FLOPs):

• MuonH+HyperP: 1.58× CEL over Muon baseline, lowest irreducible floor ($C_0 = 0.85$)
• MuonH+HyperP MoE: 3.38× CEL over dense baselines, comparable floor ($C_0 = 0.87$)
• MuonH without HyperP: Declines to 0.70× CEL, showing LR transfer mismatch compounds into substantial penalty

Transferable Stability

All six monitored instability indicators (attention Z-values, router Z-values, attention output RMS, MoE output RMS, attention outlier %, MoE outlier %) are bounded and non-increasing with scale. Router Z-value peaks decrease monotonically with depth (from 56 at d=8 to 33 at d=20).

LR Estimation Sensitivity

• Optimal loss is 50–140× less sensitive than optimal LR estimation
• With n=5 points, worst-case LR error is 4.1%, but corresponding loss error is only 0.04% (~0.001 nats)
• Five points suffice for reliable estimates

Theoretical and Practical Implications

Theoretical Contributions

1. Weight decay elimination: Proves that Frobenius-sphere optimization makes weight decay a first-order no-op, reducing hyperparameter search from the joint $(\eta, \lambda)$ plane to a single dimension $\eta$.

2. Universal scaling exponent: The "magic exponent" 0.32 for data scaling matches previous AdamW results, suggesting universality across optimizers.

3. Depth transfer correction: Refutes the claim that MuonH is inherently depth-transferable, showing Depth-µP remains necessary due to cumulative angular drift from residual connections.

4. Bounded logits: For any weight matrix with $\|W\|_F = C$ and input with $\|x\|_{rms} = O(1)$:

$$\|Wx\|_{rms} \leq C\sqrt{\frac{d_{in}}{d_{out}}}\|x\|_{rms} \tag{10}$$

This provides structural stability guarantees that standard training lacks.

Practical Implications

1. Single-scale tuning: One LR sweep at small scale (d=8, 208M parameters) determines learning rates for the full scaling trajectory up to 3.8B active parameters.

2. MoE load balancing: Hypersphere optimization enables substantially larger auxiliary load-balancing weights, achieving both strong performance and good expert balance simultaneously.

3. Architecture stability: While loss improvements from QK-Norm, Gated Attention, and SqrtGate diminish with scale, their stability benefits (preventing RMS spikes, controlling Z-values) become increasingly important for long-horizon training.

4. Compute efficiency: The advantage of HyperP over baselines grows monotonically with scale, implying even larger gains at frontier compute budgets.

Conclusion

HyperP is the first framework for transferring a single optimal learning rate across model width, depth, training tokens, and MoE granularity under Frobenius-sphere optimization. Key achievements include:

• Theoretical advances: Weight decay elimination proof, depth scaling derivation, and the discovery of a universal data-scaling exponent (0.32)
• Practical gains: 1.58× compute efficiency over Muon baseline at $6 \times 10^{21}$ FLOPs, with MoE models reaching 3.38× CEL
• Stability guarantees: All monitored instability indicators are non-increasing with scale under transferred hyperparameters
• Novel MoE mechanism: SqrtGate preserves output RMS across granularities with 5× reduction in router Z-value peaks

Future Directions

• Theoretical derivation of the magic data scaling exponent 0.32 for universality guarantees
• Extension to other architectures (hybrid models, linear recurrent models)
• Batch size scaling: The exponent 0.56 deviates from SDE-predicted 0.5, warranting further investigation
• Optimizer momentum analysis: Extending transfer laws to momentum terms for better scaling performance
• Re-fitting Chinchilla law per training dataset, as the current assumption may not be universally optimal