PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimal Scaling Needs Optimal Norm

Oleg Filatov, Jiangtao Wang, Jan Ebert, Stefan Kesselheim

arXiv:2510.03871v2cs.LGcs.AIstat.ML

TL;DR

Although norm guidance is established, how to choose optimal learning-rate and batch-size combinations for a horizon, and what suffices for optimality, remains unclear. The paper studies output-layer norms and hyperparameter scaling, finding that the output norm guides joint optimal scaling and that dataset-dependent rules yield optimal loss.

  • Problem

    Despite norm guidance, the optimal learning-rate and batch-size combination for a given horizon, and the sufficient condition for optimality, remain unclear.

    Source: §3.3 · p. 7

  • Method

    The authors profile training loss across learning rates for each batch size, fit loss against output norm, and measure the resulting optimum across dataset horizons and model scales.

    Source: §2.3 · p. 4

  • Results

    The output-layer operator norm guides joint optimal scaling across model and dataset dimensions, and the derived hyperparameter scaling rules yield optimal loss.

    Source: §5 · p. 10

  • Takeaways & Limitations

    The output norm is a necessary condition for transfer, while optimal dataset-dependent learning-rate and batch-size scaling rules yield optimal loss.

    Source: §5 · p. 10

    The authors do not yet know why the scaling rules take their observed form or how they connect to the constant-norm scaling condition.

    Source: §5 · p. 10

Abstract

from arXiv · show

Despite recent progress in optimal hyperparameter transfer under model and dataset scaling, no unifying explanatory principle has been established. For Adam and Scion optimizers, we discover that joint optimal scaling across model and dataset sizes is conditioned on a single invariant: the operator norm of the output layer. Across models with up to 1.3B parameters trained on up to 138B tokens, the optimal learning rate/batch size pair $(\eta^{\ast}, B^{\ast})$ consistently has the same operator norm value - a phenomenon we term norm transfer. This constant norm condition is necessary but not sufficient: while for each dataset size, multiple $(\eta, B)$ reach the optimal norm, only a unique $(\eta^{\ast}, B^{\ast})$ achieves the best loss. As a sufficient condition, we provide the first measurement of $(\eta^{\ast}, B^{\ast})$ scaling with dataset size for Scion, and find that the scaling rules are consistent with those of Adam. Tuning per-layer-group learning rates also improves model performance, with the output layer being the most sensitive and hidden layers benefiting from lower learning rates. We provide practical insights on norm-guided optimal scaling and release our Distributed Scion (Disco) implementation with logs from over two thousand runs to support research on LLM training dynamics at scale.

1 INTRODUCTION

Optimal scaling across model and dataset sizes lacks a unifying principle. The paper identifies output-layer norm transfer and measures optimizer and layer-group learning-rate rules.

Unifying invariant for optimal scaling

  • The output-layer operator norm at optimal (η, B) transfers across model and dataset scaling for both Scion and Adam, providing a necessary but insufficient condition for optimality.Multiple non-optimal (η, B) pairs can reach the same norm value.Source: §1 · pp. 1–2

Scaling rules for Scion

  • For Scion, η*(B, D) ∝ B^0.62 · D^-0.56, B*(D) ∝ D^0.45±0.07, and η*(D) ∝ D^-0.28±0.07, matching Adam’s square-root scaling rules.At fixed D, η and B can trade off within a low-sensitivity region around the optimal norm, allowing larger batches with little performance sensitivity.Source: §1 · p. 2

Optimal per-layer-group learning rate

  • Per-layer-group tuning improves relative loss by up to 6%, with ηinput : ηhidden : ηoutput = 1 : 1/8 : 1 consistently optimal across dataset and batch sizes.The output layer is most sensitive to tuning, followed by hidden layers and then the input layer.Source: §1 · p. 2

Distributed Scion/Muon and experimental logs

  • The authors release Disco, a distributed Scion/Muon implementation, and norm logs from over two thousand training runs to support research on large-scale training dynamics.Source: §1 · p. 2

2 METHODOLOGY

The methodology combines norm-based optimization concepts with controlled language-model training and grid searches over learning rate and batch size. It estimates optimal output norms across dataset horizons and model scales.

2.1 BACKGROUND & TERMINOLOGY

  • µP scales hyperparameters with model width so differently sized models share optimal hyperparameters, enabling zero-shot transfer; Scion assigns layer-specific operator norms whose duality maps transform raw gradients before weight updates.Source: §2.1 · pp. 2–3

2.2 TRAINING SETUP

  • Experiments use Llama 3 with a 69M-parameter proxy, scale models to 1.3B parameters in width or 168M in depth, and train on Nemotron-CC with a 4096-token context window.A norm-everywhere setup normalizes each linear layer’s input to RMS norm 1; runs generally use a constant learning rate without warmup or decay.Source: §2.2 · p. 4

2.3 OPTIMAL NORM MEASUREMENT

  • For each model and dataset horizon, the authors grid-search (η, B) and use training loss to identify the optimal Scion-assigned output-layer norm, then repeat across horizons and model width and depth.They focus on the RMS→∞ output norm and treat training loss in the non-repeating infinite-data regime as reflecting performance and generalization.Source: §2.3 · p. 4
  • Because the lowest observed loss across learning rates is statistically noisy, each batch size’s loss-versus-output-norm distribution is fitted to estimate the optimal norm and corresponding learning rate.The learning rate is selected from the data point nearest the fitted optimum.Source: §2.3 · p. 4

3 RESULTS

The output-layer norm evolves with learning rate and batch size, and its optimal value transfers across dataset horizons and model scales. Norm transfer alone is not sufficient to select the best learning-rate/batch-size pair; measured scaling rules and per-layer-group tuning further characterize optimal training.

3.1 OUTPUT NORM DYNAMICS

  • The output norm grows with training and its growth is piecewise linear in log-log scale, with slope changes near norms 2^6–2^7 and 2^9–2^10; learning rate shifts the curves while batch size affects their separation.After the curves decouple, they converge to a shared slope lower than the initial one.Source: §3.1 · p. 5

3.2 OPTIMAL NORM TRANSFER

  • The optimal output norm transfers across token horizons and model scaling: it is about 2^7 for data scaling and remains shared when width increases ×12 or depth increases ×32.The width-scaled models also show lower loss as width grows; depth transfer is observed without established depth-transfer techniques, while the authors speculate it may relate to their norm-everywhere approach.Source: §3.2 · pp. 5–6
  • Norm transfer also appears with Adam and Scion variants, but the norm condition is necessary rather than sufficient: many learning-rate/batch-size pairs can reach it, while loss distinguishes the optimum.Momentum and learning-rate decay reduce sensitivity to batch-size and learning-rate choices, respectively; removing the output normalization layer lowers the optimal norm by about a factor of 10.Source: §3.2 · pp. 6–7, §3.3 · p. 8

3.3 OPTIMAL (η, B) SCALING RULE

  • For Scion, the fitted optimal learning-rate rule is η*(B,D) ∝ B^0.62 · D^-0.56, consistent with Adam’s square-root scaling exponents; the optimal batch size grows as D^0.45±0.07.For each horizon, a low-norm-sensitivity region permits trading learning rate against batch size, including larger batches, without loss in performance; a unique optimal batch size per horizon also bounds device scaling.Source: §3.3 · p. 8

3.4 OPTIMAL PER-LAYER-GROUP LEARNING RATE

  • Per-layer-group tuning finds the output layer most sensitive to learning-rate mistuning, while uniform rates perform similarly to the best configurations within uncertainty.With input and output rates tied, all horizons favor a V-shaped layout with the hidden-layer rate lower by ×1/8; the paper reports up to 6% relative-loss improvement and transfer across dataset sizes.Source: §3.4 · p. 9

4 RELATED WORK

Prior work developed methods for transferring hyperparameters across model sizes and measuring dataset-size scaling, but these strands remain distinct. Research on learning-rate and batch-size scaling has also approached the problem through several theoretical frameworks.

Hyperparameters with model scaling

  • µP and related work provide principled hyperparameter transfer across model width and depth, while empirical laws describe scaling with compute, loss, or model size.The section also notes extensions to other parameterizations and empirical scaling rules.Source: §4 · p. 10

Hyperparameters with data scaling

  • Dataset-size hyperparameter scaling remains poorly understood theoretically and has largely been studied by experimentally measuring scaling rules.Smith and Le (2018) modeled SGD optimization trajectories with a stochastic differential equation; later work measured dataset-size scaling experimentally.Source: §4 · p. 10

(η, B) scaling rules

Norm-based optimization

  • Norm-based optimization developed from spectral conditions and gradient-update transformations, leading to optimizers including Muon and Scion and subsequent improvements and benchmarks.Related work also explores constraining weights, manifold-learning connections, and optimizer design.Source: §4 · p. 10

5 CONCLUSION AND DISCUSSION

The paper presents output-layer norm as a guide to joint optimal scaling, with a progression from hyperparameter choices to norms, optimal loss, and dataset-dependent scaling rules. It also identifies unresolved questions about why the rules arise, why norm transfer holds, and which norm is optimal.

5 CONCLUSION AND DISCUSSION

  • The paper frames joint optimal scaling as a progression from hyperparameter choices affecting layer norms to optimal loss requiring an optimal norm and scaling rules yielding optimal loss.The authors describe empirical study of norm changes and tuning, a predefined output-layer norm for transfer across data and model scales, and hyperparameter scaling rules for optimal loss.Source: §5 · p. 10
  • The origin of the observed scaling rules remains unexplained, despite the authors' confidence that they hold at larger scales.They note that the rules resemble square-root and 1/4-power laws.Source: §5 · p. 10
  • Why the optimal scaling trajectory stays on a constant-norm manifold, and what defines that manifold's structure, remain open questions.Source: §5 · p. 10
  • The paper leaves open which norm is exactly optimal, having focused primarily on the output-layer norm.Source: §5 · p. 10
  • The authors suggest the constant-norm condition may be leveraged as an inductive bias for optimizing training, but report that these questions remain unanswered.Source: §5 · p. 10

A.1 LLM USAGE

LLMs assisted with prose polishing and clarity, and code-assistant tools provided occasional minor programming support.

A.1 LLM USAGE

  • LLMs were used to polish writing and improve exposition clarity, while code assistants occasionally supported minor programming tasks such as completion and syntax suggestions.They were not used to design algorithms, generate experiments, or implement the proposed methods from scratch.Source: §A.1 · p. 18

A.2 MODEL TRAINING CONFIGURATION

The experiments use a 69M-parameter proxy model and wider or deeper variants, with specified architectural, initialization, normalization, and implementation settings. Random seeds vary across proxy runs, scaled-up runs, and layout scans.

A.2 MODEL TRAINING CONFIGURATION

  • The 69M-parameter proxy has four 256-dimensional hidden layers; width-scaled models reach 314M and 1.3B parameters, while depth-scaled models reach 91M and 168M.The width-scaled variants use d_model values of 1024 and 3072; the depth-scaled variants use 32 and 128 layers.Source: §A.2 · p. 18
  • The models use semi-orthogonal hidden-layer initialization, row-wise normalized Gaussian input/output embeddings, disabled dropout, untied embeddings, and norm-everywhere RMSNorm without learnable parameters.The norm-everywhere setup adds V-, O-, and MLP-normalization to the Pre-LN setup with QK-norm; residual connections remain intact.Source: §A.2 · p. 18
  • Proxy runs use 30 seeds, scaled-up runs interleave seeds 30 and 3034, and layout scans use additional seed averages; implementation uses torchtitan, FSDP2, and FlashAttention-2.Source: §A.2 · p. 18

A.3 OPTIMIZER CONFIGURATION

The optimizer configuration uses a common set of hyperparameters except in dedicated ablations, including an unconstrained version and fixed learning-rate and schedule choices. Scion orthogonalizes hidden-layer gradients with Newton-Schulz iterations.

A.3 OPTIMIZER CONFIGURATION

  • Except for dedicated ablations, the optimizer is unconstrained without weight decay, uses momentum μ = 0, no warmup, a constant learning-rate schedule, and ε = 1e−20 for gradient normalization.Learning rates are scanned on a 2^0.5 step for the proxy model and a 2^1 step for width/depth-scaled models.Source: §A.3 · p. 18
  • Scion orthogonalizes hidden-layer gradients under an RMS-to-RMS norm assumption using five Newton-Schulz iterations and the stated Muon coefficients.Source: §A.3 · p. 19

A.4 OPTIMAL NORM FITTING & LOSS SMOOTHING

The authors estimate optimal loss and norm by fitting loss-versus-output-norm scans, then smooth selected loss measurements to improve robustness. They vary the fitting procedure to estimate uncertainty in downstream scaling results.

A.4 OPTIMAL NORM FITTING & LOSS SMOOTHING

  • A constrained quadratic fit in log-log space estimates optimal loss and norm from seven points around each empirical optimum.The fit fixes its intercept to the initialization loss; the optimal learning rate is taken from the run nearest the fitted optimum.Source: §A.4 · p. 19
  • For small batches and long horizons, loss estimates average each point with its neighboring evaluations to reduce variance and avoid bias.Smoothing is used for B ≤128 and D ≥233, where loss variance is high and training is largely converged.Source: §A.4 · p. 19
  • Six fitting and smoothing variants track changes in optimal norm, loss, and learning rate to estimate uncertainty without multiple random seeds.The variants differ in fitting, smoothing, and whether the initialization-loss constraint is applied.Source: §A.4 · p. 19

A.5 DISTRIBUTED SCION

The section introduces a distributed Scion/Muon implementation and assumes momentum-buffer updates occur before weight updates.

A.5 DISTRIBUTED SCION

  • Disco performs the vectorized momentum-buffer update before applying the corresponding weight update.The section introduces a distributed version of Scion/Muon.Source: §A.5 · p. 20

A.5.1 DDP-DISCO

The appendix details Disco’s distributed update paths and reports supplementary scaling and optimizer ablations. Results include a Scion batch-size scaling close to square-root growth, alongside several conditions under which norm transfer persists or becomes less sensitive.

A.5.1 DDP-DISCO

  • The DDP implementation supports PP without extra changes, while TP requires an additional all-gather in the local update loop.The implementation does not explicitly overlap communication and computation for DDP.Source: §A.5.1 · p. 20

A.5.2 FSDP-DISCO

  • For MoE experts, Disco transposes weights when needed so output dimensions come first, and batched SVD or Newton–Schulz methods must handle the resulting 3D gradients.The FSDP setup assumes one tensor or matrix per layer; fused QKV and W13 tensors are unsupported.Source: §A.5.2 · p. 21

A.8 OPTIMAL B∗(D) MEASUREMENT

  • B∗ scales as D^0.45±0.07 for Scion, consistent with square-root scaling, based on an OLS fit across nine horizons.The extended horizon set is used because four horizons alone would not precisely establish the scaling.Source: §A.8 · p. 25, Figure 7 · p. 25

A.10 ABLATION OF DEPTH TRANSFER TECHNIQUES

  • Across depths, Scion’s optimal learning rate remains near 2^-2 regardless of initialization or residual scaling, with negligible improvement from total-depth initialization plus identity residual scaling.The ablation compares proxy models with 2, 16, and 64 layers.Source: §A.10 · p. 26

A.11 ABLATION WITH FINEWEB-2 DATASET

  • FineWeb-2’s Russian partition shares an optimal norm across nested horizon curves, while the Thai partition’s first horizon deviates, possibly from early training or statistical fluctuation.The individual learning-rate scans are provided for completeness.Source: Figure 9 · p. 28, Figure 10 · p. 28

A.12 ABLATION WITH ADAM OPTIMIZER

  • Adam shows norm transfer both without momentum, at Scion’s optimal norm, and with momentum, at a higher optimal norm.With momentum, norm transfer remains but sensitivity to batch-size choice is reduced.Source: Figure 11 · p. 29

A.15.1 MOMENTUM & LEARNING RATE DECAY

  • With momentum and learning-rate decay, a factor ×4–8 range of learning rates can yield the same loss, so the apparent norm drift is not significant.For long horizons, decay smooths the norm optimum and reduces sensitivity to its precise value.Source: §A.15.1 · p. 31

A.15.2 NORM CHOICE

  • Both the output RMS-to-RMS norm and input 1-to-RMS norm also exhibit norm transfer in the ablation.The default comparison uses the output RMS-to-infinity norm.Source: §A.15.2 · p. 32
Loading 2510.03871v2…