A Theory of Initialisation's Impact on Specialisation

Authors: Devon Jarvis, Sebastian Lee, Clémentine Carla Juliette Dominé, Andrew M Saxe, Stefano Sarao Mannelli
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: machine learning theory, teacher student setup, initialisation, specialisation, statitistical mechanics of 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.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730152327448

thread_type:
official_review

<thread_object>
creator: reviewer_i4vx

<content>
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.

The 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.

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.
The 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.

questions:
Please see weaknesses point above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730681785134

thread_type:
official_review

<thread_object>
creator: reviewer_yusa

<content>
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.

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. 

Quality: Figure 3&4 provide good insights on how imbalance leads to specialization/disentanglement both theoretically and experimentally. 

Clarity: Notations are clear and maths are largely clear (except Eq. 9). 

Significance: 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) 
		
Clarity:  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.  
		
Significance: Apart from the other mentioned comments, one key aspect the paper could gain on its significance is to make more clear conclusions (see questions).

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? 

b. 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.  

c. 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. 

d. 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? 

e. 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."

f. 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.  

g. 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. 

h. 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).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730682415684

thread_type:
official_review

<thread_object>
creator: reviewer_csfh

<content>
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.

strengths:
- The paper is informative on the theoretical setup and its implication in empirical experiments such as disentangled learning and continual learning.
- The conclusion about network initialization is insightful.
- 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.
- Although the paper is theoretical, the main paper does not have any formal theorem/claim to explain the relation between initialization and specialization.
- 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.

questions:
See weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1731274132217

thread_type:
official_review

<thread_object>
creator: reviewer_spqy

<content>
summary:
- The authors introduce a theoretical framework that demonstrates how initialization controls whether networks develop specialized or shared representations
- This work contains a rigorous mathematical analysis demonstrating that weight imbalance between layers drives specialization
- This paper presents a novel analysis of how initialization affects continual learning methods, particularly demonstrating the dependency of EWC on specialization

strengths:
- This work takes a new perspective on the class catastrophic forgetting problem and provides theoretical justification for why initialization is important
- 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.

questions:
How would your initialization schemes scale to modern architectures with skip connections?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732157465814

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to the review

comment:
Thank you to the reviewers for your detailed and thoughtful feedback. We really appreciate the time and effort you put into reviewing our manuscript—your input has been key to improving its quality. We’ve carefully addressed each of your comments and included linked responses for any overlapping points. To summarise the key updates we will make to the revised manuscript:

**Clarity**
We agree with the reviewer’s comments regarding clarity. A key challenge in presenting our contributions lies in integrating insights from two theoretical approaches that are typically not combined, yet are complementary in that they can address each other’s limitations. We have revised all sections highlighted by the reviewers for clarification and corrected all identified typographical errors. We hope these updates will make the paper’s message clearer, as detailed in the individual responses.

**Message**
This paper aims to show that the initialisation has a fundamental importance in achieving a specialised or unspecialised solution. Our paper starts by showing that specialisation is severely impacted by initialisation. Fig.1b serves as a motivation figure and contrasts with results from the literature (Goldt et al. 2019) that concluded that sigmoidal networks specialise while ReLU networks do not. However, we show that ReLU networks can also specialise (and sigmoidal networks do not always specialise) depending on weight initialisations. The rest of the paper then tries to understand in more detail “what” aspects of the initialisation scheme lead to a specialised solution. Initialisation is not a single variable to analyse, and our investigation has to consider concepts like weight imbalance and entropy; as a result, there is not necessarily one single punchline result. Finally, we demonstrate the importance of understanding this effect for downstream tasks, including continual learning. Specifically, our analysis shows that not only can initialisation lead to bad forgetting, but it can also disrupt standard mitigation strategies like Elastic Weight Consolidation EWC (Kirkpatrick et al. 2017).

**Experiments**
We intended this work to be mainly a theoretical contribution that characterises the impact of initialisation on specialisation, beginning in the linear network regime in Sec. 3.1. Subsequently, Section 4 aims to use the new insights on initialisation’s role in specialisation to explain two established forgetting profiles empirically observed in the literature. Namely, the Maslow’s Hammer profile, observed empirically first in Ramasesh et al. (2021), and the monotonic forgetting profile, more typically assumed and observed in Goodfellow et al. (2014). Hence, our work aims to explain empirical results theoretically, and there is merit to our approach without empirics (see, e.g. Lee et al. 2021). Nevertheless, we agree with reviewers on the importance of empirical validations beyond the assumptions of simple models. Indeed, in the paper, we provided some results in the disentangled representation learning problem, Sec.3.2. The aim of this section is to empirically supplement the theoretical results of Sec. 3.1. 

We hope these revisions enhance the clarity and impact of our work and that you find our manuscript a valuable contribution to the field. We look forward to any further feedback you may have!
</content>
</thread_object>

</thread>
<thread>
creation date:
1732483474888

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Discussion period ending soon

comment:
Dear Area Chair and Reviewers,

We have submitted detailed responses to all the reviewers' concerns, along with a general response summarizing the changes made. With only one day remaining in the discussion period, we kindly ask for confirmation on whether our replies have addressed the reviewers' concerns, and, if so, to consider adjusting your score accordingly.

Your feedback is crucial to improving the quality of the paper, and we would greatly appreciate your engagement before the deadline.

Thank you for your time and consideration.

Best regards,
The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1732732378977

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revised Draft from Reviewer Feedback

comment:
We thank the reviewers again for their feedback. An updated version of the manuscript has been uploaded where we *addressed the clarity concerns* of the reviewers and *added two novel experiments on MNIST whose results support our theories*. Here we provide a summary of the findings of the two novel experiments:
1. To support the findings on the forgetting profiles in the student-teacher setting, we turn to a more realistic continual learning task constructed around the MNIST dataset. We show that the forgetting profiles (in terms of their monotonocity or Maslow’s Hammer patterns) qualitatively match those observed in the theoretical toy problems discussed in the paper.
2. The second experiment supports our finding that a large imbalance between hidden-to-output and input-to-hidden weights leads to higher specialization, derived from the linear network theory. As noted in the main text, the concept of specialization in this study closely aligns with activation sparsity. We compared two kinds of autoencoders trained on MNIST: a first. where sparsity was enhanced using the weight imbalance; a second autoencoder where sparsity was induced via a regularisation approach recently used in the large language model literature [1].

Furthermore, to the best of judgement we have provided detailed responses to all reviewers’ concerns and adjusted the revised draft accordingly. For easy comparison, the main changes in the revised draft have been written in blue. In particular Reviewers **SPqY** and **yUsa** asked explicitly for more experiments to support the theory. We believe the additional experiments should address these concerns.

We are committed to making the most of the discussion period. Thus, if the reviewers have any further questions which would strengthen their evaluation or confidence in the merit of our work, we would gladly aim to answer these. Thank you for your time and consideration.

[1] Templeton, Adly. Scaling monosemanticity: Extracting interpretable features from Claude 3 sonnet. Anthropic, 2024.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper looks at how initialisation affects a network's ability to 'specialise' its representation, adopting a theoretical approach, and with an application in continual learning (the impact on forgetting due to task similarity). 

Reviewers agree that the theory is useful and has interesting implications, and the conclusions are interesting. However, reviewers also agree upon a key limitation: most reviewers found it hard to understand the paper, commenting on the paper's overall clarity and message. During the author rebuttal and subsequent discussion, some of the reviewers agree that the paper is now clearer, and I agree. I also think that the authors can still improve readability to focus on the key messages of the paper, and this will further improve the reach of this paper. Some examples have been provided by reviewers, such as Reviewer yUsa: 
- "For example, in the unchanged abstract, instead of saying 'assessing the implications on EWC', maybe say something like 'specialization by weight imbalance is beneficial for EWC to reduce forgetting'. The same applies to defining Q and R. I would encourage the authors to spend more efforts on improving the readability of their paper so that their important message on weight imbalance / specialization / sparsity could actually be delivered to the broader community."