Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer
Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, Jianfeng Gao
NeurIPS 2021
cs.LG, cond-mat.dis-nn, cs.NE
2022-03-07
Under μP, HPs tuned on a 13M proxy beat Megatron BERT-large (test loss 1.683 vs 1.731); a 40M proxy tunes 6.7B GPT-3 at 7% of pretrain cost.
Bad hyperparameters make large nets untrainable. Under standard parametrization the best learning rate moves with width, and a wider net can lose to a narrower one. Earlier Tensor Programs papers introduced Maximal Update Parametrization (μP), which keeps each layer's effect on activations from vanishing or exploding as width grows. This NeurIPS 2021 paper, from Microsoft and OpenAI, turns that into an engineering recipe.
The recipe is μTransfer: put the target model in μP, tune a smaller proxy, copy the hyperparameters, and never search on the large model itself.
μP rescales init variance and learning rate with layer width. For SGD, the input-layer learning rate grows with fan-out, the output-layer rate shrinks with fan-in, and hidden layers use 1/fan-in; output init variance is 1/fanin², smaller than the usual 1/fanin. Adam converts differently: hidden learning rates stay width-independent, the output rate still divides by fan-in. Transformers also replace 1/√d attention with 1/d, because trained q and k correlate, so the logit scales as d by the law of large numbers, not √d.
Optimization hyperparameters transfer: learning rate, Adam betas, schedules, init variance, parameter multipliers. Regularizers do not: dropout and weight decay depend on both capacity and data. Scale axes you can transfer across include width, and empirically depth, batch, sequence length, and training steps. Depth transfer holds for pre-LN Transformers only, and init standard deviation is unstable across depth, so practice often freezes init and tunes the rest.
The proxy has to be large enough to resemble the target. Reported rules of thumb: width around 256, depth around 4, batch around 32, sequence length around 128, training steps around 5000. Below that the optimum drifts.
On IWSLT14 De-En, a 4M 0.25× proxy copied onto a 40M 1× model scored BLEU percentiles 35.27 to 35.33, above the fairseq default 33.62 to 35.00. Same-budget tuning on 1×, and naive copy under standard parametrization, diverged or did worse. On WMT14 En-De, a 15M proxy copied onto 211M scored BLEU 25.94 to 26.34 versus fairseq 25.69.
For BERT, a 13M prototype served both base and large; total tuning FLOPs matched one BERT-large pretrain. Versus Megatron defaults, μTransfer BERT-base test loss was 1.970 vs 1.995, MNLI 84.3/84.8 vs 84.2/84.2, QQP 90.8 vs 90.6. BERT-large test loss 1.683 vs 1.731, MNLI 87.0/86.5 vs 86.3/86.2, QQP 91.4 vs 90.9. Naive transfer diverged both times.
GPT-3 6.7B (32 layers, width 4096) transferred from a width-256 proxy of about 40M parameters, 168× smaller, at 7% of pretraining cost. Validation loss 1.98 versus 2.03 for a re-run with the original hyperparameters. LAMBADA zero-shot 73.49 versus 70.3 (published 6.7B) and 72.5 (13B). HellaSwag zero-shot 72.0 versus 67.4 and 70.9. Read the caveats together: the μP run used FP32 after FP16 divergences, while the original and the re-run used FP16; the re-run also used absolute attention by mistake. Perplexities use a custom tokenizer and are not comparable to the literature.
Under μP, wider is better throughout training. Under standard parametrization a wider net can be worse. A 4-layer GPT-3-style Transformer swept width to 32,768; except a brief stretch near 1e8 tokens, wider stayed lower in training loss.
Billion-parameter models cannot host their own hyperparameter grids. μTransfer moves the search back onto single-GPU proxies, and one tuned point can serve a whole width family. pip install mup made the recipe copyable. It covers optimization hyperparameters, not regularizers, and not transfer after you change data or task.
Init does not transfer cleanly across depth, and post-LN depth transfer fails. Small proxies still show a slight optimum shift. Regularization is explicitly out of scope, especially for small-data finetuning. The GPT-3 comparison mixes precision and attention implementation, so 6.7B beating 13B on some tasks is not a clean ablation. The paper optimizes training loss; test gains appear because regularization was not the bottleneck in these setups.