Unit-Scaled μP: Unified Scaling for Neural Nets

• Unit-Scaled μP is a neural network parametrization scheme that synthesizes Maximal Update Parametrization and Unit Scaling to ensure width-invariant hyperparameters and stable low-precision training.
• It defines precise scaling laws using the abc-parametrization framework, enabling consistent feature learning and effortless hyperparameter transfer across model sizes via μTransfer.
• Empirical and theoretical analyses demonstrate that u-μP achieves faster convergence, enhanced numerical stability, and effective function-preserving upscaling in deep learning models.

Unit-scaled μP (u-μP) is a neural network parametrization scheme synthesizing Maximal Update Parametrization (μP) and Unit Scaling. It enforces both width-invariant learning-rate scaling and unit variance for all tensors at initialization. This ensures that optimal hyperparameters, especially learning rate, transfer across model sizes with no manual rescaling, while simultaneously rendering the entire model robust to low-precision numerics (e.g., FP8). u-μP theory underpins principled model upscaling and enables consistent feature learning in the infinite-width limit, overcoming key pathologies of earlier parametrizations. Its formalism is grounded in Tensor Programs theory, yielding precise scaling laws for parameters, updates, and optimizer settings that generalize across architectures (Blake et al., 2024, Yang et al., 2020, Ma et al., 11 Feb 2026).

1. Foundations: μP, Unit Scaling, and the Motivation for u-μP

μP was introduced to decouple neural network training dynamics from architecture width. In the abc-parametrization formalism, each weight update is scaled by three width-dependent factors: $A_W$ (forward), $B_W$ (init stddev), and $C_W$ (update size). μP prescribes their scaling such that all layers learn at the same maximal rate, yielding width-independence for the learning rate and other hyperparameters in the infinite-width limit.

Unit Scaling, on the other hand, statically rescales every tensor (weights, activations, gradients) at initialization to have variance $1$, centering their magnitudes within the numeric range of low-precision formats. This stabilizes training under aggressive quantization (FP16/FP8), minimizing underflow, overflow, and the need for dynamic scaling.

u-μP combines these schemes. Whereas μP by itself leaves unconstrained multiplicative constants that can impact practical training, and Unit Scaling by itself does not guarantee width-invariant hyperparameters, u-μP inherits both robustness to numerics and provable size-independent optimal hyperparameters. This unification implies both ergonomic and theoretical advantages: significantly simpler hyperparameter sweeps, function-preserving upscaling, and guaranteed feature learning at infinite width (Blake et al., 2024, Ma et al., 11 Feb 2026).

2. Formal Definition and Scaling Laws

u-μP operates within abc-parametrization (Yang et al., 2020), tracking three nonnegative scalars per weight tensor:

• $A_W$: Forward scaling
• $B_W$: Initialization stddev
• $C_W$: Step size per-parameter

For a weight of shape $d_\text{out}\times d_\text{in}$, u-μP prescribes:

• Initialization: $B_W = 1$ for all tensors; $W^{(\ell)} \sim \mathcal N(0, 1/d_\text{in})$
• Forward scaling:
  • Embedding, attention-input: $B_W$0
  • Hidden matmul: $B_W$1
  • Output: $B_W$2 (fwd), $B_W$3 (bwd)
• Parameter update (learning rate):
  • Embedding, hidden: $B_W$4
  • Output: $B_W$5
  • Residual: $B_W$6

These rules ensure that, at initialization, every weight, activation, and gradient is $B_W$7, and the effective gradient scale at each width is preserved. Nonlinearities and layernorms then also exhibit width-independent dynamics. All operations (matmuls, softmax, gated-SiLU, residual-add, etc.) are unit-scaled—accompanied by static forward and backward multipliers for both activations and gradients (Blake et al., 2024).

3. Parameter Initialization, Update Dynamics, and Infinite-Width Theory

Under u-μP, each hidden layer's weight matrix $B_W$8 is initialized by drawing each entry i.i.d. from $B_W$9, while all bias vectors are zero. The forward pass computes:

$C_W$0

ensuring that if $C_W$1, then $C_W$2. The backward pass maintains $C_W$3 gradients for all parameters.

Learning rate per-parameter is set by $C_W$4, e.g., for AdamW:

$C_W$5

where $C_W$6 is the Adam-corrected gradient, and per-step changes $C_W$7 scale as $C_W$8, remaining width-stable.

Tensor Programs theory, specifically the Master Theorem, establishes that as width tends to infinity, the joint distribution of activations, gradients, and parameters converges to a deterministic, width-invariant limit. In u-μP the dynamics are in the “feature-learning regime” (update exponent $C_W$9), meaning features evolve $1$0 (change nontrivially) throughout training, contrasting with NTK-style regimes (where $1$1 and features are frozen) (Yang et al., 2020, Blake et al., 2024).

4. Hyperparameter Transfer and Model Upscaling

A critical practical application of u-μP is μTransfer, which enables direct transfer of optimal hyperparameters and even optimizer state across model sizes, including function-preserving upscaling. All width dependence is absorbed in the prescribed scaling rules, so global learning rate, Adam $1$2, and weight decay determined on a small model continue to be optimal for any upscaled version, whether widened by duplication or initialized de novo (Ma et al., 11 Feb 2026).

For $1$3-fold model width expansion:

• Matrix weights: $1$4
• Biases: $1$5
• Inject i.i.d. noise at u-μP variance to break symmetry
• Duplicate optimizer states accordingly

Theorems in (Ma et al., 11 Feb 2026) guarantee that after static rescaling, the widened model is dynamically equivalent: its function-space trajectory and learning dynamics are identical to the base model. This result applies to all architectures amenable to tensor-program analysis.

5. Practical Implementation and Default Settings

u-μP’s practical implementation involves:

• Substituting every operation (linear, softmax, activation, add, etc.) with a unit-scaled variant for both forward and backward passes.
• Tagging modules or parameters by μP type, so that optimizers scale each gradient appropriately.
• Using static multipliers: no per-tensor dynamic rescaling is necessary for low-precision safety.
• For FP8, insert casting operators (.to(E4M3), etc.) for all layers, except specific exceptions (FFN-down, attention output) which use alternate exponent-mantissa splits.

Empirically-validated default constants, for transformers and LLaMA-style models:

• $1$6 (learning rate)
• Weight decay for independent AdamW: $1$7
• All auxiliary $1$8-multipliers (softmax, residual, output, etc.): $1$9

This configuration enables direct, robust, and loss-preserving transfers from proxy (small) models to large target models in any floating-point format up to and including FP8 (Blake et al., 2024, Ma et al., 11 Feb 2026).

6. Empirical Results and Comparative Analysis

u-μP demonstrates superior hyperparameter transferability, stable low-precision training, and improved scaling efficiency:

• Hyperparameter independence: Average transfer error is $A_W$0 for u-μP, order-of-magnitude smaller than $A_W$1 for μP.
• Hyperparameter search: η–only sweep (9 runs) suffices for u-μP versus multidimensional grid (>100 runs) for standard μP.
• Learning-rate transfer: η is strictly constant across width 128 up to 4096 for u-μP; μP exhibits drift.
• Loss scaling: For equivalent loss, u-μP requires half the width of μP (e.g. width-2048 u-μP versus width-4096 μP).
• Low-precision (FP8) training: μP suffers from numerical instability (gradient underflow); u-μP achieves near full-precision loss (Δ ≈ 0.01).
• Large scale: Up to 7B parameters, stable FP8 convergence, matching BF16 accuracy (Blake et al., 2024, Ma et al., 11 Feb 2026).

Additionally, upscaling experiments with MLPs, ResNets, and GPT-2-style transformers confirm that u-μP upscaling with μTransfer retains all training dynamics from the small model, allows tuning hyperparameters only once on the proxy, and achieves same or better convergence on large target models.

Model + Optimizer	Setup	Convergence Speed	Final Train/Val Metric
MLP + SGD	Tabular, n=500→2000	×2–3 faster	Same train loss, same val acc
ResNet + SGD	CIFAR-100, 1×→4×	×1.5–2 faster	Comparable train/val loss
GPT-2 + AdamW	8B→32B tokens, 2×→4×	×2 faster	Same NLL

7. Theoretical and Empirical Implications

u-μP unifies two independent stability principles—μP’s scaling for width-invariant learning rates and Unit Scaling’s guarantee of numeric safety into a rigorously grounded, architecture-agnostic scheme. The approach collapses hyperparameter tuning to a single learning-rate sweep, interprets all auxiliary scaling constants as softmax temperature or gain, and extends straightforwardly to both new architectures and new low-precision number formats.

In the infinite-width limit, u-μP is proven to admit nontrivial feature learning—the kernel $A_W$2 evolves during training—unlike NTK or standard parametrizations where features are frozen and learning is purely kernel gradient descent. This ensures pretraining and transfer learning regimes (e.g., BERT, Word2Vec) are faithfully captured (Yang et al., 2020).

A plausible implication is the broad applicability of u-μP to function-preserving upscaling, zero-shot optimizer transfer, and out-of-the-box low-precision deployment at scale, with formal guarantees of static and dynamic equivalence. The result is substantially improved efficiency in model design, engineering, and hyperparameter tuning for modern deep learning systems (Blake et al., 2024, Ma et al., 11 Feb 2026).