Tensor Programs VI: Feature Learning in Infinite Depth Neural Networks

We thank the reviewer for the thoughtful and insightful review! We would like to first address the major concern about the cause of hyperparameter transfer.

What causes hyperparameter transfer in general? We have added a new section (Appendix I), particularly for discussing what leads to hyperparameter transfer. The fundamental idea is that the optimal limit implies the transfer of optimal hyperparameters. In the scenario of widthwise scaling (discussed in Tensor Program IV and Tensor Program V paper), $\mu$P has the optimal limit. In other words, all other widthwise scalings have obvious defects when compared with $\mu$P, so they must perform worse than $\mu$P in the limit. Thus, optimal hyperparameters transfer under $\mu$P.

What causes hyperparameter transfer across depth? Back to our scenario where we consider the depthwise scaling, the question is whether there exists a parameterization (represented by $\alpha$ and $\gamma$) that is always better than others. Without the concept of feature diversity, we are unable to identify one because a parametrization is stable, nontrivial, faithful, and induces feature learning if $\alpha+\gamma=1$ and $\alpha\in [1/2, 1]$. The good news is that there is one unique parametrization (i.e., Depth-$\mu$P, or $\alpha=\gamma=1/2$) that maximizes feature diversity. The feature diversity exponent for Depth-$\mu$P is $1/2$, while the feature diversity exponent for $\alpha\in (1/2,1]$ is always $0$.

Why is feature diversity important? Conceptually, low feature diversity prevents the intermediate features from changing quickly across layers, and hence the network is wasting the capacity. We have added another section (Appendix J) to explore more details and applications of feature diversity. We have also added experiments in Appendix L.2 showing that maximal feature diversity is necessary for hyperparameter transfer, as the so-called ODE parametrization ($\alpha=1,\gamma=0$) induces feature learning but fails to transfer hyperparameters. In Appendix L.4, we also empirically verify our claims about feature diversity exponents for different parametrizations and different nonlinearities.

Below we address other comments one by one.

Missing citations.

We will provide a more comprehensive reference section in the final version.

How much do the optimal hyperparameters (e.g. learning rate) shift as a function of training time t? It would be nice to comment on this in the main paper, as practitioners might be interested in tuning the architecture with minimal amount of training steps.

In general, the shift across time depends on the task, model, and training recipe. For example, the optimal learning rate becomes smaller when trained longer with a constant learning rate schedule. However, one would expect using a proper learning rate decaying strategy can reduce the shift. Note that with the same reasoning, we expect the whole learning rate schedule to transfer with depth.

Why use the mean-centered version of Resnet? For instance, in https://arxiv.org/abs/2302.00453 the weights are applied after the nonlinearity, thus achieving centering.

In our revised version, we provide the reason during our discussion about feature diversity in Appendix J. The short answer is that with nonlinearity posed before the weights, the gradients of nearby layers have little diversity, hence effectively wasting capacity.

… How is the infinite depth limit handled? … Is the order of the limits still commutative during training?

This work focuses on the infinite-width-then-depth limit. In other words, we consider the case where width >> depth, which is exactly why we can apply the Tensor Program framework while the depth goes to infinity.

However, we still expect that the infinite width and depth limits commute during training, especially because the width is sometimes smaller than the depth in our experiments. Nevertheless, it is difficult to prove without more advanced theoretical tools, and we are currently investigating this for future work.

Thanks for your valuable time and further comments! We provide our response below:

Then why in Figure 11, 12 does the ODE limit have a similar training loss to the ($\alpha=1/2$) model (despite not having hyperparameter transfer?).

Recall the "real" block multiplier in the experiment is $aL^{-\alpha}$ and the "real" learning rate is $\eta L^{-\gamma}$, so no matter what $\alpha$ and $\gamma$ are chosen, for fixed $L$, one can always sweep $a,\eta$ extensively such that the training loss is minimal over all possible "real" block multiplier and "real" learning rate choice for this particular $L$. So directly comparing the minimal training loss for $L=2^{10}$ is not meaningful. Instead, we should fix $a,\eta$ for all $L$ by sweeping them extensively at some small $L$. In our experiments, we choose to find the optimal $a,\eta$ when $L=2^3$ (you can see that the $L=2^3$ curves in Fig. 11 and 12 have the same shape). Now back to Fig. 12, if we fix $\eta$ to be the optimal one at $L=2^3$, then the log of training loss at $L=2^{10}$ is around -3.1; while for $\alpha=\gamma=1/2$ in Fig. 11, the number is around -3.5. I hope this also convinces you that having a better usage of the weights (i.e. model capacity) should lead to better performance at least when $n\to\infty, L\to\infty$.

... how does maximizing feature diversity imply that the optimal learning rate converges to an optimal in $L, N$ and it does so at a faster rate than for other values of $\alpha$ (e.g. $\alpha=1$)?

For better illustration, we introduce function $f^L(\bar a, \bar\eta)$ that represents the training loss after $T$ steps of training with "real" block multiplier $\bar a$ and "real" learning rate $\bar\eta$ with depth $L$ and infinite width. Let $\bar a^{*L}, \bar\eta^{*L}\triangleq\arg\min_{\bar a, \bar\eta} f^L(\bar a,\bar\eta)$. (For simplicity, let us assume the minimum is unique.) Then we are essentially looking for how $\bar a^{*L}, \bar\eta^{*L}$ changes when $L$ grows, i.e., if $\bar a^{*L}=\Theta(L^{-\alpha}), \bar\eta^{*L}=\Theta(L^{-\gamma})$, then optimal hyperparameters converge (only) under $(\alpha,\gamma)$ parametrization.

It is clear that $\bar a^{*L}$ must not be $\Theta(L^{-0.49})$. $\alpha=0.49$ makes the network unstable when $L\to\infty$, so $\lim_{L\to\infty} f^L(aL^{-0.49},\bar\eta)$ cannot be small for any $\bar\eta$.

Similarly, if $\lim_{L\to\infty} f^L(aL^{-1/2}, \eta L^{-1/2}) < \lim_{L\to\infty} f^L(a'L^{-1}, \eta')$ is always true, then $\bar a^{*L}, \bar\eta^{*L}$ cannot be $\Theta(L^{-1})$ and $\Theta(1)$.

If the reviewer agrees that having a better usage of the weights (i.e. model capacity) should lead to better performance when $L\to\infty$, then we have $\lim_{L\to\infty} f^L(aL^{-1/2}, \eta L^{-1/2}) < \lim_{L\to\infty} f^L(a'L^{-\alpha}, \eta'L^{-\gamma})$ holds true for any $(\alpha,\gamma)\neq (1/2,1/2)$. Thus, if $\bar a^{*L}=\Theta(L^{-\alpha}), \bar\eta^{*L}=\Theta(L^{-\gamma})$ for some $\alpha, \gamma$, then $(\alpha,\gamma)$ must be $(1/2,1/2)$. Of course, one can argue that $\bar a^{*L}$ and $\bar\eta^{*L}$ may not scale like real powers of $L$, but they can be rejected by a straightforward generalization of our claims in the paper.

We would like to thank the reviewer for encouraging review and thoughtful questions! Please see our response below.

The resnet archtiecture is too simplifeid. The muP initialization can't work when the residual block is two-layer. How about the theory for linear resnet when the residual block is two-layer?

We thank the reviewer for pointing this out! This is indeed an important question. We have added discussions about using two-layer residual blocks in Appendix K, where we provide both theoretical and empirical arguments as to why the optimal parametrization of the general architecture is impossible with the current setup. Our argument in Appendix K.3 indeed uses the linear resnet as an example.

the tensor program framework can deal with the activation function and normalization layer but this paper dropped the part there. Can even just analyze one step of gd？ (or linear net+noramlization layer)

In our general case (see Sec. 5 and 6), the activation function and mean subtraction are included. Replacing mean subtraction with the normalization layer will not change our results because normalization does not change the scaling of the activation for any stable parametrization. To avoid unwanted complexity, we decided not to include it in the paper for better presentation. Empirically, we verify that adding normalization layers still exhibits hyperparameter transfer under Depth-$\mu$P (see Appendix L.2).

the convergence in the tensor program is just loss function convergence and can't obtain the convergence of optimal hyperparameters. Can the author show that the optimal hyperparameter have faster convergence speed?

We have included a section on “What causes hyperparameter transfer?” in Appendix I. By classifying all the parametrizations, we found Depth-$\mu$P to be the unique optimal parametrization in the sense that it guarantees both feature learning and maximal feature diversity, and that is why we claim the convergence of optimal hyperparameters.

The convergence speed of the optimal hyperparameter indeed decides whether hyperparameter transfer is useful in practice, but is almost impossible to prove or estimate without empirical evaluation.

We thank the reviewer for the detailed review. We provide our responses to the comments and questions below.

“The stacking of many residual blocks causes an obvious issue even at the initialization — the norm of $x^l$ grows with $l$,...” does not match the real design of the neural network. As there will be a normalization layer …

We thank the reviewer for pointing this out. We will include the following discussion in the final paper.

Firstly, this motivation still stands for the pre-LN architecture, where the normalization is applied inside the residual branch. In particular, if we consider the following residual network $x^l=x^{l-1}+W^{l2}ReLU(W^{l1} LN(x^{l-1}))$, then when the width $n$ goes to infinity, we have $||x^l||_2^2=||x^{l-1}||_2^2 + 1/2$ at initialization since inner product between $x^{l-1}$ and $W^{l2}ReLU(W^{l1} LN(x^{l-1}))$ is zero and $||W^{l2}ReLU(W^{l1} x)||_2^2 =1/2$ for any $x$ with $||x||_2=1$. So with pre-LN, the norm of $x^l$ still grows with $l$, but with a slower speed.

Then one might ask the follow-up question: since people apply LN on $x^L$ in practice, would that solve the problem? The answer is still “no”. In particular, at the initialization of the above residual network, the last-layer features of any two different inputs are always independent (i.e., the inner product of the last-layer features is always zero) when $n, L\to \infty$. It means the network is not aware of any correlation between the inputs! This observation is followed by a simple calculation: let $x_1^{l}$ and $x_2^l$ be the outputs of $l$-th layer for two inputs respectively, then $\langle x_1^{l}, x_2^l\rangle=\langle x_1^{l-1}, x_2^{l-1}\rangle+\rho^{l-1}(\pi-arccos(\rho^{l-1}))/2\pi$ where $\rho^{l-1}=\frac{\langle x_1^{l-1}, x_2^{l-1}\rangle}{||x_1^{l-1}||_2||x_2^{l-1}||_2}$. One can easily see that $\rho^{l}$ decreases when $l$ grows, and $\rho$ converges to $0$ in the end.

A similar calculation can be applied to post-LN architectures as well, and lead to the same observation on the independence between last-layer features.

No performance comparison between this structure and commonly used structure.

We are not sure what the reviewer means by “modified structure”. Does it refer to adding the branch multiplier? Or $g$ contains only one layer?

We aim to deliver the theory of depthwise scaling with block depth 1 in this paper, so we care more about the performance comparison between different scaling strategies. Our experiments serve as a verification of the theory and we do not expect it to perform exactly the same on convnets or transformers where the block depth is not 1. (We have included a discussion about using deeper blocks in Appendix K, and provided experiments on Megatron trained on Common Crawl in Appendix L.3.)

It is not so easy for me to follow the theoretical analysis…why the analysis needs the width of the network to go to infinity is still not clear.

We have added more intuition to $\Gamma$ and $C$ in Sec. 4.

Our analysis relies on Tensor Programs, which helps us better characterize the training dynamics, but it requires taking the width to infinity first. It is believable that our results hold even when the width is not very large compared to the depth, especially because the width is sometimes smaller than the depth in our experiments. However, it is difficult to prove without better theoretical tools.

Thanks again for your valuable time in reviewing our paper!

Since the discussion period is coming to an end, we would like to know if our response addresses your main concerns. We are happy to address any further comments!

We thank the reviewer for the detailed review. We believe there are some misunderstandings that we would like to clarify first:

• Depth-$\mu$P is unique because it not only exhibits feature learning but also maximizes feature diversity. All parametrizations with $\alpha+\gamma=1$ and $\alpha> 1/2$ also enable feature learning, but they lack feature diversity (Claim 6.4). We have emphasized this in the paper and added more discussion about feature diversity in Appendix J.
• The purpose of this work is mostly theoretical and focuses on the optimization side. The experiments on CIFAR-10 serve as a verification of the theory and we do not expect one can find perfect hyperparameter transfer for transformers where the block depth is not 1. (We have included a discussion about using deeper blocks in Appendix K, and provided experiments on Megatron trained on Common Crawl in Appendix L.3.) We neither expect hyperparameter transfer for the late stage of training on image datasets where the noise caused by SGD is prominent. (This is also true for the original $\mu$P.)
• Sections 3, 4, C, and D in the paper serve as a warmup for readers to understand why Depth-$\mu$P is a good scaling strategy, i.e., why it is stable, why it exhibits feature learning, and how to derive its training dynamics in the limit. We compare Depth-$\mu$P with other parametrizations in Section 6.

Thank you for your thorough reply. It has largely cleared up my concerns and confusion, and I believe that the readability has greatly improved. I have raised my score to an acceptance side. Since the options are only 6 and 8, I chose 6, but if allowed, I would prefer a middle evaluation 7.

There are still some parts that concern me slightly, although none are critical flaws. They just seem preferable to furhter enhance the significance of this study.

On question (i)

My understanding is that the derivation of Depth-$\mu$P is based on the linear activation case. As the authors remarked, for non-linear cases (the general case), the derivation remains heuristic. I mean, it seems quite hard to derive the Depth-$\mu$P (shown in Section G.4) from only watching the tensor program lines of the general case (shown in Section G.1-G3). Certainly, in the revised version, the authors have added Section H and clarified that the tensor program lines certainly give the correct Depth-$\mu$P for the linear case. However, without such an example, it seems not easy to evaluate the concrete values of ($\alpha, \gamma$). From a theoretical standpoint, readers may wonder about any potential discrepancy arising from non-linearity.

On concerns that width is not so large in experiments

Initially, I was concerned about the lack of experimental verification for Depth-$\mu$P, but the revised manuscript now includes both relu nets and transformers, significantly enhancing the theory's validity. My only minor dissatisfaction concerns how closely a finite width (=256) approximates the infinite limit. While the argument about limits commute is persuasive, there's still some uncertainty about whether the success of Depth-$\mu$P in such a finite setting is coincidental or inherent. Therefore, a more explicit demonstration of width dependence would be advantageous. I look forward to possibly seeing it after the decision.

Thanks again for your valuable time in reviewing our paper!

Since the discussion period is coming to an end, we would like to know if our response addresses your main concerns. We are happy to address any further comments!

We thank the reviewer for the thoughtful and encouraging review! We provide our response below.

Additional feature diversity experiments?

We have added more discussions on feature diversity in Appendix J, and corresponding experiments in Appendix L.4. Comparison between Fig. 11 and 12 demonstrates the empirical advantage of Depth-$\mu$P ($\alpha=\gamma=1/2$) over the so-called ODE parametrization ($\alpha=1, \gamma=0$). Fig. 18 and 19 show their feature diversity difference. Depth-$\mu$P with higher feature diversity performs better.

… section 4 … I encourage the authors to either provide some more intuition on the "physical" meaning of the different equations, or to shorten the section and defer details to the appendix. …

Thanks for the suggestion. We have shortened Sec. 4 and added more details of the technical road map.

Also, how does this related to the AMP and DMFT literature? some more comparison to related literature would be helpful (in particular, this concurrent paper seems highly related and should be cited and compared to in the final version)

Thank you for pointing out this paper, which was also recently submitted to ICLR. We have indeed been in touch with the authors of that paper, and we intend to include a comparison with that work in the final version. In short, our work introduces the notion of feature diversity and we argue that it is essential in obtaining hyperparameter transfer (optimal limit), this was not discussed in the reference mentioned above. Secondly, regarding optimal parameterization with Adam, it seems that our conclusions are in contradiction with the conclusions of that work: Depth-muP suggests that both block multiplier and learning rate should scale as $1/\sqrt{depth}$, while in that work, they suggest only scaling the blocks. In our experiments with depths up to 2^10, we showed that Depth-muP yields hyperparameter transfer with Adam.

mean subtraction: could you clarify a bit more the role of this? Is MS specifically needed because of the activation? Would adding an MLP output layer as in transformers drop this requirement? Note that in practice LLM practitioners often drop the mean subtraction part of layer-norm, using RMSnorm instead

Yes, as we mentioned in the last paragraph of Section 5, MS is needed due to the activation, which leads to the exploding behavior of intermediate features because of the non-zero mean of the block outputs at the initialization. Therefore, adding an MLP output layer in the residual blocks can solve the problem as it makes the features zero-mean, so it is not surprising that dropping MS in transformers works in practice. However, the analysis in this paper does not trivially apply to 2-layer blocks. We have included a discussion of 2-layer blocks in Appendix K and experiments in Appendix L.3.

scaling: do you expect any different behaviors in different limits, e.g. in the proportional deep-wide limit studied e.g. by Hanin, or what do you expect would change if the number of GD steps can grow with the number of layers? any comparison or intuition here would be useful.

This is an interesting question. Actually, we expect (and can give heuristic arguments) that the infinite width and depth limits commute, in the sense that no matter how you take that limit, the result is the same. For instance, you get the same limiting dynamics if you take infinite-width-then-depth limit or the proportional limit. This was proven rigorously at initialization in Hayou and Yang (2023). However, proving the same result during training is challenging and we are currently investigating this for future work.

on feature diversity: how does the quantity in Def 6.6 behave at initialization vs later in training? what do you mean by "maximal" in claim 6.5? any more intuition beyond the footnote on how to get smaller exponents than 1/2?

The behavior does not change during training. $\kappa=1/2$ is maximal for all stable and non-trivial parametrization. Specifically, combining Claim 6.4 and Claim 6.5, there are only two choices of $\kappa$: $0$ for $\alpha\in(1/2,1]$, and $1/2$ for $\alpha=1/2$. Regarding exponents smaller than 1/2, this is closely related to rough path theory in stochastic calculus. Consider Depth-muP parameterization: the weights are generally initialized to be iid, and therefore when this noise adds up (because of the skip connection), it yields something similar to a Brownian motion, which has a continuity exponent of 1/2. This exponent can be further reduced by correlating the weights of different layers in which case we obtain something similar to a fractional Brownian motion which can have a general continuity exponent $H \in (0,1)$, depending on the correlation level.