Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer

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 can not 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 ] . 2 Parametrization Matters : A Primer . 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 . 3 Hyperparameters Don ’ t Transfer Conventionally . 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 .

Summary: 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.