A Theory of Initialisation's Impact on Specialisation
Authors: Devon Jarvis
Sebastian Lee
Clémentine Carla Juliette Dominé
Andrew M Saxe
Stefano Sarao Mannelli
Venue: ICLR 2025 Poster (2025)
Keywords: machine learning theory
teacher student setup
initialisation
specialisation
statitistical mechanics of learning
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: We examine in theoretical frameworks how different initialisation schemes influence specialisation in neural networks and explore their impact on downstream tasks in settings such as continual learning.

Abstract: Prior work has demonstrated a consistent tendency in neural networks engaged in continual learning tasks, wherein intermediate task similarity results in the highest levels of catastrophic interference. This phenomenon is attributed to the network's tendency to reuse learned features across tasks. However, this explanation heavily relies on the premise that neuron specialisation occurs, i.e. the emergence of  localised representations. Our investigation challenges the validity of this assumption. Using theoretical frameworks for the analysis of neural networks, we show a strong dependence of specialisation on the initial condition. More precisely, we show that weight imbalance and high weight entropy can favour specialised solutions. We then apply these insights in the context of continual learning, first showing the emergence of a monotonic relation between task-similarity and forgetting in non-specialised networks. Finally, we show that  specialization by weight imbalance is beneficial on the commonly employed elastic weight consolidation regularisation technique.

Published as a conference paper at ICLR 2025
A THEORY OF INITIALISATION'S S IMPACT ON SPECIALISATION
Devon Jarvis Sebastian Lee Clémentine Carla Juliette Dominé*.3 Andrew M Saxe, 3,6 & Stefano Sarao Mannelli
ABSTRACT
Prior work has demonstrated a consistent tendency in neural networks engaged in continual learning tasks, wherein intermediate task similarity results in the highest levels of catastrophic interference. This phenomenon is attributed to the network's tendency to reuse learned features across tasks. However, this explanation heavily relies on the premise that neuron specialisation occurs, i.c. the emergence of localised representations. Our investigation challenges the validity of this assumption. Using theoretical frameworks for the analysis of neural networks, condition. More precisely, we show that weight imbalance and high weight entropy can favour specialised solutions. We then apply these insights in the context of continual learning, first showing the emergence of monotonic relation between task- similarity and forgetting in non-specialised networks. Finally, we show that specialization by weight imbalance is beneficial on the commonly employed elastic weight consolidation regularisation technique.
show strong dependence of specialisation on the initial
INTRODUCTION
Theories of representation in biological neural networks span from highly localised representations in single neural units (Barlow [972) to fully distributed shared representations (Hopfield 1982). While shared representations offer greater resilience, specialised representations allow for more ef- ficient encoding of information. Experimental evidence supports both ends of this spectrum, with different brain areas and tasks exhibiting distinct forms of representation (Blakemore et al. 1973] Quiroga et al., 2005) Georgopoulos et al.) 1986 Ishai et 2000
Averbeck et al. 2006). Similarly, artificial neural networks display both shared (LeCun et al. 1989 Erhan et al. 2010 Yosinski et al.)
2014) and specialised representations (Zeiler & Fergus 2014] Voita Voita et al. 2019), ) where a recent ad- vancements in explainable AI, such as the Golden Gate Claude model (Templeton) 2024). exemplify an extreme of the spectrum.
Given the trade-off between shared and specialised representations, critical research challenge lies in understanding how to guide neural networks towards one form or the other. This tension is especially relevant in contexts like disentangled representation learning (Bengio et al.) 2013) and
multi-task learning (Caruana 1997). including continual learning and transfer learning. Specialised representations facilitate faster adaptation and reduce catastrophic forgetting (McCloskey & Cohen 1989: Ratcliff 1990), as they allow networks to rewire efficiently (Suddarth & Kergosien 1990). Rich Caruana's 's seminal work on multi-task learning (Caruana 1997) emphasised the value of specialisation in enhancing performance across multiple tasks. In disentangled representation learning [Locatello et al.] (2019) highlighted that, despite the potential success of unsupervised approaches, disentanglement does not emerge naturally without an explicit inductive bias, un- derscoring the need for supervision or regularisation to enforce such structures. Thus, recent efforts to mitigate catastrophic forgetting (Parisi et al., 2019 De Lange et al., (021) have led to the development of regularisation strategies that promote specialisation, such as elastic weight
School of Computer Science and Applied Mathematics, University of the Witwatersrand; "Center for Com- putational Neuroscience, Flatiron Institute, Simons Foundation; *Gatsby Computational Neuroscience Unit & Sainsbury Wellcome Centre, UCL; 'Data Science and AI, Computer Science and Engineering, Chalmers Uni- versity of Technology and University of Gothenburg: 'Machine Intelligence and Neural Discovery Institute, University of the Witwatersrand; CCIARR Azrieli Glohal Scholar, CIFAR; "Equal cont ri rutuiiiuui on in random order.

---

Published as a conference paper at ICLR 2025
consolidation (Kirkpatrick al. 2017), synaptic intelligence Zenke al. 2017). and learning without forgetting (Li & Hoiem| 2017).
This paper aims to show that initialisation has fundamental importance in achieving specialised lutions, providing a complementary perspective on both the lazy (Jacot et al. 2018) and rich learn-
ing regimes (Mei et al., 2018 Chizat & Bach 2018 Rotskoff & Vanden- Eijnden 2018). Previous research (Chizat et al.
2019] Geiger et al.] [020 Bordelon & Pehlevan 2022) has showns that by
represeenatiion-hhhaacttreeeee by interpolating between these regimes, we can transition from shared representations characterised
random projections in the neural tangent kernels-to effective feature learning (Tarmoun et al., [20TT]
Kunin al., 2024, Xu & Ziyin, 2024] Dominé al., 2024; Varre et al., (024). While our analy- sis remains within the feature learning regime, it adopts a distinct theoretical approach compared
to these studies, concentrating specifically on the impact of initialisation within standard synthetic frameworks for neural networks. This exploration reveals how initialisation can skew the learning dynamies towards either specialised or shared representations, thereby adding new dimension to the study of learning dynamics in over-parameterised networks.
Our work makes the following main contributions:
We study the impact of initialisation on specialisation through two theoretical frameworks:
We utilise the dynamies of deep linear networks to investigate the evolution of spe- cialisation (Saxe et al. (2033;;
We extend this analysis to nigh-dimensional mean-field neural networks learning with stochastic gradient descent (Saad & Solla 1995b a Biehl & Schwarze, 1995).
We identify specific initialisation schemes that promote specialised solutions by increasing the entropy of the readout weights and creating an imbalance between the first and last layers, akin to the findings of Dominé et al. (2024).
We demonstrate that there are two regimes of forgetting profiles contingent on neuron spe- cialisation, reconciling recent findings regarding the non-monotonie relationship between task similarity and catastrophic forgetting (Ramasesh et al., 2020 Lee et al. 2021 2022)
with the traditional belief in monotonic forgetting (Goodfellow et al. (003)).
Finally, we demonstrate the practical implications of our results on regularisation strate- gies, specifically analysing how Elastic Weight Consolidation (EWC) (Kirkpatrick et al. 2017) is influenced by specialisation dynamics, highlighting potential pitfalls associated with regularisation methods in continual learning.
Fig. of Sec. 2 serves as a motivating figure and contrasts with results from the literature (Goldt et al. 2019) that sigmoidal networks specialise while ReI.U networks do not. We show that ReL. networks can also specialise (and sigmoidal networks do not always specialise) depending on weight initialisations. Sec. 3 then aims to understand in more detail "what" aspects of the initialisation scheme lead to specialised solutions. This is achieved theoretically through the lens of deep lin- network dynamics (Saxe et al.) 2013) and validated empirically
canonical disentanglement task with a variational autoencoder (VAE) in Sec. 3.2 Having established an initialisation strategy then apply this strategy control specialisation when studying
which promotes specialisation, continual learning in Sec. 4 Consequently, we reconcile the two distinct forgetting profiles observed in practice which determine how an increase in task similarity leads to forgetting of earlier tasks: 1. the Monotonic profile (Goodfellow et al.] (003)) 2. the Maslow's Hammer profile (Ramaseshi let al., 2020, Lee et al., 2021; 2022). We demonstrate that the Maslow's 's Hammer profile results from neuron specialisation, incurs less interference as tasks become significantly different and en- ables regularisation techniques like Elastic Weight Consolidation (EWC) Kirkpatrick et al.) (2017). Finally, in Sec.5) we reflect on the limitations of our work and propose future directions for research.
2 SPECIALISATION IN THE TEACHER-STUDENT FRAMEWORK
The teacher-student framework is a generative model that allows for the controlled creation of syn- thetic datasets (Gardner & Derrida 1989). The framework involves two classifiers: the teacher and the student, for instance represented as neural networks as exemplified in Fig. 1a The teacher, has fixed randomly drawn weights and maps random inputs æ from a given distribution labels, pro- viding a rule for generating data. The student, on the other hand, updates its parameters through learning protocols like stochastic gradient descent (SGD) to approximate the teacher's outputs.
2

---

Published as a conference paper at ICLR 2025
a)
b)
Qik
Rin
-2
Sigmoidal
Generalisation Error, log
-
6
C(y."">
2
4 Step
6
8
10
Qik
R
W
h
ReLU
-5
0 = Iw (Ow) h [0] In (9n)
-10
2
6
Step
10"
Figure I: Initialisation impacts specialisation. a) In the teacher-student setup student network is trained with labels generated by fixed teacher network. Previous work established relationship between the activation function & and the propensity for the student nodes to specialise to teacher nodes. However show in this work that this is overly simplistic description; other factors including student weight initialisations /w. lh, parameterised by θw, θ₁ arguably play a stronger role. b) Generalisation error curves for two simulations of the teacher-student setup, one with ReLU activation function and one with a scaled error activation function. Ow and θ, are chosen to achieve solution with ReLU that specialises as indicated by sparser overlap matrices on the bottom right, and a scaled error function solution that does not specialise-as indicated by denser overlap matrices on the top right. sparse (dense) Q matrix shows few (many) student nodes are active, while a sparse (dense) R matrix shows student nodes are representing teacher nodes in targeted (redundant) manner. Further details for the quantities described can be found Sec.
While detailed quantitative characterisation of specialisation follows in the next sections, we briefly introduce the concept within the teacher-student framework. [Saad & Solla (1995b) showed that, when both teacher and student are modelled as committee machines, each student neuron specialises by aligning with specific teacher neuron. Similarly, Goldt al. (2019) observed that for certain activation functions in two-layer networks, over-parameterised student will selectively use only a subset of those units to replicate the teacher's outputs. This phenomenon, termed specialisation, stands in contrast to student redundantly sharing representations of the teacher across neurons. this work we present a more comprehensive account of the factors underlying specialisation. In contrast to (Goldt et al. 2019), we argue that initialisation-aa not the activation function-is is chiefly responsible. We highlight this in Fig. 16] by showing that with carefully chosen initialisations we can train highly specialised ReLU student (bottom panels), and non- specialising sigmoidal student (top paholn)-shown by the sparser Q and I matrices of the ReLU network-which represents the opposite of the conclusions presented in (Goldt et al. 2019). We be- gin by aiming to establish what properties of an initialisation promote specialisation. This question well suited the deep linear network theory (Saxe et al.) (013) and we turn to this strategy now.
SPECIALISATION EXPLAINED USING LINEAR DYNAMICS
Here we construct a synthetic setup. to study the influence of initialisation on specialisation using the deep linear network framework (Saxe et al. 2022 2019). While deep linear networks can only represent linear input-output mappings, they showcase intricate fixed point structure and nonlinear learning dynamics reminiscent of phenomena seen in nonlinear networks (Baldi & Hornik 1989 Fukumizu 1998 Arora al. 2018 Lampinen & Ganguli 2019). We consider specialisation adher- ing to the definition proposed by the statistical physics literature (Goldt et al. 2019) which considers whether neuron will account for all of the variance associated to feature, while the others remain inactive (a phenomenon close to activation sparsity and reminiscent how initial conditions can lead to minimal subnetworks in the "Lottery Ticket Hypothesis" (Frankle & Carbin 2018))
3

---

Published as a conference paper at ICLR 2025
Original Network
Pathway Representation
SV Representation
h
h
h2
h
h2
w.=hWW
W
W2
W
W
W2
Figure 2: Summary of our setup, notation and strategy. a) The original network with two hidden neurons learning the regression task. b) We split the network into two separate pathways and con- sider their dynamics individually. Since both networks are learning the same task simultancously. their dynamics are coupled. c) To obtain the dynamics of the two pathways and calculate their es- caping and hitting time we track the pathway dynamics in terms of the network's effective singular values. The closed form dynamics for the pathway singular value are given in Eq. 3
This is in contrast to other work on modularity (Jarvis et al. 2023 such as Neural Module Networks (Andreas et al.] 2016; Hu et al. |017,2018] Andreas (2188). mixture-of-expert models (Masoudnia
& Ebrahimpour 2014] Bengio et al. 2015] Shazeer et al. (017), tensor product networks (Smolen- 2019), which consider specialisation as a subset of a network or module performing a single "task" " or only being activated by one in- terpretable feature in the dataset. Thus, these works are focusing on specialisation to imply feature sparsity (Dasgupta et al. 2022), while we are concerned with the learning mechanism that leads to sparse input-output mappings.
sky et al.. 2022). among others (Chang al. 2018| al.
3.1 SPECIALISATION THE DEEP LINEAR NETWORK FRAMEWORK
To connect this framework to specialisation we use the notion of the "neural race" from Saxe et al. (022). The neural race hypothesis says that the pathways through a network are racing to explain the variance in the dataset (L.c. to perform the onput-output mapping). Thus, we consider the limited case of a network with two hidden neurons and one output neuron. Fig. Fig.Zdepice depicts the setup, notation and strategy for this section. By defining "hitting time" as how long it takes a pathway to reach its final converged value, and "escappig time" as how long it takes the pathway to begin learning, ask the question: "when will one pathway finish learning (reach it's hitting time t*) before the other begins learning (reaches it's 's escaping time i)" In cases when this occurs, the network would have specialised as only one pathway will have any activity and will explain all of the data. Similar to Sec. 4 we generate data by sampling the elements of a data point from a Gaussian distribution (a;; ~ (O, 1)) with i = 1, 1,....dl dl We then define ground-truth mapping (Wr) and generate labels y WT x. We only consider regression tasks in this section, thus y E R. For inputs we can form the input matrix E _ and row vector of scalar outputs E Rlxn, The dataset statistics which drive learning are collected in the input and input-output correlation matrices, Σ² and y respectively. For the task described above the singular value decomposition of these matrices are:
D = =[Xx']==V
E' E'yX usv
Here, u E 1-1,1), u is a vector such that v = and V is an orthogonal singular vector matrix. Correspondingly, is the singular value for the rank task and D diagonal matrix of singular values. Note that as in [axx et al. (2013) 2022) - we assume that the correlation matrices are mutually diagonalisable (share the same V) up the rank of Eør.
For this task we consider a single hidden layer network (Fig. leftt computing output I hW* x with h IR" and W Rpxd in response to an input E Rd The network trained to minimise the mean squared error loss using full batch gradient descent with a small learning rate To identify when specialisation will occur in this network, we split the network into two pathways with one hidden neuron each. The input and output dimensions remain the same (Fig. 2 middle). Finally we obtain the linear dynamics (ultimately depicted as Eq. 3 for each pathway (the full details and assumptions of the derivation are given in Appendix AA. In this setting, the pathway's input-output mapping after t epochs of training is (t(t) w (t). Notice that, the pathways have one hidden neuron and so h is scalar and w is the vector of input respective hidden neuron weights, to once again alleviate notation do not denote which pathway. Assuming that the pathway weights align to the singular vectors of the dataset from early in training, as described by the "silent alignment effect" (Atanasov et al. 2021
4

---

Published as a conference paper at ICLR 2025
b/
100
10000
50
9000
Network Singulan
60-
15000
A (Faet)
prodicted
10
20-
2000
1000
2000 3000 4000
10.0
Epoch number
(Slo
Figure 3: Linear Dynamics from imbalanced initialisation leads to specialisation. Panels a- b) Show agreement between our theoretical curves and simulations for the training dynamics of: (a) the network's 's singular value dynamics, escaping times (verticals towards left) and hitting times (verticals towards right) for varying scales of weight imbalance (depicted by colour), (b) and the network's s movement in weight space depicted by the sequence of dots over weight space. Colour depicts the loss of the network configuration at point. Panel c) shows a phase diagram representing how pathways with different initial weight imbalances lead to specialisation. The two axis represent the weight imbalance the two pathways our broader network (A2 on the x-axis for the slower pathway and Å the y-axis for the faster pathway). The colour represents how close the slower pathway is to reaching its escaping time at its closest point throughout training (in log scale). We that the more inbalanced the fast pathway relative to the slower pathway, the more likely the network will specialise. The white region represents when the inbalance is equal or reversed.
perform change of variables and write the mapping in terms of the dataset singular vectors:
h(t)w(()) ut)o"
where (t) is the pathway's scalar effective singular value and the only time-dependent component of the decomposed mapping. While the alignment assumption is strong, linear paradigms with these
assumptions have been used successfully in the past (Saxe et al.] 2019 Lampinen & Ganguli [219]]
Braun et al.] 2022 Jarvis et al. 2023 Dominé et al.) 2024) Jarvis et al. (025). By appropriately changing variables, we obtain a closed form equation describing how evolves through time
k (eexp) (gn(A)kk)
-1)
Ad (eepp egg((Akk) 1) + +1)
A (1)) = sinh 2 tanh-1 -1
(3)
2s cexp
(gm((Ak t +
where e is defined constant, shows that & the variable interacting with time (1) and consequence determines how quickly the network will learn. Three factors affect & fastening learning: 1. s the input-output correlation matrix singular value, 2. the input correlation matrix singular value, and = Λ² ww" which denotes the imbalance between the weights of the network. Notice that-as shown in Appendix A a conserved quantity and constant throughout training. Thus, given a and d matrices-the only property which can promote faster learning in the network the imbalance parameter. For our experiments to remove one of the interactions within k. With the training dynamics of a singular value defined as in Eq. 3 we can formally define the escaping time as f f such that wtt)
is the learning time constant and 4 + Eq.
determine the increase whiten the input data such that k: / + X
for small IR Similarly, ly, we define the hitting time as t = such that ((o) - (() = 8 for a small d E I.
Fig. 3 aa-b) show confirmation of the validity our theory by comparing with simulations. Instead, Fig. 3C represents the main result of this section. We consider both network pathways and vary the weight imbalance for each (Astow for the pathway with the lower imbalance and Ahigh for the pathway with the larger imbalance). We place these two values on the axes and in colour depict how close the slower pathway comes to reaching its escaping time across its training (min(fstom (t))) When zero, it means that during training there is timestep where the network is less than one epoch from its escaping time (so it will learn). In this case there will not be specialisation as both pathways will learn some part of the input-output mapping. When the colour is positive it means there will be specialisation as the slower pathway is always at least a full epoch away from learning. It is important to note that the slower pathway's escaping time is moving constantly as the faster pathway accounts for variance the data. This decreases the input-output singular value in l for this pathway and makes leaming slower. Due to this coupling we are also unable to obtain completely closed form
5

---

Published as a conference paper at ICLR 2025
a) 0.38
b)
Loss reconstruction
37131
Gain 0.3
0.28
3607
0.18
3501
Gain
0.08
0.3
0.7
1.0
2.0 3.0
3394
Gain
Gain
Figure 4: Violin plots of a) the Disentanglement, Completeness, and Informativenes (DCI) East- wood & Williams 2018)) score and b) the reconstruction loss against gain. The disentanglement score decreases as the gain increases while the reconstruction loss remains steady, c) Example traver- sals of models with gains 2 and 0.3, respectively, highlighting disentangled dimension for gain 0.3
and mixed dimension for gain 2. Experimental details can be found in appendix D
equations for the slower pathway term's of the faster pathway's effective singular value. However, this phase diagram would not be computationally feasible without the closed-form escaping time, hitting time and training dynamics (see Appendix B for our process on constructing this plot). Fi- nally, we only consider imbalances where the output scale is larger than the input scale. Recent work (Kunin et al., 2024 Dominé et al., 2024) has shown that having larger input weight scale pushes the network towards lazy learning while output heavy imbalance promotes feature learning. From Fig. F we see that there a clear phase transition from non-specialised representations to specialised ones. This occurs with increasing imbalance of the faster pathway. Increasing the imbalance of the slower pathway can similarly combat this specialisation pressure. Thus, the relative imbalance of the two pathway initialisation will dictate whether specialised representations are learned.
3.2 INCREASING SPECIALISATION DISENTANGLEI REPRESENTATION LEARNING
To empirically support the linear network theory, we extend the results inbalanced initialisation and apply them, beyond the limited setting of our framework, in the context of disentangled repre- sentation learning, where the goal is to separate latent factors into both feature and activity sparse neural representations. For further empirical support of our theoretical results on Sparse Autoen- coders (SAEs) see App. H Bengio et al. (2013) introduced the importance of disentanglement for interpretability and generalisation. A seminal contribution to this domain came with the 3- model, where Higgins al. (2017) demonstrated how increasing the KI.-divergence term can en- force disentanglement by encouraging specialised latent representations. Many studies have built upon these foundational frameworks to enhance disentanglement performance, exploring different training regimes (Locatello et al. 22200] Fumeer et al. 202T) and loss functions (Chen et al. 2019) Kim & Mnih 2019) Kumar et al. 2018). Here we contribute to this literature by applying our theoretical insights and examining the impact of initialisation disentanglement performance.
Specifically, we examine how initialisation impacts specialisation in disentanglement learning on the 2018) using the B-VAE model- adopted for such tasks Higgins et al. 2017: 7 Burgesss et al. 2018). We implement a 3--AE model, employing the "Deep- GaussianL architecture for the decoder and the "DeepLinear" architecture for the encoder, as specified in Locatello et al.] (2019). Both architectures are composed of five fully connected layers with ReLU activations. The model is trained using the Adam optimiser, optimising loss func- tion that combines KL divergence and binary cross-entropy- based reconstruction loss. Additional details are given in Appendix [D] In these experiments, we adjust the variance of the weights deep fully-connected encoder, by varying the constant gain of the Xavier initialisation (Glorot & [Bengio| 2010). Specifically,, the first block of layers was initialised with gain g while the readout layer received a gain 1/g. Notice that g represents the standard initialisation scheme.
3DShapes dataset (Burgess & Kim
Results are shown in Fig. despite very similar levels reconstruction loss, networks initialised with smaller gains improved disentanglement in the 3- 3-VAE network, as reflected in higher Disentan- glement, Completeness, and Informativeness (DCI) scores (Eastwood & Williams 2018). This re- sult confirms that modulating the initialisation gain can increase the network 's disentanglement. Al- though the scope of these experiments is limited, they provide preliminary validation of our theoreti- cal framework in more realistic contexts, encouraging further investigation into alternative initialisa-

---

Published as a conference paper at ICLR 2025
tion schemes with varying levels of balance. Having investigated the role initialisation plays in pro- moting specialisation, we return to the original setting of Sec. Aec.2to understand the role of initiaaisatiin in governing network behaviour during continual learning. In the light of these results aim to visit the two established forgetting profiles empirically observed in the continual learning literature: Namely, the Maslow's Hammer profile, observed empirically first in Ramasesh et. (2020), and the monotonic forgetting profile, more typically assumed and observed in Goodfellow el al. (2013).
4 CONTINUAL LEARNING
As Caruana (1997) noted, multi-task learning benefits significantly from task-specific specialisa- tion, allowing the network to better preserve performance across multiple domains. In the context of continual learning, Ramasesh et al. (2020) Lee et al. (2021) observed that forgetting does not monotonically increase with task similarity. Lee et al.) (0222) provided mechanistic explanation, showing that this phenomenon is due to the interplay between re-use of specialised neurons and acti- vation of unused ones. In this section, we build on these findings and show that this phenomenology can be disrupted by initialisation schemes that disincentives specialisation.
4.1 CONTINUAL LEARNING IN THE TWO-LAYER TEACHER-STUDENT SETUP
We use teacher-student framework, introduced in Sec. 2 which has been analysed in Lee et al. 202T 2022). This model consists of two randomly initialised teacher networks-one for an up- stream task and one for downstream task. Each teacher is represented by two-layer neural networks with p" hidden units and weights W
x A. for the upstream task, and W T 222 for the down- stream task. Given a random input E IRd, drawn i.i.d. from a Gaussian distribution 2i ~ No0, 1), the teachers generate labels according the equation:
W Vd
h
for 1,2,
(4)
where & is non-linear activation function, chosen here ((x) = erf (z/ 2). This setup allows us generate two datasets D(1) and D(2), with controlled similarity between the tasks by manipulating the teacher weights. Specifically, we generate W . and h. with i.i.d. Gaussian entries, while
( is generated as: T
W (2) T
W (1) T
VI
(aux)
where W (1ix)) an auxiliary weight matrix, and controls the correlation between tasks. The student is a two-layer neural network with p hidden units, using the same non-linearity é. It is trained using online stochastic gradient descent on a squared error loss, with a shared first-layer weight matrix W and task specife readout weights (!) and (2)), For both layers, the initial weights are sampled i..dd from a Gaussian distribution, with the first-layer weights W having standard deviation ow- While most previous studies follow a similar scheme for the readout weights, we introduce a novel initialisation scheme using polar coordinates. The updates for W and (t) at iteration e, under SGD on the squared error loss, are given by:
W| 1] [(e]
h
Wx Vd
(0))
0'
Wx Vd
1
(6)
Vd
0 [e +
(09)
?
h (t)
s
W x
(1)
Wx x
d
Vd
Vd
where 77 is the learning rate and y is the target output from the teacher network for task t. In the large input dimension limit d oc, key observables, such the generalisation error, can be captured by a few order parameters:
WWT. Q d
R 1 (w
T(,e) 1 -w - w('T
h(0,
h4; (8)
where t E {1,2) refer to the two tasks. The generalisation error for task t is then:
(1) Izn (n(( + Inn IennnT (T(1.)) .1 2
IIz(Q. ( ((t.)) h (0) 1
7

---

Published as a conference paper at ICLR 2025
a/16
x/8
3w/11
-a/4
m
ReLU
00.6
0.6
.0,6
20.G
g0.G
.35
0.4
0.4
0.4
0.4
0.4
-3
-2 -1
-3
-2
-1
-3
-2
-1
-3
-2
-1
-3
-2
-1
log
log
log
log
log
b)
0
1//16
x/a
3///10
/r/4
Sigmoidal
.0.6
0..6
_0.6
L0.6
L0.6
.35
0.4
0.4
0.4
0.4
0.4
-3 -
-3
-
-3
-
-3
-
-
-3
-
log
Log
Log
log
Log
Figure 5: Phase diagrams show significance of initialisation for specialisation. The phase di- agrams show with colour the aggregated entropy Eq. 10 evaluated for different initialisations. On the x-axis span the standard deviation of the first layer. The second layer is initialised us- ing polar coordinates, and the y-axis represents the norm while the different panels give the angle spanning from orthogonal units (0 0) to identical units (0 x/4). Specialisation is achieved by blue-leaning initialisations, while yellow-leaning ones exhibit high entropy and therefore non- specialised solutions. Additional results can be found in Appendix E
The evolu- tion of these parameters throughout training can be tracked to study the learning dynamics, as first shown in (Saad & Solla (1995a); Bichl & Schwarze (1995); Goldt et al. (2019). For the specific case continual learning, Lee et al. (2021) derived the governing ordinary differential equations (ODEs), provided in Appendix
where /21 and 122 are explicit functions of the order parameters, detailed in Appendix C
4.2 SPECIALISATION RELEVANCE FOR CONTINUAL LEARNING
The continual learning results in the teacher-student setup, including the non-monotonic relationship between catastrophic forgetting and task similarity, often implicitly assume that the student has specialised to the teacher in the first task. This assumption allows for spare capacity to represent the second task. However, as shown in Fig. lb there are regimes where this assumption of specialisation is violated. Here, we expand on these findings and their implications for forgetting.
student can effectively ignore unit in two ways: either the unit's post-activation is near (inac- tive), or the corresponding second-layer weight is O. This motivates three measures for specialisation based on the definition of entropy-over the hidden units, head weights, and the product of both:
H Σ 1iu log [hul.
HQ
Ed Qu, H - -n=--aa -n
Qulha log(Qu]hal): (10)
where the tilde denote normalisation, i.e. |A, ha EF l., and Q.
Qu Maximum entropy ΣS Q.
corresponds to no specialisation, while minimum entropy corresponds to maximum specialisation.
We can investigate how these measures vary function of different properties of the problem setup, in particular those related to initialisation. To simplify the analysis, we begin with the case where the optimal number of tasks is p" and the network has p 2 output units. This allows us to initialise the second layer weights in polar coordinates, with precise and interpretable control over scale and asymmetry of weights. Formally
parameterise our readout initialisations according to h (+) (0; (t), g(()] (f-(t) cos ((t), p((t) sin g(t)). Fig. contain phase diagrams showing how the entropy measures in Eg. 10 vary with the initialisation parameters p((), ((t), and ow. We can make several observations: (i) the strongest determinant
specialisation is the asymmetry the second layer weights, i.e. the 0 parameter. (ii) this is the for both ReLU and sigmoidal activation functions, reinforcing the point made in the example from Fig. lb (iii) the seale of initialisations (parameters ow, r)) are also important.
8

---

Published as a conference paper at ICLR 2025
a)
0.3
Final Node Norm
Forgetting
0.2
0.5
0.1
0
0.5
0.5
Similarity, 7
Similarity,
(0) = 0.01, r = 0.1
(()) = 0.01, 113)
(()) = 0.01, (2) = 0.1
((0) (
000 x/4, (12)
( 0, ((9 x/4
Figure 6: Initialisation and specialisation properties can influence profile of forgetting vs. sim- function of task similarity be both monotonic, shown here for the cases of specialisation after both tasks (blue), and no-specialisation large, asymmetric second head initialisation (orange); or non-monotonic (green, characterised by Maslow's hammer Lee iet al.] (2022)). (b) the final norm of the two nodes (one solid and one dashed), i.e. at the end of training on both tasks, as a function of task similarity. In the cases that lead monotonic forget- ting. nodes are fully re-used, either because the corresponding new head is initialised large (orange) or because the new head is symmetrically initialised and the nodes continue to represent redundant information during the second task (blue). Params: N 10000, 17 1, = 1, 2, ow = 0.001.
ilarity. (a) forgetting
4.2.1 SPECTALISATION UNDERLIES MASLOW'S HAMMER.
The phase diagrams in 5 demonstrate that initialisation can drastically change the type of so- lutions found by the student after training on one teacher. While this may be inconsequential if the generalisation eтror remains unaffected, in many cases, the precise nature of the learned representa- tion can significantly impact downstream tasks.
In the worst case scenario, the student undergoes specialisation during the first task. During the second task there is no notion of the trade-off between node re-use and node activation discussed in Lee et al. (202);; rather the student continues to find non-specialised solution to the second teacher, effectively fully re-using it's entire representation for the second task. Consequently, the amount of forgetting with respect to the initial task decreases monotonically with task similarity, thereby breaking the inverted U-shaped pattern characteristic of Maslow's 's hammer that has been observed in various continual learning setups (Ramasesh et al. 2020). This extreme case is illus- trated in Fig. 6) Further, even with specialisation after the first task, large asymmetric initialisation in the second task readout weights induce this monotonic relationship. again by pushing the student into re-use rather than activation. In Appppedii G we complement these results with experiments on a task constructed around MNIST and find qualitatively similar results.
In a broader context, rich diversity of behaviours can emerge, driven by factors such as the ini- tialisation schemes, the scale of weights the first layer, and the readout heads for both tasks. A glimpse of this behavioural diversity is provided in [Appendix F] where we further explore the interaction between these factors and their impact on forgetting in continual learning.
4.2.2 SPECIALISATION UNDERLIES EWC.
The findings relating specialisation to forgetting from See. 4.2.1 have direct consequences for inter- ference mitigation strategies such as EWC. EWC is regularisation-but based method that computes measure "importance" for each weight with respect to a task via the Fischer information (Kirk- patrick et al. 2017). Subsequently squared penalty scaled by this importance is applied to deviation of this weight during learning of future tasks as follows: LEwc(W) L(W) + & E, F. (W; W?), where the Fischer information matrix, E is a regularisation strength parameter, and W are the weights the end of training the first task.

---

Published as a conference paper at ICLR 2025
a)
Specialised Solution
b)
Non-Specialised Solution
-
log
-4
loge
-4
-6
0
0.2
0.4
0.6
0.8
0
0.2
0.4
0.6
0.8
1
Step
-107
Step
-10
==1
= 10
= 10
5 =
10
Figure 7: EWC is strongly reliant on specialisation. We show the generalisation error in the first (solid line) and second (dashed) task for different EWC regularisation strengths. (a) When the student finds specialised solution the first task, there is range of EWC regularisation strength & for which the activated units can remain fixed and spare capacity can be used to learn the second task-leading
low generalisation error in both tasks (g = 10- (b) When the student does not specialise in the first task, EWC reduces to an inflexible regulariser that either penalises plasticity everywhere-leadin
10-4 perform very well).
little forgetting but no further learning (e.g. & 1). or does not penalise any plasticity-laadin to catastrophic forgetting (e.g. & = 10-
In cases where the network does not specialise, i.e. multiple student nodes learn redundant repre- sentations for given teacher node, the nodes have equal importance. Consequently EWC cannot distinguish between these sets of weights and depending on the regularisation parameter À either lets these nodes move during training on the second task (under-regullarises) leading forgetting, or lets none move (over-regularises) leading to no transfer. We show results illustrating this behaviour in the teacher-student setup in Fig. In particular we show the regime of intermediate task similarity, wherein (Lee et al. 2022 previously argued that EWC should perform better than methods such as replay.
5 LIMITATIONS AND PERSPECTIVES
This work operates within simplified frameworks, which-while widely used in the analysis of neu- ral networks do not fully capture the complexity of modern architectures and real-world data. Our experiments rely on Gaussian input data and simplified input-output relations, which are far from the intricacies of real-world scenarios. A natural next step is to extend our analysis to more realistic generative models, such as the hidden manifold model (Goldt et al., 2020) or the superstatistical generative model (Adomaityte et al. 2023), which offer more structured data distributions and bet- ter capture observations from real data experiments Another promising direction is to complement analytical approaches with numerical experiments on controlled real-world datasets. While this may sacrifice some analytical tractability, it brings us closer to addressing practical challenges. For in- stance, transfer learning settings, such as those explored in Gerace al. (2024), provide a useful benchmark for testing our theoretical findings in more complex environments. While we focused on continual learning, other ML domains are affected by specialised representations. An interesting direction concerns the emergence of compositionality. [epori et al. (2023); Driscoll et al. (2024) reported the emergence of compositional representations in neural networks and theoretical frame-
works are now available to investigate the phenomenon (Lee et al. 2024).
While the current work remains theoretical in nature, focusing on simplified models for analyti- cal tractability, thorough exploration of the practical implications of our findings, particularly in disentangled representation learning, is beyond the scope of this paper. However, we aim to ad- dress this in future work by shifting towards a more experimental approach. Specifically, we plan to explore a broader range of network architectures, datasets-sun as Car3D (Du et al. 2024) and dSprites (Matthey et al. 2017) and evaluation metrics satricsssuchaae as SAP (Kumar et al.) 2018 Higgins al. 2017). This future study will allow us to validate our theoretical insights and fully assess their relevance in real-world settings.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This research was funded in whole, in part, by the Wellcome Trust [216386/2//9/22] For the purpose of Open Access, the author has applied a CC BY public copyright license to any Author Accepted Manuscript version arising from this submission. C.D. and A.S. are supported by the Gatsby Charitable Foundation. A.S. was supported by the Sainsbury Wellcome Centre Core Grant (219627/2/19/7) and A.S. is CIFAR Azrieli Global Scholar in the Learning Machines & Brains program. S.S.M. was supported by the Wallenberg AI, Autonomous Systems, and Software Program (WASP). D.J. a Google PhD Fellow mentored by Dr. Gamaleldin F. Elsayed and Commonwealth Scholar.
REFERENCES
Amir H. Abdi, Purang Abolmaesumi, and Sidney Fels. Variational learning with disentanglement- pytorch. arXiv preprint arXiv: 19/2.05844 2019.
Urte Adomaityte, Gabriele Sicuro, and Pierpaolo Vivo. Classification of superstatistical features in high dimensions. In 2023 Conference on Neural Information Procecessing Systems, 2023.
Jacob Andreas. Measuring compositionality in representation learning. In International Conference on Learning Representations, 2018 8.
Jacob Andreas, Mareus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 39 39.49,,448 2016.
S. Arora, N. Cohen, and E. Hazan. On the optimization of deep networks: Implicit acceleration by overparameterizatton. 35th International Conference on Machine Learning, ICML 2018, 1: 372-389, 2018. arXiv: 1802.06509 ISBN: 9781510867963.
Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. arXiv preprint arXi::2111.00034 2021.
Bruno Averbeck, Peter Latham, and Alexandre Pouget. Neural correlations, population coding and computation. Nature reviews neuroscience, 7(5):358-366, 2006.
Baldi and K. Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Newral Networks, 2(1):53-58, January 1989. ISSN 08936080. doi: 10.
0116/8893-6800889988
URL
eeaabbellleee elsevier. come remm///etrieee//
p117089360009900114
Horace B Barlow. Single units and sensation: neuron doctrine for perceptual psychology? Per- ception, 144)3371-394, 1972.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for sgd: Effective dynamics and critical scaling. Advances in Neural Information Processing Systems, 35:25349 2.225249-25362. 2022.
Emmanuel Bengio, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup. Conditional computation in neural networks for faster models. arXiv preprint 15 2015.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions pattern analysis and machine intelligence, 35(8) 1828, 2013.
Michael Bichl and Holm Schwarze. Learning by on-line gradient descent. Journal of Physics A: Mathematical and general, 28(3):643, 1995.
Colin Blakemore, James PJ Muncey, and Rosalind M Ridley. Stimulus specificity in the human visual system. Vision research, 13110):1915-1931 1973.
Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. Advances in Neural Information Processing Systems, 35:32240-3222 2022.
11

---

Published as a conference paper at ICLR 2025
Lukas Braun, Clémentine Dominé, James Fitzgerald, and Andrew Saxe. Exact learning dynamics of deep linear networks with prior knowledge. Advances in Neural Information Processing Systems, 35:6615-6629, 2022.
Chris Burgess and Hyunjik Kim. 3d shapes dataset. https///shhhbbooooeeeemmmmlissss 2018.
Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Des- jardins, and Alexander Lerchner. Understanding disentangling in f-vae. arXiv preprint 1804.03900 2018.
Rich Caruana. Multitask learning, Machine learning, 28:41-75, 1997.
Michael B Chang, Abhishek Gupta, Sergey Levine, and Thomas L Griffiths. Automatically composing representation transformations as a means for generalization. arXiv preprint 2018.
Ricky T. Q. Chen, Xuechen Li, Roger Grosse, and David Duvenaud. Isolating sources of disentan- glement in variational autoencoders, 2019.
942
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- lems, 31, 2018
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Ishita Dasgupta, Erin Grant, and Tom Griffiths. Distinguishing rule and exemplar-hane general- ization in learning systems. In International Conference on Machine Learning, pp. 4816-4830.
PMLR, 2022.
Matthias De Lange, Rahaf Aljundi, Mare Masana, Sarah Parisot, Xu Jia, Aleš Leonardis, Gregory Slabaugh, and Tinne Tuytelaars. continual learning survey: Defying forgetting classification tasks. IEEE transactions on pattern analysis and machine intelligence, 44(7):3366-332) 2021.
Clémentine C. Dominé, Nicolas Anguita, Alexandra M. Proca, Lukas Braun, Daniel Kunin, Pedro A. M. Mediano, and Andrew M. Saxe. From lazy to rich: Exact learning dynamics in deep linear networks, 2024. htth::::: / arxiv. org/ abs/ 2409. 14623
Laura N Driscoll, Krishna Shenoy, and David Sussillo. Flexible multitask computation in recurrent networks utilizes shared dynamical motifs. Nature Neuroscience, 27(7):1349-1363 2024.
Xiaobiao Du, Haiyang Sun, Shuyun Wang, Zhuojie Wu, Hongwei Sheng, Jiaying Ying, Ming Lu, Tianqing Zhu, Kun Zhan, and Xin Yu. 3drealcar: An in- wild rgb-d car dataset with 360-degree views, 2024. URL htts:::::/ wwq.iv.or./.o... abs/ 2200 04875
Cian Eastwood and Christopher KI Williams. A framework for the quantitative evaluation of disen- tangled representations. In 6th International Conference on Learning Representations, 2018.
Dumitru Erhan, Aaron Courville, Yoshua Bengio, and Pascal Vincent. Why does unsupervised pre- training help deep learning? In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 201-208 JMLR Workshop and Conference Proceedings, 2010.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv: 1803:033330333 2018
K. Fukumizu. Effect of Batch Learning In Multilayer Neural Networks. In Proceedings of the 5th International Conference Neural Information Processing, pp. 67-70, 1998.
Marco Fumero, Luca Cosmo, Simone Melzi, and Emanuele Rodola. Learning disentangled rep- resentations via product manifold projection. In Marina Meila and Tong Zhang (eds.), Pro- ceedings of the 38th International Conference Machine Learning. volume 139 of Proceed- ings of Machine Learning Research, pp. 3530-3540 PMLR 18-24 Jul 2021. URL hetps:
7/proceed ngs. ml www.pppeasssem. / www.://ff/mmmmoommm htm
12

---

Published as a conference paper at ICLR 2025
Elizabeth Gardner and Bernard Derrida. Three unfinished works the optimal storage capacity of networks. Journal of Physics A: Mathematical and General, 22(12):1983, 1989.
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, 2020 (11):113301, 2020.
Apostolos P Georgopoulos, Andrew B Schwartz, and Ronald E Kettner. Neuronal population coding of movement direction. Science, 233(4771): 1986.
Federica Gerace, Diego Doimo, Stefano Sarao Mannelli, Luca Saglictti, and Alessandro Laio. How to choose the right transfer learning protocol? a qualitative analysis in a controlled set-up. Trans- actions on Machine Learning Research, 2024.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256. JMLR Workshop and Conference Proceedings, 2010.
Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dy- namics of stochastic gradient descent for two-layer neural networks in the teacher-studens setup. Advances in neural information processing systems, 32, 2019.
Sebastian Goldt, Marc Mézard, Florent Krzakala, and Lenka Zdeborová. Modeling the influence of data structure learning in neural networks: The hidden manifold model. Physical Review X, 10(4):041044, 2020.
Ian Goodfellow, Mehdi Mirza, Da Xiao, Aaron Courville, and Yoshua Bengio. An empiri- cal investigation of catastrophic forgetting in gradient-based neural networks. arXiv preprint 132.611112..621,, 2013.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. arXiv preprint arXiv:1909 10893, 2019.
Irina Higgins, Loic Matthey, Arka Pal, Christopher P Burgess, Xavier Glorot, Matthew M Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. ICLR (Poster), 2017.
John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 295522554-2558. 1982.
Ronghang Hu, Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Kate Saenko. Learning to reason: End-to-end module networks for visual question answering. In Proceedings of the IEEE International Conference Computer Vision, pp. 604-8133 2017.
Ronghang Hu, Jacob Andreas, Trevor Darrell, and Kate Saenko. Explainable neural computation via stack neural module networks. In Proceedings the European conference computer vision (ECCV), pp. 53-69, 2018
Alumit Ishai, Leslie G Ungerleider, Alex Martin, and James V Haxby. The representation of objects in the human occipital and temporal cortex. Journal of cognitive neuroscience, 12(Supplement 2):35-51, 2000.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kerne:: Convergence and gen- eralization in neural networks. Advances in neural information processing systems, 31, 2018.
Anchit Jain, Rozhin Nobahari, Aristide Baratin, and Stefano Sarao Mannelli. Bias motion: Theo- retical insights into the dynamics of bias in sgd training. arXiv preprint arXiv:2405. 2024.
Devon Jarvis, Richard Klein, Benjamin Rosman, and Andrew M Saxe. On the specialization of neural modules. In The Eleventh International Conference Learning Representations, 2023.
13

---

Published as a conference paper at ICLR 2025
Devon Jarvis, Richard Klein, Benjamin Rosman, and Andrew M Saxe. Make haste slowly: A theory of emergent structured mixed selectivity in feature learning reLU networks. In The Thirteenth International Conference Learning Representations, 2025. https:://ppeneeviee
net dd= 85x
Hyunjik Kim and Andriy Mnih. Disentangling by factorising, 2019. URL https://arxiv.
org/aba/1800 05983
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska. et al. Overcom- ing catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114(13):3521-3526, 2017.
Abhishek Kumar, Prasanna Sattigeri, and Avinash Balakrishnan. Variational inference disentan- gled latent concepts from unlabeled observations, 2018. URL https://acxvv.orr////
1711. 00848
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
06158
A.K. Lampinen and S. Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. In T. Sainath (ed.). International Conference on Learning Representations, 2019. ISBN 0311-5518. doi: 10. 00.100008111511988811999999 URL http: L/axxiv.org/abs//
10374
arXiv: 1809.10374.
Yann LeCun, Bernhard Boser, John S Denker, Donnie Henderson, Richard Howard, Wayne Hub- bard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 1(4)5441-551, 1989.
Jin Hwa Lee, Stefano Sarao Mannelli, and Andrew Saxe. Why do animals need shaping? theory of task composition and curriculum learning. Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024, 2024. URL https:
//openrevient net. httt:/formmmi-=SDDPpC. d=SDDPEE7t
Sebastian Lee, Sebastian Goldt, and Andrew Saxe. Continual learning in the macherrstudent setup: Impact of task similarity. In International Conference on Machine Learning, pp. 6109-6119.
PMLR, 2021.
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 arXiv:2205.09029 2022.
Michael Lepori, Thomas Serre, and Ellie Pavlick. Break it down: Evidence for structural compo- sitionality in neural networks. Advances in Neural Information Processing Systems, 36:42623- 42660, 2023.
Zhizhong Li and Derek Hoiem. Learning without forgetting. IEEE transactions pattern analysis and machine intelligence, 40(12):2935-22)) 29452,29772222 2017.
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In international conference on machine learning, pp. 41 4114-4124 PMLR, 2019.
Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael Tschannen. Seakly-supervised disentanglement without compromises, 2020. URL https:
//arxiiv.org/bb///002222.. 02886
Saeed Masoudnia and Reza Ebrahimpour. Mixture of experts: literature survey. Artificial Intelli- gence Review, 42(2):275-293, 2014.
Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerehner. dsprites: Disentanglement testing sprites dataset, 2017.
14

---

Published as a conference paper at ICLR 2025
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 Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape two- layer neural networks. Proceedings of the National Academy of Sciences, 11(33)):77665--767), 2018.
Genman Parisi, Ronald Kemker, Jose L Part, Christopher Kanan, and Stefan Wermter. Continual lifelong learning with neural networks: A review. Neural networks, 113:54-71, 2019.
Quian Quiroga, Leila Reddy, Gabriel Kreiman, Christof Koch, and Itzhak Fried. Invariant visual representation by single neurons in the human brain. Nature, 435(70455 2005.
Vinay V Ramasesh, Ethan Dyer, and Maithra Raghu. Anatomy catastrophic forgetting: Hidden
representations and task semantics. arXiv preprint arXivv2007.07400. 2020.
Roger Ratcliff. Connectionist models of recognition memory: constraints imposed by learning and forgetting functions. Psychological review, 97(2))285, 1990.
Grant M Rotskoff and Eric Vanden-Fijnden Neural networks interacting particle systems: Asymptotic convexity of the loss landscape and universal scaling of the approximation error. stat, 1050:22, 2018.
David Saad and Sara Solla. Dynamics on-line gradient descent learning for multilayer neural networks. Advances in neural information processing systems, 8, 1995a.
David Saad and Sara A Solla. On-line learning in soft committee machines. Physical Review E, 52 (4):4225, 1995b.
Andrew Saxe, Shagun Sodhani, and Sam Jay Lewallen. The neural reduction: Dynamics of abstraction in gated networks. In International Conference on Machine Learning, pp. 19287- 19309, PMLR, 2022.
Andrew M Saxe, James McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynam- ics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013.
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 123):11537-11566, 2019.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv: wxXiv:1701.06538 2017.
Paul Smolensky, R Thomas McCoy, Roland Fernandez, Matthew Goldrick, and Jianfeng Gao. Neu- rocompositional computing: From the central paradox of cognition to a new generation ai systems. arXiv preprint arXiv:2205.01 2022.
Steven C Suddarth and Kergosien. Rule-injection hints as means of improving network perfor- mance and learning time. In European association for signal processing workshop, pp. 120-129. Springer, 1990.
Salma Tarmoun, Guilherme Franca, Benjamin D Haeffele, and Rene Vidal. Understanding the dynamies of gradient flow in overparameterized linear models. In International Conference on Machine Learning, pp. 10153-10161. PMIR, 2021.
Adly Templeton. Scaling monosemanticity: Extracting interpretable features from claude sonnet.
Anthropic, 2024.
Aditya Vardhan Varre, Maria-Luiza Vladarean, Loucas Pillaud-Vivien, and Nicolas Flammarion. On the spectral bias of two-layer linear networks. Advances Neural Information Processing Systems, 36, 2024.
15

---

Published as a conference paper at ICLR 2025
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Anna Korhonen, David R. Traum, and Lluís Märquez (eds.), Proceedings of the 57th Conference of the Association for Computational Linguistics, ACI. 2019, Florence, Italy, July 28- August 2, 2019, Volume I: Long Papers, pp. 5797-5808. Association for Computational Linguistics, 2019. doi: 10.18653/
V1/P19- 1580. /tLb::t:: ttt.:/ii.ooogg////.. htt.:/...1631111///
1580
Yizhou Xu and Liu Ziyin. When does feature learning happen? perspective from analytically solvable model. arXiv preprint arXiv:2401.07083, 2024.
Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson. How transferable are features in deep neural networks? Advances neural information processing systems, 27, 2014.
Matthew D Zeiler and Rob Fergus. Visualizing and understanding convolutional networks, In Computer Vision-ECCCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part I 13, pp. 818-833, Springer, 2014.
Friedemann Zenke, Ben Poole, and Surya Ganguli. Continual learning through synaptic intelligence. In International conference machine learning, pp. 3987-3995, PMLR, 2017.
A HYPERBOLIC-LINEAR DYNAMICS
For convenience we will derive the general dynamics here as it requires less notion. Consider linear network performing a regression task with one hidden layer computing output Y = hW X in response to an input batch of data X, with n datapoints, and trained minimize the quadratic loss using gradient descent:
L(W,h) = Σ3 lI AWx;1 1
This gives the learning rules for each layer with learning rate as:
AW myhh" (Σ hWE*);
Ah mm hWE"
These equations can be derived for batch of data using the linearity of expectation, where y E[X x is the input correlation matrix and V x"" is the input-output correlation matrix, as follows:
AW (WWh))
2 (Y; (Y: AWX,)
dW
dw
n M hT - IWX,) = mnh (E(, -NEE(XXX())
mnh (2 hWE"
d Ah "1 dh
L(W, h) 1) dh Σ
d
"
1
(Y, AWX,) -W
2
Σ ( XWX xx)"" (m = XW
nyn m F0
WWE WT
By using small learning rate and taking the continuous time limit, the mean change in weights is given by:
d W = =1²² 1E AWE");
d
dl # (E -WWE W
di
16

---

Published as a conference paper at ICLR 2025
where T Ps is the learning time constant. Here, t measures units of learning epochs. It is helpful note that since we using small learning rate the full batch gradient descent and stochastic gradient descent dynamics will be the same.
Saxe al. (2019) has shown that the learning dynamics depend on the singular value decomposition of:
E" USV
aau on
=== = =VVVVV
Sau o
Here ry and Tn. denote the ranks of the matrices. To solve for the dynamics we require that E and E. are mutually diagonalizable such that the right singular vectors V of Ene are also the singular vectors of " We verify that this is true for the tasks considered in this work and assume it to be true for these derivations. We also assume that the network has at least Ty hidden neurons (the rank of Σne which determines the number of singular values in the input-output covariance matrix) so that it can learn the desired mapping perfectly. If this is not the case then the model will learn the top ny singular values of the input-output mapping where nh is the number of hidden neurons (axx et al. (203). To ease notation for the remainder of this section we will use nı to denote both the number of hidden neurons and rank I S and D then are diagonal matrices of the singular values of the input-output correlation and input correlation matrices respectfully.
We now perform a change of variables using the SVD of the dataset statistics. The purpose of this step decouple the complex dynamics of the weights of the network, with interacting terms, into multiple one-dimensional systems. Specifically we set:
h= UKR:;
W RWVT
where R is an arbitrary orthogonal matrix such that Rd R = I. Substituting this into the gradient descent update rules for the parameters above yields:
d
W =h" AWE")
dt
(RWV's ") =RhU (USVT- - UUR RWV*VDVV)
dt
d
(RWVT) -Rh(SVT - -WWDVT)
di
7 d
W -TiS hWD)
dt
and
T d
h = "w - WT
dt
(UhRT) -USSV -UNR RW RVVTVDVTTVWN' dl
(UTRR) -(US - UNWDWWR
dt
d T TT - hW D)
dt
Here have used the orthogonality of the singular vectors such that VTV 7 and UTU I. Importantly, all matrices in the dynamics are now diagonal and represent the decoupling of the network into the modes transmitted from input to the hidden neurons and from hidden output neurons. In practice we do not initialize the network weights to adhere to this diagonalisation and so it is not guaranteed that the matrices will be diagonal at initialisation. However, empirically has been found that the network singular values rapidly align to this required configuration (Saxe et al.) 2013 2019).
The derivative then for the full-network input-output mapping can be obtain by using the product rule:
d
hW =(T h)W
+ h(T dt w) = (W(s - hWW)) W (T(S hWD))
dt
dt
17

---

Published as a conference paper at ICLR 2025
+ (S hWD) (W + h² (S - hWD)
This means that at a minimum: 5 S-hWD or DW W and h. As result we can use the change variables: W vAsinh and parametrized by 0.
h. This defines a hyperbolic space between
vA 2
We note that there is conserved quantity between the singular values of the weight matrices:
(P==(VAsiml) (VAsinhg -~=i==(Assinn 9) - -(VIcosh 2 2 =
W h²-h²= -
This is known X-Balanced weights (Kunin et al. 2024) and for given initial value for À this quantity will be conserved for all times during training. Aiming to write the network dynamies in terms of this quantity to understand its effect on learning speed and initialisation and with the change of variables to hyperbolic coordinates we begin with:
(w² + +h²) = =(W²)² + (Th3)
W 2
Substituting this into the network dynamics equation and defining the network singular value as = =/W we obtain:
d
W + +h (S wD)
di w
w = ( - 32 4W (S- WD)
dt
Now applying the change of variables to hyperbolic coordinates with W v V and V & parametrized by 0:
0 т A cosh sinh
d
0 - --V(Assini)
()= (A (Acosh? ) + 4A2(cosh,, sinh ) A cosh 0
sinh D)
dt 2
2
2
2 2
We can then apply the identities: cosh @ sinh 0 = sinh 0 and Asinh? 0 > cosh 2 = A:
2
dt sinh (0) X2 + =422( 1
d A
sinh (0))² sinh (0)
r dt 2
sinh (0) = | cosh (B) 2
sinh (0)
Now applying the derivative on the left:
A
2 cosh (0) 0 \ cosh (0) X sinh ())
(11)
dt
d 0 Sgn AD sinh (0))
(12)
dt
(13)
This is separable differential equation in
80 (2S - AD sinh (0) d9 sgn(A)
9; 1
dt
T
0
log (/25 tanh (2) + V X2D + >D( log (1 (|2S Lanh (2) + + XD
sgn sgn(A)
4 + D
80
18

---

Published as a conference paper at ICLR 2025
25 tanh
41 2
+
V X² + AD
log
V45a X2
25 tanh
0/
V
XDD2 AD
2
22 tanh
(8u) 2
+
v
x )
+
sgm(1))
log
22 tanh
V45²
7
If we let:
2S tanh
(80
+ V452 + X2D2 AD
C
K V X2D2 X2D2
2S tanh
%0
V452 + X2D2 + AD
2
then:
1
25 tanh
14
+ + AD
sgn(A)
log
log(C))
K
25 tanh
+ AD
We can further simplify this expression by writing y in terms of
0,
sgn(A)K
01
25 tanh
+ K + AD -Cep
(K 25 tanh
AD)
2
2
-K (1 11-Cexp Cexp (sn(())
AD (1 Cexp
sgn(A)
[
tanh
2
28
+ Cexp >gn(A)K
7
K (Cexp) sgm(A))
-1) - -AD AD (Cexp
(gn(A)) +
8 tanh -1
2S
C expp
(gg(A)K( +
To obtain the dynamics for the singular value of mode the network we use:
0
=Asinh
2 cosh
A
sinh
A
2
K (Cexp
sgm(A) -1) AD (Cexp
sinh 21anh-1
2S
C Ceep (gm(A) +1)
In the 1-dimensional case studied in Sec.
3 this equation becomes:
lo (cexp)
sgn())A
-1 11
Ad (cexp" egn(())k +1
> sinh
w(t) 2
2 tanh-
(14)
2s cexp commhik ) +
with:
2s tanh
00
+
4
2
12d2 + Xd
/45² + X²d2
c
2s tanh
80
x2d + >d
With the linear network dynamics we can now derive a network's hitting time (L*) for each mode. Let vn be a sufficiently small value, using Eq. 14 on the relation $
D w = v² we obtain
K (Cexp agn(A)K) 1) -1) AD (Cexp
(amM()) +1)
1 tanh ( sinh -1 12
25 2Du
AD ))
2S Ceep agn(NN 1) +
19

---

Published as a conference paper at ICLR 2025
Let T tanh sinh ( -2==2"=)) then
K (Cexp
com(A)K)
-1 - AD (Cexp
)9m(A)K
1) +1)
T
Cexp (Sn(()K +1
28
sgn(X)K
+1)= =K x(Ceepp sgn(X)K
-1) AD (Cexp) (ggn()) 1) +1)
25T
Cexp
T
T
Sgn(X)K
exp
(2ST C- KC + XDC) -2TT" --AD AD
sgm
(
-2ST* K AD
= log
T
2ST - KC + ADC
T
2ST + AD
log
sgm(A)K
C
Similarly we derive the escaping time for each mode with sufficiently small as:
w
K (C exp
sgn(A]K,
-1) AD (Cexp
sgn(N)K +
X
sinh
tannh
25
C Cexp sgn(2) 1)
+
K (Cexp) (2mm()) 1) -1) -AD (Cexp
+ +1) T
tanh ( sinh -1 (20)
20
25 Cexp (om(A)(k 1) +1)
1
Let f = tanh (2 sinh -1 (20)) then
K (Cexp (gg(()( t) -1) AD (cexp *gn())K t) +1)
T
2S
Cexp (gn())( t) +1)
2s (Coxp
(sgn(A)K T
1) +1) = (Cexp (sgn()) t -1)
- -AD (Cexp
(sgn(A)K&
+)+1
T
Thus, the escaping time can summarised as:
T log
+ + AD
(15)
sgn(A)K
KC -C--CCC ADC
with the escaping time constant:
sinh (20
= tanh
(16)
(
Similarly the hitting time is summarised as:
K + 2ST +DD
sgm (A)K log KC-227 2ST XDC
(17)
with the hitting time constant:
T = tanh( ( sinh -1 25 2Dv*
(18)
AD
Fig. 8 depicts the accuracy of these closed-form equations.
20

---

Published as a conference paper at ICLR 2025
a)
b)
simulation
Network Singular Value
100
1.0
10000
predicted
1.0
80
Quadratic Loss
8000
60
simulation
0.6
6000
0.6
predicted
40
0.2
0.2
20
0
1000
2000
3000
4000
1000
3000
Epoch number
Epoch number
Figure 8: Comparison between the predicted and simulated linear network dynamics. depicts the singular value trajectories for varying levels of weight imbalance (depicted in various colours) while gray represents the corresponding predicted trajetectories. derived from the singular values. We see exact agreement between the simulations and predictions. Vertical bar depict the escaping time (on left) and hitting time (one right).
The same plot for the loss dynamics
B METHOD FOR WEIGHT IMBALANCE PHASE PLOT
Given Eqn. and 17 we can discuss our method for construction of the phase plot in Fig. Bt). For each combination of weight imbalanced for the two pathways we aim to find how close the slower pathway is begin learning its closest point. We note that merely simulating the full network is not enough as this would merely tell us whether the slower pathway learns something, but with no additional precision. Further, we can also constrain our search space over time by noting that the slower pathway will never be quicker than its initial escaping time. Finally, we share the notation in this section with those in the main text and App.AA Anand omit the notation definitions here. Thus, the algorithm for constructing the phase plot, which we reproduce in Fig. 9 is as follows:
Algorithm 1 An algorithm for constructing Fig.
3 J).
Hyperparameters used: S
105, A [0, 1000,,
[0, 20),
5.0, e*
1.0, 1E
Require: s. 2, c², A₁ and A2
Require: a(0) > 0
Phase = [xAalx
for A₂ in A₃ do
for A2 in A2 do
if A2 then
break
end if
b(0): V(1 + ag)
b(0)2 v( + ag)
arcsin (2a(0))(0))/>>
0g arcsin (2aa())(()))>>))
= HittingLime((I, (aatingiiimm(.............
> Apply Eqn. Apply Equ.
17
w(t Dynamics 01, s, E 0,
for & E [0, do Slow SlowThetasst + Theta Deriu X2, omega(t), T, end for goupled
& Numerically integrate coupled slow dynamics using Eqn. EscapeT Time (wwbSmmeeeeeaass X2, s T,),,T,)) e [0, & Apply Eqn.
15
Phase (1,, Ae) min, min, gooupted (t)
end for end for
21

---

Published as a conference paper at ICLR 2025
22
60
A (Fast)
10
20
0
1. (Slemr)
Figure 9: Reproduction of Fig. 3p). A phase diagram representing how pathways with different initial weight imbalances lead to specialisation. The two axis represent the weight imbalance of the two pathways in our broader network (A2 on the x-axis for the slower pathway and a on the y-axis for the faster pathway). The colour represents how close the slower pathway escaping time at its closest point throughout training (in log scale). We see that the more inbalanced the fast pathway relative to the slower pathway, the more likely the network will specialise. The white region represents when the inbalance is equal or reversed.
reaching its
C MEAN-FIELD THEORY OF THE DYNAMICS
As outlined in Sec 4 the key observation for the mean-field analysis is that the main properties of the learning dynamics can be expressed functions of the order parameters-Fqs. By combining these definitions with the update rules-Eqss 6 7-we can derive closed-form expressions for the evolution of the order parameters, enabling us to track the key observables throughout the train- ing process. In the high-dimensional limit (d + 20), these discrete update equations converge to ordinary differential equations (ODEs), which can be integrated either numerically or analytically in certain cases (Jain et al. (024).. As is often the case in the statistical physics of disordered sys- tems, this approach was first derived non non-rigorously by (Sadd Solla (995b) and Biehl & Schwarze (1995), with later works laying down a mathematical foundation showing concentration of the ODEs Goldt et al. 2020 Ben Arous et al. 2022).
Following these prescriptions, we obtain the update equations in Lee et al. (2021) Let us define the pre-activations of the student and task-t teacher given an input x from task as
X,
W, dd
x,
W T3
x,
(19)
Vd
and denote the difference between the teacher and student predictions by A = h()) ())(((())) - h
(p)). The corresponding ODEs for the order parameters in the limit d - oo given by:
dQu
nh (d' Aa) (p'o)) + ( (d(A)) (20)
dr
dR!!') nhh ( (2) p
km
(21)
dr
dh
(22)
dr
where epoch/d represents continuous time in the high-dimensional limit, and t, E 1,2 denote the task indices. The angular brackets indicate an average over the pre-activations. The pre- activations themselves are centered Gaussian random variables with covariances determined by the order parameters Q, ((1) and T.
These averages can be computed analytically for certain activation functions. For instance, in the case of rescaled error function introduced in the main text (Saad & Solla 1995b Biehl &
22

---

Published as a conference paper at ICLR 2025
Schwarze 1995), the relevant averages given by:
E13
(+(3)((7)) arcsin
(23)
((1 + Σ11) + E22)
(d(())86(n)) 2Eaa + Σ11) -1312213
(24)
VAa(1 + +E1)
4
Ao
aresin
(25)
7²
A
where the Greek letters represent arbitrary pre-activations with covariance matrix E, and the auxil- iary quantities A; are given by:
Ao A4 E2s (1 Eni) E13 + Ema) + E3 +
E12314
(26)
A Aa(1 + Mss) (1 + End 21s + M22) 2112211328.
(27)
A2 A4(1+ + Σ②) 21 En1) 21 (1 + M2)) 22112214222
(28)
A3 (1 En)(1 E33) EZg-
(29)
These expressions provide comprehensive analytical framework for tracking the dynamics of the student network and the evolution of specialisation across training.
D DISENTENGLEMEENT
We conduct our experiments using open-source frameworks Locatello al. (2019); Abdi et al. (019).. Specifically, we implement a beta-VAE with the "DeepGaussian architecture for the decoder and "DeepLinear" for the encoder. We modify the Xavier initialisation where the weights of the linear layers will have values sampled from U ("(a,, a) with
a = gain x fan + fan_out
We vary the gain between 0.3 and 3 and run each experiment over seeds. All network parameters are set to their default values as provided by the respective open-source frameworks. We run the experiments for 20 Epochs and 157499 iterations.
These experiment illustrate the impact of initialisation on network specialisation. Although the scope of these experiments is limited, they provide preliminary validation of our theoretical frame- work in more realistic contexts. We advocate for further investigation into alternative initialisation schemes with varying levels of balance. Moreover, we highlight the need for future research to extend these experiments by considering wider variety of datasets (Car3D Du et al. (024)). dSprites Matthey et al.] (2017).,, network architectures (Conv.Lincar). initialisation strategies Gaussian Xavier Initalisation) and different metric (SAP et al. (2088) Higgins et al. (017). to fully explore the implications of our findings.
DCI Disentanglement Eastwood & Williams (0188) define three key properties of learned repre- sentations: Disentanglement, Completeness, and Informativeness. To assess these, they calculate the importance of each dimension of the representation in predicting factor of variation. This can be done using models like Lasso or Random Forest classifiers. Disentanglement is computed by subtracting the entropy of the probability that a representation dimension predicts a factor, weighted by its relative importance. Completeness is similarly measured, focusing on how well a factor is captured by the dimensions. Informativeness is evaluated as the prediction error of the factors. We use the implementation in [ncaaatello et al. (2019). this implementation, we sample 10,000 training and 5,000 test points, then use gradient-boosted trees from Scikit-learn obtain feature importance weights. These weights form an importance matrix, with rows representing factors and columns representing dimensions. Disentanglement is calculated by normalizing the columns of this matrix, subtracting the entropy from for each column, and then weighting by each dimension's 's relative importance.
23

---

Published as a conference paper at ICLR 2025
E ADDITIONAL ENTROPY PHASE DIAGRAMS
In Fig. 5 we showed phase diagrams of the aggregate entropy as function of initialisation param- eters, for both ReLU and sigmoidal networks. In Fig, I0 below, we show additional plots with
the individual entropy terms (Ha defined over the unit activations, and Hn, defined over the head weights).
(a)
0 9/16
x/8
3///16
0 x/4
Ba.
0.8
0.8
0.8
0.8
0.8
ReLU
.0.6
00.6
_0.6
.0.6
L0.6
.35
0.4
0.4
0.2
0.2
0.2
0.2
-3
-3
2
-1
-3
-1
-3
2
log
log
7/16
0= x/8
3 16
m/4
0.8
0.8
0.8
0.8
0.8
ReLU
.0.6
.066
.0.6
.0.6
.0.6
.35
0.4
0.4
0.4
0.4
0.2
0.2
0.2
0.2
0.2
-3
-2
-1
-3
-2
-1
-3
-2
-1
-3
-2
-1
-3
-2
-1
log
log
log
(b)
8 n/16
x/8
3w/ 16
///4
0.5
0.8
0.8
Sigmoidal
.0.6
.0.6
.0.6
.0.6
10.6
.35
0.4
0.4
0.4
0.4
0.4
0.2
0.2
0.2
0.2
0.2
-3
-2
-1
-3
-2
-1
-3
-2
-1
-3
-2
-1
-3
-2
-1
log
log
log
log
log
n//1G
x/8
3// /G
///4
Sigmoidal
0.6
.0.6
.0.6
.0.6
.0.6
.35
0.4
0.4
0.4
0.4
0.4
-3
22
-1
-3
-2
-1
-3
-2
-1
-3
-2
-1
-3
-
-1
log
log
log FW
log
log
Figure 10: Additional Phase Diagrams. Here show the equivalent phase diagrams from Fig. for entropy measures over the unit activations and head weights.
F DIVERSITY OF FORGETTING CURVES
In Fig. 6 we show that initialisation can lead to variety of specialisation profiles in contrast with what observe in previously.
G FORGETTING CURVES IN MNIST
In order to support the findings presented in Sec. 4.2.1 we turn now to continual learning task constructed around the MNIST dataset This dataset has previously been adapted to continual learning benchmarks e.g. most famously in the permuted MNIST task, Here construct slightly different continual learning task to encode a notion of task similarity.
We begin by considering only one half of the 10 class MNIST dataset, such that are left with only data in the first classes. Our first task in the sequence of two tasks consists simply of clas- sifying these 5 digits. Our second task is also to classify 5 digits and ranges from classifying the same 5 digits (maximum task similarity) to classifying new digits, i.e. those that were discarded construct the first task (orthogonal tasks minimum task similarity). In 10 class dataset like
24

---

Published as a conference paper at ICLR 2025
MNIST this gives us only very coarse grip task similarity, but this is enough to robustly elicit behaviour analogous to what we observe in the toy teacher-student models.
We use two-layer, multi-head, feed-forward architecture with sigmoidal activations to mirror the models used in the teacher-student setup. The hidden dimension of our networks needs to be larger to properly learn the classification task; we therefore lose the elegance and control afforded by the polar coordinate initialisations of [Sec. 4.2.1 to vary entropy and scale of initialisations. The method we use to generalise this notion initialisation (e.g. a uniform distribution), and a relatively low entopy initialisation (e.g- Normal or Laplace distribution). is straightforward (but important) to ensure the scale of the samples generated is consistent across this interpolation.
interpolate between two initialisations: high entropy
In Fig. 12, we show forgetting profiles for three different initialisation schemes (analogous those shown in Fig. 6) for the continual MNIST task described above. It is clear that in the case of low entropy and specialisation in the first task along with high entropy second head initialisation, we get behaviour characteristic of Maslow's s hammer. However when we initialise the second head with low entropy, we recover the monotonic relationships found in the equivalent initialisations from the toy models. At this stage these are primarily qualitative results, i.e. we are comparing the shapes of these forgetting profiles and not the relative magnitudes or detailed forgetting metrics.
H SPARSE AUTO ENCODER EXPERIMENT
To further verify the applicability of the linear network theory presented in Sec. we we experimentally verify a prediction from the theory. Sec. 3 finds that networks which are initialised with larger hidden-to-outpu weights compared to the input-to-hiden As mention in the main text, the notion of specialisation in this work is very similar to activation sparsity. As a consequence, we predict that by leveraging an output heavy initialisation scheme we can improve the sparsity of an autoencoder.
weights will have a specialisation benefit.
We conduct the following experiment in two phases: Phase /: We train a standard VAE (similar to Sec. 3.2) on MNIST which initialised with small weights ensure the network is in the feature learning regime (Geiger et al., 2020) (we sample from a Gaussian with standard deviation 0.001). Importantly, the latent dimension of this VAE is smaller than the input and forms an entangled latent space. Phase 2: In a similar manner to the recent approach on the Claude line of Large Language Models (Templeton] 2024). we train a sparse autoencoder (SAE) from the latent space the VAE, with the aim of improving the sparsity and disentanglement of the latent space. In our experiments the VAE has 16 hidden neurons. These 16 neurons become the input (and output) to the SAE The SAE then projects this up to latent space of dimension 2048 which has ReLU activation function. For our baseline, we train the SAE with the typical L2 reconstruction loss and LI regularisation on the hidden activity. For our model, we train in exactly the same manner, except we do not the L. regularisation on the hidden activity. Thus, for our model there is no explicit pressure on the autoencoder embed representations sparsely. For ease we will refer to this model as an implicit Sparse Autoencoder (iSAE). We repeat this process with varying degrees of initialisation imbalance and track the sparsity of the SAE and iSAE. Denoting the hidden layer activity of the networks for the entire MNIST dataset as H we define an indicator function in Eqn. 30 for a single neuron responding to a single datapoint:
[1,
Hy >
HHH))
(30)
0. otherwise
We calculate the sparsity across the dataset as the average number of datapoints the hidden neurons respond to, over the 60000 datapoints:
2048 60000
Σ 1(H.))
(31)
2048
To initialise the layers of the iSAE and SAE we define an imbalance parameter U (note that this is not the same hyper-parameter as the À notation employed in the main text and is defined purely for
25

---

Published as a conference paper at ICLR 2025
practicality in this experiment). The encoder weights are initialised by sampling from Gaussian v. The decoder weights are sampled from a Gaussian with 0.001v. Thus, increases the decoder is initialised with increasingly
with standard deviation σ = 0.001 2 standard deviation large weights compared to the decoder.
The results of this experiment shown in Fig. 13 We clearly from these results that the ini- tialisation imbalance is pushed towards the hidden-to-outpur weights such that they are larger than the input-to-hidden weights, then the sparsity of the iSAE latent space improves dramatically, This corresponds to a positive X-balance in the theoretical results and, thus, our empirical and theoretical results are consistent. This is in spite of there only being an implicit bias towards sparsity. Con- versely the SAE with explicit sparsity regularization does not change in response to varying degrees of initialisation imbalance. Importantly, this provides empirical support for the findings from the linear network dynamics and verifies our prediction resulting from this theory.
26

---

Published as a conference paper at ICLR 2025
(0C3),
((3)) (0,8) 8)
(o(1),
((2)
(1,8)
(o(t), ((2)) (2.8)
(801),00a) (3.8)
(6(1), (())) (4.
Regeting
0.3
Foraatin
Rengrting
0.2
Rigoningg
Raganing
0.2
0.1
0.5
0.5 Similanty
0.5 Simiharity
Similaniyy
Similamiy
Similarity
(0(1), ( (0. 7)
(o(1) g(3) (1. 7)
(o(1) (2, 7)
(ot1), (02)) (3,
(o()) (( (4, 7)
Fagelling
Regening
Regeting
Rogating
0.2
0.1
0.5 Similanay
Similant
Similanay
Similarit
Similanity
(oC3)_ ((2)
(o(1), g(
(o(1), ((2) (2.
(o(1), ((2)
(@t1), g(2) (4.
Fungetiigg
Fogeeting
Fongeting
0-2
15
Fagetting
0.1
0.5
0.5 Similarity
0.5 Similarity
0.5 Similarity
0.5 Similarity
Similarity
(et3) (2 00. 6)
(e(1), (C2) (1, 5)
(0(1), ((2)
(2, 5)
(6(1), (( (3. 5)
(of1), ((2)) (4.5)
Rergeing
Forgeting
Fongnting
Ragenting
0.1
0.1
Similarity
Similariy
Similarity
Similaaityy
Similarity
(of1), (2)) ((2)
4)
(o(3)
() (1.4)
(001) 2( (2,
(0t)) ((2)) 4)
(o()) [) (4,
Furgtting
Furaniig
Fungeting
Forgeting
0.3
0
U.S
0.5
0.5 Similanty
0.5 Similarity
0.5 Simularity
Similariy
Similarayy
(()) (0.3)
(o(*), ((2)
(1.3)
(0(1)_00)) (2.3)
(6(1), ((2) (a. a)
(6(1), (()) (4.
Pogeting
Forattin
0.3
Rogering
0.2
Rongtting
0.5
0.5 Similaniy
0.5 Similarity
0.5 Similority
0.5 Similuity
Similariy
) (0.2)
(o(1)
(1.2)
(2.2)
(o(1), (( (3. 2)
(g()) ) (4.2)
Fargettigg P
Forgeting & 0
Forgetin
Rogettnn
0.1
0.
Similarit
Similanit
Similaniy
Similaity
Similaity
(o(1) ((3))
(o(1), g(2))
(o(L) g(2)) 12.1)
(o(1) ((2)
(o(1) g(2) (4.
Fungeting
Fagetinn
Fongeting
Fagetting
0.5
0.5
0.5
0.5
0.5 Similarity
Similbritt
Similarit
Similanty
Similarity
(0 (())) (0.
(o(*), ((2) (1,0)
(0(1) of
0)
(o(1), (( (3, 0)
(o(1) ((2)) (4.0)
Regeting
0.3
Porgttigg 9.
Regentng
0.2
Regating
0.1
0.1
Similanity
Similarit
Similarity
Similarity
Similarity
Figure 11: Initialisation can lead to a diversity of specialisation dynamics and a diversity of relationships between forgetting and task similarity. H, ow fixed, ((1, ((2) measured incre- ments of T/ 16. Scaled error function, p 1, P
27

---

Published as a conference paper at ICLR 2025
3
Forgetting
2
1
2
Task Similarity
(1) 0.01, (1) o 1,, 0.01
(1) = 0.01, (1) = 0.01, (
= = 0.01
(0) 0.01, (1) 0.01, 0.01, =1
Figure 12: Forgetting profiles on MNIST continual learning problem. Forgetting vs. task sim- ilarity a continual learning task using the MNIST dataset. Here similarity is defined as the the number classes that are the same ond, i.e. corresponds to 5 new classes and 5 corresponds the same 5 classes. The green line achieved by initialising with low entropy and small weights in the first head followed by low entropy and small weights in the second, while the blue and orange lines have low entropy second head initialisations with high and low entropy initialisations in the first head respectively. These forgetting profiles (in terms of their monotonocity patterns) qualitatively match those observed in the theoretical toy problems discussed in Sec. 4.2.1 (see Fig. 6). Note (f)) denotes the scale of the tt head initialisation (equivalent in Fig. 6) and ((()) the relative entropy (plays similar role to
5-way classification problem from the first task to the sec-
Fig.
28

---

Published as a conference paper at ICLR 2025
iSAE (No Regularisation)
2500
SAE (Regularisation)
2000
Sparsity Value
1500
1000
500
-3.0
-0.7 0.0 0.7
3.0 3.7
4.6 5.3
Imbalance (Log Scale)
Figure 13: Implicit regularisation from initialisation imbalance: We track the sparsity of the iSAF and SAE for varying degrees of initialisation imbalance (x-axis). The imbalance the x-axis depicts the natural log of the imbalance parameter (v). Thus, 0.0 depicts balanced initialisation typically used in practice. The y-axis depicts the corresponding sparsity calculated using Eqn. 31 Clearly, as the imbalance increases the sparsity of the iS decreases (which is consistent with the findings from the linear network theory of Sec. } while the SAE does not respond due to its explicit regularisation. Results depict the average over ten runs with two standard deviations on either side of the mean.
29


Reviews:
[{"review_id": "sMVraZhGMe", "reviewer": "Reviewer_SPqY", "summary": "- The authors introduce a theoretical framework that demonstrates how initialization controls whether networks develop specialized or shared representations\n- This work contains a rigorous mathematical analysis demonstrating that weight imbalance between layers drives specialization\n- This paper presents a novel analysis of how initialization affects continual learning methods, particularly demonstrating the dependency of EWC on specialization", "questions": "How would your initialization schemes scale to modern architectures with skip connections?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- This work takes a new perspective on the class catastrophic forgetting problem and provides theoretical justification for why initialization is important\n- I believe continual learning researchers will be interested in this approach and in building on this work", "weaknesses": "- The primary weakness is the lack of empirical experiments. I would like to see different initialization methods compared on several simple continual learning or class incremental learning benchmarks. Overall, I think this work still has value as a theory paper.", "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": "brqFRIDQg7", "reviewer": "Reviewer_CSfH", "summary": "The paper studies how the initialization of a network can affect its ability to “specialize” where representations are localized and disentangled. The paper discusses a simple student-teacher framework with extensions to disentangled attribute learning and continual learning. I think the paper looks solid but it is highly likely that I didn’t understand the paper.", "questions": "See weaknesses", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is informative on the theoretical setup and its implication in empirical experiments such as disentangled learning and continual learning.\n- The conclusion about network initialization is insightful.\n- The figures in the paper are nicely organized.", "weaknesses": "- Clarity: I found the paper not very clear. First it builds off from some prior knowledge on “specialization” and its definition of specialization is “one pathway finish learning (reach it’s hitting time t ∗) before the other begins learning (reaches it’s escaping time tˆ)”. However, “hitting” and “escaping” are two new concepts without definition. And why does the leading time matter for specialization? Can’t a MoE like architecture also specialize? In other words, the two neurons could simultaneously specialize into two different aspects of the input data, where the lead time of one neuron does not matter. Perhaps this is because I am very new to the literature, but without properly understanding the concept of specialization, I cannot have a confident judgment on the paper.\n- Although the paper is theoretical, the main paper does not have any formal theorem/claim to explain the relation between initialization and specialization.\n- The paper lacks a core claim. The first part attributed specialization to the “imbalance” in initialization (I also didn’t understand why h-ww^T is imbalance). The second part focuses on how “gain” of initialization affects disentanglement. And lastly, in continual learning, initialization was mentioned but unclear which part of initialization can cause the failure of continual learning. Basically, it seems that the whole paper is talking about initialization but the actual mechanism is quite different in each part and lacks a unifying explanation.", "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": "TL8jAVNqJ1", "reviewer": "Reviewer_yUsa", "summary": "The paper investigates how weight initialization, specifically imbalance, affect the arising of specialization in the network representation and the subsequent consequences on forgetting wrt. task similarity.", "questions": "a. Conflicting results? - Figure 6 orange is no specialization and with lowest forgetting (good?) but Figure 7 trying to say no specialization is not good? Is specialization good or bad for continual learning? \n\nb. Can the authors give any explicit conclusion on weight imbalance and their impact on continual learning, apart from initialization has an effect? This seems to be the most important message of the paper but I can't find any statement in the text. Eq 9 gives the generalization error but was not explained.  \n\nc. Fig 6 is very confusing. Two factors: layer-wise imbalance (r) vs. pathway imbalance (theta). But they are both changing across scenarios; without controls, I don't see how conclusions can be drawn. \n\nd. Figure 6 blue - why is it 'specialization after the first task' since theta=pi/4 and according to Figure 5 that's yellow = no specialization? \n\ne. Unclear conclusion: For example, Sec 4.3 on continual learning, its conclusion only mentions the concerned factors without making clear their consequences on CL so standing more like an introduction than a conclusion \"In a broader context, a rich diversity of behaviours can emerge, driven by factors such as the initialisation schemes, the scale of weights in the first layer, and the readout heads for both tasks.\"\n\nf. Could the authors somehow materialize the concept of 'specialization' in Figure 2? I could vaguely see how two pathways of different lambda could lead to specialization as one simply wins the race but it was after much back-and-forth comprehension.  \n\ng. Figure 1b main message = activation function is subsidiary to weight imbalance? If so, it seems to be disconnected from the rest of the paper and seems like better belonging to the appendix than as a killer figure. \n\nh. Flow problem: Q and R in Figure 1 are not really introduced until in Sec 4.1 but even then they are only defined without explained (l352).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "Originality: Although specialization and continual learning have been widely studied in the context of modularity, this paper attempts to take a more thorough investigation from the perspective of weight imbalance at initialization which is interesting and an important direction. \n\nQuality: Figure 3&4 provide good insights on how imbalance leads to specialization/disentanglement both theoretically and experimentally. \n\nClarity: Notations are clear and maths are largely clear (except Eq. 9). \n\nSignificance: I think weight imbalance is an interesting and important angle for continual learning thus the paper has great potential.", "weaknesses": "Quality: 1) As mentioned by the authors, simplified setting is a limitation of the paper: all theoretical results are on single-hidden layer networks (linear for learning dynamic, nonlinear for student-teacher); All experiments including the ones on continual learning are on toy examples except Figure 4. 2) I find the overall main message unclear - is specialization good or bad for continual learning?  (see question) \n\t\t\nClarity:  The current flow of the paper can be problematic, it required significant back-and-forth for me to understand the messages (example see questions).  The manuscript seems incomplete, for example the caption for Figure 4.  \n\t\t\nSignificance: Apart from the other mentioned comments, one key aspect the paper could gain on its significance is to make more clear conclusions (see questions).", "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": "LnnhasiQ2d", "reviewer": "Reviewer_i4VX", "summary": "The paper presents an analysis of specialization in neural networks without nonlinearities, with focus on the determinants of specialization and its influence on forgetting.\n\nThe analytical methods of the paper fall outside my expertise hence I cannot give a meaningful evaluation of the paper as a whole as they make up the majority of the argument. The below ratings and comments are limited to a superficial evaluation of the main points and should not be regarded as reliable.", "questions": "Please see weaknesses point above.", "limitations": "", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The analysis and discussion about the relevance of specialziation to forgetting, as well as the determinants of specialization, is an interesting area on inquiry.\nThe analysis is provided with a rigorous mathematical framework.", "weaknesses": "The analysis is performed on linear networks, while nonlinearity is the defining feature of NNs. Authors are clear about that and apparently it is not an unprecedented way to conduct a similar analysis, but this would render the conclusions of the analysis questionable nonetheless. I would suggest the authors to include a more detailed discussion regarding what sort of guarantees can be given when transferring conclusions obtained from linear networks to conventional, nonlinear ones.", "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)