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