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(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)

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 3.9e-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, rngs=None):
        super().__init__(width, rngs=rngs)
        self.m = width / base_width

    def __call__(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):
        def labels(params):                   # rule 2: hidden LR / m
            return jax.tree_util.tree_map_with_path(
                lambda path, _: 'hidden'
                if 'fc_h' in jax.tree_util.keystr(path)
                and 'kernel' in jax.tree_util.keystr(path) else 'rest',
                params)
        tx = optax.multi_transform(
            {'hidden': optax.adam(lr / self.m), 'rest': optax.adam(lr)},
            labels)
        return nnx.Optimizer(self, tx, wrt=nnx.Param)

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.