# Rethinking Language Model Scaling under Transferable Hypersphere Optimization

> HyperP transfers a single learning rate across model width, depth, tokens, and MoE granularity, achieving 1.58x compute efficiency over Muon baselines.

- **Source:** [arXiv](https://arxiv.org/abs/2603.28743)
- **Published:** 2026-08-17
- **Permalink:** https://picx.dev/p/FyDUWn
- **Whiteboard:** https://picx.dev/p/FyDUWn/image

## Summary

## 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

| Method | Fitted η* | Best Val Loss | Weight Decay |
|--------|-----------|---------------|--------------|
| Muon (best λ=10⁻³) | 0.0222 | 2.479 | 10⁻³ |
| MuonH (λ=0) | 0.0155 | 2.475 | 0 |

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.

| γ | Best Val Loss | Mean MaxVio |
|---|---------------|-------------|
| 10⁻³ | 2.334 | 0.848 |
| 10⁻² | 2.336 | 0.132 |
| 10⁻¹ | 2.332 | 0.086 |

**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

---

_Markdown view of https://picx.dev/p/FyDUWn, served by PicX — AI-generated visual whiteboard summaries of research papers._