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

Scale each layer's initialization and learning rate with width, and the best hyperparameters of a small model also work for a large one.

The paper that introduced μTransfer: parametrize a network in μP (the Maximal Update Parametrization), tune a copy that is 10 to 168 times smaller, and copy the settings to the full model without trying a single configuration on it.

GPT-3 6.7B was tuned with 467 training runs of a 40-million-parameter model, at 7% of the cost of training it once.

Tuning a model you can train once

A hyperparameter is a setting you choose before training rather than learn: the learning rate, the scale of the random initial weights, the learning-rate schedule, the Adam decay rates, constant multipliers on the outputs of some layers. The loss you end up with depends strongly on them, and the usual way to find good values is a search: train many copies of the model with different settings and keep the best one. For a model that costs millions of dollars to train once, that search is not affordable.

The obvious workaround is to tune a small version of the model and reuse its settings. Practitioners have always done this, and the paper describes the results as "hit-or-miss at best." Published model families were configured size by size. GPT-3's eight models share one recipe, yet the peak learning rate falls tenfold from $6\times 10^{-4}$ for the 125M model to $0.6\times 10^{-4}$ for 175B. Each size got its own value, and a sweep at 175B parameters would have cost several full training runs of the largest model.

This paper makes the small-model workaround reliable. Its claim is that the optimal hyperparameters drift with width only because of how networks are conventionally parametrized: how the initialization variance and learning rate of each weight tensor are set as a function of its shape. The authors use a different rule, the Maximal Update Parametrization (μP), derived by Greg Yang and Edward Hu in an earlier paper of the Tensor Programs series. Under μP, the optimal learning rate, initialization scale and several other settings stay approximately fixed as the model gets wider, meaning more units per layer (the hidden size of an MLP, $d_{model}$ and the feed-forward and head dimensions of a Transformer). The procedure built on it, which the paper calls μTransfer, is three steps:

1. Parametrize the target model in μP.
2. Tune a smaller version of it (narrower, and optionally shallower, trained on fewer tokens).
3. Copy the tuned hyperparameters to the target model and train it once.

The word "parametrization" here means a rule for scaling hyperparameters with width ("how should my initialization and learning rate change when my width doubles?"), which the paper distinguishes from choosing values at one width. μP and the default parametrization can produce exactly the same network at one width; they differ in what happens when you make it wider.

A few ideas explain the method, in this order: a toy problem where a rescaled constant transfers and the raw one does not; what one training step does to a wide network under the default setup; the two different ways a sum of $n$ terms grows, which fixes every rule in μP; the rules themselves and how to check an implementation of them; then the transfer experiments and their limits.

The right coordinates for a hyperparameter

Section 2 of the paper starts from the central limit theorem. If $x_1, \ldots, x_n$ are independent with mean 0 and variance 1, their sum has a standard deviation of $\sqrt n$, so $(x_1 + \cdots + x_n)/\sqrt n$ converges to a standard Gaussian as $n$ grows. Dividing by $n$ instead sends the sum to zero, and dividing by nothing lets it grow without bound. Now treat the scale as a hyperparameter $c$ and minimize

over $c$, for some bounded function $f$. Here $n$ plays the role of width, the $x_i$ the random initial weights, $c$ a hyperparameter such as the learning rate, and $f$ the loss after training. The best $c$ depends on $n$: the sum has size about $\sqrt n$, so the best $c$ shrinks like $1/\sqrt n$. Write $c = \alpha/\sqrt n$ instead, and the central limit theorem makes $F_n$ converge to a fixed function of $\alpha$:

So the best $\alpha$ found at a small $n$ is close to the best $\alpha$ at any larger $n$. Tuning $\alpha$ on the small problem and copying it is a correct transfer procedure; tuning $c$ and copying it is not. The paper calls $c = \alpha/\sqrt n$ the correct parametrization for this problem.

The figure computes this exactly for one concrete case. Each $x_i$ is $+1$ or $-1$ with equal probability, so the sum is a shifted binomial and $F_n$ is a finite sum with no sampling noise. The loss is $f(y) = -y^2 e^{-y^2/2}$, which is bounded, and for Gaussian input it gives the closed form $-\alpha^2/(1+\alpha^2)^{3/2}$, minimized at $\alpha = \sqrt 2 \approx 1.414$. With the proxy at $n = 16$, the best raw constant is $c^* = 0.349$. Copied unchanged to $n = 4096$, it gives a loss of $-0.045$ against the achievable $-0.385$, 88% short. Copying $\alpha^* = 1.395$ instead lands within 0.1% of the optimum.

The smallest widths also show a feature of the real experiments: the optimum in $\alpha$ is not exactly constant. It moves from 1.265 at $n=4$ to 1.395 at $n=16$ and then barely changes. Transfer is accurate once the proxy is wide enough to be near the limit, and the paper reports the same thing for Transformers (a minimum width around 256).

A network has many hyperparameters, and Appendix C extends the toy to several constants whose effects add:

If you reparametrize only $c^1$ as $\alpha^1/\sqrt n$ and leave the others as raw constants, the best $\alpha^1$ is

which still grows with $n$. The correctly parametrized hyperparameter has to compensate for the incorrectly parametrized ones, and its optimum drifts anyway. In a network the other constants are the initialization scales, the output multipliers and the per-layer learning rates, so μP has to rescale all of them, not only the global learning rate.

The default network after one step

The default in PyTorch and other frameworks is what the paper calls the standard parametrization (SP): every weight matrix is initialized with variance proportional to $1/\text{fan\_in}$ (the number of inputs to each unit, as in LeCun or He initialization), and one global learning rate is applied to every tensor. For the paper's 2-hidden-layer MLP of width $n$:

This initialization is designed so that every activation has entries of size about 1 at the start, at any width. The imbalance appears at the first update. Section 5 of the paper shows it on the smallest possible network: a 1-hidden-layer linear model $f(x) = V^\top U x$ with scalar input and output and hidden width $n$, so $U$ and $V$ are vectors of length $n$. In SP, $U_\alpha \sim \mathcal N(0,1)$ and $V_\alpha \sim \mathcal N(0, 1/n)$. One SGD step with learning rate 1 moves each vector along the other, $V' = V + \theta U$ and $U' = U + \theta V$, where $\theta$ is a scalar that does not depend on the width $n$, set by the input, label and loss. The new output is

The term $\theta\,U^\top U$ is a sum of $n$ squares of standard Gaussians, so it is about $\theta n$. At $n = 4096$ and $\theta = 0.01$, the output jumps by about 41 after one step; at $n = 64$, by about 0.65. The same step moves the wider model's output about 63 times as far. That term comes from updating $V$, the output layer. The term from updating $U$, the input layer, is $\theta\,V^\top V \approx \theta$, the same at every width.

Dividing the learning rate by $n$ tames the output layer, but it also divides the input layer's contribution by $n$, so at large width the input layer stops learning. No single global learning rate works for both layers once the model is wide. The paper sees the same split in a Transformer (its Figure 5): after one Adam step, the output logits and attention logits change by an amount that grows with width, while the word embeddings change by a width-independent amount. A learning rate small enough to keep the logits stable leaves the embeddings nearly frozen in a wide model.

μP fixes this per layer. It initializes $V_\alpha \sim \mathcal N(0, 1/n^2)$, keeps $U_\alpha \sim \mathcal N(0,1)$, and uses learning rates $\eta_V = 1/n$ and $\eta_U = n$. After one step,

where $n^{-1}U^\top U \approx 1$, and $nV^\top V \approx 1$ because $V^\top V$ sums $n$ squares of size $1/n^2$. Both layers now move the output by about $\theta$ at every width, and the initial output $V^\top U$ has standard deviation $1/\sqrt n$, so it starts near zero in a wide model.

Two kinds of sums

Every rule in μP comes from one question about a matrix-vector product. A layer computes $Wx$, and each entry of the result, $\sum_j W_{ij} x_j$, is a sum of $n$ products. Assume the entries of $x$ have size about 1. How large must the entries of $W$ be for the sum to also have size about 1? The answer depends on whether the terms of the sum have random signs or a common sign (Appendix J.1, Table 14):

• Random initialization. $W_{ij}$ is drawn independently of $x_j$, so the products have random signs and mostly cancel. The sum grows like $\sqrt n$ (central limit theorem), and entries of size $1/\sqrt n$ keep it at size 1. Frameworks already initialize this way, with variance $1/\text{fan\_in}$.
• A training update. One SGD step on one example changes a weight matrix by $\Delta W = -\eta\, \delta\, x^\top$, where $\delta$ is the gradient arriving at the layer's output. When the next forward pass sees a similar input, the product $\Delta W_{ij}\, x_j$ equals $-\eta\,\delta_i\, x_j^2$, which has the sign of $-\eta\,\delta_i$ for every $j$. Nothing cancels, the sum grows like $n$ (law of large numbers), and entries of size $1/n$ are needed to keep it at size 1.

A concrete case: a hidden layer of width $n = 4096$, with input entries of size about 1. At initialization the weights have standard deviation $1/\sqrt{4096} = 1/64$, and each output entry is a sum of 4096 terms of random sign, giving a Gaussian with standard deviation about 1. After one update with entries of size $1/64$, the change in each output entry is 4096 aligned terms of size $1/64$, about 64. An update entry of size $1/4096$ brings it back to about 1.

With the entry scale at $1/\sqrt n$ and $n = 4096$, the figure reads about 64 for the aligned sum and about 1 for the random one. So the initial weights and the change in the weights need different scalings with width. SP scales only the initialization, so at a fixed learning rate its hidden-layer updates come out $\sqrt n$ times too large with SGD and $n$ times too large with Adam, and shrinking the global rate to compensate makes the input-layer updates too small.

The optimizer determines which setting controls the update size. SGD's update is the raw gradient times the learning rate, and the gradient's size depends on the layers around it. Adam divides each gradient entry by its own running RMS (see the Adam explainer), so every entry of the update has size about the learning rate, whatever the gradient's scale. For Adam, "update entries of size $1/n$" therefore means a learning rate proportional to $1/n$ for hidden weights, which is the Adam row of the μP table below. The original μP paper covered only SGD, and deriving the Adam version is one of this paper's contributions.

The same arithmetic changes attention. The Transformer divides the query-key dot product $q^\top k$ by $\sqrt d$, which is right when $q$ and $k$ are independent, as at initialization. After training, queries and keys that should match are correlated, the $d$ products share a sign, and $q^\top k$ grows like $d$. μP divides by $d$ instead.

The μP rules, layer by layer

The paper sorts weight tensors by how many of their dimensions grow with width. A dimension is "infinite" if it scales with width (in a Transformer: $d_{model}$, $d_{ffn}$, $d_{head}$, $n_{head}$) and "finite" otherwise (vocabulary size, context length, number of classes). Input weights map finite to infinite, output weights map infinite to finite, and hidden weights map infinite to infinite. Biases count as input weights with a constant input of 1. Applying the two sums above to each kind gives the paper's Table 3, reproduced here with SP in parentheses where it differs:

Each entry follows from the two sums. Output weights ($1 \times n$ per output) sum $n$ terms that become aligned with the input during training, because the change in the last hidden layer is driven by gradients that flowed back through these same weights. Both the weights and their updates need entries of size $1/n$: variance $1/n^2$ and learning rate $1/n$ for SGD and Adam alike. Hidden weights keep the $1/\sqrt n$ initialization (random sum) and need updates of size $1/n$ (aligned sum). With Adam that means learning rate $1/n$. With SGD the learning rate stays 1, because the gradient reaching a hidden layer is already of size $1/n$ per entry once the output weights are that small. Input weights sum over a fixed number of inputs $d_{in}$, so their entries and updates should be of size 1: Adam learning rate 1, and SGD learning rate $n$ to undo the $1/n$-sized gradient.

In practice nobody wants $n$ itself in the formulas, because then no width would match the standard setup. The paper inserts a base width $n_0$ and writes each factor of $n$ as the width multiplier $\tilde n = n/n_0$. For the MLP with SGD, μP in basic form and with a base width are

At $n = n_0$ every $\tilde n$ is 1 and (4) is exactly the SP network of (2). You can take a model whose hyperparameters are already tuned at some width, declare that width the base, and nothing about it changes.

The same network also admits other bookkeeping. A weight $W$ with a constant multiplier $A$ in the forward pass (the layer computes $AWx$), initialization standard deviation $B$ and learning rate $C$ trains to exactly the same function if you replace them by $A\theta$, $B/\theta$ and $C/\theta^2$ for SGD, or $C/\theta$ for Adam (Lemma J.1). The product $AW$ starts the same, and its SGD update scales as $CA^2$ while its Adam update, which ignores the gradient's scale, scales as $CA$. Taking $\theta = 1/\text{fan\_in}$ on the output layer moves the $1/n$ out of the initialization and into a multiplier and gives the paper's Table 8. The released package implements this form because it lets the input embedding and output unembedding share one matrix as most Transformers do. For Adam with base width 128 and target 4096 it looks like this:

# muP for Adam by hand (the paper's Table 8 form; what the mup package does)
# base width n0 = 128, target width n = 4096
m = n / n0                        # width multiplier = 32

W1 ~ N(0, 1/d_in);  lr = eta      # input weights: same as the default
b1, b2 = 0;         lr = eta      # biases: same as the default
W2 ~ N(0, 1/n);     lr = eta / m  # hidden weights: learning rate / 32
W3 ~ N(0, 1/n0);    lr = eta      # readout: variance frozen at the base value
logits = W3 @ (h2 / m)            # ... and its input divided by 32
# transformer only: attention logits use 1/d instead of 1/sqrt(d)
score = alpha_attn * sqrt(d_head0) / d_head * (q @ k)

In the official mup package the same thing is three changes to a PyTorch model: replace the output layer with MuReadout, call set_base_shapes with a base model and a second model that differs only in width (so the package can tell which dimensions scale), and use MuAdam. In the package source, MuAdam puts every tensor with two width dimensions in its own parameter group with learning rate divided by fan_in / base_fan_in, and MuReadout divides its input by the width multiplier and multiplies its initial weights by the square root of it. Attention is not automatic: you change the $1/\sqrt d$ to $1/d$ yourself.

from mup import MuReadout, set_base_shapes, MuAdam

class MLP(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.l1 = nn.Linear(32, n)
        self.l2 = nn.Linear(n, n)
        self.out = MuReadout(n, 10)        # was nn.Linear(n, 10)

model = MLP(4096)
set_base_shapes(model, MLP(128), delta=MLP(256))  # tag which dims are "width"
opt = MuAdam(model.parameters(), lr=2**-7)       # hidden lr becomes 2**-7 / 32

Checking an implementation

If one layer is left in the default parametrization, transfer breaks and nothing raises an error. Appendix D.1 proposes a test the authors compare to gradient checking for autograd, called a coordinate check. Train models of several widths for a few steps, and for each activation vector record the standard deviation of its entries' change from initialization. In a correct μP implementation every such curve is flat in width. A curve that rises or falls with width identifies the layer that is scaled wrong.

We ran a coordinate check on a 2-hidden-layer ReLU MLP (32 inputs, width $n$, 10 outputs) with Adam at learning rate $2^{-7}$, at widths 128 to 8192, in three settings: PyTorch defaults (SP), SP with the global learning rate divided by $n/128$, and μP through the mup package with base width 128.

In SP the hidden-layer change after one step grows 54-fold from width 128 to 8192 (the aligned sum, with Adam update entries of size equal to the learning rate), and the logit change grows about 730-fold. Dividing the learning rate by width flattens the hidden layer, but the input layer's change falls 64-fold, exactly the ratio of the widths: the input layer stops learning, as the linear example predicted. Under μP all three curves stay within 15% of their width-128 values across a 64-fold range of width. The paper recommends running this check before any tuning; its second use is debugging, since "wider is better" failing early in training is also a symptom of a scaling bug.

Tune small, copy, train once

A flat coordinate check says each layer moves by a width-independent amount per step. The transfer claim goes further: the whole loss-versus-hyperparameter curve converges as width grows, so its minimum stops moving. The paper's Figure 1 shows it for Transformers on WikiText-2 trained with Adam: in SP the best learning rate shifts with width and wider models are not always better, and the width-8192 SP model at its own best learning rate still loses to the width-8192 μP model. In μP the minimum stays put and loss improves monotonically with width. For an MLP on CIFAR-10 trained with SGD (the paper's Figure 3), the SP optimum moves by roughly an order of magnitude between widths 256 and 8192.

Figure 4 repeats that experiment at laptop scale with the same MLP: widths 128 to 2048, trained with Adam for 1200 steps at batch 128 on a fixed synthetic 10-class dataset of 32,768 examples (Gaussian inputs, labels sampled from a small random teacher network), learning rates from $2^{-10}$ to $2^{-3.5}$ in half-octave steps, 2 seeds each. Following Appendix D.2, the output layer is initialized to zero in both parametrizations, so at the base width 128 they are the same network and their curves coincide exactly.

In our run the SP optimum moves from $2^{-7.5}$ at width 128 to $2^{-9.5}$ at width 2048, two octaves (a factor of 4) in the direction the paper reports. The μP optimum moves one octave the other way, to $2^{-6.5}$. Copying the width-128 learning rate to width 2048 gives a loss of 1.052 in SP, 0.159 above the best SP loss at that width, and 1.002 in μP, 0.074 above the best μP loss. At the fixed rate $2^{-7.5}$, μP's loss falls at every doubling of width, from 1.107 to 1.002, while SP's rises again at width 2048: the "wider is better" behavior of the paper's Section 8.

Our run does not reproduce the paper's performance claim. Here SP, tuned separately at every width, reaches a lower loss than μP (0.893 against 0.928 at width 2048), and μP's optimum drifts more than in the paper's figures. The setting is far from the paper's: a 32-input MLP trained for 1200 steps, starting at width 128, below the minimums the paper gives for reliable transfer in Transformers (width 256 and 5000 steps). The paper also reports that the optimum still shifts slightly at small width and names finite-width corrections as open work.

Width is the only dimension the theory covers. The paper also tests transfer across depth, batch size, sequence length and number of training steps on Transformers, empirically. Its Figure 4 sweeps learning rate, the output multiplier, the initialization scale and the learning-rate schedule on a 2-layer pre-layernorm Transformer on WikiText-2, varying one scale dimension at a time. The rule of thumb that comes out: optimal hyperparameters transfer once the proxy has width at least about 256, depth 4, batch size 32, sequence length 128 and 5000 training steps, and the target stays within the range tested. Table 1 of the paper groups what transfers:

The asterisks mark dimensions validated only empirically, and only on Transformers. Regularizers are excluded because the right amount of regularization depends on how much data the model has relative to its size, and a rule that depends only on width has no term for the data-to-model ratio. The paper's experiments measure training loss, or test loss in settings where overfitting is not the bottleneck, such as pretraining.

Section 9 points out a question the method raises and does not answer. For transfer to be useful, the proxy must be wide enough that its optimal hyperparameters have converged to the large model's, yet narrow enough that its loss has not converged to the large model's; otherwise there would be no reason to train the large one. The experiments show such widths exist. The authors describe the optimum as a coarse variable that converges quickly with width and the trained function as a fine one that converges slowly, and leave the explanation to future work.

Translation, BERT, GPT-3

The paper then tunes real models only through a proxy and compares against baselines tuned directly. In every experiment it also tries naive transfer: the same small-model tuning, but in SP.

IWSLT14 German-English translation. The 40M-parameter fairseq Transformer (post-layernorm) is the target; the proxy has a quarter of its width, 4M parameters. Random search tunes the learning rate, the output multiplier and the attention multiplier. At equal total compute, 64 samples on the proxy compete with 5 samples on the full model, and each whole search is repeated 25 times. μTransfer's median BLEU is 35.33 against 35.00 for direct tuning, its 25th percentile 35.27 against 33.62, and its best 35.53 against fairseq's presumably well-tuned default of 35.40. Naive transfer diverged. On the paper's compute-versus-BLEU Pareto frontier μTransfer dominates at every budget; per sample, direct tuning is slightly better, and the gap disappears as samples increase.

WMT14 English-German. Target 211M parameters, proxy 15M. With the budget of 3 direct samples, direct tuning reached 25.69 BLEU at best and diverged at worst; μTransfer from the proxy scored 25.94 / 26.34 / 26.42 (worst, median, best of 3 searches) against the fairseq default of 26.40.

BERT. One 13M-parameter proxy (10 layers, $d_{model} = d_{ffn} = 256$, 8 heads of dimension 32) is scaled in both width and depth to BERT-base (110M) and BERT-large (350M). The search samples 256 combinations of six hyperparameters, each trained for $10^5$ steps. That totals 256 × 10⁵ × 13M ≈ 3.3 × 10¹⁴ parameter-steps, about the same as one BERT-large run of 10⁶ steps × 350M = 3.5 × 10¹⁴. BERT-large's pretraining test loss falls from 1.731 with the published Megatron settings to 1.683; MNLI goes from 86.3/86.2 to 87.0/86.5 and QQP from 90.9 to 91.4. BERT-base improves by less (loss 1.995 to 1.970). Naive transfer diverged for both.

GPT-3 6.7B. The target has 32 layers at width 4096. The proxy keeps the depth and shrinks the width to 256, about 40M parameters, 168 times smaller. The random search covered learning rate, initialization scale, attention and output temperatures, and embedding and relative-position multipliers, with 350 runs on 4B tokens and 117 on 16B (286 and 80 of them finished without diverging). Learning rate and initialization scale mattered most, and the chosen values were 0.006 and 2.5, with an embedding multiplier of 10 and the rest left at 1. The search also found that a linear learning-rate decay beat the default cosine on the proxy, and the transferred model used it. Counting parameters times tokens as FLOPs, the search cost

of one full training run, which the paper rounds to 7%. The μTransfer model reached validation loss 1.98 against 2.03 for a rerun of the original settings, 73.5% zero-shot LAMBADA accuracy against 70.3% for the published 6.7B model (and 72.5% for the published 13B), and 72.0% zero-shot HellaSwag against 67.4% (13B: 70.9%). It was not better everywhere: few-shot LAMBADA dropped to 74.7% against the published 79.1%. The comparison is also not fully controlled. The rerun baseline accidentally used absolute position attention while the μTransfer model used relative attention, and the μTransfer model diverged in FP16 and was trained in FP32, which the authors attribute to more aggressive transferred hyperparameters.

Appendix I adds a use the authors did not plan for. Because μP ties a narrow model's behavior to a wide one's, a large run's instability can be reproduced cheaply by transferring its hyperparameters backwards to a narrow model. In one case a width-8192 Transformer that diverged mid-training was reproduced at width 256, where the cause (attention logits blowing up) was found much faster.

What does not transfer

The paper lists its own gaps. Depth transfer worked only for pre-layernorm Transformers; post-layernorm models transferred across width but not depth. The best initialization scale did not transfer across depth even in pre-layernorm models, so for BERT the authors held it fixed and tuned everything else. The optimum still moves slightly at small width, which the authors suggest finite-width corrections to μP might fix. Regularization settings for fine-tuning on small datasets are outside the method.

Appendix D collects practical conditions under which transfer degrades. Small attention heads make the proxy's hyperparameter landscape noisy: at width multiplier 0.0625 with 4 heads, $d_k = 8$, and the attention multiplier did not transfer until $d_k$ was pegged at a minimum (32 generally works). Squashing activations such as tanh saturate more in narrow models and bias the proxy, so the authors recommend ReLU-like activations. Zero-initializing the output layer and the query projection removes a difference between narrow and wide models at initialization, where the network's initial random output has variance of order $1/\text{width}$ and so is much larger in the proxy.

The theory has limits too. The claim that μP is the only parametrization that allows transfer across width is argued in Appendix J.3 from the earlier paper's classification of parametrizations: any other stable rule either reaches a kernel limit, where the network stops learning features and hyperparameters stop mattering, or leaves some layer's updates vanishingly small at large width, so its learning rate has no effect in the limit. Everett et al. (2024) later reported that other parametrizations, standard included, also transfer once each layer's learning rate is scaled correctly, and that Adam's $\epsilon$ must scale with width too, which Appendix B.3 of this paper already notes. Depth got its own treatment in Tensor Programs VI (Depth-μP), which adds a depth-dependent multiplier on residual branches.

In the paper's own runs, the learning rate, initialization and multipliers of a model family were tuned once on a proxy that fits on one GPU: one 13M proxy for both BERT sizes, and a 40M proxy for GPT-3 6.7B at 7% of its training cost.