Exact learning dynamics of deep linear networks with prior knowledge
Authors: Lukas Braun
Clémentine Carla Juliette Dominé
James E Fitzgerald
Andrew M Saxe
Venue: NeurIPS 2022 Accept (2022)

Abstract: Learning in deep neural networks is known to depend critically on the knowledge embedded in the initial network weights. However, few theoretical results have precisely linked prior knowledge to learning dynamics. Here we derive exact solutions to the dynamics of learning with rich prior knowledge in deep linear networks by generalising Fukumizu's matrix Riccati solution \citep{fukumizu1998effect}. We obtain explicit expressions for the evolving network function, hidden representational similarity, and neural tangent kernel over training for a broad class of initialisations and tasks. The expressions reveal a class of task-independent initialisations that radically alter learning dynamics from slow non-linear dynamics to fast exponential trajectories while converging to a global optimum with identical representational similarity, dissociating learning trajectories from the structure of initial internal representations. We characterise how network weights dynamically align with task structure, rigorously justifying why previous solutions successfully described learning from small initial weights without incorporating their fine-scale structure. Finally, we discuss the implications of these findings for continual learning, reversal learning and learning of structured knowledge. Taken together, our results provide a mathematical toolkit for understanding the impact of prior knowledge on deep learning.

Exact learning dynamics of deep linear networks with prior knowledge
Lukas Braun t.1 lukas. braunOpsy. ox ac uk
Clémentine C. Dominé f.2 cl clementine. domine. 200uc1. ac uk
James E. Fitzgerald 3 fitzgeraldj@jamell thmi.org
Andrew M. Saxe 2,1,5
a. sax.euxlc.......... acc uk
Abstract
Learning deep neural networks is known depend eritically on the knowledge embedded in the initial network weights. However, few theoretical results have precisely linked prior knowledge to learning dynamics, Here we derive exact solutions to the dynamics of learning with rich prior knowledge in deep linear networks by generalising Fukumizu's 's matrix Riccati solution [ We obtain explicit expressions for the evolving network function, hidden representational similarity, and neural tangent kernel over training for a broad class of initialisations and tasks. The expressions reveal class of task-independent initialisations that radically alter learning dynamics from slow non non-linear dynamics to fast exponential trajectories while converging to global optimum with identical representational similarity, dissociating learning trajectories from the structure of initial internal representations. We characterise how network weights dynamically align with task structure, rigorously justifying why previous solutions successfully described learning from small initial weights without incorporating their fine-scale structure. Finally, we discuss the implications of these findings for continual learning, reversal learning and learning of structured knowledge. Taken together, our results provide a mathematical toolkit for understanding the impact of prior knowledge on deep learning.
Introduction
hallmark of human learning is our exquisite sensitivity to prior knowledge: what we already know affects how subsequently learn (2). For instance, having learned about the attributes of nine animals, we may learn about the tenth more quickly [3] [ 6). In machine learning, the impact of prior knowledge
learning is evident in range of paradigms including reversal learning [ 10 [T] continual learning [12 [3] [4]. curriculum learning [3]. and meta learning [16) One form of prior knowledge in deep networks is the initial network state, which is known to strongly impact learning dynamics [17] 18 19]. Even random initial weights of different variance can yield qualitative shifts in network behaviour between the lazy and rich regimes (20), imparting distinct inductive biases on the learning process. More broadly, rich representations such those obtained through pretraining provide empirically fertile inductive biases for subsequent
transfer learning 8
9
First authors, random order Department of Experimental Psychology, University of Oxford, Oxford, United Kingdom
Galsby Computational Neuroscience Unit, University College London, London, United Kingdom
3. Howard Hughes Medical Institute, Janelia Research Campus, Ashburn, USA
4. Sainsbury Wellcome Centre, University College London, London, United Kingdom
5. CIFAR Azrieli Global Scholar, CIFAR, Toronto, Canada
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

New Network
Training Data
Gradient Descent
h E Rt%
C
D
Network Output
3%
Training Steps
Training Steps
ning Steps
Figure 1: Learning with prior knowledge. In our setting, a deep linear network with N; input, N₁, hidden and N. output neurons is trained from particular initialisation using gradient descent. B-D Network output for an example task over training time when starting from B small random weights, C large random weights, and D the weights of previously learned task. The dynamics depend in detail on the initialisation. Solid lines indicate simulations, dotted lines indicate the analytical solutions we derive in this work.
fine-tuning (11) Yet while the importance of prior knowledge learning is clear, theoretical understanding remains limited, and fundamental questions remain about the implicit inductive biases of neural networks trained from structured initial weights. A better understanding of the impact of initialisation on gradient- based learning may lead to improved pretraining schemes and illuminate pathologies like catastrophic forgetting in continual learning 22
Here, we address this gap by deriving exact solutions to the dynamics of learning in deep linear networks function of network initialisation, revealing intricate and systematic dependence. We consider the setting depicted in Fig. IA, where a network is trained with standard gradient descent from potentially complex initialisation. When trained the same task, different initialisations can radically change the network's ' learning trajectory (Fig. B-D). Our approach, based on a matrix Riccati formalism (1]. provides explicit analytical expressions for the network output time (Fig. B dotted). While simple, deep linear networks have a non-convex loss landscape and have been shown recapitulate several features of nonlinear deep networks while retaining mathematical tractability.
1.1 Contributions
We derive an explicit solution for the gradient flow of the network function, internal repre- sentational similarity, and finite-width neural tangent kernel of over- and under-complete two-layer deep linear networks for a rich class of initial conditions (Section 3).
We characterise a set of random initial network states that exhibit fast, exponential learning dynamics and yet converge to rich neural representations. Dissociating fast and slow learning dynamics from the rich and lazy learning regimes (Section 4
We analyse how weights dynamically align to task-relevant structure over the course of learning, going beyond prior work that has assumed initial alignment (Section 5).
We provide exact solutions to continual learning dynamics, reversal learning dynamics and
to the dynamics of learning and revising structured representations (Section 6
1.2
Related work
Our work builds on analyses of deep linear networks 23
which have shown that this simple model nevertheless has intricate fixed point structure and nonlinear learning dynamics reminiscent phenomena seen nonlinear networks. A variety of works has analysed convergence
28 (29,, generalisation [10, 30 31, and the implicit bias of gradient descent [2 33 34 [3). These works mostly considers the labula rusa case of small initial random weights, for which exact solutions are known [17]. By contrast our formalism describes dynamics from a much larger class of initial conditions and can describe alignment dynamics that do not occur the tabula rasa setting. Most directly, our results build from the matrix Riccati formulation proposed by [. Connecting this

---

formulation and matrix factorisation problems yields better characterisation of the convergence rate [27]. We extend and refine the matrix Riccati result to obtain the dynamics of over- and under- complete networks; to obtain numerically stable forms of the matrix equations; and to more explicitly reveal the impact of initialisation.
A line of theoretical research has considered online learning dynamics in teacher-student settings [36 37 38] deriving ordinary differential equations for the average learning dynamies even in nonlinear networks. However, solving these equations requires numerical integration. By contrast, our approach provides explicit analytical solutions for the more restricted case of deep linear networks.
Other approaches for analysing deep network dynamics include the Neural Tangent Kernel (NTK) describe nonlinear
While the former networks but not the learning dynamics of hidden representations, the later yields a description of representation learning dynamics in wide networks in terms of partial differential equation. Our work is similar in seeking a subset of more tractable models that are amenable to analysis, but we focus on the impact of initialisation representation learning dynamics and explicit solutions.
39
40 41 and the mean field approach 42 43
large body work has investigated the effect different random initialisations on learning in deep networks. The role of initialisation in the vanishing gradient problem and proposals for better initialisation schemes have been illuminated by several works drawing on the central limit theorem [45] [17 [46 [8] 47). reviewed in [4 19 49]. These approaches typically guarantee that gradients do not vanish at the start of learning, but do not analytically describe the resulting learning trajectories. Influential work has shown that network initialisation variance mediates transition from rich representation learning to lazy NTK dynamics (20]. which we analyse our framework.
Preliminaries and setting
Consider supervised learning task in which input vectors x,, R.N. from set of P training pairs Yn),, have to be associated with their target output vectors Ya E RRN, We learn this task with two-layer linear network model (Fig. [TA) that produces the output prediction
y W Wix,,
(1)
with weight matrices W1 R. and W₂ R,*,,, where N₂. is the number of hidden units. The network's s weights are optimised using full batch gradient descent with learning rate v) (or respectively time constant 1/1) on the mean squared error loss
((g.y) = (1IS - y
where (-) denotes the average over the dataset. The input and input-output correlation matrices the dataset are
1 P
1 P
S Σ X E N.M. and
S
vnXI E N..N,
(3)
Finally, the gradient optimisation starts from an initialisation W2(0), (O).). Our goal is to under- stand the full time trajectory the network's 's output and internal representations function of this initialisation and the task statistics.
Our starting point is the seminal work of Fukumizu M. which showed that the gradient flow dynamics could be written as a matrix Riccati equation with known solution. In particular, defining
W W2
E
and F
(4)
Euz
the continuous time dynamics of the matrix QQ" from initial state Q(0) is
QQ (+) c (0) I+
1 Q(0) (f
-FFFFF -F FF1) Q(0)
(5)
the following four assumptions hold: Assumption 2.1 The dimensions of the input and target vectors identical, that is N; No.
3

---

Assumption 2.2 The input data is whitened. that is Σ.d = I.
Assumption 2.3 The network's weight matrices are cero-balanced at the beginning of training. W2(0)"W2(0)) Wz(0). If this condition holds at initialisation, it will persist
that is W. W1(0)W1(0))
throughout training 17 24
Assumption 2.4 The input-output correlation of the task and the initial state of the network funttoou have full rank, that is rank (ank(())= rank W (0) = N; No. This implies that the network not bottlenecked, ie. N₁, 2 min(N, N.)-
For completeness, we include a derivation of this solution in Appendix
Rather than tracking the weights' dynamics directly, this approach tracks several key statistics collected in the matrix
W (t) W/
QQ" T
(6)
W (t)
W
which can be separated into four quadrants with intuitive meaning: the off-diagonal blocks contain the network function
(()) W₂ Wa (t)X,
while the on-diagonal blocks contain the correlation structure of the weight matrices. These permit calculation of the temporal evolution of the network's 's internal representations including the task- relevant representational similarity matrices (RSM) [50), i.c. the kernel matrix o(x)² 6(x'), of the neural representations in the hidden layer
RSM, x w/'wTw WTW, (t) RSM Y (W, (T)
where + denotes the pseudoinverse; and the network's finite-width neural langent kernel
40
NTK = IN. @ X'W? (t) W (t) @ x'x,
where is the identity matrix and @ is the Kronecker product. For a derivation of these quantities see Appendix B Hence, the solution in Equation deseribes important aspects of network behaviour.
However, this form, the solution has several limitations. First, relies on general matrix exponen- tials and inverses, which are barrier to explicit understanding. Second, when evaluated numerically, is often unstable. And third, the equation only valid for equal input and output dimensions. In the following section we address these limitations.
Implementation and simulation Simulation details are in Appendix H Code replicate all simulations and plots are available online" under a GPLv3 license and requires <6 hours to execute on single AMD Ryzen 5950x.
3 Exact learning dynamics with prior knowledge
In this section we derive an exact and numerically stable solution for QQ that better reveals the learning dynamics, convergence behaviour and generalisation properties of two-layer linear networks with prior knowledge. Further, we alter the equations to be applicable to equal and unequal input and output dimensions, overcoming Assumption 2.1).
To place the solution in a more explicit form, we make use the compact singular value decom- position. Let the compact singular value decomposition of the initial network function and the input-output correlation of the task be
(0) (0)) USVT and SVD(E"*) ÜSVT.
(10)
Here, U and Ü e N, N... denote the left singular vectors, S and IR N.. x N... the square matrix with ordered, non-zero eigenvalues on its diagonal and V and V E N, the corresponding right singular vectors. For unequal input-output dimensions (N; * N.) the right and left singular vectors are therefore not generally square and orthonormal. Accordingly, for the case N; > No. we define Ü, R, (N.-NN. as a matrix containing orthogonal column vectors that complete the basis, i.e., make [Ū Ü orthonormal. Conversely, we define V_ E (NN, for the case of N; > No.
https://gghuu..cc///sccaabbeee networks www.:wtttthhh--tooeeeeeee knowledee

---

A
1.0
1.0-
1.0
100-
WWW-yo
W/W,t [t)
WsW11t)
NTK(C)
10-
Loss
0.5
0.0
MSE(2,
10
-0.5
0.0
-0.5
10-14
D
0.0005
n
1.0
1.0
Loss
0.5
WSW6
w(W,tt
0.5
WsWfttt
0.5
NTK(()
0.7
-0.5
0.0
0.4
0.0
00.5
10 10 10
10a
10*
10a
10a
10
10
10" Training Steps
10a
Training Steps
Training Steps
Training Steps
Training Steps
Training Steps
Figure 2: Exact learning dynamics A The temporal dynamics of the numerical simulation (coloured lines) of the loss, network function, correlation of input and output weights and the NTK (columns 1-5 respectively) are exactly matched by the analytical solution (black dotted lines) for small initial weight values and B large initial weight values. C Each line shows the deviation of the analytical loss C from the numerical loss L for one of n 50 networks with random architecture and training data (details in Appendix H across range of learning rates r/ E [0.05, 0.0005). The deviation mutually decreases with the learning rate. D Numerical and analytical learning curves for five randomly sampled example networks (coloured x in C).
Assumption 3.1 Define B Ü VTV and C U-Ú VTV. B is non-singular.
Theorem 3.1 Under the assumptions of whitened inputs, 2.2 3.1 the temporal dynamics of QQ are
zero-balanced weights 2.3 full rank
2.4 and B non-singular!
QQ"(1) Z
$ B T-'-s-1
(B"))'e ---S-BBB)
3 I
25 8-1
B
25 -
S- (B -
(B")
e
(11)
+ + B - U
ZT
with
+2 - V I $ Cz
(12)
ú I+e CT (B") 2Ú ÚHU(")) -te
For proof of Theorem 3.1 please refer to Appendix
With this solution and NTK (Fig. ZA, B). As the solution contains only negative exponentials, it is numerically stable and provides high precision across wide range of learning rates and network architectures (Fig. Zc. D).
calculate the exact temporal dynamics of the loss, network function, RSMs
note that solution for the weights (1) and W2()), i.e., Q(1), can derived up to a time varying orthogonal transformation demonstrated in Appendix C Further, time-dependent variables only occur in matrix exponentials of diagonal matrices of negative sign, the network approaches steady state solution.
functrrfuncttinn converges to the global minimum USVT and acquires a rich task-specific internal representation, that WT W, = VSVT and W₂ WT ÜSUT
Theorem 3.2 Under the assumptions Theorem 3.1 the network function
The proof of Theorem 3.2 is in Appendix C We now turn several implications of these results.
4 Rich and lazy learning regimes and generalisation
Recent results have shown that large deep networks can operate in qualitatively distinct regimes that depend on their weight initialisations 20 5II, the so called rich and lazy regimes. In the rich regime, learning dynamics can be highly nonlinear and lead to task-specific solutions thought
5

---

to plants. Likewise, the representational similarity of the task's target values (ÜSUT) reveals the for example predict Gaussian Small
A
B
C
Goussian
E
Balanced
Imput
0.8
Weight Initialisaiin
Loss
Gaussinn Large
Fish
Bird Algan
Flower
VSVV
-0.
Balanced Large
1-04 0.0
Input
Oucput
Input
Oucput
Input
Output
Training Stops 2000
Figure 3: Rich and lazy learning. Semantic learning task, B SVD of the input-output correlation of the task (top) and the respective RSMs (bottom). Rows and columns in the SVD and RSMs are identically ordered as the order of items in the hicrarchical tree. C Final QQ" matrices after training converged when initialised from random small weights, D random large weights (note how the upper left and lower right quadrant differ from the task's RSMs) and E large zero-balanced weights. Learning curves for the three different initialisations as in C (green), D (pink) and E (blue). While both large weight initialisations lead to fast exponential learning curves, the small weight initialisation leads to slow step-like decay of the loss.
51]. By contrast, the lazy regime exhibits simple exponential learning dynamics and exploits nigh-dimensional nonlinear projections of the data produced by the initial random weights, leading to task-agnostic representations that attain zero training error but possibly lower generalisation performance [39 40 41). Traditionally, the rich and lazy learning regimes have been respectively linked to low and high variance initial weights (relative to the network layer size).
to lead to favourable generalisation properties 20 25
To illustrate these phenomena, we consider a semantic learning task in which a set of living things
have to be linked to their position in a hierarchical structure (Fig. [TA) [17]. The representational
similarity of the input of the task (VSVT) reveals its inherent structure (Fig.) BB). For example, the
representations of the two fishes are most similar to each other, less similar to birds and least similar
primary groups among which items are organised. As consequence, from an object being a fish that it is an animal and from an object being a plant that it is not a bird. Reflecting these structural relationships in internal representations can allow the rich regime to generalise in ways the lazy regime cannot. Crucially, QQ (4) contains the temporal dynamics of the weights' representational similarity and therefore can be used study a network finds a rich or lazy solution.
When training two layer network from random small initial weights, the weights' input and output RSM (Fig.Br BC, upper left and lower right quadrant) are identical to the task's 's structure at convergence. However, when training from large initial weights, the RSM reveals that the network has converged to a lazy solution (Fig.BP), We emphasise that the network function in both cases is identical (Fig. D, lower left quadrant). And while their final loss is identical too, their learning dynamics evolve slow and step- wise in the case of small initial weights and fast and exponentially in the case of large initial weights (Fig. 3 F) predicted by previous work 20
However, from Theorem 3.2 it directly follows that our setup is guaranteed to find a rich solution in which the weights' RSM is identical to the task's RSM, i.c., W w, = VSVT and W₂ W = ÜSÜT, Therefore, as zero-balanced weights may be large, there exist initial states that converge to rich solutions while evolving as rapid exponential learning curves (Fig. 3 F). Crucially, these initialisations are task-agnostic. in the sense that they are independent of the task structure (cf. (52)). This finding applies any learning task with well defined input-output correlation. For additional simulations see Appendix D Hence our equation can describe the change in dynamics from step-like to exponential with increasing weight scale, and separate this dynamical phenomenon from the structure of internal representations.
5 Decoupling dynamics
The learning dynamics of deep linear networks depend on the exact initial values of the synaptic weights. Previous solutions studied learning dynamics under the assumption that initial network weights are "decoupled", such that the initial state the network and the task share the same singular vectors, i.e. that U- = Ü and V = V 17.. Intuitively, this assumption means that there is no cross-coupling between different singular modes, such that each evolves independently. However, this

---

A
D
Small
Intermediate
Large
W2W- (C)
A'ALO)
Diggnel
400
300000
100000
40000
Training steps
Diagoral
Off Diagonal
B
C
91
US
5 1
1
F
A'A(C)
34-
USV butioo
0
40000
10000
10300
Training steps
Training steps
Training steps
Figure 4: Decoupling dynamics. A Analytical (black dotted lines) and numerical (solid lines) of the temporal dynamies of the on- and off-diagonal elements of A A in blue and red, respectively, B Schematic representation of the decoupling process. C Three target matrices with dense, unequal diagonal, and equal diagonal structure. D-F Decoupling dynamics for the top (D), middle (E), and bottom (F) tasks depicted in panel C. Row F contains analytical predictions for the time of the peak of the off-diagonal (dashed green). The network is initialised defined in E with small, intermediate and large variance.
assumption is violated in most real-world scenarios. As a consequence, most prior work has relied on the empirical observation that learning from iabula rasa small initial weights occurs two phases: First, the network's input-output map rapidly decouples; then subsequently, independent singular modes are learned in this decoupled regime. Because this decoupling process is fast when training begins from small initial weights, the learning dynamics are still approximately described by the temporal learning dynamics of the singular values assuming decoupling from the start. This dynamic has been called silent alignment process [26). Here we leverage our matrix Riccati approach to analytically study the dynamics this decoupling process. We begin by deriving an alternate form of the exact solution that eases the analysis.
Theorem 5.1 Let the weight matrices of two layer linear network be initialised by W1 A and W, (O)T", where (0) E & xN; is an arbitrary, invertible matrix. Then, under the assumptions of equal input-ouipul dimensions 2.1 whitened inputs [2.2 zero-balanced weights 2.3
and rank 2.4 the temporal dynamics of QQ are fully determined by
AT
88
8 (A(O)) A(0) -5 -I-
25 S 1-11
(13)
For proof of Theorem 5.1 please refer to Appendix We remark that this form is less general than that in Theorem 3.1 and in particular implies UV ÜV. Here the matrix A represents the dynamics directly the SVD basis the task. Off-diagonal elements represent counterproductive coupling between different singular modes (for instance, [A" A]21 is the strength of connection from input singular vector to output singular vector 2, which must approach zero to perform the task perfectly), while on-diagonal elements represent the coupling within the same mode (for instance, [A L Al is the strength of connection from input singular vector to output singular vector 1, which must approach the associated task singular value perform the task perfectly). Hence the decoupling process can studied by examining the dynamics by which A 2 A becomes approximately diagonal.
The outer inverse in Equation 13 renders it difficult to study high dimensional networks analytically. Therefore, we focus on small networks with input and output dimension N; and N, for which lengthy but explicit analytical solution is given in Appendix E In this setting, the structure of the weight initialisation and task are encoded the matrices
A TOO (()
Ta (0)
b(0)
and $
81
0
(14)
b(0)
a2(()]
0
82
7

---

A
Task
Task 2
T
B
D
0.2
0.25
0.6
-.2 0.0
F Forgeting
5-
0.0
Numerical Forgetting
0.00
[|W []] w/Wyll+
-0.2
34
Sergetting
2.5
2-
5
0.6
0-
0.5
10-
0.0
0
5000
10000
15000
20000
3
Training Steps
Analytical
After Task
After Task
After Task
Figure 5: Continual learning. A Top: Network training from small zero-balanced weights on sequence of tasks (coloured lines show simulation and black dotted lines analytical results). Bottom: Evaluation loss for tasks of the sequence (dotted) while training on the current task (solid). As the network function is optimised on the current task, the loss of other tasks increases. B Comparison of the numerical and analytical amount of catastrophic forgetting on first task after training on second task for n = 50 linear (red), tanh (blue) and ReLU (green) networks. C Weight aligmment before and after training on a sequence of two tasks for n = 50 networks in linear (red), tanh (blue) and ReL.U (green) networks. Shaded area shows + std. D Evaluation loss for each of 5 tasks during training a linear (red), tanh (blue) and ReLU (green) network. Same data es in D but evaluated as relative change (i.e. amount of catastrophic forgetting). The top half of each square shows the pre-computed analytical amount of forgetting and the bottom half the numerical value.
where the parameters (0) and a2 represent the component of the initialisation that is aligned with the task, and ((0) represents cross-coupling, such that taking b(0) 0 recovers previously known and more restricted solutions for the decoupled case [17]. We use this setting to demonstrate two features of the learning dynamics.
Decoupling dynamics. First, we track decoupling by considering the dynamics of the off-diagonal element b(t) (Fig. 4 red lines). At convergence, the off-diagonal element shrinks zero as shown in Appendix E However, strikingly, b(t) can exhibit non- monotonic trajectories with transient peaks or valleys partway through the learning process. In particular, in Appendix E we derive the time In e(8-a)-a2)
of the peak magnitude as Epeak = 4s approximately with the time at which the on-diagonal element is half learned. If initialised from small random weights, the off dff-oiaaonal remains near-zero throughout learning, reminiscent of the silent alignment effect (26). For large initialisations, no peak is observed and the dynamics are exponential. At intermediate initialisations, the maximum of the off-diagonal reached before the singular mode is fully learned (Appendix (E). Intuitively, particular input singular vector can initially project appreciably onto the wrong output singular vector, corresponding to initial misalignment. This is only revealed when this link is amplified, which point corrective dynamics remove the counter- productive coupling, as schematised in Fig. 4B. We report further measurements of decoupling in Appendix B
2102 o102-b(0)² (Fig. 4F green dotted line), which coincides
Effect of initialisation variance. Next, we revisit the impact of initialisation scale for the on-diagonal dynamics, As shown in Fig. 4D-F, as the initialisation variance grows the learning dynamics change from sigmoidal to exponential, possibly displaying more complex behavior at intermediate variance (Appendix E). In this simple setting we can analyse this transition in detail. Taking 81 = = 8 as in Fig. F and [a (0)), a2 (0)), [b(0) <
we recover a sigmoidal trajectory,
801 (0)
a1
(15)
-2st an (0) - a2())) + a (0) + (2
while for |a (0)L az (0)), [b0 > the dynamics of the on-diagonal element a close to expo- nential (Fig. 4D F--Faa left and right columns). We examine larger networks in Appendix B
8

---

A
H
Arimal
Loss
0.2
W.W. W,
Dim 2
0.0
Flah
and Algao
Flower
2
-2-
B
16a
10
2000
2000
-2.5
0.0
2.5
255
0.0
2.5
2-
2
0.0
Primary
Pianary
Pamary
Prinsry
Loss
W W,
Dim 2
0
0
0.0
10 10
10'
500
0.0
255
-2
0.0
255
Training
Training Step
Training Step
Figure 6: Reversal learning and revising structured knowledge. Scale of x-axis varies in top and bottom rows. A Analytical (black dotted) and numerical (solid) learning dynamics of reversal learning task. The analytical solution gets stuck on a saddle point, whereas the numerical simulation escapes the saddle point and converges to the target. B In a shallow network, training on the same task as in A converges analytically (black dotted) and numerically (solid). C Semantic learning tasks. Revised living kingdom (top) and colour hierarchy (bottom). D SVD of the input-output coreclation of the tasks and respective RSMs. E Analytical (black dotted) and simulation (solid) loss and learning dynamics of first training on the living kingdom (Fig. BAA and subsequently on the respective task in C. The analytical solution fails for the revised animal kingdom as it gets stuck in saddle point, while the simulation escapes the saddle (top, green circle). Initial training on the living kingdom task from large initial weights and subsequent training on the colour hierarchy have similar convergence times (bottom) G Multidimensional scaling (MDS) of the network function for initial training on the living kingdom task from small (top) and large initial weights (bottom). Note how despite the seemingly chaotic learning dynamics when starting form large initial weights, both simulations learn the same representation. H MDS of subsequent training on the respective task in C.
6 Applications
The solutions derived in Sections 3 and 3 provide tools to examine the impact of prior knowledge on dynamics in deep linear networks. So far we have traced general features of the behaviour of these solutions. In this section, use this toolkit to develop accounts of several specific phenomena.
Continual Learning Continual learning (see 12] for review) and the pathology of catastrophic
forgetting have long been a challenge for neural network models (22 [53 54. variety of theoretical work has investigated aspects of continual learning (55] 56 57 58, 59 In this setting, starting from an initial set of weights, a network trained on sequence of tasks with respective input-output correlations Ti = ΣNT To Σ. Ta ΣN As shown in Fig. 5/, our dynamics immediately enable exact solutions for the full continual learning process, whereby the final state after training on one task becomes the initial network state for the next task. These solutions thus reveal the exact time course of forgetting for arbitrary sequences of tasks.
Training on later tasks can overwrite previously learned knowledge, a phenomenon known as catastrophic forgetting [22] 53) [54]. From Theorem 3.2 it follows that from any arbitrary zero- balanced initialisation] 2.3 the network converges the global optimum such that the initialisation
is completely overwritten and forgetting is truly catastrophic. In particular, the loss of any other
task T7 after training to convergence on task T, is (T,) 1/2 1V + c, where is constant that only depends on training data of task T, (Appendix F). As a consequence, the amount of forgetting, i.c. the relative change of loss, is fully determined by the similarity structure of the tasks and thus can be fully determined for a sequence of tasks before the onset of training (Fig. 5 Appendix FD. For example, the amount of catastrophic forgetting in task T,, when training on task Ta after having trained the network on task Tú is La (Te) - L. (Ta). As expected, our results depend on our linear setting and tanh or ReLU nonlinearities can show different behaviour, typically increasing the amount of forgetting (Fig. (.E. Further, in nonlinear networks, weights become rapidly unbalanced and forgetting values that are calculated before the onset of training do not predict the actual outcome (Fig. BBB). In summary, our results link exact learning dynamics with catastrophic forgetting and thus provide an analytical tool to study the mechanisms and potential counter measures underlying catastrophic forgetting.

---

Reversal learning During reversal learning, pre-existing knowledge has to be relearned, overcom- ing previously learned relationship between inputs and outputs. For example, reversal learning occurs when items of class are mislaheled and later corrected. We show analytically, that reversal learning in fact does not succeed in deep linear networks (Appendix (G). The pre-existing knowledge lies exactly on the separatrix of saddle point causing the learning dynamics to converge to zero (Fig. [6A). In contrast, the learning still succeeds numerically, as any noise will perturb the dynamics off the saddle point, allowing learning to proceed (Fig. (A)) However, the dynamics still slow in the vicinity of the saddle point, providing a theoretical explanation for catastrophic slowing in deep linear networks [60]. We note that the analytical solution requires adaptation of Theorem 3.1 B is generally not invertible in the case of reversal learning (Appendix G). Further, as is revealed by the exact learning dynamics (Appendix G). shallow networks do succeed without exhibiting catastrophic slowing during reversal learning (Fig. 6)).
Revising structured knowledge Knowledge is often organised within underlying. shared structure, of which many can be learned and represented in deep linear networks [25]. For example, spatial locations can be related to each other using the same cardinal directions, varying semantic knowledge can be organised using the same hierarchical tree. Here, we investigate if deep linear networks benefit from shared underlying structure. To this end, a network is first trained on the three-level hierarchical tree of Section [ (eight items of the living kingdom, each with a set of eight associated features). and subsequently trained on revised version of the hierarchy. The revised task varies the relation of inputs and outputs while keeping the same underlying tree structure. If the revision involves swapping two neighbouring nodes on any level of the hierarchy, e.g. the identity of the two fish on the lowest level of the hierarchy (Fig. 6C, top), the task is identical to reversal learning, leading to catastrophically slowed dynamics (Fi.. top). When training the network on a new hierarchical tree with identical items but a new set of features, like colour hierarchy (Fig. [ bottom), there is no speed advantage in comparison to a random initialisation with similar initial variance (Fig. 66-F, bottom). Importantly, from Theorem 3.2 it follows, that the learning process can be sped up significantly by initialising from large bbalanedeee weights, while converging to a global minimum with identical generalisation properties as when training from small weights (Fig. (fG-- In summary, having incorporated structured knowledge before revision does not speed up or even slows down learning in comparison to learning from random acro-balanced weights. Notably, that is despite the tasks' structure being almost identical (Fig. 3B and Fig. 6D).
Discussion
We derive exact solutions to the dynamics of learning with rich prior knowledge in tractable model class: deep linear networks. While our results broaden the class of two-layer linear network problems that can be described analytically, they remain limited and rely on set of assumptions 2.1 In particular, weakening the requirement that the input covariance be white and the weights be zero-balanced would enable analysis of the impact of initialisation on internal representations. Nevertheless, these solutions reveal several insights into network behaviour. We show that there exists large set of initial values, namely zero-balanced weights 2.3] which lead to task task-specific representations; and that large initialisations lead to exponential rather than sigmoidal learning curves. We hope our results provide a mathematical toolkit that illuminates the complex impact of prior knowledge on deep learning dynamics.
Acknowledgments and Disclosure of Funding
L.B. was supported by the Woodward Scholarship awarded by Wadham College, Oxford and the Medical Research Council [MR/N013468/1], C.D. and A.S. were supported by the Gatsby Charitable Foundation (GAT3755). Further, A.S. was supported by a Sir Henry Dale Fellowship from the Wellcome Trust and Royal Society (216386/Z/19/2) and the Sainsbury Wellcome Centre Core Grant (219627/2/1/ZZ) A.S. is a CIFAR Azrieli Global Scholar in the Learning in Machines & Brains program. J.F. was supported by the Howard Hughes Medical Institute.
10

---

References
[1] Kenji Fukumizu. Effect of batch learning in multilayer neural networks. Gen, 10-)): 1998.
[2] S.E. Carey. Conceptual Change In Childhood. MIT Press, Cambridge, MA, 1985.
[3] James L McClelland, Bruce L McNaughton, and Randall C C'RReilly Why there are comple- mentary learning systems in the hippocampus and neocortex: insights from the successes and failures connectionist models of learning and memory. Psychological review, 102(3):419, 1995.
[4] Gregory Murphy. The big book of concepts. MIT press, 2004.
[5] James L McClelland. Incorporating rapid neocortical learning of new schema-consistent information into complementary learning systems theory. Journal of Experimental Psychology: General, 142(4): 1190, 2013.
[6] Timo Flesch, Jan Balaguer, Ronald Dekker, Hamed Nili, and Christopher Summerfield. Com- paring continual task learning in minds and machines. Proceedings of the National Academy of Sciences, 1
2018.
[7] Burak Erdeniz and Nart Bedin Atalay. Simulating probability learning and probabílistic reversal learning using the attention- gated reinforcement learning (agrel) model. In The 2010 International Joint Conference on Neural Networks (IJCNN), pages 1-6. IEEE, 2010.
[8] Sebastian Thrun and Lorien Pratt. Learning to learn. Springer Science & Business Media, 2012.
[9] Matthew E Taylor and Peter Stone. Transfer learning for reinforcement learning domains: A survey. Journal of Machine Learning Research, 10(7), 2009.
[10] Andrew K Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. arXiv preprint arXiv: 1809. 10374, 2018.
[11] Federica Gerace, Luca Saglietti, Stefano Sarao Mannelli, Andrew Saxe, and Lenka Zdeborová. Probing transfer learning with a model of synthetic correlated datasets. Machine Learning: Science and Technology, 2022.
[12] German I Parisi, Ronald Kemker, Jose L Part, Christopher Kanan, and Stefan Wermter, Continual lifelong learning with neural networks: review. Neural Networks, 113:54-71, 2019.
[13] James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Srrabska-Barwinnkee et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 1 114(13):35213352, 2017.
[14] Friedemann Zenke, Ben Poole, and Surya Ganguli. Continual learning through synaptic intelligence. In International Conference on Machine Learning, pages 3987-3995. PMLR, 2017.
[15] Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In Proceedings of the 26th annual international conference on machine learning, pages 41-48,
2009.
[16] Khurram Javed and Martha White. Meta-learning representations for continual learning. In Advances in Neural Information Processing Systems, pages 8830-1830, 2019.
[17] Andrew Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In 2nd International Conference Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Truck Proceedings, 2014.

---

[18] Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. Resurrecting the sigmoid in deep learning through dynamical isometry; theory and practice. Advances in neural information processing systems, 30, 2017.
[19] Yasaman Bahri, Jonathan Kadmon, Jeffrey Pennington, Sam S Schoenholz, Jascha Sohl- Dickstein, and Surya Ganguli. Statistical mechanics of deep learning. Annual Review of Condensed Matter Physics, 11:501-528, 2020.
[20] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. Advances in Neural Information Processing Systems, 32, 2019.
[21] Maithra Raghu, Chiyuan Zhang, Jon Kleinberg, and Samy Bengio. Transfusion: Understanding transfer learning for medical imaging. Advances in neural information processing systems, 32, 2019.
[22] Michael McCloskey and Neal J Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of learning and motivation, volume 24, pages 109-165. Elsevier, 1989.
[23] Pierre Baldi and Kurt Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural networks, 2(1):53-58, 1989.
[24] Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization.. In International Conference on Machine Learning, pages 244 253. PMLR, 2018.
[25] Andrew M Saxe, James L McClelland, and Surya Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 116 (23):11537-11546, 2019.
[26] Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. In International Conference on Learning Representations, 2022.
[27] Salma Tarmoun, Guilherme Franca, Benjamin D Haeffele, and Rene Vidal. Understanding the dynamics of gradient flow in overparameterized linear models. In International Conference on Machine Learning, pages 10153-10161. PMLR, 2021.
[28] Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. arXiv preprint (rXivv/810002220 2018.
[29] Simon Du and Wei Hu. Width provably matters in optimization for deep linear neural networks. In International Conference on Machine Learning, pages 1655-1664. PMLR, 2019.
[30] Tomaso Poggio, Qianli Liao, Brando Miranda, Andrzej Banburski, Xavier Boix, and Jack Hidary. Theory iib:: Generalization in deep networks. arXiv preprint arXiv: 11379, 2018
[31] Dongsung Huh. Curvature-corrected learning dynamics in deep neural networks. In Interna- tional Conference on Machine Learning, pages 4552-4560. PMLR 2020.
[32] Thomas Laurent and James Brecht. Deep linear networks with arbitrary loss: All local minima are global. In International conference on machine learning, pages 2902-2907. PMLR, 2018.
[33] Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. arXiv preprint arXiv: 181000001. 2018
[34] Suriya Gunasckar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent linear convolutional networks. Advances in Neural Information Processing Systems, 31, 2018.
[35] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
[36] Michael Biehl and Holm Schwarze. Learning by on-line gradient descent. Journal of Physics A: Mathematical and general, 28(3):643, 1995.
12

---

[37] David Saad and Sara A Solla. Exact solution for on-line learning in multilayer neural networks. Physical Review Letters, 74(21):4337, 1995.
[38] Sebastian Goldt, Madhu Advani, Andrew Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems, 32, 2019.
[39] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[40] Jaehoon 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.
[41] Sanjeev Arora, Simon S Du, Wei Hu, Zbiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances in Neural Information Processing Systems, 32, 2019.
[42] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33): E7665-E77671 2018.
[43] Grant Rotskoff and Eric Vanden-Eijnden. Parameters interacting particles: long time convergence and asymptotic error scaling of neural networks. Advances in neural information processing systems, 31, 2018.
[44] Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications, 130(3):1820-133) 2020.
[45] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedfor- ward neural networks. In Proceedings of the thirteenth international conference artificial intelligence and statistics, pages 249-25422 29--566.11 JMLR Workshop and Conference Proceedings, 2010.
[46] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification, In Proceedings of the IEEE international conference on computer vision, pages 1026-1034, 2015.
[47] Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein. Samuel Schoenholz, and Jeffrey Pen- nington. Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks. In International Conference on Machine Learning, pages 1393-5402 PMLR, 2018
[48] Giuseppe Carleo, Ignacio Cirac, Kyle Cranmer, Laurent Daudet, Maria Schuld, Naftali Tishby, Leslie Vogt-Maranto, and Lenka 7deborová. Machine learning and the physical sciences. Reviews of Modern Physics, 91(4):045002, 2019.
[49] Raman Arora, Sanjeev Arora, Joan Bruma, Nadav Cohen, Simon Du, Rong Ge, Suriya Gunasekar, Chi Jin, Jason Lee, Tengyu Ma, Benhnam Neyshabur, and Zhao Song. Theory of deep learning, 2020.
[50] Nikolaus Kriegeskorte, Marieke Mur, and Peter A Bandettini. Representational similarity analysis-connecting the branches of systems neuroscience. Frontiers in systems neuroscience, 2:4, 2008
[51] Timo Flesch, Keno Juechems, Tsvetomira Dumbalska, Andrew Saxe, and Christopher Summer- field. Orthogonal representations for robust context-dependent task performance in brains and neural networks. Neuron, 2022.
[52] Dmytro Míshkin and Jiri Matas. All you need is good init. arXiv preprint 1rXiv:1511.000 2015.
[53] Roger Ratcliff. Connectionist models of recognition memory: constraints imposed by learning and forgetting functions. Psychological review, 97(2):285, 1990.
13

---

[54] Robert M French. Catastrophic forgetting in connectionist networks. Trends in cognitive sciences, 3(4):128-1335 1999.
[55] Haruka Asanuma, Shiro Takagi, Yoshihiro Nagano, Yuki Yoshida, Yasuhiko Igarashi, and Masato Okada. Statistical mechanical analysis of catastrophic forgetting in continual learning with teacher and student networks. Journal of the Physical Society of Japan, 2021.
[56] Thang Doan, Mehdi Abbana Bennani, Bogdan Mazoure, Guillaume Rabusseau, and Pierre Alquier. theoretical analysis of catastrophic forgetting through the ntk overlap matrix. In International Conference on Artificial Intelligence and Statistics, pages 1072-1080. PMLR, 2021.
[57] Sebastian Lee, Sebastian Goldt, and Andrew Saxe. Continual learning the teacher-student setup: Impact of task similarity. In International Conference on Machine Learning, pages 6109-6119. PMLR, 2021.
[58] Gal Shachaf, Alon Brutzkus, and Amir Globerson. A theoretical analysis of fine-tuning with linear teachers. Advances in Neural Information Processing Systems, 34, 2021.
[59] Nilesh Tripuraneni, Michael Jordan, and Chi Jin. On the theory of transfer learning: The importance task diversity. Advances Neural Information Prucessing Systems, 33:7852- 7862, 2020.
[60] Sebastian Lee, Stefano Sarao Mannelli, Claudia Clopath, Sebastian Goldt, and Andrew Saxe. Maslow's hammer for catastrophic forgetting: Node re-use vs node activation. arXiv preprint arXiv:2205.09029 2022.
[61] Wei- Yong Yan, Uwe Helmke, and John B Moore. Global analysis of oja's flow for neural networks. IEEE Transactions on Neural Networks, 5(5): 683, 1994.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See section 7
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See section H
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See section
4. If you are using existing assets (c.g., code, data, models) or curating/relcasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
14

---

(d) Did you discuss whether and how consent was obtained from people whose data you' re using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
15


Reviews:
[{"review_id": "y7K9o35ioo", "reviewer": "Reviewer_rRZa", "summary": "This contribution introduces a reformulation of an existing matrix Riccati formulation of learning dynamics in linear networks.\nThis reformulation is numerically more stable and can be used to explain various known observations.\nThe analysis of the dynamics is used to enable rich learning with fast alignment using high variance weights.\nFinally, the insights provided by the reformulated dynamics are used to predict the learning dynamics and amount of forgetting in the context of continual learning.", "questions": "- Would it be possible to generalise this analysis to deeper networks?\n- Can you provide sufficient conditions for the invertibility of $\\mathbf{B}$?\n- Could you verify that equations 7 and 8 could be simplified by using $$\\mathbf{Q}\\mathbf{Q}^T(t) = \\mathbf{Z} \\bigg[\\mathbf{S}^{-1} + \\mathbf{B}^T \\Big(e^{\\tilde{\\mathbf{S}}} - \\mathbf{I}\\Big)\\tilde{\\mathbf{S}}^{-1} \\mathbf{B} + \\mathbf{C}^T \\Big(e^{\\tilde{\\mathbf{S}}} - \\mathbf{I}\\Big)\\tilde{\\mathbf{S}}^{-1} \\mathbf{C}\\bigg]^{-1} \\mathbf{Z}^T,$$ where $\\mathbf{Z} = \\begin{bmatrix}\\tilde{\\mathbf{V}} \\Big(e^{\\frac{1}{2} \\tilde{\\mathbf{S}}} \\mathbf{B} + e^{-\\frac{1}{2} \\tilde{\\mathbf{S}}} \\mathbf{C} \\Big) \\\\\\\\ \\tilde{\\mathbf{U}} \\Big(e^{\\frac{1}{2} \\tilde{\\mathbf{S}}} \\mathbf{B} - e^{-\\frac{1}{2} \\tilde{\\mathbf{S}}} \\mathbf{C} \\Big) \\end{bmatrix},$ $\\mathbf{B} = \\tilde{\\mathbf{V}}^T\\mathbf{V} + \\tilde{\\mathbf{U}}^T\\mathbf{U}$ and $\\mathbf{C} = \\tilde{\\mathbf{V}}^T\\mathbf{V} - \\tilde{\\mathbf{U}}^T\\mathbf{U}$?", "limitations": "I think it is great that the authors added a section to the appendix to warn that they discovered a mistake in their proof for the invertibility of $\\mathbf{B}$.\nApart from that, I do not think that this paper would have a direct impact on society", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### originality\n\nThis paper introduces a novel point of view on known dynamics.\nAlso, most contributions are framed in the context of relevant/related work.\nHowever, I noticed a few sections where references could be improved:\n\n  - The experiment with hierarchical data looks very similar to the one used in [17],\n    which was introduced in (Saxe et al., 2013).\n  - I believe that theorem 5.1 is more similar to the setting in [17] than the paper seems to acknowledge.\n    After all, in [17,&nbsp;section 1.3] it is clearly stated that $\\mathbf{W}_1 = \\mathbf{R} \\operatorname{diag}(\\mathbf{b}) \\mathbf{V}^T$ and $\\mathbf{W}_2 = \\mathbf{U} \\operatorname{diag}(\\mathbf{a}) \\mathbf{R}^T$ for some orthogonal matrix $\\mathbf{R}$.\n    In this sense, theorem 5.1 has very similar assumptions as [17], \n    in contrast to what is claimed in section 5 of the submission.\n\nSaxe, A., McClelland, J., & Ganguli, S. Learning hierarchical category structure in deep neural networks. In Proceedings of the 35th Annual Conference of the Cognitive Science Society, 2013.\n\n### quality\n\nAlthough I like the idea, the mathematical aspects of this paper appear to be sloppy.\nThere is also one critical flaw: $\\mathbf{B}$ might not be invertible.\nAssuming that reasonable, sufficient conditions for the invertibility of $\\mathbf{B}$ can be formulated, I believe most results would be correct.\nHowever, often the proofs are not well formulated:\n\n  - The proof for theorem 3.1 starts from equation 5.\n    Therefore, theorem 3.1 would still require assumption 2.1.\n  - The proof for theorem 3.2 has an unexplained gap in it.\n    After all, there are no explicit expressions for $\\mathbf{W}_1(t)$ and $\\mathbf{W}_2$.\n  - Both theorem 3.1 and a proper proof for theorem 3.2 depend on the invertibility of $\\mathbf{B}$,\n    which is only guaranteed for $N_i = N_o$ according to the appendix.\n  - Is Assumption 3.1 an actual assumption? To me, this implies that the results only hold if $N_i \\neq N_o$, which should not be necessary if I understand it correctly.\n    I think it would be better to simply state that assumption 2.1 is not needed anymore.\n  - I would argue that assumption 3.3 is not necessary, because it is implied by assumption 3.2.\n    After all, if $N_h < \\min(N_i, N_o)$, then both $\\operatorname{rank}(\\mathbf{W}_1(0)) \\leq N_h$ and $\\operatorname{rank}(\\mathbf{W}_2(0)) \\leq N_h$ and therefore, $\\operatorname{rank}(\\mathbf{W}_2(0) \\mathbf{W}_1(0)) \\leq \\min(\\operatorname{rank}(\\mathbf{W}_2(0)), \\operatorname{rank}(\\mathbf{W}_1(0))) \\leq N_h$, which leads to a contradiction, given assumption 3.2.\n  - Also, _sufficient rank_ is a term I am not familiar with.\n    I believe rectangular matrices are still said to be of _full rank_ if their rank is the minimum of its dimensions.\n    In this sense, Assumptions 2.4 and 3.2 are equivalent.\n  - The proof of theorem 3.1 introduces a limit, that (I believe) is unnecessary.\n    This leads to additional terms that should always be zero, making the entire formulation harder to understand.\n  \n\n### clarity\n\nOverall, the paper is a nice read and figures are clearly annotated.\nHowever, I think there are a few concepts that would benefit from further explanation:\n\n  - I do not understand what the RSA is.\n    The reference to [49] mentions multiple steps involved in the analysis.\n    However, in the text it is used as something that should be straightforward to compute.\n    Figure 4A seems to indicate it should be related to the covariance,\n    but I do not know if this intuition would be correct.\n  - I fail to understand what it means for a solution to be _rich_ or _lazy_.\n    After all, _rich_ and _lazy_ are used to describe dynamics.\n  - Figure 5C introduces target matrices, which are not explained in the text.\n    I would have assumed that these are the targets for $\\mathbf{A}^T \\mathbf{A}(t)$, but then the off-diagonal elements should converge to non-zero values.\n\n### significance\n\nThe results of this paper could be used to establish possibly interesting connections between different types of analyses that already exist.\nMoreover, the expression for the dynamics might make it possible to enable rich learning when using high variance for the initial weights.\nAlthough it can be argued that the proposed contribution is incremental, I believe that it could be valuable for the community.\nOn the other hand, this approach seems to be limited to networks with a single hidden layer.\n\n---\n\n### Minor comments\n - on line 11: the slow dynamics are described to be _step-like_, in contrast to _exponential_.\n   However, the dynamics in figure 4E are still exponential (albeit after a plateau) and I would describe the dynamics in figure 5A as _sigmoidal_.\n - on line 49-50: I tend to read _This work_ as a reference to the manuscript,\n   rather than to the (variety of) work listed in the previous sentence.\n - line 63: the _but_ seems to be out of place.\n - possible typo on line 77: I would have expected $\\mathbf{W}_2 \\in \\mathbb{R}^{N_o \\times N_h}$.\n - on line 120, it might be useful to explicitly mention that this is actually the _reduced_ SVD.\n - In equation 7, I believe $\\mathbf{A}^{-1}$ should be $\\mathbf{S}^{-1}$.\n - on line 224, the term _troughs_ is used, but I would argue that it is more common to talk about _valleys_.\n - on lines 35 and 247, the term _catastrophic interference_ is used. \n   To the best of my knowledge, this is actually called _catastrophic forgetting_, which can be considered a form of _destructive interference_.\n - on line 238, there is a reference to D, which might not belong there.\n - it would be nice to cite the conference papers rather than their pre-prints\n   (e.g. [17] and [26] were presented at ICLR 2014 and 2022, respectively).\n\n---\n\n**update**: \nConsidering the revision, the answers to my questions and further discussion with the authors significantly, my original evaluation is no longer accurate. \nTherefore, I decided to raise my score from 3 to 6 (as well as the soundness score from 1 to 4).\n\nMy main motivation for increasing the score:\n - A proper proof for the main theorem has been provided in the revision\n - My misunderstandings concerning the generalisation of the Fukuzima results were resolved in the discussion.\n - The main presentation issues were tackled in the revision or were promised in the discussion.\n\nMy main motivation for not increasing the score even further:\n - The analysis is inherently limited to (linear) networks with a single hidden layer.", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "uSpzzZK4nq_", "reviewer": "Reviewer_Rrbu", "summary": "The paper studies the training dynamics of deep (2-layer) linear networks. It derives the full gradient flow dynamics under mild conditions, in a form that is suitable for numerical evaluation. It then applies these results to study several deep learning phenomena including alignment, catastrophic forgetting, and knowledge revision.", "questions": "My main comment is that I do not understand parts of the proof of Theorem 3.1, and have listed several questions below. Looking at the code, these issues seem to be resolved (in particular, SVD computes both $\\tilde{U}$ and $\\tilde{U}_\\perp$ etc.) which makes me think that they have more to do with how the theorem and proof are presented and do not indicate an actual flaw with the proof. I will be happy to revise my score if the authors can answer these questions and clarify the presentation in the paper.\n\n1. Line 123: \"We note that if $N_i < N_o$, the left and right singular vectors are of shape $N_o \\times N_i$ and $N_i \\times N_i$\". If $A$ is a matrix with dimensions $N_o \\times N_i$, its singular value decomposition is given by $A = U S V^T$ where $U$ has dimensions $N_o \\times N_o$, and $V$ has dimensions $N_i \\times N_i$. In other words, the left and right singular vector matrices are square, and so I don't understand the statement in the paper. Also, the matrix $S$ has the same dimensions as $A$ and is not square. Much of the proof revolves around completing partial bases to full ones, but it is not clear to me what this means, given that the singular vectors already span their respective spaces.\n\n2. In the proof (Appendix B.1), equation (26) implies that $e^{F t/\\tau}$ depends on $V_\\perp$ and $U_\\perp$ through $M$. However, according to (20), $F$ itself does not depend on these matrices, and therefore its exponent cannot depend on them either. This is clearly true in eq. (24), since if we take the limit we see that this dependence on $M$ disappears. However, in going to (25) it seems the $\\epsilon$ limit was exchanged with the limit in the exponential series, leading to this result. (I believe the issue here is related to the previous one, as $F$ should already depend on complete bases obtained through SVD.)", "limitations": "This work has no obvious societal impact.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "[Score updated after authors addressed all of my questions]\n\nThe paper improves on previous works in that is presents a closed-form expression for the gradient flow of deep linear networks that works under fairly mild assumptions, and that can be readily evaluated numerically. In particular, equations (7-8) do not include any non-trivial matrix exponentials (i.e. exponentials that include non-diagonal matrices), and only involve decaying exponents (which are numerically stable).\n\nThe numerical experiments confirm the results presented in the paper. The applications of the theoretical results to alignment, catastrophic forgetting, and other effects show the promise that these results hold in shedding some light on interesting deep learning phenomena. The main limitation of this work is that it applies to deep linear networks, and it is not clear to what extent these results will extend to non-linear networks.\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "ME5UoULX8BJ", "reviewer": "Reviewer_vKXk", "summary": "This paper, entitled \"Exact Learning Dynamics of Deep Linear Networks With Prior Knowledge\", \"derives exact solutions to the dynamics of learning\" for deep linear networks. The authors leverage these solutions to understand how different regimes ", "questions": "I have no additional questions, beyond those asked in the Strengths and Weaknesses section above, for the authors. ", "limitations": "The authors do a great job making clear what the assumptions they rely on are and how they are limited.  The featuring of nonlinear networks (Fig. 6) was helpful to this regard and makes this point clear. ", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "STRENGTHS:\n1. The paper is very clearly written and motivated. This is true both of the conceptual and technical discussion. The authors should be commended on their use of clear figures and their commentary on how their expressions should be interpreted. \n\n2. The goal of the paper - to provide theoretical understanding and tools for studying how the specific initialization of networks affects their learning dynamics - is a worthy one. The authors make a substantiative step in this direction by overcoming a difficulty that had previously limited the development of theory. \n\n3. The framework the authors develop provides novel insight. For instance, it was from the theory that weight matrices initialized with large values, that are in the rich regime and have exponential loss, was discovered. Similarly, the connection between input-output correlation and forgetting in deep linear networks was made. These illustrate the potential of the framework. \n\nWEAKNESSES:\n1. It is not clear to me why the task shown in Fig. 4A has the input and output correlation that is shown. Looking at the hierarchy, the fish are equally similar to each other as the flowers, but looking at the input correlation matrix, the fish are more similar to each other than the flowers. Similarly, I did not understand why the flowers have lower output correlation than the fish. This made understanding the results a little difficult. \n\n2. That the alignment dynamics are different for different sized weights is very interesting (Fig. 5). But there was not explanation as to why some tasks should see transient increases in the alignment part way through learning. Is there some intuition/explanation that can be gleaned from the framework?\n\n3. The forgetting metric used in Fig. 6 should be defined. What does negative forgetting correspond to? Additionally, in Fig. 6B, tanh networks look like they have considerable difference in the analytical and numerical forgetting, but in Fig. 6C, this differences look quite small. Indeed, it seems like ReLU networks have larger discrepancy when looking at Fig. 6C, but that they have smaller discrepancy when looking at Fig. 6B. Is this a matter of the color scheme? Or are the two subpanels computing different things?\n\n4. I found it a little hard to keep track of all the nice findings the authors had (possibly because there framework \"bounced around\" as applied to multiple different ideas). The impact of the paper would be increased if a table (or clear discussion in Sec. 7) was included that really spelled out exactly the theoretical contributions that were made. \n\nMINOR POINTS:\n1. The referencing of figures, equations, and assumptions varies throughout the paper (eq., Eq., equation / Fig 5D, Fig. 5D, figure 5D, 5D / assumption 2.3, 2.3). These should be made consistent \n\n2. There is a typo in Eq. 8. It should be $S^{-1}$ and not $A^{-1}$ (this is correct in the appendix just not in the main text). \n\n3. The curves in Fig. 3C should be color coded/marked by their learning rate. As is, it is difficult to determine whether there is any pattern in the error of the approximation of the loss. \n\n4. Is the convention $-T$ for the transpose of the inverse a common one? To me it looks like the quantity is being raised to the negative $T$ power. \n\n5. It should be made clear in the main text that Fig. 4B-D shows the RSA computed once the network has converged. \n\n6. Is there a typo in Fig. 5C? Is $S_{top}$ supposed to be the same as $S_{centre}$? \n\n7. The example of learning a different hierarchy in the subsection of Sec. 6 \"Revising structured knowledge\" did not seem to fit in with the theme of the subsection (although I did think the result was interesting). I would suggest either dropping it or making it its own subsection. ", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "AjrmHg2EZoR", "reviewer": "Reviewer_zYQs", "summary": "The paper studies the impact of initialization settings on gradient based learning and describes a mathematical toolkit for two layer linear networks for understanding the impact of prior knowledge on learning dynamics. The paper derives the exact solutions to these dynamics in the two layer linear network regime as a function of network initialization, utilizing the Matrix Riccati solution.", "questions": "None.", "limitations": "As you point out in your summary, \"While our results broaden the class of two-layer linear network problems that can be described analytically, they remain limited and rely on a set of assumptions\", the paper by itself may come across as too restricted. However, I applaud stating this fact, as it makes it clear that the paper should be better considered as a first step in a direction that potentially gives tooling to analyze more generic network types.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is overall well written and a lot of effort has been placed in making the presentation of the topic very accessible. I specifically want to point out the high quality and informativeness of the included figures.\n\nThe paper continues a longer research direction that tries to understand learning dynamics in (linear) networks. By taking into account so called \"rich learning regimes\" the paper covers a larger class of settings (including pretraining settings hereby), which are also practically more relevant. A core contribution my understanding is picking up a older idea (Matrix Riccati formulation for learning dynamics) and re-establishing it in the deep learning setting.", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept