Paper: Depthwise Hyperparameter Transfer in Residual Networks: Dynamics and Scaling Limit (ICLR-2024)

Reviewer comment (Theory):
There is a gap between the stated goal of the theory and its practical implementation. The theory predicts stable results following a scale-up; however, spending more compute and getting equal performance is of little practical value. In practice, one expects to get better results, which seems to be the case at least for the larger datasets considered. This implies that in examples where HP-transfer is worthwhile, one is by definition well away from the scaling limit.

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

Reviewer comment (Theory):
Another gap is that the DMFT framework is based on gradient flow but is then used to transfer the learning rate. While gradient flow may be a good approximation for SGD in the asymptotically low learning rate regime, typically the optimal learning rates are far from this regime.

Author response:
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).

Reviewer comment (Evaluation):
For CIFAR-10, the 1/√depth scaling leads to poorer performance compared to muP scaling.

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

Reviewer comment (Novelty):
Reference to current literature could be broadened and sharpened. For instance, muP parameterization is also a unique parameterization having maximal updates, not just a stable limit. This fuller characterization also partially escapes the first gap described above. Citing further alternatives to muP scaling and adaptive kernel approaches to feature learning is also desirable.

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

Reviewer comment (Writing):
Since that the $\mu P$ parameterization is not a commonly used method, the authors should provide a more comprehensive introduction.

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

Reviewer comment (Novelty):
The $1/\sqrt{depth}$ trick has been extensively investigated in previous works [1-3]. The novelty of the proposed method appears to be limited. Please clarify the differences. Moreover, many recent works propose initializing ResNet using "rezero"-like methods, which achieve similar performance to the $1/\sqrt{depth}$ trick. The authors should provide some comparisons. In particular, I would like to see some discussions on the scaling parameter. Why use $1/\sqrt{depth}$? What would happen if the scaling parameter is set as $1/depth$ or 0?

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

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

Reviewer comment (Experiments):
As the authors admitted in the "Limitation" paragraph, most of the experiments are confined to 10-20 epochs. This limitation makes the empirical evidence less convincing. In fact, I don't understand the difficulty in training a ResNet with more epochs if one can train it with 10-20 epochs.

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

Reviewer comment (Evaluation):
Fig. 1 Loss levels reached seem to be higher in case of proposed parameterization compared to alternatives. Could authors add some comments on the topic? Especially, is a proposed parameterization capable of reaching an optimal value of hyper parameters in scope, e.g., learning rate in practical settings? I rate paper 'accepted' conditioned on this issue will be alleviated in the camera-ready version.

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

Reviewer comment (Evaluation):
Limited applicability(?) - Could authors argue otherwise? How could it be improved in the paper? Could authors elaborate on computational costs of proposed method vs. alternatives?

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

Reviewer comment (Writing):
Paper heavily depends on $\mu P$ parameterisation properties derived in Yang & Hu (2021). In my view a short primer on $\mu P$ parameterisation (even in Appendix) would improve self-consistency and readability of the paper.

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

Reviewer comment (Theory):
Proposition 2, Assumption $\gamma_0 \rightarrow 0$ renders ODE from Proposition 1 solvable due to time invariant kernels (and thus given by Gaussian initialization) and the same goes for the second example where linearity of network makes time invariance explicit. To what extend are conclusion transferable to any realistic non-linear scenario?

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

Reviewer comment (Experiments):
The experimental section is quite limited: only a single task and a very simple dataset are considered.

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