From Lazy to Rich: Exact Learning Dynamics in Deep Linear Networks
Authors: Clémentine Carla Juliette Dominé
Nicolas Anguita
Alexandra Maria Proca
Lukas Braun
Daniel Kunin
Pedro A. M. Mediano
Andrew M Saxe
Venue: ICLR 2025 Poster (2025)
Keywords: Deep learning
Learning theory
Learning Regime
Rich
Lazy
Primary area: learning theory
TL;DR: The paper offers explicit solutions for gradient flow in two-layer linear networks under various initializations, modeling the shift between lazy and rich learning, with applications to neuroscience and machine learning.

Abstract: Biological and artificial neural networks develop internal representations that enable them to perform complex tasks. In artificial networks, the effectiveness of these models relies on their ability to build task specific representation, a process influenced by interactions among datasets, architectures, initialization strategies, and optimization algorithms. Prior studies highlight that different initializations can place networks in either a lazy regime, where representations remain static, or a rich/feature learning regime, where representations evolve dynamically. Here, we examine how initialization influences learning dynamics in deep linear neural networks, deriving exact solutions for lambda-balanced initializations-defined by the relative scale of weights across layers. These solutions capture the evolution of representations and the Neural Tangent Kernel across the spectrum from the rich to the lazy regimes. Our findings deepen the theoretical understanding of the impact of weight initialization on learning regimes, with implications for continual learning, reversal learning, and transfer learning, relevant to both neuroscience and practical applications.

Published as a conference paper at ICLR 2025
FROM LAZY TO RICH: EXACT LEARNING DYNAMICS IN DEEP LINEAR NETWORKS
Clémentine C. Dominé'.1, Nicolas Anguita2.1, Alexandra M. Proca", Lukas Braun",
Daniel Kunin', Pedro A. M. Mediano 2:5.1, Andrew M. Saxe
Gatsby Computational Neuroscience Unit, University College London, UK
2 Department of Computing, Imperial College London, UK
3 Department of Experimental Psychology, University of Oxford, UK
Institute for Computational and Mathematical Engineering, Stanford University, USA
5 Division of Psychology and Language Sciences, University College London, UK
6 Sainsbury Wellcome Centre, University College London, UK
7 CIFAR Azrieli Global Scholar, CIFAR, Toronto, Canada
1 Co-first firsttt author
$ Co-senior author.
cl ement nne ine. domi ne. 2..0mine.e0Quul.......
na na6580cam.ac..
ABSTRACT
Biological and artificial neural networks develop internal representations that able them to perform complex tasks. artificial networks, the effectiveness these models relies their ability to build task-specific representations, pro- cess influenced by interactions among datasets, architectures, initialization strate- gies, and optimization algorithms. Prior studies highlight that different initializa- tions can place networks in either a lazy regime, where representations remain static, or rich (or reature-learning regime, where representations evolve dy- namically. Here, we examine how initialization influences learning dynamics in deep linear neural networks, deriving exact solutions for 'X-balanced" initializa- tions-de by the relative scale of weights across layers. These solutions cap- ture the evolution of representations and the Neural Tangent Kernel across the spectrum from the rich the lazy regimes. Our findings deepen the theoretical understanding of the impact of weight initialization on learning regimes, with im- plications for continual learning, reversal learning, and transfer learning, relevant to both neuroscience and practical applications.
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 repre- sentations from comparable stimuli, suggesting shared information processing mechanisms (Yamins)
iet al. 2014). This similarity, though not fully understood, has drawn interest from neuroscience, Al.
and cognitive science (Haxby al. 200T Laakso & Cottrell 2000) Morcos et al., 2018 Kornblith
et al., 2019 Moschella et al.) 2022). 2). The success of neural models relies on their ability to ex- tract relevant features from data build internal representations, a complex process that machine
learning is defined by two regimes: lazy and rich (Saxe et al. 2014 Pennington et al. 2017 Chizat
al. 2019 Bahri al. 2020).
Lazy regime. The lazy regime follows from a fundamental phenomenon in overparameterrize neu- ral networks: during training, these networks frequently remain near their linearized form, undergo- ing 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 (Jacot

---

Published as a conference paper at ICLR 2025
A
9-
9-
Analytical
C
wyW: W.
W1
Numerical
22
Hidden dimentions
5 Training Steps
c RRK.
Training Steps
Training Steps
R.
A =
A W W. W.W! AI
A= 2
x kE
Hidden dimentions
Figure I: A minimal model of the rich and lazy regimes. A. We examine a deep and wide linear network trained using gradient descent starting from an initialization characterized by a relative scale parameter À - which characterizes the difference in the weight covariance between the first and second layers. B. Network output for an example task over training time, starting from a range of relative scale values. The dynamics are influenced by the initialization. Solid lines represent simulations, while dotted lines indicate the analytical solutions derived in this work. C. A network with LeCun weight initialization (LeCun et al.) 1998) in the infinite width limit becomes X-balanced, as W W2 - W1 W approaches the scaled identity matrix.
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 dynamies that navigate saddle points (Baldi & Hornik 1989 [axe et al. 2014 2019b; Jacot et al. (021). This regime features sigmoidal learning curves and simplicity biases, such as low oow-rankness (Li et al., (020) 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). How- ever, 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) - highlighting the significance of both absolute scale (initialization variance) and relative scale (difference in weight magnitude between layers) in generating diverse learning dynamics (Atanasov et al., 2022). 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. (2023b). layer-specific initialization variances Yang Hu (2020): Luo et al. (202T): Yang et al. (2022), and the use of large learning rates Lewkowycz et al. bp; 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.
Despite significant advances, these learning regimes and their characterization are not yet fully un- derstood and would benefit from clearer theoretical predictions, particularly regarding the influence of prior knowledge (initialization) on the learning regime. In this work, we address this gap by deriving exact solutions for the learning dynamics in deep linear networks
function of network initialization, providing one of the few analytical models of the rich and lazy regimes in wide and deep neural networks (Xu & Ziyin| 2024 Kunin et al. 2024 Tu et al. 2024). To illustrate the dra- matic effect of initialization and the kind of phenomenon we build a theory for, we consider a two- layer linear network parameterized by an encoding layer W, and decoding layer Wa (Fig [A). This network can be initialized with different relative scalings, such that W, WT > W 2 W 2- W1 WT < W W2. or W, WT w W2, while maintaining the same absolute scale. As shown in BB, the choice of relative scaling can result in drastically different learning trajecto- ries and representations and the theory we develop over the course of this paper describes these effects. Through these solutions, we aim to gain insights into the rich and lazy regimes, as well as the transition between them during training, by examining the impact of relative scaling, As shown in [C and further proved in Appendix A.3] initialization methods used in practice, such as Le- Cun initialization in wide networks, approximate the relative scaling initialization explored this paper, making it relevant to machine learning community as further demonstrated by Kunin et al. (2024). We consider applications relevant to machine learning and neuroscience, including continual learning (Kirkpatrick et al. 2017 Zenke et al. 2017 Parisi et al. 2019), reversal learning (Erdeniz
2

---

Published as a conference paper at ICLR 2025
& Atalay 2010) and transfer learning (Taylor & Stone 2009 Thrun & Pratt 2012 Lampinen & Ganguli 2018 Gerace et al. 2022).
Our contributions.
We derive explicit solutions for the gradient flow, internal representational similarity, and finite-width NTK in unequal-input-outpuu two-layer deep linear networks, under broad range of X-balanced initialization conditions S
We model the full range of learning dynamics from lazy to rich, showing that this transition is influenced by complex interaction of architecture, relative scale, and absolute scale, extending beyond just initialization absolute scale (Section[5).
We present applications of these solutions relevant to both the neuroscience and machine learning communities, providing exact solutions for continual learning dynamics, reversal learning dynamics, and transfer learning (Section 6).
2 RELATED WORK
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). Previous research
has extensively analyzed convergence (Arora et al. [0188 Du & Hu 2019), generalization
properties (Lampinen & Ganguli, 2018; Poggio et al. 2018; Huh (020). and the implicit bias of
gradient descent (Arora et al., 2019al Woodworth et al. 2020 Chizat & Bach, 2020 Kunin et al.)
2022) in linear networks. These studies have also revealed that deep linear networks have intricate fixed point structures and nonlinear learning dynamics in parameter and function space, reminiscent of phenomena observed in nonlinear networks (Arora et al., 2018b Lampinen & Ganguli] [018)) Seminal work by [axe et al. (2014) laid the groundwork by providing exact solutions to gradient flow dynamics under task-aligned initializations, demonstrating that the largest singular values are learned first during training. This analysis has been extended to deep linear networks (Arora al. 2018b 2019a| Ziyin et al. [002) with more flexible initialization schemes (Gidel et al. 2019 Tarmoun et al. 2 Gissin et al. 2019). This work directly builds on the matrix Riccati formulation proposed by Fukumizu (1998) and Braun et al. (2022) which extends these solutions to wide networks. We extend and refine these results to obtain the dynamics for wider class of X-balanced networks to more clearly demonstrate the impact of initialization on rich and lazy learning regimes also developed in TTu et al. (2024) for a set of orthogonal initalizations. Our work extends previous analyses (Xu & Ziyin 2024; Kunin et al.) 2024) of these regimes 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 ayser-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 popula- tion dynamics of two-layer nonlinear neural networks initialized using a mean-field parameterization (Mei et al. 2018 Rotskoff & Handen-Fiinddes 2018)) CCizat & Bach, 2018 [Sirignano & Spiliopou- los] 2020; Rotskoff& & Vanden- Vanden-Eijnden [022]] Sirignano & Spiliopoulos/ 2020). Other methods for analyzing deep network dynamics include the NTK limit, where the network effectively performs
kernel regression without feature learning (Jacot et al. 2018 [ee et al. 2019 Arora et al., 2019b). Furthermore, these approaches typically require numerical integration
operate within limited learning regime, and are unable to describe the learning dynamics of hidden representations. In- stead, our work provides explicit analytical solutions for the dynamics of the network and its NTK in the finite-width case (Jacot et al. 2021 Xu & Ziyin 2024 Kunin et al. 2024 Chizat et al. 2019).
3 PRELIMINARIES
Consider a supervised learning task where input vectors X E R.V, from a set of P training pairs (Xx,, yn)) J need be mapped their corresponding target output vectors yn IR.,. We learn this task with a two-layer linear network model that produces the output prediction
yx W₂ WiX,.
3

---

Published as a conference paper at ICLR 2025
with weight matrices W1 R RM and W2 E RNN,,,, where N₁, is the number of hidden units. The network's s weights are optimized using full batch gradient descent with learning rate the mean squared error loss L(
(or respectively time constant = ))
Y (Il -yll?).
where (-) denotes the average over the dataset. Our objective is describe the entire dynamics of the network's s output and internal representations based on the input covariance and input-output cross-covariance matrices of the dataset, defined as
Σ X ER.
R.N.N. and ΣNE
1 P y ERN.NN,
and the initialization (0).) (0). We employ approach first introduced in the foundational
work of Fukumizu (1998) and extended in recent work by Braun et al. (2022), which instead of
studying the parameters directly, considers the dynamics of matrix of the important statisties. In particular, defining Q = [W1 W 2 matrix
E IN N.)x.N,, N.)..N we consider the (N₂ + No) x (Nr + No)
WTW
WT W W2
(1) QQ _
W, (t)
which is divided into four quadrants with interpretable meanings, and where E R represents training time. The approach monitors several key statistics collected in the matrix. The off- blocks contain the network function (t) W, (t) X. The on-diagonal blocks capture the correlation structure 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 al. (2022).
RSM, x" w/ W, (t) RSMo Y (W, W (t))+Y,
where + denotes the pseudoinverse; and the network's finite-width NTK (lacot et al. 2018 Lee
ct al. 2019 Arora et al. (019b)
NTK IN. & x WF + W2 x x
where IN. is the No x No identity matrix and @ is the Kronecker product. Hence, the dynamics of QQ" describes the important aspects of network behaviour.
4 EXACT LEARNING DYNAMICS
We derive an exact solution for QQ" offering insight into the learning dynamics, convergence behavior, and generalization properties of two-layer linear networks with prior knowledge.
Assumptions. To derive these solutions we make the following assumptions:
A1 (Whitened imput). The input data is whitened, i.e. 5 I.
A2 (X-Balanced). The network's 's weight matrices are X-balanced at the beginning of train- ing, i.e. W Wz (2 - W, W, (0)² AI. If this condition holds initialization, will persist throughout training (Saxe et al. 2014 Arora et al., 2018a). For completeness, prove this in Appendix [A]
A3 (Dimensions). The hidden dimension of the network is defined as N₁, min(N, No), ensuring the network is neither bottlenecked (N,, < min(N,, N.)) nor overparameterizee (N, min(NV, NN.)).
These assumptions are strictly weaker than prior works (Fukumizu 1998 Braun et al., 2022
Kunin al.
2024 Xu & Ziyin| 2024). The main distinction between our work and prior works is that both Fukumizu (1998) and Braun et al. (2022) assumed zero-balanced weights (W1 (O)
(2( W (O), while we relax this assumption to X-balanced. The zero- balanced condition restricts the networks to rich setting. We develop solutions to explore the continuum between the rich and the lazy regime. While some works, such as Tarmoun et af.) (2021) have considered removing this constraint, their solutions remain in an unstable and mixed form. Other studies, such as Xu & Ziyin (002) and Kunin et al. (2024), have similarly relaxed the balanced assumption but were limited to single-output neuron settings. See Appendix A.2 for further discussion on each of these works' assumptions and their relationship to ours.
Lemmas and definitions. To derive exact solutions we start by presenting the main lemmas which prove in the appendix.
4

---

Published as a conference paper at ICLR 2025
the gradient flow dynamics of QQ" QQ" (0) = Q())Q(0)) can be written as a differential matrix Riccati equation
Lemma 4.1. Under assumptions I and
(t), with initalization
d (QQ")= = FQQ" + = where F
IN,
(6)
dt
Σ
2 IN.
As derived in Fukumizu (1988 and extended in Braun et al. 2022), whenever F is symmetric and diagonalizable such that F PAP" where P an orthonormal matrix and A is diagonal matrix, then the unique solution to this matrix Riccatti equation is given by
- 2A
2A -
PTQ(o)
= e - Q(0) I+ Q
P
Q
In Appendix B.2 we prove that this equation is the unique solution to the initial value problem derived in Lemma 4. I for any value A. However, as discussed in Braun et al. (2022), the solution in this form is not very useable interpretable due to the matrix inverse mixing the blocks of QQ". Additionally, we need to diagonalize F. To do so we consider the compact singular value decomposition SVD(É** = ÜSVT. Here, Ú E RN, X.N, denotes the left singular vectors, Š e R., xN, the square matrix with ordered, non-zero eigenvalues on its diagonal, and V E N.x. the corresponding right singular vectors. For unequal input-output dimensions (N; & No). the right and left singular vectors are not square. Accordingly, for the case N; > N₁ = No. we define U. E R.N. N.,INN.NN -Nsl as a matrix containing orthogonal column vectors that complete the basis for U. i.e., make [Ü Ū1] orthonormal, and V1 E R MMM,IN,-N, as matrix of zeros. Conversely, when N; = Nv. < No, then V1 is a matrix containing orthogonal column vectors that complete the basis for V and Ü_ matrix of zeros. Using this SVD structure we can now describe the eigendecomposstion of F.
Lemma 4.2. Under assumptions the eigendecoomposition of PAPT is
S
p= VV HG) V2 Ü HG) -Ü - -HG)
VIG HG)
V2V V2Ú
A
0
(8)
0
0
A1
where the matrices Sx, Y1. H. and G are diagonal matrices defined as:
S A = sgu(NN -N) 21) A ===sen sgg(A)) S
A
S. = 1
(9)
VI+ H²
Main theorem. Thanks to the cigendecomposition of F we can separate the solution provided in equation 7 into four quadrants. Following an approach used in Braun et al.) (2022), we will find it useful define the following variables of the initialization that will allow us define the product pT (0) more succinctly,
-HG) E NAXNN - + ERNNNN, D = W2 (O)V E -.
B
(10)
(11)
(12)
Using these variables of the initialization, this brings us to our main theorem: Theorem 4.3. Under the assumptions of whitened inputs f X-balanced weights 2 and no bot- tleneck 3 the temporal dynamics of QQ" are
Z (1)Z ())
(t)
(())) (1) ()))
with the time- demedeepnnden variables (t) E N Zz(t)= E N.XN, No, and A(t) E N.NN.
(t)
1 ; viG 5.
v(ā HHGe S. 5. + + D..
(13)
(2(t) = Ú( + S
SAA BT + Ú(G-HG)
- S. CT + +
(14)
2
25 -
c [c-Sst+ -1
A -I
A(I) I+ B
BT C
CT DD
DT. (15)
45X
45
A
5

---

Published as a conference paper at ICLR 2025
A 150
B 150
C
The proof of Theorem 4.3 is in Appendix B With this solution we can calculate the ex- act temporal dynamics of the loss, network function, RSMs and NTK (Fig.ZA, C) over range of X-balanced initializations.
Loss
Loss
Loss
Colored
Numerical
Analytical
WsWi(t W(W1(0) W2W1(t)
W2W10)
W2W (6)
Implementation and simulation. One is- sue with the expression we derived in The- irem[4.3] that it can be numerically unsta- ble when simulating it for long time > 0 as involves taking the inverse of terms that involve exponentials that are diverging with f. If we make the additional assump- tion that B is invertible, then can rear- range this expression to only use exponen-
10
(TW1(0)
WTWi
WoWI(t)
(c
tials with negative coefficients, which derive in Appendix [B.], In the next sec- tion we will discuss the significance of B being invertible at initialization on the con- vergence of the dynamics. Simulation de-
NTK(t)
50-
NTK(t)
50-
NTK(e)
50
tails are in Appendix E.7
-50
550
20000
20000
20000
Training Steps
Training Steps
Training Steps
RICH AND LAZY LEARNING
x=-2
A=
x
Next, we use these solutions to gain deeper understanding of the transition be- tween the rich and lazy regimes by exam- ining the dynamics as function of X the relative scale - as it varies between positive and negative infinity. We investigate four key indicators of the learning regimes: the dynamics of singular vectors, the structure and robustness of the representations, and the evolution of the NTK.
Figure 2: A. The temporal dynamics of the numerical simulation (colored lines) of the loss, network func- tion, correlation input and output weights, and the NTK (row 1-5 respectively) the analytical solution (black dotted lines) for X = -2. B. 0 Large initial weight values. C. X initial weight values initialized as described in E.7
exactly matched by
Dynamics of the singular values. Here we examine X-balanced linear network initialized with tusk-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. This setting offers a valuable opportunity to build intuition about the impact of À on the dynamics of learning regimes, extending beyond previous solutions Tarmoun et al., 2021 Varre et al. 2024).
Theorem 5.1. Under the assumptions of Theorem 4.3 and with a task-aligned initialization, as de- fined in Saxe et al. (201)), the network function is given by the expression W2 (1)= ÜS(() VT where 5(1) E NR is diagonal matrix of singular values with elements (t) that evolve according to the equation.
So. so) =o) + Ja (t; (5a sa))
(16)
where S. is the singular value of S and Ya. (t: A) is X-dependent monotonic transition function for each singular value that increases from Yaa X) = 0 to lim,- To (t; X) = 1 defined explicitly
in Appendix C.I. We find that under different limits of X, the transition function converges pointwise
to the sigmoidal (2 + 0) and exponential (A + =00) transition functions,
28n
lim To (1; A
lim (4(4) + 11 x
(17)
25
The proof for Theorem 5.1 can be found in Appendix C.I As shown in Fig. 3 B, as À approaches zero, the dynamics resemble sigmoidal learning curves that traverse between saddle points, charac- teristic 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. 3A 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

---

Published as a conference paper at ICLR 2025
Exponential Dynamics
Sigmoidal Dynamics
C
Exponential Dynamics
Singular Valies
15
15
0
0
4000
4000
4000
Training Steps
Training Steps
Training Steps
-=
==
x=
Numerical I...
Analytical
Figure 3: Simulated and analytical dynamics of the singular values of the network function with relative scale of A. X = 2, B. A 0, or C. À 2, initialized as described in Appendix E.7
dynamics resemble that of shallow network. Notably, similar effects have been observed previ- ously in the context of large absolute scales (Saxe et al. 2019a) independently of the relative scale. Overall, our results highlight the critical influence the relative scale À has in shaping the learning dynamics, from sigmoidal to exponential, steering the network between the rich and lazy regimes.
The dynamics of the representations. We now consider how the representations of the individual parameters W and W2 change through training. We note that under X-balanced initializations there is a simple structure which persists throughout training that allows to recover the dynamics of the parameters up to a time-dependent orthogonal transformation from the dynamics of QQ (I).
Theorem 5.2. Under assumptions - if the network function W2W: (f) (tt) (t)
is full rank, then we can recover the parameters Wz(t) = S(t) (t) RT(() and Wi(t)
( (1) up to time-dependent orthogonal transformation R(1) E NNN,, where
(t) -/ 5²() + Xº1 Ind and
S, (t) = ((sx()) - 21 0m.x((...-NN.) (2(t) - ((Sx(t) ((sx))((s(())))))))))))))))))))) + 21)
:m..) (18)
The effective singular values Sx of the corresponding weights are either up-weighted or down- weighted depending on the magnitude and sign of X, splitting the representation into two parts. This division is reflected in the network's internal representations. With our solution, QQ"(t), which captures the temporal dynamics of the similarity between hidden layer activations, can analyze the network's internal representations in relation to the task. This allows us to determine whether the network adopts rich lazy representation, depending on the value of A. Assuming convergence to the global minimum, which is guaranteed when the matrix B is non-singular, the internal representation satisfies WFW, = V and W,W3 US3Ū" with W, ÜSVT. Theorem C.3 in the Appendix provides a detailed proof of this limiting behavior. To illustrate this, we consider a hierarchical semantic learning task]) introduced in Saxe et al. 2014):; [Braun et al.] (2022), where living organisms are organized according to their features (Fig. A). The representational similarity of the task's inputs (VSVT) reflects this hierarchical structure (Fig Similarly, the representational similarity of the task's target values (ÜSÜT) highlights the primary groupings of items. When training a two-layer network with relative scale A equal to zero and task-agnostic initialization (Mishkin & Matas 2015), the input and output representational similarity matrices (Fig(4B) match the task's structure upon convergence. As derived in Theorem C.4 the network is guaranteed to find rich solution regardless of the absolute scale, meaning W W, VŠVT and W, WT ÜŠÜT shown in Fig. 4c.. Hence, the network learns task-specific representations. We also show that as X approaches either positive or negative infinity, the network symmetrically transitions into the lazy regime. As demonstrated in Theorem C.4 and illustrated in Fig. 4 the representations converge to an identity matrix for both large positive and large negative values of X- emerging the output representations for large positive À and input representations for large negative X. This convergence indicates that the network adopts task-agnostic representations. Meanwhile, the other respective RSMs become negligible, with scales proportional to 1/X. Therefore, as shown in Theorem C.5 the NTK becomes static and equivalent to the identity matrix in the limit as À approaches infinity. However, the downscaled representations of the network remain structured and task-spec Intuitively, in this
In this setting, the network has equal input and output dimensions
7

---

Published as a conference paper at ICLR 2025
A
Task specific
B
Lazy
Rich
Lazy
Lazy
E
Input Noise
Simulations
==
A=
Gaussian Large
W/W:
W,
WEW:
W.W]
WFWL
WoW
WFW,
W.W:
Anma
Poet
Haly
ded Alpas
Rowwrr
Parameter Noise
WTW;
W,wy
W,Wg
WFW,
W, W
Theory
Scale
Relative
scale
À
0
A X
Figure 4: A. A semantic learning task with the SVD of the input-output correlation matrix of the task. (top) U and V represent the singular vectors, and S contains the singular values. (bottom) The respective RSMs as USUT for the input and VSVT for the output task. B. Simulation results and C. Theoretical input and output representation matrices after training. showing convergence when initialized with values of X equal to 22, O. and 2, according to the initialization scheme described in Appendix E.7 Final RSMs matrices after training converged when initialised from random large weights. E. After convergence, the network's sensitivity to input noise (top panel) is invariant to X, but the sensitivity to parameter noise increases as becomes smaller (or larger) than zero.
setup, the larger weights function as an identity-like projection, while the smaller weights adapt and align to the task. However, due to their relative scale compared to the larger weights, their contribution to the NTK remains negligible. This property could be beneficial if the weights are later rescaled, such as during fine-tuning, potentially enhancing generalization and transfer learning, as we will demonstrate in Section 6 We compare this to the scenario where both weights are initialized with large Gaussian values, leading to lazy learning that maintains fixed NTK but lacks any structural representation, as illustrated in Fig. 4D. We further discuss the relationship between the scale and the relative scale in appendix A..4 Furthermore, 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 learn 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. Consequently, we propose a new lazy regime, which we refer to as the semi-structured lazy regime. We note that these existing regimes preserve only the input or output representation, resulting in partial loss of structural information. 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. (0224). In that case large relative weight (large positive X) leads to a rapid rich regime. All together, we find that initialization will determine which layer in the network the task specification features resides in: layers initialized with large values will be task-agnostic, while those initialized with small values will be task-specific.
Representation robustness and sensitivity to noise. Here we examine the relationship be- tween the learning regime and the robustness of the learned representations to added noise in the inputs and parameters. The expected post-convergence loss with added noise to the inputs is determined by the norm of the network function (Braun et al., 2025), which in our setting is independent of X. Specifically, if we add zero-centered noise {x with variance xx xx D.
to the inputs, then the expected loss is(C) depends solely on the statistics of the training data (Figure 4E, Appendix C.3) However, if instead noise is added to the parameters. the expected loss scales quadratically with the norm of the weight matrices (Braun al.) (2025), which in our setting depend on x. In particu- lar, zero-centered parameter noise Ew, and Éw, with variance on results in an expected loss of(C) (W) = 1. W 2l + NN.wW 1 + NN..N + e, with norms
No S + c, where e is constant that
(V433 4S + This ||Wi = 1 NN (V 15? + and W = 2 M. No
implies that, under the assumption of equal input-output dimensions, networks initialized with weights such that À= = 0, corresponding to the rich regime, converge to solutions that are most robust to parameter noise (Figure 4 Appendix C.3 In practice, parameter noise could be interpreted as
8

---

Published as a conference paper at ICLR 2025
0.7
Kernel Distance
Kernel Distance
Sqaare rebwort usSAeriieheers
Anti-funnel network
Funnel network
Square network
9
Value of
Value of
Figure 5: A. Schematic representations of the network architectures considered, from left to right:
funnel network, square network, and inverted-funnel network. B. The plot shows the NTK kernel distance from initialization, as defined in Fort et al.] (2020) across the three architecture depicted schematically. C. The NTK kernel distance away from initialization over training time.
the noise occurring within the neurons of a biological network. Hence, a rich solution may enable 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 of rich and lazy learning. Here, we explore how network's s learning regime is also shaped by the interaction of its architecture and the sign of the relative scale. We consider three types of network architectures, depicted in Fig. 5AA funnel networks, which narrow from input to output (N; > Nı = No): inverted-finnnee networks, which expand from input to output (N, N₁, No): and square networks, where input and output dimensions are equal (N; = N₁ = N.). Our solution, QQ"", captures the dynamics of the NTK across these different network architectures. To examine the NT'''s evolution under varying À initializations, we compute the kernel distance from initialization, as defined in Fort et al. (2020). As shown in Fig. 5B, we observe that funnel networks enter the lazy regime as À + 0o, while inverted- so X + -00 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 X, these networks transition from 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. SC. This behavior is also quantified in Theorem C.6 which describes the rate of learning in this network. In Appendix A.4 we further explore the impact of the architecture as a function of the absolute and relative scale. Intuitively, the delayed onset of the rich regime occurs because no least-squares solution exists within the span of the network at initialization. In such cases, the network enters a delayed rich phase, where À tends toward infinity, with the magnitude of determining the length of the delay. At first, the network exhibits lazy dynamics, striving to approximate the solution. However, constraints necessitate adjustments in its directions, the network gradually transitions into the rich phase. For square networks this behavior is discussed in Section 5 Across all architectures, as À + 0, the networks consistently transition into the rich regime. Altogether, we further characterize the delayed rich regime in wide networks.
networks do
6 APPLICATIONS
In this section, apply the exact solutions for the learning dynamics in deep linear networks de- scribed in Section 4 to illustrate several phenomena relevant to machine learning and neuroscience.
Continual learning. Continual learning, as thoroughly reviewed by Parisi et al. (019), has long posed significant challenge for neural network models in contrast biological networks, particularly due to the issue of catastrophic forgetting (McCloskey & Cohen 1989 Ratclifi) 1990 French 1999). Similarly to the framework presented by Braun et al. (202), our approach describes the exact solutions of the networks dynamics trained across a sequence of tasks describing the entire continual learning process. As detailed in Appendix [..] we demonstrate that, regardless of the chosen value of X, training subsequent tasks can result in the overwriting of previously acquired knowledge, leading to catastrophic forgetting.
Reversal learning. During reversal learning, previously acquired knowledge must be relearned, necessitating the overcoming of an earlier established relationship between inputs and outputs. As demonstrated in Braun et al. (2022), reversal learning theoretically does not succeed in deep linear networks the initalization aligns with the separatrix of saddle point. While simulations show

---

Published as a conference paper at ICLR 2025
that the learning dynamics 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, they avoid passing through the saddle point due to the initialization scheme. This is both theoretically proven and numerically illustrated in Appendix D.2 We also present a spectrum of reversal learning behaviors controlled by the relative scale X, 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 toneuroscience experiments.
Transfer learning. We consider how different À initializations influence generalization to new feature after being trained on initial task. As detailed in Appendix D.3 we first train each network on the hierarchical semantic learning task described Fig. + We then add new feature to the dataset, and train the network specifically on the corresponding item while keeping the rest of the network parameters unchanged. Afterwards, we evaluate the generalization to the other items. We observe in Appendix Fig. D.3 that the hierarchical structure of the data is effectively transferred to the new feature when the representation is task-specific and À is zero Conversely, when the input feature representation is lazy (A S 0). meaning the hidden representation lacks adaptation, no hierarchical generalization is observed. Strikingly, when X is positive, the hierarchical structure in the input weights remains small but structured, while the output weights exhibit a lazy representation and the network generalizes hierarchically. Specifically Fig. D.3 shows that the generalization loss on untrained items with the new feature decreases as a function of increasing A. Therefore, as 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 indicates that the lazy regime structure (large À values) can be beneficial for transfer learning.
Fine-tuning It is common practice to pre- train neural networks on large auxiliary task before fine-tuning them on a downstream task with limited samples. Despite the widespread use of this ap- proach, the dynamics and outcomes of this method remain poorly understood. Our study establishes a theoretical basis for the success of fine-tuning, particularly how changes in X-balancedness initial- isation after pre-training affect performance on new datasets (see Appendix D.4). We consistently find that finetuning performance improves and converges more quickly as networks are re- re-balanced to larger values of XFT and, conversely, decreases XFT approaches 0 as shown in Fig. D.4 Our work examines the fine-tuning dynamics of two-layer linear networks. While simple, these architec- tures are commonly utilized for fine-tuning large pre-trained language and vision models, notably in Low-Rank Adapters (LoRA) (Hu et al., 2022) as further discussed in Appendix D.4] While detailed exploration of fine-tuning performance in practice as a function of initialization lies beyond the scope of this work, it remains an important direction for future research.
DISCUSSION
We derive exact solutions the learning dynamics of deep linear networks. While our findings ex- tend the range of analytically describable two-layer linear network problems, they are still limited by of assumptions. In particular, relaxing the assumptions that input covariance must be white and that initialization must be X-balanced could bring the analysis closer to practical applications in ma- chine learning and neuroscience. Moving towards the nonlinear setting would also make the findings more applicable in real-world scenarios. Despite these limitations, our solutions provide valuable insights into network behavior. We examine the transition between the rich and lazy regimes by analyzing the dynamics as function of X the relative scale across its full range from positive to negative infinity. Our analysis demonstrates that the relative scale, X, is pivotal in managing this transition. Notably, we identify a structured lazy regime that promotes transfer learning. Building on previous work (Kunin et al. 2024) that shows these findings extend to basic nonlinear settings and practical scenarios, theory suggests that further exploration of unbalanced initialization could optimize efficient feature learning. We leave for future work, the extension of this initialization to deep networks. Future work will focus on extending the application of the solution to the dynamics of fine-tuning mun-ttning and linear autoencoders.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This research was funded in whole, or in part, by the Wellcome Trust [216386/7/19/72] For the purpose of Open Access, the author has applied a CC BY public copyright licence to any Author Accepted Manuscript version arising from this submission. C.D. and A.S. were supported by the Gatsby Charitable Foundation (GAT3755). Further, A.S. was supported by the Sainsbury Wellcome Centre Core Grant (219627/2/19/2) and A.S. is CIFAR Azrieli Global Scholar the Learning in Machines & Brains program. A.M.P. was supported by the Imperial College London President's PhD Scholarship. L.B. was supported by the Woodward Scholarship awarded by Wadham College, Oxford. and the Medical Research Council [MR/N013468//] I]. D.K. thanks the Open Philanthropy AI Fellowship for support. This research was funded whole, or in part, by the Wellcome Trust [216386/2/19/2].
REFERENCES
Zeyuan Allen-Zhi Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparame- terized neural networks, going beyond two layers. Advances neural information processing systems, 32, 2019a.
Zeyuan Allen-Zhu, Yuanzhi Li. and Zhao Song. A convergence theory for deep learning via over- parameterization. International conference on machine learning, pp. 242-252. PMLR, 2019b.
Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. arXiv preprint
201 8a.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In International Conference on Machine Learning, pp. 244-253. PMLR, 2018b.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019a.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances in Neural Information Processing Systems, 32, 2019b.
Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. In International Conference Learning Representations, 2022.
Shahar Azulay, Edward Moroshko, Mor Shpigel Nacson, Blake E Woodworth, Nathan Srebro, Amir Globerson, and Daniel Soudry. On the implicit bias of initialization shape: Beyond infinitesimal mirror descent. In International Conference on Machine Learning, pp. 468-477. PMLR, 2021.
Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37944 2022.
Yasaman Bahri, Jonathan Kadmon, Jeffrey Pennington, Sam Schoenholz, Jascha Sohl-Dickstein, and Surya Ganguli. Statistical mechanics of deep learning. Annual Review of Condensed Matter Physics, 1 11:501-522 -528, 2020.
Pierre Baldi and Kurt Homík, Neural networks and principal component analysis: Learning from examples without local minima. Neural networks, 2(1):53-58, 1989.
Lukas Braun, Clémentine Dominé, James Fitzgerald, and Andrew Saxe. Exact learning dynamics of deep linear networks with prior knowledge. Advances Neural Information Processing Systems, 35:6615-6629 12 2022. URL
nips cc/paper Files/paper!
html
Lukas Braun, Erin Grant, and Andrew M Saxe. Not all solutions created equal: An analytical dissociation of functional and representational similarity in deep linear neural networks. Under review, 2025.
11

---

Published as a conference paper at ICLR 2025
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing sys- tems, 31, 2018.
Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on learning theory, pp. 1305-1338, PMLR, 2020.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Hugo Cui, Luca Pesce, Yatin Dandi, Florent Krzakala, Yue M Lu, Lenka Zdeborová, and Bruno Loureiro. Asymptotics of feature learning in two-layer networks after one gradient-step. arXiv preprint arXiv:2402.n 04980, 2024.
Sander Dalm, Joshua Offergeld, Nasir Ahmad, and Marcel van Gerven. Efficient deep learning with decorrelated backpropagation. arXiv preprint arXiv:2405.02338 2024.
Simon Du and Wei Hu. Width provably matters in optimization for deep linear neural networks. In International Conference on Machine Learning, pp. 1665-1664, PMLR, 2019.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning, pp. 1675- 1685. PMLR, 2019.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes owerpparameterized neural networks. arXiv preprint arXiv:1810.02054, 2018.
Burak Erdeniz and Nart Bedin Atalay. Simulating probability learning and probabilistic reversal learning using the attention-gatet reinforcement learning (agrel) model. In The 2010 International Joint Conference on Neural Networks (IJCNN), pp. 1-6. IEEE, 2010.
Matthew Farrell, Stefano Recanatesi, and Eric Shea-Brown. From lazy to rich to exclusive task rep- resentations in neural networks and neural codes. Current Opinion in Neurobiology, 83:102780- 102780, 12 2023. doi: 10. 10.1016/j.conb.200 202b. 102780.
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study loss landscape geometry and the time evolution of the neural tangent kernel. Advances in Neural Information Processing Systems, 33:5850-5861, 2020.
Robert French. Catastrophic forgetting in connectionist networks. Trends in cognitive sciences, 344):128- 1999.
Kenji Fukumizu. Effect of batch learning in multilayer neural networks. IEEE Transactions on Neural Networks, 11:17-26, 1998. doi: 10. 109/77.822506..
Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, 13301, 11 2020. doi: 00.1088/142255468/bbbeee URL httttss://axiiv..rggg... 1906. 08034
Federica Gerace, Luca Saglietti, Stefano Sarao Mannelli, Andrew Saxe, and Lenka Zdeborová. Probing transfer learning with model of synthetic correlated datasets. Machine Learning: Sci- ence and Technology, 2022.
Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of discrete gradient dynamics in linear neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Daniel Gissin, Shai Shalev-Shward and Amit Daniely. The implicit bias of depth: How incremental learning drives generalization. arXiv preprint arviv: /909, ptt.://09..........
/2051, 2019.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. pp. 249-256. proccedings.mlr..ree. JMLR Workshop and Conference Proceedings, 03 2010. URL htttt:::tppppeeeeeeeeeeee
pttp./pp/rs///11111 orot 0a html
12

---

Published as a conference paper at ICLR 2025
James V Haxby, M Ida Gobbini, Maura L Furey, Alumit Ishai, Jennifer L Schouten, and Pietro Pietrini. Distributed and overlapping representations of faces and objects in ventral temporal cortex. Science, 293(5539):2425 22225-239)2222-2330, 2001.
Soufiane Hayou, Nikhil Ghosh, and Bin Yu. The impact of initialization on lora finetuning dynamics. arXiv preprint 011::2440000047777.,. 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Sur- passing human-level performance on imagenet classification. 2015 IEEE International Con- ference on Computer Vision (ICCV), 12 2015. doi: 10. 10.11099icc.22115.22.. URL https:
//1eeexplore. ieee. org httg://douumtt/7/110000
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Con- ference on Learning Representations, 2022. URL
net.eewww.ee/eeeeeeeeeeeeeeeeeee forum?
id=nZev/eeeFff9
Dongsung Huh. Curvature-corrected learning dynamics in deep neural networks. In International Conference on Machine Learning, pp. 4552 3552-456. PMLR, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- cralization in neural networks. Advances in neural information processing systems, 31, 2018.
Arthur Jacot, François Ged, Berfin Şimşek, Clément Hongler, and Franck Gabriel. Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. arXiv preprint arXiv:2106. 1oXY::2106.19333,, 2021.
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska- Barrskikka, al. Overeom- ing catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114(13):351-33566 2017.
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In International conference on machine learning, pp. 3519- 3529. PMLR, 2019.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. Imagenet classification with deep con- volutional neural networks. Communications of the ACM, 60:84-90, 05 2012. doi: 10. 145/ 3065386. URL http:///////////////ce//o/oooors
fi files//aper/2012/
Daniel Kunin, Atsushi Yamamura. Chao Ma, and Surya Ganguli. The asymmetric maximum margin
bias of quasi- houggehhoooeseouus neural networks. arXiv preprint arXiv:2210.03820. 2022.
Daniel Kunin, Allan Raventós, Clémentine Dominé, Feng Chen, David Klindt, Andrew Saxe, and Surya Ganguli. Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning, 06 2024. URL /bttpp:://axxxv wttpxxxiiv.rg/aa////aaaa/
06158
Aarr Laakso and Garrison Cottrell. Content and cluster analysis: assessing representational simi- larity neural systems. Philosophical psychology, 13(1):47-76, 2000.
Andrew K Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. arXiv preprint arXiv: 10374, 2018.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-25 2324,, 1998. URL https: //hal www.scciineeeccc.e 0eeeaall032282////00
Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
13

---

Published as a conference paper at ICLR 2025
Sebastian Lee, Stefano Sarao Mannelli, Claudia Clopath, Sebastian Goldt, and Andrew Saxe. Maslow's 's hammer for catastrophic forgetting: Node re-use vs node activation. arXiv preprint TrXii:2205509002, 2022.
Aitor Lewkowyzz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.022 02218, 2020.
Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning, arXiv preprint arXiv:2012.099 0.0899, 2020.
Yiheng Liu, Tianle Han, Siyuan Ma, Jiayue Zhang, Yuanyuan Yang, Jiaming Tian, Hongping Hc, Antong Li, Mengshen He, Zhengliang Liu, Zihao Wu, Dajiang Zhu, Xiang Li, Qiang Niu, Dingang Shen, Tianming Liu, and Bao Ge. Summary of chatgpt-related research and per- spective towards the future of large language models. Meta Meta-Radiolog) 1, 04 2023a. doi: (tt.:jmmetrad2223.11000
Yuhan Helena Liu, Aristide Baratin, Jonathan Cornford, Stefan Mihalas, Eric Shea-Brown and Guillaume Lajoie. How connectivity structure shapes rich and lazy learning in neural circuits. ArXiv, 2023b.
Tao Luo, Zhi-Qin John Xu, Zheng Ma, and Yaoyu Zhang, Phase diagram for two-layer relu neural networks at infthite-width limit. Journal of Machine Learning Research, 22((1):1-47, 2021.
Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora. A kernel-based view of language model fine-tuning. In International Conference on Machine Learning, pp- 23610-23641. PMLR, 2023.
Sibylle Marcotte, Remi Gribonval, and Gabriel Peyré. Abide by the law and follow the flow: conservation laws for gradient flows, 12 2023. URL
c 7bbeeh/h66e211166522222266666660
Michael McCloskey and Neal J Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of learning and motivation, volume 24, pp. 109-165. Elsevier, 1989.
Song Mci, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape two- layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7665-7777), 2018.
Dmytro Mishkin and Jiri Matas. All you need is a good init. arXiv preprint 1rXiv:1511.06422 2015.
Ari Morcos, Maithra Raghu, and Samy Bengio. Insights representational similarity in neural networks with canonical correlation. Advances in neural information processing systems, 31, 2018.
Luca Moschella, Valentino Maiorca, Marco Fumero, Antonio Norelli, Francesco Locatello, and Emanuele Rodolà. Relative representations enable zero-shot latent space communication. arXiv preprint arXiv:2209. 15430, 2022.
Genman Parisi, Ronald Kemker, Jose Part, Christopher Kanan, and Stefan Wermter. Continual lifelong learning with neural networks: A review. Neural Networks, 113:54-71, 2019.
Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. Resurrecting the sigmoid in deep learn- ing through dynamical isometry: theory and practice. Advances neural information processing systems, 30, 2017.
Tomaso Poggio. Qianli Liao, Brando Miranda, Andrzej Banburski. Xavier Boix, and Jack Hidary. Theory iibb: Generalization in deep networks. arXiv preprint a8Xiv: /379, 2018.
Shahryar Rahnamayan and Gary Wang. Toward effeectve initialization for scage search spaces. 01 2009. URL hhtpps: /wwww. s ff /ttp...c///////////////000000000
14

---

Published as a conference paper at ICLR 2025
Roger Ratcliff. Connectionist models of recognition memory: constraints imposed by learning and forgetting functions. Psychological review, 97(2):285, 1990.
Grant Rotskoff and Eric Vanden- Dijnen-Eijjden... Parameters interacting particles: long time convergence and asymptotic error scaling of neural networks. Advances in neural information processing systems, 31, 2018.
Grant Rotskoff and Eric Vanden-Eijndee Trainability and accuracy artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics, 75 (9):1889-1935, 2022.
Andrew Saxe, James L.. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dy- namics of learning in deep linear neural networks. openreview.net, 12 2013. URL https: //openrevi net. forum? perrmmid==eee _tmm:/dd==wZ/wDTFFFFFFFF 9C
Andrew M. Saxe, James I.. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dy- namics of learning in deep linear neural networks. In 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceed- ings, 2014.
Andrew M. Saxe, James L. McClelland, and Surya Ganguli. A mathematical theory of semantic development deep neural networks. Proceedings of the National Academy of Sciences, 116 (23):11537-11566, 2019a. doi: 10. 10.1073/pnas. 1r53/pnas.12022661 1820226116. httpst/:: betttss:/www.ppaa /ww///ww.... pnas org/
[ontent/116/23/1155
Andrew M Saxe, James L McClelland, and Surya Ganguli. A mathematical theory of semantic development deep neural networks. Proceedings of the National Academy of Sciences, 116 (23):11537-11546, 2019b.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM Journal on Applied Mathematics, 80(2):725-752, 2020.
Salma Tarmoun, Guilherme Franca, Benjamin Haeffele, and Rene Vidal. Understanding the dynamics of gradient flow in overparameterized linear models. In International Conference Machine Learning, pp. 10153-10161. PMLR, 2021.
Matthew E Taylor and Peter Stone. Transfer learning for reinforcement learning domains: A survey. Journal of Machine Learning Research, 10(7), 2009.
Sebastian Thrun and Lorien Pratt. Learning to learn. Springer Science & Business Media, 2012.
Zhenfeng Tu, Santiago Aranguri, and Arthur Jacot. Mixed dynamics linear networks: Unifying the lazy and active regimes. arXiv preprint MrXiv:2405.17580 2024.
Aditya Vardhan Varre, Maria-Luiza Vladarean, Loucas Pillaud-Vivien, and Nicolas Flammarion. On the spectral bias of two-layer linear networks. Advances in Neural Information Processing Systems, 36, 2024.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference Learning Theory, pp. 3635-3673. PMLR, 2020.
Xiangxiang Xu and Lizhong Zheng. Neural feature learning in function space Journal of Ma- chine Learning Research, 25:1-76, 2024. //attp imlr. orgg httg/pppers///e///o catsssvolummee volume
23-1202/23-1222 pdf
Yizhou Xu and Liu Ziyin. When does feature learning happen? perspective from analytically solvable model. arXiv preprint arXiv:2401.07085, 2024.
Daniel LK Yamins, Ha Hong, Charles F Cadieu, Ethan A Solomon, Darren Seibert, and James J DiCarlo. iiftommance-otiiiiae hierarchical models predict neural responses in higher visual cortex. Proceedings of the national academy sciences, 111((3):8619-88) 2014.
15

---

Published as a conference paper at ICLR 2025
Greg Yang and Edward J Hu. Feature learning in infinite-wide neural networks. arXiv preprint arXiv:2011.1452,, 2020.
Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ry- der, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466, 2022.
Friedemann Zenke, Ben Poole, and Surya Ganguli. Continual learning through synaptic intelligence. In International Conference on Machine Learning, pp. 3987-3995, PMLR, 2017.
Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Catapults in sgd: spikes in the training loss and their impact generalization through feature learning. arXiv preprint arXiv:2306.048111 2023.
Liu Ziyin, Botao Li, and Xiangming Meng. Exact solutions of a deep linear network. Advances in Neural Information Processing Systems, 35:24446-24458, 2022.
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep relu networks. Machine learning, 109:467-492, 2020.
16

---

Published as a conference paper at ICLR 2025
A PRELIMINARIES
A.1 APPENDIX: BALANCED CONDITION
Definition A.1 (Definition of X-balanced property (Saxe et al. (2013). Marcotte al. (2023)). The
weights W1, W2 are X-balanced if and only if there exists Balanced Coefficient À E R such that:
BW W2) W22 W2 W, W XI
(19)
where B is called the Balanced Computation.
For À = we have Kero-Balanced given as A4 O. (1) = Wa(0)"W2(0),
Theorem A.2. Balanced Condition Persists Through Training Suppose at initialization
W2(0)*2((0) (0)))))) = II
(20)
Then for all t
= II
(21)
Proof. Consider:
d -- = "att" Wa(t) Wa(t) + Wz(t)
d
(-d Wa(t)
Wa(t)Wa(t)
(T
dt
(1d (T dt
w2(t)) Wi(t)* Wi(1) T dt
Wi(t)
T
Wi(1) ($S"
W2 ( (1)
W2(1)
+ W2((1T
STT (())))
Wi(t)
W2(1)² EN' (1(
Wi(t)
Wi(()
E (()))
W₂(1)
Note that (2(+) W2 Wi (t) is conserved for any initial value A.
A.2 DISCUSSION ASSUMPTIONS
Whittened Inputs. Although the whitened input assumption is quite strong, it is commonly used in analytical work obtain exact solutions, and much of the existing literature relies on these so- lutions Fukumizu 19988;; Braun et al. (2022); Kunin et al. (2024). While relaxing this assump- tion prevents the exact description of network dynamics, Kunin et al.) (2024) examine the implicit bias of the training trajectory without relying on whitened inputs. If the interpolating manifold is one-dimensional, the solution can be solved exactly terms of X. Their findings demonstrate similar quantitative dependence on X, 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 more optimal learning condition. This approaches emphasize the potential advantages of input whitening for downstream tasks, reinforcing the validity of our assumption.
Dimension. Previous works imposed specific dimensionality constraints. For example: Fuku-
mizu (1998) assumed equal input and output dimensions (N; N.) while allowing a bottle- neck in the hidden dimension (N₂, < N; -N).. Braun et al. (2022) extended these solutions to cases with unequal input and output dimensions (N; * N.) but restricted bottleneck networks (N,, = min(N, .N)) and introduced an additional invertibility condition the B. In our work allow for unequal input and output N; * N, and do not introduce an additional invertibility assump- tion. This flexibility expands the applicability of our framework wider range of architectures.
17

---

Published as a conference paper at ICLR 2025
Full rank Previous work by Braun al. 20022), imposed full-rank initialization condition, de- fined as (an)((2(00)Wi = N₂ No. However, this assumption is not necessary our framework.
Balancedness Assumption significant departure from prior works the relaxation of the bal- ancedness assumption: Earlier studies, such as Fukumizu (1998) and Braun al. (2022) assumed strict saro-balancedness (W](0)W;
= (2 w O)), which constrained the networks to
the rich regime. Our approach generalizes this to X-balancedness, enabling exploration of the tinuum between the rich and lazy regimes. While some efforts, such as Tarmoun et al. (2021), have explored removing the zero-balanced constraint, their solutions mixed forms. In contrast, our methodology systematically studies different learning regimes by varying initialization properties, particularly through the relative scale parameter. This allows con- trolled transitions between regimes, advancing understanding of neural network behavior across the spectrum. Other studies, such as Kunin et al.] (2024) and Xu & Zheng (2024) have also re- laxed 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, in the averaging limit, standard network initializations (e.g., LeCun initialization [LeCun et al. (1998), He initialization) lead to X-balanced behavior in the infinitewwidth limit. Specifically, the term Wi (O)) studies have demonstrated that the relative scaling significantly impacts the learning regime in practical scenarios, highlighting the crucial role of dynamical studies of networks function of this parameter.
limited unstable or
W2( W22) converges
scaled identity matrix. Furthermore, previous
A.3 RANDOM WEIGHT INITIALISATIONS AND A- A-BALANCED PROPERTY
Throughout this work, we assume that initial weights are X-Balanced. However, practice, weights are not initialized with that goal in mind. Usually, weight matrix W is initialized with some random distribution centered around O, with variance inversely proportional to the number of layers on which W has direct effect (Glorot & Bengio (2010). LeCun et al. (1998). He al. (2015)) In this section, we show that many common initialization techniques lead to X-Balanced weights in expectation. Furthermore, as the size of network tends infinity, these random weights A- --Bllanced probability.
We do this by first finding the expectation and variance of the balance computation for two adjacent weight matrices, W and W., initialized under normal distribution with zero mean. Subsequently, we describe how network structure and size can impact the expectation and variance of the balance computation.
Theorem A.3. [Random Weight Initialization Leads to Balanced Condition] Consider a fully connected neural network with L layers. Each layer has N; neurons, and the weights of each layer W; is a matrix of dimension (N;, Ni+1). The matrix W; = (ww,.) where wx,,. ~ - NN0,, of), where of is determined based on the initialization technique. Then the following hold for all e [,, 1):
1. IE [WT, Wi+1 W.W = (Ni420 Nooţ)I
2. Var [WT ++1 Wi+1 W. W = (N+20 + N₂O²) where is a square matrix with fours across the diagonal and ones everywhere else.
Note that in the case = 3, Ng = i, N₁ = h N₂ with i, h, being the input, hidden and output dimensions respectively defined in the main text.
Proof of Theorem A.3
W1,1
W1,2
WI.N442
W2,1
W2.2
W2 W242
Let Wi+1
(22)
WN++11
WN..1.2
WN No
(W)
W2
W N.+42
with wy (Wi1j W2.j1-.- - WN. 3T
18

---

Published as a conference paper at ICLR 2025
Then,
W
2
W
i41
W+11=
W
W2
WN; 2)
N.+2
(w) Wi
(wi, W2) (w2, W2)
(w) WNiia
(W2, wi
(w2, W
N442 wi
No4a W2)
(w Na W N
Now, consider (w,, my)) with + j,
No42 W wz)
Wke Wki
k=1
2
(,., wyyl = (Kk.WWkk
No+a
Elwr, wayl
k=1
N442
E
JE [[wks] = 0 (by independence)
k=1
Var [T, myy)
E
W
E (T |!
N140 Σ WKk.
E
N:
N:
E
w Wk
+
Wk, WW,
k=
k= 1>
Nota
Neta
Elw k
+
Wkke I. Wwkk TF w1 E[wr.
k=1
No22
Σ
= (1142)
Similarly, consider ww),
19

---

Published as a conference paper at ICLR 2025
Nia2
1 W.
K,,
No4z
E [ws, wi:)
IE
Wkii
2442 E [w N4120N,11
Var Kww,, wx:) = IE
(!Ws, wx))²
EE
[wi ws))
Ni+2
E
w
2 No11
Ni+2
E
2
22 Ne11
k=1
Nit2
Ni42 No+2
Wki + 2
W k Wi,s
- 22 N.+1
k=1
k=
1=k+1
N14a
Ni42
Ni
Σ
E
+ 2
E
E
[w].]]
No 2
1=k+1
Ne42(30
+ (N
N442)o Ma11 - N2
N:
142
=4N420
Hence
E B W4+1] (N4+20441)
Var W W++1] (Ni42) of
For the case for W. notice we can express W.WT (WT)) Hence we can use the proof above, with W WTr. In this case the matrix W has shape (N, N1+1), and each element of the matrix has variance of We have:
E [W. N,o2
Var W.W N.O
By assumption, W,, W4+1 are independent. Hence Coe W. Wi+1) = 0. We can use this prop- erty together with linearity of expectation:
E
W T 4+1 W - W W. (N1420711 Nio,)) I
Var [ 1+
W i+1 W w.] (N4420 + Noo,)) B
This completes the proof.
In neural network training, proper weight initialization is crucial for ensuring stable gradients during backpropagation, which helps to avoid issues such as vanishing and exploding gradients. The goal of weight scaling is to maintain appropriate variance across layers, enabling efficient and effective learning (Goort & Bengio (2010). The weights are typically initialized to be random and centered around 0 to break symmetry and ensure that different neurons learn different features.
Some of the most commonly used initialization methods detailed below:
20

---

Published as a conference paper at ICLR 2025
LeCun Initialization (LeCun et al. (1998)) Weights initialized using normal distri- bution with a mean of 0 and a variance of No where N; is the number input units in the layer. Mathematically, the weights drawn from N/O, )
Glorot Initialization (Glorot Bengio (2010): Weights are initialized using a nor- mal distribution with mean of () and variance of N,+N.y1' where N; is the number of input units and N is the number of output units. This method balances the vari- ance between layers with different widths. Mathematically, the weights are drawn from ((, N.+N.
2
He Initialization (He al. (0155) Weights are initialized using normal distribution with mean of 0 and a variance of 2 N, where N; is the number of input units in the layer. This method is particularly suited for layers with ReLU activation functions. Mathemati- cally, the weights drawn from N(0, N.).
Scaled Initialization (Rahnamayan & Wang (000))): Weights are initialized using normal distribution with mean of 0 and variance of 2s N, where N; is the number of input
units in the layer and a is parameter specific to each layer. Mathematically, the weights w are drawn from N(0, N;
These initialization methods help ensure that the network starts with weights that facilitate stable and efficient learning, avoiding the common pitfalls of poorly initialized neural networks. Using (A.3), we can calculate the respective Expectations and Variances of the Balanced Computation under the different initialisations:
Initialization
Var(ar pm)
Var(w
E.WF W W.W!]
Var Var W W.W?)
LeCun
Nesi
Ni
1
Nima -1 Ni-1
NNaz Ne
Ni
B
Glurol
Nost Noaz
N:
Nist
Myp
(Nita
+ Nr
NitNies
He
Ni12 Ness
N7.1
Scaled
a²
Nes af of)
(Neaz
[o Mist
No
Table 1: Comparison of Variance and Expectation of Balanced Computation for Different Weight Initializations
Fig. 6 shows numerical example of how the Balanced computation would look like for initialising
weights with LeCun, He, Scaled and Glorot initialisations using N; = 160,NNN44 NN+1 = 80, No = 120. With these numbers of nodes in each layer on the weights is visually similar to a scaled identity matrix. In Fig. We show how the sign of the scaled identity changes with the dimentions N₁, N N442-
appreciate how the Balanced Computation
21

---

Published as a conference paper at ICLR 2025
Balanced Computation for LeCun Initalisation
Balanced Computation for He Initialisation
2.0
1.00
-15
0.73
-10
200
0.25
-1.0
Balanced Computation Scaled Initialisation
Balanced Computation Glorct anibaaisttion
0.8
0.4
ialue
-0
C
Figure 6: Balanced computation different weight initializationsusing N 160, N441 80, N1+2 120, a1 ,02 2. The figure compares LeCun, He, Scaled, and Glorot initial- izations, showing how the balanced computation on the weights visually resembles a scaled identity matrix.
22

---

Published as a conference paper at ICLR 2025
0250 Gent
S
22
Figure 7: Impact of Network Structure sign of balanced coefficient for different initialisations. Lecun, He and Scaled initialisations depend solely on the ratio of the sizes of the output and hidden layers. Scaled initialisation's 's dependence is affected by the parameters 0,+1, 04. Glorot initialisa- tion's s sign depends on both ratios.
State-of-the-an models such as (Krizhevsky al. 2012), Liu et al. (2023a)) use more than 10,000 nodes in each hidden layer. This implies that if wished to perform some mathematical analysis on these models, assuming the initial random weights are A-Balanced would be a very close depiction to reality. In addition, these models often have a different number of nodes per layer, so understanding the effect of the relationship between N,, N(41, and N,42 is crucial.
Specifically, we aim to understand how changes in the relative width of layers i, + 1, + affect the Balanced Computation: suppose N,+1 kN; for some R+ and N,42 rN; for some R Then we can express the table above in terms of k, r, and N; only:
Initialization
Var(w
Var(w
B i+1 Wi+1 W.W?]
VarlW 6+1 Wi+1 W.W.]
LeCun
1 AN;
1
(A 11-
N. ( +1)
Ni
Glorot
2
2 N.(++1)
r htr k+1
r + (k+1)²
N((k+r)
N:
(r+k)a
He
2.
2
( 1) B)
KNN
24
Scaled
a2
(fasn) oZ)
ra R
a² IB
Table 2: Comparison of Variance and Expectation for Different Initializations
From Table 2 we can observe that as the number of nodes in each layer tends to infinity, while the ratios between the number of nodes in each layer (r, k) are maintained, the variance of the Balanced Computation tends to zero. Hence the Balanced Computation converges in probability to X-Balanced weights.
Further, Table 2 shows that a larger value of r will lead to a higher value À in every one of the initializations displayed. Moreover, larger & has the opposite effect. In addition, we observe that in some initializations there are limits as to what value X can take (such as in LeCun 0).
Some special cases of and k are interesting to consider to gain an intuition of how chang- ing these values influences the Balancedness of the Weights. In the table below, we consider the cases:
1. == h: the three layers of the network have the same number of nodes.
23

---

Published as a conference paper at ICLR 2025
2. 0, k fixed: N; >> N442,NNis No+1, the first of the three layers is much larger than the other two layers.
3. + co, it fixed: 2 >> Ni, Ni+1, the last of the three layers much larger than the
other two layers (k: is fixed).
4. k > 0, fixed: the middle layer exceedingly small, N << No, N442
5. k + fixed: the middle layer is much bigger than the other two layers, Ni-1 >>
N, N442
of
6.
di the inner and outer layers have ratio proportional to the scale factors of each
weight layer. This case is important for Scaled initialization.
Initialization
0
h
k: -
k
LeCun
1
I
Glorot
21
0
K++
He
21
(E)
I
-
2
I
Scaled
(a
oz))
Fo
i e
Table 3: Comparison of Variance and Expectation under Different Conditions
From Table above one can appreciate the impact network structure can have on the Balanced Computation of the weights of each layer. One can also see that there are many cases when the Balanced computation does not equal O, both the limit of r, and not in the limit.
We have showed that although the Balanced property is only preserved in Linear Networks, it occurs at initialisation for large networks which utilise some of the most common weight initialisation techniques.
These findings provide motivation to better understand the relation between the Balanced Computation of a Network, its structure and the regime will learn in. If are able understand the relation between X-Balanced weights and Rich and Lazy Learning in Linear Networks, one might be able to approximate these results to the nonlinear case.
A possible future application might be the ability heavily influence network's learning regime by altering the relative width of its layers, its activation functions or weight initialisation techniques used for each layer.
In order to better understand the effects of on the learning dynamics and learning regime of the network, first study aligned A- A-Balanced networks.
A.4 SCALE VS RELATIVE SCALE
straightforward intuition for the scale and the relative scale can be gained by considering the scalar case where N, N₁. = N, 1. In this scenario, it is easy to ensure that w w² satisfies while allowing for different absolute scales. For instance, Wi wa 0.001 or w1 we 5. In such cases, the absolute scale is clearly decoupled from the relative scale. However, in more complex settings, 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). However, this influence 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 We investigate first how X, the relative seale parameter, governs the transition between sigmoidal and exponential dynamical regimes. A similar argument applics to absolute scale, which appears explicitly as s, (0) these equations. Consider the case when À = 0, the dynamics of 8a reduce to the classical solution of the Bernoulli differential
24

---

Published as a conference paper at ICLR 2025
equation. In the limiting case where sa (0) - O, the system exhibits classic sigmoidal dynamics (characteristic of the rich regime), whereas the limit sa + oc yields exponential dynamics (characteristic of the lazy regime).
We performed an experiment to further examine the interplay between relative weight scale, absolute weight scale, and the network's s learning regime in the general setting. We define the absolute scale of the weights as the norm of W2 W1. We generate random initial weights of a given relative and absolute scale and train the network on a random input-output task. We compute the logarithmic kernel distance of the NTK from initialization and the logarithmic loss throughout training. We plot these values in heat map for À in [-9, 9] and relative scale in (0, 20). We repeat this procedure for three architectures: square network (N; = 2, No. 2, No 2). funnel network (N; 4, N₁. = 2, Nr, = 2), and an anti-funnel network (N; = 2, N₁ = 2, No = 4). These are the same architectures as in Figure in the main text.
In all three different types of networks, at the start of training, the model loss depends entirely on an absolute scale, not the relative scale. Throughout training across networks, the learning dynamics are intricately influenced by both the absolute and relative scales and fully captured by our theoretical solution. In the square network, the loss increases with absolute scale but decreases with relative scale, as shown in Fig [ Similarly, the kernel distance phase plot exhibits an intricate relationship with the relative and absolute scales, as illustrated in Fig Strikingly, for large imbalanced X, even at small scales, the network transitions into a lazy regime. The funnel and anti-funnel network architectures demonstrate antisymmetrie behavior shown in Fig. and Fig. IO Here, we focus on the anti-funne network for brevity. The evolution of the loss reveals that negatively À imitializations first converge, whereas positively initializations retain larger loss values. Additionally, the kernel distance attains its maximum for positive X, aligning with the results outlined in section 5 At convergence, the loss across all networks stabilizes uniformly, irrespective of initial conditions, confirming consistent convergence. This outcome aligns with the theoretical expectation for linear networks under gradient flow, which predictably converge to the same solution. Furthermore, in square networks, the kernel distance depends exclusively (results corresponding to the green curve in Fig. (B.) This observation illustrates that the regime at is consistently rich, independent of the absolute scale as predicted by our theoretical results in For funnel and anti-funnel networks, the kernel distance exhibits an antisymmetric pattern. In the anti-funnel network, the kernel distance depends mostly on X, achieving high values for positive A and approaching zero for negative À (matching the results in Fig. 5B pink line)). Conversely, in the funnel network, the kernel distance high for negative À and approaches zero for positive X, corroborating the results in Fig. 5B. (blue line). These results emphasize the interplay between relative and absolute scales, highlighting their critical roles determining the system's behavior. Altogether, the absolute scale and relative scale of the weights play critical role in describing the phase portrait of the learning regime, as first demonstrated in the Kunin et al. (2024) paper on ReLU networks.
the relative scale À and peaks À 0
C-3
25

---

Published as a conference paper at ICLR 2025
2.50
Kem2 Disaaoo, 1720
Lo0G, 1722
2.25
2,00
1.75
1.50
Dreek
1.25
Kerrel Disance,
Loss 1-100
1,00
Figure 8: Square network: Phase plots showing (left) the logarithmic kernel distance of the NTK from initialization and (right) the corresponding logarithmic loss as functions of the relative scale and the absolute scale. (Top to bottom) Different time steps during training = 1, t = 20, = 100.
Garnll Diabane,
Toanre 2006
Kamel Distanoe. 1-102
Loss, MC3
2.50
2.25
2.00
1.75
1.50
Figure 9: Anti-funnel network: Phase plots showing (left) the logarithmic kernel distance of the NTK from initialization and (right) the corresponding logarithmic loss as functions of the relative scale and the absolute scale. (Top to bottom) Different time steps during training t = 1, = 400,
I 4000.
26

---

SW? dt SW! Published as a conference paper at ICLR 2025
Komel Distance,
Locs, 21
Rodie
26
Cattares,
Konel Detane
22
1.8
Figure 10: Funnel network: Phase plots showing (left) the logarithmic kernel distance of the NTK from initialization and (right) the corresponding logarithmic loss as functions of the relative scale and the absolute scale. (Top to bottom) Different time steps during training
400,
( = 4000.
B APPENDIX: EXACT LEARNING DYNAMICS WITH PRIOR KNOWLEDGE
B. APPENDIX: FUKUMIZU APPROACH
Lemma B.1. We introduce the variables
W W
W W, W2W,
WTWT W. W
Q
and
QQ
(23)
Defining
(EuRyT 1
F
(24)
the gradient flow dynamics of QQ)((t) can be written as a differential matrix Riccati equation
d (QQ")
FQQ + =Q=F=(QQ0))
(25)
di
Proof. We introduce the variables
W W₂
WTW, W/WT W,W1 W?
Q
and QQ"
(26)
We compute the time derivative
aw? dt dW
Wi + WT aW, dt W + W2 dW, de
W + WT dW, W
(QQ")
de WT dW:
(27)
di
de
Using equations 18 and 19, we compute the four quadrants separately giving
dWT
dWi dt
W
+
(28)
dt
27

---

Published as a conference paper at ICLR 2025
[I W2W1) W W.W1+ W W2W1)
(29)
= - Wi + W W/ -
W W, - (W2W1) W
(30)
= (2 W w W W. W. W/ W/W. W. W.
(31)
w + W M dWT
(32)
dt
dt
34 W,W, WWW + W/ W. (D W2W,)²
= (laaa W2 W/ W W/W. WWW.W.W.
(33)
(34)
[W W + W
dw1
(35)
dt
dt
= -5 WWW,) + W₂ WEWT W2
(E
(36)
= = Wo W2W W2 W W WT.
(37)
(W2 di
dw?
WT + W₂
(38)
di
# W.W, W WT W.W, (Ite W2W,)²
(39)
S W W + W.W T W, W/ W/ W W. W. (W
(40)
= = W + W W OF W. W. W W W. WWWWWwW
(41)
5t W. W,W, (D W. W.W.W/W?
- W. W.W, AW,W7
(42)
Defining
F l
(Σve)?
(43)
y
the gradient flow dynamics of QQ (t) can be written as a differential matrix Riccati equation
d (QQ") FQQ + QQ"F (2Q").
(44)
dt
We write _ (QQ") for completeness
d
(ÉWR)T) I/
W W W
W/w, W.W?
WTw,
wTw, A W2W, W
(ÉWR)T
di
(Q07) = Syve
+
Syve
2
WW W2
W W W₂ WT
2
(45)
(Σ T
W W,
W
W
WT
T
-
(É
3/4
W.W,
W. W
+
W2
Wow,
(46)
W/W,
W/W.
W
W2
W W,
W W
W
W
W/W: (5 W.
W W. + +SSNE)
WF W A W
LNT W/ WT + W
AW WT + WT WIWEE
AW W. + (E"RIT
(47)
2 SW WT + W.W, (SUR)T
W, W + +W W. (S
WTW, W W
W W W2
WT W W.W.
W.W,
W W
28

---

Published as a conference paper at ICLR 2025
W. + (E S"" Wi + A 2 W.W,
WWW,, + WWWW
E W| WE WWW!
AW W + WT W (SSN)T
WT 2 + WT W (2
(45)
W. + W W1 (I"")T
W2W2 + W2W, (I
WT', W. WTW,W W. W/W,WT W1 W
W2W1W1 W1 + W2
W2W,W? W WWW:
The four quadrants of
27 27 are equivalent to equations
31
34
37 and 42 respectively.
B.2 QQ DIAGONALISATION
Lemma B.2. If F PAPT is symmetric and diagonalizable, then the matrix Riccati differential
equation di (QQ") = FQQ" + QQ" - -QQ with initialization QQ (0) = Q(0)((0)
has unique solution for all 2 0, and the solution is given by
2
1) PTQ(0)
(0)
I P
Q
(48)
2A
This is true even when there exists A; = O.
Proof. First we show that there exists a unique solution to the initial value problem stated. This
true by Picard-Lindelöf theorem. Now we show that the provided solution satisfies the ODE. Let
2 PT Q(0) such that solution QQT
L = c Q())((0) and C I+)(
(t) LC-
2A
The time derivative of QQ" is then given by (QQ")
d
d
- ( + LC-1 d (I")
+ L d
dt (L)C-
(49)
dt
dt
Solving for these derivatives individually, we find
d
d F Q Fc F Q(0) FL
dt (L)
(50)
dt
d - [C) -
d c2 -1 PTQ(O)C-1
(51)
dl
di
2A
We consider the derivative of the fraction serpately,
2A
24
(52)
dt
2A
this is true even in the limit as A + 0. Plugging these derivatives back in we see that the solution satisfies the ODE. Lastly, let t = O, see that the the solution satisfies the initial conditions.
B.3 F
DIAGONALIZATION
Lemma B.3. The eigendecompositionn F = PAPT where
S
P= VVG HÂ) V(++GG
V2V (2Ú
V2 Ú HG) -Ü HG)
A
x
(53)
and the matrices Sx, A1, A. and & are the diagonal matrices defined
12 \ = sgn(No -., > N))
Sx S G=
1
S^ 1
H H sgg((A))
(54)
S
S
VI+ H²
Beyond 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 computing
= PAPT
(55)
n
29

---

Published as a conference paper at ICLR 2025
B.4 SOLUTION UNEQUAL-INPPU-OUTRUU
Theorem B.4. Under the assumptions of whitened inputs, the temporal dynamics of QQ" are
lambda-balanced weights 2 no bottle-
neck 3
ZIA-IZ
Z Z2A-1
QQ
Z2
where the variables Z1 R.N.N, Zz NoNN, and A RNNNN are defined
(1 v H + B"
2 v H c + +VLe
1
(56)
(2 2
ú(( + HG)e
& B" +2 HGJE c + +
(57)
2 -1
25- I 1)
c" + Wa(0)'01 x+
A(t) -
B" -c
U W2(0)
4SA
4SA
W1
(W)
(58)
Proof. We start and use the diagonalization of F to rewrite the matrix exponential of F and F. Note that PTP= = PP = I and therefore p -P-1.
P."P"
ViG HG)) [ HG)
V Hd) -Ü HG)
V V2U
ViG Hτ) UGG HG)
V
-Ü HG)
V2U
A17
ViG- Hб) / [0(G HG)
V + HC) G HG)
0
-
V2
1
ViG HC)
UGG + HG)
- HG
VIG+ HQ)
+2
V
01
A
v2
=0 2M M
(99
F
F
F
OA '0 + - 1)AA
MM
(60)
F = OAOT + 2MA M
(61)
Where M
Placing these expressions into equation 48 gives
(t) Oe 4 Om + 2M MT Q(0)
I+ Q(0) (0(22 -1 1) A' - M( -M 1)AI'M²) Q(0)
(62)
Q(0) [oA+OO + +
V(G- - -GG
V(GHHG) T WT
Q²Q(0)= = V2 Ú(C+HG) -Ú -ÚG--HČ W2(0)
1 ( -
22 (G W - -G HG)
1 B
(63)
V2 -CT
where
B = W HG) E NAXN,
(64)
C = W HG) - HŌ) NANN.
(65)
30

---

Published as a conference paper at ICLR 2025
V HG
(55 0
OgAI/T
V2 U(G+HG (G-HH)
c
1 V(C-HG) S.
V HGG+ S - -GG-HG)e 8
(66)
V2 U
OC O 1 VV ((G-HM -
+ V S.
B"
(-cr)
- - -GG-HGJe Sx -
2
ViG - H S Bd V(G+ HG) -x - 2
(67)
S. + B" +Ü -AUGG-HCH Sx CT
W (0)
2 Q00 = 2 V2 U 0
e
V2
U
W2(0)
VICLL:V
0
Ú
(W2(0)
VI
Uie U (2)
(68)
Putting it together we get the expressions for (t) and Zz(t)
Oo 2M Q(0)
//i6 H $ Bd VV V(G++ãCe -
V Vi W
Ü HG)e S.t +T+UGG HH)e - S.:
+ U (
(69)
(1) 1 viā H 8 B" v Hā)e S. V م V
(70)
(1 Ú H
S. B" + ū HOJE C. Ú A U?w:(
(71)
2
We now compute the terms inside the inverse
(( TMA''M"QUA
[
A1
V
(O)
Wi(0) W2(0)T]
V2
0
cxi4
0
1
v2
W2(0)
A
(0))
(W] W2(0)']
= =
A x ( Ú A
(72)
Q(0) = 2 (W) (2)
1
V
2
V
W. (
22
0
V2
Ú
W2(0)
V
V
(0)7
= =WWr(0) W2(0)
U
A
U '0 (2
(73)
31

---

Published as a conference paper at ICLR 2025
Now
2 Q00000 (22A A ~-
= - 4
1
[B - -cl -0((()e (cA=-1) A (1B²)
-
(Cca)
B
1 ( B ( 25x -1) ($x) -1 -BBT - -c(e-25, - (Sx) (5.)
1
(74)
Putting it all together
A(t) -I+B
25
I
BT CC
(e- (e-25.4 -1
CT W2
(A
W2(0)
4S
4S
A_
W1(0)V1
VIW,(0)
(75)
T
So, final form:
QQ (1)
VGG H 5.t BT - H&)e S. CT + V
Uū + H 5. 1BT++U( + -U(G-HC -HGje --))e- CT +
+
1 I+
2
e e-2 CT
B
BT C
4
Sx
Sx
cdi: I
+ +2200)
U/ V
V V
A1
A1
1(
LVVIG- Sx
S BT - -
J Hā)e S CT + +VIeL:
)"
jÜG HG) S+ + f
HG)e S. CT + 1 O
(76)
B.5 STABLE SOLUTION UNEQUAL-INPUT-OCTPUU
Theorem B.5. Given the assumptions of Theorem 4.3 further assuming that B is invertible and defining + sgu(N. - N) the temporal evolution of QQ" is described as follows:
QQ" (t) = Z e : B-1
B
(77)
I- -28
S. -1
Sx C
C-T
T
4SA
45
S. B W2 -, A T S
S. B W,
U0 T SA
+e S1t B W (WVV VV -T -..
- S W W1 V/ -
- 3,11-1
ZT
1/ - (G-HG)- - HG)e S B ;
+ -V
-2
z
1ū (G+ H.) + (ā- H&)E S. $ 0
B 5.
+
(78)
32

---

Published as a conference paper at ICLR 2025
Proof. We start from
QQ"(1=
1 L HG B" -V 3. + V A : V/
GG + H e S + L HG) - + -e A A (2((0
2 -1
1 I+ 4
B
BT -C
CTT
Sx
Sx
Ed+t
+ +22
U -V
V V(
A
A
-VGG-HG) : (V( HÕ)e S CT + V VI
]²
T
HG)+ + + - HG)e $ OT + -e UT
(79)
We extract B B-Te- -5.4 from all terms as exemplified bellow
V (G (G-HHG)-CC - - (
QCT)) 2 ú ( + $ S. C B T -
Bd S (80)
and rewrite the dynamis as
QQ"(1)
LiG HG) - S:: S Sxt VieL
V W B
T 5
LU(G+ + HG) + -0 H 5.
C B-T + + + (0) B B- S.
1- 25 -2..
2 -1
B. S 1
S. - B-CC
BT
S^
S^
tt- $ B T.
UT TT -S
T
te BBB-WW(O0VV
V BO -.
A
V - H - V HH)e S. : B T V e T
S,
CT T S.
U Sx
(81)
33

---

Published as a conference paper at ICLR 2025
QQ
V G Hद) (G Sx B VWV
j (G+ HG) + - S. + C 7 B Sx
Ü B_ S2
e : B- T :
25x
E - -1 I
Sx C
C C²BBBT-S.
4Sx
x
S.4 B (W
8.
S
W/ V VV S
B VIW. -
HG) - (G H 5 B $
+ WW
Ú (G + Hद) + (G-HC (G-HG) S C B B- S.4
+ Ú
(82)
where edit : sgn( No No) N/ is a scalar
B.5.1 PROOF EXACT LEARNING DYNAMICS WITH PRIOR KNOWLEDGE UNEQUAL
DIMENSION
We follow a similar derivation presented in Braun et (2022) and start with the following equation
QQ" (t)
1 + 2 M Q(0)
I Q(0) (0 245 -
-1)-111AA A -
+ - -AA'M' 'M') Q(0)
(83)
C-1
Q00) OT + i
R
-LC-'R. RR
(84)
Substituting our solution into the matrix Riccati equation then yields
QQ" = FQQQ + QQQ (QQ7)²
(85)
dt
T dt
LC R - F R RF LC RRR- R.
(86)
Using the chain rule ((AB) (OA)B + A(B)) and the identities
d
1) = -1 d
A)A and d
d
(c!A) = AA = C!A A
(87)
dt
dt
db
d
d
QQ = dt
LC-'R R
(88)
dt
d
d C-1R
=
L
C R + rL
(89)
dt
d
L
c
R
+++1--11-
+
R
7L
d
c-1 R.
(90)
dt
dt
34

---

Published as a conference paper at ICLR 2025
Next, we note that
0 --y 1 VVG- HC) VI
/2 1 H - -UG- HG)
(91)
(V( HG) V(G+HG)
1 -
V(
o²o=
// ( HG) -
3 ( Ú -
-UG-
(66)
(92)
=1
(93)
VG- Hद) HG V
O
(94)
V2 UG
-U /2 Ü
1
+G+
2 ( 1
(95)
=0
(96)
and
úن (V(G-HG) V( HG)
M
V2 MV
(97)
((G+
-G-HH))
-
(98)
2
0.
(99)
get
d
d
T
QQ"
(LC-RR
(100)
dt
dt
d
=
L
C-RR
R
C--RR
(101)
dt
dt
d
(d
d
L
C
R
R
rL
c-1
R.
(102)
dt
dt
dt
with
T
(
c C-
R = (o- + 2M
e M²)
Q(0)C-'R
(103)
2T
-OACA+OO + +MALICALMM Q(OCC-RR
(104)
= O + 22 (OCAS (O.A:
+M 4Mm Q(0)-
(105)
=FLC-'R,
(106)
LLC-1
d
R
-LC- Q(()"
0 - A$ AOT + 2M
& Trr 2r
M
(107)
dl
= ((0)
0 A AOT + 1
TT
MT
(108)
2r
= =CC-IRF
(109)
and
35

---

Published as a conference paper at ICLR 2025
d
TL
c-1
R = -
- C C-'
(110)
dl
di
--IC
( Q 2n AA O'Q(0))())
(111)
1 +
+T20(0) (0(0))4-Mas) M (AL)
C-
--IC-
Q00) 21 0 Q(0) 200000
C-'R
(112)
-CCC
QFQ(()
+2Q(0) M A'M
(113)
+ +2Q(0) M O(Q(()
+ MQ
M (
C-R
= RCC R.
(114)
Finally, substituting equations ity.
103
107 and 110 into the left hand side of equation 86 proves equal-
C RICH-LAZY
C.1 DYNAMICS OF THE SINGULAR VALUES
Theorem C.1. Under the assumptions of Theorem 4.3 and with task-aligned initialization given by W1 (0) RS, VT and Wz(0) US, where R E NN x.N is an orthonormal matrix, then the network function is given by the expression W. W (1) ÚS)) where S(1) E Nr x.N, is a diagonal matrix of singular values with elements 8 that evolve according to the equation,
sa so + Ya. (t; A) (3a. so))
(115)
where Ba is the a singular value of S and Yo (1; X) is a X-dependent monotonic transition function for each singular value that increases from V/a (0; A) 0 to lim Ya(1; defined
Sx sinh (28x. :) + 8a + x² SA sinh (25x {) + (5aSa + ) cosh (28x.a 4) + Sa (3a. sa)
4 cosh (28x.a :) - Sa + 12)
Yo (t:
(116)
where 3x,a V + 8Aa V (a)2 4. and (0). We find that under different
limits of X, the transition function converges pointwise to sigmoidal (A - 0) and exponential (A + =00) transition functions.
2Fo +-++
as + 0,
(t:
1+
(117)
+ =
Proof. According to Theorem
4.3 the network function given by the equation
W2Wi(t) (t) (t).
(118)
which depends on the variables the initialization B and C. Plugging the expressions for task- aligned initialization Wi (0) and (0))) into these variables get the following simplified expres-
36

---

Published as a conference paper at ICLR 2025
sions,
B R S2 HG) + Si(G HG))
(119)
Dn
C C= R (2 HG) - S + HG).
(120)
De
where we defnne the diagonal matrices DB and De for ease of notation. Using these expressions, we now get the following time-dependent expressions for (t). A (t), and (t).
(1) v (i-- H 5x Dn - -G Hā)e Sa S..De) R.
(121)
(f) _ (G H
Sx DB + (
3.t De
R
(122)
2
25
I
25x
A(t) R I+
D
D
R
(123)
45
4SA
Plugging these expressions into the expression for the network function, notice that the R terms cancel each other resulting in following equation
4 Da (G+ HG)e (tWi(t) = (I- (G)
Sxt De
(G H Da DD (G H&)e 5. (e)
4
+
2 $
41
D
D2
S
Sn
S(t)
(124)
Notice that the middle term is simply a product of diagonal matrices. We can factor the numerator of this expressions as,
(G 28. D3 HG) (ā ( DrDc- H e 25
-25
D
(125)
We can further factor this expression as,
GI # 25 D - 25. De -
(126)
Putting all together we find that S(t) can be expressed as,
GI - H²)
2 D - - 25, : 4G HD
S(t)
(127)
255
-5,
41 +
D3
D
Now using the relationship between H and & we use the following two identities:
$
A
G (I - H
46
(128)
Sx¹
S
Plugging these identities into the previous expression and multiplying the numerator and denomina- tor by Sx gives,
25 D3 - 225 D2) - DD
S(t)
(129)
4S, + 28 D3 - 28 Dÿ + DD- D
Add and subtract $ (45s D2 D from the numerator such that
5 255 + D3 - -D² DDDDC
S(1) = 5
(130)
45 25 D2 c 25 D + +D -D?
37

---

Published as a conference paper at ICLR 2025
Using the form of Dg and De notice the following two identities:
DBDC From the second identity we can derive a third identity,
3/2 Sx
5 (5-S.
D
[ + =11
(131)
S^
S
45, + D - D
(
( S.Si)
(132)
S
Plugging the first and third identities into the numerator for the previous expression gives,
(15*+
($ S,S,
St) S
SA
(133)
48, 25 D - e 2 D2 + D' - D
Multiply numerator and denominator by Sa and simplify terms gives the expression,
SA
(t 5 5
($ S.5,)))
(134)
Sx $
$1 4
25 D3 - 25. D2
S1 (D} - -D -
+
Thus have found the transition function,
$. 28 D - 25 D
(e28)
$ 7 (Di - D
T
+
((;
(135)
S.
D3
e
28. D
+ $, (8xx D D
4
We will use our previous identities and the definitions of D and D2 simplify this expression. Notice the following identity,
Sx ( D;
X² 25x+ SSSA sinh (2555)) (SS(()) +
cosh 25 t
sinh
- I
4
7 (136)
Putting it all together get
SxSx sinh sinh (25x.) + (55((0) 1) cosh (25x - ($s(()) + 1
>1: A)
(137)
S,SA sinh 2 + ($S(0) + 1 cosh 225x 4 + 5 S S())
We will now show why under certain limits of À this expression simplifies the sigmoidal and exponential dynamics discussed in the previous section.
Sigmoidal dynamics. When À = O, then S^ = Š and SA - (0). Notice, that the coefficients for the hyperbolic functions all simplify to (0). Using the hyperbolic identity sinh(x) + cosh(z) = e", we can simplify the expression for the transition function to
(t; A) = 25 55(0)
(138)
- -S((0) - + 52
Dividing the numerator and denominator by 55((0) gives the final expression.
Exponential dynamics. In the limit À + too the expressions Sx +
[Nl
2 and Sx
2
Al
X²I I. As result of Additionally, in these limits because A'1 > 55(00 then ($S(()) 1
these simplifications the coefficients for the hyperbolic functions all simplify to can again use the hyperbolic identity sinh(x)) cosh(z) e simplify the expression
X I. As a result we
2 Lel
A2 7
A1;
(139)
M + $ (5-S(0)
X² Dividing the numerator and denominator by 3 results in all terms without coefficient proportional to X2 vanishing, which simplifying further gives the final expression.
38

---

Published as a conference paper at ICLR 2025
C.2 DYNAMICS THE REPRESENTATION FROM THE LAZY TO THE RICH REGIME
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. (2023)
The NTK intuitively measures the movement of the network representations through training. As
shown in (Braun et al. (2022), , in specific experimental setup, can calculate the NTK of the network in terms of the internal representations in straightforward way:
NTK IN. @ x' WT W (()X + W. (
(140)
In order to better understand the effect of À on NTK dynamics, we first prove some theorems in- volving the Singular Values of the X-balanced weights, and the representations of X-balanced network.
C.2.1 LAMBDA-BALANCED SINGULAR VALUE
Theorem C.2. Under a X-Balanced initialization
2
if the network function W2W1() =
U())(())VT (t) is full rank and we define (t) = (t) + 2I. then we can recover the parameters (t( = (2t))2(t) (t), Wi(t) (t)) (t) up to time-dependent orthogo-
(t)
nal transformation R(t)) of size N₂ No where
(1( - ((5 - ) M,
S2(t) ((Sx(t) + 2) 00axx(,, Mo N₂)
(141)
Proof. We prove the case N; I No and N₁ = (N, same structure. Let USVT W2(t) W1 (t) be the Singular Value Decomposition of the product of the weights at training step I. We will use W₂ W2()), W1 = Wi(() as shorthand.
NN, No). The proof for No I N₁ follows the
We write the initialisation for our setting W2 W, USVT We therefore can write without loss of generality the weight matrices W₂ = US and W; = V. In this case we requiere that G2 G. and S,S= =
We assume the balanced property such that W W2 W WT XI. know this holds for any since this is a conserved quantity in linear networks. The matrices W, wf and W W, are symmetric, which consequently implics that their singular vectors are orthogonal. Consequently, in specific scenario, require that G1 and G2 orthogonal matrices and it follows that
G2 = G GT G2 = G R²
(142)
(143)
where We can write W2 = US,R", W RS, V where R is an orthonormal matrix and S2, S are diagonal (possibly rectangular) matrices.
Hence
RS SRR" RS,, AI
(144)
S S,S, =
(145)
The matrices S,, S2, have shapes (N,, N.), (N., NN,) respectively. We introduce the diagonal ma- trices $1 of shape (Nn, N.). $2 of shape (N;, NN) such that the zero matrix has size (No N,, N₂)
S = S = (
(146)
Hence
ST S1S, AI
(147)
39

---

Published as a conference paper at ICLR 2025
From the equation above and the fact that S,S²
S derive that:
3
1
VA2I+
4.2
+
VA 452
11
S2 =
S =
(148)
2 : RR,
2
Hence
VAPI+-SS+AAI
W2 = U
W1
R
Omax(o.NN.NN.) V
2
(149)
C.2.2 CONVERGENCE PROOF
With our solution, QQ"(t), which captures the temporal dynamics the similarity between hidden layer activations, we can analyze the network's internal representations in relation the task. This allows us to determine whether the network adopts rich or lazy representation, depending on the value of A. Consider a J-Balanced network training on data g ÜSVT. We assume that the convergence is toward global minima and B is invertible
Theorem C.3. Under the assumptions of Theorem B.5 the network function converges USVT
and acquires the internal representation, that is W Wi VSTVT and W W U
Proof. As training time increases, all terms including a matrix exponential with negative exponent in Equation 77 vanish to zero, as Sx = Sx a diagonal matrix with entries larger zero
As training time increases, all terms in the equations vanish to zero. Terms Equation 77 decay
s2 x
lim e-
O.
(150)
and
lim (+000
52 O.
(151)
where Sx Sx is diagonal matrix with entries larger zero Therefore, in the temporal limit, eq. 77 reduces to
W (() W2 (t)
WT WT (()
lim QQ" (t) (-000
lim 1-00
(152)
VVG HG)
[5x " ( (U(HG+
(153)
(U(HE+G)
V(G (( VGG-HGIS,HG++
(154)
(G - HG-HGIS + Øद) = Sx H
(155)
1+ H²
((G-HÃC) 51 #2)
S2H2
V 452 X AI
(156)
12
14 H2
2
5x (2) S Hद²
S2(2))
4S + III
(157)
1 H²
1+ H
2
WF' WtW
WT (8)
lim QQ"(1)= lim 11000
(158)
t-
W W.W. (
V V ÜSVT ÜS:U
(159)
40

---

Published as a conference paper at ICLR 2025
C.2.3 REPRESENTATION IN THE LIMIT Theorem C.4. Under the assumptions of Theorem B.5 the representation tends to
training on data T0 ÜSVT.
as À + 00
W W2W2 = Ü
AI (mmxx(0.NN -.)
Omax(0.N.-
W V
$2
O0. Va
00axx(0,M N.-))
As + -0
W2W2" -
g
w wf ( AI O0. U..
(,, V
O0, X
(0. -N.)
As A -00
W2W,
- X
g Omax(o.N.- NN)) ÚT.
w Wi - v
(
IYI
0 0max(0.N)-No) Vi
Omaax
N0.
Proof. We start from the representation derived in C.3 and using the Taylor expansion of f(x:)
vI + x² we compute
I + 45$
AI
+
28
+ AI
(160)
2
2
|
1 +
28 A
+ 0(A-1)
+ +XI
| + A
$2
+
+
+0
(161)
2
2
|>|
Hence
A I + 45 + AI
A81 +
4S
+ AI
$2
52
lim X000
AI,
lim
(162)
2
20-00
2
|
X
Similarly,
V V - XI
A - X 2
$2
0(1-3)
(163)
A
V 2 + 45
AI
52
X2I + 452
AI
52
lim X
lim
-X
(164)
X
A
Since U, V are independent of A:
lim x
W2WT =Ü (5(,))
lim S2 :T
(165)
lim W. W1 V ( lim S₁ VT A
(166)
As |X| - co, of the network representations approaches scaled identity matrix, while the other tends toward zero. Intuitively, this suggests that the representations shift less and less as [X] increases. Next, we demonstrate that the NTK becomes progressively less variable |>] grows and ultimately converges to zero.
C.2.4 NTK MOVEMENT
Relationship between À and the NTK of the network
Theorem C.5. Under the assumptions of Theorem B.5 consider a linear network training on data
E ÜŠT At any arbitrary training time 2 0, let (2(1)W = U* V' Then,
41

---

Published as a conference paper at ICLR 2025
1. For any X R.
VA-1+1S-2-1
grx
OT = IN. & X
(167)
O Oxxxx
+U
2
0
SAT444 AI 0 gv
NTK(t) IN. @ X
(168)
-A-T4-55
0 U.T X"X
U
0
2. As À - 00:
(TK(t) - NTK(O) + A (IN. & V'V IN. 8 x V=V +
1
(169)
3. As X 000
1 (TK(L) NTK(O) + x (US*U & X"X - UU x"x) +0
(170)
Proof. Follows by substituting the expressions for the network representations in terms of À from (Braun al. 2022))'s 's expression for the NTK of linear network. Similarly, follows from substi- tuting the limit expressions for the network representations and the fact that the Kronecker product linear in both arguments.
The theorem above demonstrates that X + Jo, the NTK of X-Balanced network remains constant. This indicates that the network operates in the lazy regime throughout all training steps. The X-balanced condition imposes relationship between the singular values of the two weight matrices. Specifically, if W2 and W are X-balanced and satisfy W2W1 = Eyx, then for arbitrary singular values az bi, and 8;, the following relations hold:
a3 - D2 = X, a; b; = 84-
As À increases, the value of b, must decrease. In the limit as X + oc, aa > and b2 + 0. From + x. Since these equations apply to all singular values of the
the first equation, when 6? + 0, a matrices, it follows that for all i, a X, leading to the conclusion that:
W W₂ AI,
as expected. Consequently, the task representation becomes task-agnostic W,. The intuition here is that the weights are constrained by the need to fit the data, which bounds their overall norms. The X-balanced condition further specifies a relationship between these norms, and as [X] increases, this constraint tightens, driving W2 toward the identity matrix. In this regime, the network behaves sim- ilarly to a shallow network, with À acting as a toggle between deep and shallow learning dynamics. This finding is significant highlights the impact of weight initialization learning regimes.
C.3 REPRESENTATION ROBUSTNESS AND SENSITIVITY TO NOISE
As derived in (Braun et al. 2025). the expected mean squared error under additive, independent and identically distributed input noise with mean μ and variance of is
1
W, (x, +xx yv11
W.
(171)
2P
&x
where T3")
1 Tr EE EVeT) is noise independent constant that only depends we show that the network function converges to
the statistics of the training data. In Theorem C U and therefore
W2W
af
USV
(172)
N₂
02
$?
=
42

---

Published as a conference paper at ICLR 2025
As derived in (Braun et al. 2025). under the assumption of whitened inputs (Assumption D. in the case of additive parameter noise with μ = 0 and variance ow the expected mean squared error is
(W2 + Ew,) (Wi + Ew,) x; -yllé
2P
Sw,
(173)
5 N. W2ll
1 Noo [Wili + 1
+ NNNNNN.o' +c.
2
Using Theorem
we have
W Will = Tr(WT W1)
12T + 482 AI
T
(174)
NA
1
X2 + 45p + A
2
and
W2 = Tr( W
API + 452 AI
=
2
(175)
X2 + 48,
To find the À that minimises the expected loss, we substitute the equations for the norms, take the partial derivative with respect and set it to zero
(C) Ew. Swz 1
0X
N
NoW W
X2
+
4S?
2)
-
1
Noow
à
Σ
X2
487 +1
=0
Ox
4
Ox
i=1
NR.
Nn
(176)
A
N₂N, + +No
+ NNoN, 0
X2 45²
X2 48p
No
Nu N₂ No
N; + No
X2 + 45²
It follows, that under the assumption that N₂ = No, the equation reduces to
N
= 0.
(177)
We note, that the denominator always positive and therefore, that the left-hand side of the equation is always larger zero for any À > O, and smaller than zero for any À < 0. The euqation is therefore only solved for = 0.
C.4 EFFECT OF THE ARCHITECTURE FROM THE LAZY TO THE RICH REGIME
Theorem C.6. Under the conditions of Theorem B.5 when À > the network enters regime referred to as the delayed-rich phase. In this phase, the learning rate is determined by two competing exponential factors:
A
22
and
$24221
43

---

Published as a conference paper at ICLR 2025
As X increases. various parts of the network display different learning dynamics: some components adjust rapidly, converging exponentially with X, while others adapt more slowly, with their corver- gence rate inversely proportional to X. leading slow adaptation.
Proof. The solution to Theorem B.5 is governed by two time-dependent terms:
and e e 1
The first term exhibits exponential decay approaching zero time progresses:
52 231 0.
lim
In the limit as lambda gets large the rate of learning is given by
A + 4S2
AI 2
lim 1000
(178)
The second term also decays over time
lim e
A :
52 22
0.
but the limit as lambda gets large the rate of learning is given by
X2I + 45 AI $2
lim A000
(179)
Thus, À increases, the convergence rate slows for certain parts of the network, while other components continue to learn more quickly. This explains the delay observed in the delayed-rich regime.
D APPENDIX: APPLICATION
D.I APPENDIX: CONTINUAL LEARNING
We build upon the derivation presented in Braun al. (2022 incorporate the dynamics of contin- ual learning throughout the entire learning trajectory. Utilizing the assumption of whitened inputs, the entire batch loss for the ith task is
1
L. (T,)
W2W,X - Y.Il2
2P
1
TF (W2 - (.2 -
2P
1
- 1 - TWW + Tr (Y.Y?)
Tr XX
2P
2P
1
TF (W:W,
TT
W2
Tr E!!!
2
1
Th ((W2W1 - -S,0) (W:Wi- E000
EN Sym
+
2
1 W2W
Sye0 - Tr Σ Sqre"
1
+
Σ uo
2
Hence, the extent of forgetting, denoted as for task Tr during training on task Ta subsequent to training the network on task T, specifically, the relative change in loss, is entirely dictated by the similarity structure among tasks.
Fi (T,, Th) = C. (Tx) - C, (Tx)
2
Σ
$.0
+
W2W W.
Syz
2
F
2
F
|S
SuzZ
2
W, W,
Eq
2
44

---

Published as a conference paper at ICLR 2025
It is important note that the amount of forgetting is function of the weight trajectories. Therefore, we have analytical solutions for trajectories forgetting as well.
A
Task
Task
Task 3
0.5
W2W1(t)
0.0
-0.5
5
Loss
0
5000
10000
15000
20000
Training Steps
Figure 11: Continual learning. A Top: Network training from small zero-balanced weights across sequence of tasks (colored lines represent simulations, and black dotted lines represent analytical results). Bottom: Evaluation loss for the tasks in the sequence (dotted lines) while training on the current task (solid lines). As the network optimizes its function on the current task, the loss on previously learned tasks increases.
Figure. D.I panel was generated by training linear network with N; = 5, N₁ 10, No = 6 subsequently on four different random regression tasks with 25. The learning rate was r) = 0.05 and the initial weights were small (a 0.0001).
D.2 APPENDIX: REVERSAL LEARNING
As first introduced Braun 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 u, ü, and V,, V., 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 V, with the exception of one more columns of the left singular vectors, where the direction is reversed: -u, = ū; It is important note that if reversal occurs in the right singular vectors, such that -v, V, 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, both B S Hद) S HG) and C S HG) S.V HG) are diagonal matrices.
In the case where lambda is zero, the same argument given in Braun et al. (0022) follows, the diagonal entries of C are zero if the singular vectors are aligned and non zero if they are reversed. Similarly, diagonal entries B are non-zen con-eero the singular vectors are aligned and zero if they are reversed. Therefore, in the case of reversal learning. B is diagonal matrix with 0 values and thus not invertible. As a consequence, the learning dynamics cannot be described by Equation (66] However, as B and C are diagonal matrices, the learning dynamics simplify. Let b, C₁ S; and 8;
45

---

Published as a conference paper at ICLR 2025
denote the i-th diagonal entry of B. C, S and S respectively, then the network dynamics can be rewritten as
W2W1(t) ū (G e + (G - -
2
S + 1
c 28 -1) , B 1 c 28x ( -1)
S -
B
(180)
GG HG
8x b (ā H S.t() VT
N;
bie 2 c.e 26
Σ
ū.V.
(181)
4sxi
+ b3e
bク²yi cie
25 T + c,
1=1
bx,b 33,C,3. 45.
n.
45
2 + b,,
1 - e 25x1
+ s2,c, (e (e-28rit
e-18xi
(182)
follows, that in the reversal learning case, i.c. b 0, for each reversed singular vector, the dynamics vanish to zero
SA,005,e-1Bx Bric
lim 10000 43x,ee 25x + s, --axit- 43xe -4āxe
ū. 0.
(183)
Analytically, the learning dynamics are initialized and remain along the separatrix of saddle point until the corresponding singular value of the network function decreases to zero and stays there, indicating convergence to the saddle point. 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 time- consuming Lee et al. (2022). In contrast, when the singular vectors governing temporal dynamics, as described in Saxe al. (2014), is recovered. Under these con- ditions, training succeeds, with the singular value of the network function converging to its target value.
aligned (c 0), the equation
N
S
Sxib,5xi V.
lim (-00
ū.V.'
(184)
A
28 + syib? (1 c22x
28x
siise
=f.ū.V..
(185)
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 point due to numerical imprecision, but the learning dynamics are catastrophically slowed the vicinity of the saddle point as shown in figure D.2 non-zero, the diagonal of C are also non-zero: this is true regardless of
In the case where 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 diagonal matrix with elements that are zero. In figure D.2
46

---

Published as a conference paper at ICLR 2025
20:00
steps Convergencn deps leaming
15200
Shop to Corneggero0
10200
5500
2500
Figure 12: 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 X = 0.
D.3 APPENDIX: GENERALIZATION AND STRUCTURED LEARNING
We 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 associated 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 S X S 10) on the hierarchical semantic learning task in Section 5 and then add new feature (e.g.. 'eats worms') a single item (e.g.. the goldfish) (Fig. D.3A), 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₂ and train on the single item with the new feature, while keeping the rest of the network frozen. Thus, only change the weights from the hidden layer to the new feature which may produce different behavior depending on how the hidden layer query the
representations vary based X. After training the new feature-item association, 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 projects the feature onto items based on their distance to the trained item (Fig, D.3, For example, after learning that goldfish eats worms, the network can extrapolate the hierarchy infer that another fish, or birds, may also eat worms, instead, plants are not likely to eat worms. Alternatively. as A 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 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 new feature association only depends the first layer W1. Let ÿ, denote the vector of the representation of the new feature ff across items in the dataset. Additionally, let un (/)? be the new row of weights appended W2 which map the hidden layer to the new feature. Following Saxe ct al. (2019b), (f)T is initialized with small random weights and trained on item H,, it will converge to
w (/VT A. WTTW//N / H.
(186)
ÿ = (H W/ /|
(187)
From this we can see that differences in the representations of the new feature across items و across are only influenced by W,.
In the case of the rich learning regime where À = O. the semantic relationship between features and items is distributed across both layers. Instead, when > O, the second layer W₂ exhibits lazy learning, yielding an output representation W2W? of a weighted identity matrix. However, the first layer W, 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 new feature (which only depends on the first layer). Thus, in this case, the 'shallowing'
47

---

Published as a conference paper at ICLR 2025
of the network into the first layer is actually beneficial. Finally, we also observe the opposite case when À < Here, rich learning happens in the second layer, while the first layer is lazy and learns represent weighted identity matrix. As such, these networks do not learn transfer the hierarchy of different items to the new feature.
Animal
Plant
Fish
Bird Algae
What
9,
== -10
A. 0
A 10
5
135
20
Figure 13: Transfer learning for different X. A A new feature (such as 'eats worms') is introduced the dataset after training on the hierarchical semantic learning task (Section 5 randomly initialized row is added to W₂ and trained on single item with the new feature (for example, the goldfish), with the rest of the network frozen. The network 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.
D.4 APPENDIX: FINETUNING
It is a common practice to pretrain neural networks on a large auxiliary task before fine-tuning them on a downstream task with limited samples. Despite the widespread use of this approach, the dynamics and outcomes of this method remain poorly understood. In our study, we provide a theoretical foundation for the empirical success of fine-tuning, aiming to improve our understanding of how performance depends the initialization. We're interested in understanding how changing X-balancedness after pre-training may impact fine-tuning on new dataset. We use ApT to denote how networks are first initialized prior to pretraining, and XPT to how they are re-balanced after pre-training and before fine-tuning on a new task. Similar to the previous section, we first train each network (for different values of -10 APT 10) the hierarchical semantic learning task. We then change the >-balancedness each network (for different values of -10 I AFT S 10) and retrain
dataset to observe how this impacts fine-tuning for different values and compare to networks that are not re-balanced to some AFT (AFT = 0) after initial pre-training. the >-balancedness of pretrained network lar values of each layer (S1, S2) using the singular values of the entire network function (S UTW, W, V). while keeping the left and right singular vectors of the network unchanged.
In particular, to
ÄFT, we rescale the singu-
We consider three different tasks to fine-tune the networks on. In the first, we add an existing feature from one item another item the hierarchy in order to disrupt the structure of the left and right singular vectors. In the second task, we consider the same reversal learning task discussed in Appendix D.2 where one column of the right singular vectors are reversed such that -v; d, Finally, we consider a scaled version of the hierarchy where each singular value is scaled by 2. Across all the tasks we consider, we consistently find that fine-tuning performance improves as net- works re- re-balanced to larger values of XPT and, conversely, decreases XFT approaches O. Networks re- balanced to AFT = 0 also learn more slowly compared to |APT| > 0. Interestingly, when studying networks that not re-balanced prior to finetuning (ArT 0: but are first initial-
48

---

Published as a conference paper at ICLR 2025
ized prior to pretraining to XpT), that they perform similarly on the new tasks to networks that are re-balanced to XFT = Xpr.
A
B
C
1/-
500
Figure 14: Fine-tuning performance on three tasks for different re-balancing ÄFT. A After training on the hierarchical semantic learning task (Section 5). networks are re-balanced and trained on one of three tasks: adding existing feature from one item to another item in the hierarchy (left), the reversal learning task in Appendix D.2 (center), or a scaled version of the hierarchy where each singular value is scaled by (right). B Change in loss on the task across different ÄPT for different Apr. As XFT approaches O, the loss on the new task increases across all Xpr. Interestingly. networks that are re-balanced to the same values (AFT ApT). C Dynamics of the loss across the first pre-training task and the new fine-tuning task. Networks re-balanced to AFT consistently learn slower across all tasks compared to networks that are re- re-balanced to larger magnitude values ((Arr) > 0
re-balanced prior fine-tuning (APT 0) perform similarly to networks that
In this work, we derive the precise dynamies of two-layer linear. While straightforward in design, these architectures are foundational in numerous machine learning applications, particularly in the implementation of Low Rank Adapters (LoRA Hu et al. (022)). A key innovation in LoRA is to parameterize the update of large weight matrix W IR axa within a language model as AW AB, the product two low-rank matrices A ddd and E grxd where only and are trained. To ensure A WW 0 at initialization, it is standard practice to initialize A (O0 o²) and B ; Hayou et al. (2024). is noteworthy that this parameterization, AW- = AB, effectively embeds two-layer linear network within the language model. When d, this initial- ization scheme approximately adheres to our X-balanced condition, with a² playing the role of the balance parameter Investigating how the initialization scale of and B influences fine-tuning dynamics under LoRA, and connecting this to our work on X-balanced two-layer linear networks and their role in feature learning, represents an intriguing avenue for future exploration. This per- spective aligns with recent studies suggesting that low-rank fine-tuning operates in a "lazy" regime, as well work examining how the initialization of A B affects fine-tuning performance Malladi al. (2023); Hayou et al. (2024). Our framework offers potential bridge understanding these phenomena more comprehensively. While detailed exploration of fine-tuning performance lies beyond the scope of this work, it remains important direction for future research.
(Hu et al. 2022);
IMPLEMENTATION AND SIMULATIONS
The details of the simulation studies are described as follows. Specifically, N,, Nj. and N, 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 1)
49

---

Published as a conference paper at ICLR 2025
E. LAMBDA-BALANCE BALANCED WEIGHT INITIALIZATION
In practice, to initialize the network with algorithm, serves scaling factor that controls the variance of the weights, allowing for adjust- ments between smaller and larger weight initializations.
weights, we use Algorithm E.I In this
Algorithm Get X-balanced
1: function GET_LAMBDA_RR GETLLAMBDDEELLNCDD in.lim, hhiden_dim,, oul_dim, a 1)
2: if out_dim > in in_im and À < 0 then
3: raise Exception("
Exception"LLambda Lambda must be positive if out dim ú in dim')
4: end
5: if in in_dim out ,ii and À > 0 then
6: raise (eception("Lamblee must be positive if in didiim out_dim))
7: end if
8: if hidden_lingg < muininn_iinn out plim) then
9: raise Exception("Neee ("Network cannot be bottlenecked")
10: end if
if hidden dim > max(in dim, out dim) and + then
12:
raise Exception( hidden_dim cannot be the largest dimension if lambda is not 0*)
13:
end
14:
W1 <- random normal matrix(hidden dimm, in dn.dmm
15:
W2 < a random normal matrix(out
sidden_dim,
16:
S. Vt) t SVD(W2 W1)
17:
R f random orthonormal matrix (hirix dim)
18:
S2 S2egual dim
(VX' + + St + +1) 1
19:
S deg
(VX2 52. 2) /2
20:
if dum > in dim then
S2equal
0 Ohidden dim dém
21:
S2 +
0
S1 Slequallddin
22:
51 +-
23:
else if in_dim out_dim then
51 S1 Slequal.dim
0 S2 + [S2equaLdimm
0 hidden dim cann-oou dim
25:
26:
end if
init. W₂ + U. S2 RT
28:
init_W, +- R S1 Vt
29:
return (inil_W,, (nnll_WWW))
30:
end function
E.2 TASKS
In the following, we describe the different tasks that are used throughout the simulation studies.
E.2.1 RANDOM REGRESSION TASK
In the random regression task, the inputs X E N:N are generated from standard normal dis- tribution, X ~ [ip = 0, = 1). The input data is then whitened to satisfy + X I. The target values E N. are independently sampled from a normal distribution with variance scaled according the number of output nodes, Y N(p 0, >xx)) Consequently, the network inputs and target values are uncorrelated Gaussian noise, implying that linear solution may not always exist.
E.2.2 SEMANTIC HIERARCHY
We use the same task in Braun et al. (022 and modify it to match the theoretical dynamics. The modification ensures that the inputs are whitened. In the semantic hierarchy task, input items represented one-hot vectors, i.e., X g- The corresponding target vectors, y,, encode the
50

---

Published as a conference paper at ICLR 2025
item's position within the hierarchical tree. Specifically, value of 1 indicates that the item is left child of a node, -1 denotes a right child, and 0 indicates that the item is not a child of that node. For example, consider the blue fish: it is blue fish, left child of the root node, left child of the animal node, not part of the plant branch, right child of the fish node, and not part of the bird, algae, flower branches, resulting in the label [1, 1, 1, -,,0,,,,000 0,0, The labels for all objects in the semantic tree, as shown in Figure A, are given by:
1
1
1
1
1
1
-1
-1
-1
1
1 0 1
-1
-1
0
0
0
0
1
-1
-1
Y 8
(188)
1
0
0
0
0
0
0
-1
0
0 0
0
0
0
1
-1
0
0 0
0
The singular value decomposition (SVD) of the corresponding correlation matrix, Ene is not unique due to identical singular values: the first two, the third and fourth, and the last four values are the same. To align the numerical and analytical solutions, this permutation invariance addressed by adding small perturbation to each column y:, for i E 1, N, of the labels:
y1 = yı 1+ 0.1
(189)
i
resulting in singular values that are nearly, but not exactly, identical.
E.3 FIGURE 1
Panels B illustrates three simulations conducted on the same task with varying initial X-balanced weights respectively X = 2, = 0, 2. The regression task parameters were set with (a VIO). The network architecture consisted of N; = 3. N₁. = 2, N,. = 2, with a learning rate of The batch size is N 10. The zero-balanced weights are initialized with variance = 0.00001. The lambda-balanced network are initialized with sigmary v1 of random regression task with same architecture.
7 0.0002.
On Panel C, we plot the ballancedness W W2( (0) W (0) w (0)* for two layer network initialised with Lecun initialization with dimension N; 40 ,N,, 120 -N,=255
E,4 FIGURF 2
Panel A, B, C illlustrates three simulations conducted on the same task with varying initial X-balanced weights respectively
-2, 0, À 2 according the initialization scheme described in E.7 The regression task parameters were with (a v10). The network architecture consisted of 0.0002. The batch size is N = 10. The
N; = 3, N = 2. N. 2 with a learning rate of n zero-balanced weights are initialized with variance o 0.00001. The balanced network are initialized with sigmarry v1 of random regression task with same architecture.
E.5 FIGURE 3
Panel A. B, illustrates three simulations conducted on the same task with varying initial A-balanced weights respectively À= = 2, A 0, X = 2 according to the initialization scheme described in E.7 The regression task parameters were set with (a v12). The network architecture consisted of with learning rate n 0.0002. The batch size is N The zero-balanced weights are initialized with variance a = 0.0009. The lambda-balanced network are initialized with sigmary = VI2 of random regression task with same architecture.
N; 3, N₁, 3, N₂
E.6 FIGURE 4
In Panel A presents a semantic learning task with the SVD of the input-output correlation matrix of the task. U and V represent the singular vectors, and S contains the singular values. This decomposition allows us to compute the respective RSMs USU for the input and VSVT for the output task. The rows and columns in the SVD and RSMs are ordered identically to the items in the hierarchical tree.
51

---

Published as a conference paper at ICLR 2025
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. ( Saxe et al. (2013). These matrices are generated using varying initial X- balanced weights set at À = -2, x 0, and À = 2, following the initialization scheme outlined in E.. The network architecture includes N, = 8, N₂ 8, and No 8 with a learning rate of 7 = 0.001 and a batch size of N = 8. Zero-balanced weights are initialized with a variance of a = 0.00001, while X-balanced networks are initialized with azy - vI based on 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; = Nn. = N, = 4 on a random regression task for À E [-5. -, 33,- -2, -1, 0, 1, 2,3, 4, 5] to convergence. Subsequently, we added Gaussian noise with μ 0, E [0, 0.5, 1] to the inputs (top panel) or synaptic weights (bottom panel) and calculated the expected mean squared error.
E.7 5
Panel A illustrates schematic representations of the network architectures considered: from left to right, a funnel network (N; = 4, NN, = 2. .N, = 2). square network (N; = 4, No. = 4, N. 4), and an perted-funnee network (N; 2, N₂
N₂ = 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:
(Ko, Kz)
S(t) ||Ko|l||KKell
The simulations conducted the same task with eleven varying initial X-balanced weights in [-9,9]. 9]. The regression task parameters were set with (ar v3). The task has batch size N 10. network are initialized with
The network has with learning rate of y) = 0.01. The azy- = /1 of random regression task. Panel C shows the Neural Tangent Kernel (NTK) distance from initialization for the funnel archi-
tectures shown schematically with dimensions N; = 3, N₂, = 2, and N, 2. The simulations conducted the same task with twenty one varying initial X-balanced weights in [-9, 9). The
gression task parameters were with (a v3). The task has batch size N 30. The network has with learning rate 17 0.002. The mambda-balanced network initialized with ary /1 of random regression task.
52


Reviews:
[{"review_id": "1GT3ilfh7A", "reviewer": "Reviewer_6KfS", "summary": "The paper theoretically derived the exact solutions for the learning dynamic of $\\lambda$-balanced initializations in two-layer linear networks and discussed the implications in continual learning etc.", "questions": "a. Could the authors comment on how they would expect the results to generalize to more complex tasks?\n\nb. Although it is mentioned some choice of $\\lambda$ may be beneficial to transfer learning due to induced lazy learning in the shallow layer (D3), it would be at least better to show a performance increase.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Originality: The paper built on top of existing work by extending from zero-balanced condition into $\\lambda$-balanced scenarios, leading to a continuum from lazy to rich regime in terms of weight structure of two consecutive layers. \n\nQuality: Solid theoretical derivations backed with rich simulation results. \n\nClarity: The paper is very well-written. \n\nSignificance: The $\\lambda$-balanced discussed in the paper covers from architecture shapes to initialization schemes and relate these structural properties to the learning regimes, which will be of interest to the neuroscience community as well as continual learning and beyond communities.", "weaknesses": "The only weakness of the paper is on the weak demonstration of applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YD9mTnzzzN", "reviewer": "Reviewer_ifg2", "summary": "This work extends previous results on the exact dynamics of deep linear networks to include the lambda-balanced initialization conditions, where lambda is a tunable parameter that allows for exploration of a full range of dynamics from lazy to rich learning. By detailed analysis of the result, the authors showed that the transition from the lazy regime to the rich regime depends on the complex interaction of multiple factors, including the balance parameter lambda and the network architecture. This paper then also discusses potential application of the results on more complicated learning paradigms such as continual learning, reversal learning and transfer learning.", "questions": "1.\tA minor point: over-parameterization is commonly referring to the relation between the amount of parameters in the network (capacity of the network) vs. the amount of parameters needed to achieve zero training loss. I believe this may be different from ‘N_h > min(N_i, N_o)’?\n2.\tThe assumption that the input data is whitened is relatively strong and not very applicable to realistic settings. Do you have any empirical evidence demonstrating that your observation at least qualitatively still holds for benchmark datasets without necessarily whitening the data?\n3.\tThis may be just a typo. In Fig.2 b, in the caption it says lambda=0.001 but in the figure lambda=0.\n4.\tIn previous works concerning lazy and rich dynamics, especially in infinitely wide networks, there’s usually the assumption that N goes to infinity and depending on how initialization scales with N the network resides in either the lazy or the rich regime. Here there doesn’t seem to be any assumptions about large N or the scaling of the initialization. Specifically since Nh=min(Ni, No), it seems that N is finite. However, it seems that the authors are able to get fixed NTK in the lazy regime without the assumption of large N in section 5. Is this specific to the linear network? Can you comment on how your results may scale to larger dimensional networks, and on the relation of your results and the lazy and rich regimes in the infinite width limit? \n5.\tIn Fig.4 b the simulation w1w1 and w2w2 labels seem to be reversed compared to the theory. \n6.\tIn Fig.5, how does the ratio between Nh and Ni (or Nh and No) play a role here?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1.\tThe paper is very well structured and the theoretical results are clearly written. \n2.\tIt is a novel and interesting idea to model the range of dynamics from lazy to rich with the balance parameter \\lambda. \n3.\tThe authors analyzed their theoretical results on interesting semantic task examples, and showed multiple implications of their theory on other learning paradigms.", "weaknesses": "1.\tThe technical and theoretical novelty of this work is limited, the method used in this paper has mostly been established in Braun et. al. 2022, this work extends to the \\lambda-balanced case, which is more general. However, the extension is achieved by enforcing a stricter assumption that the input is whitened, limiting the applicability of the results. \n2.\tThis paper makes several very strong assumptions such as task-aligned initialization, whitened inputs, linear networks, etc. It would be nice if the authors can show even some empirical evidence that these assumptions may be somewhat relaxed while the main results still qualitatively hold. \n3.\tThe applications are only briefly discussed. I think the paper would benefit from making the theoretical results in section 4 more concise and moving some of the application results (transfer learning for example I think is very interesting) to main text.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EVuxknZhAI", "reviewer": "Reviewer_5Z1o", "summary": "This study derives exact solutions (Theorem 4.3) to examine how initialization --- particularly the relative scale of weights across layers --- influences rich versus lazy learning in linear feedforward neural networks. The study relaxes several assumptions from prior work, notably the balanced assumption in Braun et al., which is a special case in here when $\\lambda=0$. Extending this, the authors show in Theorem C.4 that the network transitions into a rich regime as $\\lambda$ approaches 0, and into a lazy regime as $\\lambda$ approaches infinity. Overall, this work significantly contributes to the ongoing investigation into how various network parameters influence rich versus lazy learning regimes, with demonstrated relevance to both neuroscience and deep learning.", "questions": "•\tIt would be interesting to understand more how this relative scale interact with other knobs that impact learning regimes. For instance, could you provide more discussion on how the relative scale and absolute scale of the weights interact? See the weakness point above. \n\n•\tCould you offer more intuition on how \\(\\lambda \\to \\infty\\) results in an identity representation (i.e., a task-agnostic representation)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "•\tThis work is novel and timely.\n\n•\tStrong theoretical foundation: the approach extends a previous theoretical framework (Braun et al., 2022) to derive exact solutions for NTK, representation similarity, gradient flow, and loss. This allows for precise determination of how initialization (specifically, the $\\lambda$ parameter) influences representation and learning dynamics under the given assumptions. However, I did not have the bandwidth to verify the proofs in the appendix, which affects my confidence score.\n\n•\tInsightful experiments demonstrating how $\\lambda$ interpolates between learning regimes, with applications to both neuroscience and machine learning.\n\n•\tThe manuscript is well-written, with a comprehensive literature review.", "weaknesses": "•\tThis work relies on a list of assumptions and focuses on a simple two-layer linear feedforward network, which deviates from real-world settings. However, this didn’t significantly impact my score, as the assumptions are already more relaxed compared to previous works, and the authors have adequately addressed these limitations in the discussion.\n\n•\tI wish there were more intuitive explanations of how \\(\\lambda\\) interpolates between learning regimes. The authors provide some results in this direction, such as Theorem 5.1, which shows that under different limits of \\(\\lambda\\), the transition function converges to sigmoidal (\\(\\lambda \\to 0\\)) and exponential (\\(\\lambda \\to \\pm \\infty\\)). However, it would be helpful to have more insight into how the limit of \\(\\lambda \\to \\pm \\infty\\) leads to a task-agnostic identity representation. \n\n•\tI’m not sure if I have a good understanding of how the relative and absolute scales interact. While the authors demonstrated that in the extreme case where \\(\\lambda \\to 0\\), the rich regime occurs regardless of the absolute scale, I’m less certain about the case where \\(\\lambda \\to \\infty\\), as the weight norm of one of the layers would presumably become very large. Large initialization can lead to lazy learning, so I’m unsure if it’s possible to decouple the contributions of \\(\\lambda\\)-balanced initialization from large weight initialization. I believe some control experiments with large \\(\\lambda\\) values using random (unbalanced) initialization but with the same \\(W_1\\) and \\(W_2\\) norms as the balanced case would help distinguish the two.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)