Scaling Up

Dive into Deep Learning · §9.11

Transferring small-scale tuning to larger models
the drifting optimum · muP · the coordinate check · reported large-scale approaches

Hyperparameter transfer to large models

Sweeps are practical when each run takes seconds.

  • Target-scale runs can require thousands of accelerators for months, making a full hyperparameter grid infeasible.
  • A common approach is to tune a smaller proxy and transfer the settings.
  • We examine width-dependent failure under standard parametrization, muP’s scaling rules, and diagnostics for transfer.

A family of widths

Fashion-MNIST MLP, 784 \to n \to n \to n \to 10, width n from 128 to 1,024. Only the two middle matrices are n \times n: that is where scale lives. Initialization pinned explicitly: variance 1/\text{fan-in}, biases zero — standard parametrization (SP).

class MLP(nn.Module):
    """A three-hidden-layer ReLU network under standard parametrization."""
    def __init__(self, width):
        super().__init__()
        self.fc_in = nn.Linear(784, width)
        self.fc_h1 = nn.Linear(width, width)
        self.fc_h2 = nn.Linear(width, width)
        self.fc_out = nn.Linear(width, 10)
        for lin in (self.fc_in, self.fc_h1, self.fc_h2, self.fc_out):
            nn.init.normal_(lin.weight, std=lin.in_features ** -0.5)
            nn.init.zeros_(lin.bias)

    def features(self, X):
        h = F.relu(self.fc_h1(F.relu(self.fc_in(X))))
        return F.relu(self.fc_h2(h))

    def forward(self, X):
        return self.fc_out(self.features(X))

    def configure_adam(self, lr):
        return torch.optim.Adam(self.parameters(), lr)

400 Adam steps at batch 512; score = final loss on the whole training set.

Learning rate as width changes

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 7.8e-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×.
  • In these fixed-seed runs, tuned wider models match or improve on the narrower models, but the width-128 optimum destabilizes width 1,024.
  • Under standard parametrization, the preferred learning rate depends on model width as well as on the task.

Dependence of the first update on fan-in

Single-example gradient of a hidden matrix = outer product \mathbf{g} = \boldsymbol{\delta}\mathbf{h}^\top; Adam’s first update is a sign step, and signs of outer products factorize:

(\Delta\mathbf{W}\mathbf{h})_i = -\eta \operatorname{sign}(\delta_i) \sum_{j=1}^n |h_j| \approx -\eta\, n\, \overline{|h|}.

All n terms add coherently in this single-example first-step calculation: doubling width doubles the activation change and gives an inverse-width stability scale for \eta.

Per layer: input weights (fan-in 784, fixed) and biases don’t scale. A global \eta must accommodate layers with different width scaling.

muP: the rules (Adam, width multiplier m)

Yang & Hu et al., 2022 — Tensor Programs V

parameters init Adam LR forward
input weights, biases unchanged \eta unchanged
hidden matrices \propto 1/\text{fan-in} (unchanged) \eta/m unchanged
output matrix unchanged \eta logits \times 1/m

Derived from the infinite-width limit where activations stay O(1) and every layer keeps learning — the maximal update. Embeddings count as input-like; attention needs 1/d; SGD has its own column.

Implementation of the muP rules

class MuMLP(MLP):
    """The same network under muP, relative to a width-128 base."""
    def __init__(self, width, base_width=128):
        super().__init__(width)
        self.m = width / base_width

    def forward(self, X):
        # rule 1: scale the output matrix's logits; the bias is untouched
        return self.fc_out(self.features(X) / self.m)

    def configure_adam(self, lr):
        hidden = [self.fc_h1.weight, self.fc_h2.weight]
        rest = [p for p in self.parameters()
                if not any(p is q for q in hidden)]
        return torch.optim.Adam([              # rule 2: hidden LR / m
            {'params': rest, 'lr': lr},
            {'params': hidden, 'lr': lr / self.m}])

At the base width m=1: muP changes nothing about the model you tune.

The coordinate check

We measure mean |activation| per layer after one Adam step, across widths 128 → 4,096. Under standard parametrization:

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 grow by roughly 100×, exposing the first-step width dependence.

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.

This diagnostic can reveal missing multipliers, mislabeled layers, and unintentionally restored framework defaults.

Learning-rate transfer

We repeat the learning-rate 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 7.8e-03
width  256: best lr 7.8e-03
width  512: best lr 7.8e-03
width 1024: best lr 7.8e-03
  • The grid minima remain within one or two steps across 8× width; under SP, they moved three steps. At width 1,024, reusing the base selection finishes within a few percent of the grid minimum (SP: 15–20% above it).
  • Tensor Programs V reports tuning GPT-3 6.7B from a 40M proxy at about 7% of the pretraining cost and improving on the original model.
  • At a width that can be swept directly, retuned standard parametrization matches muP; muP avoids repeating that sweep at large scale.

The spectral view

Healthy layer scale (:numref:sec_muon): \|\mathbf{W}\|_2 \asymp \|\Delta\mathbf{W}\|_2 \asymp \sqrt{n_{\text{out}}/n_{\text{in}}}.

  • muP \equiv this spectral condition (Yang–Simon–Bernstein, 2023): the per-layer LRs make Adam’s updates land at the right spectral scale.
  • muP rescales Adam’s coordinatewise updates to meet this spectral condition; Muon directly controls matrix-update shape, although its RMS-matching convention differs from \sqrt{n_{\text{out}}/n_{\text{in}}}.

Hyperparameter transfer in large runs

  • Cerebras: muP in production; family tuned from a ~40M proxy.
  • DeepSeek: fit the observed drift with power laws for the preferred learning rate and batch size versus compute, then extrapolate.
  • Meta: “MetaP” per-layer LRs/init for Llama 4 (undisclosed).
  • Moonshot / Kimi K2: no parametrization; every Muon update scaled to an RMS matched empirically to AdamW’s. 15.5T tokens.
  • Live debate: weight decay, not muP, may drive transfer in long runs (Kosson et al., 2025). muP: one mechanism, not settled law.

Recap

  • Best LR drifts with width: one Adam step perturbs a layer \propto fan-in.
  • muP: init unchanged, hidden LR \eta/m, logits \times 1/m — optimum transfers from the base width.
  • Coordinate check: verify that activation scales follow the prescribed width dependence without unintended growth.
  • Spectral view ties muP to Muon; labs mix parametrization, scaling-law fits, and matched update sizes. Shared goal: tune small, transfer big.