Make Haste Slowly: A Theory of Emergent Structured Mixed Selectivity in Feature Learning ReLU Networks
Authors: Devon Jarvis
Richard Klein
Benjamin Rosman
Andrew M Saxe
Venue: ICLR 2025 Poster (2025)
Keywords: Gated Deep Linear Networks
Feature Learning Dynamics
Structured Mixed Selectivity
ReLU Networks
Primary area: learning theory
TL;DR: We establish an equivalence between finite feature learning ReLU networks and Gated Deep Linear Networks to analyse the full learning dynamics of the ReLU network. We find a bias towards structured mixed selective representations on a set of tasks.

Abstract: In spite of finite dimension ReLU neural networks being a consistent factor behind recent deep learning successes, a theory of feature learning in these models remains elusive. Currently, insightful theories still rely on assumptions including the linearity of the network computations, unstructured input data and architectural constraints such as infinite width or a single hidden layer. To begin to address this gap we establish an equivalence between ReLU networks and Gated Deep Linear Networks, and use their greater tractability to derive dynamics of learning. We then consider multiple variants of a core task reminiscent of multi-task learning or contextual control which requires both feature learning and nonlinearity. We make explicit that, for these tasks, the ReLU networks possess an inductive bias towards latent representations which are *not* strictly modular or disentangled but are still highly structured and reusable between contexts. This effect is amplified with the addition of more contexts and hidden layers. Thus, we take a step towards a theory of feature learning in finite ReLU networks and shed light on how structured mixed-selective latent representations can emerge due to a bias for node-reuse and learning speed.

Published as a conference paper at ICLR 2025
MAKE HASTE SLOWLY: A THEORY OF EMERGENT STRUCTURED MIXED SELECTIVITY IN FEATURE LEARNING RELU NETWORKS
Devon Jarvis Richard 122 Benjamin Rosman & Andrew M. Saxe School of Computer Science and Applied Mathematics, University of the Witwatersrand Maachnne Intelligence and Neural Discovery Institute, University of the Witwatersrand "Gatsby Computational Neuroscience Unit & Sainsbury Wellcome Centre, UCL 'CIFAR Azrieli Global Scholar, CIFAR
{devon. jarvis, richard..... klein, jen @wits wwitssac.2.. za a. saxe@ucll ac.
ABSTRACT
In spite of finite dimension ReLU neural networks being a consistent factor behind recent deep learning successes, theory of feature learning in these models re- mains elusive. Currently, insightful theories still rely on assumptions including the linearity of the network computations, unstructured input data and architectural con- straints such as infinite width or a single hidden layer. To begin to address this gap we establish an equivalence between ReLU networks and Gated Deep Linear Net- works, and use their greater tractability to derive dynamics of learning. We then con- sider multiple variants of core task reminiscent of multi-task learning or contextual control which requires both feature learning and nonlinearity. We make explicit that, for these tasks, the ReLU networks possess an inductive bias towards latent representations which are not strictly modular or disentangled but are still highly structured and reusable between contexts. This effect is amplified with the addition of more contexts and hidden layers. Thus, we take step towards theory of feature learning in finite ReLU networks and shed light on how structured mixed-seles -exeetselective latent representations can emerge due bias for node-reuse and learning speed.
INTRODUCTION
When contrast against the rapid empirical progress of neural networks (Mnih et al. 2015 Amodei al. 2016 Andreas et al. 2016 He al. 2016)] (iiveer et 2017 Vaswani al. 7 Baevski et al.
2020 Reid et al. (2224) it can appear that theoretical understanding of these models is not keeping pace.
However, great theoretical progress has been made with a number of new paradigms being proposed
and investigated (Jacot et al. 2018 Goldt et al. 2020 Saxe et al. 2022 Li & Sompolinsky 2022).
While these paradigms illuminate many aspects of deep network behaviour, they typically simplify feature learning a key ingredient thought to underlie the success of deep networks (Krizhevsky et al.
2017 Jing & Tian (020). For instance, while statistical physics provides a theoretical paradigm with
clear insight for unstructured data from Gaussian distributions (Saad & Solla 1995 Riegler & Biehl
1995 Goldt et al. 2019;; Advani et al. 2020), treating more intricately structured domains remains challenging. Consequently. these theories can certainly explain phenomena observed in practice (Ramasesh et al., 2020 Lee et al., 2022 2024) but do not explicitly capture dynamics with richer data structure (Goldt et al. (020). Similarly, the Neural Tangent Kernel (NTK) (Jacot al. [0188] Lee et al.) 2019) offers exact learning dynamics for network architectures trained in the "lazy regime" of infinite hidden layer width or large initial weights (Chizat et al., 2019 Alemohammad et al., 2021). Yet, lazy networks tend to generalize worse than those in the "feature learning regime" (Bietti & Mairal 2019 Geiger et al. 2020) of small initial weights and finite width. Conversely, deep linear networks (Saxe et al. 201 2019) offer a theory of feature learning and have successfully characterised continual learning (Braun et al. 2022) and systematic generalisation (Jarvis et al. (023). Yet, these models
"Corresponding author

---

Published as a conference paper at ICLR 2025
do not incorporate nonlinear activation functions which are necessary for the network to perform
real world tasks ranging in complexity from the simple XoR (Minsky & Papert 1969) to Context
Generalization (Sodhani et al. [221]] Dahl et al. 2011: Beukman et al. (224), Other useful approaches consider a static analysis of stable points (Seung et al., (92 & Krzakala 2016) or require
infinite input dimension with infinite data (known the thermodynamic limit) (Mignacco et al. 2020
Goldt et al. 2020; Li & Sompolinsky, 2222). Hence an explicit theory of feature learning in ReLU
networks (Fukushimal 1969; Nair& & Hinton 2010 with finite width has not yet been presented, in
spite of their consistent use in the empirical successes of deep learning (Nwankpa et al., 2018).
In this work we take a step towards a theory of feature learn- ing in finite ReLU networks. We do this by drawing an equivalence between ReLU networks and Gated Deep Linear Networks (GDLN) (Saxe et al., (022) in Section 3 As shown in Figure I GDLNS provide a class of network architectures which perform nonlinear computations on their inputs using a composition of linear modules and are amenable to a full characterisation of their training dynamics. We show that for a ReLU network it is possible to create an equivalent GDLN which has the same output as the ReLU network at all points in time. We term this the "Rectified Linear Nel- work (ReLN)" and obtain the full training dynamics of both networks trained on an adapted XoR task. Section 4 then applies this technique to a significantly more realistic set- ting reminiscent of multi-task learning or contextual control (Zamir al. 2018). We provide full training dynamics for the ReLU network in this setting via the ReLN and prove that this case the equivalence is unique. In this task, we find that ReLU networks present an inductive bias towards structured mixed selective latent representations which are reused across contexts. Section 5 and Section 6 then consider the effect of adding contexts and multiple hidden layers re- spectively. We find that the bias towards structured mixed selectivity and node reuse is amplified by these changes, indicating that this may be a consistent emergent strategy for tasks of increasing difficulty. Thus, we summarize our main contributions as follows:
he We
Cpp c)
9y
s(p.c)
9
Figure 1: GDLN Formalism and notation. a) The GDLN applies gating variables to nodes (gv) and edges (g,) in an otherwise linear network. b) The gradient for an edge (using the red edge in (a) for example) can be written in terms of paths through that edge (colored lines). Each path is broken into the component preceding (s(p, e)) and following (f(p, e)) the edge.
We introduce Rectified Linear Networks (ReLNs) as the subset of GDLNs which each imitate a ReLU network. Through this we obtain the full training dynamics for finite, feature learning ReLU networks
We demonstrate that, in a contextual task, ReLU networks have an implicit bias towards structured mixed selectivity even during feature learning as this minimizes the loss the fastest.
Finally, we demonstrate that the bias towards node reuse is exasperated with the addition of more contexts and hidden layers on this task.
Thus, while performing "slow" feature learning mixed-selective functional modules emerge in ReLU networks which exploit node-reuse for learning speed (or "haste").
2 BACKGROUND
We briefly review the GDLN paradigm here as our work builds directly from it. However, the paradigm itself builds on the linear neural network theory (Saxe et al. 2014 2019). Thus, for readers being introduced to linear neural network theory from this work we provide a more thorough review in Appendix A The GDLN notation and formalism is also depicted in Figure I GDLNs are class of neural network which compute nonlinear functions of their inputs through the composition of linear network modules (Saxe al. 2022). The nonlinearity is achieved by gating of the network computation for different inputs. Formally, let I denote a directed graph with nodes V and edges E. Each node E V represents a layer of neurons with activity h2 E R ol where |u| E N denotes the number of neurons in the layer. Each edge e E E connects two nodes with weight matrix W. of size [t(e]] x [s(e] where s, E + V return the source and target nodes of the edge, respectively. It is also helpful to generalize this notion of edges to paths p, which are sequence
off portions
2

---

Published as a conference paper at ICLR 2025
of edges Wp. Thus, paths also connect their source node s(p) to their target node t(p). We denote the portion a path p preceding particular edge as s(p, e) and the portion of the path subsequent to the edge E(p, e) such that Wr ep.
We collect nodes with only outgoing edges into the set In([') C V and call them input nodes. Similarly, output nodes only have incoming edges and are collected in the set Out(1') C V. Activity is propagated through the network for given datapoint which specifies values 2p E Ri for all input nodes E In()). When concatenated the entire data point is denoted æ. The activity of each subsequent layer is then given by:
h = go
Σ VEE:())
9n W ((
Here Iv is node gate and 94 is an edge gate. Thus, the gating variables modulate the propagation of activity through the network by switching off entire nodes (ge) and edges between nodes (gg).
We train the GDLN to minimize the L2 loss averaged over the full dataset of N datapoints using gradient descent. The i-th datapoint then a triple specified by the input æ", output labels y and specified gating structure g", where vi E Rlo for v E Out(T'.
L((W)) 2N
Σ
Σ eEOui(T)
l1 1
(2)
Note that in Saxe et al. (2022) the gates are treated as part of the dataset. As the main results of this work do not depend on how we find the gates for our GDLN, we follow this precedent. However, the case of finding the gates is an important consideration for the usability of the paradigm we introduce here. Thus, in Appendix B we discuss a very simple clustering approach finding the gates from a ReLU network.
If we denote Pee) as the set of paths which pass through edge e, T(u) as the set of paths terminating at node and Ne then the update step for single weight matrix in the GDLN is:
d W. OL((W)) dW.
Σ
W. t(p.e)
(VE
Σ U.D jET((())
W s(nee YEEE (3)
dt
pEP(e)
Thus, the update to a weight matrix Wa is determined by the error at the end of a path which contributes to (Ev(()) EjET(p)) W (j.p)) summed over all paths it is a part of E P(e).
Notably, all dataset statistics which direct learning are collected into the correlation matrices:
N
E (p)
Σ N
ST p(p)
ET(i.p) N Σ "f 9
(4)
These dataset statistics depend on the gating variables specific to datapoint g', indicating that each path sees its own effective dataset determined by the network architecture. From this perspective. each path resembles the gradient flow of deep linear network (Saxe et al. 2014: 2019) which have been shown to exhibit nonlinear learning dynamics observed in general deep neural networks (Baldi & Hornik 1989 Fukumizul 1998| Arora et al.) 2018 Lampinen & Ganguli 2019). The final step then is to use the change of variables employed by the linear network dynamics (Saxe et al., 2014 2019). Assuming the effective dataset correlation matrices are mutually diagonalizable (such that V is the same matrix for y and Σ²), we write the correlation matrices and pathway weights in terms of their singular value decompositions:
E"G.p)) - "Vi) D(. Wo(1) Re(e) Be
(5)
where Spp) and (p)) are diagonal singular value matrices, (e()) are the new dynamic variables, R satisfies R. R, = 7 for all v E V, and for input and output nodes Ra(p) = Va(p) and Re(p) = Un()) respectively. For Raip) = Va(p) and Repp = Ut(p) we assume that the pathway mappings align to the same singular vectors which diagonalise their effective datasets and apply the same change of variables to the pathway weights. This assumption has been used successfully for feature learning linear networks previously (Saxe al.) 2014 [019 Jarvis al.) (023) and the trend for networks to align in this manner has be termed the "silent alignment effect" (Atanasov et al., (021). This change of variables removes competitive interactions between singular value modes along path such that
3

---

Published as a conference paper at ICLR 2025
the dynamics of the overall network is described by summing several "ID networks" (one for each singular value) resulting in the "neural race reduction" " (Saxe et al. 2022
B dt
Σ PEP(e)
Boke
S(p) - Σ BjD(j.P) VEEE
(6)
JET((())
From this reduction we note that the update to layer of weights is proportional to the input-output singular values of the effective dataset along a pathway it contributes to, and the number of these pathways (shown by the summation over p P(e)). Thus, we can summarize the assumptions of the GDLN paradigm from Saxe al. (2022) and note that we do not introduce any new assumptions:
Assumption 2.1. The assumptions of the GDLN paradigm are:
1. The dataset correlation matrices are mutually diagonalizable.
2. The neural network weights align to the singular vectors of the dataset correlation matrices.
Importantly, Assumption 2.1.1 1 is necessary for the full tractability and interpretability offered by the GDLN framework. However, if it is violated we can still continue the derivation to a dynamics reduction which is the same as the neural race reduction (Saxe et al. 2022) but includes additional matrix multiplication, which is insightful and likely to be enough for many cases. Assumption 2.1.2 stricter, however in the present context this is the same assuming the network successfully feature learning. We elaborate more on both assumptions in Section Aaaad note that Section depicts the utility of our framework in a case where the alignment of the singular vectors not guaranteed.
THE RECTIFIED LINEAR NETWORK (RELN)
We denote the output of GDLN for the i-th data point x² in dataset of N data points as (W)) G(z', =") where Wit) is the set of weights and G(x', ((x',t) is the set of gates that are dependent on the data point. Both the gates and weights can be time dependent, however in this section and Sections 4 and 5 the gating structure is constant from the beginning. Thus, G(z',t)= = G(z 0) Git
V E in these sections. We denote the ReLU network's output as weights. We formally define ReLN as:
ReLU ReLU(W(() =!), with (tt) as the
Definition 3.1. Rectified Linear Network (ReLN): The GDLN with G(x") = G*(x") such that B+ and V E {0, NJ..
GWLN(W G (x*),
(t),
This definition is not specific to the tasks considered in this work and it is always possible obtain ReLN for given ReLU network and task. Letting h, ΣK W ox be the pre-activation of the j-th hidden neuron in our ReLU network for any given data point x' (K is the number of input features), then the post-activation from the ReLU activation can be written as: (h;) step(h, (h,) Miy. Thus, in GDLN with exactly the same architecture we could gate the corresponding neuron with G(x²), styp(Tr)). In other words we could gate the GDLN hidden neurons based on the activity of the ReLU network explicitly. This extends multiple hidden layers well. However, this is not desirable and demonstrates the extremity of the ReLN framework. Instead, we aim to find as few unique gating patterns G(z²) across the dataset as possible exposing the functional "modules" of the network. We now illustrate this approach with an example.
Transition to nonlinear separability. ReLU networks have the ability to perform nonlinear tasks. Yet when faced with linearly separable problem, they can also possibly learn linear decision boundary. Here we examine a task which permits both of these possibilities, to examine the implicit bias in gradient descent. As depicted in Figure 2b. the task has XoR structure in its first two dimensions; and linearly separable with margin 2A in the third. Hence for A = 0, the task is not linearly separable, but for A 0 it is. To identify ReLN on this task, consider two intuitive gating structures. The first attempts to exploit the linearly separable structure and contains two pathways, one active for positive examples and
active for negative examples, depicted in Figure 2p. The second GDLN (Figure 2 contains four pathways, each active on exactly one example. This GDLN can solve the task even when A = 0, because it uses different linear model for each datapoint.
Given the problem, these gating structures are intuitive. It remains to show that they behave like a ReLU network trained on the same task. Figure 2 shows the loss dynamics of each GDLN and for
4

---

Published as a conference paper at ICLR 2025
a)
b) Linear
XoR
Leaminng Time (Epochs)
Liner
8g
10
RaLU
101
10
30
Separation A
04
04
0.4
04
04
Linarr
Xot
02
02:
02
02
00
0.0
00
00
150
100
Epachs
Epochs
Epochs
Epochs
Epochs
A=
A=2
Figure 2: Dynamics of ReLU networks during transition to nonlinear separability, (a) A simple dataset which has XoR structure in the first two dimensions, and is linearly separable with margin 2A in the third. (b) GDLN exploiting linear structure. The network contains two pathways, one gated on only for positive examples and one gated on only for negative examples. Blue arrows in panel (a) depict ReL.U weight directions that achieve this gating.
(c) GDLN exploiting XoR structure. The network contains four pathways, each active on exactly one example. Orange arrows in panel (a) depict ReLU weight directions that achieve this gating. (d) Time to learn to a fixed criterion (loss=.2) calculated analytically for GDLNs with linear and XoR gating structure (blue and orange, respectively), and in simulation of ReLU networks (green). The ReLU network behaves like the faster of the two gated networks. Which gating structure is fastest changes at A v 2/3 (grey dashed). (e) Analytical loss trajectories for the gated networks, and simulated ReLU networks for several values of A. The full trajectories of the simulation match the faster of the gated networks. Parameters: learning rate 1/T = A, N₂ = 128 hidden units, initialization variance 4 10 8/nn.
a 128 hidden unit ReLU network trained on the same task, for a range of values of A. We see that across the spectrum of As, one of these GDLNs always closely matches the ReLU network dynamies which is the ReLN in that setting. Figure 21 shows the learning time of these networks as function of A, showing that ReLU networks behave like the faster GDLN (the ReLN).
Because of the simple structure in these ReLNs, it is possible to solve their dynamics analytically (see Appendix [c. These dynamics yield a prediction for the behaviour of the ReLUU network. Using these solutions, we can calculate the value of A when the speed of the linear GDLN crosses over that of the XoR ReLN, This crossover point is A 22/3 (Figure 2 grey dashed line), which marks the approximate point at which a ReLU network switches from a linear decision boundary to nonlinear boundary (see Appendix C). Perhaps surprisingly, this analysis shows that the ReLU network transitions to a nonlinear boundary well before the problem becomes nonlinearly separable.
Hence this example demonstrates how ReLNs with appropriate gating can describe the behaviour of ReLU networks. When the ReLN has a simple structure, its greater tractability permits analytical so- lutions that closely describe ReLU dynamics gradient-trained ReLU networks do have a bias toward linear decision boundaries when the margin large, they can adopt nonlinear decision boundaries even when problem is linearly separable.
the same task. Finally, this example shows that while
EMERGENT STRUCTURED MIXED SELECTIVITY IN RELU NETWORKS
Having introduced ReLNs in Section 3 we now apply this theory to understand the inductive biases of ReLU networks in a more challenging setting, summarized in Figure 3 Similar to prior work we create a task where the network is provided a object index as input and required to produce a set of corresponding properties for the object (Saxe et al., 2019 Braun et al. 2022). however we include a context feature to the input similar to previous connectionist models of controlled semantic cognition (Rogers et al. 2004) and more recently in multi-task learning (Sodhani et al. 2021). Each item is queried with a one-hot representation and one-hot context feature, such that all items are present in each context, forming the input matrix X (X, denotes the contextual portion of the dataset). The outputs, corresponding to predicted features the items, then impart structure into the dataset based on the similarity of items. For example, the first set of output labels (rows) form hierarchy structure.
5

---

Published as a conference paper at ICLR 2025
Rel.U Network
Dellwny
GDLN
Sarresy ARE
Hergchool Camorr
Fis
Ring
622
Chipi
Cross
X
Rel.1 Artivation
Comsed
Context independent hierarchy with context dependent variance for each item in each context.
Figure 3: Dataset used to train the ReLU network (left) and ReLN architecture used to imitate it (right). Inputs (middle matrix) are created by appending a one hot vector encoding object identity to a one hot vector encoding context such that each item appears in all contexts. Target outputs (left and right matrices) contain some contextindeependent (top block) and some context specific properties (bottom three blocks). These datasets broadly follow hierarchical structure across items (hierarchical tree depicted in the middle over input datapoint along the columns), but with some variation in each context-specific block. All structures taken from Saxe al. (2019). The analysis in this work shows that the ReLU network dynamics arise from four implicit modules, made explicit by the ReLN pathways towards the right, which receive different subsets of inputs and generate different subsets of outputs. Together these graded mixed-selectivg pathways couple together to produce the correct output labels for each object. While each context-specific pathway only on in two contexts (blocks of columns) they still produce labels for all three contexi-specific parts of the output space (blocks of rows). This creates errors which other pathways learn remove. If this fine balance of activity between pathways is broken then errors will be incurred.
This block of features should be activated regardless of the queried context. In contrast the three other blocks of output labels need to be activated only in one of the contexts. This requires contextual control and a nonlinear network mapping to learn. We train the networks with full-batch gradient descent and L2 loss (Equation 2 to perform this task. We use a linear output layer, assume the model is peeemarraetriieed (a pathway needs to have at least l hidden nodes to learn a rank h effective dataset) and do not regularize or bias the networks towards context specificity in its hidden neurons. The only difference between the two networks is that the gating of the ReLN's hidden neurons is explicit, while the ReLU network learns the gating pattern using the nonlinearity. Remarkably, in this setting the gating of the ReLN is constant from the start.
To obtain the closed-form dynamics of the ReLN, and by extension ReLU network, it is necessary determine the effective datasets of each pathway through the network. We identify four pathways which imitate functional modules in the ReLU network. The first "common" pathway is composed of hidden neurons which are always active, forming a linear subnetwork responsible for learning the aver- age activation across the dataset. While this pathway does not use context for gating, it still uses this as feature similar bias variable. The three remaining pathways are all active for two out of the three contexts and collectively learn the residual of the common pathway. All pathways map onto the full output space. The ReLN and corresponding pathways are depicted on the right of Figure 3 Figure a) shows the loss dynamics of the ReLU network compared to the ReLN and the predicted dynamics with samples of the ReLU and ReLN outputs at certain points during training") We see excellent agreement across all dynamics and outputs. Further, comparing Figure ( ()) which show the multi-dimensional Scaling (MDS) plots of the ReLU network and ReLN respectively, we see near exact correspondence between the relative positions of the latent representations of the dataset across both architectures. The derivation of the dynamics using the race reduction from Section is presented in full in Appendix F
By interpreting the corresponding ReLN, find that there no neurons in the ReLU network which are context specific. Instead, the ReLU network employs a strategy of mixed selectivity where neurons are active across multiple contexts. These pathways are still doing feature learning, however, as evidenced by the effective datasets which we uncover and the corresponding singular
"ode using the Jax Library (Bradbury
2018 is htttttttt//ttnrrrrrree www.:/ll.oooooo//o////

---

Published as a conference paper at ICLR 2025
a)
Comparison of Loss Dynamics
b)
Common Partway Dynamics
Connext Pathways Weon Dynamics
ReLU
100
ReLN (GDLN)
GDLN Single
Quadratic Loss
80
Race Dynamics
1900
Tpech narber
40
C)
RoLU Network Hiddon Layer MDS
Ret D4 Hitten Layer MDS
20
2000
4000
6000
8000
Epoch number
Figure 4: Summary of results for the RcLN imitating a ReLU network on a contextual nonlinear task. (a) Comparison of loss trajectories between the ReLU, empirical and predicted ReLN, and alternate GDLN which does not imitate the ReLU network loss trajectory. We find that a GDLN (here called "GDLN Single") which has contextual pathways active for individual contexts unable to imitate the loss trajectory of the ReLU network. This is used in Proposition 4.2 to prove the uniqueness of the ReLN which we have identified. Example outputs from the ReLU network and ReLN are also shown and we see exact agreement between these output samples. (b) Singular value dynamics for the ReLN architecture using the neural race reduction dynamics. (c) kulti-dimensional Scaling which compares the relative latent representations of the ReLUU network and ReI.N over time. Both architectures demonstrate an equivalent latent representation at all points in time.
value trajectories for each summarized in Figure 4b). Finally, Figure 4 Honly shows that the loss of the ReLN matches that of the ReLU network, which is weaker condition than producing the same output at all times, as in Definition 3.1 A natural question may also emerge on whether another ReLN exists which provides the same dynamics. To answer this, we present Proposition 4.2 which proves that the loss trajectory of this ReLN is unique (accounting for symmetric gating patterns such as permuting hidden neurons (Simsek al. [021] Martinelli et al.) (023)). In addition, the
ReLN which emerges is the fastest learner of all GDLN architectures implementable by ReLU network, winning the "neural race" (Saxe et al., 2022). Thus, while the ReLU network is doing "slow" feature learning, it does so by utilizing the most efficient network substructures. The structure emerges from feature learning, while mixed-sclectivity and node reuse remains as strategy for learning speed, commonly found in the "fast" lazy learning regime (Jacot et al., 2018: Geiger et al., 2020). Importantly, mixed-sslectiiityy is an established concept within the neurosciènce literature (Anderson 2010, Rigotti et al. 2013), however this balancing of learning speed and feature learning reflects a potential mechanism for its emergence. Before we present Proposition 4.2 we first require Lemma 4.1 which limits the possible gating strategies implementable by a ReLU network with single hidden layer on this dataset and drastically compresses the possible set of GDLN architectures which could be potential ReLNs. The full proof can be found in Appendix D
Informal Lemma 4.1 Any gating strategy employed by trained ReLU network with a single hidden layer on the input X (as defined in Section 4 can be implemented by GDLN with linear modules that gate based on the context features in isolation or are only active for a single datapoint.
Lemma 4.1. Given the definition of GDLNs and ReLU Networks, for any gating pattern imple- mentable by the ReLU network there is G(rz)) such that G(æe); step(h;) or G(z); = 1 1xx) V
{0, I,.,, H} where h is the pre-activations of the ReLU network's 's hidden layer with H neurons and Te is the portion of the data point x with contextual features.
We note that a ReLU network must use the same set of weights to perform the gating and forward propagation of information. Without loss of generality then consider the simple case of two objects in two contexts forming a dataset of four datapoints. We show by contradiction that a hidden neuron is unable to gate inconsistently for both object and context, for example by being active for item 1 in context 1 and item in context 2 while being off on the remaining datapoints. This means all hidden neurons must gate consistently using either context or object features. We then consider a strategy where a neuron gates based on an item 's features. In this case the readout from the neuron will produce outputs across the entire label
Proof Sketch:
7

---

Published as a conference paper at ICLR 2025
space agnostic to the requested context as it will be activated by all context features. Thus, it is unable to perform the nonlinear mapping from context to output. Consequently the only viable gating strategy for the network is partition the dataset using the context features in isolation or allocate subnetworks to each datapoint.
With the search space of the ReLN within the space of GDLNs reduced we may prove that the loss trajectory of the ReLN in this section is unique. The full proof of Proposition 4.2 can be found in Appendix Note that LGDLN and RELU are the loss of the GDLN and ReL.U networks respectively.
Informal Proposition 4.2 The ReLU network in this section: 1) has a unique corresponding ReLN which imitates the loss dynamics and output at all times, 2) finds the neural race winning strategy of mixed selectivity.
Proposition 4.2. There is unique (X) (up to symmetries such as permutations) such that LGDLN (t), GO (X)) = LROLU (W t E IR+ and for alternative gating strategies G' (X) (W()),
then LODEN (GDLN(W(() G (X)) < (IDDLN(W()) ((),G"(X))
We first make use of generalization of the Cauchy Interlacing Theorem to non-square matrices, closely following the strategy from Thompson (1972). We show that if an output label or input feature is removed from the dataset, the first singular value of the input- output correlation matrix (E*") can only decrease. Since the singular values of the correlation matrix alone determine the time-course of learning in linear networks and pathways (Saxe et al., 2014 2022) this means network can only learn slower when removing these input or output dimensions. Thus, pathways using all available features will win the neural race, consistent with our ReLN architecture, and any other strategy would result in clear deviation from the ReLU loss trajectory. We then consider whether gating along datapoints and not features could provide alternative architecture. Here note that the output of the common pathway is element-wise positive. Thus, removing a datapoint with only positive terms can only decrease the correlation caleulation. Thus, the common pathway which uses all features and is active for all datapoints is the fastest possible module the network could implement. Lastly, we consider if alternative gatings along the context specific pathways exist. Using Lemma 4.1 we note that the network can only implement gates which partition the input consistently based on the contexts or activate separate modules for each datapoint. For the latter case, once the common pathway has finished learning there is no remaining correlation between the context features and residual. Thus when a single datapoint allocated to a single module activity would need to be learned by the item feature. In our dataset, removing all other items in a context would then be the same as removing their individual one-hot input features. From the Cauchy Interlacing Theorem, we know this can only slow learning. Thus, there are two remaining strategies which gate consistently based on context, making a proof by exhaustion feasible, and so we simply simulate both and plot the trajectories in Figure 4 Clearly the architecture with pathways active in two contexts learns faster and matches the loss trajectory of the ReLU network. Thus, only one ReLN in the space of implementable GDLNs matches the ReLU network loss trajectory, and this is the fastest possible ReLN which wins the neural race.
Proof Sketch:
5 THE EFFECT OF ADDITIONAL CONTEXTS
Having identified the inductive bias towards mixed selectivity in ReLU networks we now consider whether this trend continues as we increase the number of contexts. Once again we construct a ReLN which imitates the ReLU networks and identify the effective datasets for each of its pathways. In this case we consider datasets with three, four and five contexts. Similar to Section [3] the output labels contain a block with hierarchical structure which is active for all contexts. This is appended with context specific blocks which are active for a single context. Here we use a hierarchy structure for all contexts to isolate the effect of scaling the number of contexts. Additionally, with this repeated structure it is possible to continue the derivation of the dynamics beyond the neural race reduction and obtain fully closed form equations for the singular value trajectories of the linear pathways. The derivations for the three, four and five context cases are presented in Appendix F and G We plot the closed form singular value trajectories and corresponding loss curves for each context Figure 5 and once again see excellent agreement. Moreover we find that the trend identified in Section Sattionn]conn continues as the ReI.N which imitates the ReL.U network in each case uses linear pathway that is active for all
8

---

Published as a conference paper at ICLR 2025
contexts. The ReLN then has C' context specific pathways which are active for C - 1 contexts. Thus, as the number of contexts increases, so too does the bias towards node reuse and mixed selectivity.
Context
Context
Context
100
Rebi a Cortoot
170
Riil Corted
Contect
1.00
sebd Corsoc
Rest Cartest
Guard Cynamics Carmod
100
120
Loss Vajectories
Quadrake
60-
Quadirate
100-
20
20
2000
4000
5000
8000
3010
4000
1002
1003
2000
5000
6000
8000
Epoch number
Epsch number
Epoch number
20-
2.5
2.0
Common Futhway 5V Dyramica
1.2
Singula
0.5
0.5
5228
.000
2005
4000
Fpoch number
Fpor number
Fpoch nummee
2.0
2.0-
1.5
Contect Pethways Dynsmice
1.3
1.0
1000 2000 3000 4523 5020 2006 T000 3000 1000 2023 2000 4000 5000 5000 7300 1003 3000 2000 3000 4000 5290 E0CS TOCO 0000
Figure 5: Effect of increasing the number of contexts on mixed selectivity: Summary of the loss trajectory (first row) for the ReLU network and corresponding ReLN as the number of contexts increases (columns). Due the symmetry of the task we are also able to continue the derivation of the ReLN dynamics beyond the neural race reduction and obtain closed form trajectories for the networks' singular values. The singular value trajectories for the common pathway (second row) and mean trajectories for the contextual pathways (bottom row) match simulations exactly. Consequently we can closed form derive the loss trajectory for each architecture and see perfect agreement to both the ReLN and ReLU networks. All together these results demonstrate that the number of contexts increase, so too do the number of contexts a pathway is active for.
THE EFFECT OF ADDITIONAL HIDDEN LAYERS
Finally, we consider the effect of adding second hidden layer with ReLU activation to the network. Once again we are able to identify the corresponding GDLN which imitates the network. This GDLN architecture is summarized in Figure 6 We that in this case the ReLU network does not use the first layer for gating and instead all neurons are active for all datapoints. However, this layer does still benefit the network as it enables it gate using a combination of item and context variables in the second hidden layer. This removes the constraint from Lemma I and allows a wider range of gating strategies. However, we see that the gating strategy of the network remains relatively similar as the contextual pathways are active in two out of the three contexts. Similarly the common pathway remains. Thus, the trend towards mixed-selectivity and node reuse remains with the addition of another hidden layer, which is then used as a linear layer.
In Figure oo we plot the trajectories from 100 runs of the ReLU network and GDLN respectively. We note that there is now variance in the trajectories of both networks, which was not the case for the single-hidded hiddee--ae--ee architectures. This is due to an inconsistency in the timing at which the network singular vectors align with the singular vectors of the dataset correlation matrices. As consequence there is variance in the loss trajectories as the contextual pathways do not align perfectly to consistent set of singular vectors from early in training (violating Assumption| 2 We see that the GDLN is still informative about the behaviour of the ReLU network. In Figure 6 we find the loss trajectories closest to the mean trajectory for the ReLU network and GDLN separately. We then plot these stereotypical loss trajectories for comparison and plot the faded remaining trajectories around this. We see close agreement between the ReL.U network and GDLN stereotypical trajectories, but
9

---

Published as a conference paper at ICLR 2025
100
ReLU
RoLN (GDLN)
Quadratic Loss
40
Creport indeppseeeer beerethy for lors COMEEE
5000 10000 15000 20000 25000 30000 Epoch number
Figure 6: Summary of architecture and results for the addition of depth: With the addition of depth the ReLU network is no longer limited to gating consistently using context variables. However, we find that the network still uses a common pathway and contextual pathways active in two contexts. The first layer now is completely shared across all contexts and does not the nonlinearity. This architecture is imitated by the ReLN depicted on the left. We note that there is now variance between runs of the ReLU network and the ReLN. Interestingly the ReLN shares the same characteristics in variance as the ReLU network. We plot 100 runs of each architecture and highlight one trajectory for the ReLU network and ReLN which could be considered stereotypical (using the sum of L2 distances across all timepoints). Overall the ReLN still provides clear insight into the ReLU network and its inductive bias towards mixed selectivity and node reuse.
also the same variance profiles across the various runs. We note that pathways within the GDLN become active in stages. For example, the common pathway align around epoch 500, while the contextual pathways only appear epoch 12000. The relative timing of the drops in loss and positions of the plateaus in loss are all consistent following the alignment. Moreover, when comparing samples from both networks at plateaus in the loss trajectories see exact agreement between the outputs. In summary, while the addition of a hidden layer allows the ReLU to have a more intricate gating structure, this comes with an additional difficulty of perfectly aligning the dataset singular vectors. Thus, the ReLU networks and GDLNs fail to learn the exact optimal features on all runs.
7 DISCUSSION
In this work we have introduced the Rectified Linear Network which is a GDLN designed imitate the gating structure of a specific ReLU network. We have used this to provide a theory of feature learning in finite ReL.U networks and identified an implicit bias towards structured mixed selectivity which results due to node reuse. Our approach has several limitations. Most notably, while some GDLNs can imitate any ReLU, a reduction based on this idea will only be insightful if the resulting GDLN has few pathways. Otherwise, the many pathways of the GDLN become hard to interpret as the many neurons in a ReLU network. A second limiting factor is how easily identifiable the gating structure is. As first step towards addressing this we propose very simple clustering based algorithm in Appendix B which we show is able to retrieve the correct modules from a ReLU network in the setup of Section 4 In this appendix we also briefly touch on the connection meta-learning and polysemanticity (neurons being activated by multiple semantic unit or features) (Olah et al.) 2017a Lecomte et al. 2024), and promote these as directions of future work.
In the reasonably complex yet structured cases we examine here, we have shown that subnetworks emerge within ReLU network due the learning speed benefits offered by mixed-selectivity. Moreover, we find that the bias towards node reuse increases with the addition of more contexts and greater depth. Thus, our findings are in agreement with prior empirical and theoretical findings on disentanglement which demonstrate that neural networks are not implicitly biased towards dis- entangled representations that select for semantic factors of variation in the input (Locatello et al.. 2019, Michlo Michloo et al., 2023). However, we find that mixed selectivity of neural representations does
not imply that the network not learning structured features. Instead the network exploits reusable structure between contexts, forming functional modules which we make explicit with pathways in the ReLNs. This demonstrates that node reuse and modularity are not as incongruent as previously thought in both machine learning (Andreas al. [2016] Andreas 2018 Masoudnia & Ebrahimpour 2014) and neuroscience (Anderson] 2010 Rigotti et al. 2013)) and we shed light on the competing factors influencing their emergence from learning.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Joseph Warren and Mingda Xu for useful discussions and feedback on this work. For the purpose of Open Access, the author has applied a CC BY public copyright license to any Author Accepted Manuseript version arising from this submission. A.S. is supported by the Gatsby Charitable Foundation and Sainsbury Wellcome Centre Core Grant (219627/2/1//2), DJJ. is Google PhD Fellow mentored by Gamaleldin F. Elsayed and a Commonwealth Scholar. A.S. and B.R. are CIFAR Azrieli Global Scholars in the Learning in Machines & Brains program. Computations were performed using a Nvidia GTX 1080 from the High Performance Computing infrastructure provided by the Mathematical Sciences Support unit at the University of the Witwatersrand.
REFERENCES
Madhu S Advani, Andrew M Saxc, and Haim Sompolinsky. High- High-dimensional dynamics of general- ization error neural networks. Neural Networks, 132:428-4444 2020.
Sina Alemohammad, Zichao Wang, Randall Balestriero, and Richard Baraniuk. The recurrent neural tangent kernel. In The International Conference on Learning Representations, 2021.
Dario Amodei, Sundaram Ananthanarayanan, Rishita Anubhai, Jingliang Bai, Eric Battenberg, Carl Case, Jared Casper, Bryan Catanzaro, Qiang Cheng, Guoliang Chen, et al. Deep speech 2: End-to- end speech recognition in english and mandarin. International conference on machine learning. pp. 173 182. PMLR, 2016.
Michael L. Anderson. Neural reuse: A fundamental organizational principle of the brain. Behavioral and brain sciences, 33(4):245 266, 2010.
Jacob Andreas. Measuring compositionality in representation learning. In International Conference on Learning Representations, 2018.
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 39-48, 2016.
S. Arora, N. Cohen, and E. Hazan. On the optimization of deep networks: Implicit acceleration by overparameterivaaton. 35th International Conference on Machine Learning, ICMI. 2018, 372-389, 2018. arXiv: 1802.06509 ISBN: 9781510867963.
Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan, Neural networks as kemel learners: The silent alignment effect. In International Conference on Learning Representations, 2021.
Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vee 2.0: A framework for self-supervised learning of speech representations. Advances in neural information processing systems, 1112242492666 2020.
P. Baldi and K. Horiik. Neural networks and principal component analysis: Learning from examples without local minima. Neural Networks, 2(1):53 58, January 1989. ISSN 08936080, doi: 10.1016/ 0893 (893-6000003)900044- 0893608089900112
URL http:::://innigghubb wtwpssubbbeseeeeee elsevier wttp///eeeteeeeeee com/retrieve/pii.
Michael Beukman, Devon Jarvis, Richard Klein, Steven James, and Benjamin Rosman. Dynamics generalisation in reinforcement learning via adaptive context-aware policies. Advances in Neural Information Processing Systems, 36, 2024.
Alberto Bietti and Julien Mairal. On the inductive bias of neural tangent kernels. Advances in Neural Information Processing Systems, 32, 2019.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milun and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http:// /tttt/u com/ coon/google/jee
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.
11

---

Published as a conference paper at ICLR 2025
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
George E Dahl, Dong Yu, Li Deng, and Alex Acero. Context-dependent pre-traine pre-trained deep neural networks for large-vocabulary speech recognition. IEEE Transactions on audio, speech, and language processing, 20(1):30 2011.
K. Fukumizu. Effect of Batch Learning In Multilayer Neural Networks. In Proceedings of the 5th International Conference Neural Information Processing, pp. 67-70, 1998.
Kunihiko Fukushima. Visual feature extraction by a multilayered network of analog threshold elements. IEEE Transactions on Systems Science and Cybernetics, 5(4):322-333, 1969.
Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy training deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2020 (11):1 2020.
Sebastian Goldt, Madhu Advani, Andrew M Saxe. Florent Krzakala, and Lenka Zdeborová. Dynamies of stochastic gradient descent for two-layer neural networks in the teacher-student 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 on learning in neural networks: The hidden manifold model. Physical Review X, 10 (4):041044, 2020.
Kaiming Hc, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
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. 804 0144811,777 2017.
Minyoung Huh, Hossein Mobahi, Richard Zhang, Brian Cheung, Pulkit Agrawal, and Phillip Isola. The low-rank simplicity bias in deep networks. arXiv preprint arXiv:2103.10422 2021.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Devon Jarvis, Richard Klein, Benjamin Rosman, and Andrew M Saxe. On the specialization of neural modules. In The Eleventh International Conference Learning Representations, 2023.
Longlong Jing and Yingli Tian. Self-supervised visual feature learning with deep neural networks: A survey. IEEE transactions on pattern analysis and machine intelligence, 43(1 1):4037 4058, 2020.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. Communications of the ACM, 60(6):84-0))0)) 90, 2017.
Brenden M Lake. Compositional generalization through meta sequence-tor -oweencee---seuencee learning. Ad- vances neural information processing systems, 32, 2019.
Brenden M Lake and Marco Baroni. Human-like systematic generalization through a meta-learning neural network. Nature, 15--11, 2023.
A.K. Lampinen and S. Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. In Sainath (ed.), International Conference on Learning Representations, 2019. ISBN 0311-5518. doi: 1809. 10374
URL http: http///xxvvvooo/ www.aaiww.....aaa....../aaa
arXiv: 1839.10374.
Victor Lecomte, Kushal Thaman. Rylan Schaeffer, Naomi Bashkansky. Trevor Chow, and Sanmi Koyejo. What causes polysemanticity? an alternative origin story of mixed selectivity from incidental causes. In ICLR 2024 Workshop Representational Aligument, 2024.
12

---

Published as a conference paper at ICLR 2025
Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
Jin Hwa Lee, Stefano Sarao Mannelli, and Andrew M Saxe. Why do animals need shaping? a theory of task composition and curriculum learning. In International Conference Machine Learning, pp. 26837-26855. PMLR, 2024.
Sebastian Lee, Stefano Sarao Mannelli, Claudia Clopath, Sebastian Goldt, and Andrew Saxe. Maslow's hammer for catastrophic forgetting: Node re-use vs node activation. arXiv preprint (rXiv:2205.09029 2022.
Qianyi Li and Haim Sompolinsky. Globally gated deep linear networks. Advances in Neural Information Processing Systems, 35:34789-34801, 2022.
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentan- gled representations. In international conference on machine learning, pp. 4114-4124. PMLR, 2019.
Flavio Martinelli, Berfin Simsek, Wulfram Gerstner, and Johanni Brea. recovery of neural networks. arXiv preprint arXiv:2304. 12794,, 2023.
Parame-
Saced Masoudnia and Reza Ebrahimpour. Mixture of experts: a literature survey, Artificial Intelli- gence Review, 42:275-293, 2014.
Nathan Michlo, Richard Klein, and Steven James. Overlooked implications of the reconstruction loss for vae disentanglement, In Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence, pp. 4073 4081, 2023.
Francesca Mignacco, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Dynamical mean-field theory for stochastic gradient descent in gaussian mixture classification. Advances in Neural Information Processing Systems, 33:9540-9999 95000,9550000 2020.
Marvin Minsky and Seymour Papert. An introduction to computational geometry. Cambridge tiass., HIT, 1794400):104, 1969.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 518(7540):529-333 2015.
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th international conference machine learning (ICML-10), pp. 807-814, 2010.
Chigozie Nwankpa, Winifred Ijomah, Anthony Gachagan, and Stephen Marshall. Activation functions: Comparison of trends practice and research for deep learning. arXiv preprint arXiv:1811.03378, 2018.
Chris Olah, Alexander Mordvintsev, and Ludwig Schubert. Feature visualization. Distill, 2(11):e7, 2017a.
Vinay Ramasesh, Ethan Dyer, and Maithra Raghu. Anatomy of catastrophic forgetting: Hidden representations and task semantics. arXiv preprint arXiv:2007.07400. 2020.
Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap. baptiste Alayrac. Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.055300 2024.
Peter Riegler and Michael Bichl. On-line backpropagation in two-layered neural networks. Journal of Physics A: Mathematical and General, 28(20)1 1995.
13

---

Published as a conference paper at ICLR 2025
Mattia Rigotti, Omri Barak, Melissa R Warden, Xiao-Jing Wang. Nathaniel D Daw, Earl K Miller, and Stefano Fusi. The importance of mixed selectivity in complex cognitive tasks. Nature, 497 (7451):585-590, 2013.
Timothy T Rogers, James L McClelland, al. Semantic cognition: A parallel distributed processing approach. MIT press, 2004.
Laura Ruis and Brenden Lake. Improving systematic generalization through modularity and augmen- tation. arXiv preprint arXiv:2202.10745, 2022.
David Saad and Sara A Solla. Exact solution for on-line learning in multilayer neural networks. Physical Review Letters, 74(21):437,, 1995.
1995.
A.M. Saxe, J.I. McClelland, and S. Ganguli. Exact solutions the nonlinear dynamics of learning in deep linear neural networks. In Y. Bengio and Y. LeCun (eds.), International Conference on Learning Representations, Banff, Canada, 2014. Oral presentation. arXiv: 1312.6120v3.
Andrew Saxe, Shagun Sodhani, and Sam Jay Lewallen. The neural race reduction: Dynamics of abstraction in gated networks. In International Conference on Machine Learning, pp. 19287-19309. PMLR, 2022.
Andrew M Saxe, James L McClelland, and Surya Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 116(23): 11537-11546, 2019.
Hyunjune Sebastian Seung, Haim Sompolinsky, and Naftali Tishby. Statistical mechanics of learning from examples. Physical review A, 45(8):6056, 1992.
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676):54-359) 2017.
Berfin Simsek, François Ged, Arthur Jacot, Francesco Spadaro, Clément Hongler, Wulfram Gerstner, and Johanni Brea. Geometry of the loss landscape in overparameterized neural networks: Symme- tries and invariances. In International Conference on Machine Learning, pp. 9722-9732. PMLR, 2021.
Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context- based representations. In International Conference on Machine Learning, pp. 9767-9779, PMLR, 2021.
Robert C Thompson. Principal submatrices ix: Interlacing inequalities for singular values of submatrices. Linear Algebra and its Applications, 5(1):1-12, 1972.
Ankit Vani, Max Schwarzer, Yuchen Lu, Eeshan Dhekane, and Aaron Courville. Iterated learning for emergent systematicity in vqa. arXiv preprint arXiv:2105.01119, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Sifan Wang, Xinling Yu, and Paris Perdikaris. When and why pinns fail train: A neural tangent kernel perspective. Journal of Computational Physics, 449: 10768, 2022.
Amir R Zamir, Alexander Sax, William Shen, Leonidas J Guibas, Jitendra Malik, and Silvio Savarese. Taskonomy: Disentangling task transfer learning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3712-3722, 2018.
Lenka Zéleorová and Florent Krzakala. Statistical physics of inference: Thresholds and algorithms. Advances in Physics, 65(5):453-552, 2016.
14

---

Published as a conference paper at ICLR 2025
A LINEAR NEURAL NETWORKS BACKGROUND
In this section we review the deep linear neural network theory paradigm. Importantly, we do not review other theories such as the Neural Tangent Kernel (Jacot et al. 2018) which linearize neural networks and establish the regimes when this linearization is valid. These theories are highly insightful in their own right and have shed light on the training of various neural networks in general (Jacot et al. 2018] Bietti & Mairal [219 Wang et al., [0022). However, we do not build directly on these works and so limit the scope this review the linear neural networks paradigm (Saxe et al.. 2014
2019) in Appendix A.1 and GDLNs (Saxe et al. 2022) in Appendix A.2
A.1
DEEP LINEAR NEURAL NETWORKS
The primary theoretical strategy in this work is to calculate the training dynamics of linear neural networks modules formed from these networks. A first important distinction is between deep and shallow linear networks. While deep linear networks can only represent linear input-outpur mappings, the dynamics of learning change dramatically with the introduction of one more hidden layers (Fukumizu 1998 Saxe et al. 2014 2019 Arora et al. 2018: Lampinen & Ganguli| 2019), and the learning problem becomes non-convex (Baldi& & Hornik 1989). They therefore serve as tractable model of the influence of depth specifically on learning dynamics, which prior work has shown to impart a low-rank inductive bias on the linear mapping (Huh et al., (02T). The exact solutions the dynamics of learning from small random weights in deep linear networks have been derived in [Saxe] ctt al. 2014 2019) and as result possible to obtain the full learning trajectory analytically for number of representative tasks.
To review this paradigm consider linear network with one hidden layer computing the predicted output y w X in response to input batch of data X, with P datapoints, and trained to minimize the quadratic loss using gradient descent:
L(W1, W3) = Σ I x'II
This gives the learning rules for each layer with learning rate e as:
AW 2P 2" (I where Σ is the input correlation matrix and E is the input-output correlation matrix defined as:
W²W'
AW <P W'W'
1 -xx;
D = - YX
3 =
These equations can be derived for batch of data using the linearity of expectation as follows:
AWl c dwi
I(W)
P 1
d
Σ
-V:
----WW''x, w
dwi
i=1
W WV-wwW''xxxx
1 P
EP
Σ
W (Yr - -X--W w
P
1 P
2P
Y;X, - waw1
1
Σ x,x?)
P 1=1
p
1=1
= cPW2 (N - W2W15T)
15

---

Published as a conference paper at ICLR 2025
AW
L(W2 W/2)
dW2
p
2 (Y, (x
c dW2
(Y W x.)(xx)))
1m1
eP
(Y W² x;)x
P
EP( YXX w 1 Σ
P [xx,)w
=P w Wo
By using small learning rate and taking the continuous time limit, the mean change weights is given by:
d W w ( W²W
d W == (EV² W2
dt
dt
where 7 1 is the learning time constant. Here, measures units of learning epochs. It is helpful to note that since we are using a small learning rate the full batch gradient descent and stochastic gradient descent dynamics will be the same. |aaxe al. (2019) has shown that the learning dynamics depend on the singular value decomposition of the correlation matrices, where w and e are singular vectors grouped into the matrices U and V respectively. Further, 8 and Xa the singular values of E² and gRe respectively and grouped into the matrices D and S. Then the correlation matrices decompose
Ps
| yV' =UV = Σ Aau
XXX 3 VVV = Σ 8 a ya?
Here |- denotes the rank of the matrix. To solve for the dynamics require that V and En are mutually diagonalizable such that the right singular vectors V of E are also the singular vectors of Σr. 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 Y Y hidden neurons (the rank of y which determines the number of singular values in the input-output correlation matrix) can learn the desired mapping perfectly. If this is not the case then the model will learn the top ni singular values of the input-output mapping where ny is the number of hidden neurons (Saxe et al.) 2014). To ease notation for the remainder of this section we use nı to denote both the number of hidden neurons and rank of yur
that it
We now perform a change of variables using the SVD of the dataset statistics. The purpose of this step is to decouple the complex dynamics of the weights of the network, with interacting terms, into multiple one-dimensional systems. Specifically we set:
2 -W2R RT:
W RW VT
where is an arbitrary orthogonal matrix such that Rd R 1. Substituting this into the gradient descent update rules for the parameters above yields:
W 2 t0 W2
dt
(RW VT) =RW2 UT (USV- - -UW² RTT RW VTVDVT)
dt
(RW) VT) =R (SVT -W" DVT)
W - (s - ( D)
dt
16

---

Published as a conference paper at ICLR 2025
and
W2 - - WW
dt
T (UW R²)) =UUSV - UW RCT RW V V DV VVW-"VVWWWWW" RT
di
(UW R²) =US -UW D)W Ra
W WW- dt
Here we have used the orthogonality of the singular vectors such that VTV 7 and UTU Importantly, if the assumptions that the weight matrices align and the correlation matrices are mutually diagonalisable hold, then 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 initialization. However, empirically it has been found that the network singular values rapidly align to this required configuration (Saxe eet al. 2014 2019) and this has been termed the "silent alignment effect" (Atanasov al. (021)) However, if this alignment does not occur fully then the weight matrices after the change of variables will not be diagonal. The alignment assumption corresponds the assumption that the network is feature learning perfectly. While the need for perfect feature learning is strong, it is valid start for the questions around feature learning considered in this work. the assumption of the correlation matrices being mutually diagonalisable does not hold then the update equations simplify to:
d W' W2 -W" w2"'D D) T di
d W² de
and
D)
where V now denotes the right singular vectors of V and are different to from L We note that the remainder of the derivation requires both assumptions, however these two equations provide valid dynamics reduction which could be sufficient for many interesting cases. If the weight alignment holds then this is especially true as VTV would become a small interpretable matrix used to align the singular vectors from the correlation matrices. Thus, if we let da be the a-th mode of S and define 5n similarly for D, w for W and w for 2 then we can write the individual mode dynamics as:
d dt W. wā(Aa (
d w wa(2n 5.
dt
In general wt and w² can be different but if they are initialized with small values then they will be roughly equal. We study this balanced setting and assume to be true for all dynamics calculations in this work. Thus let w = " and track the dynamics of an entire mode as w, = wdw!. Using the product rules this gives the separable differential equation (we drop the dependence on a now for notational convenience):
w - (T d = + (T n )w²
d
dt d dt d di d
dt
=w (w/ (2 w(w' 8)) + (w) (w=(2- www/ 8))c²
-ww )²(1 - w2 2) + (w²)2() w2w'd)
=u( w)) w(1 wd)
dt d
22 wo)
dt
This is a separable differential equation and integrating to solve for t then yields:
T du
dt =
2 w(1 wo)
T pr
dw
t =
2 we
w(1 - wo)
17

---

Published as a conference paper at ICLR 2025
T wt (2 w00) 2X wO(A w 6)
1 In
where is the time taken for the mode to reach value ((t) w from the initial strength w(0) a. By re-arranging the terms we can obtain the dynamics of the mode for all points in time:
Aa/ba
Ta (t)
1
m exp( 23
All together this means that, given the SVDs of the two correlation matrices, the learning dynamics can be described explicitly as:
|XX*
= UAA())+ Σ
na (a)
where A(t) is the effective singular value matrix of the network's mapping and the trajectory of each singular value in A(t) described by (t) which begins at the initial value no when 0 and increases to T. = A 6a, as f + co. From these dynamics it is helpful to note that the time-course of the trajectory is only dependent on the Σn singular values. Thus, E affects the stable point of the network singular values but not the time-course of learning.
To make this setup more concrete, consider the example shown in Figure 7 For simplicity in this case we pick the input to be the identity matrix where each datapoint is I-hot vector. This 1-hot vector then depicts the index of an item in the dataset. The task then is for the network to produce set features or properties for each item. For example, when asked about the "Canary" item the network must predict that it can "Grow", "Move" and "Fly" but must not indicate that it has "Roots". Importantly, this dataset has a hierarchical structure defined by the output labels. For example, all items grow, animals can the hierarchy. Move and roots discriminate between animals and plants. Then each item has its own feature unique to it. This discriminates types of animals (canary salmon) and types of plants (oak vs rose).
and plants have roots. Thus, the grow item depicts the top level of
The SVD of this dataset produces the singular vector matrices U and VT depicting output and input "concepts" respectively. Due to the hierarchical nature of the dataset the input singular vectors discriminate between levels of the hierarchy. For example the second row of V has positive elements detect animals and negative elements detect plants. The last two rows identify types of animals and types of plants. The output concepts then are the features corresponding to each level of the hierarchy. For example, the second column of U which corresponds with the animal vs plant discriminator of V has positive elements for the move attribute and negative elements for roots. Taken together U and V depict that animals can move and do not have roots. The matrix then depicts the strength of association from input to output concepts. It is important note that these interpretable input and output singular vector matrices and the mode strengths arise purely from the dataset statistics.
Based on the dataset statistics in terms of the association it has between concepts- its mode strength when written in terms of the dataset SVD. Thus, we again assume that the singular vectors align very early on in training and it is sufficient to consider just the decoupled modes. If this is the case then the full neural network training dynamics can be obtained in closed-from using Equation 7 In this example we see near exact agreement between the closed-form predicted dynamics and the actual dynamics of training the neural network.
then determine the linear neural network training dynamics
A.2 GATED DEEP LINEAR NETWORKS
Saxe et al. (2022) extends the linear network dynamics framework above by incorporating a gating mechanism. One primary benefit of the gating is to add non- non-linearity to the network without needing to add activation functions. As result the training dynamics for this model also tractable. The
18

---

Published as a conference paper at ICLR 2025
a)
g
U
VT
Grow
Move
Modes
Propertis
Roots
Properttes
2
Fly
Swil
0
2
Canery
Salrro
Chtx
Roan
Learea
Modes
Fotals
Salmon
Casik
Output "Concepts" Object Features
Input "Concepts' Hierarchy Levels
Modes
tomms
b)
E
A(t)
VT
Grow
Move
M
Modes
X
Projentcs
Roots
Propertim
R
Fly
Swin
Comayyy nems
Cads Sainton
Leavee
Petala
Cenary
Salmon
Cate
Roaa
Training Dynamics (Equation 2.1)
Modes
c)
A
Exact Simulated
1
Ta (t) A2
222
Aa.4
3,4
0
2
Figure 7: Example dataset and linear dynamics from Saxe al. (2019). (a) Dataset with hierarchical structure defined by the output labels and the corresponding SVD. (b) The linear neural network decomposition in terms of the dataset SVD. A(t) will develop over time to learn the association between modes from S. (c) The dynamics of the linear network mode strengths over time. Initially all modes begin at but increase their final values where they have learned the dataset statistics. Dataset inspired by Rogers et al. (2004), and the exact setting and dynamics are from (Saxe et al. 2019).
framework is known as the Gated Deep Linear Network (GDLN) and defines a dynamic neural architecture by using an "architecture graph" where a particular walk through the architecture graph provides a network instantiation which is only used for subset of the training dataset. Every edge in the graph is matrix of learnable network parameters and every vertex Consequentially. vertices with only out-going edges correspond to the input layers of the network (a layer of neurons who's values are specified by the dataset, typically denoted as x; E X). Similarly, vertices with only in-coming edges correspond output layers of the network (typically denoted as &i which provides an approximation of some ground truth label ys).
layer of neurons.
When presented with an input datum a set of gating variables 9N and control the flow of information through the graph by switching off neurons within a layer (gx) and edges in the graph (ge). Thus, the pathway sequence of operations the datum follows through the network is determined by the gating operations. Importantly, once determined by the gating, the network architecture is static and acts as linear neural network. Thus, this static subnetwork is amenable to the same analysis as the deep linear networks above by utilising a layer-wise (each weight matrix connecting two vertices) singular value decomposition. The analysablity of the model is due to the linear nature of these "pathways" through the network which have no activation functions. As result it is possible to explain their training dynamics as though they are linear neural networks (Saxe et al. 2014). The benefit of the layer-wise SVD dynamics approach is that the subnetworks being tracked do not compete interfere as they are decoupled. As result, all that is required to extend the deep linear network dynamics is account for the frequency with which an edge (a set of weights) is used in the architecture graph. More frequently used edges will naturally learn quicker as a result. Thus, GDLNs are a class neural network architectures composed of number of linear neural networks which are selectively connected to different portions of the
19

---

Published as a conference paper at ICLR 2025
imput feature space, output feature space and only used for subset of datapoints. The nonlin- earity then comes in how these pathways are allocated which controlled by a discrete gating variable.
Formally, let I' denote a directed graph with nodes V and edges E. Each node v E V represents a layer of neurons with activity ht E Rell where |u| E N denotes the number of neurons in the layer. Each edge E E connects two nodes with a weight matrix W. of size [t(e)] x |s(e)] where s,t E - V return the source and target nodes of the edge, respectively. It is also helpful to generalize this notion of edges to paths p. which are sequence of edges Wp. Thus, paths also connect their source node s(p) to their target node (p)) We denote the portion of a path p preceding particular edge e s(p,e) and the portion of the path subsequent to the edge Kp,e) e) such that W, = Kip.e) W. W s(p.e). We collect nodes with only outgoing edges into the set In(T) C V and call them input nodes. Similarly, output nodes only have incoming edges and are collected in the set Out(I') C V. Activity is propagated through the network for a given datapoint which specifies values 2e E Rl for all input nodes v E In(T). The activity of each subsequent layer is then given by hy Tv Σ Bo 9y W hh where gu is node gate and 9y an edge gate. Thus, the gating variables modulate the propagation of activity through the network by switching off entire nodes (go) and edges between nodes (fq).
We train the GDLN to minimize the L2 loss averaged over the full dataset of N datapoints using gradient descent. The i-th datapoint then is a triple specified by the input æ², output labels y and specified gating structure g":
L((W))
Σ
Σ
1I - M1l12 where Yo R for v E Out(T)
2N
1=1 vEOut(T)
If we denote Pee) as the set of paths which pass through edge e and T(v) as the set of paths terminating node then the update step for single weight matrix in the GDLN is:
W. SL((W])
Σ
WT E(p.c)
(VT(p)
Σ (iP'
WT Wippe) YeEE (8)
dt
oWe
pEP(e)
JET(((M)
Thus, the update to a weight matrix We is determined by the error at the end of a path which contributes to:
(D)
Σ
(ip'(ip)
JET(((M))
summed over all paths it is a part of p E P(e). Notably, all dataset statistics which direct learning collected into the correlation matrices:
E (p)
Σ S 1=1
[b) ET N Σ of (
1
N
Crucially, these dataset statistics now depend on the gating variables g. indicating that each path sees its own effective dataset determined by the network architecture. From this perspective, each path resembles the gradient flow of a deep linear network (Saxe et al.! 2014 2019) which have been shown to exhibit nonlinear learning dynamics observed in general deep neural networks (Baldi &
Hornik 1989; Fukumizul 1998] Arora al. 2018 Lampinen & Ganguli 2019 The final step then is to use the change of variables employed by the linear network dynamics (Saxe et al. [2014] 2019). Assuming the effective dataset correlation matrices are mutually diagonalizable, we write them terms of their singular value decomposition and the path weights in terms of the singular vectors:
Σd'(p) (p) Uz(p)
(T) V
D(. s()) We(t) Re(e) (c) (9)
where S(p) and ((p) are diagonal matrices, (t) are the new dynamic variables. R, satisfies R R. = for all v E V, and for input and output nodes Ra(s) = Vs(g) and Re(;) = Unp)) respectively. This change of variables removes competitive interactions between singular value
20

---

Published as a conference paper at ICLR 2025
modes along path such that the dynamics of the overall network is described by summing several "ID networks" (one for each singular value) resulting in the "neural race reduction" Saxe et al. 2022):
Be di
Σ Bpte pEP(e)
S(p) - Σ B,D()
VEEE
(10)
JET((C)
This reduction demonstrates that the learning dynamics depend both on the input-output correlations of the effective path datasets and the number of paths which an edge contributes. Ultimately, the paths which learn the fastest from both pressures will win the neural race.
B AN ALGORITHM FOR IDENTIFYING GATING PATTERNS
In this work, our findings on ReLU networks did not rely on the manner that the ReLN (specifically the appropriate gating patterns) is obtained. This is also why it was necessary to prove the uniqueness of the ReLN, for example Proposition 4.2 demonstrates that there is no other ReLN which matches the ReLU loss trajectory perfectly in this case. Thus, we could be certain of our findings regardless of where the gates came from. However, since a contribution is the general connection between the ReLU and GDLN, it is helpful to provide simple algorithm which can identify the ReLN in the space of GDLN architectures. We note that it is always possible to find a GDLN which imitates the ReLU network,
discuss in Section 3
Our algorithm for identifying ReLNs shown in Algorithm| I and is based on a simple k-Means clustering. The primary idea is that we sample the hidden layer representations from the ReLU network at various stages throughout training. This makes it easier to detect certain structures which emerge in the gating patterns such the common pathway emerging early in training. We then the training of the network multiple times to mitigate the impact of symmetries emerging which do impact the functional mapping of the network, such as permutations and neuron splitting (Simsek et al. 2021 Martinelli et al., 2023). Each sampling from the hidden representation will have shape (H,N) where H is the number of hidden neurons and N is the number of datapoints. We stack these samples across training runs and timesteps (vertically). Thus, if C' samples are taken then the ultimate set of datapoints to be clustered will be of shape (C * H, N). Conceptually, this is C * H hidden neurons all performing some gating over the dataset. We cluster these hidden neurons (the C' H rows) with k- k-Means which results in k clusters with N didimensional centroids. As typical of k-Means, A is hyper- hyper-yaramee parameters which choose. These centroids provide the gating pattern for one type of hidden neuron and will become a gating pattern in the GDLN. Conceptually, if enough neurons of the same type, they form modules and it is these implicit modules we aim make explicit in the ReLN.
Algorithm 1 A preliminary algorithm for finding a ReLN. This follows a simple K-means clustering algorithm, but with samples taken throughout training such that is easier identify pathways through the network as they emerge.
Require: num_ trainings > O, num_epochs > 0, > 0,(X E RdxN Y E RRXNN (the dataset), H Z, K E
Ensure: a < e for sufficiently small E
for i in num cum_Iraining do
Wo E R (N(0 W, E RRXH ~ o)
for j in num epochs do
(Wo. W1} + gradient descent( (W,, X) >> Apply gradient descent update step jf mod 100 = then Sample at different times to find different structure as it emerges & Sample latent representations with ReLU activation sample_binary sample) > Threshold the sample to indicate a neuron is active samples vstack(samples.san samplebinarry))
sample = maximum(Wo,X.
Stack binary latent representations > Each sample appended vertically appears like a new neuron
end if end for end for centroids K-mmeas(samples,.)
return centroids
21

---

Published as a conference paper at ICLR 2025
The results of using this algorithm to identify ReLN for the 3-context dataset with different structure between datasets (see Section 4) is shown in Figure [ Each row in this figure corresponds to a cluster and the columns correspond to the datapoints in the dataset. For example, the green row in the 2-Clusters set reflects the dominant gating pattern of one cluster hidden neurons across the whole dataset. In this case, the hidden neurons allocated to this green cluster were typically active in context 2 and 3 but inactive in the first context (the first 8 columns have no activity which corresponds to the datapoints in the first context). We show the clustering when k: 2, all of the clusterings find a common pathway, likely due to its emergence early on in training. From there they begin to identity mixed selective modules (which are active in two contexts). However, worth noting the effect of averaging on these clusters. This is easiest to see in the 3 Clusters case where the cluster which is active for all contexts also trying to capture the cluster for the hidden neurons active in context and 3. Thus, the middle portion of that cluster is dulled compared to the others. In such cases difficult to interpret what exactly the gating pattern should be and is a sign that more clusters are needed. However, this indicates limitation of this approach it relies on the clusters to have a consistent gating pattern for each datapoint throughout training. We do not limit the GDLN such condition. We note that when using 4-clusters the exact gating pattern used in Section emerges.
and 4. We note that
Clusters
C1
13
C2
Centroid Index
03
3 Clusters
C3
C1
C2
Clusters
C3
C4
10
15
20
Neuron Index
Figure 8: Clustering find ReLNs: Clusters found with varying number of centroids for the
dataset in Section 4. Note the inconsistency in gating patterns arising from having too few clusters. Consequently, centroids have neuron activity between 0.0 and 1.0 indicating that the gating is not consistent across data points.
Since the algorithm is based on clustering, appropriately. This has clear effect on the performance of the model as the number of clusters will dictate the number of unique gating patterns our GDLN can use to imitate the ReLU network. We can, however make this less qualitative by using the typical approach of finding elbow in plot of number of clusters compared to performance metric. In this case the performance metric is very rich and can be the mean square error between the ReLU network and the GDLN with the identified gating patterns. This can give a very precise measure of whether the number of clusters was correct. We show this elbow in Figure 9 and see that the 4-Cluster model is the appropriate choice for the elbow, as expected.
still necessary to pick the number of centroids
The algorithm presented here be seen a very basic form of meta-learning for a GDLN, in the sense that the GDLN is being trained for meta-task of imitating the ReLU network, while learning to minimize its loss. However, the space of GDLNs is also larger than the space of ReLU networks. For example, ReLU networks are only able to implement gating patterns as a function of their input, while GDLNs are capable of gating based on rules, external patterns or previous datapoints. Here we are concerned with GDLNs which are restricted to gating patterns implementable by ReLU networks. However, many works have demonstrated a benefit from imposed architecturally modular networks which specialize towards solving particular and interpretable sub sub-problems within a given task (Andreas et al., 2016 Hu et al. 2017 Vani et al. (022). Indeed, one proposed direction to obtain modular networks is through meta- meta-learning (Lake) 2019 Lake & Baroni 2023). Mixed-selectivity, and the closely related concept polysemanticity (neurons being activated by multiple semantic unit or features) (Olah et al.) 7a Lecomte et al.) (024) have typically been thought to be unrelated so such modularity and the systematic generalization (Ruis Lake 2022). However, as we have shown in this work ReLU networks are able to form interpretable modules across contexts which
22

---

Published as a conference paper at ICLR 2025
Squared Difference Between RelU Loss GDLN
300
250
200
150
100
Number of Madules
Figure 9: Clustering to find ReLNs: Finding the elbow to determine the minimal number of clusters (linear modules in the ReLN).
are mixed selective and form coupled modules. Thus, a more advanced meta-learning strategy for GDLNs, than the approach of this section, which learns the gates based on entirely different rule or set features may allow a GDLN to perform computations unavailable to the ReLU network and by extension ReLN. This appears to be promising future direction.
C LEARNING DYNAMICS DURING TRANSITION TO NONLINEAR SEPARABILITY
Deep neural networks can learn representations suitable for solving nonlinear tasks. Here we investigate their learning dynamics on simple class of tasks. We will qualitatively understand the dynamics of training single hidden layer ReLU network by reducing it gated deep linear network.
Dataset. The canonical example of a nonlinear problem the XoR problem. We consider slight generalization: XoR problem in the first two input dimensions, but a linearly separable problem in third. Our dataset X R y E Rixd consists of four data points three dimensions,
X
-
-
(11)
A
A
A
A
[-1
1
1
-1]
(12)
where A 2 0 is a parameter that controls the linear separability of the points along the third dimension.
Our goal will to understand how a ReLU network transitions from a regime in which the data points clearly linearly separable (A >> 1) to the regime in which they are not (A 0, the classical XoR task).
GDLN with linear gating structure. We first solve the dynamics of the GDLN depicted in Figure 2p of the main text. This network contains two pathways, one gated on for the positive examples and one gated on for the negative examples.
The effective dataset statistics for the positive examples, which we denote as E (p), is therefore
1]
1
(13)
4
A
A
=0
+]
(14)
23

---

Published as a conference paper at ICLR 2025
1
-1
(p(p)))
=
-1
1
1
(15)
4
A
A
A
A
0
1 -1
0
(16)
0
A2
By symmetry, the loss dynamics on the negative examples (driven by effective statistics Σ.e (n)) will be identical to the positive examples. Because only one pathway is active on any example, they develop independently.
To solve the dynamics, leverage analytical solutions for deep linear networks describing each pathway. The dynamics of a deep linear network depend on the singular value dccomposition of the effective input-output correlations E (p) USV and input correlations E (p.p) p) VDVT.
The singular value decomposition is
y (p) = 0 0 -
0 ']
(17)
so the singular value is s = 2/2. The input variance in this singular vector direction is = / From the known time course of deep linear networks, the effective singular value in this pathway will be
s/d
a(t)
(18)
(1
dog )e
2st/7
1/A
(19)
(1 Dag )e
At/T
The loss from this pathway Cp therefore is
(t(t) = (p) saatt) + 1/2da(t)² = 1/4- Aa)))/22+++ + a²a
(20)
(21)
By symmetry, the contribution of the other pathway for the negative examples is identical. Hence the total loss twice this single-pathway loss, yielding
C(t) 1/2 (t( + /tt))//
(22)
the expression plotted in Figure (blue) of the main text.
GDLN with XoR gating structure. instead we have a GDLN with four pathways each gated on for exactly one input sample, then the effective dataset only contains one input sample. Taking the
pathway gated on for the first example, for instance, the singular value of (1) [1 1 A is
A
The associated input variance is d + 1e Substituting these values into Eqn. 19
16 yields the effective singular value trajectory, and similar steps yield the total loss
[()) = 1/2 4sa(t) + 2da(t)
(23)
the equation plotted in Figure 2e (orange) of the main text.
By symmetry the remaining three datapoints have identical loss dynamics. Because all pathways concern disjoint examples, they evolve independently.
We now ask for what range of will the linear separability GDLN learn faster than the XoR GDLN. The critical A separating these regimes is the one where both GDLNs learn equally quickly. The learning speed of deep linear networks of order ((1/s) where is the singular value. The crossover point is therefore
1
A
A/2
V
(24)
8
16
4A = 2 A
(25)
A = 22/3.
(26)
By testing points, we see that for 0 < A < 2/3 the XoR GDLN is faster while for A > 22/3 the linear separability GDLN is faster. Hence ReLU network does not always exploit linear separability.
The linear margin is 2A but for small enough margins, the race will be won by the XoR GDLN, implying nonlinear decision boundary.
24

---

Published as a conference paper at ICLR 2025
D PROOF OF LEMMA 4.1
On the input X (as defined in Section 4 and given the definition of GDLNs and ReLU Networks, for any gating pattern implementable by the ReLU network there is G(ze) such that G(rz), step(hy) or G(x);
Lemma 4.1:
I(x) V j E [0, H] where h is the pre-activations of the ReLU network's 's hidden layer with H neurons, j is the index of a hidden neuron in either network and ze is the portion of the data point with contextual features.
Proof: We consider which functions are implementable by a ReLU network given that it is restricted to using the same weights for both gating and association. Without loss of generality consider the simplest instantiation of two objects with two contexts. Thus, there are four datapoints, specifically 3 - [1, 0, 1, 0]. x2 [0, 1, 1, 0). 23 [1, 0, O, 1] and 24 0, 1].
First, we consider if combined gating strategies can be implemented. By a combined gating strategy what mean is that the gate is determined by both the object and context feature in combination. Thus there must be at least one object which is gated differently for the same context as another object (or else the gating could just be done based on context). Similarly there must be a context where an objects gating changes (or else the gating could be done based on object). In other words, there is no consistent gating for object or context in isolation across the whole dataset. Assume a combined gating exists. This means there is at least one hidden neuron which is on for 22 and 23 but off for Ii and T4. Thus, the network must have sufficiently negative weights connecting to this neuron such that its pre-activation is negative on It. Let these weights be (a, b, c, d]. Then a + <= 0
and b+ d <= 0. This means a <= -c and b <= -d. Conversely, b + > O and a + d 0. -d. Taken together <= -c and b > -c means that a < b. Additionally, b <= -d d and a > -d means that < a which is a contradiction. Thus, no such hidden neuron with combined gating can exist. This means that the network must gate consistently for either context or object features.
Thus, b>
and
The gating strategy based only on items can be rules out the mapping from context to output is nonlinear by construction. Thus, regardless of the number of latent neurons there is no linear mapping from contexts to output alone which will achieve no error. There is an edge case this proof which we cannot rule out, that of allocating a single neuron for each data point. Here we cannot obtain contradiction as there is no other data points contradict with. This is also the typical strategy employed by ReLU networks on XoR tasks such as this one. We will cover this case in Proposition 4.2 Thus, the only remaining gating strategies implementable by a one-layer ReLU network the dataset described in Section 4 is to gate consistently for the context variables. This leaves the common pathway, a second of contextual pathways active for two contexts, and finally contextual pathways active for single context.
E PROOF OF PROPOSITION
Proposition 4.2: For input X defined in Section 4 there is a unique G
G* (X) (up to symmetries such
as permutations) such that LEDIN (Wo (t), G (X)) (WLLLWWWWH)) Vt R and and alternative gating strategies G' ('(X) then LGDLN
(aonnnn(())) (W(),, (X)) < LODEN (W),,, (X))
Proof: To begin we prove generalization of the Cauchy Interlacing Theorem for non-square matrices
which closely follows the strategy of Thompson (1922). The aim of this is to show that if an input or output feature is removed from a pathway's effective dataset that the first singular value will decrease. Since the singular values determine the loss trajectory of the network, this will then mean that the loss curve will no longer fit the ReLU network. Thus:
For any matrix M lets denote the submatrix of rows i1..... ēp and columns J1...., Jg as Mii..... éplii, ..., Jy]- To ease notation let μ (i..... }p} and
{ji,..., Jg... Then B
[ppp] since it is a submatrix of Egr. If U = Im and V = I, (the identity matrices of m and dimensions) then B can also written as B = [M] 1, 1,.., m Eyyz V[1,.. .., n/u). Thus:
BBT= = U[[ 1, ...,,
..l......]. m V[L, 12 null [v|L, n n UT ( I,,.,, m|]
U l...,,.., m] m|]]
[I,.....]]] m|14
where A = Ens [I,, n|u] E mx with ordered, non-zero singular values of a1 022 .... (minn(m..)) Thus BBT is principal p-square submatrix of the m-square symmetric matrix U UA" Since
25

---

Published as a conference paper at ICLR 2025
BB" is guaranteed to be symmetric we know it will be diagonalizable. Thus, there exists orthonormal basis of eigenvectors (61....,min((p
corresponding to eigenvalues {,37, P min(p..) We define Gj- - span ..., by] (the span of the first j eigenvectors of BB") for S min(p; q) and Sj = span [i;; -. mm.)()).)) (the span of the last min(m, q)
eigenvectors of AAT). We also
define the subspace:
H = 1t [1, [I.... mlugg.9 E G2} There exits a unit length vector 2 = U
mlp]z for E Gj which lies in H, n Sy, as if there is not then the dimension of H2 nS, would be + min(m, g) + which is impossible in Rmiin(m,.), Additionally:
[1. 1
2
r;bs
and
(2, = -T
i=1
(since the eigenvectors of symmetric matrices orthogonal). Rayleigh Bitliigh-Ritzeaeo quotient for BB" and a unit length vector et
We begin by considering the
min(p.g)
REg (v) (BB" u. u) BB Σ B
(v,0)
Here rē acts as weighting each eigenvalue for how much it contributes to the sum. Thus to minimize the quotient we place all the weighting on the lowest eigenvalue 53 by picking e such that r; = for E j } and r, = This result is known as the Min-Max Theorem. similar result can be shown where the j-th eigenvalue also results from maximizing the quotient over the span of the bottom eigenvector beginning with the j-th eigenvector. Thus:
83
min
RBB (u) by the Min-Max Theorem
-EG,E|1-|1-1
min
(BB"M) v)
< (BB"z, 2, since Gj does not necessarily minimize the quotient
(U1m41, AAA [1, ml))2,22) 2)
(AAT (UT [1,.,,,,ll ml/1=)) [ [1, ma..m//uii mlu]z) by the linearity of the inner product
(AATZ, 2,
I max aES,il|v|==1
(AA" U. ū) since ă Sy does not necessarily maximize the quotient like
= max RAAT(A)
aESgillo|==1
aa by the Min-Max Theorem
Thus 3j ay. Similarly, AT A of az. a min((..)) Thus, AT is a principal q-square submatrix of the n-square symmetric matrix VTΣT Eye Since A is guaranteed to be symmetric know it will be diagonalizable. Thus, there exists orthonormal basis of cigenvectors {a1, ..., Ominn(m.q)) corresponding to eigen- values {a? .... o min(m.g)) }- We overload the notation and define Gj span (....., az] (the span of the first j eigenvectors of AT for j < min(m, q) and S, = span[o,, .... mi) (the span of the last min(m, n) + eigenvectors of Σx Dyrr). We also define the subspace:
// V[1,..n|vv) with eigenvalues of
H, = (V[1, nlvlg. E Gj}
There exits a unit length vector V[I, ., n|p]= for E G, which lies H, n nSj. as there not then the dimension of Hy S.S would be j + min,m, n)) - j + which is impossible mm))n Additionally:
T:a;
and
(2,z) = TTZ
i=
26

---

Published as a conference paper at ICLR 2025
(since the cigenvectors of symmetric matrices orthogonal). Rayleigh- Brtteigh-Ritz..oo quotient for A A and a unit length vector u:
We begin by considering the
min(m..) Σ air?
RAT (v) = {A Ar, = A
(e,0) U)
Here r² acts as weighting each eigenvalue for how much it contributes to the sum. Thus to minimize the quotient we place all the weighting on the lowest eigenvalue a² by picking e such that r; = O for E (1,-..j- -0 1}
and rj = 1:
vECgille| RRTA(y) by the Min-Max Theorem
min
min (AT Au, v)
vEGyillo| =1
(A Az, since E G2 does not necessarily minimize the quotient
=(VT(11,, n1 I yut
Syz [I1, n 1)2,,)
(2 Vyn (V[1, n 1/22).
by the linearity of the inner product
(I Eye 2, 2)
I max
( Eyz ū) since 2 Sj does not necessarily maximize the quotient like
DESyillu|==1
= max Ryy E (7)
DEStillo|==1
=o by the Min-Max Theorem
Thus a; S J,. Putting both inequalities together we obtain the result: By S a, S o, with the particularly important case of 31 S Ø1- This means that shared pathway will always be learned faster than a pathway which considers only a subset the input features or output labels (which input features is determined by μ and which output labels determined by v).
Next we consider whether a different gating strategy exists which gates along datapoints instead of features. We note that the output of the common pathway is element-wise positive for all datapoints and all input features are positive. Consequently, let E = YXT be the correlation matrix of this dataset. If we factor out Y as Y Y + Z where Y is Y with a datapoint removed (its column set to 0) and Z is the matrix of zeros everywhere except for the column removed from Y where equal to Y. Then D = YXT = (Y + Z)XX Y + ZXT. Now let be the eigenvector corresponding to the largest eigenvalue of Y and v be the top eigenvector for the full dataset E Then YEE v 2 I""y"=" 6 = X + Zx") + Z Y + Y xx"xX'X positive orthant where the entire dataset lies. Thus, the common pathway can only train faster with the addition of more datapoints. Thus, the common pathway is the fastest possible subnetwork for the given dataset, and any other module will train slower (there is an edge case here where the added datapoint is orthogonal to all previous datapoints, in which case X 0, however no such datapoint exists in this dataset).
XX"XÝTo 8 X > Y X since will be in the all
We can now rule out the edge case from Lemma 4.1
Once the common pathway has finished learning there is no correlation between the context features and labels. As a consequence in a case when single module all activity would need to be learned by the item
single datapoint allocated feature. In our dataset, removing all other items in a context would then be the same as removing their individual input features. From the Cauchy Interlacing Theorem, we know this can only slow learning.
Finally, we consider whether alternatives to the context sensitive pathways exists. We note that from Lemma [41 with the addition of common pathway, there are only two available options to the remaining pathways, as these pathways can only gate consistently using context features. Thus, either the modules must be active for one or for two contexts. In this case we can easily simulate this to determine which is quickest and see in Figure 4 that the architecture which uses modules that are active for two contexts wins the race. Indeed even if all six possible modules are trained simultancously, the pathways active for a single context remain inactive. Thus, the pathways active for two contexts have a unique loss trajectory and when paired with the common pathway have a unique and fastest loss trajectory. Thus, the ReL.U network uncovers the neural race winner.
27

---

Published as a conference paper at ICLR 2025
F GDLN DYNAMICS WITH THREE CONTEXTS
We now derive the dynamics when there are three context specific portions of the output space. Specifically, we have dataset D (Z, C, Y) with object, context and output feature matrices I E RNN C' E 3x and Y E 2 Ix respectively. Here N denotes the number of datapoints. We concatenate the object and context features vertically to form the input matrix: X E R 3 Every object (depicted by a one-hot input vector in Z) can be encountered in every context (depicted by another one-hot vector C). is convenient to partition the input matrix into three blocks with the same one-hot context vector:
X2, X3] where X; E R I Similar partitions follow for I and C separately. Both models are told which of the objects to produce the features for as well as in which context the object is encountered in. The models must then produce the correct output features, some of which are appropriate for all contexts and others only in one specific context.
Our ReLN trained to imitate the ReLU network loss trajectory and dynamics composed of four pathways. All pathways have linear layer of 100 neurons resulting in the same number of total hidden neurons as the ReLU network. The first pathway receives all objects and contexts and maps to the full output matrix. We call this the common pathway and denote by with weight and output mapping: W. = (W?, W and ÿ = W X. Thus, to model this pathway's 's loss trajectory we apply the linear dynamics equations of [aaxe et al. (2014) to the input-output correlation: E' YX = U.S.V V.. The common pathway's 's dynamics then follow W?W! = U.A, V?" and will have the mapping Y (t) U. A=)((t) V X at all points in time. It is assumed that the network aligns to the dataset singular vectors immediately with U. and VT, The dynamics of the network can then be obtained through the decoupled singular value dynamics of:
Sa/Do
(1) = (1 Damo ) exp( 25.
where Sa is the a-th mode the input-output correlation singular value matrix S, Da is similarly the th mode of the input correlation matrix (XX*) and To is the initial state of the network singular values Ta. Thus, To begins with a value of To and increases until its convergence point at Sa/ Da. We can then define the residual of this model as Y = (co).
The three remaining (context specific) pathways are then trained on this residual. The context specific pathways are each gated on for two context settings and connect to all output features. In addition, while these pathways are used in two contexts they are not shown the input context as feature. The result is that these pathways can only use the input features Z, in their mapping when they are gated on. Since these pathways are trained at once and have dynamics which depend on other pathways we now use the GDLN framework (Saxe et al. 2022) to determine their dynamics (we review this setup in Section 2 in the main text). The learning rule for a weight matrix in a GDLN follows the update equation (we use the brackets (.) to denote averaging as typically used in statistical physics to simplify notation):
d
We = SL((W))
T dt
Ve E
OW.
Σ PPP
9r Wp Mp.c) (1 = WT Wa(pre)
3.2.0
We note that for our architecture there is only one pathway that each edge is involved in. Thus we may drop the summation over pathways and denote the one relevant pathway with p. We also apply the usual linearity of the expectation:
= Wi
(Io
WT a(p.e)
9.3.9
y.5.g
The next step is to determine the output of the network which impacts pathway p. Specifically we substitute in for h(6) the sum of all pathways mapping onto t(p) the terminal point for pathway p:
h((p)
Σ
JET((()
28

---

Published as a conference paper at ICLR 2025
Here [tt(p) is the set of all pathways leading to the same terminal node as pathway p. Thus denotes each of these pathways with corresponding mapping Wj (note W, here is the product of all matrices along pathway j). Finally ((+(j)) is the input pathway j since s(j) denotes the source of j. Substituting this into the dynamics equation we obtain:
-(80
= Wiip.ee
9v
Σ JET(U(P)
9 (WFSs))
W s(p.c)
u.s.g
M.
W t)p.e)
E
W3 (1)7515))
W s(p,e)
JET(((M)
v.2.g
= W e(p.c)
(p)
E (i jET(UP))
W
Thus, the dynamics for one layer of weights (W:) in single pathway (p) depends on the input- output correlation for that pathway (2** (p)), the input correlation for the pathway itself (E* (p,p)) and the input correlation for the pathway with the input to all other pathways sharing the terminal point (E* (j. # p). Thus rewrite the dynamics in terms of these three kinds of correlations separately:
W H(p.e)
E (p) Wp
(jP)) *+PET(U())
E
WT s(p.c)
We note then that for the particular setting we are working in " (p. p) and substitute this relationship into the dynamics above. Now the input correlation terms do not depend on the alternative pathways taken to the terminal points.
(p. p). Thus, we denote
WT t(p.e)
EV (p) W2S
Σ
W
Dr
WT s(p.e)
W2 dip.e)
EV (p) -W
Σ
W
W a(p.e)
2
(*PET(((p))
T W. WT l(p.e)
(V"
Wp +
1
Σ
W
Σx
W a(p.e)
iAPET(((p))
We now apply the change of variables obtain the dynamics reduction:
d
(Reee) (e = Rete))
[UGn)
di
Ui( ByVT
1
Σ
B,y Vo D(.
(Rece) Be(p.e) VoTo) salp)) )²
s(p)
(+PET((())
We can then simplify the expression by multiplying matching eigenvectors and noting that Va(p) = Vigi since all pathways share the same input correlation:
T B. = =Bip.) S(p) dt
Bp +
Σ (tpE)(
Unn D(p) Bx(p.e)
This is usually where we would need to stop for the GDLN reduction and Section where the output matrix has different structure, the "race" dynamics are obtain using this reduction. However, in Section because the structure is symmetric
exploit some more information in this task to
29

---

Published as a conference paper at ICLR 2025
continue the derivation further to closed form solution. Having determined the effective datasets for each pathway we know that U UL() = - Thus:
((p)
T Be Bi(p.c) dt
S(p)
Bp+
Σ
2 B, D(p)
Bs(p.e)
(*PET(1(p)
Finally, note that due to the symmetry of the task between contexts, all pathways will learn identical singular values. Thus B, = By V + p T((())) Thus, we substitute this equality into the dynamics reduction:
d Be Bi(p.e) dt
S(p)
Bp
Σ
Bp
D(p)
Bs(p.e)
)+pET(t(p)
S(p) BD) Bu()))
We note that for each pathway our network has two layers: E {0, 1}. Thus:
d Bu Bi(n.0)
(s()) - 1 BBD(()) Bs(
dt
BB ((() =B.(s()
Bp D())
and
d
Bi Bi(n.)) 5(p) B3 dt
2 Bstp.1)
s(p) =1(s(p)
=
5 (D) Bo
Assuming balanced solutions, which is reasonable from small initial weights we know that Bg = B We may also then switch consider the dynamics of an entire pathway and not just one layer in the pathway: By B, Bo. The dynamics of the pathway can be obtain by the product rule:
d d Bp BulT B1) + B1 Bo) dt
dt
dt
Bo S(p) Bp (())
+ +BB
S(p)
B, D())
Bp 8(p) ] BDD(
1
BBp 8(p) ((D
d Bp 2B, S(p) 2 B(D))) dl
This is separable differential equation which can solved as per the linear dynamics (Saxe et al. 2014 2019). Thus the full learning trajectory for the a-th mode of context dependent pathway is (we have removed the dependence on to lighten notation):
(2Sa/Da)
(a) (1 DB
exp((S..
This is the equation used to obtain the "closed" dynamics in Section for the three context case.
G GDLN DYNAMICS FOR FOUR AND FIVE CONTEXTS WITH HIERARCHICAL STRUCTURE
G.1 GDLN DYNAMICS FOR FOUR CONTEXTS
We now derive the dynamics when there are four context specific portions of the output space. This closely follows the derivation of the three context case from Appendix however present it here
30

---

Published as a conference paper at ICLR 2025
in full. Specifically, we have dataset == (Y) with object, context and output feature matrices E RNVA C' C R4xN and Y E 1 respectively. Here N denotes the number of datapoints. We concatenate the object and context features vertically to form the input matrix: X R N Every object (depicted by a one-hot input vector in 2) can be encountered in every context (depicted by another one-hot vector in C). It is convenient to partition the input matrix into three blocks with the same one-hot context vector: X = X2, X3, Xa] where X, R N Similar partitions follow for I and C separately. Both models are told which of the objects to produce the features for as well in which context the object is encountered in. The models must then produce the correct output features, some of which are appropriate for all contexts and others only in one specific context.
Our ReLN trained imitate the ReLU network loss trajectory and dynamics is composed of four pathways. All pathways have a linear layer of 100 neurons resulting in the same number of total hidden neurons the ReLU network. The first pathway receives all objects and contexts and maps to the full output matrix. We call this the common pathway and denote it by *, with weight and output mapping: W. = (W.. W and O W?W. X. Thus, to model this pathway's 's loss trajectory we apply the linear dynamics equations of [axe et al. (2014) to the input-output correlation: Ex Y = U.S.V²" The common pathway's is dynamics then follow W?W! = U. A. (.), and will have the mapping y' (t) = U. U_A(t).VV* VT X at all points in time. It is assumed that the network aligns to the dataset singular vectors immediately with U. and VT. The dynamics of the network can then be obtained through the decoupled singular value dynamics of:
Sa/Da
(t) (1
N., expl 25
D
where So is the a-th mode of the input-output correlation singular value matrix S, Da is similarly the a - th mode of the input correlation matrix (Xx²) and To is the initial state of the network singular values Va- Thus, To begins with value of mo and increases until its convergence point at Sa/ D.. We can then define the residual of this model Y = (oo).
The three remaining (context specific) pathways are then trained on this residual. The context specific pathways are each gated on for two context settings and connect to all output features. In addition. while these pathways are used in two contexts they are not shown the input context as a feature. The result that these pathways can only use the input features I, in their mapping when they are gated on. Since these pathways are trained at once and have dynamics which depend on the other pathways we now use the GDLN framework (Saxe et al. 2022) to determine their dynamics (we review this setup in Section 2 in the main text). The learning rule for weight matrix in GDLN follows the update equation (we use the brackets (.) to denote averaging as typically used in statistical physics to simplify notation):
d We 5L([WH) óW.
Ve E E
dt
Σ P(e)
W. tip.e) (Mco)tztie) h Watp.e)
gn
y.2.9
We note that for our architecture there is only one pathway that each edge is involved in. Thus we may drop the summation over pathways and denote the one relevant pathway with p. We also apply the usual linearity of the expectation:
( (908)19)=
WT (()ee)
WT s(p.e)
v...9
y.a.g
The next step is to determine the output of the network which impacts pathway p. Specifically we substitute in for he(p) the sum of all pathways mapping onto t(p) the terminal point for pathway
ha(p)
Σ
(ET(((P))
Here T(t(p)) is the set of all pathways leading to the same terminal node as pathway p. Thus j denotes each of these pathways with corresponding mapping W, (note W, here is the product of all
31

---

Published as a conference paper at ICLR 2025
matrices along pathway j). Finally z(+()) is the imput pathway j since s(j) denotes the souree of j. Substituting this into the dynamics equation we obtain:
((l.sn) ((p.c)
9p
Σ jET((())
96 W,W
WT 8(p,e)
yor.g
Mi
= -Wippe
(DDh(t)) u.a.g
E JET((p))
W, (1)
WT s(p.e)
(ip.e)
= Wiip.e)
2
(p)
E (jP) (TT((()))
W s(p.c)
Thus, the dynamics for layer of weights (We) in single pathway (p) depends the input- output correlation for that pathway (2**(p) (p)), the input correlation for the pathway itself (2" (p.p) and the input correlation for the pathway with the input to all other pathways sharing the terminal point (j.(jp)): # p). Thus we rewrite the dynamics in terms these three kinds of correlations separately:
W.T i(p,e)
E (j. j+PET(U())
W a(p.e)
We note then that for the particular setting we are working in Σ (j.pp) 3 (p, p). Thus, we denote Σ Ed (p. p) and substitute this relationship into the dynamics above. Now the input correlation (p.
terms do not depend on the alternative pathways taken to the terminal points.
W f(p.e)
Σ (p) W
E
W
D
W s(p.e)
j+PET(((p))
WT f(p.c)
2
Σ
W
WT s(p.c)
j+PET(((p))
2 Wp +
W. WT ((p.e)
[V"(p)-
Σ
W
Σ
WT a(p.e)
3
}*PET(U())
We now apply the change of variables to obtain the dynamics reduction:
d dt (B(e)BeR((o)))
T [UnG))
2
Σ
U,n D(..)) s(p)
(Rete) Be(r.e) VoTp) )z VoTp))
s(p)
i*PET("())
We can then simplify the expression by multiplying matching eigenvectors and noting that Va(p) Vagi since all pathways share the same input correlation:
T B. = -Bip.e) S(p)
B
2
Σ
Unn) D(p) Ba(r.e)
}*PET(((p)
Once again because the dataset structure is symmetric task to continue the derivation further to a closed form solution. Having determined the effective datasets for each pathway know that U. Ung) - V + p. Thus:
exploit some information in this
T B. B((p.c) S(p) dt
B₁ +
2
Σ
-B D(p) Ba(.e)
(
32

---

Published as a conference paper at ICLR 2025
Finally, we note that due to the symmetry of the task between contexts, pathways will learn identical singular values. Thus Bp B, V # T(t(p)). Thus, substitute this equality into the dynamics reduction:
d Be Bi(rno) dt
S(p)
Bp- -p-
Bp (+PET(U())
Σ
D(p) Bs(p.e)
- 3 (1,D(o))
Bippee
S(p)
Ba(p.e)
We note that for each pathway our network has two layers: e E {0, 1j. Thus:
d
Bu Bi(m.0) S(p) 3
D(D Bs(p.0)
dt
Bı S(p) BB, D()
and
1
-- B₁ = Br(21)) S(p) 3 BB)D(o) dl
Bz(r.1)
=1 I 5(p) -
3 B()) Bo
Assuming balanced solutions, which reasonable from small initial weights know that Bo B We may also then switch to consider the dynamics of an entire pathway and not just one layer in the pathway: By Bg. The dynamics of the pathway can be obtain by the product rule:
B1) + B1 Bo) - Bp = Bo(T d
d
dt
dt
dt
BoBo S(p) 3 B)
+ BBB
S(p) //2 B,D(P)
B, 8(p)
BDD( 3
1
B, ((()) B,
d Bp 2Bp S(p) 3 B.D())
dt
This is a separable differential equation which can solved as per the linear dynamics (Saxe et al. 2014 2019). Thus the full learning trajectory for the a-th mode of context dependent pathway is (we have removed the dependence on to lighten notation):
(3Sa/Da)
(t) 1- (1 - Di exp (2S
This is the equation used to obtain the "closed" dynamics in Section for the four context case.
G.2 GDLN DYNAMICS FOR FIVE CONTEXTS
An almost identical derivation for the three and four context cases holds for the five context case. We omit the initial portion for brevity we once again train linear pathway and the context specific pathways on the residual. In the five context case there are now five residual pathways each active for four contexts. We will pick up the derivation where the specifics of the dataset are used to obtain closed form equations from the neural race reduction.
We rewrite the neural race reduction dynamics in terms of these three kinds of correlations separately:
W.T
f(p.e) 1T - W,'
Σ
(.
W.lp.e)
33

---

Published as a conference paper at ICLR 2025
We note then that for the particular setting we working in E) (p. p). Thus, denote
p) and substitute this relationship into the dynamics above. Now the input correlation terms do not depend on the alternative pathways taken to the terminal points.
Σ" E"
(p: (p.
WT
E (p) -,
Σ
W
WT s(p.c)
t(p.c)
}+PE7(U(p)
W dp..))
[2'(P) -W
3
Σ
W
W sip.e)
4
j*PET(A(p)
T We = W
(p)
W +
3
Σ
W
W 8(p.c)
dt
t(p.c)
4
H+PET(U())
We now apply the change of variables to obtain the dynamics reduction:
T ddl (Bule = (ippet T
d
(UGn)
VoT Volp) )
U1i))BV2 a(p)
Σ
Unn
VsU) D(I.)) s(m))
(oaee) B Be(p.e)
4
j*PET(U())
We then simplify the expression by multiplying matching eigenvectors and noting that Va(p)) Va()) since all pathways share the same input correlation:
-- Be = Bi(p.e) S(p) dt
3 Bp
+
Σ
Uan D(p) Bs(p.e)
H+PET(UiP)
Once again because the dataset structure is symmetric we can exploit some more information in this task to continue the derivation further to closed form solution. Having determined the effective datasets for each pathway we know that k(p)
U7, Un) 12
J #
Thus:
d T B Bi(p.e) dt
S(p)
Bp +
Σ
1
Bj
D(p)
Bs(p.e)
12
(FPET((())
Finally, due the symmetry of the task between contexts, all pathways will learn identical singular values. Thus Bp = Bj v # T(t(p)). Thus, we substitute this equality into the dynamics reduction:
T
B. B(ip.e) dl
d
S(p) - By 16
Σ
Bp D(p) Bz(p.e)
((PET(U())
= Bi(g.e)
S(p) Bp)
Ba(n.e)
We note that for each pathway our network has two layers: {0,1), 13. Thus:
d dt
B, = Bi(n.0) S(p) 4 B(D) Ba(r
= Bi S(p) 1 11, I
and
- Bi = B1(D1))
S(p) - ( D())
Bs(
dl
= S(p) (() Bo
34

---

Published as a conference paper at ICLR 2025
Assuming balanced solutions, which is reasonable from small initial weights we know that Bo B1- We may also then switch to consider the dynamics of an entire pathway and not just one layer in the pathway: Bp B, The dynamics of the pathway can be obtain by the product rule:
d Bp = -Bll d
B2) + BT d dl
Bu) dl
di
BoBo S(() 1 BrD(
+ +BiBi S(p) B,,))
1 D())
+Bp ((() - - B,D
Bp ((p) 4
d di Bp 2B, ((p) {DD
This is separable differential equation which can be solved per the linear dynamics (Saxe et al., 2014 2019). Thus the full learning trajectory for the a-th mode of a a context dependent pathway is (we have removed the dependence on p to lighten notation):
(4Sa/Da D..)
(1) (1 DE
(2
This is the equation used obtain the "closed" dynamics in Section for the five context case.
H HYPER-PARAMETERS FOR CONTEXTUAL DATASETS
Here we provide the modules needed for each type of GDLN referenced in the main text. Importantly, these modules would still need be connected to the appropriate portions of the input and output space, as described there. As described in Appendix A we only require many hidden neurons as the rank of the input-output correlation matrix the linear pathway is aiming to solve. Finally, as long as the ReLU network and GDLN begin with the same hyper-parameters the mapping to ReLN will be valid (assuming the correct allocation of linear pathways).
used Sections 4 and 6 We provide the number of
Hyper-parameter
Value
Num-Hidden (GDLN)
100
Num-Modules (GDLN Single Modules, Contexts)
4
Num-Modules (GDLN Double Modules, 3 Contexts)
4
Num-Modules (GDLN Single and Double Modules, Contexts)
7
Num-Modules (GDLN Triple Modules, 4 Contexts)
5
Num-Modules (GDLN Quad Modules, Contexts)
6
Num-Hidden (ReLU)
700
Num-Modules (ReLU)
NA
Init-Scale
1 10-7
Num- Jum-Epochs
8000
Step-Size
0,001
Table 1: paaeee-paraeeters used for the contextual tasks in Sections and
35


Reviews:
[{"review_id": "OAo9nrlLAX", "reviewer": "Reviewer_sUq5", "summary": "This work builds provides a step towards theory of feature learning in finite ReLU neural networks by building on an equivalence between ReLU networks and Gated Deep Linear Networks (GDLNs). The authors introduce \"Rectified Linear Networks\" (ReLNs), which are GDLNs specifically designed to imitate ReLU networks, allowing them to derive exact learning dynamics.\nThe key contributions are:\n\nThe introduction of ReLNs as a theoretical framework to analyze ReLU networks, providing tractable dynamics for finite-width networks during feature learning. A demonstration that ReLU networks exhibit an implicit bias towards structured mixed selectivity - where neural units are active across multiple contexts but in structured, reusable ways. \nEvidence that this bias towards mixed selectivity and node reuse is amplified when: 1) more contexts are added to the task\nand 2) additional hidden layers are included in the network\n\nThe authors support their theoretical findings with analytical solutions and empirical demonstrations on several tasks, including an adapted XOR problem and multi-context classification tasks. They show that while ReLU networks aren't biased towards strictly modular or disentangled representations, they do learn structured features that can be somewhat reused across contexts.\nThe work takes a step towards understanding how and why structured representations emerge in ReLU networks during learning, bridging a gap in theoretical understanding between linear networks and modern deep learning architectures.", "questions": "1) Can you comment a bit more on Assumption 2.1 - I know that the mutally diagonalizable structure is quite restrictive, in particular, can you comment on why the tasks you chose follow these assumption(s). Which tasks can you not study, give these assumptions? \n\n2) Can you give a bit more experimental results, especially when training your networks. Maybe a hyperparameter table in the appendix is nice. Can you give more details how you derived hyperparameters when training networks, you for example mention some in Figure 2. I missed if these hps are analytically derived. \n\n3) I find the discussion wrt to compositionality and modularity, you mention \"strictly modular\" in the abstract, a bit unclear. Can you clarify, or even define in the work, what you mean with this - and then contrast this to your findings. I guess the same applies for an up-front (maybe even in the intro) definition of what you mean with this. I would find it easier to follow the paper, having these things more clearly explained.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "Disclaimer: I did not study the proofs of the paper, as I was very late reviewing this paper. I will try to find more time in the coming weeks. \n\nI find the paper clearly writtin and well presented, the authors make an effort presenting the dense results in a clear manner. The authors, to the best of my understanding, extend the theory around GDLNs and allow for a more in-depth study of the training dynamics of ReLU networks. I find the exposition of quite toyish tasks enlightning and make the paper more easy to follow.", "weaknesses": "I can not judge fairly the novely of the paper, especially its relation to Saxe et a., 2022. For me, it would have been helpful to highlight the novelty a bit more.", "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": "Ad5tQ9W8xK", "reviewer": "Reviewer_YprU", "summary": "This paper introduces Rectified Linear Networks (ReLNs), a subset of Gated Deep Linear Networks (GDLNs) designed to capture the training dynamics of finite-dimensional ReLU networks. By drawing an equivalence between ReLU and ReLN, the authors aim to provide theoretical insights into feature learning and structured mixed selectivity in ReLU networks, especially in tasks involving multi-contextual inputs.", "questions": "1. Could the main findings (e.g., mixed selectivity) be directly observed in the ReLU network without using ReLNs?\n2. Does the method extend efficiently to larger datasets (e.g., ReLU networks trained on MNIST), and can a ReLN adequately explain such networks?\n3. Could the authors clarify their definition of \"feature learning\" and what they mean by a pathway \"doing feature learning\" (line 302)?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The introduction of ReLN as a GDLN variant to study ReLU dynamics is innovative and offers a new angle on understanding feature learning in finite ReLU networks.\n2. The paper provides theoretical insights into inductive biases in ReLU networks, particularly regarding structured mixed selectivity and node reuse.", "weaknesses": "1. It remains unclear why ReLU’s training dynamics and mixed selectivity properties cannot be derived directly from ReLU networks, as they’re already nonlinear.\n2. The approach is demonstrated on relatively simple tasks, raising concerns about scalability. For instance, applying ReLNs to realistic datasets like MNIST remains unaddressed.\n3. Terms like \"feature learning\" and \"pathway doing feature learning\" (line 302) lack precise definitions. More clarity is needed to distinguish “feature learning” in this theoretical context.", "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": "Fr1l0GkwDB", "reviewer": "Reviewer_zKcn", "summary": "The paper introduced \"Rectified Linear Network\" (ReLNs) as a subset of Gated Deep Linear Networks (GDLN), and showed that for single-hidden layer network, for a ReLU network, it's possible to find a ReLN that have same loss and outputs at all time-steps (provided that the assumption 2.1 at L159 are satisfied) in a simple synthetic dataset. Using the ReLN as the equivalence of ReLU, the authors provided an analytical solution for the training dynamic for the finite single-hidden layer network with synthetic dataset, and also demonstrated that the predicted loss of ReLN matches with the empirical loss of ReLU networks. In this specific synthetic dataset (which includes hierarchical structure (animals, plants, fish, birds), and multiple contexts), the papers show that the equivalent ReLN network employs an implicit bias for structured mixed selectivity (e.g, one gating pathway can encode multiple context).", "questions": "**Questions**\n1. [L159] Would be helpful to comments on cases for each of the mentioned assumptions in `Assumption 2.1` since these seem to be not very general assumptions and can be easily violated.\n2. [L183] What is the network architecture for the XoR dataset? Is it a 2-layer network? How would the gating structure look like? It would be helpful to explicitly write down the equation or numerical examples instead of vague description of “2 pathways” and “4 pathways”.\n3. [L257] How do these pathways get identified? Are they identified manually based on the prior knowledge about the task? Since identifying gates would be a major bottleneck for the applications of ReLN, it would be helpful to provide more information on how to identify the pathways.\n\n**Suggestions**\n1. In the Contribution section, clearly state the models and dataset in which the theoretical and empirical results are obtained (single-layer network for theoretical results, and simple synthetic dataset with hierarchy and context (maybe would be helpful to give this dataset a name for easy reference?) to provide a clear expectation for the readers.\n2. Spending more space in the main text to explain how to identify the gating structure and pathways from ReLU to ReLN.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "**Originality**\n1. The paper presents a novel framework Rectified Linear Network to analytically investigate the dynamic of finite-width single-hidden layer neural networks via an extension of Gated Deep Linear Network.\n\n**Quality**\n1. The paper provides theoretical proof for the equivalence between ReLN and ReLU for the case of finite-width single-hidden layer network and verified that the predicted loss matches with the empirical loss in a synthetic dataset.\n\n**Clarity**\n1. The paper clearly present the background relevant work, including the formulation of Gated Deep Linear Network, decomposition of eigenmodes, and derived of learning dynamic based on these eigenmodes.\n2. The theoretical proofs are presented clearly with both informal and formal versions, along with a proof sketch in the main paper, and a detailed proof in the appendix, facilitating the reader's understanding\n3. The empirical results (including the dataset and figures) are clearly presented with clear figure and caption, along with descriptions in the main text.\n\n**Significance**\n1. The paper showed that there's an equivalent ReLN for ReLU for a single-hidden layer network, and it's possible to identify the gating structure of this equivalent ReLN via a clustering algorithm.\n2. By investigating the gating structure of the equivalent ReLN network, the paper shows that there's an implicit bias for structured mixed selectivity (e.g, one gating pathway can encode multiple context).", "weaknesses": "While the paper proposed a novel framework ReLN to study learning dynamic and feature learning in finite-width neural network, a significant topic in the machine learning community, the paper has several limitations in both theoretical and empirical results, which I would clarify below:\n\n1. Theoretical results: The main theoretical results of the paper requires very strong assumptions (L159, *Assumption 2.1*) in both (1) input dataset (*The dataset correlation matrices are mutually diagonalizable*) and (2) model training trajectory (*The neural network weights align to the singular vectors of the dataset correlation matrices.*). These are very strong assumptions, and the authors did not offer any analysis on specific scenario in which these assumption holds or not hold. As we see in section 6, assumption (2) is violated in the case of 2-layer hidden network. As the paper only uses a very simple synthetic dataset (one-hot vector input and sparse binary-vector output), it's difficult to tell whether assumption (1) can hold for realistic dataset with complicated distribution, even in the 1-layer hidden network case.\n\n2. Empirical evaluations: The empirical experiments that supports the claim and theoretical results only includes single-layer hidden network on a simple synthetic dataset (the paper did include 2-layer hidden network but this experiment did not align with the theoretical results, due to violation of the theoretical assumption). While it's reasonable to use this simple synthetic dataset as a proof of concept to verify the theoretical results and illustrate the mixed-selectivity phenomenon, it would be helpful and more convincing if the authors can demonstrate that the proposed approach work in a realistic image dataset (MNIST, CIFAR, etc.) with a more realistic and variety of model architectures (multiple layer perceptron, convolutional neural nets, etc.). The inclusion of real dataset and variety of architecture is even more important since the theoretical results require such strong assumptions.\n\n3. Identification of gating structure: The gating structure identification is a central bottleneck of this framework, since gating *g* is treated as inputs, and needed to be identified before training. Identifying a fixed gating structure, or varying gating structure through training, would be one of the major difficulty of the framework. The paper did propose a clustering algorithm to identify the gating structure of ReLN from the representation of ReLU models. However, since the paper only operates with very small synthetic dataset and models, it's unclear whether this clustering algorithms can scale with realistic dataset and larger models. Therefore, this is another reason that it is critical to evaluate this proposed framework on more realistic dataset, instead of only on the simple synthetic dataset.", "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": "qChOWK0WU5", "reviewer": "Reviewer_H7x8", "summary": "This paper shows that, under strong alignment assumptions of the data and network weights, finite feature learning ReLU networks learn modules that can be represented by Gated Deep Linear Networks (GDLNs), for which the exact learning dynamics are analytically tractable. For a synthetic task with hierarchical structure, the training dynamics of 2-layer ReLU nets are shown to exactly match those of a GDLN constructed for the task. Through this equivalence, it is shown that ReLU networks learn mixed-selective representations due to a bias for optimal learning speed.", "questions": "**Critical questions (order: most important to least important):**\n\n1. Your analysis looks very particular to the structure of your synthetic task. Do you think it will be feasible to find unique or representative GDLNs for more complex datasets and architectures?\n\n    a. Are the assumptions more generally broken, even for 2-layer ReLU nets beyond this toy task of hierarchical structure without noise? Atanasov et al. (2021) show that non-whitened data can weaken the silent alignment effect.\n\n    b. ‘Once the common pathway has finished learning there is no correlation between the context features and labels’. But in practice, likely learning is not perfectly separated. Which impact does this have on the generalizability of your results?\n\n2. You claim that the learning dynamics of the ReLN exactly describe those of the ReLU network. Why are the curves in Figure 4 then not exactly matching?\n3. How do your results depend on the number of neurons? You never explicitly mention this.\n4. Your proposed clustering algorithm in Appendix B evaluates the model at different stages of training. Is this not dangerous in cases where the learned function evolves in a more complex way than in your toy task and systematically changes? Then the algorithm tries to cluster outputs from intermediate functions that are rather unrelated to the final learned function. I would like to see whether clear gating mechanisms can be identified when applied to real data.\n\n**Questions out of curiosity:**\n\n1. In Figures 4 and 5, why is not the largest SV learned first?\n2. Where does the bias towards node reuse come from? Can you see this in the learning dynamics equations?\n3. Can your analysis be extended to SGD? Would you expect fundamentally different dynamics?\n4. Can this theory also inform how to maybe slow down learning to learn disentangled representations?\n\n**References:**\n\nAtanasov, Alexander, Blake Bordelon, and Cengiz Pehlevan. \"Neural networks as kernel learners: The silent alignment effect.\" *arXiv:2111.00034* (2021).", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The approach of identifying the linear submodules learned by ReLU networks and using this for tractability of the analysis is interesting. At least on toy datasets, natural modules to learn the structure in the data are exactly recovered by ReLU networks together with the learning speed. This provides mechanistic understanding how the bias toward learning speed explains why 2-layer ReLU networks learn common and context-specific but entangled representations.", "weaknesses": "I cannot strongly recommend acceptance as some key questions remain unaddressed. While this is a nice and tractable toy model for the learning dynamics of 2-layer ReLU networks, generalizability of the findings is questionable. The results only hold for the synthetic and very specific structure considered. The relevance on real data sets or practical architectures remains elusive, as it is not evaluated. The authors show and acknowledge that the silent alignment Assumption 2 already does not hold for 2-hidden-layer ReLU nets. In Figure 6, against the authors’ claim of ‘sharing the same characteristics’, ReLU networks appear to make larger, sharper steps in the loss. What happens in more practical architectures and real datasets, and whether an appropriate and interpretable GDLN can still be identified remains completely unclear. Experiments that show that this approach works on real data would greatly alleviate these concerns. See other critical questions below.\n\n**Typos in lines:** 30, 392, many in 462-463, 473, 476, 477, 503, 527", "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)