Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Hyperparameter Selection › Hyperparameter Transfer

Michael BrenndoerferJanuary 29, 20269 min read

Language AI HandbookAll chapters

Part 32 of 67Training Optimization

Chapter 238Hyperparameter Selection › Hyperparameter Transfer

Reading section 5 of 14 · Hyperparameter Transfer · 9 minSection 5 of 14 · 9 min

Section 5 of 14 · 9 min read

Hyperparameter Selection › Hyperparameter TransferOne section · Sections in this chapter5 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Even efficient search strategies are costly when experiments take days. Hyperparameter transfer sidesteps the problem by asking: can we find good hyperparameters on a small model and use them for a large model without re-searching?

The intuitive hope is that the optimal learning rate for a 125M parameter model should tell us something about the optimal learning rate for a 1.3B parameter model. This hope is sometimes true and sometimes badly wrong, depending on the hyperparameter and the scaling strategy.

The Transfer ProblemLink Copied

Standard training wisdom has been that hyperparameters do not transfer cleanly across model scales. The commonly observed phenomenon is that larger models require smaller learning rates. Practitioners scaling from 125M to 1.3B to 6.7B parameters would reduce the learning rate at each step, often by trial and error. This made scaling expensive: every new model size required its own hyperparameter search.

The root cause is that standard parameter initialization does not have stable statistics across width. When you double the width of a layer (its hidden dimension), the total gradient flowing through that layer changes in a way that depends on the initialization. If initialization variance does not scale correctly with width, the effective learning rate at a neuron level changes even when the optimizer's nominal learning rate stays the same.

To see why, think about what happens when a layer receives a gradient and updates its weights. The parameter update changes the layer's outputs, and the size of that output change determines how much the model changes at each step. If the output change is too large relative to the current activations, training becomes unstable. If it is too small, learning is inefficiently slow. Maintaining a stable ratio between the update magnitude and the activation magnitude is what makes the learning rate transferable.

Maximal Update Parametrization (muP)Link Copied

Maximal Update Parametrization (muP), introduced by Yang et al. (2022), addresses this directly. It is a specific set of initialization scales and learning rate scaling rules that makes feature updates approximately constant as model width increases. The key insight is:

Maximal Update Parametrization (muP)

muP is a reparametrization of neural network weights that ensures the optimal learning rate is the same across model widths. With standard parametrization (SP), the optimal learning rate decreases as width grows. With muP, the optimal learning rate is width-independent, enabling hyperparameter transfer from small to large models.

To understand why this matters, consider a single linear layer with input dimension $d_{in}$ , output dimension $d_{out}$ , and weight matrix $W \in \mathbb{R}^{d_{out} \times d_{in}}$ . In standard Xavier initialization, each weight is drawn from:

W_{ij} \sim \mathcal{N}\left(0, \frac{2}{d_{in} + d_{out}}\right)

where:

$W_{ij}$ : the $(i,j)$ entry of the weight matrix
$d_{in}$ : input dimension of the layer
$d_{out}$ : output dimension of the layer

When a gradient update of magnitude $g$ is applied, the change in the output activations is approximately:

\Delta \mathbf{h} \approx \eta \cdot g \cdot \sqrt{d_{in}}

where:

$\Delta \mathbf{h}$ : change in output activations
$\eta$ : learning rate
$g$ : gradient magnitude per weight

As $d_{in}$ grows, the effective activation update magnitude grows as $\sqrt{d_{in}}$ , even at the same nominal learning rate. To keep activation updates stable, the learning rate must shrink as $1/\sqrt{d_{in}}$ . This is the source of the need to retune learning rates when changing width.

muP resolves this by modifying the initialization and the effective learning rate applied to each parameter type:

Input layer weights: Scale initialization by $1/d_{in}$ , use learning rate $\eta$
Hidden layer weights: Scale initialization by $1/\sqrt{d_{in}}$ , scale learning rate by $1/d_{width}$ where $d_{width}$ is the hidden dimension
Output layer weights: Scale initialization by $1/d_{out}$ , use a smaller learning rate $\eta / d_{width}$

Under muP, the contribution of each weight to the output remains constant in expectation as width grows, and the optimal learning rate becomes width-invariant. This means:

Tune hyperparameters (especially learning rate) on a small proxy model using muP
Scale up to the target model width using muP
Use the same learning rate without retuning

In[7]:

Code

import numpy as np

# Demonstrate how initialization variance changes with width under SP vs muP


def sp_init_std(d_in, d_out):
    """Standard Xavier initialization std."""
    return np.sqrt(2.0 / (d_in + d_out))


def mup_hidden_init_std(d_in, d_out):
    """muP hidden layer initialization: 1/sqrt(d_in)."""
    return 1.0 / np.sqrt(d_in)


def effective_activation_change_sp(d_in, d_out, lr=1e-3, grad_magnitude=1.0):
    """
    Approximate magnitude of activation change per step under SP.
    This simplified scaling model isolates the sqrt(d_in) growth described
    above; it is not a full optimizer simulation.
    """
    return lr * grad_magnitude * np.sqrt(d_in)


def effective_activation_change_mup(
    d_in, d_out, base_width, lr=1e-3, grad_magnitude=1.0
):
    """
    Approximate width-invariant activation change under muP, anchored to the
    same value as SP at base_width.
    """
    return lr * grad_magnitude * np.sqrt(base_width)


# Vary width from 128 to 4096
widths = [128, 256, 512, 1024, 2048, 4096]
base_width = 128
lr = 1e-3

sp_changes = [effective_activation_change_sp(w, w, lr=lr) for w in widths]
mup_changes = [
    effective_activation_change_mup(w, w, base_width, lr=lr) for w in widths
]

Out[8]:

Console

   Width   SP activation change   muP activation change
--------------------------------------------------------
     128               0.011314                0.011314
     256               0.016000                0.011314
     512               0.022627                0.011314
    1024               0.032000                0.011314
    2048               0.045255                0.011314
    4096               0.064000                0.011314

SP ratio (width 4096 vs 128): 5.66x
muP ratio (width 4096 vs 128): 1.00x

Under standard parametrization, the effective activation change grows with width even at the same learning rate. Under muP, it stays constant regardless of width, which is why the same learning rate works across scales.

Out[9]:

Visualization

Line plot on a logarithmic width axis. The standard-parametrization curve rises from below one to about 5.7, while the muP curve remains flat at one. — A simplified scaling model compares activation updates relative to width 128. Standard parametrization grows with the square root of width, reaching 5.7 times the baseline at width 4096, while maximal update parametrization (muP) remains constant by construction.

A Worked muP Transfer ExampleLink Copied

To make the muP workflow concrete, consider a team training a 1.3B parameter model and wanting to find the optimal learning rate efficiently. Without muP, they would need to run several training experiments directly at 1.3B scale, each costing substantial GPU time. With muP, they can work as follows.

First, they select a proxy width of 256 (versus the target hidden dimension of 2048, a ratio of 8). They train the proxy for 5,000 steps (proportional to the proxy's compute-optimal budget) with learning rates spanning $\{3 \times 10^{-4},\ 6 \times 10^{-4},\ 1 \times 10^{-3},\ 2 \times 10^{-3},\ 4 \times 10^{-3}\}$ . They initialize using muP rules: hidden layer weights scaled by $1/\sqrt{d_{in}}$ , output layer weights scaled by $1/d_{width}$ .

After the proxy sweep, the optimal learning rate at width 256 is found to be $1 \times 10^{-3}$ . Under muP, this is the same optimal learning rate that should work at width 2048. They set the 1.3B model's learning rate to $1 \times 10^{-3}$ and begin the full training run. The first few thousand steps confirm that loss decreases smoothly, consistent with the proxy's behavior. No retuning at the 1.3B scale was needed.

Without muP, scaling from width 256 to 2048 (an 8x increase) under standard Xavier initialization would require reducing the learning rate by roughly $\sqrt{8} \approx 2.8 \times$ to account for the activation update scaling. The team would either guess this scaling (risking a suboptimal learning rate) or run additional full-scale experiments (expensive). muP eliminates this uncertainty.

What Transfers and What Does NotLink Copied

Hyperparameter transfer via muP works primarily for learning rate across model width. Other hyperparameters have different transfer properties, and understanding these differences prevents false confidence about which settings you can simply copy from a proxy run.

Learning rate (with muP): transfers almost perfectly across width. Transfer across depth is less reliable but generally acceptable within a 2x depth range. If you double both width and depth, transfer the learning rate for the width change and monitor carefully for the depth change.
Batch size: transfers well when adjusting learning rate proportionally. Two scaling rules are in common use. The linear scaling rule (double batch size, double learning rate) holds in the data-rich, short-horizon regime and was empirically validated for ImageNet training by Goyal et al. (2017). The square root scaling rule (double batch size, multiply learning rate by $\sqrt{2}$ ) is more conservative and tends to be safer for very large batch sizes or longer training runs where noise reduction from larger batches matters more.
Warmup steps: does not transfer directly as an absolute count. A 1000-step warmup for a 10,000-step proxy run represents 10% warmup, far too long for a 100,000-step target run. Express warmup as a fraction of total steps (e.g., 2% warmup) so it transfers naturally.
Weight decay: generally transfers well. The optimal weight decay coefficient is relatively insensitive to model width and moderately insensitive to scale. It is more sensitive to dataset size: when fine-tuning or domain-adapting on a small dataset, the optimal weight decay is often lower than the pretraining default.
Beta parameters ( $\beta_1$ , $\beta_2$ for Adam): very robust; near-universal defaults of $0.9$ and values between $0.95$ and $0.999$ apply almost everywhere.
Gradient clipping threshold: transfers moderately. Larger models can produce larger gradient norms during early training, especially without muP. The gradient norm typically stabilizes after warmup; monitoring it during the first few thousand steps reveals whether your clipping threshold is too aggressive or too permissive.

The practical implication is that for a new model, you should:

Use muP if you plan to search learning rate (it is increasingly the default in research settings)
Search learning rate and weight decay on a proxy model at 1/10th to 1/100th the target parameter count
Use principled scaling rules for batch size
Keep beta parameters at universal defaults
Monitor gradient norms to set a clipping threshold