class MLP(nnx.Module):
"""A three-hidden-layer ReLU network under standard parametrization."""
def __init__(self, width, rngs=None):
rngs = nnx.Rngs(0) if rngs is None else rngs
init = nnx.initializers.variance_scaling(1.0, 'fan_in', 'normal')
self.fc_in = nnx.Linear(784, width, kernel_init=init, rngs=rngs)
self.fc_h1 = nnx.Linear(width, width, kernel_init=init, rngs=rngs)
self.fc_h2 = nnx.Linear(width, width, kernel_init=init, rngs=rngs)
self.fc_out = nnx.Linear(width, 10, kernel_init=init, rngs=rngs)
def features(self, X):
h = nnx.relu(self.fc_h1(nnx.relu(self.fc_in(X))))
return nnx.relu(self.fc_h2(h))
def __call__(self, X):
return self.fc_out(self.features(X))
def configure_adam(self, lr):
return nnx.Optimizer(self, optax.adam(lr), wrt=nnx.Param)The sweep: the optimum moves
Eight learning rates × four widths, 32 runs, about a minute:
widths = [128, 256, 512, 1024]
lrs = [2 ** k for k in range(-12, -4)]
sp_loss = {w: [train_mlp(MLP, w, lr) for lr in lrs] for w in widths}
for w in widths:
print(f'width {w:4d}: best lr {min(zip(sp_loss[w], lrs))[1]:.1e}')
d2l.plot(lrs, [sp_loss[w] for w in widths], 'learning rate',
'training loss', xscale='log', ylim=[0.25, 0.8],
legend=[f'width {w}' for w in widths])
width 128: best lr 3.9e-03
width 256: best lr 2.0e-03
width 512: best lr 2.0e-03
width 1024: best lr 9.8e-04- Every width: a U. The minima do not line up — the best learning rate falls about 8× as width grows 8×.
- Wider is never worse at its own optimum; the narrow model’s optimum sits on the wide model’s unstable branch.
- The best learning rate is a property of the model size, not the task.
The coordinate check
Mean |activation| per layer, one Adam step, widths 128 → 4,096. Correct scaling = flat curves. Under SP:
check_widths = [128, 256, 512, 1024, 2048, 4096]
sp_acts = coord_check(MLP, check_widths)
d2l.plot(check_widths, list(sp_acts), 'width', 'mean |activation|',
xscale='log', yscale='log',
legend=['layer 1', 'layer 2', 'layer 3', 'logits'])
Fixed fan-in layer flat; the layers behind square matrices grow, compounding with depth; logits blow up ~100× — the first step already writes the instability.
The coordinate check, under muP
mup_acts = coord_check(MuMLP, check_widths)
d2l.plot(check_widths, list(mup_acts), 'width', 'mean |activation|',
xscale='log', yscale='log',
legend=['layer 1', 'layer 2', 'layer 3', 'logits'])
- All layer curves flat: no width-dependent update scale left.
- Logits fall by design (the 1/m multiplier); they grow to O(1) by learning, not by size. Growth is the unambiguous failure — the check reads activation size, not feature learning.
Run this before trusting any muP integration — it catches missed multipliers, mislabeled layers, framework defaults sneaking back.
Learning-rate transfer
The payoff: the same sweep, under muP:
mup_loss = {w: [train_mlp(MuMLP, w, lr) for lr in lrs] for w in widths}
for w in widths:
print(f'width {w:4d}: best lr {min(zip(mup_loss[w], lrs))[1]:.1e}')
d2l.plot(lrs, [mup_loss[w] for w in widths], 'learning rate',
'training loss', xscale='log', ylim=[0.25, 0.8],
legend=[f'width {w}' for w in widths])
width 128: best lr 3.9e-03
width 256: best lr 7.8e-03
width 512: best lr 7.8e-03
width 1024: best lr 7.8e-03- The optimum stops sliding: scatter within a grid step or two across 8× width (SP: slid three). Reusing the base optimum at width 1,024: within a couple % of that width’s best loss (SP: misses by 15–20%).
- Tune small, ship big: TP5 tuned GPT-3 6.7B from a 40M proxy at ~7% of pretraining cost — and beat the original.
- Caveat: at a width you can sweep, retuned SP matches; muP buys the sweep you cannot afford.