μpscaling Small Models: Principled Warm Starts and Hyperparameter Transfer

TL;DR

Most training pipelines ship multi-size model suites by re-training from scratch at each size. μpscaling reframes this as a function-preserving widening problem: given a trained base model of width $n$ , construct a wider model of width $kn$ whose weights, optimizer hyperparameters, and optimizer state are coordinated so that the widened model parametrizes the exact same function and follows the exact same training trajectory as the base model. The authors prove this static + dynamic equivalence for any architecture expressible in the Neor tensor-program framework (MLPs, ResNets, transformers, …) under any entrywise optimizer with a homogeneous update function (SGD, Adam, AdamW). To actually exploit the widened model’s extra capacity they inject Gaussian noise at initialization, with variance scaled by width in exactly the same way μP scales fresh init — which lets them encode the entire upscaling procedure (warm-start + noise injection + continued training) inside μP, recover μTransfer-style zero-shot hyperparameter transfer from cheap upscaling-to-narrow-targets to expensive upscaling-to-wide-targets, and characterize the infinite-width limit of the whole upscaled training dynamics. Empirically: MLP on Forest Cover, ResNet on CIFAR-100, and GPT-2 on FineWeb all show faster convergence and competitive-or-better terminal loss vs from-scratch — with the ResNet validation accuracy as the explicit negative-result counterexample.

Key claims

Duplicating + rescaling weights of a base MLP produces a widened MLP that parametrizes the exact same function for any input, with explicit per-weight multipliers $1/(\mu_i \mu_{i-1}^a)$ where $\mu_i$ are width multipliers [§2.1, Prop. 2.1].
Dynamic equivalence: with appropriately rescaled per-layer learning rates, the widened MLP under SGD follows an identical training trajectory in function space to the base MLP [§2.1, Prop. 2.2]. The hidden-weight learning rate is unchanged under uniform width multipliers; only the input-layer and output-layer LRs need to be rescaled.
The static + dynamic equivalence extends to any entrywise optimizer with a homogeneous update function (covering SGD, momentum-SGD, Adam, AdamW, with or without coupled/decoupled weight decay) with explicit rescaling rules for $\eta$ , $\lambda$ , and additional hyperparameters [§2.2, Prop. 2.4].
The equivalence further extends to any “standard” neural network architecture expressible as a Neor tensor program (MLPs, RNNs, convolution, attention, pooling, residual connections, batch/layer norm), with width-axis-aware rescaling rules for scalar-like, vector-like, and matrix-like weights [§2.3, Thm. 2.5, Table 1].
The proposed widening rules in Table 1 are exactly compatible with μP: if the base model is trained under μP, then the widened model — initialized via the widening operation and continued with the same base constants — automatically inherits the correct μP-rescaled hyperparameters [§2.3, Table 2].
Zero noise gives a function-equivalent widened model that trains as if no widening occurred — the upscaled model never exploits its extra capacity. Injecting Gaussian noise with variance scaled like μP’s fresh init breaks the symmetry, escapes the narrower model’s low-dimensional subspace in parameter space, and unlocks the widened model’s representational capacity [Meta-Algorithm 1, §3.1, Fig. 1(a)].
The noise-injection step admits a Neor-program encoding via partitioned variables, which means the entire upscaling procedure (including the noise) lives inside μP’s analysis framework and inherits μTransfer’s zero-shot hyperparameter transfer guarantees across widths [§3.2].
Hyperparameters (noise std $\sigma$ , base learning-rate constant $\eta$ ) tuned on a cheap small-target upscaling task $n \to k_{\text{small}} \cdot n$ transfer directly to the expensive large-target task $n \to k_{\text{large}} \cdot n$ — verified on GPT-2 (Fig. 3 a/d), MLPs under SGD (3 b/e), and MLPs under AdamW (3 c/f).
Upscaling under the proposed μP-compatible procedure puts training dynamics in the optimal feature-learning infinite-width regime: hidden activations stay at $\Theta(1)$ scale and their updates stay at $\Theta(1)$ scale throughout training; pre-activations converge to i.i.d. blocks of $k$ -dimensional Gaussian vectors (vs i.i.d. scalars without upscaling) whose covariance structure is explicitly tracked via a modified Tensor Program [§3.2, App. E].
On three realistic settings — MLP/AdamW on Forest Cover Type, ResNet/SGD on CIFAR-100, GPT-2/AdamW on FineWeb — upscaled models converge faster and achieve comparable or lower training loss than from-scratch baselines at the same target width [§4, Fig. 2].
Negative result: the upscaled ResNet generalizes worse than the from-scratch baseline on CIFAR-100 validation accuracy, even though it trains faster and reaches lower training loss. The theory accounts for training dynamics only; generalization behavior under upscaling is unanalyzed [§4, Fig. 2 bottom-middle].
The framework explicitly handles upscaling along width axes only; depth upscaling is left to future work [§1, end].

Method

The core construction is a recipe for converting any trained base model into a widened model that is mathematically equivalent at initialization and remains equivalent throughout training under a compatible optimizer.

Static equivalence comes from duplicating each weight matrix along its width axes (Kronecker product with a matrix of ones) and rescaling each entry by $1/\mu^a$ for a fixed exponent — for an MLP, taking $W_i^{\text{wide}} = (W_i \otimes \mathbf{1}_{\mu_i \times \mu_{i-1}}) / \mu_{i-1}$ produces a wider MLP that computes the exact same function on any input. Dynamic equivalence requires the optimizer to update the duplicated weights together at every step — achieved by rescaling per-layer learning rates (and, for general optimizers, weight decay and additional hyperparameters like Adam’s eps) according to the entry in Table 1.

The general result (Theorem 2.5) reframes this via tensor programs: any architecture expressible in the Neor program admits a width-axis classification of its learnable weights into scalar-like, vector-like, and matrix-like, and each class has a fixed widening operation + hyperparameter rescaling rule. Standard sum readouts must be replaced by width-normalized mean readouts to make the rescaling clean. Crucially, the resulting rules match exactly what μP would prescribe for fresh initialization at the wider width — so a μP-trained base model upscales into a μP-correct widened model with no additional bookkeeping.

The upscaling algorithm (Meta-Algorithm 1) is then: (1) widen by Table 1; (2) inject Gaussian noise to the duplicated weights with variance scaled by width exactly as μP scales fresh initialization variance; (3) continue training with the same base constants used for the small model. The noise is what lets the widened model escape the duplicate-subspace and exploit its extra capacity — without noise, the duplicates remain identical throughout training and the model behaves as if no widening happened.

The non-trivial theoretical contribution is showing that even after noise injection, the procedure can be re-encoded inside a (slightly extended) Neor program by introducing partitioned variables, so Yang et al. (2022)‘s μTransfer guarantees apply across widths. This means $\sigma$ and $\eta$ tuned on a small target upscaling task transfer zero-shot to a large target upscaling task — exactly mirroring μTransfer for from-scratch training.

Results

MLP on Forest Cover Type (AdamW). Upscaled MLP converges faster than from-scratch and reaches comparable terminal training loss; validation accuracy also matches [§4, Fig. 2 left].

ResNet on CIFAR-100 (SGD). Upscaled ResNet shows the cleanest training-loss acceleration of the three settings — but is the negative generalization case: validation accuracy is worse than from-scratch despite the faster training-loss convergence [§4, Fig. 2 middle]. The authors flag this explicitly as a case where upscaling does not help final task performance and where the framework’s theory (training dynamics only) does not predict the outcome.

GPT-2 on FineWeb (AdamW). Upscaled GPT-2 converges faster than from-scratch and achieves lower validation loss at fixed step count — the clearest positive result [§4, Fig. 2 right]. This is the setting where the procedure most cleanly delivers compute savings on a frontier-relevant architecture.

Hyperparameter transfer. For GPT-2 specifically (Fig. 3a/d), training base models at widths $n \in \{128, 256, 512, 1024\}$ and upscaling each by $k = 4$ gives optimal $(\sigma, \eta)$ that coincide across all four base widths — i.e. tuning on the cheapest $128 \to 512$ upscaling gives the right hyperparameters for the expensive $1024 \to 4096$ upscaling. Same pattern for MLP/SGD and MLP/AdamW [Figs. 3b/e, 3c/f].

The headline practical claim is that the additional hyperparameter tuning needed by the upscaling procedure ( $\sigma$ + $\eta$ ) can be done at small target widths, so the total compute cost of “train a multi-size suite via upscaling” is dominated by the from-scratch base model plus a single warm-started training run per target size, vs the standard from-scratch-at-each-size cost.

Why it’s interesting

This is the cleanest filed paper on function-preserving model expansion done right under μP. The wiki’s existing hyperparameter-scaling-laws cluster has three positions on how to set LR at scale — Fitting (How to Set the Learning Rate for Large-Scale Pre-training?), Transfer (μTransfer + μA from Learning Rate Scaling across LoRA Ranks and Transfer to Full Finetuning), and Enforce (Controlled LLM Training on Spectral Sphere) — but none address the warm-start axis: what if you don’t want to start from scratch at the target width? μpscaling extends the Transfer branch onto exactly that axis, proving that μTransfer’s zero-shot hyperparameter transfer carries through model expansion when the widening operation is the μP-compatible one.

The result also bears directly on the Training stability at scale question of where stability lives. How to Set the Learning Rate for Large-Scale Pre-training? argues μP’s motivation is stale under QK-Norm; Learning Rate Scaling across LoRA Ranks and Transfer to Full Finetuning (μA) shows μP-style analysis is still load-bearing for LoRA finetuning even where it’s eroded for from-scratch pretraining; this paper adds a third finetuning-adjacent regime — function-preserving widening warm-starts — where μP-style analysis remains essential. The pattern (μP analysis survives in the adaptation/expansion regimes even where it’s stale for from-scratch) is starting to be visible across three papers. The associated GitHub release (yuxinma98/mupscaling) extends the existing μP software package, which is the same dependency the SSO line (Controlled LLM Training on Spectral Sphere) and μA build on — so all three converge on a single shared tooling stack.

The negative ResNet generalization result is also worth flagging: this is one of the few filed papers that explicitly carves out a regime where its proposed recipe gives faster training but worse validation. The theory is upfront that it predicts dynamics, not generalization — which is the same disclaimer Yang et al. (2022) attach to μTransfer itself, and a useful caveat for anyone planning to use upscaling as a production cost-saver.