Make Haste Slowly: A Theory of Emergent Structured Mixed Selectivity in Feature Learning ReLU Networks

We thank the reviewer for their time and consideration of our work. We will begin by addressing the weaknesses in order:

1. We kindly direct the reviewer towards the general comment where we discuss the assumption of our paradigms. We thank the reviewer for the question as we recognise that more justification of our assumptions must be added to the text and will strengthen the applicability of our work.
2. We thank the reviewer for pointing out the typos and will fix these on a subsequent draft.

Turning to the critical questions:

1. We kindly direct the reviewer to the general comment. However, to be more specific to the questions: we show clearly for our setting that the addition of depth breaks the assumption and can conclude that this is due to Lemma 4.1 no longer holding. To our knowledge this clear demonstration that multiple layers of ReLU neurons disrupts feature learning is new to our work. Thus, it is possible that other factors could impact the alignment and we provide a framework for similarly exploring these factors. In Appendix B we discuss our views on how future work could make finding the network gates more feasible for complex datasets and architectures. Similarly it is likely that a dataset could be constructed where the common pathway is not learned first. Once again, the framework we provide here could be of use in identifying such cases and providing clear insight into why the behaviour changes. A similar phase shift from linearity to nonlinearity is shown for the XoR task in Section 3. Thus, it appears reasonable that future applications of our paradigm could lend similar insight. Importantly, the only general claims of this work is that for every ReLU network there exists a ReLN - which we prove, and that the ReLN framework could be useful to future theorists and shed light on neural network behaviour more broadly. Indeed, one of the largest strengths of the linear neural network paradigm is that it demonstrates how the dataset structure affects learning.
2. The ReLU network and ReLNs are initialised with different weights (they are still small as per the linear network framework). This can lead to very minor discrepancies especially as the ReLU network is learning to identify relevant nonlinear pathways which the ReLN is provided from the beginning of training. A similar effect where minor deviations in the randomly initialised values of a linear neural network can lead to minor variation from the closed form dynamics can be seen in Saxe et al. (2019). Finally, on Lines 263 to 267 we do not say we see exact correspondence. Indeed we say “near exact” and “excellent agreement”. In the caption of Figure 4 we say that we see exact agreement between the output of the ReLU network and ReLN at three points in time. We do not mean to be pedantic in our wording here but take the accuracy of our claims seriously and believe these to be accurate descriptions of the results depicted in Figure 4. Similarly, for the concern on Figure 6 saying that the ReLN “shares the same characteristics” as the ReLU network we believe there may be a slight misunderstanding here as the full sentence says “shares the same characteristics in variance”. In other words the points along the trajectories where the drops in loss become inconsistent are the same and have a similar degree of variance. We agree that the ReLU network drops in loss are more stage-like. This point is reiterated on Lines 472 where we say “we also see the same variance profiles”. Indeed, our main claim in Section 6 is one Line 468: “we see that the GDLN still informative about the behaviour of the ReLU network”. We will make clearer our discussion around the variance of the trajectories and commonalities between the dynamics in Section 6. We apologise for the misunderstanding we may have caused and thank the reviewer for pointing this out.
3. We thank the reviewer for raising this potential point of confusion and will fix this in a subsequent draft. We direct the reviewer to Lines 252 and 253 of Section 4 in the main text. We say “We use a linear output layer, assuming the model is over-parametrised (a pathway needs to have at least $h$ hidden nodes to learn a rank $h$ effective dataset)”. We will be clear that this is the only consideration for the number of hidden neurons. In addition we direct the reviewer to our revision of the deep linear network paradigm in Appendix A, specifically Lines 808 to 813 where we say: “We also assume that the network has at least $|\Sigma^{yx}|$ hidden neurons (the rank of $\Sigma^{yx}$ which determines the number of singular values in the input-output correlation matrix) so that it can learn the desired mapping perfectly. If this is not the case then the model will learn the top $n_h$ singular values of the input-output mapping where $n_h$ is the number of hidden neurons (Saxe et al., 2014).”

We thank the reviewer for their time and consideration of our work. We will begin by kindly directing the reviewer to the general comment as we discuss all of the noted weaknesses there. Considering the questions:

1. We thank the reviewer for the question as we recognise that this justification of our assumptions must be added to the text and will strengthen the applicability of our work. We kindly direct the reviewer towards the general comment where we discuss Assumption 2.1.
2. We thank the reviewer for raising this and will aim to make the network architecture more clear. We direct the reviewer to Appendix C for the full derivation of the setting and dynamics of Section 3. We also note that Figure 2b and 2c depict the network architecture of the ReLNs which imitate the ReLU network with one hidden layer. Each square in these diagrams corresponds to a layer of neurons as described in Section 2, specifically Figure 1. We will be clear in the main text that the ReLU network has a single hidden layer. Further, if the reviewer could perhaps guide us on why Figure 2b and 2c did not have the desired effect (as these aim to depict the gating structure and make clear what is meant by 2 pathways and 4 pathways) then we would be happy to incorporate those changes and suggestions.
3. In this case it is possible to identify the gates manually. For the larger networks the clustering algorithm can be used to assist with finding the gates (for the tasks we consider in this work). Please see the general comment and Appendix B for more discussion on how to identify gating structures for other architectures and tasks. Importantly, the manner in which the gates are found is not the focus of Section 3. Rather the insight which the ReLN paradigm offers in demonstrating clearly how a ReLU network rapidly changes gating strategy due to subtle change in dataset structure (forming a phase transition), is the primary aim and contribution of that section (as well as introducing the ReLN paradigm in an intuitive setting).

We thank the reviewer for their suggestions and will certainly incorporate these into a subsequent revision. We thank the reviewer once again for their review and assistance in making the utility of our framework clear. In particular we are thankful for the suggestion to elaborate on the nuance of our assumptions as we believe this addition greatly strengthens the broader applicability of our work. We are happy to engage further to address any remaining questions during the discussion period.

We thank the reviewer for their time and consideration of our work. We will begin by addressing the weaknesses in order:

1. The nonlinearity of the ReLU network is exactly why the dynamics are intractable (for general datasets and in the feature learning regime - see our response to Weakness 3 for the meaning of feature learning in this work). We direct the reviewer to Appendix A, and specifically Line 856. Here we are reviewing the deep linear network paradigm and are using the fact that for a linear network $Y =W^2W^1X$ we are able to compute the matrix multiplication up front $W=W^2W^1$ providing the network mapping $Y=WX$ (we note that the dynamics of $Y=WX$ and $Y =W^2W^1X$ are different which is why we must begin the derivation from $Y =W^2W^1X$ and cannot just start with $Y =WX$ - we note this fact on Lines 136 to 139). The ability to resolve $W=W^2W^1$ is due to the associativity of matrix multiplication. However, when we introduce a nonlinearity (denoted as $f(\cdot)$ here) then $Y = W^2f(W^1X)$ and to evaluate $W^2f(W^1X)$ we must know the datapoints $X$. Thus, we cannot summarise the entire network mapping in terms of its singular values any longer. This is why we believe the correspondence we introduce in this work between ReLU networks and ReLNs is useful - it provides a nonlinear network where we are able to map the ReLU network onto individual linear pathways where the linear network dynamics are once again valid.
2. We kindly direct the reviewer towards the general comment for the broader discussion on the scalability and utility of our approach.
3. We will certainly be clearer on this in the background and direct the reviewer to the example in Appendix A reviewing the linear network dynamics. Specifically Lines 905 to 914. To be clear here, the alignment of the network singular vectors to the singular vectors of the dataset correlation matrices corresponds to feature learning in a dense network. The dot product between two vectors provides a measure of the similarity between the vectors. Thus, when a layer of the network aligns its singular vector to a singular vector of the dataset it means that the layer is looking for that feature (higher similarity in the dot product will result in more activation being propagated through the network down the corresponding hidden neurons). This is conceptually identical to a convolution kernel responding to an appropriate feature as it is applied to an image, except the metric of similarity here is the dot product. The terminology of the feature learning regime versus lazy learning was introduced on Lines 42 and 43 and refers to the two identified training regimes for deep neural networks in the literature (Chizat et al., 2019). The lazy regime occurs when the network begins with large weights and the randomly initialised weights are used to learn the task. We operate in the feature learning regime where the randomly initialised network weights change significantly in aligning with the dataset singular vectors. We note that some prior theoretical work such as the NTK only work explicitly in the lazy regime (Jacot et al., 2018; Lee et al., 2019). Thus, the fact that our work applies to the feature learning regime is useful, and in addition justifies Assumption 2.1.2. Near line 143 we will emphasise that this alignment of the network singular vectors constitutes feature learning and thank the reviewer in helping with the clarity of our work.

We thank the reviewer for their consideration of our work and positive review. To address the one weakness, we will certainly make our contributions from Saxe et al. (2022) more clear. We note that we aimed to highlight our main contributions on Lines 70 to 81 and will work to improve these points. To recap briefly and at a high level our contributions are:

• Noting the connection between a ReLU network and a subset of GDLNs which we introduce as the ReLNs and the proof of existence of a ReLN for any dense ReLU network.
• The design of the tasks in Sections 3 to 6 and derivation of the ReLN dynamics in these settings.
• The subsequent insight which follows into the behaviour of ReLU networks, and the source of structured mixed-selectivity in these cases.

To answer the questions:

1. For Question 1 we direct the reviewer towards our general comment above. We do emphasise that if the assumption is violated there are means of obtaining a dynamics reduction rather than a closed form solution, which is a strategy used in GDLNs (Saxe et al. 2022). Thus, there are strategies for mitigating the effect of breaking the assumption and continuing the analysis with some loss in tractability and interpretability. For many cases, however, this will be sufficient. We thank the reviewer for the question as we recognise that this justification of our assumptions must be added to the text and will strengthen the applicability of our work.
2. We will certainly add a hyper-parameter table in the appendix. In addition, the hyper-parameters are discussed more in Appendix A. For example, the initialization variance and learning rate need to be small and the number of hidden units must be greater than the rank of the input-output correlation matrix. The small learning rate is needed to take the continuous time-limit when solving the dynamics (Line 795) and small initialization variance ensures that the two weight matrices begin with roughly balanced singular values (Line 856). The number of hidden neurons is discussed on Lines 811 to 814 and this just needs to be large enough to have sufficient rank to learn the dataset. Importantly, as long as the same hyper-parameters are used for both the ReLU network and ReLN, and these conditions of the linear dynamics are maintained, then the dynamics in Sections 3, 4 and 5 will match. Thus, the hyper-parameters are fairly arbitrary. We still mention them for reproducibility.
3. We thank the reviewer for raising this point and will clarify it in a subsequent draft. Modularity in this work refers to the architectural property where a certain set of hidden neurons are responsive to the same set of particular stimuli. When multiple modules exist they can be composed. For example, take a module which can identify red objects and another which can recognise squares. Through their composition (by using both modules) we can identify red squares. By strict modularity we refer to modules of this nature where an individual module can also be used in isolation. In other words it is not coupled with any other modules. In the case of the context specific pathways through our ReLNs in Sections 4 to 6 none of the pathways can be used in isolation - they all require another contextual pathway to be active at the same time to obtain the appropriate output response. This is in spite of the fact that the pathways form clear and interpretable subnetworks which respond consistently to a given set of contexts. We will clarify that this decoupling is a property of interest to us and the literature (Andreas et al., 2016; Andreas, 2018) and thank the reviewer for helping to improve the clarity of our work.

We thank the reviewer again for their review and assistance in making the utility of our framework clear. We are happy to answer any further questions or concerns the reviewer may have.

The point where both assumptions, from Assumption 2.1, are used in the deep linear dynamics framework can be seen in Appendix A from Line 824 to 841. Firstly, we note from the equations that the assumptions are independent of the task and necessary for the deep linear network framework as a whole (Saxe et al. 2014). From these equations it is also clear what will occur when the assumptions fail:

For Assumption 2.1.1 on the mutual diagonalisability of the input-output and input correlation matrices:

• If this does not hold then $\Sigma^{yx} = US\hat{V}^T$ and $\Sigma^x = VDV^T$. Note that now $\hat{V}$ and $V$ are different matrices. In this case the update equations in terms of the singular vectors will be: $\tau \frac{d}{dt} \bar{W}^1 = \bar{W}^2(S\hat{V}^TV - \bar{W}^2\bar{W}^1D)$ and $\tau \frac{d}{dt} \bar{W}^2 = \bar{W}^1(S\hat{V}^TV - \bar{W}^2\bar{W}^1D)$. Thus, there is now a matrix $\hat{V}^TV$ relating the two correlation matrices. It is important to note that these equations are still highly useful reductions to the gradient descent dynamics and the matrix $\hat{V}^TV$ is potentially interpretable (especially in comparison to interpreting weights of the network). However, from the perspective of the deep linear network framework we can no longer obtain fully closed form equations for the training dynamics of the network. Secondly, the exact features learned by the linear network are no longer as clear. Thus, there is merit toward aiming to maintain the full tractability and interpretability of the linear dynamics. For the GDLN framework specifically, the violation of Assumption 2.1.1 is even less severe. For the GDLN framework the closed form dynamics are not always obtainable (in Section 4 we use the reduction dynamics, while in Section 5 we use closed form dynamics and we derive these cases in Appendices F and G) and in fact we will still be able to derive the neural race reduction to the same point shown in Equation 12, just with the added term of $\hat{V}^TV$. Thus, some loss of interpretability for the linear pathways is the consequence of violating the assumption. To summarise - two of the primary strengths of the deep linear network paradigm is its tractability and interpretability. Assumption 2.1.1 is key for both properties. However, if it is violated we are still able to make very meaningful progress in deriving a dynamics reduction in a similar manner to the original GDLN paper and in many cases this is likely to be sufficient (Saxe et al. 2022). In the revised draft we will be clear in the main text that Assumption 2.1.1 is needed for full tractability and analysability but there are means of making clear progress towards a reduction if it is violated. We will also add a more in-detail discussion to Appendix A where the assumptions are introduced in the linear network dynamics review. We believe these two corrections will greatly add to the clarity, but also ensure that we do not make our framework overly restrictive or prescriptive. We thank the reviewers for their assistance with these points.

For Assumption 2.1.2:

• A failure of this assumption can lead towards variability in the loss trajectory of the model. We clearly demonstrate such a case in Section 6. Further, if this assumption is violated then a reduction, while possible, is still quite complex as seen on Lines 824 to 842. However, in spite of a violation of this assumption a GDLN is still interesting from two perspectives: Firstly the GDLN will likely shed light on some difficulty involved in full feature learning (the network fully aligning to the dataset correlation matrices). Secondly, the GDLN is still able to be used as a comparative model or baseline for the ReLU network - particularly as it is more interpretable with the explicit gating pattern. Finally, Assumption 2.1.2 is necessary as we aim to study feature learning. Interesting cases can arise where there is significant movement towards alignment in feature space but not convergence (we believe the demonstration of such an effect in Section 6 and the insight it lends to a ReLU network is a contribution of this work), however, if this assumption is severely broken (the modes align very little or not at all) then this would mean that the analysis is no longer considering feature learning. This defeats the original purpose of our proposed framework - to begin to address the gap in theoretical paradigms for finite size, feature learning, nonlinear (with ReLU activation) models. As we motivate in the Introduction, we believe finite size, feature learning, ReLU network are of practical interest to the field and should be considered theoretically, which would necessitate Assumption 2.1.2 or something of a similar nature. We will clarify the purpose of the assumptions in the main text and add a similar discussion to the one here into our review of the linear dynamics framework. We thank the reviewers again for assisting with the clarity of our assumptions.

The limitations of Assumption 2.1:

• Assumption 2.1 is necessary for the full tractability and interpretability offered by the deep linear network dynamics. However if it is violated we can still continue the derivation to a dynamics reduction similar to the GDLN framework, which is insightful and likely to be enough for many cases. The fact that Assumption 2.1 being violated does not stop a useful derivation being obtained was not clearly mentioned and will be added to a revised draft.

The limitations of Assumption 2.2:

• In the context of the deep linear network paradigm, Assumption 2.2 is equivalent to saying that the network is successfully feature learning. In Section 6 we show a case where the vectors do not align fully and the consequences this has. The network on Section 6 is certainly doing feature learning as the vectors are still aligning to a significant degree, but from the linear network dynamics perspective there is an element of lazy learning. Being able to provide dynamics which can accommodate imperfect feature learning of this nature without relying on high dimensional limits is an important direction of future work. However, all theories of feature learning will have a similar assumption. We will add this description to the background as it was omitted. This also aids in clarifying that this alignment is our notion of feature learning.

The absence of larger scale empirics:

• Mixed selectivity is an established phenomena in the machine learning and neuroscience literature, especially with recent interest in similar concepts like polysemanticity (Olah et al., 2017a; Lecomte et al., 2024). In our case distilling this complex phenomenon into a parsimonious setup where theory is manageable is a contribution. Thus, the tractable setting reproducing empirical phenomena is a strength. We will be clearer in Section 3 that mixed selectivity is an established concept but we provide a mechanism from which mixed selectivity may emerge. Specifically, node reuse speeds up learning and can be used to quickly minimise loss. We will also draw this conclusion more clearly from Equations 3 and 6 (shown by the summation over pathways).

The difficulty of findings gates:

• While we provide an algorithm to find the gates in the tasks considered in this work, we agree that the algorithm we present is simple. Hence we do not claim this to be a large contribution of our work. We include the algorithm to demonstrate our recognition of its importance and provide direction for future work. We do note that this is an advance on the GDLN framework which assumed gates were given as part of the dataset (Saxe et al. 2022). We will move some of the discussion on finding gates from Appendix B into Section 7 and discuss in more detail in the main text how gating structures may be found in our framework.

The scalability of our paradigm and experiments:

• On the paradigm generality: The assumptions and limitations of our work are common in prior, accepted theoretical work on the deep linear dynamics paradigm which have been presented at high-impact venues (Saxe et al. 2014,Saxe et al. 2019,Saxe et al. 2022,Jarvis et al. 2023). More importantly the simple datasets which they study have all led to valuable insight into larger models. The same is true of all other paradigms of theory we note in the Introduction, which have various limitations but equally have been highly insightful. We argue that the same will be true of our paradigm, especially considering that our paradigm satisfies an explicit set of assumptions not possible with any other paradigm.
• On the experiment generality: The tasks we consider here are motivated by prior work which has proven insightful to the field. The extended XoR task we present is a generalisation of the seminal XoR task which is of broad interest to the machine learning community. More importantly it shows a clear phase transition in the use of nonlinearity by a ReLU network on a highly interpretable domain. This makes it ideal for introducing the paradigm and this behavioural phase transition resulting from a gradual change in dataset structure is an insight which to our knowledge has not been shown before. Similarly the contextual task builds on prior work. The individual output structures are all from Saxe et al. (2019) - which shed light on semantic learning in artificial neural networks and children. The input structure, while simple, was based on the manner context was incorporated in older connectionist models (Rogers et al., 2004), but also more recently with larger deep networks (Sodhani et al., 2021). Ultimately, the datasets presented in Sections 4,5 and 6 are significantly more complex and realistic than any previous datasets used in the deep linear network paradigm. Similar to adding more information on mixed-selectivity in the literature contextualising our results and methods, we will elaborate more on these points in the main text of Section 3.

To begin, we kindly direct the reviewers towards Lines 119 to 123, as well as Section 7 and Appendix B for more discussion on obtaining the gates for our model. Specifically, Lines 119 to 123 make clear that using a fixed gating strategy is motivated by the original GDLN framework (Saxe et al. 2022). Section 7 then clearly states that finding the gating structure is a limiting open problem but emphasises the utility of our paradigm as it is. Finally, Appendix B provides a simple algorithm for taking a step towards addressing this limitation, but also points to meta-learning as a future direction for addressing this limitation. Appendix B also notes the complexity which meta-learning brings and the possibility to learn gating patterns unavailable to a ReLU network. We argue that for these reasons, a proper treatment of meta-learning should be left to future work, placing it outside the scope of our present study. Furthermore, Section 6 demonstrates that we are sufficiently testing the paradigm and even shows an interpretable case where the paradigm is not able to fully explain the complexity of training a ReLU network. The tractability and control of our setting even make it clear that the addition of more depth is the cause of the networks difficult to completely learn the correlation-aligned features and the mechanism is the violation of Lemma 4.1. Thus, even in this case our model lends insight into the behaviour of nonlinear neural networks and points towards helpful directions of future work. Additionally, the tasks considered in Sections 4 to 6 are significantly more complex than any tasks previously considered in the deep linear network paradigm. In conclusion, we believe scaling to more complex settings appears premature, but we still provide a step forward for the deep linear network paradigm of theory and for theoretical approaches towards obtaining full training dynamics of neural networks in general.

We conclude with a final point on the broader utility of our findings made with the ReLN framework - mixed selectivity and similar phenomena have been identified in multiple, much more naturalistic and large-scale settings (Olah et al., 2017a; Lecomte et al., 2024; Locatello et al., 2019). As we mention in our discussion, mixed-selectivity is highly relevant (if not overlapping terminology) with polysemanticity (and monosemanticity) which is a concept that has garnered much attention recently with the massive LLM Claude (Templeton, et al. 2024). We draw this comparison to demonstrate that our findings are of interest to the broader machine learning community and not distinct from other large-scale experimental findings. We cite more works in both machine learning and neuroscience (which also has a long history discussing mixed-selectivity and its emergence in neural codes) which our work speaks to in the discussion (Anderson, 2010; Rigotti et al., 2013). In conclusion, while the tasks presented in this theoretical work necessarily remain tractable for analysis, in the context of the literature we believe our work provides a helpful contribution to a broader discussion of far more general applicability and impact to both machine learning and neuroscience.