The three things muP keeps under control

muP is defined by three desiderata, one for each quantity that could blow up or vanish as the width n grows. First, every preactivation coordinate should be Θ(1) — neither growing nor shrinking with n. Second, the network’s output (the logits) should be O(1). Third, every layer should be updated as much as possible without diverging — the ‘maximal’ in the name.

The first two are familiar: ordinary initialization schemes already arrange for sensible activations and outputs at initialization. But they say nothing about what happens after a gradient step, and that is where wide models silently misbehave once training starts. muP’s contribution is enforcing all three conditions at once and throughout training, so that behaviour at width 256 genuinely predicts behaviour at width 16k.

Advertisement

The abc-parametrization

muP is expressed through three exponents attached to every weight tensor. Write the effective weight used in the forward pass as W = n^(-a) · w, initialize the trainable part w with variance n^(-2b), and give it a learning rate η · n^(-c), where n is the width. The triple (a, b, c), chosen per layer, is the parametrization.

W_eff = n^(-a) * w        # forward multiplier
w     ~ Normal(0, n^(-2b))  # initialization variance
lr(w) = η * n^(-c)         # per-layer learning rate

Standard parametrization is one choice of (a, b, c); the NTK parametrization is another; muP is the specific choice that satisfies all three desiderata simultaneously. Nothing here is exotic: pick an exponent for the multiplier, one for the init, and one for the learning rate, separately for each kind of layer. Everything below just reads off those exponents for the input, hidden, and output layers.

Advertisement

Three kinds of layer, three sets of exponents

muP treats a weight matrix differently depending on which of its dimensions scale with width. Input / embedding weights have a fixed fan-in (the vocabulary or input dimension) but a fan-out that grows like n. Hidden weights have both fan-in and fan-out growing like n — these are the bulk of a transformer’s parameters. Output / readout weights have a fan-in that grows like n but a fixed fan-out (the vocabulary or class count).

The distinction matters because a sum over a width-n dimension accumulates √n worth of noise, while a sum over a fixed dimension does not — so the three families genuinely need different initialization and learning-rate scaling. Treating them all the same is exactly the habit that makes hyperparameters shift when you change width, and exactly what muP corrects by handling each family on its own terms.

The initialization rule and its quiet surprise

For hidden weights, muP initializes with variance Θ(1/fan_in), i.e. variance ∝ 1/n. That is the quiet surprise: it is exactly ordinary Kaiming / fan-in initialization, the same rule competent implementations already use. muP does not invent a new hidden-layer init at all.

Input / embedding weights keep Θ(1) variance — their fan-in is fixed, so there is nothing to shrink against. The readout is the real exception: muP scales its output down by a 1/fan_in multiplier (the ‘MuReadout’ treatment) and often initializes it small, or even to zero, so that at the start of training the logits are dominated by learned features rather than random initialization noise. So the initialization story is almost entirely standard: muP’s real signature lives not in the init but in the learning rate, the half ordinary training gets wrong.

Why activations stay bounded: the central-limit argument

Take a hidden preactivation z_i = Σ_j W_ij x_j, a sum of n = fan_in terms. If the inputs x_j are Θ(1) and the weights W_ij are independent with variance 1/n, each product has variance about 1/n, and the n independent terms sum to a variance of n · (1/n) = 1. So z_i is Θ(1) regardless of width: the √n growth from adding more terms is cancelled exactly by the 1/√n shrink built into each weight.

A worked number: at n = 1024 with unit inputs and weight std 1/√1024 ≈ 0.031, the 1024-term sum has std ≈ √1024 × 0.031 ≈ 1.0; double the width to 2048 and it is still ≈ 1.0. That width-invariance at initialization is desideratum one, and precisely why the fan-in rule is not arbitrary.

The learning-rate rule

Initialization keeps activations sane at step zero; the learning rate keeps them sane at every step after. Under Adam — what transformers actually train with — muP scales the hidden-layer learning rate by 1/fan_in, i.e. η ∝ 1/n. The reference mup library implements this literally as globalLR / (fan_in / base_fan_in).

Input and embedding layers keep a constant Θ(1) learning rate, and the readout’s learning rate is scaled down with width too. (Under plain SGD the pattern differs, but Adam is the case that matters for language models.) One transformer-specific twist: muP scales attention logits by 1/d rather than the usual 1/√d, because the query–key dot product is itself a width-n sum.

The maximal part: keeping updates order-one

Why 1/n on the hidden learning rate, and not 1/√n or a constant? Because of the third desideratum. A single Adam step changes a weight by roughly the learning rate times a Θ(1) normalized gradient. When that weight change propagates through a width-n sum into the next layer’s preactivation, its naive contribution scales with n. To keep the change in each preactivation Θ(1) — a real, order-one update to the features, not a vanishing or exploding one — the learning rate must absorb a factor of 1/n.

‘Maximal’ means this is the largest learning-rate scaling that still keeps updates bounded. Any larger and the activation changes grow with width until training destabilizes; any smaller and the features barely move. muP sits exactly on that knife-edge — the biggest stable step — for every layer at once.