Maximal Update Parametrization (muP): How Tensor Programs Enable Zero-Shot Hyperparameter Transfer in LLM Pre-Training

Pre-training a frontier large language model requires hundreds of thousands of GPU hours and millions of dollars in compute. At that scale, traditional hyperparameter tuning is financially and operationally impossible: teams cannot sweep learning rates, weight initializations, or optimizer betas across multiple 70B parameter runs to find the loss minimum. Historically, practitioners relied on ad-hoc heuristic extrapolation or manual guesses from small runs, often leading to sub-optimal loss curves or catastrophic mid-training loss spikes.

Maximal Update Parametrization (muP), formulated by Greg Yang and collaborators across Microsoft Research and OpenAI in the Tensor Programs V paper, solves this scaling bottleneck. By deriving the infinite-width limit of neural network computation, muP guarantees that representations across every layer update at a consistent, non-vanishing, and non-exploding scale as width increases. This mathematical framework enables zero-shot hyperparameter transfer (muTransfer): engineering teams can optimize learning rates, initialization scales, and optimizer parameters on a small 10M to 100M parameter proxy model and directly apply those exact values to billion-scale pre-training runs without tuning at scale.

Maximal Update Parametrization and Zero-Shot Hyperparameter Transfer Architecture

The Hyperparameter Scaling Failure in Standard Parametrization

In the Standard Parametrization (SP) used by default across deep learning libraries like PyTorch, weight tensors are initialized according to He/Kaiming or LeCun initialization rules (scaling weight variance inversely with input dimension, $\sigma^2 \propto 1/d_{in}$ ), and all layers share a uniform learning rate $\eta$ .

While this setup behaves stably at fixed network dimensions, it fails when scaling model width $d_{model} \to \infty$ . Under Standard Parametrization, the dynamics of forward activations and backward gradient updates scale inconsistently across different layers:

Early layers (such as token embeddings) receive gradient signals whose update magnitude shrinks as model width grows, effectively freezing feature representations in initial layers.
Late layers (such as hidden linear projections and output heads) experience activation variance and update magnitudes that scale with $O(\sqrt{d})$ or $O(d)$ , causing activations to blow up and gradients to explode unless the global learning rate is aggressively reduced.
If the practitioner reduces the global learning rate to prevent divergence (for example, scaling $\eta \propto 1/d$ ), the hidden layers update too slowly, pushing the network into the "lazy training" or Neural Tangent Kernel (NTK) regime where features do not adapt and the model behaves as a linear kernel machine over fixed representations.

Because the optimal learning rate shifts unpredictably with width in Standard Parametrization, a learning rate tuned on a 125M parameter model will destabilize or severely underfit a 13B or 70B parameter model of the same architecture family.

Theoretical Foundations: Tensor Programs and the Maximal Feature Learning Limit

To resolve the divergence between width scaling and feature learning, Greg Yang introduced the theoretical framework of Tensor Programs. By expressing any neural network computation, forward propagation, loss computation, and backpropagation as a structured sequence of matrix multiplications and coordinatewise nonlinearities, Tensor Programs allow researchers to compute the exact asymptotic limits of all intermediate representations as width tends to infinity.

Tensor Programs proved that there is a unique parameterization that maximizes feature learning across every layer without causing representation collapse or gradient explosion. This parameterization is Maximal Update Parametrization (muP).

Under muP, the parameter update $\Delta W$ from optimization interacts with the layer input $x$ such that the update to post-activations $\Delta h = (\Delta W) x$ remains $\Theta(1)$ with respect to width for every layer simultaneously.

By enforcing this condition:

Every layer learns representations at an equal, non-trivial rate throughout training.
Forward activations and backward adjoints maintain stable $O(1)$ variances across infinite width.
The loss landscape geometry remains invariant to width, ensuring that the optimal learning rate $\eta^$ on a narrow network is identical to the optimal learning rate $\eta^$ on a wide network.

The muP Scaling Rules Across Transformer Layers

Achieving maximal update scaling requires adjusting initialization variance, learning rate scaling, and output multiplier constants differently depending on the role of each matrix in the Transformer architecture.

The core scaling rules compare Standard Parametrization (SP) and Maximal Update Parametrization (muP) across a network scaled from a base width $d_{base}$ to target width $d$ :

1. Input Embeddings

Standard Parametrization: Weights initialized as $\mathcal{N}(0, 1)$ , learning rate scaled as $\Theta(1)$ .
Maximal Update Parametrization: Weights initialized as $\mathcal{N}(0, 1)$ , learning rate scaled as $\Theta(1)$ (no width penalty).

Because the embedding table maps discrete one-hot token indices to continuous vectors, the input dimension is fixed (vocabulary size), while the output dimension is width $d$ . An $O(1)$ learning rate preserves constant coordinate variance.

2. Hidden Linear Layers (Attention Projections and MLP Weights)

Standard Parametrization: Weights initialized as $\mathcal{N}(0, 1/d_{in})$ , learning rate scaled as $\Theta(1)$ .
Maximal Update Parametrization: Weights initialized as $\mathcal{N}(0, 1/d_{in})$ , learning rate scaled as $\Theta(1/d_{in})$ via a parameter multiplier $(d_{base} / d)$ .

In hidden matrix multiplications where both input dimension $d_{in}$ and output dimension $d_{out}$ scale with width, standard gradient descent updates the weights by an amount proportional to $d$ . To keep the activation change $\Delta h = (\Delta W) x$ at $\Theta(1)$ , the effective learning rate must scale inversely with width.

3. Attention Logits

Standard Parametrization: Dot-product attention scaled by $1/\sqrt{d_{head}}$ .
Maximal Update Parametrization: Dot-product attention scaled by $1/d_{head}$ (or an explicit $1/d$ attention multiplier).

In standard scaled dot-product attention, query and key vectors of dimension $d_{head}$ produce dot products whose variance grows linearly with dimension, leading practitioners to divide by $\sqrt{d_{head}}$ . However, under maximal feature learning, query and key coordinates become correlated during training, which causes the dot product to scale as $O(d_{head})$ rather than $O(\sqrt{d_{head}})$ . Dividing by $d_{head}$ prevents the softmax distribution from degenerating into a hard one-hot argmax or suffering entropy collapse at large head dimensions.

4. Output Unembedding Head (lm_head)

Standard Parametrization: Weights initialized as $\mathcal{N}(0, 1/d_{in})$ , learning rate scaled as $\Theta(1)$ .
Maximal Update Parametrization: Output projection scaled by an explicit multiplier $1/d_{in}$ , weights initialized with variance $\sigma^2 \propto 1/d_{in}^2$ (or zero initialization), and learning rate scaled as $\Theta(1/d_{in})$ .

The output head maps hidden vectors of dimension $d$ back to the fixed vocabulary dimension. Without an explicit $1/d$ scaling factor, the logits would scale as $O(d)$ , saturating the cross-entropy loss and halting gradient flow early in training.

Practical Implementation: The muTransfer Workflow

The operational power of muP lies in zero-shot hyperparameter transfer (muTransfer). In practice, the pre-training workflow proceeds through distinct stages:

import torch
import torch.nn as nn

class MuPLinear(nn.Module):
    def __init__(self, in_features, out_features, base_width=256, is_output_head=False):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.base_width = base_width
        self.is_output_head = is_output_head
        
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        self.bias = nn.Parameter(torch.zeros(out_features))
        self.reset_parameters()
        
    def reset_parameters(self):
        if self.is_output_head:
            # Output head initialized with 1/d^2 variance or zero
            nn.init.normal_(self.weight, std=1.0 / self.in_features)
        else:
            # Hidden weights initialized with standard 1/sqrt(d_in)
            nn.init.normal_(self.weight, std=1.0 / (self.in_features ** 0.5))
            
    def forward(self, x):
        if self.is_output_head:
            # Explicit 1/d scaling on output logits
            width_mult = self.base_width / self.in_features
            return nn.functional.linear(x, self.weight * width_mult, self.bias)
        else:
            return nn.functional.linear(x, self.weight, self.bias)

In production pipelines, developers utilize the official Microsoft mup open-source library, which wraps model definitions to calculate layer-specific width multipliers and assigns separate learning rate scaling factors into optimizer parameter groups:

Base Model Definition: Configure the target Transformer architecture with a compact base width (e.g. $d_{base} = 256$ or $512$ , around 13M to 100M parameters).
Proxy Sweep: Perform grid search or Bayesian optimization on the small proxy model across learning rate (sweeping over 5 to 6 orders of magnitude), warmup steps, Adam $\beta_1$ and $\beta_2$ , weight decay, and gradient clipping thresholds.
Coordinate Checks: Run a few hundred training steps on multiple intermediate widths (e.g., 256, 512, 1024, 2048) using mup.coord_check to verify that activation update norms remain flat across width.
Full-Scale Transfer: Apply the identical optimal hyperparameters discovered on the small proxy directly to the 7B, 13B, or 70B+ target model.

Empirical Validation Across Frontier Models

Empirical studies have confirmed that muP yields consistent performance gains and substantial compute savings:

Cerebras-GPT: Cerebras trained a full family of open models ranging from 111M to 13B parameters using muP. Hyperparameters were tuned solely on the 111M model and transferred zero-shot to all larger models, achieving Pareto-optimal pre-training compute curves without hyperparameter tuning runs on the 13B checkpoint.
Microsoft Phi Series: Models including Phi-1, Phi-2, and Phi-3 leveraged muP scaling to ensure stability and optimal learning dynamics during intensive synthetic data training.
Compute Efficiency: Across extensive benchmarks in the original NeurIPS 2022 paper, muTransfer reduced total pre-training compute budgets by an estimated 15% to 30% by eliminating exploratory scaling sweeps and converging to a lower final validation loss compared to Standard Parametrization models tuned with heuristic learning rates.

Depth Scaling and Modern Architectural Extensions

While the original muP formulation focused on horizontal width scaling ( $d_{model}$ , MLP dimension, attention heads), subsequent research addressed vertical depth scaling:

Depth-muP and Residual Scaling: As detailed in Tensor Programs VI, scaling depth $L \to \infty$ introduces vanishing gradient and representation collapse issues if residual branches are unscaled. Implementing depth-aware parameterization requires scaling residual branch additions by $1/\sqrt{2L}$ (or using DeepNorm / Pre-LayerNorm scaling) to ensure that deeper layers contribute equally to the residual stream.
Compatibility with Modern Primitives: muP integrates cleanly with modern Transformer components, including Rotary Position Embeddings (RoPE), RMSNorm, SwiGLU activations, and Grouped-Query Attention (GQA). Because RMSNorm normalizes activation norms per token, it complements muP by maintaining layer input variance, while muP governs the parameter learning rates and update dynamics.