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

Authors: clémentine c j dominé
nicolas  anguita
alexandra m proca
lukas  braun
daniel  kunin
pedro a m mediano
andrew m saxe
Venue: NeurIPS_cc_2024_Workshop_UniReps
Published: 2024-09-24 13:44:28

Abstract: Biological and artificial neural networks create internal representations for complex tasks. In artificial networks, the ability to form task-specific representations is shaped by datasets, architectures, initialization strategies, and optimization algorithms. Previous studies show that different initializations lead to either a lazy regime, where representations stay static, or a rich regime, where they evolve dynamically. This work examines how initialization affects learning dynamics in deep linear networks, deriving exact solutions for $\lambda$-balanced initializations, which reflect the weight scaling across layers. These solutions explain how representations and the Neural Tangent Kernel evolve from rich to lazy regimes, with implications for continual, reversal, and transfer learning in neuroscience and practical applications.

Main paper: 
 introduction: 
 Biological and artificial neural networks learn internal representations that enable complex tasks such as categorization, reasoning, and decision-making. Both systems often develop similar representations from comparable stimuli, suggesting shared information processing mechanisms Yamins et al. (2014) .Although not yet fully understood, this similarity has garnered significant interest from neuroscience, AI, and cognitive science Haxby et al. (2001) ; Laakso & Cottrell (2000) ; Morcos et al. (2018) ; Kornblith et al. (2019) ; Moschella et al. (2022) . The success of neural models relies on their ability to form these representations and extract relevant features from data to build internal representations, a complex process that in machine learning is defined by two regimes: lazy and rich Saxe et al. (2014) ; Pennington et al. (2017) ; Chizat et al. (2019) ; Bahri et al. (2020) . Despite significant advances, these learning regimes and their characterization are not yet fully understood and would benefit from clearer theoretical predictions, particularly regarding the influence of prior knowledge (initialization) on the learning regime. We discuss related works in the appendix A.Our contributions. (1) We derive exact solutions for the gradient flow in unequal-input-output two-layer deep linear networks, under a broad range of lambda-balanced initialization conditions (Section 2). (2) We model the full range of learning dynamics from lazy to rich, showing that this transition is influenced by a complex interaction of architecture, relative scale, and absolute scale, (Section 3). (3) We present applications relevant to both the neuroscience and machine learning field, providing exact solutions for continual learning dynamics, reversal learning dynamics, and transfer learning (Section 4). 
exact learning dynamics: 
 Setting Consider a supervised learning task where input vectors x n ∈ R Ni , from a set of P training pairs {(x n , y n )} P n=1 , need to be mapped to their corresponding target output vectors y n ∈ R No . We learn this task with a two-layer linear network model that produces the output predictionŷ n = W 2 W 1 x n , with weight matrices W 1 ∈ R N h ×Ni and W 2 ∈ R No×N h , where N h is the number of hidden units. The network's weights are optimized using full batch gradient descent with learning rate η (or respectively time constant τ = 1 η ) on the mean squared error loss L(ŷ, y) = 1 2 ||ŷ -y|| 2 , where ⟨•⟩ denotes the average over the dataset. The dynamics are completely determined by the input covariance and input-output correlation matrices of the dataset, defined as Σxx = 1 P P n=1 x n x T n ∈ R Ni×Ni and Σyx = 1 P P n=1 y n x T n ∈ R No×Ni , and the initialization W 2 (0),W 1 (0). Our objective is to describe the entire dynamics of the network's output and internal representations based on this initialization and the task statistics. We consider an approach first introduced in the foundational work of Fukumizu Fukumizu (1998) and extended in recent work by Braun et al. (2022) , which rather than consider the dynamics of the parameters directly, we consider the dynamics of a matrix of the important statistics. In particular, definingQ = W 1 W T 2 T ∈ R (Ni+No)×N h , we consider the (N i + N o ) × (N i + N o ) matrix QQ T (t) = W T 1 W 1 (t) W T 1 W T 2 (t) W 2 W 1 (t) W 2 W T 2 (t), which is divided into four quadrants with interpretable meanings. The approach tracks several key statistics collected in the matrix. The offdiagonal blocks contain the network function Ŷ(t) = W 2 W 1 (t)X, which can be used to evaluate the dynamics of the loss as shown in Fig. 1 . The on-diagonal blocks capture the correlation structure of the weight matrices, allowing for the calculation of the temporal evolution of the network's internal representations. This includes the representational similarity matrices (RSM) of the neural representations within the hidden layer, as first defined by Braun et al. (2022) ,RSM I = X T W T 1 W 1 (t)X, RSM O = Y T (W 2 W T 2 (t)) + Y, where + denotes the pseudoinverse; and the network's finite-width NTK Jacot et al. (2018) ; Lee et al. (2019) ; Arora et al. (2019b) NTK = I No ⊗ X T W T 1 W 1 (t)X + W 2 W T 2 (t) ⊗ X T X, where I is the identity matrix and ⊗ is the Kronecker product. Hence, the dynamics of QQ T describes the important aspects of network behaviour.We extend previous solutions Fukumizu (1998) ; Braun et al. (2022) ; Kunin et al. (2024) and derive exact solutions for the dynamics of QQ T in unequal-input-output under a broad range of lambdabalanced initialization conditions. See Appendix B.2 for a further discussion of the assumptions and their relation to previous works. The proof of the Theorem and lemma leading to the theorem is in Appendix C. With this solution we can calculate the exact temporal dynamics of the loss, network function, RSMs and NTK (Fig. 1A, C ) over a range of lambda-balanced initializations. Implementation and simulation. Simulation details are in Appendix F.7. 
rich and lazy learning: 
 In this section we aim to gain a deeper understanding of the transition between the rich and lazy regimes by examining the dynamics as a function of lambda -the relative scale -as it varies between positive and negative infinity.Here we examine a lambda-balanced linear network initialized with task-aligned weights. Previous research Saxe et al. (2019a) has demonstrated that initial weights that are aligned with the task remain aligned throughout training, restricting the learning dynamics to the singular values of the network. As shown in Fig. 4 B, as λ approaches zero, the dynamics resemble sigmoidal learning curves that traverse between saddle points, characteristic of the rich regime Braun et al. (2022) . In this regime the network learns the most salient features first, which can be beneficial for generalization Lampinen & Ganguli (2018) . Conversely, as shown in Fig. 4 A and C, as the magnitude of λ increases, the dynamics become exponential, characteristic of the lazy regime. In this regime, all features are treated equally and the network's dynamics resemble that of a shallow network. Overall, our results highlight the critical influence of both the absolute scale and the relative scale λ has in shaping the learning dynamics, from sigmoidal to exponential, steering the network between the rich and lazy regimes. The proof can be found in Appendix D.1. The dynamics of the representations. We examine how the representations of the parameters W 1 and W 2 evolve during training. With lambda-balanced initializations, the structure persists throughout training, allowing us to recover the dynamics up to a time-dependent orthogonal transformation. The singular values S λ of the weights are adjusted based on λ, splitting the representation into two parts (Theorem D.2). Using QQ T (t), we capture the temporal dynamics of hidden layer activations and analyze whether the network adopts a rich or lazy representation, depending on λ. Upon convergence, the internal representation satisfiesEQUATIONW 2 W T 2 = ŨS 2 2ŨT , with detailed proof in Theorem D.3. For a hierarchical semantic task Saxe et al. (2014) ; Braun et al. (2022) , the representational similarity of inputs ( ṼS ṼT ) and targets ( ŨS ŨT ) aligns with the task structure. When training a two-layer network with relative scale λ = 0, the representational similarity matrices match the task upon convergence (Theorem D.3). As λ approaches positive or negative infinity, the network transitions to the lazy regime, adopting task-agnostic representations (Theorem D.4, Fig. 2 ). The NTK becomes static and identity-like, while downscaled representations remain structured. This property can enhance generalization during fine-tuning, as shown in Section 4. In contrast, large Gaussian initializations result in lazy learning, lacking structural representation. We propose a new semistructured lazy regime, where initialization determines whether task-specific features reside in input or output layers. Representation robustness and sensitivity to noise. In appendix D.3 we examine the relationship between the learning regime and the robustness of the learned representations to added noise in the inputs and parameters. In practice, parameter noise could be interpreted as the noise occurring within the neurons of a biological network. We find that a rich solution may enable a more robust representation in such systems.The impact of the architecture. Thus far, we have found that the magnitude of the relative scale parameter λ determines the extent or rich and lazy learning. Here, we explore how a network's learning regime is influenced by the interaction of its architecture and the sign of the relative scale. We consider three types of network architectures, depicted in Fig. 3A : funnel networks, which narrow from input to output (N i > N h = N o ); inverted-funnel networks, which expand from input to output (N i = N h < N o ); and square networks, where input and output dimensions are equal (N i = N h = N o ). Our solution, QQ T , captures the NTK dynamics across these different network architectures. To examine the NTK's evolution under varying λ initializations, we compute the kernel distance from initialization, as defined in Fort et al. (2020) . As shown in Fig. 3B , we observe that funnel networks consistently enter the lazy regime as λ → ∞, while inverted-funnel networks do so as λ → -∞. The NTK remains static during the initial phase, rigorously confirming the rank argument first introduced by Kunin et al. ( 2024) for the multi-output setting. In the opposite limits of λ, these networks transition from a lazy regime to a rich regime. During this second alignment phase, the NTK matrix undergoes changes, indicating an initial lazy phase followed by a delayed rich phase. We further investigate and quantify this delayed rich regime, showing the NTK movement over training in Fig. 3C . This behavior is also quantified in Theorem D.6, which describes the rate of learning in this network. For square networks with equal input and output dimensions, this behavior is discussed in Section 3. Across all architectures, as λ → 0, the networks consistently transition into the rich regime. Altogether, we further characterize the delayed rich regime in wide networks.  
application: 
 Continual learning. In line with the framework presented by Braun et al. (2022) , our approach describes the exact solutions of the networks dynamics trained across a sequence of tasks. As detailed in Appendix E.1, we demonstrate that, regardless of the chosen value of lambda, training on subsequent tasks can result in the overwriting of previously acquired knowledge, leading to catastrophic forgetting McCloskey & Cohen (1989) ; Ratcliff (1990); French (1999) .Reversal learning. As demonstrated in Braun et al. (2022) , reversal learning theoretically does not succeed in deep linear networks as the initalization aligns with the separatrix of a saddle point. While simulations show that the learning dynamics can escape the saddle point due to numerical imprecision, the process is catastrophically slowed in its vicinity. However, when λ is non-zero, reversal learning dynamics consistently succeed, as they avoid passing through the saddle point due to the initialization scheme. This is both theoretically proven and numerically illustrated in Appendix E.2. We also present a spectrum of reversal learning behaviors controlled by the relative scale λ, ranging from rich to lazy learning regimes. This spectrum has the potential to explain the diverse dynamics observed in animal behavior, offering insights into the learning regimes relevant to various neuroscience experiments.Transfer learning. We consider how different λ initializations influence generalization to a new feature after being trained on an initial task. We observe in Appendix figure 7 that the task specific structure of the data is effectively transferred to the new feature when the representation is taskspecific and λ is zero. Conversely, when the output feature representation is lazy, meaning the hidden representation lacks adaptation, no task specific generalization is observed. Strikingly, when λ is positive, the task specific structure in the input weights remains small but structured, while the output weights exhibit a lazy representation and the network generalizes to the task specific features. This suggests that the lazy regime structure can be beneficial for transfer learning. 
discussion: 
 We derive exact solutions to the learning dynamics within a tractable model class: deep linear networks. We examine the transition between the rich and lazy regimes by analyzing the dynamics as a function of λ-the relative scale-across its full range from positive to negative infinity. Our analysis demonstrates that the relative scale, λ, plays a crucial role in managing the transition between rich and lazy regimes. learning in the rich and lazy settings. Clementine took the lead on the reversal learning analysis. She was primarily responsible for leading the manuscript writing, with assistance from Daniel. Initial work on this project was carried during Nicolas masters supervised by Clementine and Pedro. All authors contributed to writing the appendix and refining the final manuscript.Lazy regime. Extensive research has identified a fundamental phenomenon in overparameterized neural networks: during training, these networks frequently remain near their linearized form, undergoing minimal changes in the parameter space Chizat et al. (2019) . Consequently, they adopt learning dynamics akin to kernel regression, characterized by the Neural Tangent Kernel (NTK) matrix and exhibiting exponential learning behavior Du et al. (2018) ; Jacot et al. (2018) ; Du et al. (2019) ; Allen- Zhu et al. (2019a,b) ; Zou et al. (2020) . This behavior, known as the lazy or kernel regime, typically occurs in infinitely wide architectures and can be triggered by large variance initialization at the start of training Jacot et al. (2018) ; Chizat et al. (2019) . While the lazy regime offers valuable insights into how networks converge to a global minimum, it does not fully account for the generalization capabilities of neural networks trained with standard initializations. It is, therefore, widely believed that another regime, driven by small or vanishing initializations, underpins some of the successes of neural networks.Rich regime. In contrast, the rich feature-learning regime is characterized by a NTK that evolves throughout training, accompanied by non-convex dynamics that navigate saddle points Baldi & Hornik (1989) ; Saxe et al. (2014 Saxe et al. ( , 2019b)) ; Jacot et al. (2021) . This regime features sigmoidal learning curves and simplicity biases, such as low-rankness Li et al. (2020) or sparsity Woodworth et al. (2020) . Numerous studies have shown that the absolute scale of initialization drives the rich regime, which typically emerges at small initialization scales Chizat et al. (2019); Geiger et al. (2020) . However, it's also been shown that even at small initialization scales, differences in weight magnitudes between layers can induce the lazy learning regime Azulay et al. (2021); Kunin et al. (2024) . This highlights the significance of both absolute scale (initialization variance) and relative scale (difference in weight magnitude between layers) in generating diverse learning dynamics. Beyond absolute scale and relative scale, additional aspects of initialization can profoundly affect feature learning, including the effective rank of the weight matrices Liu et al. ( 2023), layer-specific initialization variances Yang & Hu (2020) ; Luo et al. (2021) ; Yang et al. (2022) , and the use of large learning rates Lewkowycz et al. (2020) ; Ba et al. (2022) ; Zhu et al. (2023) ; Cui et al. (2024) . These findings illustrate the effect of initialization on inducing complex learning behavior through the resulting dynamics. Here we develop a solvable model which captures these diverse phenomena.Rich and lazy regimes in the brain. The distinction between rich and lazy learning may also hold implications for neuroscience, where neural representations have been argued to have task-specific or task-agnostic characteristics in different settings Farrell et al. (2023a) ; Ostojic & Fusi (2024) ; Tye et al. (2024) . The lazy regime can be linked to the non-linear mixed selectivity of neurons, where task variables are represented in a high-dimensional space which mixes various potentially relevant variables Raposo et al. (2014) ; Tang et al. (2019) ; Rigotti et al. (2013) ; Bernardi et al. (2020) . Conversely, the rich regime aligns with linear mixed selectivity Tye et al. (2024) and the manifold learning regime, where the brain encodes tasks on a structured, low-dimensional, taskspecific manifold, as observed in grid cells within the entorhinal cortex Chaudhuri et al. (2019) ; Bernardi et al. (2020); Flesch et al. (2022) .Linear networks. Our work builds upon a rich body of research on deep linear networks, which, despite their simplicity, have proven to be valuable models for understanding more complex neural networks Baldi & Hornik (1989) ; Fukumizu (1998) ; Saxe et al. (2014) . 2019). This work directly builds on the matrix Riccati formulation proposed by Fukumizu (1998) ; Braun et al. (2022) which extends these solutions to wide networks. We extend and refine these results to obtain the dynamics for lambda-balanced initialization dynamics of networks to more clearly demonstrate the impact of initialization on rich and lazy learning regimes also developed in Tu et al. (2024) for a set of orthogonal initalizations. Our work extends previous analysis Xu & Ziyin (2024) ; Kunin et al. (2024) of these regime to wide networks. Previous studies leveraged these solutions primarily to characterize convergence rates; however, our work goes beyond this by providing a comprehensive characterization of the complete dynamics of the system Tarmoun et al. ( 2021).Infinite-width networks. Recent advances in understanding the rich regime have largely stemmed from examining how the initialization variance and layer-wise learning rates must scale in the infinite-width limit to maintain consistent behavior in activations, gradients, and outputs. Several studies have employed statistical mechanics tools to derive analytical solutions for the rich population dynamics of two-layer nonlinear neural networks initialized using the mean field parameterization Mei et al. (2018) ; Rotskoff & Vanden-Eijnden (2018); Chizat & Bach (2018); Sirignano & Spiliopoulos (2020); Rotskoff & Vanden-Eijnden (2022); Sirignano & Spiliopoulos (2020) 2019). Furthermore, these approaches typically require numerical integration or operate within a limited learning regime, and are unable to describe the learning dynamics of hidden representations. Instead, our work focuses on the impact of initialization on representation learning dynamics and derives explicit analytical solutions within tractable models. 2023))). The weights W 1 , W 2 are λ-balanced if and only if there exists a Balanced Coefficient λ ∈ R such that:B(W 1 , W 2 ) = W T 2 W 2 -W 1 W T 1 = λI (1)where B is called the Balanced Computation.For λ = 0 we have Zero-Balanced given as A1 (). W 1 (0)W 1 (0) T = W 2 (0) T W 2 (0). Theorem B.2Suppose at initializationW 2 (0) T W 2 (0) -W 1 (0)W 1 (0) T = λI (2)Then for all t ≥ 0EQUATIONProof. Consider:τ d dt W 2 (t)W 2 (t) T -W 1 (t)W 1 (t) T = τ d dt W 2 (t) T W 2 (t) + W 2 (t) T τ d dt W 2 (t) -τ d dt W 1 (t) W 1 (t) T -W 1 (t) τ d dt W 1 (t) T = W 1 (t) Σyx -W 2 (t)W 1 (t) Σxx T W 2 (t) + W 2 (t) T Σyx -W 2 (t)W 1 (t) Σxx W 1 (t) -W 2 (t) T Σyx -W 2 (t)W 1 (t) Σxx W 1 (t) -W 1 (t) Σyx -W 2 (t)W 1 (t) Σxx W 2 (t) = 0 Note that W 2 (t) T W 2 (t) -W 1 (t)W 1 (t)T is conserved for any initial value λ.• A2 (Whitened input). The input data is whitened, that is Σxx = I.• A3 (Lambda-balanced). The network's weight matrices are lambda-balanced at the beginning of training, that isW 2 (0) T W 2 (0) -W 1 (0)W 1 (0) T = λI.If this condition holds at initialization, it will persist throughout training Saxe et al. (2014) ; Arora et al. (2018a) . For completeness, we prove this in Appendix B.• A4 (Dimensions). The hidden dimension of the network is defined asN h = min(N i , N o ), ensuring the network is neither bottlenecked (N h < min(N i , N o )) nor overparameterized (N h > min(N i , N o )).• A5 (Full-rank). The input-output correlation of the task and the initial state of the network function have full rank, that is rank(Σxy ) = rank(W 2 (0)W 1 (0)) = min(N i , N o ).Whittened Inputs. Although the whitened input assumption is quite strong, it is commonly used in analytical work to obtain exact solutions, and much of the existing literature relies on these solutions Fukumizu (1998); Braun et al. (2022); Kunin et al. (2024) . Kunin et al. (2024) goes further by exploring the implicit bias of the trajectory without relying on exact solutions. When X ⊺ X is lowrank, they can only predict the trajectories in the limit as λ → ±∞. If the interpolating manifold is one-dimensional, the solution can be solved exactly in terms of λ (black dots).Dimension. Fukumizu assumed equal input and output dimensions N i = N o , but allowed for a bottleneck in the hidden dimension of the network N h ≤ N i = N o . The work by Braun et al. (2022) extended Fukumizu (1998) solutions to cases with unequal input and output dimensions N i ̸ = N o , but to so did not allow a bottleneck N h = min{N i , N o } and added an assumption on the invertibility of a statistic of the singular vector overlap between the model and the input-output statistics. In our work we allow for unequal input and output N i ̸ = N o and do not introduce an additional invertibility assumption.Balancedness. The main distinction between our work and prior works is that both Fukumizu (1998) and Braun et al. (2022) assumed zero-balanced W 1 (0)W 1 (0) T = W 2 (0) T W 2 (0), while we relax this assumption to λ-balanced. The zero-balanced condition restricts the networks to a rich setting. We develop solutions to explore the continuum between the rich and the lazy regime.While some works, such as Tarmoun et al. ( 2021), have considered removing this constraint, their solutions remain in an unstable and mixed form. EQUATIONEQUATIONthe gradient flow dynamics of QQ T (t) can be written as a differential matrix Riccati equationEQUATION)Proof. We introduce the variablesEQUATION)We compute the time derivativeEQUATION)Using equations 18 and 19, we compute the four quadrants separately givingEQUATIONEQUATIONEQUATION)the gradient flow dynamics of QQ T (t) can be written as a differential matrix Riccati equationEQUATION)We write τ d dt (QQ T ) for completenessτ d dt (QQ T ) = -λ 2 ( Σyx ) T Σyx λ 2 W T 1 W 1 W T 1 W 2 W 2 W 1 W 2 W T 2 + W T 1 W 1 W T 1 W 2 W 2 W 1 W 2 W T 2 T -λ 2 ( Σyx ) T Σyx λ 2 - W T 1 W 1 W T 1 W 2 W 2 W 1 W 2 W T 2 2 (26) = -λ 2 ( Σyx ) T Σyx λ 2 W T 1 W 1 W T 1 W 2 W 2 W 1 W 2 W T 2 + W T 1 W 1 W T 1 W 2 W 2 W 1 W 2 W T 2 T -λ 2 ( Σyx ) T Σyx λ 2 - W T 1 W 1 W T 1 W 2 W 2 W 1 W 2 W T 2 W T 1 W 1 W T 1 W 2 W 2 W 1 W 2 W T 2 (27) = -λ 2 W T 1 W 1 + ( Σyx ) T W 2 W 1 -λ 2 W T 1 W 2 + ( Σyx ) T W 2 W T 2 Σyx W T 1 W 1 + λ 2 W 2 W 1 Σyx W T 1 W T 2 + λ 2 W 2 W T 2 + -λ 2 W T 1 W 1 + W T 1 W 1 ( Σyx ) T λ 2 W T 1 W 2 + W T 1 W 2 ( Σyx ) T -λ 2 W T 2 W 1 + W 2 W 1 ( Σyx ) T λ 2 W 2 W T 2 + W 2 W T 2 ( Σyx ) T - W T 1 W 1 W T 1 W 2 W 2 W 1 W 2 W T 2 W T 1 W 1 W T 1 W 2 W 2 W 1 W 2 W T 2 (28) = -λ 2 W T 1 W 1 + ( Σyx ) T W 2 W 1 -λ 2 W T 1 W 2 + ( Σyx ) T W 2 W T 2 Σyx W T 1 W 1 + λ 2 W 2 W 1 Σyx W T 1 W T 2 + λ 2 W 2 W T 2 + -λ 2 W T 1 W 1 + W T 1 W 1 ( Σyx ) T λ 2 W T 1 W 2 + W T 1 W 2 ( Σyx ) T -λ 2 W T 2 W 1 + W 2 W 1 ( Σyx ) T λ 2 W 2 W T 2 + W 2 W T 2 ( Σyx ) T - W T 1 W 1 W T 1 W 1 + W T 1 W 2 W T 2 W 1 W T 1 W 1 W T 1 W 2 + W T 1 W 2 W T 2 W 2 W 2 W 1 W T 1 W 1 + W 2 W T 2 W 2 W 1 W 2 W 1 W T 1 W 2 + W 2 W T 2 W 2 W T 2 (45)The four quadrants of 8 are equivalent to equations 12, 15, 18, and 23 respectively.C.2 QQ T Diagonalisation Lemma C.2. If F = P ΛP T is symmetric and diagonalizable, then the matrix Riccati differential equation τ d dt (QQ T ) = FQQ T + QQ T F -(QQ T ) 2 with initialization QQ T (0) = Q(0)Q(0)T has a unique solution for all t ≥ 0, and the solution is given byEQUATION)This is true even when there exists Λ i = 0.Proof. First we show that there exists a unique solution to the initial value problem stated. This is true by Picard-Lindelöf theorem. Now we show that the provided solution satisfies the ODE. LetL = e F t τ Q(0) and C = I+Q(0) T P e 2Λ t τ -I 2Λ P T Q(0) such that solution QQ T (t) = LC -1 L T .The time derivative of QQ T is then given byEQUATIONSolving for these derivatives individually, we findτ d dt (L) = τ d dt e F t τ Q(0) = F e F t τ Q(0) = F L (31) τ d dt (C -1 ) = -τ C -1 d dt (C)C -1 = -τ C -1 Q(0) T P d dt e 2Λ t τ -I 2Λ P T Q(0)C -1 (32)We consider the derivative of the fraction serpately,EQUATIONthis is true even in the limit as λ i → 0. Plugging these derivatives back in we see that the solution satisfies the ODE. Lastly, let t = 0, we see that the the solution satisfies the initial conditions.In Appendix C.2 we prove that this equation is the unique solution to the initial value problem derived in Lemma C.2 no matter the value of Λ. However, as discussed in Braun et al. (2022) , the solution in this form is not very useable or interpretable due to the matrix inverse mixing the blocks of QQ T . Additionally, we need to diagonalize F . To do so we consider the compact singular value decomposition SVD( Σyx ) = ŨS ṼT . Here, Ũ ∈ R No×N h denote the left singular vectors, S ∈ R N h ×N h the square matrix with ordered, non-zero eigenvalues on its diagonal, and Ṽ ∈ R Ni×N h the corresponding right singular vectors. For unequal input-output dimensions (N i ̸ = N o ), the right and left singular vectors are not square. Accordingly, for the case N i > N h = N o , we define Ũ⊥ ∈ R No×|No-Ni| as a matrix containing orthogonal column vectors that complete the basis for Ũ, i.e., make Ũ Ũ⊥ orthonormal, and Ṽ⊥ ∈ R Ni×|No-Ni| as a matrix of zeros. Conversely, when N i = N h < N o , then Ṽ⊥ is a matrix containing orthogonal column vectors that complete the basis for Ṽ and Ũ⊥ is a matrix of zeros. Using this SVD structure we can now describe the eigendecomposition of F.Lemma C.3. Under assumptions of full-rank 5, the eigendecomposition of F = PΛP T whereEQUATION)and the matrices Sλ , λ ⊥ , H, and G are the diagonal matrices defined as:EQUATIONBeyond the invertibility of F , notice from the equation (Fukumizu solution) we need to understand the relationship between F and Q(0). To do this the following lemma relates the structure between the SVD of the model with the SVD structure of the individual parameters.Proof. We leave for the reader by computingEQUATIONC.4 Solution Unequal-Input-Output Theorem C.4. Under the assumptions of whitened inputs, 2, lambda-balanced weights 3, no bottleneck 4, and full rank 5, the temporal dynamics of QQ T areQQ T (t) = Z 1 A -1 Z T 1 Z 1 A -1 Z T 2 Z 2 A -1 Z T 1 Z 2 A -1 Z T 2 ,where the variablesEQUATIONEQUATIONProof. We start and use the diagonalization of F to rewrite the matrix exponential of F and F . Note that P T P = PP T = I and therefore P T = P -1 .EQUATIONEQUATION)EQUATIONWhereM = 1 √ 2 Ṽ⊥ Ũ⊥ T .Placing these expressions into equation 29 givesEQUATIONEQUATIONwhereEQUATIONPutting it together we get the expressions for Z 1 (t) and Z 2 (t)EQUATIONWe now compute the terms inside the inverseEQUATIONQ(0) T Mλ -1 ⊥ M T Q(0) = 2 W 1 (0) W 2 (0) T 1 √ 2 Ṽ⊥ Ũ⊥ λ ⊥ 0 0 λ ⊥ -1 1 √ 2 Ṽ⊥ Ũ⊥ T W 1 (0) T W 2 (0) = W 1 (0) W 2 (0) T Ṽ⊥ Ũ⊥ λ -1 ⊥ Ṽ⊥ ṼT ⊥ W 1 (0) T λ -1 ⊥ Ũ⊥ ŨT ⊥ W 2 (0) = W 1 (0) Ṽ⊥ λ -1 ⊥ ṼT ⊥ W 1 (0) T + W 2 (0) T Ũ⊥ λ -1 ⊥ ŨT ⊥ W 2 (0) (54) Now 1 2 Q(0) T O e 2Λ t τ -I Λ -1 O T = 1 4 [B -C] e Λ t τ -I Λ -1 B T -C T = 1 4 B e 2 Sλ t τ -I ( Sλ ) -1 B T -C e -2 Sλ t τ -I ( Sλ ) -1 C T (55)Putting it all togetherEQUATIONSo, final form:EQUATIONC.5 Stable solution Unequal-Input-Output Theorem C.5. Given the assumptions of Theorem C.4 further assuming that B is invertible and defining e λ ⊥ t τ = sgn(N o -N i ) λ 2 , the temporal evolution of QQ T is described as follows:EQUATIONProof. We start fromEQUATIONWe extract B -T e -Sλ t τ from all terms as exemplified bellowOe Λt/τ O T Q(0) = 1 2   Ṽ ( G -H G) -( G + H G)e -Sλ t τ C T B -T e -Sλ t τ Ũ ( G + H G) + ( G -H G)e -Sλ t τ C T B -T e -Sλ t τ   B T e Sλ t τ (61)and rewrite the dynamis asEQUATION)where e λ ⊥ t τ = sgn(N o -N i ) λ 2 is a scalarWe follow a similar derivation presented in Braun et al. (2022) and start with the following equationEQUATIONSubstituting our solution into the matrix Riccati equation then yieldsEQUATION)Using the chain rule ∂(AB) = (∂A)B + A(∂B) and the identitiesEQUATIONEQUATIONNext, we note thatEQUATIONandEQUATIONwith EQUATIONEQUATION(0) = RS 1 Ṽ T and W 2 (0) = Ũ S 2 R T , where R ∈ R N h ×N h is an orthonormal matrix, thenthe network function is given by the expression W 2 W 1 (t) = Ũ S(t) Ṽ T where S(t) ∈ R N h ×N h is a diagonal matrix of singular values with elements s α (t) that evolve according to the equation,EQUATION)where sα is the α singular value of S and γ α (t; λ) is a λ-dependent monotonic transition function for each singular value that increases from γ α (0; λ) = 0 to lim t→∞ γ α (t; λ) = 1 defined asEQUATION)where sλ,α = s2 α + λ 2 4 , s λ,α = s α (0) 2 + λ 2 4 , and s α = s α (0). We find that under different limits of λ, the transition function converges pointwise to the sigmoidal (λ → 0) and exponential (λ → ±∞) transition functions,EQUATION)Proof. According to Theorem C.4, the network function is given by the equationEQUATION99) which depends on the variables of the initialization B and C. Plugging the expressions for a taskaligned initialization W 1 (0) and W 2 (0) into these variables we get the following simplified expressions,EQUATIONEQUATION)where we define the diagonal matrices D B and D C for ease of notation. Using these expressions, we now get the following time-dependent expressions for Z 2 (t), A -1 (t), and Z 1 (t),Z 1 (t) = 1 2 Ṽ ( G -H G)e Sλ t τ D B -( G + H G)e -Sλ t τ D C R T (102) Z 2 (t) = 1 2 Ũ ( G + H G)e Sλ t τ D B + ( G -H G)e -Sλ t τ D C R T (103) A(t) = R I + e 2 Sλ t τ -I 4 Sλ D 2 B - e -2 Sλ t τ -I 4 Sλ D 2 C R T (104)Plugging these expressions into the expression for the network function, notice that the R terms cancel each other resulting in following equationEQUATIONNotice that the middle term is simply a product of diagonal matrices. We can factor the numerator of this expressions as,( G2 -H2 G2 )e 2 Sλ t τ D 2 B + ( G -H G) 2 -( G + H G) 2 D B D C -( G2 -H2 G2 )e -2 Sλ t τ D 2 C(106) We can further factor this expression as,EQUATION)Putting it all together we find that S(t) can be expressed as,EQUATION)Now using the relationship between H and G we use the following two identities:EQUATION)Plugging these identities into the previous expression and multiplying the numerator and denominator by Sλ gives,EQUATIONAdd and subtract S 4 Sλ + D 2 C -D 2 B from the numerator such thatEQUATION)Using the form of D B and D C notice the following two identities:D B D C = λ Sλ S -S 2 S 1 , D 2 C -D 2 B = - 4 Sλ SS 2 S 1 + λ 2 4 I (112)From the second identity we can derive a third identity,4 Sλ + D 2 C -D 2 B = 4 S Sλ S -S 2 S 1 (113)Plugging the first and third identities into the numerator for the previous expression gives,EQUATION)Multiply numerator and denominator by Sλ 4 and simplify terms gives the expression,EQUATION)Thus we have found the transition function,EQUATION)We will use our previous identities and the definitions of D We will now show why under certain limits of λ this expression simplifies to the sigmoidal and exponential dynamics discussed in the previous section.Sigmoidal dynamics. When λ = 0, then Sλ = S and S λ = S(0). Notice, that the coefficients for the hyperbolic functions all simplify to SS(0). Using the hyperbolic identity sinh(x) + cosh(x) = e x , we can simplify the expression for the transition function toEQUATION)Dividing the numerator and denominator by SS(0) gives the final expression.Exponential dynamics. In the limit as λ → ±∞ the expressions Sλ → |λ| 2 and S λ → |λ| 2 . Additionally, in these limits because λ 2 4 I ≫ SS(0) then SS(0) + λ 2 4 I → λ 2 4 I. As a result of these simplifications the coefficients for the hyperbolic functions all simplify to λ 2 4 I. As a result we can again use the hyperbolic identity sinh(x) + cosh(x) = e x to simplify the expression asγ(t; λ) = λ 2 4 e |λ| t τ -λ 2 4 I λ 24 e |λ| t τ + S S -S( 0)EQUATION)Dividing the numerator and denominator by λ 2 4 results in all terms without a coefficient proportional to λ 2 vanishing, which simplifying further gives the final expression. The lazy and rich regimes are defined by the dynamics of the NTK of the network. Lazy learning occurs when the NTK is constant, rich learning occurs when it is not. (Farrell et al. (2023b) )The NTK intuitively measures the movement of the network representations through training. As shown in (Braun et al. (2022) ), in specific experimental setup, we can calculate the NTK of the network in terms of the internal representations in a straightforward way:EQUATIONIn order to better understand the effect of λ on NTK dynamics, we first prove some theorems involving the Singular Values of the λ-balanced weights, and the representations of a λ-balanced network.Theorem D.2. Under a λ-Balanced initialization 3, if the network function W 2 W 1 (t) = U (t)S(t)V T (t) is full-rank 5 and we define S λ (t) = S 2 (t) + λ 2 4 I. , then we can recover the parametersEQUATIONProof. We prove the caseN i ≤ N o and N h = min(N i , N o ). The proof for N o ≤ N i follows the same structure. Let U SV T = W 2 (t)W 1 (t)be the Singular Value Decomposition of the product of the weights at training step t. We will useW 2 = W 2 (t), W 1 = W 1 (t) as a shorthand.By properties Singular Value Decomposition, we can writeW 2 = U S 2 R T , W 1 = RS 1 V T ,where R is an orthonormal matrix and S 2 , S 1 are diagonal (possibly rectangular) matrices.The Balanced property states thatW T 2 W 2 -W 1 W T 1 = λI.We know this holds for any t since this is a conserved quantity in linear networks.HenceEQUATIONEQUATIONThe matrices S 1 , S 2 , have shapes (N h , N i ), (N o , N h ) respectively. We introduce the diagonal matrices Ŝ1 of shape (N h , N i ), Ŝ2 of shape (N i , N h ) such that the zero matrix has size (N o -N i , N h ) :EQUATIONFrom the equation above and the fact that Ŝ1 Ŝ2 = S we derive that:EQUATIONHenceEQUATIONWith our solution, QQ T (t), which captures the temporal dynamics of the similarity between hidden layer activations, we can analyze the network's internal representations in relation to the task. This allows us to determine whether the network adopts a rich or lazy representation, depending on the value of λ. Consider a λ-Balanced network training on data Σ yx = Ũ S Ṽ T . We assume that the convergence is toward global minima and B is invertible Theorem D.3. Under the assumptions of Theorem C.5, the network function converges to ŨS ṼT and acquires the internal representation, that isW T 1 W 1 = ṼS 2 1 ṼT and W 2 W T 2 = ŨS 2 2 ŨT Proof.As training time increases, all terms including a matrix exponential with negative exponent in Equation 58 vanish to zero, as S λ = Sλ is a diagonal matrix with entries larger zeroAs training time increases, all terms in the equations vanish to zero. Terms in Equation 58 decay asEQUATIONandEQUATION)where Sλ = Sλ is a diagonal matrix with entries larger zero Therefore, in the temporal limit, eq. 58 reduces toEQUATIONEQUATIONEQUATIONTheorem D.4. Under the assumptions of Theorem C.5, training on data Σ yx = Ũ S Ṽ T , as λ → ∞ the representation tends toW 2 W T 2 = Ũ λI 0 max(0,No-Ni) 0 max(0,No-Ni) 0 Ũ T W T 1 W 1 = 1 λ Ṽ S2 0 max(0,Ni-No) 0 max(0,Ni-No) 0 Ṽ T As λ → -∞ W 2 W T 2 = - 1 λ Ũ S2 0 max(0,No-Ni) 0 max(0,No-Ni) 0 Ũ T , W T 1 W 1 = Ṽ -λI 0 max(0,Ni-No) 0 max(0,Ni-No) 0 Ṽ T As λ → -∞ W 2 W T 2 = - 1 λ Ũ S2 0 max(0,No-Ni) 0 max(0,No-Ni) 0 Ũ T , W T 1 W 1 = Ṽ -λI 0 max(0,Ni-No) 0 max(0,Ni-No) 0 Ṽ T Proof.We start from the representation derived in D.3 and using the Taylor expansion of f (x) = √ 1 + x 2 , we computeEQUATIONHenceEQUATIONSimilarly,EQUATIONSince Ũ , Ṽ are independent of λ:EQUATIONAs |λ| → ∞, one of the network representations approaches a scaled identity matrix, while the other tends toward zero. Intuitively, this suggests that the representations shift less and less as |λ| increases. Next, we demonstrate that the NTK becomes progressively less variable as |λ| grows and ultimately converges to zero.Relationship between λ and the NTK of the network Theorem D.5. Under the assumptions of Theorem C.5, consider a linear network training on data Σ yx = Ũ S Ṽ T . At any arbitrary training time t ≥ 0, let W 2 (t)W 1 (t) = U * S * V * T . Then, 1. For any λ ∈ R:NTK(0) = I No ⊗ X T V √ λ 2 I+4S * 2 -λI 2 0 0 0 V T X + U √ λ 2 I+4S * 2 +λI 2 0 0 0 U T ⊗ X T X (146) NTK(t) = I No ⊗ X T V * √ λ 2 I+4S * 2 -λI 2 0 0 0 V * T + U * √ λ 2 I+4S * 2 +λI 2 0 0 0 U * T ⊗ X T X (147)2. As λ → ∞:EQUATION3. As λ → -∞:EQUATIONProof. Follows by substituting the expressions for the network representations in terms of λ from (Braun et al. (2022) )'s expression for the NTK of a linear network. Similarly, follows from substituting the limit expressions for the network representations and the fact that the Kronecker product is linear in both arguments.The theorem above demonstrates that as |λ| → ∞, the NTK of a λ-Balanced network remains constant. This indicates that the network operates in the lazy regime throughout all training steps. This finding is significant as it highlights the impact of weight initialization on learning regimes.As derived in (Braun et al., 2024) , the expected mean squared error under additive, independent and identically distributed input noise with mean µ = 0 and variance σ 2x isEQUATIONwhere c = 1 2 Tr( Σyy ) -1 2 Tr( Σyx ΣyxT ) is a noise independent constant that only depends on the statistics of the training data. In Theorem D.3 we show that the network function converges to ŨS ṼT and thereforeEQUATIONAs derived in (Braun et al., 2024) , under the assumption of whitened inputs (Assumption 2), in the case of additive parameter noise with µ = 0 and variance σ 2 W , the expected mean squared error isEQUATIONUsing Theorem D.3, we haveEQUATIONandEQUATIONTo find the λ that minimises the expected loss, we substitute the equations for the norms, take the partial derivative with respect to λ and set it to zeroEQUATIONIt follows, that under the assumption that N i = N o , the equation reduces toEQUATION)We note, that the denominator is always positive and therefore, that the left-hand side of the equation is always larger zero for any λ > 0, and smaller than zero for any λ < 0. The euqation is therefore only solved for λ = 0.D.4 Effect of the architecture from the lazy to the Rich Regime Theorem D.6. Under the conditions of Theorem C.5, when λ ⊥ > 0, the network enters a regime referred to as the delayed-rich phase. In this phase, the learning rate is determined by two competing exponential factors:e λ ⊥ t τ e -S2 + λ 2 4 I t τ and e -S2 + λ 2 4 I t τ .As λ increases, different parts of the network exhibit distinct learning behaviors: some components adapt quickly and converge exponentially with lambda, while others are constrained by the singular values of the network, resulting in slower adaptation.Proof. The solution to Theorem C.5 is governed by two time-dependent terms:e -S2 + λ 2 I 4 t τand e λ ⊥ t τ e -S2 + λ 2 4 I t τ .The first term exhibits exponential decay with rate λ, approaching zero as time progresses:lim t→∞ e -S2 + λ 2 I 4 t τ = 0.The second term also decays, but at a rate governed by the singular values S, as λ tends to infinity:lim t→∞ e λ ⊥ t τ e -S2 + λ 2 4 I τ = 0. Since λ ⊥ -S2 + λ 2 4 I > 0,we havelim λ→∞ λ ⊥ -S2 + λ 2 4 I = S.Thus, as λ increases, the convergence rate slows for certain parts of the network (those governed by larger singular values), while other components continue to learn more quickly. This explains the delay observed in the delayed-rich regime.We build upon the derivation presented in Braun et al. (2022) to incorporate the dynamics of continual learning throughout the entire learning trajectory. Utilizing the assumption of whitened inputs, the entire batch loss for the ith task isL i (T j ) = 1 2P ∥W 2 W 1 X i -Y i ∥ 2 F = 1 2P Tr (W 2 W 1 X i -Y i |) (W 2 W 1 X i -Y i |) T = 1 2P Tr W 2 W 1 X i X T i (W 2 W 1 ) T - 1 P Tr W 2 W 1 X i Y T i + 1 2P Tr Y i Y T i = 1 2 Tr W 2 W 1 (W 2 W 1 ) T -Tr W 2 W 1 Σyx T i + 1 2 Tr Σyy i = 1 2 Tr W 2 W 1 -Σyx i W 2 W 1 -Σyx i T -Σyx i Σyx T i + 1 2 Σyy i = 1 2 W 2 W 1 -Σyx i 2 F - 1 2 Tr Σyx i Σyx T i + 1 2 Σyy i c .Hence, the extent of forgetting, denoted as F for task T i during training on task T k subsequent to training the network on task T j , specifically, the relative change in loss, is entirely dictated by the similarity structure among tasks.F i (T j , T k ) = L i (T k ) -L i (T j ) = 1 2 Σyx k -Σyx i 2 F + c - 1 2 W 2 W 1 -Σyx i 2 F -c = 1 2 Σyx k -Σyx i 2 F -W 2 W 1 -Σyx i 2 F .It is important to note that the amount of forgetting is a function of the weight trajectories. Therefore, we have analytical solutions for trajectories of forgetting as well. As first introduced in Braun et al. (2022) , in the following discussion, we assume that the input and output dimensions are equal. We denote the i-th columns of the left and right singular vectors as u i , ũi , and v i , ṽi , respectively.Reversal learning occurs when both the task and the initial network function share the same left and right singular vectors, i.e., U = Ũ and V = Ṽ, with the exception of one or more columns of the left singular vectors, where the direction is reversed:-u i = ũi .It is important to note that if a reversal occurs in the right singular vectors, such that -v i = ṽi , this can be equivalently represented as a reversal in the left singular vectors, as the signs of the right and left singular vectors are interchangeable.In the reversal learning setting, bothB = S 2 Ũ T Ũ ( G + H G) + S 1 V T Ṽ ( G -H G) and C = S 2 Ũ T Ũ ( G -H G) -S 1 V T Ṽ ( G + H G) are diagonal matrices.In the case where lambda is zero, the same argument given in Braun et al. (2022) follows, the diagonal entries of C are zero if the singular vectors aligned and non zero if they are reversed.Similarly, diagonal entries of B are non-zero if the singular vectors are aligned and zero if they are reversed. Therefore, in the case of reversal learning, B is a diagonal matrix with 0 values and thus is not invertible. As a consequence, the learning dynamics cannot be described by Equation 37. However, as B and C are diagonal matrices, the learning dynamics simplify. Let b i , c i , s i and si denote the i-th diagonal entry of B, C, S and S respectively, then the network dynamics can be rewritten asEQUATIONIt follows, that in the reversal learning case, i.e. b = 0, for each reversed singular vector, the dynamics vanish to zeroEQUATIONAnalytically, the learning dynamics are initialized on and remain along the separatrix of a saddle point until the corresponding singular value of the network function decreases to zero and stays there, indicating convergence to the saddle point. In numerical simulations, however, the learning dynamics can escape the saddle points due to the imprecision of floating-point arithmetic. Despite this, numerical optimization still experiences significant delays, as escaping the saddle point is timeconsuming Lee et al. (2022) . In contrast, when the singular vectors are aligned (c = 0), the equation governing temporal dynamics, as described in Saxe et al. (2014) , is recovered. Under these conditions, training succeeds, with the singular value of the network function converging to its target value.EQUATION)Figure 6 : Plot showing the steps to convergence for two tasks: (1) the reversal learning task and (2) a randomly sampled continual learning task across a range of λ values. The reversal learning task exhibits catastrophic slowing at λ = 0.In summary, in the case of aligned singular vectors, the learning dynamics can be described by the convergence of singular values. However in the case of reversal learning, analytically, training does not succeed. In simulations, the learning dynamics escape the saddle due to numerical imprecision, but the learning dynamics are catastrophically slowed in the vicinity of the saddle point as shown in figure 6 .In the case where λ is non-zero, the diagonal of C are also non-zero; this is true regardless of whether they are reversed or aligned. Similarly, the diagonal entries of B remain non-zero whether the singular vectors are aligned or reversed. Therefore, in the case of reversal learning, B is a diagonal matrix with elements that are zero. In figure 6 E.3 Appendix: Generalization and structured learningWe study how the representations learned for different λ initializations impact generalization of properties of the data. To do this, we consider the case where a new feature is associated to a learned item in a dataset and how this new feature may then be related to other items based on prior knowledge. In particular, we first train each network (for different values of -10 ≤ λ ≤ 10) on the hierarchical semantic learning task in Section 3 and then add a new feature (e.g., 'eats worms') to a single item (e.g., the goldfish) (Fig. 7A ), correspondingly increasing the output dimension to represent the novel feature. In order to learn the new feature without affecting prior knowledge, we append a randomly initialized row to W 2 and train it on the single item with the new feature, while keeping the rest of the network frozen. Thus, we only change the weights from the hidden layer to the new feature which may produce different behavior depending on how the hidden layer representations vary based on λ. After training on the new feature-item association, we query the network with the rest of the data to observe how the new feature is associated with the other items. We find that as λ increases positively, the network better transfers the hierarchy such that it projects the feature onto items based on their distance to the trained item (Fig. 7B, C ). For example, after learning that a goldfish eats worms, the network can extrapolate the hierarchy to infer that another fish, or birds, may also eat worms; instead, plants are not likely to eat worms. Alternatively, as λ becomes more negative, the network ceases to infer any hierarchical structure and only learns to map the new feature to the single item trained on. In this case, after learning that a goldfish eats worms, the network does not infer that other fish, birds, or plants may also eat worms.Interestingly, this setting highlights how asymmetries in the representations yielded by different λ can actually benefit transfer and generalization. This can be shown by observing that the learning of a new feature association only depends on the first layer W 1 . Let ŷf denote the vector of the representation of the new feature f across items i in the dataset. Additionally, let w (f )T 2be the new row of weights appended to W 2 which map the hidden layer to the new feature. Following Saxe et al. (2019b) (164) From this we can see that differences in the representations of the new feature across items ŷf across λ are only influenced by W 1 .In the case of the rich learning regime where λ = 0, the semantic relationship between features and items is distributed across layers. Instead, when λ > 0, the second layer W 2 exhibits lazy learning, yielding an output representation W 2 W T 2 of a weighted identity matrix. However, the first layer W 1 still learns a rich representation of the hierarchy, albeit at a smaller scaling. Furthermore, rather than distributing this learning across both layers, in the λ > 0 case, all learning of the hierarchy occurs in the first layer, allowing it to more readily transfer this structure to the learning of a new feature (which only depends on the first layer). Thus, in this case, the 'shallowing' of the network into the first layer is actually beneficial. Finally, we can also observe the opposite case when λ < 0. Here, rich learning happens in the second layer, while the first layer is lazy and learns to represent a weighted identity matrix. As such, these networks do not learn to transfer the hierarchy of different items to the new feature.Figure 7 : Transfer learning for different λ. A A new feature (such as 'eats worms') is introduced to the dataset after training on the hierarchical semantic learning task (Section 3). A randomly initialized row is added to W 2 and trained on a single item with the new feature (for example, the goldfish), with the rest of the network frozen. The network is then tested on the transfer of the new feature to other items, such that items closer to the goldfish in the hierarchy are more likely to have the same feature. B The generalization loss on the untrained items with the new feature decreases as λ increases. C As λ increases positively, networks better transfer the hierarchical structure of the data to the representation of the new feature.The details of the simulation studies are described as follows. Specifically, N i , N h , and N o represent the dimensions of the input, hidden layer, and output (target), respectively. The total number of training samples is denoted by N , and the learning rate is defined as η = 1 τ .In practice, to initialize the network with lambda-balanced weights, we use Algorithm F.1. In this algorithm, α serves as a scaling factor that controls the variance of the weights, allowing for adjustments between smaller and larger weight initializations.decomposition allows us to compute the respective RSMs as U SU ⊤ for the input and V SV ⊤ for the output task. The rows and columns in the SVD and RSMs are ordered identically to the items in the hierarchical tree.The results in Panel B display simulation outcomes, while Panel C presents theoretical input and output representation matrices at convergence for a network trained on the semantic task described in Braun et al. (2022) ; Saxe et al. (2013) ,. These matrices are generated using varying initial λbalanced weights set at λ = -2, λ = 0, and λ = 2, following the initialization scheme outlined in F.7. The network architecture includes N i = 8, N h = 8, and N o = 8 with a learning rate of η = 0.001 and a batch size of N = 8. Zero-balanced weights are initialized with a variance of σ = 0.00001, while λ-balanced networks are initialized with σ xy = √ 1 based on a random regression task with the same architecture.Panel D illustrates results from running the same task and network configuration but initialized with randomly large weights having a variance of σ = 1.In panel E, we trained a two-layer linear network with N i = N h = N o = 4 on a random regression task for λ ∈ [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5] to convergence. Subsequently, we added Gaussian noise with µ = 0, σ ∈ [0, 0.5, 1] to the inputs (top panel) or synaptic weights (bottom panel) and calculated the expected mean squared error.Panel A illustrates schematic representations of the network architectures considered: from left to right, a funnel network (N i = 4, N h = 2, N o = 2), a square network (N i = 4, N h = 4, N o = 4), and an inverted-funnel network (N i = 2, N h = 2, N o = 4).Panel B shows the Neural Tangent Kernel (NTK) distance from initialization, as defined in Fort et al. (2020) , across the three architectures shown schematically. The kernel distance is calculated as:S(t) = 1 - ⟨K 0 , K t ⟩ ∥K 0 ∥ F ∥K t ∥ F .The simulations conducted on the same task with eleven varying initial λ-balanced weights in [-9, 9] . The regression task parameters were set with (σ = √ 3). The task has batch size N = 10. The network has with a learning rate of η = 0.01. The lambda-balanced network are initialized with σxy = √ 1 of a random regression task.Panel C shows the Neural Tangent Kernel (NTK) distance from initialization for the funnel architectures shown schematically with dimensions N i = 3, N h = 2, and N o = 2. The simulations conducted on the same task with twenty one varying initial λ-balanced weights in [-9, 9] . The regression task parameters were set with (σ = √ 3). The task has batch size N = 30. The network has with a learning rate of η = 0.002. The lambda-balanced network are initialized with σxy = √ 1 of a random regression task. 
 
 Appendix:  


Decision: In light of the positive reviewers' feedback and relevancy of the submission, we are pleased to accept this paper for presentation at UniReps 2024. We kindly ask the authors to incorporate the reviewers' suggestions and feedback in the final camera-ready version of the manuscript.

Reviews: review: title: title: Exploring λ-balanced Initialization, Learning Dynamics, Generalization Impact, RSM Evolution, Empirical Validation, and Transfer Learning in Deep Linear Networks
main_review: main_review: Comment 1: In Section 2, the author explains the supervised learning setup and the use of deep linear networks for modeling. While the explanation of weight matrices and input-output mappings is clear, there should be more detail on how the λ-balanced initialization specifically affects the convergence behavior. 

Comment 2: The author introduces the QQT matrix to track key statistics of the learning process. This approach is insightful, but there should be a clearer explanation of how the choice of QQT over other methods improves the tracking of learning dynamics. 

Comment 3: In Section 3, the author explores the dynamics of singular values under different initializations. While the transition from sigmoidal to exponential dynamics is well-described, there is limited discussion on how these transitions impact generalization. 

Comment 4: The paper introduces representational similarity matrices (RSMs) to study the internal representations of weights W1 and W2. However, the theoretical explanation could be strengthened by adding empirical results showing how RSMs evolve during training. P

Comment 5: The results in Section 4 discuss continual and reversal learning, but the paper lacks empirical validation of these findings. The author should present detailed simulation results that show how continual learning dynamics differ across values of λ, particularly when catastrophic forgetting is more likely to occur. 

Comment 6: In the discussion of transfer learning, the author makes a compelling case for how task-specific structure can persist across tasks depending on the λ value. However, more quantitative results should be provided to support the claim that smaller λ values lead to better transfer performance.
title: title: Solid and perhaps overqualified
main_review: main_review: The authors study the effect of relative initialization scale and layer dimension on the rich/lazy learning dichotomy. They refine some previous results showing that small initialization scale (i.e. variance of random initial weights) can promote rich learning, to study the effect of imbalances between the scale of different layers. They also see how this interacts with imbalances in the dimensionality across layers. All of this is done with a full dynamical picture that allows them to highlight a hybrid "delayed rich" regime. 

While it is not really relevant for a future poster, I found this abstract hard to evaluate at times because some key information was missing or deep in the 30 page appendix. For example, one of the central variables, the "relative scale", is never defined. In the first Appendix, it is said to be the difference in weight magnitude between layers, but that leaves open the question of sign. If positive means that inputs have more weight, then it seems like, in figure 3, $\lambda$ is kind of compensating for the imbalanced dimensionality of layer layers (which would trivially produce higher magnitude depending on how weights are initialized) since funnel nets are rich for negative $\lambda$ and _vice versa_. Obviously that's not the case if $\lambda$ means the opposite. This is perhaps due to my fairly shallow familiarity with the related work, but it should still be defined. There are other rough spots presumably due to the trauma of copying from one overleaf to another, but those were mostly stylistic.

One question I had about the interesting "delayed rich" regime, is whether networks in this regime have already saturated in their loss curves. If so, it seems potentially related to the "grokking" phenomenon seen in non-linear networks (to the extent that rich representations are more robust, as briefly mentioned in the paper). I would have liked more exploration of this regime in general, seeing as it got the dignity of an italicised name. It was not clear, for example, which architecture is shown in Fig. 3c (which is otherwise a figure mostly about architecture).

My overall impression is that it is solid and interesting, but a lot of crucial detail was left out of its current presentation. In fact, this seems like fairly mature work that could have benefited from being in the proceedings track.
title: title: Good early stage work, but hampered by low quality/effort writing
main_review: main_review: **Summary**
This paper builds on the work of (largely) Braun et al., 2022 which derive exact learning dynamics for deep linear networks (DLNs) with unequal input-output dimensions using the Riccati matrix equation. Prior work demonstrated that a zero-balanced DLN can exhibit rich regime solutions; this work examines a broader class of $\lambda$-balanced networks. They replicate previous findings while also introducing new dynamics specific to $\lambda$-balanced networks.

**Strengths**
- The extension to $\lambda$-balanced networks extends our understanding of DLN dynamics to a much broader class of initializations
- Evaluation is thorough and covers all settings discussed in prior works, highlighting how $\lambda$-balanced initializations provide new behavior
- The lazy to delayed rich learning highlighted in Fig. 3C may be useful for understanding grokking phenomenon discussed in other works
- Proofs appear to be mostly correct with no major technical errors

**Weaknesses**
- The technical contributions are clear, but the implications are not well articulated.
- While new behavior for $\lambda$-balanced networks is mentioned, it is not discussed in detail. This is acceptable for early-stage results in the Extended Abstract track, but a future version of this work would strongly benefit from spending more time investigating the delayed rich behavior and reversal learning behavior when $\lambda \neq 0$.
- Section 3, "The dynamics of the representations" introduces a new regime, _semi-structured lazy_ but does not explain this in detail or discuss its implications.
- The quality is quite poor with numerous writing and mathematical errors throughout (most are minor and do not affect the overall conclusions):
    - Lines 92-93 appears to have fragments of a sentence
    - Fig 2E appears to have the titles of the subplots switched (it does not match the caption nor the conclusions reached in the corresponding appendix section)
    - The reversal learning section introduces a "spectrum" of behaviors controlled by $\lambda$, and even suggests it relates to animal behavior, but it points the reader to Appendix E.2 for details. Appendix E.2 abruptly cuts off mid-sentence before explaining any of these details.
    - Proof of Theorem B.2 is missing transpose operators (likely typo?)
    - Lemma C.2 missing factor of $\tau$
    - Cross-references (especially in appendix) are wrong or missing

**Questions**
1. What is the _semi-structured lazy regime_? What are the implications for initializing networks?
2. Prior work discusses how two-layer DLN results extend to multiple layers. Would that be feasible for the type of results in Fig. 3?
3. What are the animal behaviors referred to in the reversal learning section? What are the implications of this work on animal behavior?

**Limitations**
Should have a discussion about what initializations are not covered by the $\lambda$-balanced assumption.
title: title: meaningful extension for learning dynamics..
main_review: main_review: Summary: This work derived exact solutions for gradient flow dynamics of a few important statistics in deep linear networks under λ-balanced initializations extending the results from Braun & Dominé 2022. The authors examined learning dynamics different from lazy to rich regimes, showing how this transition is influenced by  λ, architecture, noise, and extended the discussion of the behaviors of the solution in different learning scenarios.

Minor points: 
- Repeated reference entries;
- Some lemma and theorem links are broken.

Major points:
- It is definitely looks like rigorously body of work, especially consider it is an abstract’s volume. I trust the authors’ math on finding the matrix Riccati equation for QQT, as it seems non trivial to introduce non-zero λ.
- The funnel v.s. squared network result is interesting, as this kind of breaks the symmetry in the squared network and shows the possibility of the delayed transition to rich. The magnitude of how far the NTK moved from initialization is interesting/surprising, compared to the square network.
- I wonder how the different types of layer imbalance, e.g. λ-balanced v.s. layer specific learning rate, relates to one another.