Flexible task abstractions emerge in linear networks with fast and bounded units
Authors: Kai Jappe Sandbrink
Jan Philipp Bauer
Alexandra Maria Proca
Andrew M Saxe
Christopher Summerfield
Ali Hummos
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Deep linear networks
Learning dynamics
Cognitive Science
Cognitive control
Task representations
Primary area: neuroscience_and_cognitive_science
TL;DR: We train neural networks in changing environments and show that task abstractions emerge in parameters trained with fast learning rate and heavily regularized. The task abstractions can then support cognitive flexibility.

Abstract: Animals survive in dynamic environments changing at arbitrary timescales, but such data distribution shifts are a challenge to neural networks. To adapt to change, neural systems may change a large number of parameters, which is a slow process involving forgetting past information. In contrast, animals leverage distribution changes to segment their stream of experience into tasks and associate them with internal task abstracts. Animals can then respond flexibly by selecting the appropriate task abstraction. However, how such flexible task abstractions may arise in neural systems remains unknown. Here, we analyze a linear gated network where the weights and gates are jointly optimized via gradient descent, but with neuron-like constraints on the gates including a faster timescale, non-negativity, and bounded activity. We observe that the weights self-organize into modules specialized for tasks or sub-tasks encountered, while the gates layer forms unique representations that switch the appropriate weight modules (task abstractions). We analytically reduce the learning dynamics to an effective eigenspace, revealing a virtuous cycle: fast adapting gates drive weight specialization by protecting previous knowledge, while weight specialization in turn increases the update rate of the gating layer. Task switching in the gating layer accelerates as a function of curriculum block size and task training, mirroring key findings in cognitive neuroscience. We show that the discovered task abstractions support generalization through both task and subtask composition, and we extend our findings to a non-linear network switching between two tasks. Overall, our work offers a theory of cognitive flexibility in animals as arising from joint gradient descent on synaptic and neural gating in a neural network architecture.

Flexible task abstractions emerge in linear networks with fast and bounded units
Kai Sandbrink" Exp. Psychology, Oxford Brain Mind Institute, EPFL
Jan Bauer" ELSC, HebrewU Gatsby Unit, UCL
Alexandra M. Proca" Department of Computing Imperial College London
Andrew M. Saxe Gatsby Unit, UCL
Christopher Summerfield Exp. Psychology, Oxford
Ali Hummos Brain and Cognitive Sciences MIT
Abstract
Animals survive in dynamic environments changing at arbitrary timescales, but such data distribution shifts a challenge neural networks. To adapt change, neural systems may change large number of parameters, which is a slow process involving forgetting past information. In contrast, animals leverage distribution changes to segment their stream of experience into tasks and associate them with internal task abstractions. Animals can then respond flexibly by selecting the appropriate task abstraction. However, how such flexible task abstractions may arise in neural systems remains unknown. Here, we analyze linear gated network where the weights and gates are jointly optimized via gradient descent, but with neuron-like constraints on the gates including faster timescale, nonnegativity, and bounded activity. We observe that the weights self-organize into modules specialized for tasks or sub-tasks encountered, while the gates layer forms unique representations that switch the appropriate weight modules (task abstractions). We analytically reduce the learning dynamics to an effective eigenspace, revealing virtuous cycle: fast adapting gates drive weight specialization by protecting previous knowledge, while weight specialization in turn increases the update rate of the gating layer. Task switching in the gating layer accelerates as function curriculum block size and task training, mirroring key findings in cognitive neuroscience. We show that the discovered task abstractions support generalization through both task and subtask composition, and we extend our findings to a non- linear network switching between two tasks. Overall, our work offers theory of cognitive flexibility in animals as arising from joint gradient descent synaptic and neural gating a neural network architecture.
Introduction
Humans and other animals show a remarkable capacity for flexible and adaptive behavior in the face of changes in the environment. Brains leverage change to discover latent factors underlying their sensory experience [Gershman and Niv, 2010, Yu et al., 2021, Castañón et al., 2021]: they segment the computations be learned into discrete units or 'tasks'. After learning multiple tasks, low-dimensional task representations emerge that are abstract (represent the task invariant to the current input) [Bernardi al., 2020] and compositional [Tafazoli et al., 2024]. The discovery of these
'Equal contribution, randomized order Corresponding author, ahummos@XIT.C edu
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

useful task abstractions relies on the temporal experience of change, and in fact, brains struggle when trained on randomly shuffled interleaved data [Flesch et al., 2018, Beukers et al., 2024].
In contrast, while artificial neural networks have become important models of cognition, they perform well in environments with large, shuffled datasets but struggle with temporally correlated data and distribution shifts, To adapt to changing data distributions (or 'tasks"), neural networks rely on updating their high-dimensional parameter space, even when revisiting previously learned tasks leading to catastrophic forgetting [McCloskey and Cohen, 1989, Hadsell et al., 2020]. One way to limit this forgetting through task abstractions, either provided to the models [Hummos et al., 2024] or discovered from data [Hummos, 2023]. In addition, adapting a model entirely by updating its weights data-intensive due to high dimensionality. Task abstractions simplify this process by allowing updates to a low- dimeensiioonl set of parameters, which can be switched rapidly between known tasks, and recomposed for new ones. However, despite the advantages of task abstractions, simple algorithms for segmenting tasks from a stream of data in neural systems remain an open challenge.
This paper examines a novel setting where task abstractions emerge in a linear gated network model with several neural pathways, each gated by corresponding gating variable. We jointly optimize the weight layer and gating layer through gradient descent, but impose faster timescale, nonnegativity, and bounded activity on the gating layer units, making them conceptually closer to biological neurons. We find two discrete learning regimes for such networks based on hyperparameters, a flexibee learning regime in which knowledge preserved and task structure is integrated flexibly, and a forgetful learning regime in which knowledge is overwritten in each successive task. In the flexible regime, the gating layer units align to represent tasks and subtasks encountered while the weights separate into modules that align with the computations required. Later on, gradient descent dynamics in the gating layer neurons can retrieve or combine existing representations to switch between previous tasks or solve new ones. Such flexible gating-based adaptation offers a parsimonious mechanism for continual learning and compositional generalization [Butz et al., 2019, Hummos, 2023, Qihong Lu et al., 2024, Schug et al., 2024). Our key contributions thus are as follows:
We describe flexible and forgetful modes of task-switching in neural networks and analytically identify the effective dynamics that induce the flexible regime.
The model, to our knowledge, is the first simple neural network model that benefits from data distribution shifts and longer task blocks rather than interleaved training [Flesch et al., 2018, Beukers al., 2024]. We also provide a direct comparison human behavior where task switching accelerates with further task practice [Steyvers et al., 2019].
We generalize our findings to fully cunnecco deep linear networks. We find that differential learning rates and regularization on the second layer weights are necessary and sufficient for earlier layers to form task-relevant modules and later layers to implement a gating- based solution that selects the relevant modules for each task.
We extend our findings to non-linear networks. As limited proof of concept, we embed such layer in non-linear convolutional network learning two-digit classification tasks.
2
Related work
Cognitive flexibility allows brains to adapt behavior in response to change [Miller and Cohen, 2001, Egner, 2023, Sandbrink and Summerfield, 2024]. Neural network models of cognitive flexibility frequently represent knowledge for different tasks in distinct neural populations, or modules, which then need to be additionally gated or combined [Musslick and Cohen, 2021, Yang et al., 2019]. Several models assumed access ground truth task identifiers and used them to provide information about the current task demands to the network [Kirkpatrick et al., 2017, Masse et al., 2018, Yang et al., 2019, Wang and Zhang, 2022, Driscoll et al., 2024, Hummos et al., 2024]. Indeed having access to task identifiers facilitates learning, decreases forgetting, and enables compositional generalization [Yang et al., 2019, Masse et al., 2022, Hummos et al., 2024]. Such works sidestep the problem of discovering these task representations from the data stream.
Other models train modular neural structures end-to-end, such as mixture-of-expprrt [Jacobs et al., 1991, Jordan and Jacobs, 1994, Tsuda et al., 2020], or modular networks [Andreas et al., 2016, Kirsch et al., 2018, Goyal et al., 2019]. A fundamental issue is the 'identification problem' where different
2

---

assignments of experts to tasks do not significantly influence how well the model can fit the data, making identification of useful sub-computations via specialized experts difficult [Geweke, 2007]. Practically, this results in lack of modularity with tasks learned across many experts [Mittal et al., 2020] or expert collapse, where few experts are utilized [Krishnamurthy et al., 2023]. Recent work used a surprise signal to allow temporal experience to adapt learning [Barry and Gerstner, 2022]. Our model proposes simple dynamics that benefit from the temporal structure to assign sub-tasks to modules.
Our work builds on the theoretical study of linear networks which exhibit complex learning dynamics, but are analytically tractable [Saxe et al., 2013, 2019]. Prior work examined how gating alleviates interference [Saxe et al., 2022], but gating was static and provided as data to the network. We generalize this line of work by showing how appropriate gating emerges dynamically. More recently. Shi et al. [2022] analyzed specialization of linear network with multiple paths when tasks are presented without blocking and gates, and Lee et al. [2024] studied the effects of a pretraining period. We consider continual learning with a blocked curriculum. Schug al. [2024] proved that learning a linear number of (connected) task module combinations is sufficient for compositional generalization to an exponential number of module combinations in a modular architecture similar to ours. Instead, we explicitly study the interaction between task learning and gating variable update dynamics.
Approach
3.1 Setting
We formulate a dynamic learning problem consisting of M distinct tasks. At each time step the network is presented with an input and output pair (x(t), y (xx(t)) sampled from the current task m. Tasks are presented in blocks lasting period of time TB before switching to another task sequentially
A
Task a
Task b
Task
Task
(Fig. 1A). Models are never given the task identity m or task boundaries.
B
Specifically, consider multitask teacher-student setup in which each task is defined by a teacher W+m which generates ground truth label t W* with Gaussian .i.dd. input æ at every point in time. We randomly generate the teachers to produce orthogonal responses the same input. While orthogonal tasks simplify theoretical analysis, we generalize nothoonhogonal tasks Appendix A.9.
W
y1
22
1/2
W
Ø3
The open-ended learning setting and the modeling approach. A. Example of the blocked curriculum with two tasks. B. Neural Task Abstraction (NTA) model updates W through gradient descent, but also the gating variables d², leading to task abstractions emerging in the gating layer.
Figure 1:
3.2 Model
We study the ability of linear gated neural networks [Saxe et al., 2013, 2022] to adapt to teachers sequentially. We use a one-layer linear gated network with P student weight matrices WP E Rhe Xdu, together with scalar variables c" E IR which gate a cluster of neurons in the hidden layer (Fig. IB).
The model output y E di reads
P
y = WWP..
(1. NTA)
Since the cº variables will learn to reflect which task currently active, we refer to their activation patterns as task abstractions. We refer to student weight matrix together with its corresponding gating variable as a path.
We refer to this architecture as the Neural Task Abstractions (NTA) model when the following two conditions are met during training: first, we update both the weights WP and the gating variables e via gradient descent, but on regularized loss function L Lunk + Leeg Second, we impose a shorter timescale for the gates Ta than for the weights Tuc i.e. тa < Tw (although this condition becomes unnecessary if the task is sufficiently high- nigh-dimensional, see Appendix A.2).

---

A
0.5
weight charge lldiet Will (no.)
flx
1008
tongernul
0.0
narn [WW (dashed)
gane changs [[dukte
gates
1.0
0.025
0.5
0,000
student-taathe
Eime to c-1.1
0.6
Nignment
0.5
0.3
0.0
10 15
Figure 2: Joint gradient descent on gates and weights enables fast adaptation through gradual specialization. Learning the blocked curriculum from Fig. with 0.03, Tw 1.3. and block length TB = 1.0. x-axis indicates time as multiples of тB. (Black) Flexible NTA model Eq. (1, NTA), (gray) forgetful NTA model with т. Tu and Anonney Anorm 0. Simulation averaged over 10 random seeds with standard eтror indicated. A. Loss of both models over time. B. Gate activity of flexible NTA. C. Student-teacher weight alignment W WPP normalized and averaged over rows (cosine similarity) for each student-teacher pair. D., E. Norm of updates to W and c. Dashed: norm of students correlating with update size of c. F. Time Lank 0.1 for both models over blocks.
The task loss is a mean-squared error Lusk = 1/2Ei Eim ((v y: );)²) where the average is taken components Lrey Amom frorn
a batch of samples. The regularization loss contains Anonacg nonncg weighted by coefficients >nom nonneg- The normalization term bounds gate activity, foor = 1/2 | - 1)?, where we consider = 1, 2. The nonnegativity term favors positive gates fownneg Σラ= ma -c)). Together, these regularizers incentivize the model to function as an approximate mixture model by driving solutions towards any convex combination of students without favoring specialization and reflect constraints of biological neurons (see Appendix B.3 for details).
Assuming small learning rates (gradient flow), flow), this approach implies updates of
T/a/ate - L,
Tw Twa W Vw L.
where and Tw are time constants of the model parameters. We initialize WPP as ii.d. Gaussian with small variance 0² /dmm = 0.01 and c" = 1/2.
Code for model and simulations at: http:://tthubbcom/pprrcaeeeetn task abstraction
Task abstractions emerge through joint gradient descent
We train the model with fast and bounded gates on M = 2 alternating tasks (Fig. IA) and use P P= 2 paths for simplicity (for the P $ M case, see Fig. 3 and Appendix A.3). As baseline, we compare to the same model but without gate regularization and timescales difference.
Both models reach low loss in early blocks, but only flexible NTA starts adapt task switches increasingly fast after several block changes (Fig. 2A.F). Analyzing the model components reveals what underlies this accelerated adaptation (Fig. 2C.D): in early blocks training, zero loss is reached through re- re-aligning both students WP to the active teacher W*m in every block, while the gates c are mostly drifting (Fig. 2B). Reaching low loss is furthermore only achieved towards the end of a block. Later, the weights stabilize to each align with one of the teachers (Fig. 2C,D), and the appropriate student selected via gate changes (Fig. 2B), reducing loss quickly. The rate which gates change correlates with the alignment and magnitude W/|| of the learned weights (Fig. 2C,E). Overall, this points towards transition between two learning regimes: first, learning happens by aligning student weight matrices with the current teacher, which we call forgetful because it overwrites previous weight changes. Later, as the weights specialize, the learned representations WP can be rapidly selected by the gates according to the task at hand, reflecting adaptation that is flexible. Only

---

the model equipped with fast and bounded gates (flexible NTA) is able to enter this flexible regime (Fig, 2A.F).
Next, we verify that the task abstractions in the gating variables are general, in the sense that they support compositional generalization. We consider two settings that begin by training model with three paths on three teachers A, B, and in alternating blocks, and then training on novel conditions. In task composition, the novel conditions are the teachers' additive compositions A+B, A+C, B+C (Fig. 3A), we see that the flexible NTA model trains on these combinations much faster (Fig. 3C). In subtask composition, the
1,0
1.0
flexible forgetful
loss c
0.5
0.5
novel conditions are combinations of the rows of different teachers. i.c. we break the teachers A, B, C into rows and select from these rows to compose new tasks. (Fig. 3B). In the subtask composition case, use a more expressive form of the gates in the model that control each row of the student matrices We find that, in the flexible regime, the model quickly adapts the introduction of compositional tasks, while the forgetful model with regularization removed does not (Fig. 3C,D). For more details and extended analysis, see Appendix A.8.
tasks
tasks
Figure 3: Flexible model generalizes to compositional
tasks. A. Task composition consists of new tasks that sum sets of teachers previously encountered. Subtask composition consists of new tasks that concatenate alternating rows of sets of teachers previously encountered. Loss of models trained generalization to task composition (C.) and subtask composition (D.) for the flexible (black) and forgetful (gray) NTA. 'New tasks' indicates the start of the generalization phase when the task curriculum changed to cycle through the compositional tasks.
WP individually.
We devote the next section to identifying what factors support the flexible regime of learning.
5 Mechanisms of learning flexible task abstractions
We observed in Fig. 2 that simultaneous gradient descent on weights and gates converges to flexible regime capable of rapid adaptation to task changes. But what mechanisms facilitate this solution? We here leverage the linearity of the model to identify the effective dynamics in the SVD space of the teachers, in which we describe the emergence and functioning of the flexible regime.
5.1 Reduction to effective 2D model dynamics
For simplicity, we consider the case with only M P P =3 teachers and students. We take a similar approach to Saxe et al. [2013], and project the student weights into the singular value space of the teachers for each mode a individually, yielding scalar wP = u WP Each pair of components thus reduces a 2D state vector y = c' "w' + c² - E R.2. where we stack wP along the index m and omit in the following for readability (Fig. 4A). A similar projection is possible in terms of the row vectors of both teachers (Appendix A.1.1).
The essential learning dynamics of the system can therefore be described as
d " cd Am 4) Tw
To uT (y ATD y) A V freg.
(2)
where nm describes the output of the currently-active teacher m. In Appendix A.1, we show analytically and through simulations that this reduction is exact when gradients are calculated over many samples.
5

---

B
flexible
gate change
A
effective dynamics
C
(t)
Po
subspace
out of task
v(t)
0
forgetful
trajectories
1.00
301
00
w0 (1)
0
45 0.65 0.0
30 0.5
1.0
Figure 4: Mechanism of gradual task specialization in effective 2D subspace. A. Sketch of the reduced model and dynamic feedback. -uu-of-subspace students gradually align to teacher axes. B. Trajectories of student weight matrices (blue, orange) in the teacher subspace during complete adaptation following context switch from teacher teacher in the flexible regime. Gray stripes indicate associated gate activation. The student weight matrices move little. C. Like (B), but for the forgetful regime. Student weight matrices entirely remap and gates do not turn off. D. Gradient of the task loss on cº as a function of the weight alignment. E. Trajectories in the specialization subspace as a function of gate timescale for values т, = 0.1, 0.18, 0.32, 0.56, 1.00 comparing (color) simulations and (dashed black) analytical predictions from exact solutions under symmetry in the flexible regime. Simulations begin from initial conditions complete specialization and separation wft 5pm. d 6p1 and follow complete adaptation from teacher 1 to teacher 2 over the course of block, reaching Lask 10-2 for all
5.2 Specialization emerges due to self-reinforcing feedback loops
The flexible regime is characterized by students that each attuned to single teacher (Fig. 4B), whereas in the forgetful regime, both students track the active teacher together (Fig. 4C). We can describe this difference by studying the specialization of the students. We define this by considering the difference in how represented the teachers are in the paths: for teacher 1, Wi w - and, for teacher 12 == ? _p=1 Similarly, a hallmark the flexible regime are separated gates. Together, this defines the specialization subspace
(4) The system is in the flexible regime when absolute values of t and are high (approaching 1), and in the forgetful regime when they are low (around O). this section, we study the emergence of flexible regime through self-reinforcing feedback loops, with specialized students and normalizing regularization leading to more separated gates, and separated gates in turn leading to more specialized students. In each subsection, we first describe the effect of the feedback loops the paths individually, before considering the combined effect on specialization. Without loss of generality, we consider cases where the student specializes to teacher p.
tt + W2),
T c-2
5.2.1 Specialized students and regularization encourage fast and separated gates
We first investigate the influence of up on di c². From the gate update in Eq. (2), we get
Te =1 pT =1 + E2 w 22
V
Greg:
(5)
dt
where we decomposed the error == y y into the teacher basis coefficients Em == =m The feedback between students and gates enters here in two terms, as can be seen by expressing uPP *m = | ||] cos (Z (w",, "²"")), where / denotes the angle between two vectors. As observed Fig. 2. both the alignment between students and teachers cos (4 (wp, w²m) and the magnitude of the students | control the gate switching speed.

---

As the vectors wº are formed from the students' singular values. they scale proportionally to the bare matrix entries W for random initialization (Marcenko-Pastiee distribution). Early in learning. the small initialization will therefore attenuate gate changes by prolonging their effective timescale Te/l| ee/// (or equivalently, lower their learning rate).
As we demonstrate in Fig. 4D, these effects apply to both activation and inactivation of the gates, depending on the direction of the current error E *m* z O.
The regularization in the system introduces a feedback loop between ct and c². In practice, the system quickly reaches regime where both gates c" are positive. In this case, the regularization term using the T.rmnorm becomes V Lreg x Ep c² 1, reaching minimum along the line Eg c² 1. In order to minimize the regularization loss, the upscaling of one gate c² past 0.5 will result in the downscaling of the other gate c² and vice versa. We note that this regularization term does not favor specialization by itself since the network can also attain zero loss in the unspecified forgetful solution with, for instance, ct e² 0.5.
The above dynamies mean that differences in student alignment separate the gates, as described by
da To 11 E1 02 E
We therefore see that the differences in gate activation are driven by the difference in specialization in the two components 1j and W2 and corresponding error components E and Eg. Since the error components of opposite sign following context switch, de is maximized when the students de
maximally specialized.
5.2.2 Flexible gates protect learned specialization
We now study the influence of c² not require a change in the weights. When continuing gradient descent on all parameters, however, Eq. (1) will also entail finite update to the wrong student.
di wP. The gates allow for switching mechanism that does
If we Taylor-expand second order, this update reads
d E + dd Tho wp ~ d' + 2 = dt
) dt.
(tttt
(7)
The first summand of the second term reflects the protection that arises from changes in gating d ce = wp?Te e: a task switch to nmm = 1) incurs an error E x 1-1, For specialized, but now incorrect student w" (1, 0) this term becomes de cº w < for the incorrect student. Together with the decreasing error the last term e, this reduces the student update from the leading-order first term cºe. Importantly, this protection effect grows over training as the student's contribution to the error T increases.
Alongside protection, flexible gates also accelerate specialization, as can be seen by considering w in specialization space.
Tw dı
-0 -2)
(8)
This equation shows that the students specialize through two factors: the difference in error between the two components é1 E2, and the difference in gate activation é = c² 22
5.3 Exact solutions to the learning dy namics describe protection and adaptation under symmetry the flexible regime
In this section, we study exact solutions to the learning dynamics in Eq. (6) and Eq. (8) to describe the behavior of the model as it switches between tasks when is already in the flexible regime. To solve the differential equations, we require the condition of symmetry where w condition approximately true for specialized states the flexible regime (see Fig. A.),), and its persistence follows as long é -22 holds or in the limit of strong L regularization.
W2. This
7

---

We use the method presented in Shi et al. [2022] to solve the resulting dynamics of the ratio between die de
the expressions for de and di
Tc de
(E1 E2) C ( E2
Tio dw
which separable differential equation that can be solved up to an integration constant (see Appendix A.7.2). Plugging in initial conditions that correspond to complete specialization in the flexible regime a(0) = w(0) 1, we obtain the exact dynamics of as a function of č over course of block,
1
22) 22).
(10)
Tup
This analytical solution accurately describes adaptation in the flexible regime (Fig. 4E). The relationship highlights the role shorter gate timescale тe in protecting the student's knowledge. Learning that comes from both students specializing towards the current teacher occurs outside of this specialization space and becomes more important for low Te (see Appendix A.7.3).
6 Quantifying the range of the flexible regime across block length, regularization strength, and gate speed
To assess the roles of block length, regularization, and fast gate timescale (inverse gate learning rate) in establishing the flexible regime, we run two grid searches over the gate learning rate/regularization strength and block length each task trained keeping the total time trained constant (such that models trained
1.0
10
block length talta
total alignment (W
block length
6-
4
0.7
on shorter block lengths are trained over more block switches but equal amounts of data). For each set of hyperparameters we compute the total alignment (cosine similarity) between the entire concatenated set single overall measure of specialization in the network weights at the end of learning. We identify the boundaries of the flexible regime where specialization emerges in our model, dependent on block length, gate timescale, and regularization strength (Fig. 5). A priori, the block length dependence is surprising, one might expect additional time spent in a block to be reversed by the equally-lng subsequent block. However, we show in Appendix A.6 that gating breaks this rime-reversal symmetry, and specialization grows with block length 7B for fixed overall learning time
1.25 1.66 2.5 gate leaming rate Twit.
0.05 0.1 0.15 0.2
regularization strength
Figure 5: Model specialization emerges as function of block length, gate learning rate, and regularization strength. The colorbar indicates total alignment (cosine similarity) between all sets students and teachers considered collectively.
of teachers and students
7 Inducing the flexible regime in a deep fully-connected neural network
Our NTA model uses a low-dimensional gating layer that gates computations from student networks. We sought to understand the necessity and role of this structure by considering more general form of the model in deep linear network with no architectural assumptions. Based on the analysis and results so far (Fig. 4D.5), impose regularization and faster learning rates the second layer of 2-layer fully-connected network. Behaviorally, this network also shows the signatures of the flexible regime with adaptation accelerating with each task switch experienced (Fig. A.4A).

---

To quantify specialization and gating behavior, we compute the cosine similarity between each row the first hidden layer and the teachers and use this to sort the network into two students that align to the teachers they match best. We also permute the second layer columns to match the sorting the first layer. We then visualize the specialization of the sorted first hidden layer using the same measures as the original NTA model. We also take the mean of each sorted student's s second hidden layer to be its corresponding gate. Using this analysis, we find emergent gating in the second layer (Fig. A.4B) and specialization in the first (Fig. A.4C). Adaptation to later task switches takes place primarily in the second layer (Fig. A.4E).
Task A
parameter
By visualizing the sorted second hidden layer of the fully- connected network at the last timestep two different task blocks, we indeed observe distinct gating behavior along the diagonal, specialized for each task (Fig. 6 for seed). We compare this to the same fully- cullycconeected network trained without regularization which remains the forgetful learning regime. We include visualizations of the unsorted second hidden layer for fully-connected networks arriving at both the gating and non-gating solutions (Fig. A.5 for ten seeds), as well as the sorted second hidden layer (Fig. A.6 for ten seeds) as supplement. Appendix A.4.2 discusses the potential for multiplicative gates to emerge in fully-connected architectures.
Figure 6: Task-specialized gating emerges in the second layer of a 2- layer network with faster second-layer learning rate and regularization. The sorted second layer weights at the last timestep of two different task blocks (one seed).
8 Flexible remapping of representations in nonlinear networks in two MNIST tasks
We next study whether NTA also works in larger, nonlinear systems. a proof of concept, we investigate whether NTA can help a neural network switch between two versions the MNIST dataset [Deng, 2012]. The first task the conventional MNIST task. The second is permuted version of MNIST where the image of digit is sorted based on its parity according to the function y + Ly/2] +5 x (y%2), where % the modulo operation (see Fig. 7A), We
A
B
MNIST
3
4
7
q
1
S
2
4
5
6
7
9
-
permuted
3
H
7
q
2
4
5
6
7
9
1.0
2.0
Scarrac
1.5
forgethul
0.0
200
400
600
800
1000
Figure 7: Learning flexible neural task abstractions in nonlinear character recognition setting. A. We formulate two tasks, the original and a permuted version of MNIST. B. We embed the NTA system into a larger pretrained convolutional neural network architecture. C. Accuracy reached on the MNIST test set as function of time for both (black) the NTA network and (gray) the original CNN. The two tasks are presented sequentially in blocks for both (blue shading) MNIST and (orange shading) the permuted version. The activation of the two gating units as function of time. We show mean and standard error with 10 seeds.

---

pre-train a convolutional neural network (CNN) on MNIST to learn useful representations, achieving about 90% accuracy on the test set. We then train an NTA system beginning from the final hidden layer representations that feeds into the same sigmoid nonlinearity (see Fig. 7B). We again induce the flexible regime using regularization and fast timescales, and contrast performance with forgetful model (see Appendix B.7). We find that the flexible model learns to recover its original accuracy quickly after the first task switches whereas the forgetful one needs to continuously re-learn the task, evaluated on the MNIST test set (Fig. 7C). The activity in the gating units reflects selective activity (Fig. 7D). To further test the range of NTA, we examine how much these results depend on the orthogonality of the task space by formulating two tasks based real-world groupings of clothing in fashionMNIST [Xiao et al., 2017] that have different amounts of shared structure. We find that rapid task switching occurs in both settings similar speed (Fig. A.15).
9 Relations to multi-task learning in humans
Our model captures several aspects of human learning. Humans update task knowledge in proportion to how well each probable task explains current experience [Castañón al., 2021]. Analogously in our model, weight updates are gated with the corresponding gating variable, whose activity in turn reflects how well the weights behind it capture the target computation.
Humans show faster task switching with more practice on the tasks involved. In our model, we saw that the gates change faster as weights specialize the tasks, which facilitated faster adaptation after block switches. NTA shows a qualitative fit (Fig. 8) to humans trained on alternating tasks [Steyvers al., 2019]. In contrast, forgetful model shows deceleration, possibly due to being far from optimal initialization [Dohare et al., 2024] after task switches.
A
flexible
forgetful
0.0
Trained
lass C
1922
0.55
0.2
3944
Mlacos trined
Emestops affer
Figure 8: Comparing performance after a task switch in humans and NTA model. A. Steyvers et al. [2019] report performance of humans learning two alternating tasks (CC BY-NC-ND 4.0 license). B. After a block switch, loss comparison between the flexible (left) and the forgetful (right) NTA model shows opposite trends with further training on switching speed. Bars are standard error with 10 seeds.
10 Conclusions and future work
This study demonstrates how task abstraction and cognitive flexibility can emerge in neural networks trained through joint gradient descent weights and gating variables. Simple constraints the gating variables induce gradient descent dynamics that lead to the flexible regime. In this regime, the weights self-organize into task-specialized modules, with gating variables facilitating rapid task switching. Analytical reductions revealed virtuous cycle: specialized weights enable faster gating, while flexible gating protects and accelerates weight specialization. This contrasts with the forgetful regime in which knowledge is continually forgotten and re-learned after task switches. The constraints necessary for reaching the flexible regime are appropriate regularization, differential learning rates, and sufficient task block length. These mirror properties of biological neurons and beneficial learning settings identified in cognitive science.
The mechanistic understanding of how task abstraction arises neural systems might bridge artificial neural networks and biological cognition, offering foundation for future exploration of adaptive and compositional generalization in dynamic environments. While this study focuses on simple two-layer networks, the framework is applicable other non-linear architectures such as recurrent networks or Transformer architectures. We see future work providing additional architectures and real-world applications of the framework.
10

---

Author contributions
All authors contributed to manuscript writing and conceptualization of the work. AP conceptualized and implemented the experiments, and contributed to the idea, the theory, and the model simulator. JB conceptualized and developed the theory, developed the model simulator, and contributed to idea and experiments. KS conceptualized and developed the theory, conceived of the idea and the link to cognitive flexibility, and contributed to experiments. AH directed the project, contributed to the idea and the experiments, and provided feedback the theory.
Acknowledgements
We thank Stefano Sarao Mannelli and Pedro A.M. Mediano for thoughtful discussions. We would also like to acknowledge and thank the organizers of the Analytical Connectionism Summer School, which AP, JB, and KS first met. AH is funded by Collaborative Research in Computational Neuroscience award (R01-MH132172) AP is funded by the Imperial College London President's PhD Scholarship. KS is funded by a Cusanuswerk Doctoral Fellowship. JB is supported by the Gatsby Charitable Foundation (GAT3850). This work was supported by a Sir Henry Dale Fellowship from the Wellcome Trust and Royal Society (216386/2/19/Z) AS, and the Sainsbury Wellcome Centre Core Grant from Wellcome (219627/Z/19/Z) and the Gatsby Charitable Foundation (GAT3755).
References
Samuel Gershman and Yacl Niv. Learning latent structure: carving nature at its joints. Current Opinion in Neurobiology, 20(2):551-26,, April 2010. ISSN 0959-4388. doi: 10.
1016|f.comb.20000220000 URL htttps::/www.sieeeeeeeee http///cccdeeett..c/cc/////////e htom///cienne/atiicle/..... $0959438810000309.
Linda Q. Yu, Robert C. Wilson, and Matthew R. Nassar. Adaptive learning is structure learning in time. Neuroscience & Biobehavioral Reviews, 128:270-281, September 2021. ISSN 0149-7634, doi: 10.10166/. 2e0.iii.ne 1.060024. URL https: //www//www///.// sciencedirect. www.rreeetcccccceeeeeeeeeeeeeeee www.sciieeeeeccccccc arriclee/p1/901117442212222
Santiago Herce Castañón, Pedro Cardoso-Leite, Irene Altarelli, C. Shawn Green, Paul Schrater, and Daphne Bavelier. A mixture of generative models strategy helps humans generalize across tasks. bioRriv, page
January 2021. doi: 10. 11.1111/202111111 00..002.........3150000 URL /ttp:////w.oooooeeeaaa////////// /ttp///rlyy/22/////////2/2// 02. 016.643366.... abstract.
http://biorxiv... .rg/ content
Silvia Bernardi, Mareus K. Benna, Mattia Rigotti, Jérôme Munuera, Stefano Fusi, and C. Daniel Salzman. The geometry of abstraction in hippocampus and prefrontal cortex. Cell, 183(4):
954-967.e21, November 2020. ISSN 0092-8674. doi: 00.1016f.eel1122000000333 URL https: //www. ncbi nlm. httpa//////mmmrrrrrceee////ss55/
Sina Tafazoli, Flora M. Bouchacourt, Adel Ardalan, Nikola T. Markov, Motoaki Uchimura, Marcelo G. Mattar, Nathaniel D. Daw, and Timothy J. Buschman. Building compositional tasks with shared neural subspaces. bioRxiv, 2024. doi: 10. 00.1101/2224.01.3.55722222
https: //www. biorxiv. Publisher: Cold Spring Harbor 01 31. 578263.
peapettent/eealy/22222//22 Laboratory _eprint:
f.l.l...722....ll...ff... fdf..
Timo Flesch, Jan Balaguer, Ronald Dekker, Hamed Nili, and Christopher Summerfield. Comparing continual task learning in minds and machines. Proceedings of the National Academy of Sciences, E15(44
October 2018. doi: 10. 1807/// 1800755115. URL https:/ //pp:://www... pnas. org/ /oid/o//fo//// foooffll/100/. 1073/ Academy of Sciences.
1073/pnas..... pnas. 1800755115.
Publisher: Proceedings of the National
Andre O. Beukers, Silvy H. P. Collin, Ross P. Kempner, Nicholas T. Franklin, Samuel J. Gershman, and Kenneth A. Norman. Blocked training facilitates learning of multiple schemas. Communications Psychology, 2(1):1-17, April 2024. ISSN 2731-9121. doi: 1038/ 944271-024-00079..1 URL. https: wwwpw/tt......eeeemm/tt/twt.. nature. cwtp//rricceee///2211220-0000000 Publisher: Nature Publishing Group.

---

Michael McCloskey and Neal J. Cohen. Catastrophic Interference in Connectionist Networks: The Sequential Learning Problem. Gordon H. Bower, editor, Psychology of Learning and Motivation,
volume 24, pages 109-165. Academic Press, January 1989. doi: 10.1006//55779--4211000000000000
URL https:tt///// /www. sciencedirect. htt.aaawww..a.a.eeee..e.aaaaaaaa
Raia Hadsell, Dushyant Rao, Andrei A. Rusu, and Razvan Pascanu. Embracing Change: Continual Learning in Deep Neural Networks. Trends in Cognitive Sciences, 24(12):1028- 1040, December 2020. ISSN 1364-6613. doi: 10.10166j. 202.: 004. URL wtttts:://sttt http:/wwwwcciencdrrretttt
conm/science/articlee// http://ccc.....aaaaa////////aaa1
Ali Hummos, Felipe del Río, Brabeeba Mien Wang. Julio Hurtado, Cristian B. Calderon, and Guangyu Robert Yang. Gradient-based inference of abstract task representations for generalization
in neural networks, 2024. URL https: /uap:://arxiae /ttpiixxiiii///i///aa//aaaaa org/ ht////bb//220000 17356. _eprint: 2407.17356.
Ali Hummos. Thalamus: a brain-inspiree algorithm for baalogically-llussbbee continual learning and disentangled representations. In The Eleventh International Conference on Learning Representations, 2023, URL /oppsr///peereviex wwweee/eeomdddoorrreerreoaaeeeee
Martin V. Butz, David Bilkey, Dania Humaidan, Alistair Knott, and Sebastian Otte. Learning, planning, and control in monolithic neural event inference architecture. Neural Networks, 117:135-144, September 2019. ISSN 0893-6080. doi: 10. 10.10166 https:/ ftttp::/www.scienceee http://www....cceeeeeeee/aaaaaaa
URL
Qihong Lu, Ali Hummos, and Kenneth A Norman. Episodic memory supports the acquisition of structured task representations. bioRxiv, page 2024.05.06.5927444 January 2024. doi: 10. 1101/2024.05.00 05.55594... URL http: /tbiorivbiixii http://wwwa..ooo//oo//////////// 05.06. www.59994....r......
Simon Schug, Seijin Kobayashi, Yassir Akram, Maciej Wolczyk, Alexandra Maria Proca, Johannes Von Oswald, Razvan Pascanu, Joao Sacramento, and Angelika Steger. Discovering modular solutions that generalize compositionally. In The Twelfth International Conference on Learning Representations, 2024. URL ftttt::///peenevieee het//forum/.dd+g90occcck
Mark Steyvers, Guy E. Hawkins, Frini Karayanidis, and Scott D. Brown. large-scale analysis of task switching practice effects across the lifespan. Proceedings of the National Academy of Sciences, 16636) https::/wwat pnas htt/////fff/////// 10111//1.1111//////.111 the National Academy of Sciences.
September 2019. doi: 10. 1073003/p momas.19060 19067881 16. URL
1906788116. Publisher: Proceedings of
E. K. Miller and J. D. Cohen. An integrative theory of prefrontal cortex function. Annual Review of Neuroscience, 24:167-202, 2001. ISSN 0147-006X. doi: 10. 14
Tobias Egner. Principles of cognitive control over task focus and task switching. Nature
Reviews Psychology, 2(11):702-71) November 2023. ISSN 2731-0574. doi: 10.1038/
944159-023-0022444 URL attpp://wwet http://wwttt.e..e/////////// nature. com/ Publisher: Nature Publishing Group.
Kai Sandbrink and Christopher Summerfield. Modelling cognitive flexibility with deep neural networks. Current Opinion in Behavioral Sciences, 57:101361 June 2024. ISSN 2352-1546. doi: 10. (o.10166j..obhha.200 2024. 101361. URLhttpss wwwteeeeeeeeeeeieerreeeeet..... www/sciincc///rrtaaccccoooocccc. /ariieee/t///////ee pii/$$282154424000225
Sebastian Musslick and Jonathan D. Cohen. Rationalizing constraints on the capacity for cognitive control. Trends in Cognitive Sciences, 25(9):757-775, September 2021. ISSN 1879-307X. doi: 10.1016/jjties..000 202..e..22..000000
Guangyu Robert Yang, Madhura Joglekar, H. Francis Song, William T. Newsome, and Xiao-Jing Wang. Task representations in neural networks trained to perform many cognitive tasks. Nature Neuroscience, 22(2):297-306, February 2019. ISSN 1546-1726 doi: 10.100884415993-08--00331 URL https: /wwww/ww.wc/./oo nature. com/ htmmm/rriiceeeee41e9.....11 018 0310-2.
12

---

James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A. Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska. Barwinska, Demis Hassabis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell. Overcoming catastrophic forgetring in neural networks. Proceedings of the National Academy of Sciences, 114(13):521-3322 March 2017. doi: 10. 10.1073/pmas. 16118351 14. URL https: /www. pnas. org/ /t////o/...../////////// 1073/ ppaas. 1611835114. Publisher: Proceedings of the National Academy Sciences.
Nicolas Y. Masse, Gregory D. Grant, and David J. Freedman. Alleviating catastrophic forgetting using context-dependen gating and synaptic stabilization. Proceedings of the National Academy of Sciences, 115(44)
110.//a....888991111 October 2018. doi: 10.1073/pnas. 18038391 15. URL
https://www http://wwwppp........./.aaam pnas. org htg//oi/10... 1073/ www./pnas.....aa..oo. 1803839115. Publisher: Proceedings of the National Academy of Sciences.
Dongkai Wang and Shiliang Zhang. Contextual Instance Decoupling for Robust Multi-Person Pose Estimation. pages 11000-1068, 2022. URL https://openaceess thecvf com/
Instance_ forpppllnnnooo_oooooeeeepprr _obust______ Multii-Person_Poee_HEE Pww.PprroooPPPPeeeaeamiioooorcr.
html.
Laura N. Driscoll, Krishna Shenoy, and David Sussillo. Flexible multitask computation in recurrent
networks utilizes shared dynamical motifs. Nature Neuroscience, 27(7):1349-1333 13633-3333.,.. July 2024.
ISSN 1546-1726. doi:
URI. https: //wp:::///ww..// http::/wwwwwwtwrrwee nature. com/ Publisher: Nature Publishing Group.
articles/s41593-o0 01668-6.
Nicolas Y. Masse, Matthew C. Rosen, Doris Y. Tsao, and David J. Freedman. Flexible cognition in rigid reservoir networks modulated by behavioral context, May 2022. URL https:// www. www.biorxiv.oo org/ httg//ooteettt1../ 11.00..101/2222 05. 09, 491102v2. Scction: New Results.
Pages: 2022.05.09,491110
Robert A. Jacobs, Michael I. Jordan, and Andrew G. Barto. Task Decomposition Through Competition a Modular Connectionist Architecture: The What and Where Vision Tasks. Cognitive Science, 15(2):219-2255 1991. doi: http:///oo.rrr00100020011111111/
URL
http///oi/aaa//aaaee httpaabb/10..... httts:///onlimeliiiii cattrrarrwwllllee. wiley. com www.ooooo/////oocooocoooo//.oooo
_eprint:
Michael I. Jordan and Robert A. Jacobs. Hierarchical Mixtures of Experts and the EM Algorithm. Neural Computation, 6(2):181-214, March 1994. ISSN 0899-7667.
doi: 10. 1162/ncco. 199..6...1999......... URL https: /doii org.10. 11w../eccc//c.o/// 1994.6.99.. 2. 181. _eprint: http::/dtccccaaaiccccctootttootp http://00000...119999...12111.1.
Ben Tsuda, Kay M. Tye, Hava T. Siegelmann, and Terrence J. Sejnowski. A modeling framework for adaptive lifelong learning with transfer and savings through gating in the prefrontal cortex. Proceedings of the National Academy of Sciences, 117(47):29822-90) 228822, November 2020. doi: 10. 20133/pna.2200591111 2009591117. Publisher: Proceedings of the National Academy of Sciences.
URL https: /ttps://ww// pnas /t////////b//////////// 1073//nas..
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 39-48, 2016.
Louis Kirsch, Julius Kunze, and D. Barber. Modular Networks: Learning to Decompose Neural Computation. November 2018. URL httpp://wwrr http://wwwwww..i..c.o../ www.steeaeeseeeeee.eeeeee cschollar org/ 9e980b9ee1033111777c10083
to- hew.riin--oo--eeo----eeeeeeeee aeleurlllee Bawuural-Krscc-HHRe httaa//eeeeaaaaeeaeaaeeeeeeeeee.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and
B. Schölkopf. Recurrent Independent Mechanisms. ArXiv, 2019.
John Geweke. Interpretation and inference in mixture models: Simple MCMC works. Computational Statisties & Data Analysis, 51(7):3529-3533 April 2007. ISSN 0167-9473. doi: 10. 1016/j.cdda.2006.11 1..6.. https:/ www.sccieneeiieee....... httpe//eeeeeeeeeeeeeeceeecctt//t $0167947306604506
13

---

Sarthak Mittal, Alex Lamb, Anirudh Goyal, Vikram Voleti, Murray Shanahan, Guillaume Lajoie, Michael Mozer, and Yoshua Bengio. Learning Combine Top-Down and Bottom- Signals in Recurrent Neural Networks with Attention over Modules. arXiv:2006. 16981 [cs, stat], November 2020. URL wtt.tttpp//axii/ wwtp/rr/////ooooooooo 16981. arXiv: 2606.16981.
Yamuna Krishnamurthy, Chris Watkins, and Thomas Gaertner. Improving Expert Specialization in Mixture of Experts. arXiv preprint arXiv:2302.147700 2023.
Martin Barry and Wulfram Gerstner. Fast Adaptation to Rule Switching using Neuronal Surprise, September 2022. URL https:/ //ww: biorxiv. orggcontent/10..oom 11ttt//100.11//222..0 53.77777vv. Pages: 13.5099911.777 Section: New Results.
Andrew M Saxe, James L McClelland, and Surya Ganguli, Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv: 2013. tex.creationdtt:: tex.modificatinndat: 2022-07-07T18:40:33)
0022-07-07718:40:25
Andrew M. Saxe, James L. McClelland, and Surya Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 116(23):
11537-11546, June 2019. doi: 10. 10.1003/pnas. 1177///2as.18202261111 18202261 16. URL https://wwet pnas org/doi/ 10. 11100077aaas 1820226116. Publisher: Proceedings of the National Academy of Sciences.
Andrew Saxe, Shagun Sodhani, and Sam Jay Lewallen. The Neural Race Reduction: Dynamics of Abstraction in Gated Networks. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 19287- 19309. PMLR, July 2022. URL https: htttpp:://proceeee http:///pooceeeig....e// mlr ppess/v162. stte///axx2aaa..aama html.
Jianghong Shi, Eric Shea-Brown, and Michael A. Buice. Learning dynamics of deep linear networks with multiple pathways. Advances in Neural Information Processing Systems, 35:34064-34076, December 2022. ISSN 1049-5258.
Jin Hwa Lee, Stefano Sarao Mannelli, and Andrew Saxe. Why Do Animals Need Shaping? A Theory of Task Composition and Curriculum Learning. arXiv preprint arXiv:2402. 18361, 2024.
Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, (91 2012. Publisher: IEEE.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Mashion-MNIST: novel image dataset for benchmarking machine learning algorithms,, August 2017. arXiv: cs.LG/1708.07747 [cs.LG].
Shibhansh Dohare, J. Fernando Hernandez-Garvia. Qingfeng Lan, Parash Rahman, A. Rupam Mahmood, and Richard S. Sutton. Loss of plasticity in deep continual learning. Nature, 632 78026):768-774, August 2024. ISSN 1476-4687. doi: 10. 10.1008441586-02244 1586-024-07711-77 URL
https://www, http::://wwwnnntue...m nature. http//aaiceessssss.ea/. 0atp://///55880000000000000000 Group.
Publisher: Nature Publishing
Alexander Atanasov, Blake Bordelon, and C. Pehlevan. Neural Networks as Kernel Learners: The Silent Alignment Effect. ArXiv, October 2021.
URL https:
wtt.ghghhhaaaaaoooaooooooooooooa 3A- BA-The-Silents
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 files/ Systems, 35:6615-6629, 2022. URL https://prroceedingg. wwwponngg..eeurrrr http:://pocceednn.ennee.......c html.
Alexander Atanasov, Blake Bordelon, Sabarish Sainathan, and Cengiz Pehlevan. The onset of variance-limited behavior for networks in the lazy and rich regimes. December 2022. arXiv: 2212.12147 [st tex.creationdate: 2022-12-28TOSS 223112-28T08:22:33 tex.modificationdate: 2022-12. 12- 28T08:25:21.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernel: Convergence and Generalization in Neural Networks. arXiv:1806.07577 [cs, math, stat], February 2020. URL http:// /ttp:///axxvvvve orgeiii.rrr/aai wrw./os////a//ee.oooooo abs//806/110...../.. 0757777776.00..0.11
arXiv: 1806.07572.
14

---

Appendix
Overview
We structure the Appendix as follows:
We first provide additional material on our results in Appendix A. Appendix A.1 derives the reduction to the 2D equivalent model. Appendix A.2 shows that even without differential timescale between gates and weights, high-rank students will learn more slowly compared to gates. Appendix A.3 shows that networks with more paths than teachers will split their representations across paths unless cost associated with representation. Appendix A.4 provides simulations and derivations on how gating behavior emerges in architecture without explicit pathways but with two layers, where one layer emergently takes on the role of gates, and the second layer becomes compartmentalizee. Appendix A.4 contains additional results for the fully-connected network. Appendix A.5 shows how the specialized representation incentivized by the virtuous cycle discussed the main text leads to faster reduction in loss compared to unspecialized solution. Appendix A.6 provides approximate theoretical explanation for the beneficial effect of long blocks towards specialization through symmetry breaking in an effective potential. Appendix A.7 provides approximate closed- form solutions when operating in the flexible regime. Appendix A.8 provides detail on how the model generalizes to new tasks by leveraging existing abstractions. Appendix A.9 shows that the flexible regime largely persists and slowly decays when the orthogonality assumption between teachers is relaxed. Appendix A. shows that the model can adapt few-shot fashion after a block switch, extending the results from the main text where gradients are calculated on many samples.
We then provide additional technical details in Appendix B. In Appendix B. provide notation table. Appendix B.2 lists parameters used for simulations. Appendix B.3 discusses the regularization that we use, in particular why it does not incentivize a flexible over forgetful solution. Appendix B.4 details on how we calculate alignments between teachers and students, both for the per-student and per-neuron gating models. Appendix B.5 discusses how we choose model parameters.
A Additional details on main text
A.1 Derivation of reduction 2D equivalent model
We will here show that the dynamics of the model can be reduced to effective model that acts in 2D space spanned by the singular teacher vectors across both tasks m.
Recall the task loss the mean-squared error
B
Lissk
tm
y
2Bd
B
W
Lb
Σ
c² WPxbl1
2Bd
b
for a batch X (x1) of size B.
Following the approach in Saxe et al. [2013], we assume the input data is whitened, such that the batch average XXX z Ian, and the learning rate 5 - is small (i.e., the gradient flow regime). Then, the batch gradient reads as a differential equation that simplifies as
WP 1
c" (W X EEp cd W wo')x x
(11)
dt
Bd
x (Wam- Ey d' ww
(12)
For each teacher m and singular value decomposition along a mode a (u 8 ,,,) we can project this equation to get
15

---

d
Tar
dt
11
+mT
w am
c
uua
4m
W
*m
Ex
WP va
(13)
6,a,
c
Epr
sy
m,a
The student singular vectors uº: U2 have been shown to align to those of the current teacher u vaa early in learning [Atanasov et al., 2021]. After training on both teachers, the student can therefore be fully described in terms of the coefficients {
in the basis spanned by the a-singular vectors
of both teachers, decoupling from the other singular value dimensions. all singular vectors across two teachers are pairwise orthogonal, these projections form orthogonal basis. The components outside of this projection will have finite error in all context and therefore exponentially decay to [Braun al., 2022].
This reduction allows us to study learning in a simpler and more interpretable model. For the case where M P which consider here for simplicity,
therefore reinterpret each model mmml
a-component as vectors w ( 1,0
s mm=2,a)
w²
p=2 == mm 2,a in R.:.
c'wo' + c?
(14)
$
= c
+
$1
2,a
22a
and redefine the context-dependent target vector y accordingly. The reduced model follows the update equations
d
Tw w" (y - -y)
(15)
dt d To uT (y mm y)
(16)
dt
Notably, both updates depend on the full error term e ;= (y y) with both paths entering into y. The first equation moves the student in the direction of the current total misestimation of the active teacher e. The second equation changes the gating the current path according to the alignment of the path wP to the current vectorial error, reflecting the contribution of the path to the mismatch.
In Fig. A. simulate the models side by side and show that the reduced model matches the dynamics of the full model.
A.1.1 Reduction in terms of teacher row vectors
In the main text and the previous section, we consider reduction that follows from projecting onto the eigenspace of the matrices. However, a similar reduction is possible by considering each row B independently, and considering the row vectors of the two teachers
w km
as basis for that row.
Like the projection in terms of the eigenspace, the out-of nut-oo-projectionn component of the students decays exponentially. We can then consider single row of the teacher-student system to function as a mode a above, with a row of the student path p becoming wp w w wi +2 so that we can write
y c¹ w/ + w
(17)
ct
1.
W2
W2.8
where we aggregate over rows 3. This formulation only requires pairwise orthogonality between rows w.t un = 0 to fully decouple the dynamics of the system, but does not extend elegantly to considering deeper students or low-rank solutions.
16

---

full model
equivalent model
0.5
1005
0.0
eates e
0.5
0.5
0.0
1.00
1.00
SV magnituds
0.75
0.50
0.50
0.25
0.25
0.0
5.0
10.0
12.5
15.0
17.5
20.0
0.0
10.0
12.5
15.0 17.5
20.0
time
time
Figure A.1: Simulation of dynamics of full and reduced model. The equivalent reduced model effectively captures the dynamics of the full model in terms of loss (A1., BI.), gates (Az., B2.), and singular value magnitude (Ag., B3.).
A.2 High-dimensional students learn slower
30
25
20
gate change
15
10
0
0
5
10
15
20
25
30
mintd d())
singular value sum
Σ
52
Figure A.2: High-dimensional students learn slower. Gate change To de c function of teacher dimensionallty/rank (i.c., non-zere singular values). Weight scaling is chosen such that input and output components take unit scale, y4 O(1), Zz O(1).
An intuition one might have for the model dynamics is that the weight matrices comprise of more parameters and therefore may respond more slowly under gradient descent. Here, we discuss the formal conditions under which this indeed is the case.
<Wa, with R R We now choose the scaling y: = (1), ø = O(1) on input and output, which means that entries do not depend on the respective vector dimensionalities. This is natural assumption, for example, if yı label indicators and xy are pixel brightness values of an image. Recall the model loss Lunk 1/2(y^m y)
For simplicity, consider one-path model
Then, the batch-averaged e-gradient reads
17

---

Vusk = Tr [v W WT,
c= dp
Tr W"
W " Wxx W
(0) + Tr WT WW
Tr UUS -USVV TT -T S²
min(dom dm)
(5a8a saa)
Here, we used the Gaussian i..d. initialization of x to take an expectation for large batch size ((x=T)) ) I. the SVD of W = UUVT, orthonormalization of singular vectors U V and the cyclic property of the trace Tr. We also assumed in the fourth row that the student singular vectors have already undergone Silent Alignment [Atanasov al., 2021] match the teachers, as discussed in the main text.
From the last row, observe that the updates to tend to scale with the number of nonzero singular values, i.c. the rank of the teachers.
For the fan-in scaling Wi ~ (0, o²/dml that is compatible with = y 01), we have Sa = (ar) independent of dimensionality (Marcenko- distribution). If teacher and students are initialized according to this scaling, students will respond relatively slower compared to gates as their dimension min( d dow) grows, as the student gradient Eq. (11) or the reduced form Eq. (13) does not involve a sum that scales with dimensionality.
yi
A.3 Representational cost in under-specified model
1.0
1.0
0.5
0.0
0.0
Figure A.3: Redundant paths become inactive when representation is costly. Gating variables like Fig. IB, but with more paths than teacher tasks (P cost on the weights, students that are preferably aligned due to the random initialization specialize to the . teachers, whereas other gates decay to O. B. Without representational cost, the model uses multiple paths for tasks and thus has multiple gates active at the same time for single teacher.
M 2). A. Only under representational
In our initial Eq. (1, NTA), we have introduced a model in which the number of paths of the architecture matches the number of available tasks. What happens if this match is not present? If the under-specified case P < M, the model's 's expressiveness hinders adaptation. It is however not clear what will happen in the over-specified case P M. In absence of any regularization WP, the model will not devote only P' = M < P paths to match the task. Rather, in accordance with the theory by Shi et [2022], the model will in general split its paths over the available tasks. This behavior is due to the absence of a "representational cost" of having multiple paths active the same time. We find that this effect is reduced only when introducing an LeEuuullariitto Aw
the weights
2Phh Eip (W7)2, Xw 0.77. This term additionally penalizes weight magnitude leads to the decay of inactive paths. We show this behavior in Fig. A.3.
A.4 Gating-basee solution emerges in fully-connected network
As described in the main text, we induce the flexible gating regime fully-connected network by applying regularization and a faster learning rate to the second layer and compare to forgetful
18

---

(unregularized) fully-connected network. Details of the sorting procedure used to identify and visualize this gating-based solution described in Appendix B.4.1. We find that the flexible fully- connected network exhibits behavior that is qualitatively similar to the flexible NTA (Fig. A.4). By visualizing the second layer at the end of training on different tasks in the flexible regime, we observe that the network upweights single units in each row (Fig. A.5A, A.6A), which act as gates for the first layer rows. Instead, in the forgetful regime, the network has multiple upweighted units each row and the units do not change behavior across different tasks, exhibiting a lack of task-specificity and gating-like behavior (Fig. A.5B, A.6B).
loss
10
20
30
gates C
0.1
0.0
20
30
Student-teacher allgnment (W--- W
total alignment
(W WW))))
0.5
[AW]
W
W
0.0
10
20
time
Figure A.4: Gating-based solution emerges in fully-connected network with regularized second layer weights and a faster second layer learning rate. Loss during learning. B. The dynamics of the sorted gating variables. C. Alignment between the sorted students in the first layer and the teachers. D. Total alignment between the entire set of teachers and students. E. The norm of the gradient of the first (black) and second (red) hidden layer of the fully-connected network.
19

---

Tosk A
Tosk A
Task 8
Task B
Figure A.5: Regularized, but not non-regularized, fully-connected network specializes single neurons in each row "gates' per task and exhibits specificity based task. Visualization of the unsorted second hidden layer of the flexible (left) and forgetful (right) fully-connected network for single seed.
Task A
Task A
Task 8
Task 8
0.0
0.0
Figure A.6: Second hidden layer of regularized, but not non-regularized, fully-connected network exhibits clear task-specific gating across the diagonals of the matrix. Visualization of sorted second hidden layer of the flexible (lefi) and forgetful (right) fully-connected network averaged over 10 seeds.
A.4.1 Model specialization as function of block size, gate timescale, and regularization strength in fully-connected network
We perform two hyperparameter searches to illustrate the joint effects of block length, second layer learning rate, and regularization strength on the fully-connected network, similar to that we perform on the NTA model in the main text. We run the fully-connecter network on each set of hyperparameters and report the total alignment of sorted teachers and students at the end of training as an overall measure of specialization, fixing all other hyperparameters (see Appendix B.5 for more details). We observe that the same components of block length, fast second layer learning rate, and regularization are important for specialization to emerge the fully-connected network (Fig. A.7), just as in the NTA model.
20

---

1.0
total alignment (W-W) W')
20
block length Teltw
block length ta/tw.
15
15
10
10
5-
5-
2.5
0.7
2
2.66 3.2 4
5.33
0
0.2 0.25 0.3 0.35
layer weights learning rate Tw/Tw,
regularization strength A
Figure A.7: Model specialization emerges as function of block length, second hidden layer learning rate, and regularization strength in fully-connected network. The colorbar indicates total alignment (cosine similarity) between all sets of students and teachers considered collectively.
A.4.2 Possibility of emergence of gating in two-layer network
In this work, have analyzed linear architecture with explicit architectural gating structure,
1/1 "W @c Wa),
(18)
where we have notationally stacked the students W (W p=1..P. into vector, such that W R Pxd I... @ here denotes the Hadamard (element wiime) product.
Prior work has considered deep linear networks [Saxe al., 2019, Atanasov al., 2022, Shi et al., 2022, Braun et al., 2022], which led us to study such fully-connected network in the main text,
du
dred
W
W zj =
(19)
The gated network considers gating as a multiplicative effect on each output unit (or equivalently, input unit j). whereas the deep network invokes additional all-to-all weighted summation. As such, Eq. (19) does not incorporate any modular structure, yet formally resembles Eq. (18). To further analyze how these settings connect, decompose the tasks W Ea U vmm and write the student layer matrices SVDs W = Σa U VOT W E. U a aj overall model Eq. (19) then reads
8 ah
VV The
du min min(dea this) min min(dea (hea da)
81 Σ Σ
Σ
U sa s 0'3 xy
= ( v go) vut x);-
If the minimum of weight dimensions min(dum. duia dmn) exceeds the number of task modes Σm rank (W* is possible choose/learn ( and U0) such that the second layer singular values a(2) effectively take the role of the gates c², whereas the first layer encodes the student task representations. If we put aside the question of learnability and only ask about expressivity, this argument shows that a gating structure can emerge subset of two-layer network. For the fully- connected model we have in the main text, thid 2 doat. giving the network the capacity to learn and remember solutions for both teachers.
A.5
Adaptation speed
In this section, we derive the change model output that is induced by the change in parameters depending their configuration, thereby describing the model's adaptation speed.
21

---

Neural tangent kernel Here, we briefly review the Neural Tangent Kernel (NTK, [Jacot et al., 2020]) which we then use to directly describe the adaptation speed in the output y(t). For vector- valued model y de parameterized by flattened parameter vector 6 (flatten (W ij. o²)) the output evolves as
d di yi Σ dok dl
dys dd
dy dc dok doc
dy dy dc dos dok dys
NTK.;
where used the chain rule and that the parameters update according gradient descent de aoe and have the learning rate 1 for simplicity.
This object can be understood as a matrix operating on the output space NTK (dy/dd") (dy
E Re where the inner product represents the sum across parameters Σx in the expression above. For the reduced model Eq. (2) with 0 flatten(c", (wf) pp......... we readily get
dy: do
dyA duy
w
deg c²,
where 6ig is the Kronecker delta.
We then arrive at
NTK Eee wDwoP
(20)
where we adopt standard matrix notation to imply cec = c In... as being proportional to the identity matrix Idan Rddo
A
Of contribution
Ee' contribution
both tarms
1.0
1.0
1.0
0.8
0.8
0.8
0.6
0.6
0.6
0.4
0.4
0.4
0.2
0.2
0.2
0.0
0.0
0.0
0.0
0.2
0.4
0.5
0.8
1.0
0.0
0.2
0.4
0.5
0.8
1.0
0.0
0.2
0.4
0.5
0.8
1.0
Figure A.8: Specialized students and gates accelerate adaptation. Heatmaps of the dot product (dy) contributions for different terms of the Neural Tangent Kernel (NTK), depending specialization of weight vectors w (blue), w² (orange), of which three pairs corresponding to different degrees of specialization are shown here (pairs are formed by vectors that are symmetric along the diagonal). c² c² are scaled that the sum lies the dashed black line (given by L regularization). A. shows the total contribution of both terms of Eq. (20) combined, B. isolates the contribution from the 2" u term, and C. displays the contribution from the cdc term. Dashed lines indicate possible solutions.
Accelerated adaptation through specialized weights and selective gates To study the accelerated adaptation of the loss Liask = 1/2 (y" am y)² we use the Neural Tangent Kernel the architecture that directly describes the dynamics of the model output. To this end, we study how the model output changes in response block switch nT (1, (0, 1)T, entailing error (-1,1)
22

---

that drives a change in model output. To see how this change reduces the loss, we calculate its alignment with the error term
d Link = ( (
= =e`NTKS
dt
= Σ [ + w Neeeee]
Σ |E|| (c²)² + (wwTT=)
The NTK reveals that the change in output is accelerated through two contributions which we illustrate in Fig. A.8: First, observed that student-teacher alignment which enters |wPTel increases towards the flexible regime. We note that this acceleration however does not require unique student- teacher alignment (such that two students match the same teacher); it is the joint effect of asymmetric gates which further facilitates unique specialization. Second, selective gates accelerate adaptation because sparse vector with the same norm tends to have larger sum-of-squares Σ, (c)) that enters the NTK. These factors coincide in the flexible regime.
A.6 Larger blocks enable faster specialization
1D praxy loss c
A
B
B2
1=0 50, m
1=1.50, m
3.5
0.5
A
effective avera-ae
w(block duration)
9.9
0.0
TT
101
0.5, H² 0.5
- 7 -05
Bg
B
241
-10
NOCK duration
t=3.50
t=5,50,
9.5
0.5
1222
0.0
Ioss
0.0
10-2
-0.5
-05
-0.5
0.5
0.S,
first order gi-1
Figure A.9: Larger blocks enable faster specialization. A. Two trajectories of differing block length (faint: TB, opaque: TBB 2Tp) of two students (blue, orange) in teacher subspace in Fig. 4. B. Student dynamics in an approximate loss landscape in early learning. Subpanels 1-4 are time points in simulation. Background: active context The linear first-order loss does not lead to separation, as a block switch will exactly reverse any changes to specialization. In contrast, the curvature from the second order term enables students accumulate B., but loss is in terms of the specialization variable w. The effective loss over blocks depends on block size 7n: the longer the Tn is, the longer the students will have to fall down the landscape in B.. The first-order term, corresponding to infinitely short blocks, does not prefer specialization.
initial advantage in specialization. C. Like
Here, we calculate how the block length affects the specialization v = w 2 in students, given equal total learning time l. so, we consider periods early in learning where the students have not specialized yet, w ~ 0, and the gates consequentially indifferent, 0. We then consider
23

---

period of the task, i.c. back-and-forth block switches - b, - a that last a total T = 2 We analyze the limit of small block sizes and ask how changes over T: for short blocks, the model time-reversal symmetric, i.c. any change dr ū6) тB during + 6 exactly undone during b + We therefore calculate the second-order effects to tp to analyze the dependence on TB, where we only make assumption on the approximate directions w and that apply early in learning, but leave their scale general.
Second-order derivatives for weight and gates To prepare, we first calculate the second order derivatives of the updates which we will need in what follows. By application of the product rule, we obtain for the weights w
w
ca
dl
d²
d12 w
c
E + c²
dl
di
wPT EE c² NTK
weee EE - d Ex o
whd wo'T
where we use de at (y - y) =
dt - de
NTK within block.
For the gates cº, we get
dt c² wPt E,
d2
dt c² Cde wP NTKE.
We then take the difference of the weight derivatives w = w - w² == c² - c² to get by linearity
d
8 ēe - 0
dt
d²
wt & 2 E of w wTi + wlee
diz w
where we assume that the gates have approximately equal values ē z 0 when no specialization has taken place yet.
Effect on specialization b subsequently:
The sum of second derivatives after having seen switch - b and
d
d²
T (e e² + E" E*)
di2
w
a-
di2
w
b-sa
= =| |||| (E" 8X)
2 2 [ | E,
where that io the component parallel with the error signal, implying TTee || |ee
for the errors e ea = 11/2 (1, - for some constant C' depending on weight magnitude. and small (0). We herein assumed that the errors between blocks only differ by a sign, = E = -e". In particular, we neglect the change c for simplicity. Note that this is a good approximation only if the block size тB is much shorter than the timescales of re While this limits quantitative predictions this approximation for the setting we consider, we expect that it identifies the qualitative mechanism.
24

---

Introducing the period T. which is double the block length TB = 2 (spanning two blocks of length TB), we get
w = 2. 2..7g) (0) + 2. (2 |0 ||=||
e)
=
(21)
where the factor 1/2 in the first line is due to being second order in the Taylor expansion of the update, Setting u (0) = O, this means that the cumulative change of two periods together lasting t 27 (i.e. spanning two pairs of blocks totaling four blocks) is
( 2 periods of T 2 x (2 TB))) 2 x 2 (2 | ||||| e) 22
|||| E) T
(22)
In contrast, doubling the block size T/g = 2TD thus (T" 27), but running for the same amount of time (i.e., two blocks of double the block length TB) gives
wt period of T = (2 2Tn)) 1 x 2: 2 (2 | ||=|| e) (2Tm)² = 4 (2 || ||=|| e) T
which is twice as large as the chort-block version. This explains why larger blocks can lead to faster specialization.
Mechanical analogy The importance of the quadratic term for specialization can be understood through a mechanical analogy for the weights as particles: gradient flow corresponds to the dynamics of particles undergoing an overdamped Newtonian motion in a potential. To this end, we consider the proxy loss potential ((n) that induces the respective first- and second-order time dynamies described by Eq. (21) when considering gradient flow Tue at w -n= The resulting loss potential ((") is polynomial and grows x -w (first-order) and x -|w| 18/2 (second-order), as can be verified by plugging in a solutions wI) x l,, w(() x 12.
d
To first-order, block changes will result in exactly opposite gradients, which will revert any changes from the previous block due to the lack of momentum effects in gradient flow (Fig. A.9B, dotted line). In contrast, the quadratic term in the time dynamics can be understood as resulting from an effective non-linear loss potential, breaking this time-reversal symmetry between blocks (Fig. A.9B, dashed
line). Note that the sum (( + (2) may still be monotonic.
When aggregating this effect over many block changes, it gives rise to an effective loss potential E for the specialization variable w (Fig. A.9C). As the first-order terms cancel out over blocks, the effective loss potential does not have a preferred specialized configuration. When including second-order terms however, the preferred state becomes specialized. Moreover, the speed of this specialization depends on the timescale of the blocks TB: the larger TB, the further the particles increase their advantage down the non-linear loss potential in Fig. A.9B, which manifests a steeper and thereby faster double-well loss potential in Fig. A.9C.
A.7 Exact solutions under the condition of symmetry
A.7.1 Solving the differential equations under the condition of symmetric specialization
The caltiplicativelyccoupee dynamics allow for emergent specialization of the students for one of the teachers. Due to the regularization on c as well as the orthogonality of the teachers, these dynamics are competitive, meaning that the students will, after a number of task switches, each specialize for a different teacher. We saw before that the most rapid adaptation occurs when the system is in this state where each teacher is matched by exactly one student. We formalize this state more generally under a condition of symmetric specialization, in which w = w = w w wğ w/ = W22 In this section, we present exact solutions to the learning dynamics that occur in this system under this condition. We assume without loss of generality that path in the NTA student is aligning with teacher We will see that these solutions closely match the adaptation dynamics of the full NTA model.
du = こ9 and d = -762- The learning updates of the two specialization components are di Therefore, the symmetric specialization condition is inherently preserved when E = -62. We calculate the accuracy of both relationships in Fig. A.10.
de
25

---

A
B
1.00
1.0
1.00
0.8
0.56
0.90
0.6
0
0.85
0.80
0.18
0.75
0.0
0.70
0,10
0.0
0.2
0.4
0.6
0.8
1.0
0.70 0.75
0.80
0.85
0.90
0.95
1.00
W1
Figure A.10: Accuracy of assumptions used in calculating the exact solutions, only one of which required. A. Accuracy of the assumption = E1 -2.. B. Accuracy the assumption
wi w w W2-
A.7.2 Learning occurs dominantly within but also outside of specialization subspace over the course of single block
The differential equation Eq. (9) then results from dividing the two update equations for d and d can then solve this differential equation to obtain the relationship
de dw
Te 22
= w
(23)
where C' is integration constant, which determine by plugging in the initial conditions č e 1 to represent the theoretical ideal for the flexible regime.
A.7.3 Learning the orthogonal component
As highlighted by the fact that the analytical solution is not traversed in full in Fig. 4E, learning also occurs outside of the specialization subspace. This learning be characterized by co-specialization which is characteristic of the forgetful regime
(
wh)) + (wf w
(24)
Fig. A.11 shows learning along both components vp and w. the two error components é and $2, as well as the separation of the gates over the course of single block beginning from the same initial conditions used in the differential equation.
26

---

100
0.75
0.54
0.5
0.0
0.44
0.44
-025
0.0
0.0
250
550
750
time
1000
1250
1500
2000
250
50c
150
time
1860 1250
250
500
750
1006
3250
2500
1750 2000
D
1.0
0.0-
0.5-
0.50
0.0m
0.30
0.4-
0.2
0.8
230
500
750
1750 time
1200 1250 1000
2000
0.0
0.10
Figure A.11: Learning occurs within and outside of the specialization subspace. First component of the error following task switch from task A to task B for different values of T.- B. Second component of the error across the same timeframe. C. Adaptation of weight matrices in the specialization space. D. Orthogonal component of learning that measures adaptation of both teachers for the current task. E. Gate change the specialization subspace.
A.8
Gated model generalizes to perform task and subtask composition
bock
toache
Figure A.12: Gated model generalizes to compositional tasks. A. Task composition consists of new tasks that sum sets of teachers previously encountered. B. Subtask composition consists of new tasks that concatenate alternating rows of sets of teachers previously encountered. Loss (C..D.) gating activity EE,E., and student-teacher alignment (G.,H.) of models generalization to task composition (top) and subtask composition (bottom).
As stated in the main text, we consider two settings to evaluate whether the gate layer can recombine previous knowledge for compositional generalization. We first train the NTA model with three paths on three teachers A, B, and C individually, and then change the network on task composition (Fig. A.1) or subtask composition (Fig. A. 12B). Task composition proceeds the same standard setup.
In subtask composition, to allow our model the possibility to compose not only tasks (i.e., gate the entire student matrix WP), but also subtasks (individual rows of W²), we increase the expressiveness of our gating layer by using an independent gate for each neuron (or row of WP) in the student hidden layer and allow gradient descent to update these gates individually. We call this the ner-ne gating version of our model.
27

---

In principle, the per-neuron gating NTA has Pdout independent paths modulated by gates. Thus, in order to study whether specialization and gating occurs for each teacher, we sort the Pd paths into P paths of size daut- We do this by computing the cosine similarity between each row in the first layer W and the teachers W We then sort the rows of the first layer to align with the rows of the rows of the teachers that they best match, such that we identify their respective students to visualize student-teacher alignment (Fig. A. 12H). Additionally, we permute the gating layer c to match this sorting. We take the mean of the sorted gates for each student visualize the task-specific gating (Fig. A.12F).
We find that both the per-student and per-neuron gating NTA models can solve task and subtask generalization tasks, respectively, and maintain their learned specialization after transitioning to compositional settings (Fig. A.12C-H), We additionally observe that the gating variables learn to appropriately match the latent structure of the generalization tasks by turning off the non-contributing gate and evenly weighting the two "on" gates. We again observe the rapid adaptation of the flexible NTA to compositional tasks compared to the forgetful regime.
A.9 Non-orthogonal teachers
A
B
adaptation speed (loss)
1.00
WP d0 WP -
20
0.75
0.50
10
0.25
0.00
0.0
0.2
0.4
0.6
0.8
1.0
teacher correlation Wm Wm
Figure A.13: Robustness to relaxing orthogonality between teachers. A. Illustration of changing teacher cosine similarity. B. Adaptation speed as measured by the loss after a block switch (black) and student specialization (gray), both as a function of the teacher similarity. represents the orthogonal case studied in the main text.
28

---

task
teacher
B
1.0
0.5
0.0
10
20
30
40
50
c
1.0
gates c
0.0
alignment (wm Wel)
D
student-teacher
1.0
0.5
0.0
time
Figure A.14: Flexible NTA successfully specializes underlying teachers even when trained on non-orthogonal tasks. A. Tasks are created by adding different pairs of teachers, such that each task is non-orthogonal every other task. Loss during learning. C. Gating variables learn to appropriately match the latent structure of the tasks. D. Students learn to specialize to teacher components, despite the orohoogggyyy of the tasks.
In the main text, we have worked with the assumption that different tasks are approximately orthogonal to permit our thcoretical analysis. This assumption holds for mudomml-ggnnrrae the input dimension is high. In simulations, corresponding rows of teachers to be orthogonal, w w O. Still, it is unclear what happens when the teachers not even approximately orthogonal. We investigate this question empirically in Fig. A.13 and find that specialization decays gracefully as the orthogonality assumption is relaxed.
teachers when implemented this condition by constructing the
We also design set of three non-orthogonal tasks using three orthogonal teachers, where each task is created by adding different pairs of the teachers (Fig. A. 14A). Thus, every task some similarity (and is non-orthogonal) with every other task. We find that the flexible NTA can successfully solve these tasks and identify their underlying latent teacher structure, learning to specialize and gate all teacher components which comprise the overall set of tasks (Fig, A.14B-D).
29

---

A.9.1 Experiments on fNSS dataset
Orthogonal (upper-to-lower)
Non-orthogonal (warm weather)
1.0
eccuracy
fashionMNIST
0.5
nexible tongettul
0.5
nexible forgetful
time
time C
1.0
0.5
0
0.0
time
time
Figure A.15: NTA quickly adapts across fashionMNIST for (left) an orthogonal sorting based on upper-to-lower items of clothing and (right) correlated sorting for warm-to-cold weather clothing. The panels show (top) accuracy on the test set and (bottom) activity of the gates. We show mean and standard error with seeds.
We explicitly compare the network's performance two different versions of fashionMNIST based on tasks that might appear real-world setting. The original fashionMNIST dataset has items sorted roughly by order of commonality, with the label 0 being assigned to T-shirts, and the label 9 being assigned ankle boots. We generate two different permutations of these labels representing other real-world sorting of the items that have different amounts of shared structure with the original. The close-to-ortee -orthhoo-ath ordering sorts the clothing from upper to lower body, and orders the labels 0, 2, 4, 6, 8, 3, 5, 9. The ordering with more shared structure represents warm to cold weather clothing, and orders the labels to 1, 5, 3. 7, 2, 4, 8, The results show that the stereotypical NTA-like task switching behavior and specialization emerges for both settings at similar speed despite baseline performance being higher on the task with shared structure (Fig. A.15).
A.10 Few-shot adaptation in the low sample rate regime
10-
loss
0.0
Figure A.16: Few-shot adaptation after block switches. Like Fig. 2A in the main text, but with coarsely discretized time to examine the adaptation after a single sample. As this drastically reduces signal-to-noise ratio, we average over 100 samples. Markers indicate a single step of gradient descent on one sample.
In the main text, we have considered the case of large batch size (or equivalently, small time discretization) that allows taking sample average when going towards the theoretical, equivalent model. This averaging reduces noise in the gradient signal stemming from random samples, so that it is unclear whether learning is still possible when sample rate is low. As theoretical analysis is challenging for this case, we investigate this empirically in Fig. A.16 and find that the qualitative phenomenon is preserved even if only single sample B is used for every gradient update.
30

---

B Technical details
B.1 Notation
Table 1: Overview of notation used throughout the paper.
Symbol
Description
E Rdd, d= I...B
input sample of batch
y E Idu, b = 1...B
model output
y nm Rdv,, I...M or a.
target label in task prediction error
e y² - y Rd
E R, p p= I...P
gates for each pathway p
Wr Rdaxdd
student weights for each pathway
Wm Rduxde, b,
teacher weights for each task m
w" wge R
2D vector for reduced model (for each teacher singular value singular value decomposition of weights
Wi Σe mini minldaa.da) UiaSal V.
700 1m To 111
parameter time scale (inverse learning rate)
TB
block length
C Lunk
Lrey
loss
w p=1
W0 m==1
specialization for teacher 1
m=
182 2-2
1=1
specialization teacher
17 1 [
(72)
overall specialization
separation of gates
1/
((
p=1
p=1
+ (wier w 1)/2
unspecialized learning
B.2 Hyperparameters
We perform the gradient calculations and the simulation of gradient flow using the JAX framework and make our implementation publicly available at
aastaaasta
Our hyperparameters for all experimental settings are listed in the table below. We use L for most experiments, except for generalization subtask composition and the flexible fully-connected network experiments where we use Luom-LL2 (see Appendix B.3 for discussion on this choice). For the cases where we induce the forgetful regime as an experimental control, we use the same hyperparameters, set regularization to 0 (Amomeey Anam...1, Anom.122 0), and may or may not adjust the learning rate the gating layer. Differences hyperarameters from the main model and control denoted in the tables below "main/ control.'
31

---

Table Hyperparameters.
Task specialization (Fig. 2)
Tusk composition
Subtask composittion (Fig. 3.A.12)
Reduced model (Fig. 4)
Fully-comnected network
MNISTT
Hyperparumeter
(Fig. 3,A.12)
(Fig. 6.A.H.A...A.6))
(Fig. 7)
P
2
2
da
20
20
20
20
64
20
doat
10
2
10
10
0.091/0 /
0.5 /
0.023 0
010991 /
0.2/0
0.5/0
Awerdd
0.456/0
1.25/0
0.455
0.25/0
Aman
0.011 0
0
0.1/0
1.3
0.2
0.2
5
0,06
10
te
0.03
0,005
0.7
batch size
200
200
200
200
200
100
seeds
10
10
10
number of blocks n
20
30
30
17
30
10
0.001
0.001
0.001
Table 3: Hyperparameters II.
toporpecaste
bepopersscko
20
30
20
4.3
ILS
325
3:11
0
0.5
316
1L505
borch dos
305
300
300
201
10
Table 4: Hyperparameters III.
Hyporpeamnter
Full nolunoa (Fig A.I)
Show bighhd stadenss (Fig A.2g
Rodumbni paths (Fig A.5)
Fow-doot sdppoationn (Fig A.15)
GehisMNNST (Fig 1
2
10/2
10
0.091
0.091
0.104 70.545
0.091
0,5/0
0,455
0,455
0.908 (.772
0
0
0
0,555
0,25/0
0.5
1.3
0.03/ 0100
0.1
0.01
0.005/ / 10
buch
100
ouber Books
20
20
0,001
0.001
0.001
0.02
0.001
B.3 Regularization
We use combined regularizer that is motivated by biological constraints our model parameters. The regularizers alleviate the underspecification of the solution space of our linear model and facilitate symmetry breaking to allow the model to specialize different components, while not forcing specialization (Fig. B. Here, we detail the effect of these regularizing terms. To this end, recall the definition of the reduced model, Eq. (14)
y = c' u + 2
Nonnegative neural activity The gating variables steer the model output multiplicatively. Biologically, such an interaction is mediated by firing rate, inherently positive variable. Computationally. this has implications on the solution space: with random initialization, almost all configurations of w and w² form basis of IR². By definition, this means that there will always be two coefficients c², ce R that will yield the correct solution. Nonnegativity constrains this set to lie in the positive quadrant of 2D space. In particular,
Cuwaneg
m0 -cd)
32

---

Alleviating invariance via competition invariant under
Even in the desired specialized configuration, the model is
c², W + ac", W
for any scalar a. We hence bound the norm of the vector c (2))
We consider regularizers based the L and L norm
mmom L1
1/2(l//|1 1)2
norm-LL 1/2(||||2 1)²
Crom-LL encourages sparsity in the gates which is beneficial when there are few active gating variables in the NTA model (with single teacher active in each task). However, if consider cases with potentially many active gating variables, as in the per-neuron gating NTA (see Appendix A.8) or deep fully-connected network (or many teachers active at once), favoring sparsity restricts expressivity of the model. In these cases, we instead use Lmurm-L2- practice, both regularizers facilitate specialization robustly across many settings.
Applying nonnegativity and norm regularization together has the effect of inducing competition between gating variables. While there is solution where gates are equal in magnitude (as shown in Fig. B.1). deviating from this solution while minimizing regularization loss will cause one gate to increase in magnitude and the other to decrease. Thus, this competitive effect facilitates symmetry- breaking in the gates.
Finally, we note that although these regularizers facilitate symmetry breaking through competition, they simultaneously allow for compositionality such that multiple gates can active at once. We show this in several experiments, namely our studies of nonorthogonal tasks (Appendix A.9), compositional generalization (Appendix A,8), and fully-connected networks (Appendix A.4).
A
norneg
B
rorneg
c2
83
c
Learmil
L
specialized unspecialzed
Figure B.1: The effect of regularization on gating variables. Regularization encourages competition between gates while preventing degeneracy of solutions. Importantly, regularization does force gating variables be specialized, as illustrated by the red x. This holds for regularizers we consider, A. L.arm-11 and B. L2.m.11.
B.4 Description of metrics used across experiments
Student-feacher alignment We compute a metric of alignment each student and each teacher to determine whether students are specializing and, if so, to which teachers they specialize. We do this by computing the similarity between each student WP and teacher *m More specifically, we take the mean of the cosine similarity between student and teacher vectors. We then sort each student and its gate d to the teacher has the highest cosine similarity with.
Total alignment We compute a metric total alignment of the network students and teachers to evaluate overall specialization. After computing student-teacher alignment and sorting each student to respective teacher, we concatenate all students and teachers and compute the overall cosine similarity between the set of students and teachers.
33

---

B.4.1 Description of sorting performed in per-neuron NTA and deep fully-connected network
In the cases of more expressive models, such as the per-neuron NTA and fully-connected network, there are Pdou and dou Pdow respective independent paths modulated by gates. Thus, in order to study whether specialization and gating occurs for each teacher, we sort these into P paths. To do this, we compute the cosine similarity between each row in the first layer W and each row in teachers W* We then sort the rows of the first layer to align with the rows of the teachers that they best match.
Additionally, we permute the second layer to match this sorting. In the per-neuron NTA, this corresponds scalar gates that are multiplied to each row. In the fully-connected network, this corresponds to the columns of the second layer. Finally, we take the mean of the sorted gates (set of doa columns for the fully-comnected network) for each student to visualize seacher-specife gating.
B.5 Hyperparameter search
NTA For the two hyperparameter searches we perform, we run the NTA model on each set of hyperparameters and report the total alignment of concatenated teachers and students the end of training as an overall measure of specialization. We fix all other hyperparameters.
When varying gate learning rate and block length, we fix the regularization strength to Xnonney 0.5. Amom--L1 = 1.25. When varying regularization strength, we fix gate learning rate Tw/ Te = 20. The regularization strength À is multipled separately for each type of regularizer such that Xacaneg 52/3 and A = 25A/6.
Fully-connected network We also perform two hyperparameter searches the fully-connected network. We run the fully-connected network on each set of hyperparameters and report the total alignment of sorted teachers and students at the end of training as overall measure of specialization. We fix all other hyperparameters.
When varying second layer learning rate and block length, we fix the regularization strength to Anonneg 0.23, Anomm-1.2 = 0.11. When varying regularization strength, we gate learning rate Twra Twa) = 4. The regularization strength À is multiplied separately for each type of regularizer such that Aponnee = 10A/11 and Anom--L2 = 51/11.
B.6 Flexible fully-connected network
We randomly generate two orthogonal teachers Wm E Rda We initialize our fully-comnected networks to have two weight layers, W Xdu and (222 E X2dmn We use faster learning rate and regularize the second layer during training. We treat each weight W as gate that enters into the regularization terms described in Appendix B.3. Student-teacher alignment, total alignment, and gate sorting is then performed as described above.
B.7 MNIST
The original convolutional network features a single convolutional layer with three feature maps and kernel size four, followed by MaxPool layer of kernel size 2, a ReLU nonlinearity, a and fully-connected sigmoid, ReLU and log softmax layers of size 512, 64, and 10 respectively. The networks are trained using cross-entropy loss with one-hot encoding for labels. The NTA portion is then trained beginning from the hidden layer representations of the final hidden layer with 64 units. hyperparameters for the NTA portion are given in Table 2 for MNIST and Table 4 for MNSS..
34

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: [Yes] We address each of the abstract's claims throughout our manuscript and enumerate our contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We end with section titled conclusions and limitations and mention two broad limitations.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
35

---

Answer: [Yes]
Justification: [Yes] We include derivations for analytically important quantities rather than theorems and proofs. Assumptions in model structure are stated explicitely.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We offer a full description in main text and supplementary material, we also provide repository containing the code.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
36

---

Question: Does the paper provide open access to the data and code, with sufficient instructions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: public code repository is included. Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips) cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nwpp...c////////o https/cccpullcccuuuccceooo http:/ubbbi..uuuiied/eddeeeubaia /Cooeeuboeeeeeeeeeee wwwiissssnnoiioiieiooeioooiiiiie onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: All the relevant sections give details of optimization paradigms used and learning rates. The supplementary contains further experimental details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Where ever appropriate we shaded most lines with SEM bands. Error bars as well when specific points were of interest.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
37

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [NA]
Justification: Code can generally run on any machine type, and was tested on several machine types. Our simulations are rather small and can run in 30 seconds.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: / http//nerrrrrreeee/eee https///cppppcc//.ccccccceee.ee.
Answer: [Yes]
Justification:
We reviewed the code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consideration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper presents work whose goal is to advance the field of cognitive science, with possible insights for machine learning.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
38

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e-g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The algorithm might shed light into cognition but we do not foresee any possible misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: figure uses a portion of figure from another work that is under the license "Creative Commons
License 4.0 (CC BY-NC-
ND)". This indicated in the caption.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
39

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new assets
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: No research subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
40

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
41


Reviews:
[{"review_id": "i4vpWGZDkm", "reviewer": "Reviewer_9Ed7", "summary": "This paper presents a method aimed at learning generalizable task abstractions by constraining neuron dynamics in artificial neural networks. It takes inspiration from biological neurons by constraining artificial neurons to non-negativity, forcing a faster timescale, and regularizing. \n\nThe paper goes through a problem setup, then investigates various aspects of the approach. First, it argues that task specialization emerges through joint gradient descent, showing generalization results on compositional tasks. Next, it gives extensive theoretical analysis of the mechanisms of this task adaptation. It then briefly shows evidence for emergence of specialization in the model and visualizes weights, then finally discusses applications.", "questions": "None", "limitations": "Not really", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "### Originality\n- Applying these constraints to neurons does not seem to be particularly novel. However, I have not seen this specific combination or analysis from the perspective of task adaptation.\n\n### Quality\n- Experiments are extensive. Each claim in the paper is backed up by lots of relevant evidence.\n- Motivation is interesting at a high level\n- Investigating these general, simple changes as drivers of flexible task abstractions is creative and exciting. The authors have clearly thought about various aspects of this.", "weaknesses": "### Quality\n- Though the motivation is interesting, it is confusing because the mechanisms appear to be inspired by catastrophic forgetting but then are applied to the related but distinct problem of generalizable task adaptations\n- The central claim/benefit of the paper, that this architectural adaptation might make neural networks more brain-like, is barely supported; the authors even admit to this. However, that seems to be one of the main value adds of the paper\n\n### Clarity\n- Paper setup is very confusing. It's okay to have a nontraditional structure, but it ends up being a prose-heavy laundry list of various properties of the method and results. \n\n### Significance\n- A proven (at least, reliable) model of some biological brain phenomenon is valuable. However, this paper has only shown vague similarity, not much comparative predictive power of this approach vs an unconstrained approach on neural data\n- A machine learning result can also have impact, but the paper doesn't clearly argue for that. It does it show that the representations can generalize across two training settings in MNIST, but this isn't enough to make a real ML argument.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your helpful comments. We appreciate that you found the conducted experiments relevant to our claims.\n\n> Investigating these general, simple changes as drivers of flexible task abstractions is creative and exciting.\n\nWe are glad that this contrast came across well: Despite the simplicity of the proposed mechanism, the effect on inducing flexible task abstractions is remarkable.\n\n> Paper setup is very confusing. It's okay to have a nontraditional structure, but it ends up being a prose-heavy laundry list of various properties of the method and results.\n\nThank you for this valuable feedback on the organization of our manuscript. We agree that it failed to guide the reader through. \n\nOn a high level, our paper has this structure: \n\nHypothesis: Are neuronal gates updated through gradient descent sufficient to separate tasks into dedicated representations? First, we answer through simulations whether and how implicit context switches are learned (Fig. 2). Second, we analytically address why task representations arise (Fig. 4). Third, we apply the model to cognitive science (Fig. 7). The remaining figures are not essential to this narrative, but function as controls for the necessary conditions for flexible gating. \n\nTo address this point, we now strictly keep to this structure:\n\n1\\. We reorganized the order of the mathematical analysis and grouped subsections therein. We now first precisely define flexible adaptation, and then discuss the main driver for its emergence.\n\n2\\. We moved the *Related works* section back to the main text to better state our contribution.\n\n3\\. We moved figure 3 showing how the learned abstractions generalize to the supplementary, as it does not directly support the main narrative. We clearly motivate the remaining figures as controls. \n\n##### **Contribution**\n\nThe reviewer raises concerns about the contribution of our work.\n\nWe aim to contribute to theoretical cognitive science by developing a minimal neural network model of cognitive flexibility. We build on a line of recent theoretical NeurIPS papers studying the learning dynamics in linear networks and relating them to cognition, starting with Saxe et al. (2013, *Exact solutions to the nonlinear dynamics of learning*). Later studies examined how gating alleviates interference, but their gating was static and handed to the network (Saxe et al., 2022, *Dynamics of abstraction in gated networks*). We generalize this line of work with an analytically tractable model of how appropriate gating emerges *dynamically*.\n\nOur model contributes to cognitive science:\n\n1\\. We offer, to our knowledge, the first neural network model that benefits from data distribution shifts and struggles in the shuffled data regime, similar to humans. \n\n2\\. We provide a direct comparison to humans where their task switching behavior accelerates as they practice the tasks involved. This came from both our simulations and our theoretical analysis, and we provide a mechanistic explanation for the phenomenon. \n\n3\\. Not only does the model infer tasks and retrieve the suitable task abstraction, but similar to humans, further learning or credit assignment is gated by the inferred task. I.e. only the parameters for the inferred task are changed. \n\nWe believe that such a model can make additional behavioral predictions and provide a neural basis for explaining additional experimental findings, such as Heald et al’s influential study on sensorimotor repertoires (2021, *Nature*).\n\n> A proven (at least, reliable) model of some biological brain phenomenon is valuable. However, this paper has only shown vague similarity, not much comparative predictive power of this approach vs an unconstrained approach on neural data\n\nWe aim to present a minimal hypothesis of how task abstractions in (biological) neural networks might emerge during training.\n\nThe reviewer highlights the importance of connecting directly to neural data. Unfortunately, almost no full recordings from animals during the training phase on multiple tasks exists. Recordings are mostly from *well-trained* animals switching between the tasks (discussed by Bouchacourt et al (2022), *Fast rule switching and slow rule updating*). Neuroscientists point to a technical difficulty where animals take weeks to months to learn two opposing tasks while recording electrodes slip and drift within days. \n\nStill, technical improvements especially with calcium field imaging might make recordings during learning possible soon: In addition to behavioral signatures, our model predicts gradual differentiation of both task-specific activity vectors in neural space modulated by context that arise over the course of learning. We expect this simple and interpretable signatures to be robustly observed across tasks and species, an advantage over purely Bayesian models and more expressive machine learning architectures.\n\n> A machine learning result can also have impact, but the paper doesn't clearly argue for that.\n\nThe reviewer identified that the simulated model might be extensible to an ML method for larger non-linear datasets. We view constructing an ML system that incorporates these principles to improve machine learning results as a clearly separate project, which we are considering taking on in the future after having laid the theoretical groundwork with this paper.\n\n> The mechanisms appear to be inspired by catastrophic forgetting but then are applied to the related but distinct problem of generalizable task adaptations.\n\nThe reviewer highlights that results on the generalizability of our learned task abstractions belong to a different problem and need comparison to different controls such as few-shot learning or meta-learning models. We moved these results to the appendix and clarified they only serve to show that the abstractions are functional and can be recomposed through gradient descent.\n\nWe will clarify our contributions in the introduction section in the final revision.", "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": "KNqobxmKzO", "reviewer": "Reviewer_pBhB", "summary": "This paper looks at the problem of having an agent learn a series of tasks sequentially by receiving supervised data for each task. Training a NN on this has the issue of catastrophic forgetting, and learn best with shuffled data. Their model attempts to be more similar to humans, who do best with task data presented nonshuffled. The idea is to have subnetworks which can learn each task, and then learn a high level router to the subnetworks. The router uses a gating variable trained with regularization to encourage only one subnetwork to be active at a time. As a result, when trained on the nonshuffled data, the model can pick up on the different tasks and assign them to the subnetworks. The authors show some experiments where the model does just that, with some analysis of the dynamics. They also provide some theoretical analysis of what is causing this to happen, which is possible because it is a pretty simple setup.", "questions": "Questions:\n- what is P set to for the experiments? \nSuggestions:\n- it seems like you need to explain the experiments more in the main text. the results are just shown without fully explaining the experiment setup. Are there two tasks, that are cycled repeatedly? \n- why not use some other image recognition task besides permuted MNIST? is it because the tasks have to be orthogonal? \n- Line 116 I would explain the regularization constraints are in words besides just presenting the equation. you should also link to where it is explained further in the appendix.", "limitations": "It seems like the model only works if the input tasks have orthogogonal solutions, is that true? if so, that's a huge limitation.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The paper models an important observed difference between humans and neural networks. \n- The model is remarkably simple and easy to understand. It is tempting to label the results as \"predictable\" or not surprising, but simple models are good.\n- The authors provide an impressive barrage of experiments analyzing their model. in particular, I am impressed by the experiments in section 5, where they show how specialization emerges as different hyperparameter choices reach their optimal values.", "weaknesses": "- it seems like parts of the architecture, such as the number of subnetworks and the size of the subnetworks are tuned to match those given in the system.\n- a lot of the paper structure seems different from what is typically done for neurips. for example, there is no related work discussed in the main text, but it is in the appendix. many of the experiment details are also in the appendix, which is weird because without the details, it's somewhat hard to understand the results presented in the main text. \n- as a somewhat theoretical and simplified model, it's unclear how much the approach relates to more \"real world\" tasks. the one task shown is MNIST, which is still a pretty artificial example because the second task is also MNIST, but with a weird permutation. This leads to my main question and reservation, which is that the model might only work if the input tasks are orthogonal? if so, that's a huge limitation, and dramatically changes my view of the paper. My current score is assuming that this limitation is true, but if it were not true, I would give a higher score (maybe an 8).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and thoughtful comments. We are delighted that you appreciate the simplicity of our model and the thorough experiments we provide.\n\n> The idea is to have subnetworks which can learn each task, and then learn a high level router to the subnetworks. \n\nWe thank the reviewer for the accurate summary regarding the routing function of the gating variables in our setup. In contrast to other methods that learn gating of task representations, we update weights and gates simultaneously through gradient descent, reflecting real-world conditions in which learning needs to be segmented and structured at the same time. Thus, they are fundamentally treated on equal footing. Critically, our suggested regularization (alleviating overspecification by favoring c to be of order 1\\) does not strictly encourage one task to be active at a time, since mixed solutions are given the same penalty. \n\nWe are excited about this model because it solves several computational problems with one mechanism: It is able to discover boundaries between tasks, retrieve task representations for previously learned tasks, and move representations for new tasks away from previous ones.\n\n > it seems like parts of the architecture, such as the number of subnetworks and the size of the subnetworks are tuned to match those given in the system.\n\nSo far, we have indeed only considered the case where the number of subnetworks P equals the number of tasks M (P=M=2) for simplicity. We fully agree that analysing P>M (i.e., multiple paths) improves the understanding of the model. To this end, we conduct a new experiment in the attached figure R4. The system still learns to solve both tasks and adapts flexibly to context switches. In this setting however, students only partially specialize, since they are now underconstrained. This disappears when regularization of student weight magnitude is imposed, which can physiologically be interpreted as a “representational cost”: Over time, we would expect a biological system to prune the additional components.\n\nWe note that, under certain conditions, the flexible regime can also emerge in the fully-connected architecture which is inherently free of these assumptions (Fig. 6).\n\nOn this basis, we adjust the presentation of our base model to the more general case P>M.\n\n> The paper structure seems different from what is typically done for neurips  \n> Many of the experiment details are also in the appendix, which is weird because without the details, it's somewhat hard to understand the results presented in the main text.\n\nThank you for pointing this out\\! We had deferred our related work section to the appendix for space constraints, but see how it is valuable in the main text, so that we will incorporate a shorter version there. We will move necessary experimental details back to the main text as well. \n\n> \\[why MNIST vs. a weird permutation of MNIST\\] is it because the tasks have to be orthogonal?  \n> This leads to my main question and reservation, which is that the model might only work if the input tasks are orthogonal?\n\nWe thank the reviewer for this important comment. We ran several simulations here to identify how orthogonality might impact the model. Prior related work (Lee et al., 2024\\) provided an extensive investigation of orthogonality of representations in linear networks. For theoretical simplicity, we considered only orthogonal tasks. We ran three new simulations in the PDF attached to this rebuttal. \n\nFirst, we plot how parameterizing the overlap between teachers affects student specialization (Fig. R1A). Ideally, we expect the model to maintain a shared representation according to the similarity between the tasks while separating the non-similar parts. Indeed, we see such a graded specialization in our model proportional to the overlap between the tasks. .\n\nSecond, we train our model on a set of three compositional tasks, with pairwise overlap (e.g., A+B, B+C, A+C), making them non-orthogonal. Our model learns to specialize to the underlying shared components and appropriately gate them to solve each of the three non-orthogonal tasks (Figure R1B).\n\nThird, we move slightly from MNIST to fashionMNIST to operate on natural images, and use two permutations with varying degrees of orthogonality. We use permutations because we are strictly interested in the setting of cognitive flexibility where the same input can be mapped to different responses. The new permutations are based on delineating  (orthogonal) upper-vs.-lower body or (correlated) warm-to-cold weather. We find that the system works well in both cases, with specialization occurring only marginally later in the non-orthogonal setting (Figure R2).\n\nIn summary, the studied phenomenon is robust to correlations between tasks. \n\n> My current score is assuming that this limitation is true, but if it were not true, I would give a higher score (maybe an 8).\n\nWe appreciate the reviewer stating their concern and clearly stating how it influenced their decision. We hope that our responses above and the additional simulations addressed the concern about orthogonality fully. \n\n> I would explain the regularization constraints are in words besides just presenting the equation. \n\nWe discuss two kinds of regularization:\n\n1. **Non-negativity of gates.** It is motivated by the fact that gating behavior can only amplify or attenuate a variable, but not change its sign.   \n2. **Alleviating overspecification.** Due to the linear nature of our architecture, the output *scale* can be changed both by amplifying weights or gates. However, gates should not need to account for learning the scale of the task. The regularization encourages gates to be of order 1\\. This has the effect whereby upscaling one gate above equilibrium will drive other gates to be downscaled.\n\nWe will add a short explanation to mathematical expressions in-place and emphasize the more extensive treatment in Appendix B.", "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": "Z3476w7eFz", "reviewer": "Reviewer_SwAN", "summary": "The paper investigates fast task switching/adaptation in a gated linear network. Specifically, it shows that neuron-like properties including regularization, fast learning, and non-negativity lead the model to demonstrate fast task adaptation and generalize compositionally. They provide a detailed analysis of the learning dynamics to show the components responsible for fast adaptation and the conditions in which a flexible scheme arises instead of a lazy forgetful one. They also replicate similar patterns in a more complicated task and model setting, and replicate human learning characteristics in prior studies.", "questions": "How many training instances were in the spiky loss part, in both the synthetic task and the MNIST task? E.g. does the model achieve single-shot adaptation in the task block?\n\nWould the gates adapt to a new task replacing an old task quickly?\n\nMust gating-like properties be at the 2nd layer in the deep monolithic network? What happens when you insert regularization and faster learning in layer 1 units? In general, for larger-scale applications, is it better to have gating-like mechanisms in later parts of the model?\n\nWhat happens when non-gating weights are also equipped with some amount of neuron-like properties? Are there neuroscientific reasons to maintain separate groups of task weights and gate weights?", "limitations": "Noted in the paper.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This work provides extensive model behavior/weight analyses and a thorough mathematical analysis of the learning dynamics to understand the driving forces of the observed fast adaptation.\n\nThe authors go beyond simple controlled tasks and a linear model to investigate the same phenomenon in a nonlinear model using MNIST.\n\nThis work draws a nice parallel with human cognition that is backed both theoretically and empirically.\n\nThe paper is dense but well-written.", "weaknesses": "No major weakness that I identified. One comment is that because the paper is very dense, some details are omitted (e.g. in figures) and need to be inferred.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful comments. We are glad to hear that you find the theoretical analysis of the underlying mechanism insightful. Moreover, we are encouraged that you value our aim to test our framework against behavioral experiments as well as to demonstrate a basic usefulness beyond synthetic tasks. \n\n> Q1: How many training instances were in the spiky loss part, in both the synthetic task and the MNIST task? E.g. Does the model achieve single-shot adaptation in the task block?\n\nThe reviewer brings up the important question about how quickly is this model able to adapt to previous tasks and new tasks. As we run our simulations in continuous time (gradient flow), each gradient is calculated over a large number of samples (200 for both cases). We made this choice to permit theoretical analysis where a sample average is required. Figure 1F shows the number of time steps needed until adaptation for the synthetic plot, and we supply a similar version for fashionMNIST in Fig. R2.\n\nTo investigate the sample efficiency of the model on the synthetic task, we now provide a simulation in Fig. R3 with coarsely discretized time where only a single sample is used in every update. While this makes the simulation more noisy in a single timestep, its qualitative features persist: In late blocks, the loss drops significantly on the first sample, and reaches its minimum after a few samples (few-shot adaptation). More generally, the persistence of the fast adaptation in the equivalent model in Fig. 3 (where noisy input is absent) suggests fast adaptation is indeed facilitated by the gates. \n\nWe will include both figures in the appendix of the final revision.\n\n> Would the gates adapt to a new task replacing an old task quickly?\n\nIf there are some shared components between the old task(s) and new task(s), the gates would adapt to appropriately modulate the parts of the network that correspond to the new task. For example, we illustrate this in our generalization experiments where we switch the network to a new set of tasks (A+B, B+C, A+C) composed of different combinations of the tasks previously encountered (A, B, C). In these cases, the network adapts the gates to turn on the parts of the network corresponding to the new tasks, and turns off the rest of the network, allowing it to generalize its learned knowledge rapidly.\n\n> Must gating-like properties be at the 2nd layer in the deep monolithic network? What happens when you insert regularization and faster learning in layer 1 units? In general, for larger-scale applications, is it better to have gating-like mechanisms in later parts of the model?\n\nBecause the deep monolithic model is linear, there is no meaningful difference between gating occurring in the first versus the second layer. The main distinction is that rather than gating the output of computations already performed in the network, inducing gating-like behavior in the first layer would instead gate off certain computations in advance (by zeroing their input) and only perform computations on specific paths within the network. If regularization and a faster learning rate are instead applied to layer 1 units, we would expect to see similar gating-like behavior emerge in the first layer rather than the second. Indeed, the question of whether there is an advantage to gating earlier or later in a network is an interesting open question which may only arise once additional depth and nonlinearity are added to a model.\n\n> What happens when non-gating weights are also equipped with some amount of neuron-like properties? \n\nThank you for addressing the minimal conditions needed for gating. We can distinguish two settings:\n\nFirst, what happens if we equip the weights with a fast timescale? We find that with very fast weight learning rates, specialization eventually disappears: The weight representations will be fast enough to quickly adapt after a block change, not requiring any gating (Fig. 5). This however is an unnatural assumption, as synaptic plasticity is slow.\n\nSecond, we can ask whether an architecture with two layers of weights, but without any scalar gates will show gating-like behavior. By regularizing and increasing the learning rate of one layer, we indeed observe the fully-connected weight layer to exhibit gating-like behavior, while the other layer learns to specialize to the corresponding tasks, emulating the same behavior we see in our gated model. (Fig. 6). \n\n> Are there neuroscientific reasons to maintain separate groups of task weights and gate weights?\n\nThe reviewer brings up an important question as to how the model structure maps to brain structures. There are multiple gating circuits in the brain. One is the prefrontal cortex (PFC) gating activity in earlier cortices. We think of the weights as the synaptic strengths in motor and sensory cortices as they learn a task, while the gating to switch between distinct computations is projected top-down from PFC. In addition, cognitive flexibility and switching task representations in PFC occurs through rapid changes in neural activity rather than synaptic updates. As such, we see them implemented in different regions, and also through different substrates (i.e. synaptic plasticity vs. neural activity changes). \n\n> One comment is that because the paper is very dense, some details are omitted (e.g. in figures) and need to be inferred.\n\nWe thank the reviewer for this valuable comment. We reviewed the manuscript and found some important details omitted especially regarding the details of the experiments run in each figure. The final version of the manuscript now addresses this.", "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:
We thank the reviewers for their thoughtful comments that helped us improve the manuscript. We were glad that the reviewers found “the model remarkably simple and easy to understand” (reviewers 9Ed7 and pBhB), and studying the emergence flexible task abstractions “creative and exciting” (9Ed7, SwaN). The reviewers also appreciated the significance of the work as modeling “an important observed difference between humans and neural networks” (pBhB, SwaN) with a “motivation that is interesting at a high-level” (reviewer 9Ed7). 

While the reviewers found the experiments to be “extensive” and the claims well backed up (reviewers 9Ed7 and pBhB), they also identified improvements that can be made to the conceptual structure of the paper for clarity. Additionally, reviewers pBhB and SwAN pointed out that the manuscript was dense and the experiments had some important details relegated to the supplementary. 

We discuss these insightful comments in the responses to individual reviewers. To address, we are taking the following actions to the final revision:

1) We restructured the manuscript for clarity: We rearranged the order of analytical results and grouped the points needed to state the analytical mechanism behind the model. We moved the *Related works* section to the main text to clearly state our contribution. We moved generalization results that did not support the main narrative to the appendix.   
2) We ran simulations to test how the model might handle shared task information, as our manuscript had only orthogonal tasks chosen to permit theoretical analysis. We see that tasks are separated to the extent needed, but not more (see Figure R1 in the attached PDF).  
3) We now better clarify our aim to contribute to theoretical cognitive science by developing a minimal neural network model of cognitive flexibility. In particular, the work addresses a gap in a line of previous work at NeurIPS which aimed at interpretable neural models, positioned between Bayesian approaches and more expressive machine learning architectures.  
4) Finally, we attach the necessary simulations which allow us to address the reviewers’ questions. 

Again, we thank the reviewers for their time and insightful comments. We hope our specific responses below address concerns and questions raised.

Meta-review: The reviewers all emphasized that the experiments are extensive and detailed, solidly backing up the stated claims. The simple proposed changes (faster time scales, non-negativity and regularization) on the gating layer are clearly demonstrated to induce task selectivity, enabling specialization of the subnetworks. During the discussion phase, authors answered insightful questions posed in the reviews in detail and provided new results for experiments on datasets with compositional tasks and a more complex image dataset, which provided additional insights and further evidence for the claims in the paper. The paper is a clear accept!
Decision: Accept (spotlight)