$μ$pscaling small models: Principled warm starts and hyperparameter transfer

AI summary

Overview

Research area: Deep learning training methodology — specifically model upscaling (warm-starting large models from trained small ones), hyperparameter transfer, and the theory of infinite-width limits (the Maximal Update Parametrization, μP, and Tensor Programs).

Technical level: Advanced. The paper is primarily mathematical (propositions, proofs, infinite-width limit analysis) and assumes familiarity with neural network training dynamics, optimizer update rules, and width-based scaling theory.

One-sentence scope: The paper builds a theory of "dynamic equivalence" between narrow and wide neural networks, uses it to define a principled copy-and-perturb upscaling algorithm with μP-based hyperparameter rescaling, and shows that hyperparameters tuned on small upscaled models transfer to large ones.

Affiliation: Department of Applied Mathematics and Statistics, Johns Hopkins University. Code: https://github.com/yuxinma98/mupscaling

What This Paper Is About

Organizations train and release neural networks in families of different sizes (the paper cites LLaMA and GPT as examples of such suites) to meet different hardware, latency, and cost constraints. Training each size from scratch wastes compute, so a natural shortcut is to upscale: initialize a large model from an already-trained smaller one. The obstacle is that upscaled training has different dynamics than training from scratch, so it needs its own hyperparameters — and tuning those directly at large scale is prohibitively expensive. The paper asks whether hyperparameters tuned on small upscaled models can be soundly extrapolated to large upscaled models, and answers yes by combining function-preserving weight duplication with the μP scaling rules.

Key Contributions

1. A theory of dynamic equivalence (Section 2). The paper proves that pairing a function-preserving weight transformation (duplicating and rescaling weights) with a coordinated rescaling of the learning rate, weight decay, and other optimizer hyperparameters makes a wider model compute the same function as the narrow model throughout training, not just at initialization. This is shown for a bias-free MLP trained with SGD (Propositions 2.1 and 2.2), extended to general entrywise optimizers whose update functions are homogeneous of degree m (Proposition 2.4, covering Adam, AdamW, and SGD with momentum), and extended informally to general "standard" architectures via Theorem 2.5 and the Tensor Programs framework.

2. A principled upscaling algorithm (Section 3, Meta-algorithm 1). After widening a checkpoint, the method injects a small symmetry-breaking noise so the upscaled model can escape the lower-dimensional subspace it would otherwise be confined to. Mechanically similar to Net2WiderNet (Chen et al., 2016), but the noise magnitude and optimizer hyperparameter scalings are theory-derived rather than heuristic. The paper states this is the first rigorous framework for hyperparameter transfer in the upscaling setting.

3. Infinite-width analysis via Tensor Programs (Section 3.2, Appendix E). Using the Tensor Programs machinery of Littwin and Yang (2023), the authors characterize the infinite-width limit of upscaled training for common architectures, providing the formal basis for showing the injected noise is "safe" and "useful."

4. Software and experiments (Section 4). The authors extend the μP software package (Yang et al., 2021) so upscaling applies to standard architectures including multi-layer perceptrons (MLPs), ResNets, and transformers, and report experiments across various architectures and datasets illustrating hyperparameter transfer numerically.

Main Findings

• Zero noise means exact equivalence. With zero injected noise, the widened model's training trajectory evolves exactly as if training had simply continued on the base model, up to a possible change in the learning-rate base constant. The paper says this is validated empirically across architectures and optimizers in Appendix C.4.

• Widening rule for matrix-like weights. Weights are widened as k_in^{-1} · W ⊗ (1_{k_out} 1_{k_in}^T), which in PyTorch terms is W.repeat_interleave(k_out, dim=0).repeat_interleave(k_in, dim=1) / k_in, with N_out = n_out · k_out and N_in = n_in · k_in; scalar-like weights are unchanged and vector-like weights are duplicated along the width axis.

• Required learning-rate rescaling. For general entrywise optimizers with an update function homogeneous of degree m, the widened learning rate is γ↑ = k_out^m k_in^{-1} γ, the additional hyperparameter scales as ε↑ = k_out^{-1} ε, and the weight-decay coefficient scales as k_in k_out^{-1} λ under vanilla weight decay or k_in k_out^{-m} λ under decoupled weight decay.

• Equal width multipliers simplify things. If all width multipliers are equal (k_1 = ... = k_{L-1}), then learning rates for all hidden weights are left unchanged; only the learning rates for W^(1) and W^(L) are modified.

• μP compatibility is automatic. The widening rules in the paper's Table 1 coincide with the μP scaling rules of Yang and Hu (2021) and Littwin and Yang (2023), so under μP the procedure becomes "essentially mechanical": transfer weights, and hyperparameters adjust themselves.

• Noise is provably safe and useful at infinite width. In the infinite-width limit the base-model signal is preserved, and the injected noise contributes meaningfully without vanishing or exploding.

• Zero-shot hyperparameter transfer. Hyperparameters can be tuned on narrow upscaled models and the optimal values transfer directly to large-scale upscaling. The paper reports numerically illustrating this transfer (Figure 1(b)).

• A built-in rigorous baseline. Because zero noise reproduces continued base-model training, any hyperparameter sweep over the noise level that includes zero automatically contains the "just keep training the small model" baseline. If the sweep picks a nonzero noise level, the benefit over continuing base-model training is genuine.

• Optimizer state must be transferred. For optimizers with internal state (momentum in SGD, first- and second-moment estimates in Adam/AdamW), achieving equivalence mid-training requires duplicating and rescaling the accumulated state alongside the weights; the specific rules are given in Appendix D.

• Not reported in the available text. The truncated content does not report specific dataset names, dataset sizes, benchmark scores, model configurations, or numeric experimental results. The abstract only states experiments were run on "realistic datasets and architectures."

Methodology in Plain English

The approach has three layers.

First, define what it means for two models of different widths to behave identically. The authors start from the well-known trick of building a wider network whose outputs match a narrow one exactly by duplicating each hidden unit and dividing the outgoing weights by the duplication factor. They then ask a stronger question: can the two networks stay identical through gradient updates? Working through the backpropagation recurrence for an MLP, they derive the exact per-layer learning-rate rescaling that makes the wide model's weights remain a duplicated-and-rescaled copy of the narrow model's weights at every step. They then generalize from plain SGD to a broad class of "entrywise" optimizers that includes Adam, AdamW, and momentum SGD, using the assumption that the optimizer's update function is homogeneous of degree m. Finally, they argue that the same reasoning extends to essentially any architecture expressible in the Tensor Programs language — which covers MLPs, RNNs, convolutions, attention, pooling, skip connections, and batch/layer normalization — by classifying each parameter tensor as scalar-like, vector-like, or matrix-like and giving a widening rule and hyperparameter rescaling rule for each class (Table 1).

Second, break the symmetry deliberately. A perfectly equivalent wide model is stuck in a lower-dimensional subspace of its parameter space and cannot use its extra capacity. So the algorithm adds a small noise perturbation to the widened parameters, scaled the same way μP scales random initialization, with a tunable base constant. This lets the model explore the extra dimensions while keeping the inherited signal from the small model.

Third, make tuning cheap. Because the hyperparameter scalings are pinned down by the μP rules, the only free quantities are the underlying width-independent base constants — the noise level and the learning-rate base constant. These are tuned not on the target large model but on a small upscaling system (width n_0 to k·n_0 with n_0 much smaller than the target width n), selecting the values that minimize terminal training loss. Those tuned base constants then transfer to the target width. A variant in Appendix D instead sets each layer's noise magnitude as a fraction t in [0, 1] of that layer's spectral norm.

The infinite-width behavior is analyzed with the Tensor Programs framework to show the perturbation neither destroys the inherited signal nor blows up.

Why This Matters

Impact on research. The paper turns a largely heuristic practice into something with provable guarantees. It supplies a rigorous baseline (zero noise equals continued training on the small model), a formal justification for scaling hyperparameters across widths in the upscaling setting, and a reusable theoretical toolkit that links model upscaling to the μP literature. It also opens a specific, well-posed question — whether scaling-law extrapolation of hyperparameters is sound when upscaling is involved — and answers it affirmatively.

Real-world applications (as framed by the paper):

• Training and releasing model suites at multiple sizes, as done for families like LLaMA or GPT, where each size serves a different budget.
• Reducing total compute in multi-scale training workflows that currently discard small models and retrain from scratch.
• Rapid prototyping and scaling-law analysis, where small models guide architecture choices before committing compute to large ones.
• Deploying the same model family under diverse hardware, latency, and cost constraints.

Industry relevance. The central practical claim is cost: hyperparameter tuning at the scale of modern large models is described as infeasible, and the paper's transfer result lets practitioners tune on small upscaled models instead. Because the method is implemented as an extension of the existing μP software package and is stated to work for MLPs, ResNets, and transformers, adoption requires modifying an existing training stack rather than building a new one. The paper explicitly leaves depth upscaling to future work, so the current scope is width-wise upscaling.

Future Directions

• Depth upscaling. The paper focuses on width-wise upscaling and states that extending to depth is a natural next step, citing Yang et al. (2023b) and Bordelon et al. (2024b) for hyperparameter transfer across depths.
• Empirical breadth. The available text reports that experiments span various architectures and datasets but does not give dataset names, sizes, or benchmark numbers, so the extent of the empirical validation is not determinable from the truncated content.
• Choice of noise injection. The paper offers both a μP-scaled noise with a tunable base constant and a signal-normalized variant defined as a fraction t of each widened weight's spectral norm; the practical trade-offs between these are a natural follow-up.
• Full formal treatment of general architectures. The general architecture result (Theorem 2.5) is stated informally in the main text, with the detailed version deferred to Theorem C.5 in Appendix C.2, leaving room for further theoretical tightening.

Target Audience

Researchers and engineers working on large-scale training efficiency, scaling laws, and hyperparameter transfer — particularly those already familiar with μP or Tensor Programs. It is also relevant to practitioners maintaining multi-size model families who want a principled warm-start procedure with a built-in safety baseline. Because the main text is proposition-and-proof heavy, readers without a background in training dynamics or infinite-width theory will need to work through the appendices for the full picture.

Authors’ abstract

Modern large-scale neural networks are often trained and released in multiple sizes to accommodate diverse inference budgets. To improve efficiency, recent work has explored model upscaling: initializing larger models from trained smaller ones to accelerate convergence. However, this method can be sensitive to hyperparameters that need to be tuned at the target upscaled model size, which is prohibitively costly to do directly. It remains unclear whether tuning hyperparameters on smaller models and extrapolating via scaling laws is sound in this setting. We address this with principled approaches to width-based upscaling and efficient hyperparameter tuning in this setting. Motivated by $μ$P and any-dimensional architectures, we introduce a general upscaling method that, like Net2Net, copies and perturbs weights, but uses theoretically grounded, width-dependent scalings for the perturbation noise and optimizer hyperparameters. First, we prove that under zero perturbation, the upscaled model is functionally equivalent to the base model throughout training. Second, we extend the $μ$P theory to enable infinite-width limit analysis and establish hyperparameter transfer for upscaled models, greatly reducing the tuning cost. We empirically demonstrate that this method is effective on realistic datasets and architectures.