From Lazy to Rich: Exact Learning Dynamics in Deep Linear Networks

We are grateful for your thorough review and the time you dedicated to suggesting improvements for our study. We are pleased to hear that you view our work as a timely and significant contribution to the field. Your positive remarks on the paper’s strong theoretical foundation, novelty, clarity, and relevance are highly appreciated. We will address each of the weaknesses and questions you raised individually.

Assumptions and focus on a simple two-layer linear network We acknowledge the limitations associated with the assumptions and the two-layer linear setting highlighted by the reviewer. In our general response, we aim to contrast these with prior literature, showcasing how they extend the scope of previous studies. We further demonstrate that relaxing these assumptions still leads to intriguing and qualitatively equivalent regimes.

More intuitive explanations of how ($\lambda$) interpolates between learning regimes The $\lambda$-balanced condition imposes a relationship between the singular values of the two weight matrices. Specifically, if $W_2$ and $W_1$ are $\lambda$-balanced and satisfy $W_2 W_1 = \Sigma_{yx}$, then for arbitrary singular values $a_i$, $b_i$, and $s_i$ (singular values of $W_2$, $W_1$, and $\Sigma_{yx}$ respectively), the following relations hold: $a_i^2 - b_i^2 = \lambda, \quad a_i \cdot b_i = s_i.$ As $\lambda$ increases, the value of $b_i$ must decrease. In the limit as $\lambda \to \infty$, $a_i^2 \to \lambda$ and $b_i^2 \to 0$. From the first equation, when $b_i^2 \to 0$, $a_i^2 \to \lambda$. Since these equations apply to all singular values of the matrices, it follows that for all $i$, $a_i^2 \to \lambda$, leading to the conclusion that $W_2^T W_2 = \lambda I,$ as expected. Consequently, the task representation becomes task-agnostic in $W_1$. The intuition here is that the weights are constrained by the need to fit the data, which bounds their overall norms. The $\lambda$-balanced condition further specifies a relationship between these norms, and as $|\lambda|$ increases, this constraint tightens, driving $W_2$ toward the identity matrix. In this regime, the network behaves similarly to a shallow network, with $\lambda$ acting as a toggle between deep and shallow learning dynamics. In Appendix C of the paper, we rigorously prove this behavior and further demonstrate that, as a result, the Neural Tangent Kernel (NTK) remains static in this regime, confirming that the network operates in the lazy learning regime. We included this discussion in Appendix C 2.4 to improve clarity.

Scale with relative scale You are correct that the relative scale and absolute scale interact in non-trivial ways. While our study primarily focuses on the effects of relative scale, the influence of absolute scale is inherently embedded within our framework through the definitions of $B$, $C$, and $D$ (see Equations 10, 11, and 12). However, this influence is not immediately apparent from the main theorem. A clearer distinction between the roles of relative and absolute scale can be observed in Theorem 5.1. Specifically, we investigated how $\lambda$, the relative scale parameter, governs the transition between sigmoidal and exponential dynamical regimes. A similar argument applies to absolute scale, which appears explicitly as $s_\alpha(0)$ in these equations. Consider the case when $\lambda = 0$. The dynamics of $s_\alpha$ reduce to the classical solution of the Bernoulli differential equation. In the limiting case where $s_\alpha(0) \to 0$, the system exhibits classic sigmoidal dynamics (characteristic of the rich regime), whereas the limit $s_\alpha(0) \to \infty$ yields exponential dynamics (characteristic of the lazy regime). This interplay between relative and absolute scales underscores their critical roles in shaping the system's behavior. A straightforward intuition can also be gained by considering the scalar case where $N_i = N_h = N_o = 1$. In this scenario, it is easy to ensure that $w_1^2 = w_2^2$ satisfying $\lambda = 0$ while allowing for different absolute scales. For instance, $w_1 = w_2 = 0.001$ or $w_1 = w_2 = 5$. In such cases, the absolute scale is clearly decoupled from the relative scale. Altogether, the absolute scale and relative scale of the weights play a critical role in describing the phase portrait of the learning regime, as first demonstrated in the Kunin et al., 2024 paper on ReLU networks. Thank you for this question, this is an insightful point; we will add a discussion similar to the one above in Appendix A.4.

Questions

• Question a: See above
• Question b: Please refer to the response above as well as the response to reviewer ifg2 question 4. .

I commend the authors for their additional clear explanations, expanded applications, and detailed responses to the reviewers' questions. The explanations regarding the intuition and interaction between absolute and relative scales are particularly helpful.

That said, I share reviewer ifg2's concerns about the technical and theoretical novelty of this work. The method employed in this paper has largely been established in Braun et al. (2022), with this work extending it to the $\lambda$-balanced case. Nonetheless, I believe the paper should still be accepted due to its technical soundness and the novel insights derived from the relaxation.

I would like to raise my score to a 7, but this does not appear to be an option in the ICLR score dropdown (the closest choices being 6 and 8). Thus, I will maintain my current score, which is already above the acceptance threshold. That said, I have increased the soundness score from 3 to 4 in light of the authors' rebuttal.

Questions

1. You raise a valid point. In the case of a linear network, the minimum required to learn a task in the most general sense is at least one layer with $N_i \times N_o$ parameters. As $N_h > min(N_i, N_o)$, it ensures that the network exceeds this parameter count, thus fitting the definition of overparameterization. However, we acknowledge that this term can be confusing, as it is not commonly used in this way within practical machine learning contexts. We will adopt the simpler definition $N_h > min(N_i, N_o)$ to avoid ambiguity.
2. We respond to this point above in the discussion of the assumptions. 
3. Thank you for pointing out the typo in Fig. 2. $\lambda$ should be equal to zero. We have corrected it in the updated version of the paper.
4. The claim aligns with the understanding that we achieve a fixed NTK without involving a large width. As outlined in the theorem, this is attributed to the NTK converging to an identity matrix, where the dominant singular values of the weights approach a scaled identity proportional to $\lambda$, and the corresponding weights scale inversely with $\lambda$. Intuitively, in this setup, the larger weights function as an identity-like projection, while the smaller weights adapt and align. However, due to their relatively small scale compared to the larger weights, their contribution to the NTK remains negligible. (It is important to recall that the NTK is determined by the similarity matrix of the input and output representations.) In the nonlinear setting, this behavior is not expected to hold, as an additional factor comes into play in the computation of the NTK: the activation coefficients of the nonlinearity, as demonstrated in Kunin et al., 2024. In that case, large relative weight (large positive $\lambda$) leads to a rapid rich regime. In the infinite-width regime, where weights are initialized from a Gaussian distribution with large variance, averaging effects cause both input and output representations to approximate identity matrices. In this scenario, the network learns with minimal parameter variation, operating in the lazy regime with a fixed Neural Tangent Kernel (NTK). This behavior contrasts with the dynamics observed in the current setting since both input and output representations are task agnostic.
5. Thank you for bringing the typo in Figure 4 to our attention. We have updated the figure in the manuscript.
6. As Nh​ approaches the maximum of Ni​ and No​, the importance of the delayed rich setting decreases. This happens because, in the delayed rich setting, no least-squares solution exists within the network's span at initialization. In such situations, the network shifts into a delayed rich regime where $\lambda$ approaches infinity, with the magnitude of $\lambda$ dictating the extent of the delay. Initially, the network demonstrates lazy behavior, attempting to fit the solution. However, as constraints require adjustments in its directions, the network transitions to a rich phase. Consequently, when the network's dimensions are closer to spanning the full space, the delayed rich regime becomes less relevant.

Thank you for taking the time to review our paper carefully and point out areas that could benefit further clarification. We're glad to hear your positive feedback on the paper's novelty, clarity, and significance. We will address each of the concerns you highlighted in detail to strengthen your confidence in the importance of our research.

Weaknesses

• The technical and theoretical novelty This work builds on the methodology introduced by Braun et al. (2022) while making substantial advancements and extensions. The updated version eliminates several critical assumptions that constrained the earlier study, including full-rank conditions, the invertibility of the variable B, and the extension to the $\lambda$-balanced condition. Importantly, we emphasize that our general solution does not rely on task-aligned initialization (Theorem 5.1); this assumption is only later used to gain further intuition about the dynamics of singular values and the transition between the rich and lazy regimes. By removing these limitations, this work not only extends the scope of prior research but also broadens the analytical framework. In the general response to reviewers, we detail the unique contributions that distinguish our approach from those of previous studies. Additionally, we show that multiple assumptions can be relaxed without compromising the qualitative results of the primary findings, thereby addressing the first two weaknesses identified.

• Application We extend the discussion of the applications in the new version of the paper to highlight their relevance. For instance, we delve deeper into how, in the transfer learning setting, increasing $\lambda$ enables networks to more effectively transfer the hierarchical structure to new features for untrained items. Additionally, we introduce a new application focused on fine-tuning dynamics and highlight its connection to LoRA fine-tuning techniques in large language models.

Thank you for taking the time to carefully review our paper and point out areas that could benefit from further clarification. We appreciate your positive feedback on the paper’s quality, clarity and significance. We will address each of the concerns you highlighted in detail, with the aim of strengthening your confidence in the importance of our research.

Application We extend the discussion of the applications in the new version of the paper to highlight their relevance. For instance, we delve deeper into how, in the transfer learning setting, increasing $\lambda$ enables networks to more effectively transfer the hierarchical structure to new features for untrained items. Additionally, we introduce a new application focused on fine-tuning dynamics and highlight its connection to LoRA fine-tuning techniques in large language models.

Results to generalize to more complex tasks In this work, initializations are task-agnostic, meaning they do not depend on the specific structure of the task. This observation holds true for any learning task with a well-defined input-output correlation. Beyond the linear network setting, prior work suggests that many of the phenomena we study also take place in non-linear networks. For example, the work by Kunin et al., 2024 demonstrates that positive $\lambda$ initialization enhances the interpretability of early layers in CNNs and reduces the time required for grokking in modular arithmetic tasks. Therefore, although more research is needed, we are cautiously optimistic that our results also shed some light on the dynamics of non-linear networks in complex tasks.

Transfer learning shows a performance increase. In the updated version of the manuscript, we emphasize in the main text the improvements in transfer learning performance as a function of the loss. Specifically, we highlight the results presented in Figure 11 of the appendix, which demonstrate that the generalization loss on untrained items with the new feature decreases as a function of increasing $\lambda$.  Therefore, as $\lambda$ increases, networks more effectively transfer the hierarchical structure of the network to the new feature for untrained items, leading to an increase in generalization performance. This finding effectively shows that performance improves with increasing  $\lambda$.

Further Simplified Assumptions:

• Whitened Input Assumption The assumption of whitened inputs, though strong, is a commonly used simplification in analytical studies to enable the derivation of exact solutions, as shown in previous works (e.g., Fukumizu et al., 1998; Kunin et al., 2024; Saxe et al., 2014) and notably in Braun et al., 2022. While relaxing this assumption prevents the exact description of network dynamics, Kunin et al., 2024, examine the implicit bias of the training trajectory without relying on whitened inputs. Their findings demonstrate a similar quantitative dependence on $\lambda$, governing the implicit bias transition between rich and lazy regimes. Furthermore, recent advancements, such as the "decorrelated backpropagation" technique introduced by Dalm et al., 2024, which whitens inputs during training, showing that optimizing for whitened inputs can actually be done in practice and improve efficiency in real-world applications. Importantly, this study highlights that in certain real-world scenarios, whitening can provide a more optimal learning condition. These approaches emphasize the potential advantages of input whitening for downstream tasks, reinforcing the validity of our assumption.

• Input-Output Dimensionality Previous works imposed specific dimensionality constraints. For example: Fukumizu et al., 1998 assumed equal input and output dimensions ($N_i = N_o$) while allowing a bottleneck in the hidden dimension ($N_h \leq N_i = N_o$). Braun et al., 2022 extended these solutions to cases with unequal input and output dimensions ($N_i \neq N_o$) but restricted bottleneck networks ($N_h = \min(N_i, N_o)$) and introduced an additional invertibility condition on ${ B}$. In our work, we allow for unequal input and output $N_i \neq N_o$ and do not introduce an additional invertibility assumption. This flexibility expands the applicability of our framework to a wider range of architectures.

• Balancedness Assumption A significant departure from prior works is the relaxation of the balancedness assumption: Earlier studies, such as Fukumizu et al., 1998 and Braun et al., 2022, assumed strict zero-balancedness $W_1W_1^T = W_2^T W_2$, which constrained the networks to the rich regime. Our approach generalizes this to $\lambda$-balancedness, enabling exploration of the continuum between the rich and lazy regimes. While some efforts, such as Tarmoun et al., 2021, have explored removing the zero-balanced constraint, their solutions were limited to unstable or mixed forms. In contrast, our methodology systematically studies different learning regimes by varying initialization properties, particularly through the relative scale parameter. This allows controlled transitions between regimes, advancing understanding of neural network behavior across the spectrum. Other studies, such as Kunin et al., 2024 and Xu et al., 2023, have also relaxed the balancedness assumption, though their analysis was restricted to single-output neuron settings. We emphasize the importance of this balanced quantity by rigorously proving that standard network initializations (e.g., LeCun initialization, He initialization) are $\lambda$-balanced in expectation, and in the infinite width limit they are balanced in probability (Appendix A.3). Furthermore, previous studies such as Kunin et al.,2024 have demonstrated that the relative scaling of $\lambda$ significantly impacts the learning regime in practical scenarios, highlighting the crucial role of dynamical studies of networks as a function of this parameter.

• Full Rank
  Previous work by Braun et al., 2022, imposed a full-rank initialization condition, defined as $\text{rank}(W_1(0)W_2(0)) = N_i = N_o$. However, this assumption is not necessary in our framework.

We hope this clarifies the contribution. We updated the manuscript accordingly to improve clarity.