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

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 $d_{model},d_{head}=d_{k},d_{v},n_{head},d_{ffn}$ regarding dimensions in a Transformer; one can see Fig. 11 for a refresher.

This paper is the 5th installment of the Tensor Programs series. While it is self-contained with the target audience being practitioners and empirical researchers, this paper presents the first major practical payoff of the theoretical foundation built in previous works [54, 53, 55, 58, 56, 57].

We start with a 2-hidden-layer MLP with activation function $\phi$, using the standard parametrization$^{4}$ with LeCun initialization$^{5}$ akin to the default in PyTorch:

where $W^{1}\in\mathbb{R}^{d_{in}\times n},b^{1}\in\mathbb{R}^{n}$, $W^{2}\in\mathbb{R}^{n\times n},b^{2}\in\mathbb{R}^{n}$, $W^{3}\in\mathbb{R}^{n\times d_{out}}$ and $d_{in}$, $n$, and $d_{out}$ are the input, hidden, and output dimensions. The particular MLP we use has $\phi=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.

This perhaps unsurprising observation holds for more complex architectures such as Transformer as well, as shown in Fig. 1 (left). We define width as $d_{model}$, with $d_{k}=d_{q}=d_{v}=\nicefrac{{d_{model}}}{{n_{head}}}$ and $d_{ffn}=4d_{model}$. 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.

For the MLP in Section 3, to switch to $\mu$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 is$^{6}$

Here, $\eta$ 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 $n_{0}$. Then we may insert constants as follows: For $\tilde{n}\mathbin{\overset{\mathrm{def}}{=}}n/n_{0}$,

Then at width $n=n_{0}$, all purple factors above are 1, and the parametrization is identical to SP (Eq. 2) at width $n_{0}$. Of course, as $n$ increases from $n_{0}$, then Eq. 4 quickly deviates from Eq. 2. In other words, for a particular $n$, $\mu$P and SP can be identical up to the choice of some constants (in this case $n_{0}$), but $\mu$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 $\mu$P MLPs with $n_{0}=128$. In contrast to SP, the optimal learning rate under $\mu$P is stable. This means that, the best learning rate for a width-128 network is also best for a width-8192 network in $\mu$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 $\mu$P , but not in SP.

This MLP $\mu$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 J.

We repeat the experiments with base width $n_{0}=128$ for Transformers:

The results are shown on the right in Fig. 1, where the optimal learning rate is stable, and the performance improves monotonically with width. See Appendix B for further explanation of $\mu$P.

As shown in [57] and Section J.3, in SP, the network output will blow up with width after 1 step of SGD. It’s instructive to consider a 1-hidden-layer linear perceptron $f(x)=V^{\top}Ux$ with scalar inputs and outputs, as well as weights $V,U\in\mathbb{R}^{n\times 1}$. In SP, $V_{\alpha}\sim\mathcal{N}(0,1/n)$ ad $U_{\alpha}\sim\mathcal{N}(0,1)$ for each $\alpha\in[n]$. This sampling ensures that $f(x)=\Theta(|x|)$ at initialization. After 1 step of SGD with learning rate 1, the new weights are $V^{\prime}\leftarrow V+\theta U,U^{\prime}\leftarrow U+\theta V$, where $\theta$ is some scalar of size $\Theta(1)$ depending on the inputs, labels, and loss function. But now

blows up with width $n$ because $U^{\top}U=\Theta(n)$ by Law of Large Numbers.

Now consider the same network in $\mu$P. According to Table 3, we now have $V_{\alpha}\sim\mathcal{N}(0,1/n^{2})$ in contrast to SP, but $U_{\alpha}\sim\mathcal{N}(0,1)$ as before, with learning rates $\eta_{V}=1/n,\eta_{U}=n$. After 1 step of SGD, we now have

and one can verify this is $\Theta(1)$ and thus does not blow up with width.$^{8}$

One can observe the same behavior in more advanced architectures like Transformers and optimizers like Adam; in fact, in SP, other hidden quantities like attention logits will also blow up with width after 1 step, but in $\mu$P still remain bounded, as shown in Fig. 5(middle).

One might think scaling down the learning rate with width can solve this problem in SP. However, other hidden activations like the word embedding (Fig. 5(right)) in a Transformer update by a width-independent amount for each step of training, so scaling down the learning rate will effectively mean the word embeddings are not learned in large width models. Similar conclusions apply to other models like ResNet (in fact, one can observe in the SP linear MLP example above, the input layer is updated much more slowly than the output layer). On the other hand, $\mu$P is designed so that all hidden activations update with the same speed in terms of width (see Section J.2 for why).

This is why a wide model tuned with $\mu$Transfer should in general outperform its SP counterpart with (global) learning rate tuned. For example, this is the case for the width-8192 Transformer in Fig. 1, where, in SP, the optimal learning rate needs to mollify the blow-up in quantities like logits and attention logits, but this implies others like word embeddings do not learn appreciably. This performance advantage means $\mu$Transfer does more than just predicting the optimal learning rate of wide SP models. Relatedly, we observe, for any fixed HP combination, training performance never decreases with width in $\mu$P, in contrast to SP (e.g., the $\mu$P curves in Figs. 1, 3 and 16 do not cross, but the SP curves do; see also Section 8).

Our empirical investigations focus on Transformers (here) and ResNet (in Section G.1.1), the most popular backbones of deep learning models today. We train a 2-layer pre-layernorm $\mu$P$^{9}$ 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, and the target scale is within the “reasonable range” as in our experiments. Now, there are some caveats. 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),$^{10}$ the transfer across depth only works for pre-layernorm Transformer. Nevertheless, in practice (e.g. our results in Section 7.3) we find that fixing initialization standard deviation while tuning other HPs works well when transferring across depth.

We scale up to WMT14 En-De using the large (post-layernorm) Transformer from [50] with 211M parameters. We tune on a proxy model with 15M parameters by shrinking $d_{model}$, $d_{ffn}$, and $n_{head}$. For this experiment, we tune via random search the learning rate $\eta$, the output layer parameter multiplier $\alpha_{output}$, and the attention key-projection weight multiplier $\alpha_{attn}$ following the grid in Section F.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.

Finally, we consider large-scale language model pretraining where HP tuning is known to be challenging. Using Megatron (pre-layernorm) BERT [43] 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 6 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 Section F.3.

During HP tuning, we sample 256 combinations from the search space and train each combination on BERT-prototype for $10^{5}$ steps. The total tuning cost measured in FLOPs is roughly the same as training 1 BERT-large for the full $10^{6}$ steps; the exact calculation is shown in Section F.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.

In order to further verify $\mu$Transfer at scale, we applied it to GPT-3 6.7B [7] with relative attention. This 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 Section F.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 [7]. Unfortunately, after we have finished all experiments, we found this baseline mistakenly used absolute attention (like models in [7]) when it was supposed to use relative attention like the target model. In addition, during training of the $\mu$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.$^{15}$ $^{16}$ The resulting $\mu$Transfer model outperforms the 6.7B from [7], and is in fact comparable to the twice-as-large 13B model across our evaluation suite (see Table 11). Selected evaluation results can be found in Table 7 and further details are given in Table 10 and Section F.4.

Many have sought to speedup HP tuning beyond the simple grid or random search. Snoek et al. 45 treated HP tuning as an optimization process and used Bayesian optimization by treating the performance of each HP combination as a sample from a Gaussian process (GP). Snoek et al. 46 further improved the runtime by swapping the GP with a neural network. Another thread of work investigated how massively parallel infrasture can be used for efficient tuning under the multi-arm bandit problem [18, 22]. 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 approach is distinct from all of the above in that it does not work on the HP optimization process itself. Instead, it decouples the size of the target model from the tuning cost, which was not feasible prior to this work. This means that no matter how large the target model is, we can always use a fixed-sized proxy model to probe its HP landscape. Nevertheless, our method is complementary, as the above approaches can naturally be applied to the tuning of the proxy model; it is only for scientific reasons that we use either grid search or random search throughout this work.

Many previous works explored transfer learning of HP tuning (e.g. [62, 36, 47, 15]). 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 [23] 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.

[57] introduced $\mu$P as the unique parametrization that enables all layers of a neural network to learn features in the infinite-width limit, especially in contrast to the NTK parametrization [17] (which gives rise to the NTK limit) that does not learn features in the limit. Based on this theoretical insight, in Section J.3, we argue that $\mu$P should also be the unique parametrization (in the sense of [57]) that allows HP transfer across width; in short this is because it both 1) preserves feature learning, so that performance on feature learning tasks (such as language model pretraining) does not become trivial in the limit, and 2) ensures each parameter tensor is not stuck at initialization in the large width limit, so that its learning rate does not become meaningless. At the same time, our results here suggest that $\mu$P is indeed the correct parametrization for wide neural networks and thus provide empirical motivation for the theoretical study of the infinite-width $\mu$P limit. Note, parametrization here refers to a rule to scale hyperparameters with width (“how should my initialization and learning rate change when my width doubles?”), which is coarser than a prescription for setting hyperparameters at any particular width (“how should I set my initialization and learning rate at width 1024?”).

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 $\mu$P one can fix this shift. Finally, it will be interesting to study if there’s a way to transfer regularization HPs as a function of both the model size and data size, especially in the context of finetuning of pretrained models.

In alphabetical order, we thank Arthur Jacot, Arturs Backurs, Colin Raffel, Denny Wu, Di He, Huishuai Zhang, Ilya Sutskever, James Martens, Janardhan Kulkarni, Jascha Sohl-Dickstein, Jeremy Bernstein, Lenaic Chizat, Luke Metz, Mark Chen, Michael Santacroce, Muhammad ElNokrashy, Pengchuan Zhang, Sam Schoenholz, Sanjeev Arora, Taco Cohen, Yiping Lu, Yisong Yue, and Yoshua Bengio for discussion and help during our research.

We can classify any dimension in a neural network as “infinite” if it scales with width, or “finite” otherwise. For example, in a Transformer, $d_{model},d_{ffn},d_{head},n_{head}$ are all infinite, but vocab size and context size are finite. Then we can categorize parameter tensors by how many infinite dimensions they have. If there are two such dimensions, then we say the parameter is matrix-like; if there is only one, then we say it is vector-like; if there is none, we say it is scalar-like. Then in Tables 3, 8 and 9, “input weights & all biases” and “output weights” are all vector-like parameters, while hidden weights are matrix-like parameters. An advantage of Table 8 is that it gives a uniform scaling rule of initialization and learning rate for all vector-like parameters. The multiplier rule in Table 8 can be more interpreted more generally as the following: a multiplier of order $1/\mathrm{fan\_in}$ should accompany any weight that maps an infinite dimension to a finite one. This interpretation then nicely covers both the output logits and the attention logits (i.e. $1/d$ attention).

Scalar-like parameters are not as common as matrix-like and vector-like ones, but we will mention a few examples in Section B.2. The scaling rule for their initialization, learning rate (for both SGD and Adam), and multiplier is very simple: hold them constant with width.

We did not specify the initialization mean in the tables, since most commonly the mean is just set to 0, but it can be nonzero for vector-like parameters (e.g., layernorm weights) and scalar-like parameters but must be 0 for matrix-like parameters.

The initialization scaling rules in our tables can all be trivially satisfied if the initialization variance is set to 0. This can be useful in some settings (e.g., Section D.2) but detrimental in other settings (e.g., hidden weights).

Here, input weights very specifically refer to weights that map from an infinite dimension to a finite dimension. As a counterexample, in some architectures, the first layer can actually map from a finite dimension to another finite dimension, e.g., a PCA layer. Then this is not an “input weight”; if the next layer maps into an infinite dimension, then that’s the input weight. A similar, symmetric discussion applies to output weights.

For our tables, a model is specifically a function that maps a finite dimension to another finite dimension, consistent with the discussion above. For example, for an image model on CIFAR10, it maps from $3\times 32\times 32=3072$ dimensions to 10 dimensions, and these numbers are fixed regardless of the width of the model. Likewise, for an autoregressive Transformer model, the input and output dimension are both the vocab size, which is independent of the width. In contrast, an MLP inside a Transformer is not a “model” in this sense because its input and output dimension are both equal to the width of the Transformer.

To ground the abstract description in Tables 3, 8 and 9, we walk through the parameters of a typical Transformer and discuss concretely how to parametrize each.

We assume that the user wants to replicate SP when the model widths are equal to some base widths, for example, when $d_{model}=d_{model,0}=128,d_{ffn}=d_{ffn,0}=512$, etc, as in the MLP example in Section 4. For this purpose, it’s useful to define $\tilde{d}_{model}=d_{model}/d_{model,0},\ \tilde{d}_{ffn}=d_{ffn}/d_{ffn,0}$, and so on. One can always take $d_{model,0}=d_{ffn,0}=\cdots=1$ for a “pure” $\mu$P.

Below, we introduce hyperparameters $\sigma_{\bullet},\eta_{\bullet}$ for each parameter tensor, as well as a few multipliers $\alpha_{\bullet}$. One may always tie $\sigma_{\bullet}$ (resp. $\eta_{\bullet}$) across all parameter tensors, but in our experiments, we found it beneficial to at least distinguish the input and output layer initialization and learning rates.