Background: Why Optimal Parameters for Small Models Often Fail at Scale
When training billion-parameter or larger models, directly searching for learning rates, initialization scales, output layer multipliers, and optimizer parameters at the target scale is usually prohibitively expensive. Engineering teams often use a smaller Proxy Model for hyperparameter sweeps and then replicate the results on the target model.
The problem is that under standard parameterization, changing model width simultaneously alters activation scales, gradient scales, parameter update magnitudes, and output layer behavior. The optimal learning rate on a small model may be too aggressive or too conservative on a large model. Simply increasing the number of proxy experiments doesn’t fix the root cause—the proxy and target models follow different optimization dynamics.
Maximal Update Parametrization (μP) aims to make feature learning behavior consistent across different model widths through a set of width-dependent initialization, learning rate, and output layer scaling rules. The resulting μTransfer workflow involves finding hyperparameters on a small-scale model and transferring them to a wider target model.
Key Limitation: μP primarily addresses width transfer. It is not a universal law for “zero-cost transfer across any scale, architecture, or training budget.”
Core Principles: Transferring Not a Learning Rate, But a Parameterization Contract
1. Base Shape Determines Which Dimensions Are Extensible Widths
A μP implementation typically requires three model shapes:
| Model | Role |
|---|---|
| Base Model | Defines the reference shape |
| Delta Model | Changes all width dimensions planned for scaling relative to Base |
| Target Model | The actual target width for training |
The framework uses Base and Delta to infer which parameter dimensions are finite and which are “infinite” dimensions that scale with model size, then computes scaling rules for initialization and optimizer parameter groups.
If a Transformer simultaneously scales d_model, the number of Attention Heads, or FFN width, the Delta must reflect these changes. Missing a dimension means the code may still run, but the transfer behavior will no longer be as expected.
2. Different Parameter Types Cannot Share the Same Scaling
Hidden layer matrices, input Embeddings, output Readouts, Biases, and normalization parameters play different roles. A μP implementation typically requires:
- Using μP-compatible initialization;
- Setting width-dependent learning rates for hidden matrices;
- Using
MuReadoutor an equivalent output layer; - Adjusting Transformer Attention scaling accordingly;
- Building fine-grained parameter groups with a μP optimizer.
Therefore, “just swapping AdamW for MuAdam” is not a complete conversion. Modifying only some modules creates a hybrid parameterization that is hardest to debug—loss may decrease normally, but the hyperparameter transfer curves are already distorted.
3. Coord Check Validates Coordinate Scales Across Widths
Coordinate Check records activation coordinate scales after initialization and a few training steps across multiple width models. If the implementation is correct, key layer statistics should remain roughly stable as width increases, rather than exploding or trending toward zero.
Its role is analogous to a gradient check: verifying that the implementation behaves as expected, not proving that final training will succeed. After Coord Check passes, you still need to verify the learning rate peak location, training stability, and target model short runs.
Engineering Implementation: Making μP a Deployable Training Configuration Artifact
1. Establish an Immutable Base Shape Fingerprint
Don’t just save a .bsh file. It’s recommended to generate a Manifest for each parameterization release:
schema_version: 1
model_family: decoder-transformer
code_commit: "9f83c12"
base_shape_sha256: "..."
base_width:
d_model: 512
n_heads: 8
d_ff: 2048
delta_width:
d_model: 768
n_heads: 12
d_ff: 3072
fixed_dimensions:
vocab_size: 128000
n_layers: 24
mup_contract:
readout: MuReadout
attention_scale: "8 / head_dim"
optimizer: MuAdamW-compatible
init_policy: mup_init_v2
tokenizer_fingerprint: "..."
training_horizon_tokens: 3000000000
global_batch_tokens: 1048576
The Manifest should at least record:
- Parameter names, shapes, and finite/infinite dimension classification;
- Base, Delta, and Target model configurations;
- Initialization strategy, Attention scaling, output layer type;
- Optimizer parameter groups and initial multipliers;
- Tokenizer, vocabulary, and weight tying relationships;
- Depth, token budget, batch size, and scheduler.
These fields determine whether a Proxy Sweep can establish an auditable correspondence with target training.
2. Set Base Shape Before Building the Optimizer
The official implementation requires calling set_base_shapes early, typically before custom reinitialization and optimizer creation:
from mup import MuReadout, MuAdam, make_base_shapes, set_base_shapes
base = TransformerLM(d_model=256, n_heads=4, d_ff=1024)
delta = TransformerLM(d_model=512, n_heads=8, d_ff=2048)
make_base_shapes(base, delta, "transformer-v1.bsh")
model = TransformerLM(d_model=2048, n_heads=32, d_ff=8192)
set_base_shapes(model, "transformer-v1.bsh")
model.reset_parameters_with_mup_init()
optimizer = MuAdam(model.parameters(), lr=3e-4)
⚠️ Be especially careful when restoring checkpoints. Some implementations do not fully write the
infshapemetadata attached to parameters into regular checkpoints. When re-setting Base Shape after loading trained weights, avoid rescaling existing parameters, and use unit tests to confirm parameter-by-parameter consistency before and after restoration.
3. Proxy Sweep Must Keep Non-Width Conditions Comparable
The proxy and target models should at least fix:
- Data distribution and Tokenizer;
- Optimizer type and parameter definitions;
- Scheduler shape;
- Training objective and loss normalization;
- Global Batch Tokens;
- Token Horizon or explicit transfer correction rules;
- Model depth and module topology;
- Mixed precision strategy.
If the Proxy trains for a very short token budget while the target model trains much longer, the optimal learning rate may shift systematically. Research has shown a relationship between optimal learning rate and Token Horizon, as well as coupling between learning rate, Batch Size, and critical batch size. Therefore, μP width transfer cannot replace time-dimension and batch-dimension modeling.
4. Use Multi-Width Curves, Not a Single Proxy Point
It’s recommended to prepare at least three widths:
| Width | Purpose |
|---|---|
| Small sweep model | Main search space |
| Medium validation model | Verify transfer trends |
| Target model or short-run model | Final confirmation |
For each candidate learning rate, compare training loss, gradient norms, activation statistics, and first-token efficiency. If the optimal learning rate drifts noticeably between small and medium models, do not extrapolate directly to the target scale.
More important than “which learning rate is best” is observing:
- Whether rankings are stable;
- Whether optimal intervals overlap;
- Whether the model is robust to learning rate errors;
- Whether wider models show anomalous degradation at the same step count.
5. Treat Embedding and Weight Decay as Independent Audit Items
Recent research suggests that in AdamW settings, the significant gains of μP over standard parameterization may largely come from removing the embedding layer learning rate bottleneck. Other studies challenge μP’s long-term dynamics assumptions and emphasize the importance of Weight Decay and Warmup for cross-width learning rate transfer.
Production systems should not treat μP as a black box. At a minimum, log and sweep independently:
- Embedding Learning Rate
- Hidden Weight Learning Rate
- Readout Multiplier
- Weight Decay
- Warmup Token Count
- Gradient Clipping Threshold
If transfer fails, first check module-level update magnitudes, rather than simply concluding “μP doesn’t work.”
Transfer Matrix: Which Parameters Can Be Transferred, Which Must Be Re-Validated
| Change Dimension | Recommendation |
|---|---|
| Increase Transformer width only | μP’s primary target scenario; can transfer after Proxy Sweep |
| Change model depth | Not guaranteed by default; needs independent validation or a method that explicitly supports depth transfer |
| Change total token budget | Requires Token Horizon scaling study or short-run calibration |
| Change Global Batch Tokens | Requires re-evaluation of learning rate vs. critical batch size relationship |
| Dense to MoE | Architecture and effective tokens per Expert change; do not transfer directly |
| Change Optimizer | Re-sweep; not considered the same transfer contract |
| Change Tokenizer or vocabulary | Embedding/Readout shape and data distribution change; re-validate |
| BF16 to FP8 training | Numerical stability boundaries change; requires independent precision and scale gates |
Applicable Scenarios
μP is suitable for:
- Planning to train multiple models with the same depth, topology, but different widths;
- Direct hyperparameter sweeps on large models are very expensive, but training multiple small Proxies is feasible;
- The training platform can freeze code, data, Tokenizer, and optimizer contracts;
- The team is willing to maintain Base Shape and Coord Check tests.
Not suitable as a universal cross-architecture transfer solution. For example, transferring from Dense to MoE, simultaneously changing depth and token budget, or directly applying the optimal learning rate from a short experiment to very long pretraining.
Common Misconceptions
Misconception 1: The Smaller the Proxy, the Cheaper and Better the Transfer
A proxy that is too small may operate in a different optimization regime, with noisier activation and gradient statistics. Look for the smallest model that is “cheap enough and has entered a stable transfer regime,” rather than minimizing indefinitely.
Misconception 2: A Flat Coord Check Means Transfer Success
A flat curve could mean the learning rate is too small, and the model has barely changed functionally. The official project recommends using a sufficiently large learning rate so that potential explosion issues can be exposed within a few steps.
Misconception 3: Base Shape Is Just an Auxiliary File
Base Shape actually defines the parameterization semantics. Code renaming, layer structure changes, weight tying changes, or new projection layers can invalidate an old Base Shape. It should be versioned like a database schema.
Misconception 4: The Target Model Doesn’t Need Short Runs at All
Even if Proxy Sweep results are stable, perform controlled short runs at the target scale to check for NaN/Inf, gradient norms, memory, throughput, initial loss drop, and module-level update multipliers.
Go-Live Checklist
- Base, Delta, and Target have consistent depth and module topology
- Delta covers all width dimensions planned for scaling
- Base Shape Manifest has been hashed and stored in the artifact repository
-
set_base_shapesis called before initialization and optimizer creation - Readout, Attention scaling, initialization, and optimizer all conform to a unified contract
- Multi-width Coord Check passes at initialization and after several update steps
- Proxy Sweep fixes data, Token Horizon, Batch, and Scheduler
- Optimal intervals for small and medium Proxies overlap
- Embedding LR, Weight Decay, and Warmup are logged separately
- Target scale short run passes numerical stability and loss gates
- Checkpoint restoration does not rescale parameters a second time
- On transfer failure, can fall back to standard parameterization or previous training recipe
References
- Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer — arXiv:2203.03466
- Microsoft μP Official Repository — github.com/microsoft/mup
- Scaling Optimal LR Across Token Horizons — OpenReview
- Time Transfer: On Optimal Learning Rate and Batch Size in the Infinite Data Limit — OpenReview
- Understanding the Mechanisms of Fast Hyperparameter Transfer — OpenReview
- Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate — arXiv:2605.21486
- Weight Decay may matter more than μP for Learning Rate Transfer in Practice — arXiv:2510.19093