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

Thank you for your thoughtful feedback and constructive questions. We are grateful for your positive feedback on our work and are committed to improving our manuscript by addressing each of the weaknesses you identified.

Connections to infinite width parametrizations. We agree with the review and are adding a subsection to Section 5 carefully outlining how our results can be extended to infinite width settings and acknowledging connections to existing parametrizations. In the two-layer setting, muP corresponds to the mean-field parameterization, which, for input dimension constant in width, can be written as $f(x) = 1/k a^T\sigma(Wx)$ with $a_i \sim \mathcal{N}(0, 1)$ and $W_{ij} \sim \mathcal{N}(0, 1/d)$. We note that this actually leads to the per-neuron conserved quantity being 0, or balanced, in expectation, with a non-vanishing variance. Please see our response to reviewer JJEZ for a more detailed discussion on the connection between our analysis of finite-width networks and existing works on infinite-width networks.

Initialization geometry. We chose the term initialization geometry because the per-layer learning rates and layer magnitudes collectively determine the geometry of the surface that the dynamics are constrained to. This geometry is leveraged consistently in our theory through the use of conserved quantities and is clearly visualized for the one hidden neuron model in Figure 2.

Kernel movement in downstream initialization. In Figure 1(b), the small-scale downstream initialization ($\delta<0$) starts off lazy, attempting to fit the training data by only changing the small readout weights $a$ with minimal movement in the large first-layer weights $\{w_i\}$. Up to this stage, the downstream initialization exhibits smaller kernel movement than the upstream initialization, which undergoes a rapid change in the kernel. However, in order to interpolate the training data, the network with downstream initialization needs to align its hidden neurons by a non-trivial amount, thus undergoing a change from lazy to rich learning. This alignment of the hidden neurons, which are large in scale, results in a dramatic change in the kernel. In the case of the upstream initialization, the network is able to interpolate the training data by simply aligning its hidden neurons (which are small in scale) directionally and radially as needed, leading to a smaller movement in the kernel overall. See our response to reviewer JJEZ for a detailed discussion on the dynamics of the NTK matrix.

We hope we addressed your points regarding our work. Thank you again for your constructive feedback!

Thank you for your thorough review and constructive suggestions. We appreciate your positive feedback and address the weaknesses you highlighted individually. We hope this will increase your confidence in the importance of our work.

NTK dynamics. This is a good point; we can in fact study the dynamics of the NTK matrix directly, which leads to an argument similar to the one provided in lines 239 - 264. Here we outline this analysis, which we will add to appendix A.

The dynamics of the NTK matrix $K = X M X^\intercal$ is determined by $\dot{M}$. From Equation 3 in the main text, we can write $\dot{M} = \frac{2 \|\beta\|}{\kappa} (I_d + \hat{\beta}\hat{\beta}^\intercal) \partial_t \|\beta\| + \frac{\kappa - \delta}{2} \partial_t (\hat{\beta}\hat{\beta}^\intercal)$. From this expression we see that the change in $M$ is driven by two terms, one that depends on the change in the magnitude of $\beta$ and another that depends on the change in the direction of $\beta$. As done in the main text, we consider the limits as $\delta\to\pm\infty$ and when $\delta = 0$ to identify different regimes of learning. For $\delta\to\infty$, the coefficients in front of both terms vanish, and thus, irrespective of the trajectory taken from $\beta(0)$ to $\beta_*$, the change in the NTK is vanishing, indicative of a lazy regime. For $\delta \to -\infty$, the coefficient for the first term vanishes, while the coefficient on the second term diverges. Here, the change in the NTK is driven solely by the change in the direction of $\beta$. This is why for large negative delta we observe a delayed rich regime, where the eventual alignment of $\beta$ to $\beta_*$ leads to a dramatic change in the kernel. When $\delta = 0$, the coefficients for both terms are roughly of the same order, and thus changes in both the magnitude and direction of $\beta$ contribute to a change in the kernel, indicative of a rich regime.

Adding equations for completeness. We agree that it is important to include expressions for completeness. We will include explicit solutions to the one hidden neuron model dynamics in the appendix. Specifically, we will write out the expression for $a(t)$, which is currently omitted because of its verbosity. We will also write out the expression for the kernel distance presented in figure 3c and include this in the added subsection of the appendix describing the kernel dynamics (as we discussed above).

Connections to width-dependent neural network parameterizations. We agree, discussing how our analysis of finite-width networks connects to existing analyses of infinite-width networks is important. We outline a discussion we will add to the main:

In a width-dependent parametrization, such as mean field or NTK, the random initialization of weights leads to a distribution over conserved quantities. For example, in the two-layer setting, the mean-field parametrization leads to the per-neuron conserved quantity being 0 in expectation, but with a non-vanishing variance. The NTK parametrization has the same distribution over conserved quantities, but at a larger function scale, leading to lazy dynamics. Thus, between NTK and mean-field, a change in function scale modulates the degree of feature learning. In our work, we show that relative scale between layers (i.e. $\delta$) is a separate knob that can also be tuned to influence the degree of feature learning.

When considering these parametrizations in the infinite-width limit we can recover a phase diagram analogous to our results in the finite-width setting. Reference [14] previously considered $f(x) = \frac{1}{\alpha}\sum_{i\leq k}a_i\sigma(w_i^\intercal x)$ with weights initialized as $a_i \sim \mathcal{N}(0, \beta_a^2)$ and $w_i \sim \mathcal{N}(0, \beta_W^2I_d)$ as width $k\to\infty$. They obtain a phase diagram at infinite width capturing the dependence of learning regime on the overall function scale $\beta_a\beta_W/\alpha$ and the relative scale $\beta_a/\beta_W$. The resulting phase portrait is analogous to ours in Figure 1 (b), which considers the conserved quantity $\delta$ rather than the relative scale $\beta_a/\beta_W$. In particular, there is a lazy regime, which is always achieved at large scale (just as in the large-$\tau$ regions of Figure 1 (b)), but is also achieved at smaller scale if the first layer variance is sufficiently larger than the second (as in the downstream initializations at small $\tau$ in Figure 1 (b)). On the other side of the phase boundary is the infinite width analog of rapid rich learning, with all neurons condensing to a few directions. This is induced either at small function scale, or at larger function scale if the relative scale is sufficiently large, such that $W$ learns faster than $a$.

Connections to prior work in grokking. Kumar et al. [61] explored how grokking arises as the result of the transition from lazy to rich dynamics, studying this transition as a function of overall function scale, initial NTK alignment with the test labels, and train set size. However, we add more nuance to this picture in that we show how grokking can be induced not just by modulating the overall function scale, but by simply changing the initialization geometry, that is by scaling the embedding weights for positional and token embeddings without affecting overall scale (as visualized in Figure 5 (d)). This leads to the complex phase diagram of time to grokking as a function of both scale and geometry, as presented in Figure 11. We believe this adds to the picture presented in Kumar et al. by showcasing the unique role of layer-wise initialization scales.

We hope we addressed your points regarding our work. If we have, we would appreciate it if you would consider raising your score to reflect this. Thank you again for your constructive feedback.

Thank you for taking the time to thoroughly review our paper and highlight areas needing further clarification. We appreciate your positive feedback on our work. We will address each of the weaknesses you mentioned individually, hoping this will enhance your confidence in the significance of our research.

Theorem 4.1. This theorem holds throughout training. We will make this more clear in the theorem statement by adding a time dependence to $\beta_i(t)$ and write for all $t \ge 0$.

Exact solution. The exact solutions we refer to are for the gradient flow dynamics of $a$ and $w$ in the minimal model we present in section 3. The gradient flow dynamics for this model boil down to a complex coupled system of nonlinear ODEs shown in equation 1. As stated on line 191, we solve this system of ODEs exactly, derived in appendix A in detail. See figures 2, 3, and 4; in all these figures the dashed black lines represent theoretical predictions (i.e exact solutions) while the colored lines are empirical (the results of training on a computer). We will make this more clear by referencing section 3 when we discuss exact solutions in the paragraph titled “our contributions” in section 1.

Rapid feature learning. As discussed in the related work section (lines 45-62), feature learning in the context of this work is synonymous with rich learning. We mathematically define rich learning on lines 96-101, as a change in the NTK through training measured by the kernel distance metric proposed in Fort et al. 2020. The main claim of our paper is that an upstream initialization leads to rapid feature learning relative to a balanced or downstream initialization. By "rapid" in our title, we refer to this specific claim. We demonstrate this empirically in figure 1. Sections 3 and 4 provide the necessary analysis to prove this claim in section 5. This explains why our title states “exact solutions reveal how unbalanced initializations promote rapid feature learning.” To clarify this further, we will revise the writing on lines 63-84 under "Our contributions." This section outlines our contributions and sets up the paper’s layout. We aim to make this section more concise to clearly explain why we use the approach in sections 3 and 4 to prove our main claim and its connection to the title.

Alpha in figure 5. Thank you for catching this. Yes you are correct, $\alpha$ in this context is equivalent to the scaling factor used in lines 171-172. It is also described in the caption for figure 1 and in the experimental details in appendix D, however, we forgot to thoroughly describe it in the caption for figure 5. We will revise the caption to make this more clear.

Content distribution between the main and appendix. We appreciate this comment and we agree that the paper is dense. We have tried to make the work as readable as possible, but we believe that section 3 and 4 are necessary steps to build the analysis used in section 5. With the additional page allotted for the final version of this paper we have moved up some aspects of the appendix into section 3 and 4 which should make it more readable, added clarifying transitions between the sections, and expanded the content in section 5 to clarify its connection to the previous sections and the main claims of our paper. We think these changes will significantly improve the readability.

Related work on SGD. Thank you for highlighting this reference. We will include it in our discussion on the limitations of our work (lines 405-411). In this section, we explain how SGD disrupts the conservation laws that are central to our study. As you pointed out, one of our key findings is how conserved quantities arising from symmetry affect the degree of feature learning. Other works have focused on the influence of stochasticity in SGD on these conserved quantities. We agree integrating these analyses to understand the influence of stochasticity on the degree of feature learning would be an insightful direction for future work. However, this is beyond the scope of our current paper, which is already densely packed with content.

Thank you again for your constructive feedback. We hope we addressed your points regarding our work. If we have, we would appreciate it if you could raise your score to reflect this. Thank you!

We appreciate your comprehensive review and the time you have taken to suggest improvements for our study. We are also happy to hear you think our work is a timely and important contribution to the field. We will respond individually to the weaknesses and questions you raised regarding our paper:

Notation for indices. We agree with you that using a different index for parameters, data, and layers can help with readability. This is a more comprehensive change, but something we will implement in our updated draft.

Whitened input assumption. We agree that the assumption of whitened input for our minimal model is quite strong, although a common assumption used in many prior works (see Saxe et al. 2014 for example). That said, we actually do relax this assumption throughout the analysis. As you noticed, a strength of our work is that we start with a simplified setting where we can derive exact solutions, then “gradually adding complexity and realism”, we show that the key-takeaways remain. To be specific, section 3 is composed of three parts:

1. Deriving exact solutions in parameter space. In order to solve analytically the coupled system of ODEs we need whitened input.
2. Interpreting the dynamics as preconditioned gradient flow in function space. This analysis does not require whitened input. We only require that $X^\intercal X$ is full-rank such that there exists a unique OLS solution. As shown in equation 2 we did not replace $X^\intercal X$ with $\mathbf{I}_d$.
3. Identifying the implicit bias when $X^\intercal X$ is low-rank. Actually the main theorem (Theorem A.2) in this analysis does not make any assumptions on the dimension of the null space. This theorem expresses the interpolating solution as a solution to a constrained optimization problem (which in general can only be solved numerically). We can then interpret the objective function being minimized in different limits of $\delta$. If we additionally assume the null space is one-dimensional then we can analytically express the solution to this constrained optimization problem.

In summary, we do need the whitened input assumption for the first part of section 3, however the remaining subsections relax this assumption to full-rank and then low-rank. We have added comments in this section to make this more clear and we have brought theorem A.2 up from the appendix into the main to make it very clear that the assumption on a one-dimensional null space is only needed to solve analytically the constrained optimization problem in this theorem.

Missing citation. Thank you for this additional reference we were unaware of. We have added this citation to line 398 as suggested.

Formula for the NTK with unbalanced LRs. The formula (equation 10) in the paper you referenced seems incorrect to us. Here is an explanation for the form of the NTK in our paper (we also added this derivation to the top of appendix A):

As discussed in our notation section, $K_{ij} = \Theta(x_i, x_j; \theta) = \sum_{k = 1}^p \eta_{\theta_k} \partial_{\theta_k} f(x_i;\theta)\partial_{\theta_k} f(x_j;\theta)$. For the minimal model we study $f(x;\theta) = aw^\intercal x$ implies $\partial_a f = w^\intercal x$ and $\partial_w f = ax$. Thus, we see that the term associated with the gradient of $a$, and thus the learning rate $\eta_a$, depends on $w$, while the term associated with the gradient of $w$, and thus the learning rate $\eta_w$, depends on $a$. That is why the NTK is defined by the matrix $M = \eta_wa^2 \mathbf{I}_d + \eta_aww^\intercal$ in our expression. It seems to us that the equation from the paper you referenced has a notational mistake.

Clarifying line 145. We agree this sentence is misleading. We were referencing a common simplification used in many infinite width analyses where the features before and after the nonlinearity are assumed to be of the same scale (see “A Spectral Condition for Feature Learning” by Yang et al. for example). These works will often build intuition for their analysis by replacing the nonlinearities with linear activations, without affecting their results. However, as line 145 is currently written, it sounds like the nonlinearities are always linear on the preactivation, which is not true. Thus, we have modified this sentence to read, “In this limit, analyzing dynamics becomes simpler in several respects: random variables concentrate and quantities will either vanish to zero, remain constant, or diverge to infinity.” Also, it is worth noting we have added a longer discussion on how our theory connects to infinite width limit analyses in our updated manuscript (see responses to Reviewers JJEZ and K8a2 for details).

Please let us know if you have any other questions regarding our work. Overall we are happy to hear that you enjoyed our paper and we hope you will consider our paper an important contribution to the NeurIPS community. Thank you!