PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2203.03466v2cs.LGcond-mat.dis-nncs.NE

TL;DR

Deep-learning hyperparameter tuning is prohibitively expensive for billion-parameter networks. The paper introduces µTransfer: parametrize the target with µP, tune a smaller proxy, and transfer the tuned hyperparameters zero-shot. This approach outperforms published BERT-large and GPT-3 results with substantially lower tuning costs.

  • Problem

    Hyperparameter tuning is critical but prohibitively expensive for neural networks with billions of parameters.

  • Method

    µTransfer parametrizes a large target model in µP, tunes a smaller proxy model, and copies the tuned hyperparameters to the target.

  • Results

    µTransfer outperforms published BERT-large and 6.7B GPT-3 results, with tuning costs approximately equal to one BERT-large pretraining and 7% of GPT-3 pretraining.

  • Takeaways & Limitations

    A single small model can provide transferable hyperparameters across model families and training scales, reducing large-model tuning to proxy-model tuning.

  • Takeaways & Limitations

    Transfer has caveats: initialization does not transfer well across depth, depth transfer generally fails for post-layernorm Transformers, and regularization hyperparameters remain size- and data-dependent.

Abstract

from arXiv · show

Hyperparameter (HP) tuning in deep learning is an expensive process, prohibitively so for neural networks (NNs) with billions of parameters. We show that, in the recently discovered Maximal Update Parametrization (muP), many optimal HPs remain stable even as model size changes. This leads to a new HP tuning paradigm we call muTransfer: parametrize the target model in muP, tune the HP indirectly on a smaller model, and zero-shot transfer them to the full-sized model, i.e., without directly tuning the latter at all. We verify muTransfer on Transformer and ResNet. For example, 1) by transferring pretraining HPs from a model of 13M parameters, we outperform published numbers of BERT-large (350M parameters), with a total tuning cost equivalent to pretraining BERT-large once; 2) by transferring from 40M parameters, we outperform published numbers of the 6.7B GPT-3 model, with tuning cost only 7% of total pretraining cost. A Pytorch implementation of our technique can be found at github.com/microsoft/mup and installable via `pip install mup`.

1 Introduction

Hyperparameter tuning becomes prohibitively expensive for billion-parameter networks. The paper introduces µTransfer, which uses µP to tune smaller proxy models and transfer hyperparameters to large targets without direct tuning.

  • µTransfer: µTransfer parametrizes the target model in µP, tunes a smaller proxy model, and copies its hyperparameters to the target.The procedure is summarized as target parametrization, proxy tuning, and hyperparameter copying.
  • Why µP: µP makes optimal learning rates stable across Transformer widths, unlike standard parametrization.Under standard parametrization, wider networks can have different optima and may underperform narrower models even after learning-rate tuning.
  • Transfer scope: µTransfer transfers hyperparameters across width and, with caveats, across depth, batch size, sequence length, and training time.The paper distinguishes theoretically justified width transfer from empirically supported transfer across other training-scale dimensions.
  • Benefits: µTransfer outperforms published BERT-large and GPT-3 results while reducing tuning costs to approximately one BERT-large pretraining and 7% of GPT-3 pretraining.The reported comparisons use 350M-parameter BERT-large and the 6.7B-parameter GPT-3 model.
  • Benefits: The method supports tuning one small model and reusing its hyperparameters across a fixed family of models with varying width and depth.The paper verifies the approach on Transformers and ResNets and releases a PyTorch implementation.
  • Caveats: µTransfer primarily targets training-loss transfer and may be insufficient when regularization limits test performance, such as fine-tuning on small datasets.Regularization hyperparameters are not transferred in the presented fine-tuning procedure.

2 Parametrization Matters: A Primer

The primer explains why hyperparameter transfer depends on parametrization rather than on reusing raw values alone. By scaling parameters and hyperparameters appropriately, µP yields width-stable optimization behavior that permits zero-shot transfer.

  • Scaling intuition: The central-limit-theorem example shows that choosing c = α/√n stabilizes the objective as width n grows.Other scaling choices make the summed quantity vanish or diverge in variance.
  • Transfer principle: Because the stabilized objective depends on α rather than width, an optimal α for a smaller problem can transfer approximately to a larger one.The transfer copies the appropriately parametrized hyperparameter between problem sizes.
  • Neural-network analogy: In neural networks, width-n parameters and learning-rate-like hyperparameters play the roles of the random variables and scaling factor in the example.Correct parametrization makes narrower-network optima approximately optimal for wider networks.
  • µP versus SP: µP is presented as the correct parametrization for zero-shot hyperparameter transfer, whereas standard parametrization is incorrect because its wide-network limit can blow up.The paper connects this distinction to µP’s well-defined infinite-width limit.
  • Scope of parametrization: Transferability requires correctly reparametrizing all relevant hyperparameters, not only the learning rate.The paper specifically cites parameter multipliers and input/output-layer learning rates as additional quantities that must scale correctly.

3 Hyperparameters Don’t Transfer Conventionally

Conventional parametrization makes optimal hyperparameters unstable as neural networks widen, undermining direct transfer from small to large models. This instability appears in MLPs and Transformers and is tied to training behavior beyond initialization.

  • MLP with Standard Parametrization: Figure 3 compares CIFAR-10 MLPs across hidden sizes under SP and µP, with µP showing better learning-rate stability.Both settings use SGD and 20 training epochs.
  • MLP with Standard Parametrization: In SP, the optimal learning rate shifts by roughly an order of magnitude as MLP width increases from 256 to 8192.Using the smallest model’s optimal learning rate on the largest model produces very bad performance, potentially causing divergence.
  • Why Hyperparameters Fail to Transfer: Standard parametrization with a width-scaled learning rate still does not support hyperparameter transfer because it enters a kernel regime.Thus, a well-defined infinite-width limit alone is insufficient.

4 Unlocking Zero-Shot Hyperparameter Transfer with µP

µP modifies initialization and layer-specific learning-rate scaling so that training behavior remains comparable across widths. This enables hyperparameters tuned on smaller models to transfer to larger MLPs and Transformers.

  • MLP with µP: µP modifies the MLP’s last-layer initialization and the learning rates of the first, last, and bias layers.The master learning rate η determines the width-dependent rates, including ηW3 = ηn^-1.
  • MLP with µP: At the base width n0, inserted µP multipliers equal 1 and reproduce SP; as width increases, µP defines a different optimization trajectory.This width-dependent deviation is crucial for hyperparameter transfer.
  • MLP with µP: The best learning rate for a width-128 µP MLP is also best for a width-8192 MLP, unlike under SP.The CIFAR-10 experiments evaluate performance across learning rates and widths.
  • Empirical Validation: Figure 4 evaluates stability of learning rate, αoutput, initialization standard deviation, and learning-rate schedule in pre-LN µP Transformers.Each hyperparameter is swept separately while the others remain fixed on Wikitext-2.
  • Transformers with µP: For Transformers, µP yields a stable optimal learning rate and performance that improves monotonically with width.The formulation uses base width n0 = 128 and modifies attention scaling to 1/d.

5 The Defects of SP and How µP Fixes Them

SP causes width-dependent blow-up and mismatched layer updates, while µP keeps key quantities bounded and aligns update speeds. These properties explain its stronger transfer and scaling behavior.

  • An Instructive Example: In a one-hidden-layer linear perceptron, SP outputs blow up with width after one SGD step because UᵀU = Θ(n).The divergence follows from the law of large numbers applied to the hidden weights.
  • An Instructive Example: Under µP, width-dependent initialization and learning rates make the same perceptron output Θ(1) after one step instead of blowing up.The relevant µP rates are ηV = 1/n and ηU = n.
  • Some Layers Update Too Fast, Others Too Slow: In Transformers trained with Adam, SP makes logits and attention logits blow up with width, whereas µP keeps them bounded.The same pattern is observed for other advanced architectures and optimizers.
  • Some Layers Update Too Fast, Others Too Slow: Scaling down SP’s learning rate would leave word embeddings insufficiently learned because their updates are width-independent.µP instead designs hidden activations to update at the same speed across widths.
  • Performance Advantage of µP: A wide model tuned with µTransfer should generally outperform an SP counterpart with a globally tuned learning rate.In SP, preventing logit blow-up compromises learning in quantities such as word embeddings.

6 Which Hyperparameters Can Be µTransferred?

µTransfer applies most directly to hyperparameters that remain stable across scale, while regularization settings and some depth transfers remain constrained. Experiments identify practical scale ranges and architecture-dependent caveats.

  • Hyperparameter Categories: Hyperparameters fall into transferable settings, regularization controls that transfer poorly, and training-scale variables such as width, depth, and batch size.The framework transfers other hyperparameters across the dimensions in the third category.
  • Transfer Across Scale: Width transfer is broad enough to accommodate different layer-width ratios and different numbers of Transformer attention heads.This supports practical architectural variation during transfer.
  • Transfer Across Scale: Regularization should not generally be expected to transfer when its parametrization depends only on model size, because overfitting control depends on model and data size.This is distinct from hyperparameters that govern optimization across scale.
  • Empirical Validation and Limitations: The experiments evaluate pre-LN µP Transformers by sweeping four hyperparameters across width and depth, with results averaged over five random seeds.Additional experiments examine transfer across batch size, sequence length, and training time.
  • Empirical Validation and Limitations: For Transformer language modeling, transfer generally works once minimum width 256, depth 4, batch size 32, sequence length 128, and 5000 training steps are met.The target must also remain within the experiments’ reasonable scale range.
  • Empirical Validation and Limitations: Initialization standard deviation transfers poorly across depth, and depth transfer works only for pre-LN rather than post-LN Transformers.Fixing initialization while tuning other hyperparameters can nevertheless work in the reported depth-transfer experiments.

7 Efficiency and Performance of µTransfer

µTransfer tunes hyperparameters on smaller proxy models and applies them to larger targets, aiming to improve efficiency without sacrificing performance. Across translation and language-model experiments, it matches or exceeds conventional or published results under substantially lower direct tuning requirements.

  • Experimental setup: µTransfer tunes hyperparameters on a smaller proxy model, evaluates them directly on the large target, and compares against target-model tuning.The experiments ask whether this procedure can reduce tuning efficiency costs while achieving performance comparable to traditional tuning.
  • IWSLT14 De-En: µTransfer from a 0.25x IWSLT Transformer provides better and more reliable outcomes than traditional 1x tuning at a fixed compute budget.The comparison uses 25 independent tuning trials and validation-loss-selected hyperparameters, with evaluation averaged over five random initializations.
  • Efficiency-performance frontier: The compute-performance Pareto frontier completely dominates conventional tuning across the tested compute budgets.The frontier is evaluated using median BLEU over 25 trials against relative compute budget.
  • Sample quality: With a fixed number of random samples, direct target tuning slightly outperforms proxy tuning, but the gap vanishes as the sample count increases.The proxy behaves as a noisy estimator of the wider target model, so additional samples suppress search distortion.
  • BERT pretraining: BERT hyperparameter tuning uses 256 proxy-model combinations, with total tuning cost roughly equal to pretraining one BERT-large model.The resulting BERT-large model shows sizeable improvement over the well-tuned Megatron BERT-large baseline without directly tuning the full model.
  • GPT-3 pretraining: The GPT-3 6.7B µTransfer model outperforms the published 6.7B model and is comparable to the twice-as-large 13B model.The comparison is affected by different attention implementations and FP32 precision for the µTransfer model because of numerical divergences.

8 Wider is Better in µP Throughout Training

In µP, wider models generally outperform narrower models throughout training, unlike in standard parametrization. This pattern holds in a GPT-3-style stress test across widths from 256 to 32,768, apart from a brief noisy interval.

  • Stress test: Wider models consistently match or outperform narrower models at each training point in a µP GPT-3 Transformer.The stress test scales width from 256 to 32,768 using fixed hyperparameters and finds only a brief exception around 1e8 training tokens.
  • Comparison with SP: µP shows wider-is-better behavior throughout training, whereas standard parametrization does not.The observation assumes zero-initialized output layers and allows for noise from random initialization or data ordering.
  • Implications: The authors suggest that appropriately scaled wider models are more data-efficient and that early wider-is-better checks can debug µP implementations cheaply.These are presented as implications of the observed training trajectories.

9 Useful Hyperparameter Transfer: A Theoretical Puzzle

The paper frames useful hyperparameter transfer as a separation between rapidly converging hyperparameter optima and slowly converging model functions. Why this separation occurs, and where else it applies, remains unresolved.

  • Theoretical puzzle: A useful proxy must be large enough for its hyperparameter optimum to converge, but small enough that its function and loss have not converged to the target.The paper describes the hyperparameter optimum as a coarse variable and the learned function as a fine detail that converges more slowly with width.
  • Open question: Experiments demonstrate that such an intermediate proxy size exists, even though the paper does not explain theoretically why this should happen.The authors leave the explanation and the scope of other settings to future work.

10 Related Works

µTransfer differs from prior hyperparameter-optimization and transfer methods by decoupling target-model size from tuning cost through zero-shot transfer across model scale. The paper positions these methods as complementary because they can tune the proxy model.

  • Efficient tuning: Unlike Bayesian optimization and multi-arm-bandit approaches, µTransfer does not optimize the hyperparameter-search process itself; it decouples target-model size from tuning cost.A fixed-sized proxy can probe the hyperparameter landscape regardless of target size, while existing optimizers can still be applied to that proxy.
  • Hyperparameter transfer: µTransfer is presented as the first zero-shot hyperparameter-transfer method, transferring across model scale rather than tasks or datasets.Methods such as Hyperband may use cheap model estimates but still require expensive evaluation of the large model and are therefore not zero-shot.
  • Scaling rules: Prior learning-rate and batch-size scaling laws report inconsistent behavior across datasets and models, whereas this work studies training loss across broad batch-size regimes.The paper also covers normalization, Adam and other adaptive optimizers, and transfer across depth and sequence length.
  • Scope: The method differs from other favorable initialization or parametrization schemes by focusing specifically on transferability of optimal hyperparameters across model scale.The empirical scope includes Transformers and ResNets, with derivations covering Adam as well as SGD.
  • Parametrization: The paper argues that µP is uniquely suited to width transfer because it preserves feature learning while maintaining an appropriate infinite-width limit.This theoretical argument contrasts µP with standard and NTK parametrizations, which the paper characterizes as having bad infinite-width limits for wide networks.

11 Conclusion

The paper concludes that µP stabilizes hyperparameter landscapes across width, enabling indirect tuning of large networks, while identifying important transfer limitations and open theoretical questions.

  • µP makes hyperparameter landscapes reasonably stable across width, allowing smaller models to tune hyperparameters for full-sized networks.The approach transfers across width and empirically across depth, batch size, sequence length, and training time, with caveats.
  • Regularization hyperparameters and some initialization transfers remain limited, especially for small-data finetuning and depth changes.The authors specifically note that regularization is not the bottleneck in their experiments and that initialization does not transfer well across depth.
  • Parameter multipliers affect backpropagation and therefore cannot generally be absorbed into initialization, although trained multipliers can be absorbed into weights afterward.
  • µP parametrization specifies how parameter multipliers, initialization, learning rates, and other multipliers scale as network widths change.It does not prescribe specific hyperparameter values for any one width.
  • Incorrect parametrization can force other hyperparameters to compensate increasingly with width, distorting their optima.

E.1 Further Discussions on Hyperparameter Categories

The paper distinguishes hyperparameters that transfer across training scale from those that primarily regularize, and evaluates transfer across width and other architectural or training dimensions.

  • µP theoretically justifies transfer across width, whereas transfer across depth, batch size, sequence length, and training time is supported empirically.
  • Hyperparameters That Don’t Transfer Well: Regularization-focused hyperparameters generally transfer poorly because required regularization increases with model size and decreases with data availability.Dropout probability and weight decay are examples of primarily regularizing hyperparameters.
  • Hyperparameters Transfered Across: Training-scale hyperparameters include width, depth, sequence length, batch size, and training steps, which determine computation per pass or the number of passes.
  • Number of Attention Heads: Transformer width can be represented by dmodel, dhead, nhead, or dffn, and attention-head scaling can fix dhead while increasing nhead.The relation dmodel = dhead × nhead supports models such as GPT-3 that scale by increasing the number of heads.
  • IWSLT14 De-En: A 4M-parameter IWSLT14 proxy transfers learning rate and parameter multipliers to a 40M-parameter model using the same number of attention heads.The 0.25x proxy shrinks dmodel and other dimensions fourfold while retaining nhead.
  • Finetuning: Finetuning hyperparameters are not transferred in the reported BERT procedure because regularization is essential for successful finetuning.

F.4 GPT-3

For GPT-3, the authors search hyperparameters on a reduced-width proxy over shorter training horizons, transfer the selected settings, and address numerical instability during full-scale training.

  • A 256-hidden-unit reduced-width GPT-3 proxy was used to search hyperparameters before training the 6.7B target model.The proxy inherited relative position encoding from the evolved GPT-3 codebase.
  • The search varied initialization scale, attention temperature, output temperature, embedding multiplier, and relative position embedding multiplier.
  • 4 billion and 16 billion training-token horizons were used instead of the target model’s 300 billion tokens to reduce tuning cost.The shorter horizons also changed the length of the linear learning-rate decay schedule.
  • 467 proxy-model runs were performed: 350 at 4 billion tokens and 117 at 16 billion, with 286 and 80 runs completing without divergence, respectively.
  • The 4-billion- and 16-billion-token searches aligned well; learning rate and initialization scale mattered most, while embedding scale was set to 10.The selected learning rate was 0.006 and initialization scale was 2.5; most other hyperparameters remained at 1.
  • The initial µTransfer GPT-3 training diverged because of FP16 backward-pass underflow, so the authors completed training with FP32 activations and weights.They hypothesize that more aggressive transferred hyperparameters, such as a higher learning rate, contributed to the instability.

G.1.1 ResNet on CIFAR-10

Across ResNet experiments, µP makes hyperparameter choices more stable across width and enables transfer from smaller proxies, while SP is less reliable.

  • Hyperparameter Stability: The best width-8 SP model underperforms the width-4 model, whereas wider µP models consistently perform better.The comparison uses CIFAR-10 validation accuracy across width multipliers.
  • Hyperparameter Stability: µP keeps the optimal learning rate η and output multiplier αoutput stable across width.The sweep varies width and αoutput while holding other settings fixed.
  • Hyperparameter Transfer: 0.43% ± .001%: µP outperforms SP when η and αoutput are tuned on width 0.5 and transferred to width 8.Both parametrizations use the same proxy-model grid search and transferred hyperparameters.
  • Hyperparameter Transfer: 0.41%: µP outperforms SP in top-1 validation accuracy after transferring η, αoutput, weight decay γ, and momentum β from width 0.125 to width 1.The experiment samples 64 hyperparameter combinations on ImageNet Wide ResNet models.
  • Transformer Transfer: Post-layernorm Transformer experiments show transferability across width, batch size, sequence length, and training steps, but transfer across depth is generally fragile.The transferable hyperparameters include learning rate, αoutput, initialization standard deviation, and Adam β2.

I Reverse-µTransfer for Diagnosing Training Instability in Large Models

Reverse-µTransfer maps instability-inducing hyperparameters from large Transformers to smaller models, enabling faster diagnosis of divergence and its causes.

  • Diagnostic Use: Reverse-µTransfer replicates large-model training instability on smaller Transformers, supporting its use for diagnosing instability.The paper presents two case studies involving 8192-width and 6B-parameter language models.
  • Case Studies: A width-8192 Transformer’s divergence was reproduced at width 256, where activation analysis identified exploding attention logits as the cause.The same cause was later confirmed in the original large model.
  • Case Studies: A 6B-parameter model’s repeated mid-training blow-up was replicated in a 100M-parameter model and resolved by retuning the smaller model with random search.The transferred hyperparameters came from the standard-parametrization large model.
  • Reverse Transfer: The instability curves have identical overall shapes between large models and reverse-transferred narrow models: a learning rate is unstable in one exactly when transferred back to the other.The curves’ vertical offsets reflect wider models’ greater capacity rather than the instability correspondence.
  • Theory Intuition: The scaling intuition treats Gaussian initialization as random matrices and SGD updates as sums of outer products, with correlated tensor-product updates requiring Θ(1/n)-sized entries.These matrix behaviors motivate the parametrization rules used to analyze training dynamics.

J.3 Why Other Parametrizations Cannot Admit Hyperparameter Transfer

The paper argues that hyperparameter transfer requires a feature-learning parametrization that updates all parameters maximally, which uniquely identifies µP.

  • Standard Parametrization: SP can cause infinite-width blow-up; reducing its learning rate avoids blow-up but yields a kernel-regime limit that cannot transfer hyperparameters in all situations.For Adam, Θ(1/n) learning rates avoid preactivation blow-up but produce non-maximal feature learning; for SGD, Θ(1/width) enters the kernel regime.
  • Neural Tangent Parametrization: NTP can have trivial performance in its infinite-width limit, and wider models do not always improve when feature learning is crucial.The paper therefore treats NTP, and SP with Θ(1/width) learning rates, as unsuitable for wide modern pretraining.
  • Other Parametrizations: Any nontrivial stable natural parametrization is characterized as admitting either a feature-learning limit or a kernel limit, but not both.The paper uses this dynamical dichotomy to narrow the relevant parametrization choices.
  • Other Parametrizations: Other feature-learning parametrizations leave some parameters effectively fixed at initialization in the infinite-width limit, making their learning rates useless for transfer.Their limiting hyperparameter landscapes therefore need not reflect finite-width neural networks.
  • Maximal Update Parametrization: µP uniquely updates all parameters maximally, so each parameter’s learning rate has approximately the same role at finite and infinite width.Consequently, the µP limiting hyperparameter landscape should reflect the finite-width landscape.
Loading 2203.03466v2…