Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer
Year: 2022

Abstract:
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 (μP), many optimal HPs remain stable even as model size changes. This leads to a new HP tuning paradigm we call μTransfer: parametrize the target model in μP, 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 μTransfer 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.2

text: 1 Introduction
Hyperparameter (HP) tuning is critical to deep learning. Poorly chosen HPs result in subpar performance and training instability. Many published baselines are hard to compare to one another due to varying degrees of HP tuning. These issues are exacerbated when training extremely large deep learning models, since stateof-the-art networks with billions of parameters become prohibitively expensive to tune.
Recently, [45] showed that different neural network parametrizations induce different infinitewidth limits and proposed the Maximal Update Parametrization (abbreviated µP) (summarized in Table 3) that enables “maximal” feature learning in the limit. Intuitively, it ensures that each layer is updated on the same order during training regardless of width.3 In contrast, while the
standard parametrization (SP) ensures activations are of unit order at initialization, it actually causes
†Work done partly during Microsoft AI Residency Program. ∗Equal contribution. Order is random. Correspondence to {gregyang, edwardhu}@microsoft.com 2See arxiv.org for the full, up-to-date version of this work. 3i.e., the updates’ effect on activations becomes roughly independent of width in the large width limit.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Algorithm 1 Tuning a Large Target Model via µTransfer 1: Parametrize target model in Maximal Update Parametrization (µP) 2: Tune a smaller version (in width and/or depth) of target model 3: Copy tuned hyperparameters to target model
them to blow up in wide models during training [45] essentially due to an imbalance of per-layer learning rate (also see Fig. 8). We leverage µP to zero-shot transfer HPs from small models to large models in this work – that is, we obtain near optimal HPs on a large model without directly tuning it at all! While practitioners have always guessed HPs of large models from those of small models, the results are hit-or-miss at best because of incorrect parametrization. For example, as shown in Fig. 1, in a Transformer, the optimal learning rate is stable with width in µP (right) but far from so in standard parametrization (left). In addition to width, we empirically verify that, with a few caveats, HPs can also be transferred across depth (in Section 5.1) as well as batch size, language model sequence length, and training time (in Appendix I.2.1). This reduces the tuning problem of an (arbitrarily) large model to that of a (fixed-sized) small model. Our overall procedure, which we call µTransfer, is summarized in Algorithm 1 and Fig. 2, and the HPs we cover are summarized in Tables 1 and 2.
There are several benefits to our approach: 1. Better Performance: µTransfer is not just about predicting how the optimal learning rate scales in SP. In general, we expect the µTransferred model to outperform its SP counterpart with learning rate optimally tuned. For example, this is the case in Fig. 1 with the width-8192 Transformer. We discuss the reason for this in Appendices B and C. 2. Speedup: It provides massive speedup to the tuning of large models. For example, we are able to outperform published numbers of (350M) BERT-large [9] purely by zero-shot HP transfer, with tuning cost approximately equal to 1 BERT-large pretraining. Likewise, we outperform the published numbers of the 6.7B GPT-3 model [6] with tuning cost being only 7% of total pretraining cost. For models on this scale, HP tuning is not feasible at all without our approach. 3. Tune
Once for Whole Family: For any fixed family of models with varying width and depth (such as the BERT family or the GPT-3 family), we only need to tune a single small model and can reuse its HPs for all models in the family.4 For example, we will use this technique to tune BERT-base (110M parameters) and BERT-large (350M parameters) simultaneously by transferring from a 13M model. 4. Better Compute Utilization: While large model training needs to be distributed across many GPUs, the small model tuning can happen on individual GPUs, greatly increasing the level of parallelism for tuning (and in the context of organizational compute clusters, better scheduling and utilization ratio). 5. Painless Transition from Exploration to Scaling Up: Often, researchers explore new ideas on small models but, when scaling up, find their HPs optimized during exploration work poorly on large models. µTransfer would solve this problem.
Nevertheless, µTransfer still has several limitations. For example, while it is very effective for pretraining, it cannot transfer regularization HPs,5 so it’s generally not applicable to the finetuning of pretrained models. We discuss other limitations carefully in Section 5.1.
4but possibly not for different data and/or tasks. 5It can transfer regularization HPs to the extent they help training but it may not transfer their effect on
testing.
Our Contributions
• We demonstrate it is possible to zero-shot transfer near optimal HPs to a large model from a small version via the Maximal Update Parametrization (µP) from [45].
• While [45] only covered SGD, here we derive µP for Adam as well (Table 3). • We propose a new HP tuning technique, µTransfer, for large neural networks based on this
observation that provides massive speedup over conventional methods and covers both SGD and Adam training;
• We thoroughly verify our method on machine translation and large language model pretraining (in Section 6.3) as well as image classification (in Appendix I.1);
• We release a PyTorch [27] package for implementing µTransfer painlessly. A sketch of this package is given in Appendix J.
Terminologies Sometimes, to be less ambiguous, we often refer to the “large model” as the target model, as it is the model we wish to ultimately tune, while we refer to the “small model” as the proxy model, as it proxies the HP tuning process. We follow standard notation dmodel, dhead = dk, dv, nhead, dffn regarding dimensions in a Transformer; one can see Fig. 11 for a refresher.
Tensor Programs Series This paper is the 5th installment of the Tensor Programs series. While the target audience here are practitioners and empirical researchers, this paper presents the first major practical payoff of the theoretical foundation built in previous works [41–46].

heading: 2 Parametrization Matters: A Primer
text: In this section, we give a very basic primer on why the correct parametrization can allow HP transfer across width, but see Appendices L.1 to L.3 for more (mathematical) details.
The Central Limit Theorem (CLT) says that, if x1, . . . , xn are iid samples from a zero-mean, unitvariance distribution, then 1√
n (x1 + · · ·+ xn) converges to a standard Gaussian N (0, 1) as n→∞.
Therefore, we can say that 1√ n is the right order of scaling factor cn such that cn(x1 + · · · + xn) converges to something nontrivial. In contrast, if we set cn = 1/n, then cn(x1 + · · ·+ xn)→ 0; or if cn = 1, then cn(x1 + · · ·+ xn) blows up in variance as n→∞. Now suppose we would like to minimize the function
Fn(c) def = E
x1,...,xn f(c(x1 + · · ·+ xn)) (1)
over c ∈ R, for some bounded continuous function f : R→ R. If we reparametrize c = α/ √ n for
α ∈ R, then by CLT, Gn(α) def = Fn(c)→ E f(N (0, α2)) stabilizes into a function of α as n→∞. Then for sufficiently large n, the optimal α∗n def = arg minαGn(α) should be close to α ∗ N for any N > n, and indeed, for N =∞— this precisely means we can transfer the optimal c∗n or α∗n for a smaller problem (say Fn) to a larger problem (say FN ): GN is approximately minimized by α∗n and FN is approximately minimized by c∗n √ n/N . Because the transfer algorithm is simply copying α,
we say the parametrization c = α/ √ n is the correct parametrization for this problem.
In the scenario studied in this paper, x1, . . . , xn are akin to randomly initialized parameters of a width-n neural network, c is akin to a HP such as learning rate, and f is the test-set performance of the network after training, so that Fn gives its expectation over random initializations. Just as in this example, if we parametrize the learning rate and other HPs correctly, then we can directly copy the optimal HPs for a narrower network into a wide network and expect approximately optimal
performance — this is the hyperparameter transfer we propose here. It turns out the Maximal Update Parametrization (µP) introduced in [45] is correct (akin to the parametrization in α above), while the standard parametrization (SP) is incorrect (akin to the parametrization in c). We will review both parametrizations shortly. Theoretically, a µP network has a well-defined infinite-width limit — akin to (x1 + · · ·+ xn)/ √ n having aN (0, 1) limit by CLT — while a SP network does not (the limit will blow up) [45].6 In fact, based on the theoretical foundation laid in [45], we argue in Appendix L.3 that µP should also be the unique parametrization that allows HP transfer across width.
We emphasize that, to ensure transferability of any hyperparameter (such as learning rate), it’s not sufficient to reparametrize only that hyperparameter, but rather, we need to identify and correctly reparametrize all hyperparameters in Table 2. For example, in Fig. 1, the wide models in SP still underperform their counterparts in µP, even with learning rate tuned optimally. This is precisely because SP does not scale parameter multipliers and input/output layer learning rates correctly in contrast to µP (see Table 3). See Appendix C for more intuition via a continuation of our example here. We shall also explain this more concretely in the context of neural networks in Appendix B.

heading: 3 Hyperparameters Don’t Transfer Conventionally
text: In the community there seem to be conflicting assumptions about HP stability. A priori, models of different sizes don’t have any reason to share the optimal HPs. Indeed, papers aiming for stateof-the-art results often tune them separately. On the other hand, a nontrivial fraction of papers in deep learning fixes all HPs when comparing against baselines, which reflects an assumption that the optimal HPs should be stable — not only among the same model of different sizes but also among models of different designs — therefore, such comparisons are fair. Here, we demonstrate HP instability across width explicitly in MLP and Transformers in the standard parametrization. We will only look at training loss to exclude the effect of regularization.
MLP with Standard Parametrization We start with a 2-hidden-layer MLP with activation function φ, using the standard parametrization7 with LeCun initialization8 akin to the default in PyTorch:
f(ξ) = W 3>φ(W 2>φ(W 1>ξ + b1) + b2)
with init. W 1 ∼ N (0, 1/din), W {2,3} ∼ N (0, 1/n), b{1,2} = 0, (2)
where W 1 ∈ Rdin×n, b1 ∈ Rn, W 2 ∈ Rn×n, b2 ∈ Rn, W 3 ∈ Rn×dout and din, n, and dout are the input, hidden, and output dimensions. The particular MLP we use has φ = ReLU and a cross-entropy (xent) loss function. We define the width of MLP as the hidden size n, which is varied from 256 to 8192. The models are trained on CIFAR-10 for 20 epochs, which is more than enough to ensure convergence.
As shown on the left in Fig. 3, the optimal learning rate shifts by roughly an order of magnitude as the width increases from 256 to 8192; using the optimal learning of the smallest model
on the largest model gives very bad performance, if not divergence.
Transformer with Standard Parametrization This perhaps unsurprising observation holds for more complex architectures such as Transformer as well, as shown in Fig. 1 (left). We define width
6The more theoretically astute reader may observe that SP with a Θ(1/width) learning rate induces a well-defined infinite-width limit exists as well. Nevertheless, this does not allow HP transfer because this limit is in kernel regime as shown in [45]. See Appendix L.3 for more discussions.
7i.e. the default parametrization offered by common deep learning frameworks. See Table 3 for a review. 8The key here is that the init. variance ∝ 1/fan_in, so the same insights here apply with e.g. He initialization.
Table 3: µP[45] and SP for General Neural Networks, Basic Form. This basic form emphasizes the scaling with width (fan_in or fan_out); in practice, we may insert tunable multipliers in front of fan_in and fan_out as in Eq. (4). Notations: 1) η is the “master” learning rate. 2) The fan_out of a bias vector is its dimension (whereas fan_in is 1). 3) Purple text highlights key differences from standard parametrization (SP); Gray text recalls the corresponding SP. SGD (resp. Adam) here can be replaced by variants such as SGD with momentum (resp. Adagrad, Adadelta, etc). In general, the three columns here can be interpreted as linear layers that have {finite, infinite, infinite} input dimension and {infinite, finite, infinite} output dimension in an infinite-width network; this description generalizes more readily to other parameters such as those of layernorm. Transformer µP requires one more modification (1/d attention instead of 1/ √ d); see Definition 4.1. This version of µP gets rid of parameter multipliers; for the version similar to that in [45], see Table 13. Also see
as dmodel, with dk = dq = dv = dmodel/nhead and dffn = 4dmodel. The models are trained on wikitext-2 for 5 epochs. In Fig. 18 in the appendix we also show the instability of initialization scale and other HPs.

heading: 4 Unlocking Zero-Shot Hyperparameter Transfer with µP
text: We show that µP solves the problems we see in Section 3.
MLP with µP For the MLP in Section 3, to switch to µP, we just need to modify Eq. (2)’s initialization of the last layer and its learning rates of the first and last layer as well as of the biases. The basic form is9
initialize W 1 ∼ N (0, 1/din), W 2 ∼ N (0, 1/n), W 3 ∼ N (0, 1/n2), b{1,2} = 0 with SGD learning rates ηW 1 = ηb1 = ηb2 = ηn, ηW 2 = η, ηW 3 = ηn−1.
(3)
Here, η specifies the “master” learning rate, and we highlighted in purple the differences in the two parametrizations. This basic form makes clear the scaling with width n of the parametrization, but in practice we will often insert (possibly tune-able) multiplicative constants in front of each appearance of n. For example, this is useful when we would like to be consistent with a SP MLP at a base width n0. Then we may insert constants as follows: For ñ def = n/n0,
initialize W 1 ∼ N (0, 1/din), W 2 ∼ N (0, 1/n), W 3 ∼ N (0, 1/n·ñ), b{1,2} = 0 with SGD learning rates ηW 1 = ηb1 = ηb2 = ηñ, ηW 2 = η, ηW 3 = ηñ−1.
(4)
Then at width n = n0, all purple factors above are 1, and the parametrization is identical to SP (Eq. (2)) at width n0. Of course, as n increases from n0, then Eq. (4) quickly deviates from Eq. (2). In other words, for a particular n, µP and SP can be identical up to the choice of some constants (in this case n0), but µP determines a different “set" of networks and optimization trajectory than SP as one varies n. As we will see empirically in the next section, this deviation is crucial for HP transfer.
Indeed, in Fig. 3(right), we plot the CIFAR10 performances, over various learning rates and widths, of µP MLPs with n0 = 128. In contrast to SP, the optimal learning rate under µP is stable. This means that, the best learning rate for a width-128 network is also best for a width-8192 network in µP — i.e. HP transfer works — but not for SP. In addition, we observe performance for a fixed learning rate always weakly improves with width in µP , but not in SP.
This MLP µP example can be generalized easily to general neural networks trained under SGD or Adam, as summarized in Table 3, which is derived in Appendix L.
9While superficially different, this parametrization is equivalent to the µP defined in [45].
Transformers with µP We repeat the experiments with base width n0 = 128 for Transformers: Definition 4.1. The Maximal Update Parametrization (µP) for a Transformer is given by Table 3 and 1/d attention instead of 1/ √ d, i.e. the attention logit is calculated as q>k/d instead of q>k/ √ d where query q and key k have dimension d.10
The results are shown on the right in Fig. 1, where the optimal learning rate is stable, and the performance improves monotonically as width increases.

heading: 5 Which Hyperparameters Can Be µTransferred?
text: In this section, we explore how common HPs fit into our framework. In general, they can be divided into three kinds, summarized in Table 1:
1. those that can transfer from the small to the large model, such as learning rate (Table 2);
2. those that primarily control regularization and don’t work well with our technique; and
3. those that define training scale, such as width as discussed above as well as others like depth and batch size, across which we transfer other HPs.
Those in the first category transfer across width, as theoretically justified above in Section 2. To push the practicality and generality of our technique, we empirically explore the transfer across the other dimensions in the third category. Note that µTransfer across width is quite general, e.g. it allows varying width ratio of different layers or number of attention heads in a Transformer; see Appendix G.2. This will be very useful in practice. For the second category, the amount of regularization (for the purpose of controlling overfitting) naturally depends on both the model size and data size, so we should not expect transfer to work if the parametrization only depends on model size. We discuss these HPs in more detail in Appendix G.1.
10This is roughly because during training, q and k will be correlated so q>k actually scales like d due to Law of Large Numbers, in contrast to the original motivation that q, k are uncorrelated at initialization so Central Limit applies instead. See Appendix L.2.1 for a more in-depth discussion.

heading: 5.1 Empirical Validation and Limitations
text: Our empirical investigations focus on Transformers (here) and ResNet (in Appendix I.1.1), the most popular backbones of deep learning models today. We train a 2-layer pre-layernorm µP11 Transformer with 4 attention heads on Wikitext-2. We sweep one of four HPs (learning rate, output weight multiplier, initialization standard deviation, and learning rate schedule) while fixing the others and sweeping along width and depth (with additional results in Fig. 19 on transfer across batch size, sequence length, and training time). Fig. 4 shows the results averaged over 5 random seeds.
Empirically, we find that for language modeling on Transformers, HPs generally transfer across scale dimensions if some minimum width (e.g. 256), depth (e.g., 4), batch size (e.g., 32), sequence length (e.g., 128), and training steps (e.g., 5000) are met, with some caveats discussed below. While the exact optimum can shift slightly with increasing scale, this shift usually has very small impact on the loss, compared to SP (Figs. 1 and 3(left)). However, there are some caveats. For example, the best initialization standard deviation does not seem to transfer well across depth (2nd row, 3rd column), despite having a stabler optimum across width. In addition, while our results on width, batch size, sequence length, and training time still hold for post-layernorm (Fig. 17),12 the transfer across depth only works for pre-layernorm Transformer. Nevertheless, in practice (e.g. our results in Section 6.3) we find that fixing initialization standard deviation while tuning other HPs works well when transferring across depth.

heading: 6 Efficiency and Performance of µTransfer
text: Now that the plausibility of µTransfer has been established in toy settings, we turn to more realistic scenarios to see if one can achieve tangible gains. Specifically, we perform HP tuning only on a smaller proxy model, test the obtained HPs on the large target model directly, and compare against baselines tuned using the target model. We seek to answer the question: Can µTransfer make HP tuning more efficient while achieving performance on par with traditional tuning? As we shall see by the end of the section, the answer is positive. We focus on Transformers here, while experiments on ResNets on CIFAR10 and Imagenet can be found as well in Appendix I.1. All of our experiments are run on V100 GPUs.

heading: 6.1 Transformer on IWSLT14 De-En
text: Setup IWSLT14 De-En is a well-known machine translation benchmark. We use the default IWSLT (post-layernorm) Transformer implemented in fairseq [25] with 40M parameters, which we denote as the 1x model.13 For µTransfer, we tune on a 0.25x model with 1/4 of the width, amounting to 4M parameters. For this experiment, we tune via random search the learning rate η, the output layer parameter multiplier αoutput, and the attention key-projection weight multiplier αattn. See the grid and other experimental details in Appendix H.1.
We compare transferring from the 0.25x model with tuning the 1x model while controlling the total tuning budget in FLOPs.14 To improve the reproducibility of our result: 1) we repeat the entire HP search process (a trial) 25 times for each setup, with number of samples as indicated in Table 4, and report the 25th, 50th, 75th, and 100th percentiles in BLEU score; 2) we evaluate each selected HP combination using 5 random initializations and report the mean performance.15
We pick the HP combination that achieves the lowest validation loss16 for each trial. The reported best outcome is chosen according to the validation loss during tuning. We compare against the default in fairseq, which is presumably heavily tuned. The result is shown in Table 4.
11“2 layers” means the model has 2 self-attention blocks. To compare with SP Transformer, see Fig. 18. 12in fact, post-layernorm Transformers are much more sensitive to HPs than pre-layernorm, so our technique is more crucial for them, especially for transfer across width. Fig. 1 uses post-layernorm. 13https://github.com/pytorch/fairseq/blob/master/examples/translation/README.md. 14Ideally we would like to measure the wall clock time used for tuning. However, smaller models such as the proxy Transformer used for IWSLT are not efficient on GPUs, so wall clock time would not reflect the speedup for larger models like GPT-3. Thus, we measure in FLOPs, which is less dependent on hardware optimization.
15We do not report the standard deviation over random initializations to avoid confusion. 16We find this provides more reliable result than selecting for the best BLEU score.
4 3 2 1 0 1 2 3 log2Compute
34.4
34.6
34.8
35.0
35.2
35.4
BL EU
S co
re
Method Ours Conventional
10 20 30 40 50 60 # Samples
Method Ours Conventional
Figure 5: Efficiency-performance Pareto frontier of µTransfer compared to conventional tuning, on IWSLT Transformer, using random HP search as the base method. We plot the median BLEU score over 25 trials (Left) against relative compute budget in log scale and (Right) against number of HP samples taken. While with the same number of samples, µTransfer slightly underperforms conventional tuning, this gap vanishes with more samples, and in terms of compute, our Pareto frontier strongly and consistently dominates that of conventional tuning. Note that, in larger models (e.g. BERT or GPT-3, not shown here), we believe our efficiency advantage will only widen as our small proxy model can stay the same size while the target model grows.
Performance Pareto Frontier The result above only describes a particular compute budget. Is µTransfer still preferable when we have a lot more (or less) compute? To answer this question, we produce the compute-performance Pareto frontier in Fig. 5(left), where we repeat the above experiment with different compute budgets. Evidently, our approach completely dominates conventional tuning.
Sample Quality of Proxy Model vs Target Model The Pareto frontier in Fig. 5(right) suggests that, given a fixed number of random samples from the HP space, 1) tuning the target model directly yields slightly better results than tuning the proxy model (while taking much more compute of course), but 2) this performance gap seems to vanish as more samples are taken. This can be explained by the intuition that the narrower proxy model is a “noisy estimator” of the wide target model [45].With few samples, this noise can distort the random HP search, but with more samples, this noise is suppressed.

heading: 6.2 Transformer on WMT14 En-De
text: We scale up to WMT14 En-De using the large (post-layernorm) Transformer from [37] with 211M parameters. We tune on a proxy model with 15M parameters by shrinking dmodel, dffn, and nhead. For this experiment, we tune via random search the learning rate η, the output layer parameter multiplier αoutput, and the attention key-projection weight multiplier αattn following the grid in Appendix H.2. The result is shown in Table 5: While random search with 3 HP samples far underperforms the fairseq default, we are able to match it via transfer using the same tuning budget.

heading: 6.3 BERT
text: Finally, we consider large-scale language model pretraining where HP tuning is known to be challenging. Using Megatron (pre-layernorm) BERT [32] as a baseline, we hope to recover the performance of the published HPs by only tuning a proxy model that has roughly 13M parameters, which we call BERT-prototype. While previous experiments scaled only width, here we will also scale depth, as discussed in Section 5 and validated in Fig. 4. We use a batch size of 256 for all runs and follow the standard finetuning procedures. For more details on BERT-prototype, what HPs we tune, and how we finetune the trained models, see Appendix H.3.
During HP tuning, we sample 256 combinations from the search space and train each combination on BERT-prototype for 105 steps. The total tuning cost measured in FLOPs is roughly the same as training 1 BERT-large for the full 106 steps; the exact calculation is shown in Appendix H.3. The results are shown in Table 6. Notice that on BERT-large, we obtain sizeable improvement over the well-tuned Megatron BERT-large baseline.

heading: 6.4 GPT-3
text: In order to further verify µTransfer at scale, we applied it to GPT-3 6.7B [6]. This Transformer model (the target model) consists of 32 residual blocks with width 4096. We form the small proxy model by shrinking width to 256, resulting in roughly 40 million trainable parameters, 168 times smaller than the target model. HPs were then determined by a random search on the proxy model. The total tuning cost was only 7% of total pretraining cost. Details of the HP sweep can be found in Appendix H.4.
In order to exclude code difference as a possible confounder, we also re-trained GPT-3 6.7B from scratch using the original HPs from [6]. During training of the µTransfer model we encountered numerical issues that lead to frequent divergences. In order to avoid them, the model was trained using FP32 precision, even though the original 6.7B model and our re-run were trained using FP16.17 18 The resulting µTransfer model outperforms the 6.7B from [6], and is in fact comparable to the twice-as-large 13B model across our evaluation suite (see Table 9). Selected evaluation results can be found in Table 7 and further details are given in Table 8 and Appendix H.4.

heading: 7 Related Works
text: Hyperparameter Tuning Many have sought to speed up HP tuning beyond the simple grid or random search, such as via Bayesian optimization [34, 35] or multi-arm bandits [15, 18]. There are also dedicated tools such as Optuna [4] and Talos [3] which integrate with existing deep learning frameworks and provide an easy way to apply more advanced tuning techniques. Our work is
17While we are mainly focused on the efficacy of µTransfer regardless of precision, it would be interesting to ablate the effect of precision in our results, but we did not have enough resources to rerun the baseline in FP32
18It is quite interesting that µTransfer identified a useful region of hyperparameters leading to much improved performance, which probably would be difficult to discover normally because 1) researchers usually change hyperparameters to accomodate precision and 2) there was no precise enough justification to go against this judgment until µTransfer.
complementary to the above, as they can be used to tune the proxy model. it is only for scientific reasons that we primarily did random search throughout this work.
Hyperparameter Transfer Many previous works explored transfer learning of HP tuning (e.g. [12, 28, 36, 49]). However, to the best of our knowledge, our work is the first to explore zero-shot HP transfer. In addition, we focus on transferring across model scale rather than between different tasks or datasets. Some algorithms like Hyperband [19] can leverage cheap estimates of HP evaluations (like using a small model to proxy a large model) but they are not zero-shot algorithms, so would still be very expensive to apply to large model training. Nevertheless, all of the above methods are complementary to ours as they can be applied to the tuning of our proxy model.
Previously Proposed Scaling Rules of Hyperparameters [11, 23, 31, 33] investigated the right way to scale learning rate with batch size, with sometimes conflicting proposals. which we summarize in Appendix F. [26] studied how learning rate (and batch size) should scale with width for MLPs and CNNs trained with SGD in NTK or standard parametrizations. We provide a detailed comparison of our work with theirs in Appendix F.
Many previous works proposed different initialization or parametrizations with favorable properties, such as better stability for training deep neural networks [5, 10, 13, 21, 30, 47, 48, 51]. Our work differs from these in that we focus on the transferability of optimal HPs from small models to large models in the same parametrization.

heading: 8 Conclusion
text: Leveraging the discovery of a feature learning neural network infinite-width limit, we hypothesized and verified that the HP landscape across NNs of different width is reasonably stable if parametrized according to Maximal Update Parametrization (µP). We further empirically showed that it’s possible to transfer across depth, batch size, sequence length, and training time, with a few caveats. This allowed us to indirectly tune a very large network by tuning its smaller counterparts and transferring the HPs to the full model.
Venues of Improvement Nevertheless, our method has plenty of room to improve. For example, initialization does not transfer well across depth, and depth transfer generally still does not work for post-layernorm Transformers. This begs the question whether a more principled parametrization in depth could solve these problems. Additionally, Fig. 4 shows that the optimal HP still shifts slightly for smaller models. Perhaps by considering finite-width corrections to µP one can fix this shift. Finally, it will be interesting to consider if there’s a way to transfer regularization HPs as a function of both the model size and data size.
Broader Impact Our work makes HP tuning of large models more efficient. This enables large models to be better tuned given the same compute budget, thereby increasing the performance per cost. Organizations large and small can focus their research on small models and scale up only once with reasonable confidence that the training would go well. We do not foresee any direct negative societal impact.

Review summaries:
The paper proposes an alternative method for tuning hyperparameters which is to tune hyperparameters on a small width model and transfer these sets of hyperparameters to the full network. The core idea is to utilize the maximal update parameterization so that optimization hyperparameters remain stable under model size changes. This enables a zero-shot transfer and makes the hyperparameter optimization process more efficient both in computation and memory. Empirically, the authors show that the careful hyperparameter selection on the small model with the proposed approach outperforms BERT-large on MNLI and OOP and also requires less total training costs.

Summary:  This paper makes the observation that when initializing neural network with the *maximal   update parametrization* technique, certain hyperparameters such as learning rate and learning rate schedule are transferrable across different scales of the network (e.g. width of each layer). This insight enables to authors to tune hyperpameters using a small network and directly transfer the tuned results to much larger networks, saving the cost of tuning larger networks. Although such tuning-and-transfer approach is probably already a standard practice, the authors point out that such practice is not optimal if we initialize neural networks using the standard technique and using the proposed method generally leads to much better transfer result.   I find this work to be interesting and see the potential of this submission to inspire future work. That being said, I want to point out that the *maximal update parametrization* technique is already published in a previous work and this submission should be viewed more as an add-on to this prior work. Practically speaking, I should also note that not all hyperparameters are transferrable as pointed out by the authors candidly. Certain hyperpameters such as regularization strength are not transferrable across scale. Such hyperparameters may have complex interactions with other hyperparameters and therefore make the proposed method less useful. Given these two reasons, I am leaning toward accept but have some reservation. I detail my other concerns in the following section.

Based on recently proposed Maximal Update Parametrization ($\mu$P), this paper finds that hyperparameters can be copied (i.e. transferred with zero cost) across model sizes. Therefore, it is possible to tune hyperparameters in a small model, and to apply hyperparameters in a large model.