PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking Language Model Scaling under Transferable Hypersphere Optimization

Liliang Ren, Yang Liu, Yelong Shen, Weizhu Chen

arXiv:2603.28743v2cs.LG

TL;DR

Existing transfer laws largely focus on first-order optimizers and do not structurally guarantee stability at scale. HyperP derives learning-rate transfer laws for Frobenius-sphere optimization; a base learning rate tuned at small scale transfers to larger models, with reported compute-efficiency and stability results.

  • Problem

    Existing transfer frameworks mainly cover width and depth, leaving token scaling under second-order optimizers and structural stability at larger scales insufficiently addressed.

    Source: §1 · p. 1

  • Method

    HyperP transfers learning rates across width, depth, training tokens, and MoE granularity under Frobenius-sphere optimization.

    Source: §1 · p. 2, §6 · p. 17

  • Results

    A base learning rate tuned at d=8 transfers to d=24, achieving 1.58× compute efficiency leverage over a strong Muon baseline while monitored instability indicators remain bounded and non-increasing.

    Source: §1 · p. 2, §6 · p. 17

  • Takeaways & Limitations

    HyperP supports near-optimal comparisons of architectures across scales by transferring learning rates and producing scaling curves that reflect near-optimal performance.

    Source: §1 · p. 2

    The 0.32 data-scaling exponent is empirical and lacks a theoretical derivation for a universality guarantee.

    Source: §6 · p. 17

Abstract

from arXiv · show

Scaling laws for large language models depend critically on the optimizer and parameterization. Existing hyperparameter transfer laws are mainly developed for first-order optimizers, and they do not structurally prevent training instability at scale. Recent hypersphere optimization methods constrain weight matrices to a fixed-norm hypersphere, offering a promising alternative for more stable scaling. We introduce HyperP (Hypersphere Parameterization), the first framework for transferring optimal learning rates across model width, depth, training tokens, and Mixture-of-Experts (MoE) granularity under the Frobenius-sphere constraint with the Muon optimizer. We prove that weight decay is a first-order no-op on the Frobenius sphere, show that Depth-$\mu$P remains necessary, and find that the optimal learning rate follows the same data-scaling power law with the "magic exponent" 0.32 previously observed for AdamW. A single base learning rate tuned at the smallest scale transfers across all compute budgets under HyperP, yielding $1.58\times$ compute efficiency over a strong Muon baseline at $6\times10^{21}$ FLOPs. Moreover, HyperP delivers transferable stability: all monitored instability indicators, including $Z$-values, output RMS, and activation outliers, remain bounded and non-increasing under training FLOPs scaling. We also propose SqrtGate, an MoE gating mechanism derived from the hypersphere constraint that preserves output RMS across MoE granularities for improved granularity scaling, and show that hypersphere optimization enables substantially larger auxiliary load-balancing weights, yielding both strong performance and good expert balance. We release our training codebase at https://github.com/microsoft/ArchScale.

1 Introduction

Existing transfer frameworks largely focus on width and depth, while leaving token scaling under second-order optimizers and structural stability at scale insufficiently addressed. HyperP provides transfer laws and reports efficient, stable scaling across dense and MoE models.

1 Introduction

  • HyperP transfers optimal learning rates across width, depth, training tokens, and MoE granularity, with a 0.32 data-scaling exponent matching the reported AdamW result.Weight decay is a first-order no-op under Frobenius-sphere optimization, while Depth-µP remains necessary.Source: §1 · p. 2
  • Across MoE models from 913M to 13.3B total parameters, all six monitored instability indicators remain bounded and non-increasing as training FLOPs scale.The indicators include Z-values, output RMS, and activation-outlier percentages for attention and MoE sub-layers.Source: §1 · p. 2
  • SqrtGate preserves output RMS across MoE granularities and reduces router Z-value peaks by 5×; HyperP also supports larger load-balancing weights with strong validation loss and expert balance.The reported outcomes contrast with prior findings that large load-balancing weights hurt model quality.Source: §1 · p. 2
  • At 6 × 1021 FLOPs, HyperP achieves 1.58× CEL over a strong Muon baseline with µP++ and weight decay scaling; its S = 8, k = 8 MoE reaches 3.38× CEL over dense models.The advantage over the baseline grows monotonically with scale.Source: §1 · p. 2
  • Under HyperP, SqrtGate and Gated Attention show shrinking performance gains as training FLOPs increase, while providing stability gains that can remove RMS spikes and control exploding Z-values.The comparison covers dense and MoE architectures at their optimal performance.Source: §1 · p. 2

2 Background

Hypersphere optimization constrains weight matrices to a unit sphere under a chosen matrix norm. The background describes Frobenius- and spectral-norm variants, then details MuonH's normalized updates.

2 Background

  • Hypersphere optimization re-projects weight matrices onto a unit sphere after each gradient update, using a selected matrix norm.The update is described in terms of the learning rate, a constant, the weight update, and the chosen norm.Source: §2 · pp. 2–3
  • MuonH uses the Frobenius norm with Muon, while MuonSphere and SSO use the spectral norm; SSO additionally applies steepest descent on the spectral sphere.The paper focuses primarily on matrix-wise normalization for theoretical simplicity.Source: §2 · p. 3
  • MuonH normalizes both weights and updates, while AdamH is used for the language-model head and original AdamW for remaining vectors and embeddings.Fixed relative update magnitude per layer enables the same learning-rate scale for MuonH and AdamH.Source: §2 · p. 3

3 Scaling Hypersphere Optimization

This section derives HyperP transfer rules for Frobenius-sphere optimization, covering width, depth, training data, and MoE granularity. It also shows how the constraint removes weight decay to first order and bounds logits.

3 Scaling Hypersphere Optimization

  • After Frobenius-sphere projection, only the tangent component of an update survives to first order.The section derives HyperP transfer laws across width, depth, training duration, and MoE granularity.Source: §3.1 · p. 3, §3 · p. 3

3.1 Elimination of Weight Decay

  • Frobenius-sphere renormalization makes weight decay a first-order no-op, reducing the search from jointly tuning learning rate and weight decay to tuning learning rate alone.The paper studies MuonH and leaves a formal uniqueness characterization for future work.Source: §3.1 · p. 3

3.2 Width Scaling

  • With sufficiently uniform nonzero singular values, Frobenius-sphere optimization preserves width transfer without explicit 1/w learning-rate scaling.The width-transfer argument assumes approximate isotropy on the input space.Source: §3.2 · p. 4

3.3 Depth Scaling

  • With α_L=L^-1/2, weight-only renormalization permits η_l=O(1), but normalizing updates requires η_l=O(L^-1/2); Depth-μP remains necessary.The depth-transfer result uses small-step first-order analysis; the paper reports that its bound does not rely on independence assumptions across layers.Source: §3.3 · pp. 4–5

3.4 Data Scaling

  • η* follows a training-token power law with exponent 0.32, and leave-one-out optimal-learning-rate prediction error is 1.50%.Figure 1 plots loss versus learning rate across token budgets and fitted optimal learning rate versus training tokens.Source: Figure 1 · p. 5, §3.4 · p. 5

3.5 Hypersphere Parametrization

  • HyperP combines weight-decay elimination, native width transfer, depth scaling, and η ∝ T^-0.32 into a framework for training-FLOPs scaling.Source: §3.5 · p. 5

3.6 Bounded Logit Magnitudes

  • A fixed Frobenius norm bounds logit magnitudes for inputs with O(1) RMS, naturally bounding attention and MoE-routing logits and alleviating the need for Z-loss regularization.Source: §3.6 · p. 6

3.7 MoE Granularity Scaling

  • SqrtGate preserves routed-branch RMS approximately across granularity under equal-RMS, weak-correlation conditions, whereas classical gating preserves RMS only when routing nearly collapses to Top-1.The analysis treats the equal-RMS, weak-correlation regime as natural under hypersphere optimization.Source: §3.7 · p. 7

4 Experiments & Results

The experiments evaluate MuonH and HyperP across depth, batch size, MoE sparsity and granularity, and training-FLOPs scales. The scaling comparisons use fitted loss–FLOPs trajectories and compute efficiency leverage.

4.1 Empirical Optimality of MuonH

  • MuonH achieves slightly better validation loss than Muon while eliminating weight decay as a hyperparameter; its optimal LR is approximately 1.4× smaller.For Muon, the best λ=10^-3 gives loss 2.479, while λ=4×10^-3 gives 2.500.Source: §4.1 · p. 8, Table 2 · p. 8

4.2 Parameter Scaling

  • With Depth-µP, the optimal LR stays near 0.014–0.016 from d=8 to d=24; without it, the optimum falls from 0.016 at d=8 to 0.008 at d=24.Both configurations achieve comparable best losses at each depth.Source: §4.2 · p. 9, Figure 3 · p. 9

4.3 Critical Batch Size

  • The optimal LR scales with batch size as η*=4.66 × 10^-6 · B^0.558, while minimum loss varies by at most 0.004 across tested batch sizes.The tested batch sizes are 256K, 512K, 1M, and 2M tokens; the authors conclude they are below the critical batch size for this configuration.Source: §4.3 · p. 9

4.4 MoE Scaling

  • For S=8, k=4, the largest auxiliary-loss weight, γ=10^-1, achieves the best loss (2.332) and the lowest Mean MaxVio.The experiments use γ=0.1 in the remainder of the paper.Source: §4.4 · p. 10
  • Across a 32× sparsity range, optimal LR varies only from 0.012 to 0.016, while increasing sparsity from S=1 to S=32 reduces optimal loss by 0.224.The authors estimate matching that loss reduction with dense training would require approximately 5.2× more FLOPs under the MuonH+HyperP scaling curve.Source: §4.4 · p. 11, Figure 6 · p. 11
  • Across a 32× granularity range, optimal LR varies only from 0.012 to 0.014; with SqrtGate, performance improves through k=32, where loss is 2.310.SqrtGate improves validation loss at every tested k, while the baseline saturates at k=16.Source: §4.4 · p. 11, Figure 7 · p. 11

4.5 Training FLOPs Scaling

  • Under HyperP, loss–LR curves remain aligned from d=8 through d=20, preserving the same base optimum, η0=0.02, across scales.Without HyperP, the optimal LR drifts with depth, so reusing the small-scale LR becomes increasingly miscalibrated.Source: §4.5 · p. 12, Figure 8 · p. 12
  • At 5.96 × 10^21 FLOPs, MuonH+HyperP achieves 1.58× CEL over the Muon baseline, while its MoE variant reaches 3.38× CEL over dense baselines.The dense-model CEL rises from near parity at d=8 to 1.58× at d=24; MuonH without HyperP declines to 0.70×.Source: Figure 9 · p. 12, §4.5 · p. 13

5 Analysis

The analysis examines stability transfer, the sensitivity of fitted optimal learning rates, and architectural choices under HyperP. It reports how training indicators and architecture variants behave as scale increases.

5.1 Transferable Stability

  • Across MoE depths d=8–20, HyperP’s six monitored instability metrics remain well-behaved; router Z-value peaks decline from 56 at d=8 to 33 at d=20.Attention Z-values plateau around 200–220, while output RMS and outlier percentages decrease with depth.Source: §5.1 · pp. 13–14, Figure 10 · p. 13

5.2 Sensitivity of Optimal Learning Rate Estimation

  • With only five LR sweep points, the worst-case LR error is 4.1%, while the corresponding loss error is 0.04%, about 0.001 nats.The conclusion assumes a well-fitting quadratic with R^2>0.99; when the fit is poor, the authors rely on the observed optimal LR.Source: §5.2 · p. 14, Figure 11 · p. 14

5.3 Rethinking Architecture

  • At d=8, GA QK-Norm beats QK-Norm by 0.010 nats and Baseline by 0.023 nats; MoE variants have nearly identical optimal LRs of 0.0135–0.0137.SharedExp and SqrtGate alone provide nearly identical improvements, while combining them yields the best performance.Source: §5.3 · p. 15, Figure 12 · p. 15
  • GatedAttn+QKNorm has the strongest dense scaling behavior, with compute-efficiency leverage peaking near 1.15× at intermediate scale, but architectural performance advantages shrink as training scale increases.Its small-scale advantage over Baseline is 0.023 nats.Source: Figure 13 · p. 15, §5.3 · p. 16
  • At fitted LR=0.014, SharedExp+SqrtGate has the best MoE scaling; all properly tuned variants outperform the SharedExp+SqrtGate baseline at LR=0.01.The comparison spans depths d=8, 12, 16, and 24.Source: §5.3 · p. 16, Figure 14 · p. 16
  • Without SqrtGate, router Z-values rise from about 25 to over 190 with frequent spikes; adding it keeps them below 40 throughout training.For dense variants, GatedAttn+QKNorm maintains the most stable MLP output RMS trend without spikes.Source: §5.3 · pp. 16–17, Figure 15 · p. 16

6 Conclusion

HyperP transfers a single learning rate across model scales and MoE granularities, with reported compute-efficiency and stability benefits. The authors also identify assumptions and open questions that bound these findings.

6 Conclusion

  • 1.58× compute efficiency over a strong Muon baseline at 6 × 10^21 FLOPs is reported for a learning rate transferred from d=8 to d=24.The conclusion also reports 3.38× over dense models for MoE models at d=20, with 13.3B total parameters.Source: §6 · p. 17
  • HyperP transfers one optimal learning rate across width, depth, training tokens, and MoE granularity under Frobenius-sphere optimization.The conclusion also reports that weight decay is a first-order no-op, Depth-µP remains necessary, and the data-scaling exponent 0.32 matches previous AdamW studies.Source: §6 · p. 17
  • All monitored instability indicators are non-increasing with scale under transferred hyperparameters, while stability benefits of QK-Norm, Gated Attention, and SqrtGate become more important for long-horizon training.The reported loss improvements from those architecture changes diminish with scale.Source: §6 · p. 17
  • The 0.32 data-scaling exponent lacks a theoretical derivation, and the assumed Chinchilla compute-optimality law needs refitting for each training dataset.The authors also identify extension to other architectures and larger scales, the batch-size exponent, and optimizer momentum terms as future directions.Source: §6 · p. 17

A Proof of Width Transfer under Frobenius Sphere

The proof gives conditions for attaining the Frobenius-sphere width-transfer bound and explains how renormalization acts on updates to first order.

A Proof of Width Transfer under Frobenius Sphere

B First-Order Form of Frobenius-Sphere Updates

C Depth Scaling under Frobenius-Sphere Optimization

The depth-scaling analysis proceeds from network perturbations to residual-network cases and then post-norm blocks.

C Depth Scaling under Frobenius-Sphere Optimization

C.1 First-order decomposition of the network perturbation

The subsection models changes in network output from layerwise parameter perturbations using a first-order expansion.

C.1 First-order decomposition of the network perturbation

  • F(x; W1, . . . , WL) denotes the network output as a function of all layer parameters.Source: §C.1 · p. 22
  • A first-order multivariate Taylor expansion represents the output change for perturbations to the layer weights.Source: §C.1 · p. 22
  • The first-order total perturbation is additive over layers.Source: §C.1 · p. 22

C.2 Residual network without update normalization

For residual networks without update normalization, the depth dependence of the weight update determines how the learning rate must scale to preserve an O(1) first-order function perturbation. Under the stated stable-depth assumptions, the result differs between residual scalers α_L=L^-1/2 and α_L=L^-1.

C.2 Residual network without update normalization

  • With bounded downstream transport factors, the depth exponent of the first-order perturbation can be tracked through the layer contribution b_l.The passage states that each A_l has operator norm O(1) for sufficiently small α_L.Source: §C.2 · p. 23
  • Under the stable-depth assumption, the layerwise gradient scales as ∥G_l∥F=O(α_L), so the weight update scales as ∥∆W_l∥F=O(η_lα_L).Differentiating the residual block with respect to W_l introduces one factor of α_L; the upstream signal is O(1) by assumption.Source: §C.2 · p. 23
  • With α_L=L^-1/2, a depth-independent learning rate η_l=O(1) yields an O(1) first-order function perturbation; with α_L=L^-1, η_l=O(L) is required.These results assume α_L is sufficiently small to satisfy the stated assumptions.Source: §C.2 · p. 23

C.3 Residual network with update normalization

With update normalization, each weight update has size O(η_l), and the residual contribution is analyzed by summing across depth. The post-norm block has the same depth scaler as the pre-norm block under the stated conditions.

C.3 Residual network with update normalization

  • With update normalization, ∥∆W_l∥F=O(η_l), so the residual contribution at each layer is controlled by the learning rate.The raw update is normalized before projection to the Frobenius sphere.Source: §C.3 · p. 23
  • For α_L=L^-1/2, the stated learning-rate scaling is η_l=O(L^-1/2); for α_L=L^-1, it is η_l=O(1).These scalings are given for preserving an O(1) first-order function perturbation after summing over depth.Source: §C.3 · p. 23

C.4 Post-norm residual block

  • Under post-norm, the LayerNorm Jacobian is O(1), leaving the local weight sensitivity with the same depth scaler α_L as in a pre-norm residual block.The resulting perturbation scales as ∥∆x_L∥=O(Lα_Lη_l).Source: §C.4 · p. 24

D Detailed Experimental Results

This section lists numerical results across learning-rate transfer, depth and batch-size scaling, MoE settings, and compute efficiency. Its table captions report several qualitative comparisons, including transfer with Depth-µP and increasing advantages for MuonH + HyperP at larger scale.

D Detailed Experimental Results

  • Depth-µP transfers the optimal learning rate across parameter size in the fixed-token depth comparison.Tables 5 and 6 compare validation loss versus learning rate without and with Depth-µP; Table 7 compares optimal learning rate and loss.Source: Table 5 · p. 24, Table 6 · p. 24, Table 7 · p. 25
  • Across tested dense-model batch sizes, minimum loss stays within 0.004 nats, which the caption interprets as all tested sizes being below the critical batch size.Source: Table 8 · p. 25
  • Across a 100× range of auxiliary-loss weight γ, the optimal learning rate remains η*=0.012 and achievable loss is stable.Source: Table 9 · p. 25
  • Across a 32× range of MoE sparsity, the optimal learning rate varies mildly from 0.012 to 0.016.Source: Table 10 · p. 25
  • The top-k comparison reports optimal learning rates and losses with and without SqrtGate, but gives no outcome in the caption.Source: Table 11 · p. 25
  • MuonH + HyperP’s validation-loss and compute-efficiency advantages over alternatives increase with scale, and its compute-efficiency advantage over Muon grows monotonically.The captions identify these comparisons in Tables 12 and 13 without giving the underlying values here.Source: Table 12 · p. 26, Table 13 · p. 26
  • At d=8, GA QK-Norm achieves the best loss while maintaining a similar optimal learning rate.Source: Table 15 · p. 26
  • At d=8, S=8, k=8, and 10.4B tokens, SqrtGate and the shared expert provide complementary gains.Source: Table 16 · p. 26

Related papers

Loading 2603.28743v2…