Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer

Every neural network needs hyperparameters: the learning rate controls how big each update step is, the initialization scale sets the starting point for weights, and optimizer settings like momentum determine how past gradients influence the current step. Choosing these well is the difference between a model that converges to strong performance and one that diverges or stalls.

For small models, you can afford to try hundreds of combinations — a grid search or random search over learning rates and initialization scales. But for a model with billions of parameters, a single training run might cost $100,000 in compute. Running 200 experiments is simply not feasible.

Practitioners have long observed that hyperparameters tuned on a small model sometimes work on a larger model — but sometimes they fail catastrophically. Under standard parameterization, the optimal learning rate can shift by orders of magnitude when you double the width. There is no guarantee, and no theory explaining when transfer works and when it doesn't.

To understand μP, we first need to see what goes wrong with standard parameterization (SP) — the default in frameworks like PyTorch. In SP (including Xavier and He initialization), each weight in a hidden layer is drawn from a distribution with variance proportional to $1/\text{fan\_in}$. The learning rate $\eta$ is the same for all parameters.

This setup works fine at a fixed width. But as you increase the hidden dimension $n$, trouble appears. The activations grow with width, the gradients shift in scale, and the learning rate that was perfect for width 256 becomes disastrous at width 4096. Worse, to keep the network stable at larger widths, you must shrink the learning rate — which pushes the network into the lazy training regime (also called the NTK regime), where hidden features barely change from their initialization. The network functions like a fixed kernel machine rather than learning rich features.

μP solves this by prescribing width-dependent scaling rules. The key changes for a layer with fan-in $n$ (compared to a base width $n_0$) are:

Hidden-layer initialization: variance scales as $1/n$ (same as SP).

Hidden-layer learning rate: scales as $\eta_0 \cdot (n_0 / n)$, i.e., the learning rate shrinks with width — but the per-coordinate update to activations stays $\Theta(1)$.

Output layer multiplier: the output logits are multiplied by $n_0 / n$, preventing them from growing with width.

Attention scaling (for Transformers): the softmax temperature scales as $1/d$ instead of $1/\sqrt{d}$.

The result is that every activation, every gradient, and every parameter update has a width-independent coordinate size. The network behaves the same dynamically at width 64 as at width 8192.

μP is not an arbitrary recipe — it emerges from three precise mathematical requirements. Think of them as engineering constraints: the network must satisfy all three to guarantee that hyperparameters transfer across widths.

Rule 1 — Stable initialization: The pre-activations at every layer must have $\Theta(1)$-sized coordinates at initialization. This means that if you look at any single coordinate of the hidden state, its typical magnitude does not depend on width $n$.

Rule 2 — Stable output: The network output must remain $\Theta(1)$ throughout training. If the output logits grow with width, the loss landscape changes shape and the optimal learning rate shifts.

Rule 3 — Maximal feature updates: The change in hidden features $\Delta h$ at each training step must also be $\Theta(1)$. This is the maximal part — updates are as large as possible without causing instability. If updates shrink to zero (as in the NTK regime), the network stops learning features.

The mathematical machinery behind μP is the abc-parameterization. For each weight matrix $W$ in a layer with fan-in $n$, three exponents $(a, b, c)$ control its scaling:

The layer's output is multiplied by $n^a$ (the output multiplier). The initialization variance is $n^{-2b}$ (controlling the starting scale). The learning rate is $\eta \cdot n^{-c}$ (controlling how fast the layer updates).

Different choices of $(a, b, c)$ give different parameterizations. Standard parameterization corresponds to one set of values; NTK parameterization to another. μP corresponds to the unique set of $(a, b, c)$ values that satisfy all three desiderata above.

This is perhaps the deepest conceptual insight in the paper. Neural networks can operate in two fundamentally different regimes as width grows:

In the kernel (NTK / lazy) regime, the hidden features $h^{(\ell)}$ barely change from their random initialization during training. The network effectively becomes a linear model around its initialization — its predictions are governed by a fixed kernel (the Neural Tangent Kernel). This is mathematically elegant but fundamentally limits what the network can learn, because the internal representations stay frozen.

In the feature learning regime, hidden features do evolve during training. The network discovers meaningful patterns in the data and builds increasingly useful internal representations. This is what makes deep learning powerful — and it is what μP preserves.

The choice of parameterization determines which regime the network lands in. Standard parameterization pushes wide networks toward the kernel regime. μP is the unique parameterization that keeps networks in the feature-learning regime at any width.

The paper turns the theory into a simple three-step recipe called μTransfer:

Step 1 — Parametrize in μP: Set up the target model (the large one you actually want to train) using μP scaling rules. This means adjusting initialization variances, per-layer learning rates, and output multipliers according to the abc-exponents.

Step 2 — Tune on a proxy: Build a small "proxy" model — same architecture and depth, but much narrower (e.g. 13M parameters instead of 350M). Run a large hyperparameter search on this proxy. Because of μP, the optimal learning rate, momentum, and other non-regularization hyperparameters are the same at any width.

Step 3 — Transfer: Take the best hyperparameters from the proxy and plug them directly into the large model. Train the large model once — no further tuning needed.

The authors demonstrate this concretely: a 200-sample random search on a 13M-parameter proxy yields hyperparameters that, when transferred to BERT-large (350M), outperform published BERT-large results. The total tuning cost equals one BERT-large pretraining run. Similarly, tuning on a 40M proxy produces results exceeding published GPT-3 6.7B numbers, at only 7% of pretraining cost.

How do you know your μP implementation is correct? The paper introduces a simple diagnostic called the coordinate check. The idea: train the model for a few steps at several different widths and plot the average coordinate size of activations (or logits) over training steps. If μP is implemented correctly, the curves for different widths should overlap — the coordinate size is independent of width. If they fan out (grow with width), something is wrong.

This test is cheap to run and catches implementation bugs immediately. It has become the standard verification tool for μP implementations. The open-source mup library (installable via pip install mup) automates both the parameterization and the coordinate check.

The paper validates μTransfer on two major architectures:

BERT-large (350M): Hyperparameters tuned on a 13M-parameter proxy are transferred to BERT-large. The resulting model outperforms the published BERT-large baseline, and the total cost of tuning plus one training run equals a single BERT-large pretraining.

GPT-3 6.7B: Hyperparameters tuned on a 40M proxy transfer directly to the 6.7B model. The result exceeds published GPT-3 6.7B performance, with tuning cost only 7% of pretraining. This is the first demonstration that hyperparameter transfer works across a 150× size gap.

The authors also verify that μTransfer works for ResNets on image classification, showing the approach is not specific to Transformers. Critically, they show that regularization-related hyperparameters (like weight decay and dropout) do not transfer — only optimization hyperparameters (learning rate, momentum, initialization) do.

μP and μTransfer changed the economics of training large models. Before this work, hyperparameter tuning at the billion-parameter scale was essentially guesswork combined with expensive trial-and-error. After it, teams could allocate a fixed compute budget to tuning a tiny proxy, then apply those settings with mathematical confidence.

Cerebras adopted μP for their GPT family, demonstrating that a 200-sample random search on a 40M proxy could set the learning rate for models up to 13B parameters. Multiple research labs have since integrated the mup library into their training pipelines. The framework also deepened the community's understanding of the relationship between parameterization, feature learning, and scaling laws — connecting practical recipes to deep theoretical insights from the Tensor Programs series.