Depthwise Hyperparameter Transfer in Residual Networks: Dynamics and Scaling Limit

We thank the reviewer for their careful reading and valuable feedback, to which we respond below.

Weakness 1 = Gap 1

Our article neither seeks nor claims to reduce the computational cost of training a large model with a fixed setting of hyperparameters. The core contribution of our work is a methodology for significantly recuding the cost of hyperparameter optimization itself.

Our claim was that this could be achieved by using our parameterization to optimize hyperparameters in much smaller models. We indeed found in all of our experiments that optimal hyperparameters settings were stable over many orders of magnitude with our procedure. To validate this empirically, however, it was necessary to train rather large models.

We certainly agree that it is desirable to obtain a target model accuracy while minimizing compute. However, our experiments consistently showed that larger/deeper models result in better performance. This is consistent with the literature on scaling laws and with the common practice of training very large models to obtain best SOTA performace. That said, the marginal gain from using larger models is sometimes very small (especially when very close to the limiting dynamics) and perhaps not worth the computational cost. An understanding of when this happens is not something we know how to predict.

Weakness 2 = Gap 2

The reviewer is entirely correct that we've presented in the main text the DMFT equations using optimization with gradient flow, rather than discrete time (S)GD. However, as in the prior work on DMFT (Bordelon & Pehlevan, 2022), discrete time SGD also leads to a consistent scaling limit for training dynamics and accompanying DMFT equations. These are included in the Appendix (now App. K.4). The discrete time dynamics of $h, g$ are almost identical to their continuous time counterparts but with sums instead of integrals over time. The dynamics for the predictor is no longer governed by the dynamical NTK but rather is given by the last layer feature kernel and the final layer response function (see equation 131).

Weakness 3 = CIFAR10 Performance

The reviewer is right that in Figure 1 the best learning rates found in our proposed parameterization reach somewhat higher values of training loss than the best learning rates we found with a vanilla $\mu P$ parameterization. Thank you for suggesting that we address this issue directly.

We have done so in our revision by emphasizing in the introduction that the discrepancy in training loss is actually predicted by our theory. Indeed, note that the $\mu P$ runs diverged at large depth (that is why the orange line corresponding to depth 30 doesn't appear at all in Figure 1(a)). Our theory predicts that increasing depth at a fixed learning rate in the $\mu P$ parameterization corresponds to an effectively higher learning rate. Thus, since our experiments were run for a constant but moderate number of epochs, $\mu P$ runs result in models that are closer to convergence. This is true until the effective learning rate is so large that the runs diverge.

We also note that Figure 1 seeks to illustrate hyperparameter transfer, which clearly fails for $\mu P$ but works well with our parameterization. In particular, we didn't attempt to obtain the best possible test loss for each architecture, which would have requred optimizing other hyperparameters (e.g. notably the lr schedule, how to do data augmention, momentum, batch size etc). In this sense, we do not believe that the discrepancy in test loss in Figure 1 takes away from the merit of our core proposal.

We hope that this addresses the reviewer's concern and welcome any further feedback on this point. Also, we are in the process of running several followup experiments to illustrate the first bullet point and will post then once they have finished running.

Weakness 4 = Literature

We thank the reviewer for their valuable suggestions to improve our coverage of the literature and have made the following improvements in the revision:

• We have added Appendix C outlining the original work on $\mu P$ and have included a discussion of maximal vs. stable updates as well as alternative parameterizations.
• We have included a new Appendix section (App. M) discussing other scalings which give stable feature learning infinite depth limits. We also provide more experiments showing hyperparameter transfer for zero-init block readouts (similar to rezero) and $\text{depth}^{-\alpha}$ scalings for residual branches with $\alpha=\frac{1}{2},1$. The $\alpha=1$ case gives a Neural ODE like infinite depth limit rather than the SDE-like limit at $\alpha=1/2$ (which is the focus of the rest of the paper). We hope this shows that our general methods can be applicable to derive and characterize many large depth limits which can show transfer.
• We have added to the Related Works section in our revision references to recent work on using adaptive kernel methods.

I thank the authors for addressing my comments.

Regarding the authors' response to "Gap I", I think there may have been some misunderstanding worth clearing.

My main focus in this gap is not on the matter of compute. Rather it is on the fact that, as I understand it, in the scaling limit the authors offer the DMFT equations become independent of depth and width. This seems to imply that performance should also be scale invariant. In contrast, however, in their experiments, where they employ this scaling the performance does change. For instance, in Fig. 1b and width=1024, train loss changes from about 0.3 to 0.2 as depth increases. Shouldn't this number be fixed if one is indeed in the asymptotic scaling regime? This seems to mean that in the typical regime of interest (a regime where increasing scale/compute improves performance) sizable corrections to the asymptotic still exist. These corrections could, in principle, scale differently with hyper-parameters thereby augmenting the optimal transfer of parameters based on keeping the asymptotic theory the same.

Notwithstanding, while I'm curious about this I don't consider this a major drawback. Potentially what could be going on is that the asymptotic theory is stable/invariant under scaling in early training stages with corrections to the asymptotic entering in only at later times. If so one can base the logic here on the common practice of tuning parameters based on early epochs. This is roughly consistent with the authors' Supp. figure B3a where the difference between the original and scale-up network becomes larger only at later epochs.

We thank the reviewer for their careful reading and valuable feedback, to which we respond below. We hope in light of our responses the reviewer will consider raising their score.

Weakness 1

We agree that it makes sense to add a brief description of the $\mu P$ parameterization. Thank you for the suggestion. This is now done in the Appendix C and C.2.

Weakness 2

Regarding the $1/\sqrt{\mathrm{depth}}$ scaling of residual branches:

• We agree that our work is certainly not the first to propose the $1/\sqrt{\text{depth}}$ scaling of residual branches and have emphasized this in our revised review of literature.
• As far as we are aware, ours is the first work which shows that this is precisely the residual branch scaling that gives rise to training dynamics that are both consistent across depth/width and also are capable of feature learning.
• We have included a new Appendix section (App. M) discussing other possible large depth limits which give stable feature learning infinite depth limits. We also provide more experiments showing hyperparameter transfer for zero-initialized block readouts (similar to the initialization of re-zero) and $\text{depth}^{-\alpha}$ scalings for residual branches with $\alpha = \frac{1}{2},1$. See Figure M.1. The $\alpha=1$ case gives a Neural ODE like infinite depth limit rather than the SDE-like limit at $\alpha=1/2$ (which is the focus of the rest of the paper). We hope this shows that our general methods can be applicable to derive and characterize many large depth limits which can exhibit transfer. We note that even with the zero initialization trick, the depth exponent $\alpha$ of the model is still important in controlling the training dynamics, with some different behaviors for different values of $\alpha$. We plan to explore this in greater detail in future work.

Weakness 3

We are starting to run experiments where we train models for longer and still see good hyperparameter transfer and larger "deeper is better" effects. We plan on adding additional experiments where we train these models longer (see for instance some additional experiments here https://imgur.com/a/GA6qepR).

In our original submission, we chose to train for 10 - 20 epochs in most of our experiments for these reasons:

• Our experiments were rather computationally expensive, at least with the compute resources available to us because
  • The focus of our work is hyperparameter transfer over both depth and width. We were thus bottlenecked by the computational cost of training our largest (deepest and widest) architectures, which in many experiments were Convolutional ResNets with depth 34, width 2048 and roughly 50M parameters.
  • Each experiment involved re-training the same architecture over many -- often 10 or more -- different learning rates (or other hyperparameter settings)
• We believe directly checking hyperparameter transfer after model convergence would likely require an entirely different suite of experiments, rather than simply training longer. Indeed, hyperparameter transfer for models after covergence is most interesting only after optimizing over not only the learning rate but also the learning rate schedule, the batch size, momentum coefficient, data augmention, etc. We have presented a fair amount of evidence indicating that in our proposed paramaterization all these hyperparameters individually transfer (see Figure 3 for learning rate schedules and normalization layer and Figure 4 for momentum and feature learning rate). But running large scans jointly over all these parameters was not computationally feasible.

Further Citations

We thank the reviewer for citations [1]-[3]. We have now added [2] to our list of references ([1] and [3] were already present). We have also included in our revised Related Works section the point that the main difference between [1]-[3] and our work is threefold:

• our article derives DMFT equations that govern training dynamics, rather than only the behavior at initialization
• our analysis is combined $\mu P$-type initialization in the final layer, whereas most of this prior work considered standard parameterization
• we focus on empirical investigation of whether hyperparameters transfer in this parameterization, while prior work was mainly concerned with the stability of the forward and backward pass at initialization

We thank the reviewer for their careful reading and valuable feedback, to which we respond below.

Weakness 1

The reviewer is certainly right that in Figure 1 the best learning rates found in our proposed parameterization reach somewhat higher values of training loss than the best learning rates we found with a vanilla $\mu P$ parameterization. Thank you for suggesting that we address this issue.

We have done so in our revision by emphasizing in the introduction that the discrepancy in training loss is actually predicted by our theory. Indeed, note that $\mu P$ runs diverged at large depth (that is why the orange line corresponding to depth 30 doesn't appear at all in Figure 1(a)). Our theory predicts that increasing depth at a fixed learning rate in the $\mu P$ parameterization corresponds to an effectively higher learning rate (see explanation in global response). Thus, since our experiments were run for a constant but moderate number of epochs, $\mu P$ runs result in models that are closer to convergence. This is true until the effective learning rate is so large that the runs diverge. We add a comment about this near the top of page 2 when we reference Figure 1.

We are performing new experiments where the parameterizations coincide in their dynamics at a large depth so that all SP/$\mu$P models in this range of depths are stable over the range of learning rates we consider. This can be achieved, for example by choosing a smaller $O(1)$ value for $\beta$ in the SP/$\mu$P experiments (or larger $\beta_0 = \beta \sqrt{L}$ for our proposed parameterization). In these experiments, we expect to see that the shallower SP/$\mu$P networks will train slower than their counterparts in our $\frac{1}{\sqrt L}$ parameterization, resulting in higher train loss. We will revise our sumbission and post them here if/when these expts finish during the rebuttal period.

We also note that Figure 1 seeks to illustrate hyperparameter transfer, which clearly fails for $\mu P$ but works well with our parameterization. In particular, we didn't attempt to obtain the best possible test loss for each architecture, which would have requred optimizing other hyperparameters (e.g. notably the lr schedule, how to do data augmention, momentum, batch size etc). In this sense, we do not believe that the discrepancy in test loss in Figure 1 takes away from the merit of our core proposal.

We hope that this addresses the reviewer's concern and welcome any further feedback on this point. Also, we are in the process of running several followup experiments to illustrate the first bullet point and will post then once they have finished running.

Weakness 2

While our theory was derived for a certain class of ResNets (Appendix K) and not fully theoretically analyzed for the case of transformers etc, we tried to verify experimentally the broad applicability of our proposed parameterization in at least the following ways:

• Experiments on CIFAR10 with Wide ResNets. Here, while the dataset is not very complicated, we check that not only the learning rate but also the momentum coefficient and regularization strength transfer (e.g. Figure 4).
• Experiments on both ImageNet and Tiny ImageNet with Vision Transformers. Here, we checked not only that the learning rate but indeed also the learning schedule transfers (e.g. Figure 3).
• Empirical validation of our theory in linear MLPs and ResNets in which we have analytic formulas for the NTK.
• In terms of computational savings we find, for instance, that already depth 6 and width 128, we can estimate optimal learning rates for models at depth 30 and width 1024 (Figure 1). This results in a cost savings of roughly 320x in each forward and backward pass used for hyperparameter tuning. In appendix B.2 (see Figure B.3), we saw that hyperparemeter turning after training ViTs on CIFAR-10 for only 3 epochs gives the same results as training for 9 epochs. So all together we get a computaitonal savings in this case of approximately 1000x. Moreover, we made no serious attempt to systematically investigate the minimal width and depth at which to perform hyperparameter turning.

Weakness 3

We agree that it makes sense to add a brief description of the $\mu P$ parameterization. Thank you for the suggestion. This is now done in the new Appendix C and C.2.

Weakness 4

This is an excellent question! The final Figure 7 which shows the linear network solution is actually in the feature learning regime $\gamma_0 > 0$. Figure 7b plots the change in the kernels after one step of feature learning for $\gamma_0 > 0$. The linear activations make analysis of feature learning much more tractable. As far as we are aware no one has been able to find explicit solutions for the evolution of the NTK and related feature and gradient kernels in nonlinear networks in regimes where $\gamma_0 > 0$. We hope to return to this in future work but for now have included this point in the discussion.

Thanks to authors for their "fair enough" responses, that are acceptable if the gaps are properly addressed in Discussion and Appendices, as mentioned in rebuttals. Especially, comments on computational savings are encouraging. I have no further comments and think it's an interesting paper that should be accepted. Thank you.

We thank the reviewer for their feedback and positive assessment of our work! We agree that additional experiments on larger dataset, larger models, and longer training times would be useful to explore the limits of hyperparameter transfer. However, we also wanted to respectfully point out that our original submission already included experiments on Tiny ImageNet and ImageNet (Figure 3). We plan on adding additional experiments where we train these models longer (see for instance some new experiments here where we train on Tiny-ImageNet for 100 epochs and CIFAR for 40 epochs https://imgur.com/a/GA6qepR). We plan to add these to the paper once we get the corresponding SP/$\mu$P baselines without the branch scalings.

How much utility/applicability does our proposed transfer scheme provide?

Reviewers probed the utility/applicability of our proposed transfer scheme in several ways:

• computational savings and breadth of applicability (Reviewer fkzv)
• hyperparameter transfer beyond 10 - 20 epochs of training (Reviwer SFFE)
• utility of training large (deep + wide) models in which the marginal gain in train error was small (Reviewer J7oQ)

We briefly address each of these below.

Computational Savings

We believe a core contribution of our work is a reliable way to find computational savings in hyperparameter optimization. We find, for instance that already depth 6 and width 128, we can estimate optimal learning rates for models at depth 30 and width 1024 (Figure 1). This results in a cost savings of roughly 320x in each forward and backward pass used for hyperparameter tuning. In appendix B.2 (see Figure B.3), we saw that hyperparemeter turning after training ViTs on CIFAR-10 for only 3 epochs gives the same results as training for 9 epochs. So all together we get a computaitonal savings in this case of approximately 1000x. Moreover, we made no serious attempt to systematically investigate the minimal width and depth at which to perform hyperparameter turning.

Hyperparameter transfer beyond 10 - 20 epochs

We have included in our revision experiments that check hyperparameter for models close to convergence (trained for 100's of epochs) and have found:

• For ViTs trained on Tiny ImageNet the gap in performance can be much more substantial at large depth if we train for longer. See https://imgur.com/a/GA6qepR.
• Also attached in the link above are CIFAR-10 runs for 40 epochs where more dramatic gaps in loss can be seen.

Further, in our original submission, we chose to train for 10 - 20 epochs in most of our experiments for these reasons:

• Our experiments were rather computationally expensive, at least with the compute resources available to us because
  • The focus of our work is hyperparameter transfer over both depth and width. We were thus bottlenecked by the computational cost of training our largest (deepest and widest) architectures, which in many experiments were Convolutional ResNets with depth 34, width 2048 and roughly 50M parameters.
  • Each experiment involved re-training the same architecture over many -- often 10 or more -- different learning rates (or other hyperparameter settings)
• We believe directly checking hyperparameter transfer after model convergence would likely require an entirely different suite of experiments, rather than simply training longer. Indeed, hyperparameter transfer for models fully trained models is most interesting only after optimizing over not only the learning rate but also the learning rate schedule, the batch size, momentum coefficient, data augmention, etc jointly. We have presented a fair amount of evidence indicating that in our proposed paramaterization all these hyperparameters individually transfer (see Figure 3 for learning rate schedules and normalization layer and Figure 4 for momentum and feature learning rate). But running large scans jointly over all these parameters was not computationally feasible.

Utility of Training Deeper Models

We certainly agree that whether training a larger model is worth the computational cost depends on the task. Indeed, our experiments consistently showed that larger/deeper models result in better performance, but the marginal gain from using larger models was sometimes small. In part, this was because we were not trying to reach the best possible accuracy with a given architecture in two senses:

• we did not optimize jointly all hyperparmeters (e.g. lr schedule, momentum, batch size, regularization, etc)
• we typically only trained for 10 - 20 epochs in our experiment. (Since the submission deadline, we have run several experiments https://imgur.com/a/GA6qepR, which suggest that the gap in performance can be much more substantial at large depth if we train for longer. We plan to add more of these experiments as they finish.)

We chose to set up our experiments in this way because the main focus of this paper is hyperparameter transfer, which we found could be achieved by using our parameterization to optimize hyperparameters in much smaller models. But to validate this empirically, regardless of the gains in model performance, required checking that in all of our experiments optimal hyperparameters settings were stable over many orders of magnitude in model size.