Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning

Summary

This paper analyses the lazy versus rich learning dynamics of minimal but finite neural networks by deriving exact solutions under arbitrary layerwise initialization and learning rates. The theoretical insights are successively extended to networks of increasing complexity and corroborated by several experiments.

Strengths

The paper is well written and provides an exact understanding of the learning dynamics of very simple neural networks, in increasing complexity over the course of the paper. The figures complement the derivations very comprehensively. In this way, the paper makes a significant contribution to the understanding of neural network training dynamics. I am not able to judge the claimed novelty, as I am not entirely familiar with the related work.

Weaknesses

Overall the paper leaves little room for criticism. The results being limited to shallow, often linear neural networks trained with small learning rates is justified by the detailed understanding of the training dynamics, and mitigated by experiments. It could have been acknowledged more clearly how the maximal update parameterization by Yang and Hu (2021) already shows that upstream initializations are necessary for non-vanishing feature learning in large models.

Questions

- Why do you call the layerwise initialization variances and learning rates *initialization geometry* when learning rate scalings clearly only influence the training dynamics? Similarly, the term geometry is unspecific and does not highlight the interplay between the layers. I encourage the authors to find a more suitable terminology for the choice of hyperparameters. - Why is the kernel distance in Figure 1 eventually maximized by a downstream initialization $\delta<0$, and not by an upstream initialization $\delta>0$ where feature learning occurs rapidly both directionally and radially?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The main limitations are acknowledged in the Conclusion section. It is acceptable that an exact analysis of learning dynamics does not include deep nonlinear networks trained with large learning rates.

Summary

The paper studies the dynamics of training in neural networks with the scope of identifying how the variance of weights' initialization together with layer-wise learning rates determines different learning regimes, encompassing lazy, rich, and the transition between them. The paper identifies a conserved quantity $delta$ that is preserved throughout training (i.e. $\dot{\delta} = 0$) and depends on both the weight's magnitude and the learning rates. Thus, the sign and magnitude of $\delta$ at initialization affect learning and the geometry of the landscape that is traversed through gradient flow. The paper finds three different regimes, starting from a solvable model of two-layer linear network. 1. $\delta > 0$ (named *upstream initialization*) induces lazy dynamics and corresponds to the case where $\eta_w a^2 >> \eta_a ||w||^2$, where $w$ and $a$ are the first and last layer weights, resp. 2. $\delta = 0$ (*Balanced*): corresponds to rich dynamics. 3. $\delta < 0$: initial lazy fitting regime and second rich phase. The authors then extend these findings to wider and deeper networks and piece-wise linear activations.

Strengths

1. The paper finds a simplified model that can be thoroughly analyzed and gives very precise statements on what causes rich and lazy regimes based on the single conserved quantity $\delta$, which makes the results very clean. This quantity is then naturally connected to the NTK, which makes intuitive sense given that the NTK is the ultimate factor that determines the learning regime (rich, lazy, or somewhere in between). 2. The formula for $\delta=\eta_w a^2 - \eta_a ||w||^2$ also is nicely interpretable, elucidating the interplay between learning rates and weights magnitude. 3. The authors test their theory in various settings and observe a close alignment between the experimental verification and the theoretical predictions, which makes me confident of the correctness of their theory.

Weaknesses

1. I would have appreciated more on the NTK dynamics. There is a clear connection between $\delta$ and (part of) the NTK $K$. I wonder whether this connection can be made more explicit by studying the NTK dynamics $\dot{K}$. This would make the connection between the conserved quantity and the feature learning more explicit. 2. Some equations are not included. The authors justify this in Line 1022: "We omit the solution due to its complexity, but provide a notebook used to generate our figures encoding the solution". I would just provide the formulas for the sake of completeness. Also, I do not see the equations describing the kernel distance S, which is related to the NTK movement, thus to my previous point. 3. In section 4 the paper includes various extensions, including wide networks. It would have been nice to recover known results by relating the conserved quantity $\Delta$ and specific parameterizations of the network as a function of the width. This would have connected existing results in the NTK and $\mu$P literature to the conserved quantities identified in this paper and the learning regime. 4. The connection between the scale of the initialization and the learning rate is not entirely new. This is quite obvious at grokking problems, where to exhibit grokking is often necessary to play with the initialization parameters of different layers. In this context, Kumar et al [61] already clarified that the delayed generalization was caused by the NTK dynamics transitioning from lazy to rich regimes. Overall, I am in favor of the paper's acceptance because it devised a model where the three different learning regimes are clean and crystallized.

Questions

See weaknesses

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

Addressed.

Summary

The paper studies the learning dynamics of a deep network by leveraging the conserved quantities due to symmetries

Strengths

The theoretical finding that unbalancedness in the layers drives feature learning is a novel and interesting insight The technical tool of using symmetries and conserved quantities to characterize the dynamics is also novel Finding exact solutions can help us understand better the causal aspects of the phenomena in deep learning, and should be encouraged Overall, I am positive towards this work

Weaknesses

I think the paper could improve by explaining its results in more detail. For example, I find the following points rather confusing 1. Does theorem 4.1 hold after training? Or it holds at any time in training？ 2. The paper argues that it finds an exact solution, but I am sure what this is referring to. Is theorem 4.1 the "exact solution" the paper finds? If so, what is it a solution to? 3. The title claims "rapid feature learning", but almost no place in paper discuss what it means to be rapid, nor what it means to be doing "feature learning." The only part in the paper that I find to be peripherally related to this claim is the discussion in lines 380-383. This result holds true without needing any result from sections 3-4. Then, what is the point of 3-4? If the title is the main claim of the paper, can the authors write it out in a more explicit manner, and with much greater mathematical detail? 4. In figure 5, a crucial quantity is $\alpha$, but it is not explained in the context. What is $\alpha$? It seems to me that this quantity is the scaling factor introduced in lines 171-172? Is this case? Even if so, the authors should have explained it much better in the immediate context 5. To be fair, I think the authors are putting too much content both in the appendix and in the main text. It feels to me much better to move one of section 3 or 4 to the appendix and expand what is in the main text to make it much more readable 6. A constructive suggestion is that it would be quite insightful to compare the results in section 4 to the results derived for SGD in deep linear models in https://arxiv.org/abs/2402.07193. SGD and GD have very different behaviors for these models and this discussion will be very helpful to the audience

Questions

See the weakness section

Rating

7

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

NA

Summary

The authors derive exact solutions to a minimal model that transitions between lazy and rich learning, elucidating how unbalanced layer-specific initialization variances and learning rates determine the degree of feature learning in a finite-width network. They provide evidence that this unbalanced rich regime drives feature learning in deep finite-width networks, promotes interpretability of early layers in CNNs, reduces the sample complexity of learning hierarchical data, and decreases the time to grokking in modular arithmetic.

Strengths

1. The paper is very well written: - clear visuals and exposition, going from a minimal model and gradually adding complexity and realism. - The quantities of interest (e.g. $\delta$) give an intuitive picture. - The notations are very clear and easy to follow. One small comment is that it might be better to use a different index for $\theta$ and $x$. - The contributions are distinguished from previous works and relevant connections to them are clearly stated. 2. Addressing the rich regime is arguably more interesting than the lazy regime and is much less studied, making the questions addressed here timely. 3. The experiments presented in Fig 5 show that the effects studied analytically in previous sections also appear in realistic settings, making this work relevant to a broad audience.

Weaknesses

1. To yield the exact solutions, the authors evoke the assumption of whitened input, which is rather unrealistic. The low-rank case is also addressed but then one needs to additionally assume that the interpolating manifold is one-dimensional to find the solution in terms of $\delta$ exactly. missing citations: - In line 398, when mentioning grokking and the transition from lazy to rich learning - it is worth noting Ref. [1]. [1] https://arxiv.org/abs/2310.03789

Questions

1. In Ref. [1] which you mention, the formula for the NTK with unbalanced LRs is (using your notation) $K = X (\eta_a a^2 I_d + \eta_w ww^T) X^T$ (eq. therein), i.e. the LRs have switched places. I'm not sure where this discrepancy comes from or if it is an error. 2. Line 145 - can you explain why in the infinite width limit non-linearities act linearly? It is true that the model becomes linear in the weights, but do you mean that each layer acts linearly on its pre-activation? [1] https://arxiv.org/abs/2401.07085

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

The authors adequately addressed the limitations.

Thanks for the reply

Thanks for the detailed explanation. The answer removed my concerns about the work. I will raise to 7. Although I am positive towards the work, I do not find the results significant enough for me to strongly support it.