Separating the 'what' and 'how' of compositional computation to enable reuse and continual learning
Authors: Haozhe Shan
Minni Sun
Lea Duncker
Venue: NeurIPS 2025 poster ()
Keywords: continual learning
compositionality
contextual inference
low-rank RNN
Primary area: general_machine_learning
TL;DR: A two-system RNN achieves continual learning by using a 'what' system to infer compositional structures in tasks and a 'how' system to compose low-rank components, demonstrating transfer and compositional generalization.

Abstract: The ability to continually learn new skills, retain, and flexibly deploy them to accomplish goals is a key feature of intelligent and efficient behavior. However, the neural mechanisms facilitating the continual learning and flexible (re-)composition of skills remain elusive. Here, we study continual learning and the compositional reuse of learned computations in recurrent neural network (RNN) models using a novel two-system approach: one system that infers 'what' computation to perform, and one that implements 'how' to perform it. We focus on a set of compositional cognitive tasks commonly studied in neuroscience. To construct the 'what' system, we first show that a large family of tasks can be systematically described by a probabilistic generative model, where compositionality stems from a shared underlying vocabulary of discrete task-epochs. The shared epoch structure makes these tasks inherently compositional. We first show that this compositionality can be systematically described by a probabilistic generative model. Furthermore, we develop an unsupervised online learning approach that can learn this model on a single-trial basis, building its vocabulary incrementally as it is exposed to new tasks, and inferring the latent epoch structure as a time-varying computational context within a trial. We implement the 'how' system as an RNN whose low-rank components are composed according to the context inferred by the 'what' system.
The contextual inference facilitates the creation, learning, and reuse of the low-rank RNN components as new tasks are introduced sequentially, enabling continual learning without catastrophic forgetting. Using an example task set, we demonstrate the efficacy and competitive performance of this two-system learning framework, its potential for forward and backward transfer, as well as few-shot learning via re-composition.

Separating the what and how of compositional computation to enable reuse and continual learning
Haozhe Shan'
Sun Minni'
Center for Theoretical Neuroscience Department Computer Science Columbia University hs3594@columbia.
Center for Theoretical Neuroscience
Department of Neuroscience
Columbia University
ms57240columiia edu
Lea Duncker
Center for Theoretical Neuroscience
Department of Neuroscience
Columbia University
1d8149@columbia. edu
Abstract
The ability to continually learn, retain and deploy skills to accomplish goals is a key feature of intelligent and efficient behavior. However, the neural mechanisms facilitating the continual learning and flexible (re- jeoposittitoi of skills remain elusive. Here, we study continual learning and the compositional reuse of learned computations in recurrent neural network (RNN) models using a novel two-system approach: one system that infers what computation to perform, and one that implements how to perform it. We focus set of compositional cognitive tasks commonly studied in neuroscience. To construct the what system, we first show that large family of tasks can be systematically described by probabilistic generative model, where compositionality stems from shared underlying vocabulary discrete task epochs. We develop unsupervised online learning approach that can learn this model on a single-trial basis, building its vocabulary incrementally tasks, and inferring the latent epoch structure
it is exposed varying computational context within a trial. We implement the how system as an RNN whose low-rank components composed according to the context inferred by the what system. Contextual inference facilitates the creation, learning, and reuse of low-rank RNN components enabling continual learning without catastrophic forgetting. Using an example task set, we demonstrate the efficacy and competitive performance of this two-system learning framework, its potential for forward and backward transfer, as well as fast compositional generalization unseen tasks.
time-
tasks introduced sequentially,
Introduction
Humans and animals exhibit a remarkable ability to learn and retain new skills. and flexibly deploy them to accomplish goals in an ever-changing environment. A growing literature from neuroscience and human behavior suggests that mechanisms for contextual inference and task abstraction may play a crucial role for behavioral flexibility and learning [19]. An abstract, task-relevant context could guide the selection and composition of different skills, while the maintenance of context-specific
Equal contribution.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

memories could counteract forgetting and aid continual learning M. Yet, the neural mechanisms facilitating the continual learning and flexible reuse of skills are not well understood. How do neural systems manage to expand their repertoire of skills without interference or forgetting, while maintaining the ability to access and compose them in new environmental contexts? We hypothesize that multiple interacting learning systems with different objectives may contribute to overall learning and computation. Specifically, can inferences about the higher-level compositional structure of tasks be exploited for more robust and efficient learning?
We investigate the interplay of compositionality contextual inference, and continual learning in recurrent neural networks (RNNs) using two-systems approach: one that infers an abstract computa- tional context within a task family by parsing the compositional structure with probabilistic inference, and one that implements and retains the relevant computation for each context. This framework naturally maps onto a what and how architecture, previously explored meta-learning [10]. The explicit separation into contexts (what) and computations (how) allows for flexible learning and compositional reuse. We focus our analyses on a family of cognitive and motor tasks commonly used in neuroscience [Л]]. Here, we expand upon previous work in RNNs, where continual learning is still a major challenge [2], and compositional reuse of computational building blocks has only been shown to emerge implicitly over simultaneous training multiple related tasks (3116,, rather than being explicitly implemented as a solution that facilitates sequential learning of tasks.
The paper is organized as follows: We review related work and our contributions in section 2 In section 3 we formalize task compositionality through the development of generative model for cognitive tasks. This allows us to define computational context within task family. We develop an unsupervised online learning and inference approach to directly infer context from the inputs and target responses that constitute a task, one trial at time across a number of different but related tasks. In section 4 we combine inference and learning in the task model with a contextually gated low-rank RNN. Training this architecture on a number of cognitive tasks, we demonstrate that our approach counteracts catastrophic forgetting and aids continual learning with signatures of both forward and backward transfer. We demonstrate competitive performance in comparison to continual learning approaches previously used in RNNs. We show that our architecture can rapidly learn to deploy known skills to solve novel tasks using only few examples property called compositional generalization [9] Finally, we discuss limitations and future directions in section
Our work contributes to more precise quantification of task compositionality, how it may shape computation, and how interactions across different learning systems can benefit continual learning and generalization. Ultimately, addressing these questions will help to elucidate flexibility and learning in biological systems.
2 Related work and our contributions
Statistical theories of contextual inference. We follow an influential line of work on contextual inference in cognitive sciences I1 22) Contextual inference has been proposed as model explaining many learning-related phenomena in classical conditioning (20] and motor control M. Here, latent discrete-valuee random variable is used to describe distributional changes in task- relevant, continuous observations. Given set of observations, this variable can be inferred as familiar or novel context determine whether learner may express and refine existing behaviors, or expand their repertoire using a new context. We extend this idea to the problem of learning to solve multiple common neuroscience tasks [4] To do this, we develop a formal description of these tasks in terms of a context variable: Each task transitions through several contexts, each relevant for different computation on a particular task epoch. A key aspect of our model is that these contexts can reappear in compositional manner across different tasks, and may additionally share complex temporal dependencies within a given trial. In section 3 we show how this is crucial for capturing the distributional structure of neuroscience tasks, where, for example, the animal needs to compute response based on information presented at an earlier time in the trial
-esk-depnndent modulation of neural networks. The general notion of a separate system modulat- ing computation in a neural network in a task-dependent manner has been previously studied, An example of this is the Task-Conditional Hypernetwork 2 where the weights of neural network are generated by a second, separate network based on modifiable task embeddings. Other approaches often assume the network has its own Iearnable weights but that their involvement in computation
2

---

b
traa variable 0d determies the stimutus
{
inputs and target responees, given lask epochs
task epoch id
variable
X
Z1
Zo
Z3
x= 2
Aballor
stingia
redeanso
delay
Axatton
qt
q:
inputs, targets
trial variable
== 5
stiuulus-que quided Dalay ProjAnti
mamoryyguussee guddd resporise
decision ProtAnti
Memory Pro/Ant
Figure A generative model of family of cognitive tasks used in neuroscience. a: Schematic illustration of set of cognitive tasks and their compositional nature. b: Directed Acyclic Graphical model (DAG) description the generative model capturing the distribution over cognitive tasks we consider. The time-varying inputs and target responses of each task can be modeled as mixture, where the observations for each task are composed of a set of discrete epochs, as illustrated in (a).
modulated, for instance via gating mechanism 16 that can be solved by a single subcomponent (e.g., 25
These works typically study tasks 27 or the simple linear addition of outputs from multiple subcomponents (e.g., .... [24 28 In our case, we study tasks that can only be solved via a composition of computations carried out by different subcomponents. Our choice of architecture similar that of Costacurta al. [16. where an RNN receives a low-dimensional, time-varying nnurrmoduuatoryy signal from a second RNN that adjusts how different low-rank components of its recurrent weights are combined. While Costacurta et [6] train this architecture end-to-end with backpropagation, we use the interpretable output of contextual inference algorithm as the modulatory signal. In doing so, we make the compositional reuse of RNN components across contexts explicit, thereby facilitating continual learning and compositional generalization. Beyond RNNs, the 'local module composition' scheme for feedforward networks (29 is similar to our work in spirit, except that theirs requires passing each datum to all modules in order to determine the right composition use and does not make use of an probabilistic inference procedure.
24
25
Continual learning in RNNs. Continual learning is relatively understudied in RNNs and sequence- to-sequence problems more generally [12 30 In our comparisons, we focus on two previous continual learning approaches for networks with fixed dimensionality: Elastic weight consolidation (EWC) 3T, which penalizes weight changes during future learning according an estimate of how important each individual weight is for the performance of previous tasks. EWC is representative of family of regggttregulaiizttinnbbaaee based approaches [32] Another approach. Orthogonal Weight-Space Projection (OWP), was developed in [13), which pushes the network to solve different computations with activity in orthogonal subspaces. This method is representative of a family of approaches that modify weight updates during learning 33) We found that our approach mitigates forgetting better than these methods, exhibits forward and backward transfer learning, and can rapidly learn new tasks by re-composing learned computations.
3 Formalizing task compositionality via a probabilistic formulation
3.1 A generative model of compositional cognitive tasks
Our goal is to study how the inference of a task-relevant computational context could guide the selection and composition of different skills in an RNN and aid continual learning. To address this question, we initially need to define what a task-relevant computational context is, and how may be inferred as part of RNN training. We do this for family of compositionally related cognitive tasks which have been popular both in experimental and computational neuroscience
13 16 36).
In these tasks, the learner receives sequence of time-varying inputs $:=1,..,2 and needs to produce a target response sequence y I....... following different distributions p(81.T,/ 91:l/c) depending on the task identity e [1,..., N.). These sequences are structured according to particular segments (often informally referred to as 'epochs') such as a fixation period, stimulus presentation period, or
response period. Each of these epochs defines particular distributions of inputs and target responses
3

---

presented to the learner. What makes a family of tasks compositionally related is that they are composed from the same underlying set of epochs, albeit with desk-depnndent transitions from one epoch to the next. Fig. D illustrates this schematically for the set of tasks we will consider throughout the paper. We first sought to formalize the notion of shared epochs (Fig. [T). We can describe complex dependencies in the distribution of inputs and target responses p L: M1.rlc) c) for task by introducing a discrete latent variable 2 E z denoting the task epoch at time t. We model epochs as evolving over time with task-dependent Markovian transitions p(= = - j,c) =A;, Inputs and target responses qe [ Ve.yy E R are drawn from an epoch-dependent distribution. Each epoch may involve number of different imput values, reflecting c.g. different experimental conditions like reach direction or stimulus contrast. We therefore introduce dependence on an additional latent variable, the trial variable x, which can capture this trial-specific structure, with qe p(g. =)) (Fig []]. For example, z may index stimulus, which differs from one trial to the next and determines the distribution of inputs and target responses (see Figure 2 for examples). Importantly, sharing the same z across the entire trial allows the model to capture dependencies between, e.g., the target responses in a later epoch and inputs from an earlier epoch.
Using this description of compositional cognitive tasks, a task-relevant context naturally defined as the time-varying sequence of task epochs 2,. Thus, performing contextual inference requires learning the statistical structure and dependencies of the underlying generative model, and inferring the latent task epoch from observations qr- In the continual learning setting, this is essentially a problem of online learning and inference, which we address in section 3.3] We next give an overview of the specific tasks we consider throughout the rest of the paper. More detailed descriptions of the generative model and the specific task design are provided in Appendix A
3.2 Overview of our task set
We initially focus on six commonly studied neuroscience tasks [3 14. [6] that can be imple- mented in our generative framework. The tasks are schematized in Fig. la and we will henceforth refer to them as DelayPro, DelayAnti, MemoryPro, MemoryAnti, DMPro, and DMAnti. All of these tasks require learning simple stimulus-response relationships, albeit with more complex tempo- ral structure involving e.g. delayed responses. We summarize the different tasks in terms of their general goal and epoch structure in Table [T] In practice, there is substantial inter-trial variability within a task due to variability in the amount of time spent in a given epoch, variability in the trial variable æ, and additional variability q: [st yr). All tasks involve a 5-dimensional time series of inputs 8, and 3-dimensional time series of target responses yı, which take on noisy values around a piecewise constant mean. The value of the mean is dependent on the task epoch and trial variable x, chosen to reflect the underlying task goals. We refer to Appendix A.2 for more detailed description of the tasks, their epoch structures, and distributional choices.
Note that inclusion of the trial variable has important consequences for how tasks are segmented into epochs. Considering the MemoryPro/Ant Anti tasks: During the stimulus epoch (S), the inputs reflect the trial variable x, which indexes the stimulus direction for each trial. During the 'response period' (RM, the target response along/opposite the stimulus direction, thus making their dependencies on J: different. This necessitates treating them as two different epochs, RM,P and RM,A (see TT). While this segmentation arises from our model of the complex statistical dependencies of cognitive tasks, different compositions may be possible using alternative formulations. For instance, one might imagine that all response-epochs could utilize the same context irrespective of how relates to earlier stimuli, e.g. to move an effector in particular direction. Indeed, empirical results suggest that this solution arises RNN models trained on multiple tasks
3.3 What system: online learning and inference of the compositional task structure
We next consider how the compositional structure across tasks can be learned online (one trial at a time) and continually (one task at a time). The training data contains the sequences of inputs and target responses, as well as the task identity for each trial, [[gi. Ir Ntriats is the number of trials per task. We sometimes drop the trial index r from notations for brevity. Given the training data, the goal of the what system is to obtain estimates of the parameters that describe the distribution p(q:lc) (e) and transitions between epochs. It then infers the unobserved latent variables žn and on each trial.
where

---

Task name
Description
Epoch sequence F F+3 S Hp
DelayPro
Respond towards stimulus direction 0 after a delay.
DelayAnti
Same DelayPro but respond towards 0 +
F F RA
MemoryPro
Memorize stimulus direction 0 and respond towards it.
F+5-+ M RM,
MemoryAnti
Same as MemoryPro but respond towards 0 + π.
F+5- M + RMMA
DMPro
Compare amplitudes of two stimuli in directions 0 and O' and respond towards the direction of the higher amplitude. Same as DMPro but respond towards lower amplitude stim- ulus direction.
F - SDM + RDM,
DMAnti
F SDM + RDMA
Table I: Our framework captures an ensemble of commonly used cognitive tasks with a shared compositional structure. The epochs are F: fixation; S: stimulus; Rp/A: response towards/opposite from the stimulus angle; M: memory; P/AP/A: response towards/opposite from the memorized stimulus; Spm: decision stimuli; P/A. response towards/opposite the stronger stimulus.
persmeters
b
600
C DalayAnti,
DelayAnti,
MemoryAnti,
0=1, X=3
likehood trial (pug. over x)
0=1,, X=1
450
block
300
trial index
C+2, X=4
Cm2, x-2
task
-2000
poltind
e=2, x=8
block
-0000
0:3, x=
6000
DMPro
DelayPra
DMAnii
DekyAnn
MemoryPo
prob
0=33 x=1
task 3
-3000-
block
60
100 trial index
150
Figure 2: Online continual learning of the compositional structure of tasks. a: Schematics of online continual learning of set of compositional tasks. b: Colored lines: log likelihood (LL) of tasks (averaged over trials) from the learned model over the course of learning. Background colors: the task being trained at each trial. Dashed horizontal lines represent LLs computed using ground-truth parameters. e Single-trial epoch inference by the learned model. Top row: inputs and target responses of example trials. Dashed vertical lines indicate boundaries between epochs. Trial indices are dropped for brevity. Center row: output of waiiinngtiime inference (p( (qi., e)). Bottom row: output of test-time inference (p( (1(21/112, c)), which does not require access to the target responses or information from the future. Note that, since the F and M epochs in memory-guie -uoioddd tasks have indistinguishable observation models, they are merged by the learner.
When training trials are available as batches across all trials and tasks, inference and learning can be solved using the classic ExppctationnMaxxiizaa
(EM) algorithm. EM performs a coordinate ascent on the complete data log- ligliikel via alternating updates to the model parameters and to expected sufficient statistics of the latent variables (37]. To learn online from incoming streams of trials, we maintain estimates of the full-dataset sufficient statistics and update them incrementally after each trial. pseudocode of the algorithm and details on the inference and parameter update equations are available in Appendix I To overcome convergence to bad local optima in the log log-liellhhood develop a structured incremental initialization procedure described in detail in Appendix B.2
We applied our algorithm to the problem of online continual learning of the 6 compositional tasks introduced in section 3.2 To evaluate its performance, we computed the average single-trial log likelihood (LL) of different tasks, using the estimated parameters over the course of learning. We found that the model can recover LL on par with that from ground-truth parameters (Fig. 2p), and learning the structure of later tasks does not cause forgetting of previous tasks, even though the observation models of some epochs are shared across tasks. We note that this performance robust to changes in task ordering (see examples in Appendix BB) and the algorithm does not need to know the total number of tasks or epochs in advance, recognizing unfamiliar epochs on the fly.
The learned parameters can be used to infer the epoch identity during each trial, and we considered two types of inference. The more accurate inference (used to produce Fig. 2b) is ((.T, e). This uses both the inputs 81.7 and targets F1.T for inference and produces smoothing distribution in that
5

---

wher?? id
how?
affctiis ndworr woggtts
lowraank comfoxt 2
lowrank concod 3
can
Taok monira
ip. oceveel
contentual interence
weighred by contest from task
q- 16.36
targeis
network
Figure 3: Schematic illustration of two interacting learning systems and contextual gating of RNN dynamics. a: The task model (what) processes incoming input and target pairs infers the time-varying computational context of the current task. A downstream RNN (how) learns the relevant input-to-target transformation and is modulated by the task model's 's helief pe(z) over the current context. b: The effective recurrent weights of the downstream RNN are a linear combination of low-rank components, weighted by belief over each corresponding context. e: The gating by the task model causes the network express different low-dimensional dynamics across different computational contexts, each of which can be re-expressed and reused to execute different tasks.
{8t, yt and
the inference of 2, depends on data from t' > f. Since this requires access to the target responses, we call this 'training time inference." (Fig. 23: center row). We also utilize a 'test time inference' setting, where the learner has only access to the inputs up to the current time point, 81:2. This produces filtering distribution, o( :1 c), which turns out to be highly accurate after learning (Fig. 2 bottom row). In particular, it is able to disambiguate situations (c.g., RM.P, Ras.A A at the end of memory -ennorr-guided response tasks have identical 8; statistics) by using the ground truth task label and the learned epoch transition structure for each task (c.g., knowing that MemoryPro and MemoryAnti tasks respectively end with RM.P, RM,A). Hereafter, will refer this module that learns and infers underlying task structures as the 'task model', and to the inferred epochs as 'contexts'.
4 Composing recurrent computation across multiple tasks via shared contexts
The results presented so far were concerned with using information available based on the input and response pairs of the task to infer what to do. This was achieved through the construction of probabilistic model capturing the distribution across a family of compositionally related tasks, and the development of an algorithm that can infer a computational context within that task family. We now turn the problem of using this context to guide the selection and composition of different computations in an RNN that learns how to implement the relevant stimulus/repponee transformations. The general framework for how the two systems interact is illustrated in Fig.
4.1 Recurrent neural network architecture
We chose a simple RNN architecture, where the hidden-state activity h; evolves in time according to
h = (1 - ah.- + Wres (h. + Win s, + b + 10 E,
with decay rate a, activation function di external input 8,, input bias ii and uncorrelated Gaussian noise & The network activity is read out via
0 Wout (h+) + bout
with output bias baut, Given the external input and initial state, the network trained to produce a target response time series y, by minimizing the weighted mean squared error between y: and ye (Appendix
Inspired by related works on motif execution 16 38 the connectivity weights in our RNN are dynamically modulated by the posterior belief over computational contexts pz(2) inferred by the task model (Fig. (F) More specifically, each computational context in our design corresponds to set of weights: a low-rank component of recurrent weights U_ U_VT. input weights W input bias inn output weights Wrout, and output bias bout The effective network weights at time t are a weighted

---

DelayPro
DelayAnti
MemoryPro
MemoryAnd
DMPro
DMAnti
Adam
b
EWG
OWP
log test lces
-5
porformenos
200
800
200
400
800
200
400
800
total trials (x 2580)
total trials (x 2580)
total trials (x 2860)
d
NM-RNN
Hypernetwork
ContexRNN (ours)
-5
perforance
1.0
0.5
0.0
total triats 2560)
total trials 2580)
total trials 2580)
9
h
ContexIRNN (ours)
task
order
test performence
task
1.0
order
task
order
task
order
trials Ix 2560)
trial (r 2560)
Figure 4: Continual learning. a-f: log test loss (colored lines, first row) and test performance (colored lines, second row) throughout sequential training, each plot shows results for five random seeds. Background colors denote the task being trained at each trial. Adam optimizer general RNNs. b: Elastic Weight Consolidation (EWC) on general RNNs. e Orthogonal Weight-Space Projection neuromodulated RNN. Adam optimizer Hypernetwork RNN. f: task model with context modulated RNNs (ours). g-h: Color-coded test performance during sequential training of four different task orders. Each row color-codes the average test performance across five random seeds of a specific task over training. We compare OWP (g) with method (h).
(OWP) general RNNs. d: Adam optimizer
sum of these components (Fig. 3b):
Wree Σ pi(=) Pv(=)U.V
(3)
Win
Ep(( Win,
6th -P((s)) 6DT
(4)
Wout Ep(a) Wont yuur
=
(5)
This architecture allows the network to express distinct dynamics for different computational contexts. (illustrated in Fig. Be). And computations performed by this downstream how system are selectively gated by the what system. Contrary to previous work, e.g. Duncker et [33, we do not constrain the RNN dynamics associated with one context to be non-interfering with those of another.

---

4.2 Contextual modulation enables continual learning
Given the model architecture above, we now turn to the problem of training both systems in parallel on the set of tasks introduced in section 3.2 As before, all training was performed sequentially, where (s..y) 3 pairs from only one task were shown to the network for a number of trials until the task identity switched and the previous tasks were never revisited (illustrated in Fig. 2a). During sequential training of each task, we first ran the incremental EM algorithm (Algorithm D for the task model over a whole batch, then computed contextual inference Pz(2) p(za 11:1, Vist, c) for the batch to instruct the downstream RNN. During testing, the RNN instead received test time inference Pz(2) := ((ze) 8 lit, c) from the task model, which did not make reference to the target response yr.
We compared the performance of our approach to five baselines and summarized the results in Fig. which shows the log loss and performance (definition see Appendix C.3) on test trials.
There of the baselines operating on general RNNs governed by update equations ( and (2). where Wrec, W btn, Wout bout are unconstrained and shared across tasks. The general RNNs addition- ally receive task identity input c where e is a one-hot encoding of the current task. Under our default choices of hyperparameters (Appendix C.I these unconstrained RNNs have approximately twice as many trainable parameters as our context-modulated RNN. The first baseline is to naively train general RNNs using the Adam optimizer 39 without any additional measures for continual learning. In this case, interference across tasks led catastrophic forgetting, as is readily visible from the log-loss and the task performance, which quickly degraded after a change in task identity (Fig. 4p). The second baseline is Elastic Weight Consolidation (EWC) from Kirkpatrick al. [3T] a continual learning approach which selectively slows down learning rates of single network weights deemed important on previous tasks. Due to the challenging nature of continual learning in RNNs, the simple importance-weighred approach failed to mitigate catastrophic forgetting (Fig. Fb. For a baseline that specifically addressed continual learning RNNs, we included the approach from Duncker et al. [3] which selectively slows learning rates in subspaces the network explored previously learned tasks by projecting away these dimensions from the weight update throughout learning. We refer to this method as Orthogonal Weight-Space Projection (OWP). While this approach outperformed Adam and EWC, the network suffered from capacity limitation where tasks progressively got harder to optimize. As result, it became harder for OWP to achieve proficient task performance across all tasks (Fig. 4
The other two baselines are more similar to our two-system approach in that they have dedicated component that receives the task identity input and controls the weights used by an RNN that only receives the task input 8t. One of them is neuromodulated RNNs from Costacurta et al. [16., in which a neuromodulatory subnetwork RNN receives task identity input and outputs a signal that dynamically scales the low-rank recurrent weights of output-generating RNN. Both subnetwork RNNs are jointly optimized end-to-end. While this design is similar to our two-system approach, it did not automatically solve continual learning and quickly forgot previously learned tasks when trained on new tasks (Fig. 4t). The fifth baseline uses hypernetworks for continual learning [23 32) A hypernetoork receives distinct learnable embeddings for each task and outputs the task-dependent weights for a target network (an geepurtinnnrtiingg RNN in our case). Changes to the weights for previously learned tasks are penalized when learning new tasks. This method showed limited improvement over EWC and forgot some earlier tasks at the end of training (Fig. 4 By contrast, Fig. 4 shows results from algorithm, which combines the task model with the context-modulated RNN. Selectively modulating the RNN based on the inferred context mitigated catastrophic forgetting and allowed the network to maintain high performance on all previously learned tasks throughout sequential training. This result was robust to different task training orders (Fig. 4g. h).
4.3 Contextual inference f'acilitates forward & backward transfer
We next asked whether, given the compositional structure of the task design, shared computational context across tasks could facilitate transfer learning. Positive forward transfer refers to the phe- nomenon where the test loss for a given task decreases more rapidly when it trained after another task, compared to when it is trained from scratch. We observed positive forward transfer between nearly all pairs of tasks (Fig. 5 and Appendix DD. An exception was MemoryPro/ Anti tasks, which exhibited reliable forward transfer primarily when paired with another Memory task. When instead training after Delay or DM tasks, which share fewer overlapping epochs with Memory tasks, the model's 's test loss Memory occasionally converged to slightly higher final value after rapid initial
8

---

a
the first task
as the first task
the first task
DeleyPo log tost loss
0.0
after MemoryPro
after DelayAnt
Mermry Pro log tost loss
after MemoryAnii
after DelayPro
DMPra log best loss.
0.0
after DMAnti
DelayPro
-2.5
DMPro
-2.5-
after DMPro
after MomorryPo
-5.0
-5.0
-5.0
trials (x 2563)
trials (x 2680)
triale (x 2560)
b
DelsyPro
DelayAnii
DelayPro
MemoryPro
MemoryPro
MemoryAnti
DMPro
DelayPro
-0.5
log lest loss
-1.5
20
10 triale (x 256)
trials (x 258)
trials 256)
triale (x 256)
Figure 5: Transfer learning. a: The log test loss during training of task (specified in y-axis label), either as the first task (black) or after training another task (coded by color). We observe faster training after pre-trate teeiiiiiiig tasks that share similar epochs, illustrating forward transfer. b: The log test loss during sequential training of two tasks, when each tasks is trained for fewer numbers of trials than our default choice. The loss of the previous tasks continues to decrease after switching to new task, illustrating backward transfer. We plot results from three random seeds for each task order.
decrease relative to training from scratch. This phenomenon was also observed in Duncker et al. [3]. When each task was trained for fewer trials (switching to a new task before proficiency was reached), we found that its test loss continued to decrease during training on the subsequent task (Fig. 5p and Appendix D). as long as overlapping computational contexts were revisited. This indicates positive backward transfer. Existing continual learning approaches rarely exhibit backward transfer [40] Indeed, repeating the same experiments using OWP failed to produce improvements on old tasks when learning new ones (see Appendix D).
4.4 Contextual selection facilitatess compositional generalization on new tasks
As final experiment, evaluated the ability of our model to reuse previously learned contexts when encountering new tasks. The ability to do so is called compositional generalization, and has been of major interest in neuroscience [14 [5] 366 and machine learning more generally [17H19]. Specifically, we introduced a variant of the memory-guided response task in which the memory epoch (M) is omitted. These variants are denoted P' (F + S + RM,,) and M² Anti (F S + RA (Figure note that they are different from DelayPro/Delayntt tasks in that the stimulus is not shown during response epochs). We initially trained the model sequentially on M² Pro, M² Anti, and MemoryPro (F + + M + RMIP), and then assessed whether it can rapidly learn new task, MemoryAnti (F S + RAS,A), within only a few trials. During this final stage, we froze the downstream RNN and updated only the task model, thereby testing whether the model could compose previously learned contexts solve the novel task. Note that in principle, the freezing of RNN weights was not necessary and that the number of trials needed for learning the task model was generally much smaller than a typical batch size used for training the RNN. Our method reached approximately 83% accuracy within as little as 40 trials average. In contrast, baselines (also pretrained on the MP MPPro/Ane A and MemoryPro tasks) using Adam, OWP and Hypernetworks respectively achieved only 56%, 53%, 64% accuracy after 512 trials of full model training (Figure 6p. c). These results demonstrate that our model can rapidly generalize to new tasks by recombining previously learned components.
Discussion
How does the brain learn to solve tasks through compositional computation? We approached this question by first developing formalism of compositional tasks that is sufficiently expressive to cap- ture a suite of common neuroscience tasks as compositions from a shared vocabulary of task epochs. We then showed how it is possible to learn and solve these tasks continually using a dual-system approach: First, a what system performs online probabilistic inference of the relevant compositional structure of tasks vireina,ying, low-dimensional computational context corresponding to the

---

a pre-training task epochs
compositional generalization
fseutio
adimulus
responte
foodon
xémutus
delay
response
Budion
stmuuoy
daby
response
Mr Pro/Anti
Memory Pro
Memory Anti
contestual inforence
and RNN components
only contentual inference model
b
ContexIRNN (ours)
OWP
Adam
Hypernework
performans
1.0
Memoryknti log tist tosa
Memocyllm
-2
0.5
0.0
100
400
Irials
trials
Figure Compositional generalization. a: Schematic of pre-training on set of tasks and learning via contextual inference on new task. b-c: The log loss (b) and task performance (c) on test trials of the new task as a function of trials used to train the task model (ContextRNN) compared to sequentially training the full models using OWP, Adam and Hypernetwork.
task epoch. This signal is then used to create or (re- (re-)select computational components (one for each context) in second how system. Numerical experiments suggest that incorporating compositionality and contextual inference in this way allows the network solve the challenging problem of continual learning of compositional tasks. Our approach can utilize knowledge transfer to future and previously encountered tasks, and exhibits potential for compositional generalization.
Studying interactions across different systems and how this may benefit learning, computation, and compositional generalization is also of substantial interest in neuroscience p don't 't directly map our approach to particular candidate brain areas, our work relates to an extensive neuroscience literature in this space. Previous theoretical work has proposed that contextual gating of
36 41 While we
dynamics for motor sequences can aid flexible recomposition and learning, and be implemented via thalamocortical loops modulated via selective inhibition from basal ganglia (88. A similar architecture would be possible for our model to extend these ideas from the motor domain flexible cognitive tasks (see also 27)) and determine how the inhibitory signal from basal ganglia may arise according to contextual information from the what system. More generally, there is large literature action selection, implicating basal ganglia circuits in the selection and composition of motor sequences 442 43 Other lines of work have studied interactions between hippocampus and prefrontal circuits for learning and generalization
[3] 44 45). and how abstract task states in prefrontal- [5) 8 or orbitofrontal cortex [2 may be used to guide goal-directed behavior, computation, and learning.
Limitations and future work. A limitation of our work is that we can only infer computational context correctly when this is observable from the observed input and target response pairs of task. When different epochs map onto the same observations for all x (as is the case in our setting for e.g. fixation and memory epochs), the task model cannot infer that they are different. A possible future extension to overcome this would be to incorporate feedback from the RNN to the task model. A large error produced by the RNN for familiar context could provide valuable information when different contexts are difficult to distinguish based on external observations alone. Similarly, the segmentation of tasks into epochs is not necessarily unique and depends on our modeling approach. A different notion of shared compositional structure resulting in a difffernnt task model may result in a different segmentation and opportunities for reuse. In addition, we have assumed the trial variable to be discrete variable, which simplifies the analysis by allowing us to express observation models for different z, x as a look-up table of means. Assuming continuous x may allow the model to capture richer tasks and improve scalability under complex stimulus distributions. Finally. it may be of value to model the inference of unfamiliar epochs in more principled way by using non-marametric methods (e.g.. 1 20 A key challenge is to extend existing algorithms to our case where epoch emissions are controlled by a latent factor x that needs to be consistent across epochs, and context variables are shared across multiple tasks.
10

---

Acknowledgments and Disclosure of Funding
We would like to thank Dan O'Shea and Julia Costacurta for feedback on the manuscript and helpful discussions. This work was supported by the National Science Foundation and by DoD OUSD (R&E) under Cooperative Agreement PHY-2229929 (The NSF AI Institute for Artificial and Natural Intelligence), the Gatsby Charitable Foundation (GAT3708), the Kavli Foundation, and the Simons Foundation Collaboration on the Global Brain.
References
[1] James Heald, Máté Lengyel, and Daniel M Wolpert. Contextual inference underlies the learning of sensorimotor repertoires. Nature, 600(7889)4489-993, 2021.
[2] Robert C Wilson, Yuji K Takahashi, Geoffrey Schoenbaum, and Yael Niv. Orbitofrontal cortex as a cognitive map of task space. Neuron, 81 2014.
[3] Veronika Samborska, James L Butler, Mark E Walton, Timothy EJ Behrens, and Thomas Akam. Complementary task representations in hippocampus and prefrontal cortex for generalizing the structure of problems. Nature Neuroscience, 25(10):1314-1222 2022.
[4] Yael Niv. Learning task-state representations. Nature neuroscience, 22(10):1544-155; 2019.
[5] Mohamady El-Gaby, Adam Loyd Harris, James CR Whittington, William Dorrell, Arya Bhomick, Mark E Walton, Thomas Akam, and Timothy EJ Behrens. A cellular basis for mapping behavioural structure. Nature, pages 1-10, 2024.
[6] Timo Flesch, Andrew Saxe, and Christopher Summerfield. Continual task learning in natural and artificial agents. Trends neumosciences, 46(3):199-216 2023.
[7] Jingfeng Zhou, Chunying Jia, Marlian Montesinos-Cartageen Matthew PH Gardner, Wenhui Zong, and Geoffrey Schoenbaum. Evolving schema representations in orbitofrontal ensembles during learning. Nature, 590(7847))606-6 2021.
[8] Oded Bein and Yael Niv. Schemas, reinforcement learning and the medial prefrontal cortex. Nature Reviews Neuroscience, pages 1-17, 2025.
[9] Lucas Y Tian, Kedar U Garzón, Adam G Rouse, Mark AG Eldridge, Mare H Schieber, Xiao-Jing Wang, Joshua B Tenenbaum, and Winrich A Freiwald. Neural representation of action symbols in primate frontal cortex. bioRxiv, pages 2025-03, 2025.
[10] Xu He, Jakub Sygnowski, Alexandre Galashov, Andrei A Rusu, Yee Whye Teh, and Razvan Pascanu. Task agnostic continual learning via meta learning, arXiv preprint 2019.
[11] 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-366, 2019.
[12] Andrea Cossu, Antonio Carta, Vincenzo Lomonaco, and Davide Bacciu. Continual learning for recurrent neural networks: an empirical evaluation. Neural Networks, 143:607 -27, 2021.
[13] Lea Duncker, Laura Driscoll, Krishna V Shenoy, Maneesh Sahani, and David Sussillo. Organiz- ing recurrent network dynamies by task-computation to enable continual learning. Advances in neural information processing systems, 33: 13:14187-1 2020.
[14] Laura N Driscoll, Krishna Shenoy, and David Sussillo. Flexible multitask computation in recurrent networks utilizes shared dynamical motifs. Nature Neuroscience, 2797)1349 2024.
[15] Reidar Riveland and Alexandre Pouget. Natural language instructions induce compositional generalization in networks of neurons. Nature Neuroscience, 27(5):988-999, 2024.
[16] Julia Costacurta, Shaunak Bhandarkar, David Zoltowski, and Scott Linderman. Structured flexibility in recurrent neural networks via neuromodulation. Advances in Neural Information Processing Systems, 37:1954-1972. 2024.
[17] Thaddäus Wiedemer, Prasanna Mayilvahanan, Matthias Bethge, and Wieland Brendel. Com- positional generalization from first principles. Advances in Neural Information Processing Systems, 36:6941-6966 2023.
11

---

[18] Steven M Frankland and Joshua D Greene. Concepts and compositionality: in search of the brain's language of thought. Annual review of psychology, 71(1):273-303, 2020.
[19] Ronald B Dekker, Fabian Otto, and Christopher Summerfield. Curriculum learning for human compositional generalization. Proceedings of the National Academy of Sciences, 119(41): e2205582119, 2022.
[20] Samuel J Gershman, David M Blei, and Yael Niv. Context, learning, and extinction. Psycholog- ical review, 117(1):197, 2010.
[21] Samuel Gershman, Angela Radulescu, Kenneth A Norman, and Yael Niv. Statistical compu- tations underlying the dynamics of memory updating. PLoS computational biology, 10(1);: e1003939, 2014.
[22] James B Heald, Máté Lengyel, and Daniel M Wolpert. Contextual inference in learning and memory. Trends in cognitive sciences, 27(1):43 64, 2023.
[23] Johannes Von Oswald, Christian Henning, Benjamin Grewe, and João Sacramento. Continual learning with hypernetworks. arXiv preprint 1rXivv1906.00000 2019.
[24] 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, 1 117447)29872-29822 2020.
[25] Nicolas Y Masse, Gregory D Grant, and David Freedman. Alleviating catastrophic forgetting using context-dependent gating and synaptic stabilization. Proceedings of the National Academy of Sciences, 115(44):110677-11007)
2018.
[26] Timo Flesch, David G Nagy, Andrew Saxe, and Christopher Summerfield. Modelling continual learning in humans with hebbian context gating and exponentially decaying task signals. PLoS computational biology, 19(1) 19(1)=e101088)) 1010888, 2023.
[27] Wei-Long Zheng, Zhongxuan Wu, Ali Hummos, Guangyu Robert Yang, and Michael M Halassa. Rapid context inference in a thalamocortical model using recurrent neural networks. Nature Communications, 15(1): 2024.
[28] Kai Sandbrink, Jan Bauer, Alexandra Proca, Andrew Saxe, Christopher Summerfield, and Ali Hummos. Flexible task abstractions emerge in linear networks with fast and bounded units. Advances in Neural Information Processing Systems, 37:6938-6977 2024.
[29] Oleksiy Ostapenko, Pau Rodriguez, Massimo Caccia, and Laurent Charlin. Continual learning via local module composition. Advances in Neural Information Processing Systems, 34:30298-
30312, 2021.
[30] Shagun Sodhani, Sarath Chandar, and Yoshua Bengio. Toward training recurrent neural networks for lifelong learning. Neural computation, 32(1):1-35, 2020.
[31] James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka maabska-Barwinska.. et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, (1((1):3521-3322 2017.
[32] Benjamin Ehret, Christian Henning, Maria R Cervera, Alexander Meulemans, Johannes Von Os- wald, and Benjamin F Grewe. Continual learning in recurrent neural networks. arXiv preprint 200x:20006.1109 2020.
[33] Guanxiong Zeng, Yang Chen, Bo Cui, and Shan Yu. Continual learning of context-dependent processing in neural networks. Nature Machine Intelligence, 1(8):364-37,, 2019.
[34] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. Advances in neural information processing systems, 33:5824-5836. 2020.
[35] Zeyuan Yang, Zonghan Yang, Yichen Liu, Peng Li, and Yang Liu. Restricted orthogonal gradient projection for continual learning. Al Open, 4:98-110, 2023.
[36] 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.
12

---

[37] Arthur P Dempster, Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete data via the em algorithm. Journal of the royal statistical society: series (methodological), 39 (1):1-22, 1977.
[38] Laureline Logiaco, LF Abbott, and Sean Escola. Thalamic control of cortical dynamics in a model of flexible motor sequencing. Cell reports, 35(9), 2021.
[39] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[40] Sen Lin, Yang, Deliang Fan, and Junshan Zhang. Beyond not-forgetting: Continual learning with backward knowledge transfer. Advances in Neural Information P Processing Systems, 35: 16165-16177. 2022.
[41] Weinan Sun, Madhu Advani, Nelson Spruston, Andrew Saxe, and James E Fitzgerald. Organiz- ing memories for generalization in complementary learning systems. Nature neuroscience, 26 (8):1438-1448, 2023.
[42] Mclissa J Sharpe, Thomas Stalnaker, Nicolas W Schuck, Simon Killcross, Geoffrey Schoen- baum, and Yael Niv. An integrated model of action selection: distinct modes of cortical control of striatal decision making. Annual review of psychology, 70(1):53-7,, 2019.
[43] Jeffrey E. Markowitz, Winthrop F Gillis, Maya Jay, Jeffrey Wood, Ryley W Harris, Robert Cieszkowski, Rebecca Scott, David Brann, Dorothy Koveal, Tomasz Kula, et al. Spontancous behaviour is structured by reinforcement without explicit reward. Nature, 61887776):008-117, 2023.
[44] James CR Whittington, Timothy H Muller, Shirley Mark, Guifen Chen, Caswell Barry, Neil Burgess, and Timothy EJ Behrens. The oolman-eichenbane machine: unifying space and relational memory through generalization in the hippocampal formation. Cell, 183(5):1249- 1263, 2020.
[45] James CR Whittington, William Dorrell, Timothy EJ Behrens, Surya Ganguli, and Mohamady El-Gaby. A tale of two algorithms: Structured slots explain prefrontal sequence memory and are unified with hippocampal cognitive maps. Neuron, 113(2):321-333 2025.
[46] CF Jeff Wu. On the convergence properties the em algorithm. The Annals of statistics, pages 95-103, 1983.
[47] Krishna Nathan, Andrew Senior, and Jayashree Subrahmonia. Initialization of hidden markov models for unconstrained on-line handwriting recognition. 1996 IEEE International Confer- ence on Acoustics, Speech, and Signal Processing Conference Proceedings, volume 6, pages 3502-3505. IEEE, 1996.
[48] Kenneth D Miller and Francesco Fumarola. Mathematical equivalence of common forms of firing rate models of neural networks. Neural computation, 24(1):25-31) 2012.
13

---

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: We believe that the claims are well supported by the conceptual developments and numerical evaluations we made.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in 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.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer:
[Yes]
Justification:
discussion of limitation supplied the end of the paper.
Guidelines:
The answer NA means that the paper has no 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 their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on 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 is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
14

---

Justification:
Our results do not contain proof-based based theoretical results.
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 referenced in 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 to provide short proof sketch 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
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Algorithms are stated in the appendix and we also make the code available.
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 a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make 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 can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, 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 to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
15

---

Answer: [Yes]
Justification:
Code is provided in a repository.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
7/nips.com wttpcccpppbccccceee guwdes/ddeeeeee Cooeeeooddesssseeeeeee http:///d///uuudddd/doo//e//e/e/
ssioiolicy 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 as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: We specify important parameters and settings in the text as well as provide the code for reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[NA]
Justification:
We do not include experimental results that required significance testing.
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, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
16

---

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 error bar than state that they have 96% CI, if the hypothesis of Normality 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 out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: In general, experiments have very little compute requirement. The configu- rations used are stated the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well 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 It 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 httppp::whaarrpss...cpp/l/lllllt
Answer: [Yes]
Justification: have reviewed the code ethics and can confirm that our work conforms with it.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations 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: Our research aims to better understand task compositionality and how the brain may handle it. These do not have immediate relations to MI. applications in the real world.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
17

---

The conference expects that many papers will be foundational research and not tied 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 out that an 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 to 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 being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
See previous question.
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 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. Lieenses for existing assets
Question: 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?
Answer: [Yes]
Justification: We provide citations whenever necessary.
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 is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets 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 it has changed) should be provided.
18

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide code used for our experiments but at the current moment do not expect them to be used as packages.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is 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. Crowdsourcing and research with human subjects
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)?
Question:
Answer: [NA]
Justification:
No crowdsourcing nor research with human subjects was involved.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS 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: See previous question.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor 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.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
19

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs were only used for writing and editing.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy (https: for what should or should not be described.
20

---

Separating the what and how of compositional computation to enable reuse and continual learning
Appendix
All codes available on the public repository
A Formalization of compositional tasks and task designs
We propose formalization of shared compositional structures in sequence-to-sequenee (seq2seq) tasks and show that the formulation is expressive enough to capture many commonly used tasks in the literature. In each trial of a seq2seq task, the learner receives an input time sequence tt and needs to produce a target output sequence yt=1.....T Different tasks (indexed by E =) are distinguished from each other by their input distributions as well as the underlying input-to-targed rules, together encoded in (( V1.r (c). For notational convenience let q [i, yel E RP,.
We model this distribution over task-observables by introducing of latent variables resulting in the joint distribution
p(xlc)p(2) AQUT,,RR, p((pcpp) e) II p(e II p(a: 21, p(c).
(6)
We initially provide intuition about modeling choices, before formally specifying the distributions for the tasks of the main paper in section A. and A.2 below.
To model different tasks with shared compositional structure, first observe that in many tasks used in neuroscience experiments, each trial can be temporally segmented into discrete epochs, which we denote as =t E zt. Each epoch has its distinctive input and target statisties, though they may exhibit complex temporal dependencies across epochs throughout the trial. To illustrate this, we can consider simple task testing the learner's ability memorize: During a stimulus epoch, inputs provide information about a trial-specific latent variable (e.g.. a dot on the screen indicating an angle 8). The inputs eventually turn off, and the learner needs to maintain the relevant information in memory until a cue solicits a O-dependent response (e.g., a saccade in the direction of the angle). This example corresponds to the MemoryPro task from the main paper.
One might initially attempt to model this epoch dependence as a simple Hidden Markov Model (HMM), where each epoch has its own observation model p():(ze). However, this approach would fail to capture that the yr during the response epoch is coupled to the 8, from the stimulus epoch via the shared trial-specific latent variable. In our example above, the statistics of the stimulus epoch and response epoch are coupled since they are both dependent on 8. To take this cross-epoch dependency into account, explicitly model the "stimulus condition", indexed by the trial variable z. In the specific example above, there may list of possible directions and x selects one to use as 6. Thus, in our model of cognitive tasks, the observed inputs and target responses depend only the latent epoch z but also a second latent variable x, which additionally parameterizes the observation models of all epochs.
We model tasks as being compositionally related if we can use the same underlying set of task epochs and conditional distributions over qr gr|ze, to describe them. For a simple example of different tasks sharing compositional structure, we can consider the MemoryPro task described above, and construct second task, MemoryAnti, which shares the same stimulus epoch but requires response in the opposite direction of 8. Here, the dependence of the response epoch on is different (and MemoryAnt: will therefore have different response epoch relative to MemoryPro), but the statisties of the stimulus epoch are identical across the two tasks. The differences in what epochs contribute to each task and how epoch transitions the next is captured via task-dependent Markovian transitions over epochs 1
A.1 Generative model
To make our model of neuroscience tasks tractable and applicable to training RNNs, we make some simplifying choices for its components. For all tasks, p(z c) is assumed to be uniform distribution
21

---

over some finite set of size Ng. p(21:r |c) is assumed to be a Markov process with the initial state distribution and transition probabilities determined by c. Finally, the observation model p(q. |2t, x) is assumed to be a multivariate Gaussian with a (=t, -eeeeddde mean. Under these assumptions, our model can be seen as an HMM with c-dependent latent-state dynamics and x-dependent emission models for all latent states. If there is only one task and one possible x value in all trials, the model reduces to a standard HMM. In the more general case, it can be viewed as a mixture (across tasks identities) of HMMs with Gaussian Mixture emissions (across trial variables x).
In summary, for trials r
.,...... ......,Nriiats we generate tasks as
G ~ p(c)
task identity epoch identity
G E [I,...,NN 2f (I..... No)
c (Is)
=[1=1-1 = c
(aScreteAA,..))
, ~ Discrete(
No 1)
trial variable
Iy, E No.
- = 2, = x N (g.... T²
inputs and targets
qi E IR
This model thus allows us to specify a set of N₂ tasks sharing a compositional structure as follows. Let N₂ denote the number of epochs shared among these tasks. The composition of the e-th task is specified by its epoch-transition parameters: the transition probabilities A E NgNx.N, AS. = p(z 2|22-1 c² and initial probabilities IIe e R.: II p(21 z|c). The observation model of each epoch is specified by the means corresponding to different x values, the mean of the multivariate Gaussian p(q. 2 = 2, x)). Thus, altogether an ensemble of .o tasks composed from N₂ epochs is specified by the tuple where the first component specifies the epochs and the second specifies how they are used to compose the tasks.
{q E {=.).......NN. where q
(A ,II)) [=)))=......N
A.2 Expressing common cognitive and motor tasks in our framework
Below, we provide explicit details of the distributions and parameters used to generate our task set. We follow the same epoch notation introduced in Table the main paper. We initially describe the temporal structure in terms of epochs of each task, and then define epoch-specific distributions.
All tasks start with the fixation epoch (F), which presents no stimulus and an active fixation cue. The required output is to maintain fixation without response. In all the non-response epochs (stimulus S, memory M and decision stimuli SpAS), the fixation cue is on and the re- quired output is to maintain fixation without producing a response. During the response epochs (Rp, Ra, RMIP, RM,A, RDM.P, RDM,A), the fixation cue turns off and the learner needs to stop fixation and produce a response & according to some rule, as described below. Without loss of generality, we considered 8 possible stimulus conditions per task (N, 8) and generate æ i.i.d. from a uniform distribution for the 8 possible values.
Delayed response tasks (DelayPro, DelayAnti). After F. the stimulus epoch (S) presents an angle 0, chosen from {0, T/ /4, .... 7x/4) depending on æ. The stimulus presentation stays on during the ensuing response epoch (Rp for DelayPro or RA for DelayAnti), where the target output becomes = A (R)) or Φ = f + π (R
Memory-guided response tasks (MemoryPro, MemoryAnti). After S, stimulus presentation disap- pears in the memory epoch (M). During the ensuing response epoch (RM, for MemoryPro, RM,A for MemoryAnti), there is still stimulus presentation and the learner must produce response based on the memorized 0: Ø = 0 (RM,P) or & = 0 + x (RM,A)-
Decision making tasks (DMPro, DMAnti). After F, decision stimuli epoch (Spar) presents two stimuli simultaneously. 0 = 0 with strength y in input dims 1, and = x with strength y in input dims 3, 4. The strengths (7, are determined from set of pairs by and scale the input channels. During the ensuing response epoch (RDM, for DMPro, RDM.A A for DMAnti), the stimuli persist and the required output á is the direction of the stimulus with higher strength (RDM. ,P) the one with a lower strength (RDM,A).
The epoch structure of the different tasks is encoded via A Altogether, the 6 tasks can be described as compositions using shared pool of 10 epochs. Given the epoch identity, the conditional distributions of the inputs and responses are generated as follows:
22

---

epoch
8 $
P.M .M
0,0,0,00
0, 0,
S
cos 0, sin 0, 0, 0
00, 0, 0]
Rp
cos 0, sin 8, 0,
[cos 0, sin 0
RA
cos 0, sin 0, 0, 0,0,11
cos(( T), sin (0 + mT
RMP
0, 0, 0,0,0..
cos 0, sin 0
RM,
0.0, 0, 0,
cos + sin + m) 1,0,0 cos sin Φ. 11. where á 1 1y>,, 1,
SDM
y cos 0, sin 0 7 cos 8¹ 7 sin yy11
RDM,
y cos 8, sin 0, 7 cos 0', sin 0'
RDMA
y cos 8, y sin 0 7 cos 8¹ o sin 8'
cos Φ. sin Φ, 11. where á = lye 1 + +1,+1,1 +,
Where is an indicator that takes on value the subscript true and otherwise. This fully specifies the conditional distributions p(qelza, 22, .)) for each task epoch with 9=,= = [3..x1 Vz.zl- For the DM tasks, we restrict the stimulus values to two locations 0 = 0, 0' = but pick N2 = 8 different combinations of possible (7, pairs. Note that while we generate the input and response distributions using stimulus values 0 to follow convention from the literature [1 [4], each value of A (or (r, Y²) the DM tasks) maps onto different x value. For we used 0.05; for (Y, Y'), x selects from [0.5, 1), [1, 2], [0.5, 2), [0.2, 1.5], [1, 0.5,, [2, 1J, [2, 0.5], 0.2].
Note that while have followed many conventions from previous work in the task design, these previous approaches tend to implement each task individually but with related distributional assump- tions M 14 [6]. Modeling shared structure across tasks through an explicit, shared generative framework for an entire task family is novel. While this was not the focus of the main paper, it is worth noting that access to description of shared statistical structure in the input and target response pairs of each task forms an important baseline for expectations of shared statistical structure across tasks in the solution emerging after RNN training 11 15
B
Online learning and inference of compositional task structures
B.1
Online learning and inference
In this section, provide additional details the algorithms developed for performing posterior inference over the latent variables of the task model, and online (one-trial-at-a-tine learning of the task model.
Inference. We can perform exact inference by utilizing a message passing scheme similar to that used for performing inference in classic HMM models.
Let af (2, c)
p(gf.> 2f 2 x z, c, c) denote the forward (filtering) message for given task. During a filtering pass, we compute
af(=,2,c) p(= = z/c, c)r(g)
z, 3r =
[1,1(,,,,C)) Σ
= =|æ = ,',cr = c0) pl
= =
(8)
Marginalizing the forward message at the final time-step allows us compute the marginal likelihood over observations
N₂ plq = I, = (2,
(9)
(10)
Let B8 (2, = p( 2, = v. Cr c) denote the backwards (smoothing) message.
During the smoothing pass, we compute
(2, =1
(11)
N. (2, Σ
p( = 2'127 =,,c, = c)p(g/411 = 2.,2, = =z),541 1
(2' xx, x),))
(12)
2mm1
23

---

Given these quantifies, we can compute the joint posterior over of, and Cr as
(2, c) = p(z; = 2, = C, = clq p(q' 24 22/2, 2,, -
(13)
2, 2, (,, = 2,
(14)
Ex, p(q [r, = I, C, c)p(ar (2, (2, c)pla, -2, -c) Ee,cE.
(15)
= C, c)
Note that the filtering and smoothing passes can be done for each value of and in parallel. Finally, for later use in the online learning algorithm we also compute the joint posterior over z[-1, 2f, Zr and Cr as
Ei o L. c) = p(zf-1 = 2, = =,,,,,=,,C, =,, = =cqi:r)
(16)
= af-1 x, c)p(2) = 2210f-1= = =,, =))p(g] = I' = x,,)',,,,,),,, c)p((, = 2, =c)
Σx, P I, I, G. c)p(z, = I, =c)
(17)
Online learning. Learning aims to recover parameters of the generative model, {AA, .I.............. up to permutation over 2 and x. Perform- ing parameter learning with EM requires computing the expected counts of visiting particular epochs or transitioning across epochs across all trials. When all trials are available as a batch, this involves sums over the quantities computed during inference for each trial. To make notation more compact, we denote X the of sufficient statistics needed to update the set of model parameters O Let ((..) denote the learned parameters after seeing the i-th trial and running & EM iterations. Let X².(f.k) denote the single-trial statistics trial computed using parameters (t.A), When the learning algorithm has access to all trials throughout learning (batch EM), the updates take the form
(Qq.. )====_..,,,==11..._
@(*++1) = f(S(i.)( ((x, = Σ XT.(1.))
(18)
For example, for the parameter estimates for the transition matrix, this takes the form
Σ1= -x=1 Eeres/ (L)
A 1 ez2'
(19)
No Σ - = 21 (1)
with
Li.k cazzs (t) = Σ (c c, v, 2F 2, 2 = ='1qii:r:01 0+).
(20)
In general, the batch EM parameter updates take well-known forms for HMMs and GMM-HMMs, which is why we only give an example here in the interest of brevity.
When trials are only available expected sufficient statistics across trials have to be approximated and updated after each trial instead, leading to an approximation to the batch updates. We propose the following update rule to perform parameter learning online
trial-at-a-tinn and cannot be revisited (online EM), the sums of
(1,++1) (1 ppaaams)
10 (i-1,,
+ Aporams ( ontine (X)).
(21)
cnttine (X) (1 Nstats) ontine
(()-1 ( (X) + (i.(
(22)
where K denotes the number of iterations per trial. After learning the i-th trial, we only need to store (G,K) and g(k.K) ontine! (X), taking the form
((.K)
Σ( Upprame)) Mparame (s(KK) (X)) + (1 - *pparam)) ( Oinit.
(23)
s(.,)) (X)
E Nstats) i Xx,(K.)
(24)
We summarize this online learning approach in Algorithm [ and show performance for different training orders (supplementing Figure in the main paper) in Figure S.
24

---

Log likelihood per trial
600
Log likallhood per trial
700
600
500
400
300
100
150
0
50
100
150
trial index
trial index
Figure S.I: Online continual learning of task structures with some example task orderings. See Fig for legends.
Algorithm 1 High- Overview of the Incremental EM algorithm.
θ - init.
Initialize parameter estimates.
S(X) <
> These are the estimated sums of sufficient statistics across trials
for r I...., ......,.riii Ntrials do invremental_imitiiaizeee : cr)
θ
See B.2
Ô +-6
> Create temporary copy of the parameters
S(X) + None
for k i..... K do
X getstat.s (qT.. Cr. -)
Compute sufficient statistics of this trial with current
params.
5(X) + (1 - Nstatt )((X) + X
> Decay the sums from the previous trial and add new
+ θ 3 aarams M Moate) + Mparams fSS())) @ Mgate update of params. using the modified sums. Mgate a binary mask controlling which parameters are updated (see B.3).
Incremental
end for θ - è S(X) + (XX) end for
B.2 Incremental initialization model parameters
Even for simple models such the Gaussian Mixture Model, EM is susceptible to local optima [46]. To avoid convergence to local optima, it important to obtain good initializations for the model parameters. For simple time-series models such as the HMM, this is typically achieved by collapsing sequences across time, performing clustering (e.g. with K-means (47)), and using the resulting cluster centers as the initial observation-modee parameters, given each latent discrete state.
In our model, we need initialize the underlying cluster centers of the observation model, Here, N₂ reflects the fact that the total number of epochs in the en- tire dataset is unknown a priori. We overcome this by setting up large number of 'slots' (Ñ N2), exceeding the likely total number of epochs in the task family. Our setting adds two significant challenges relative simple HMMs: First, each trial contains only few epochs and specific stimulus condition (r value). Therefore, clustering each trial can only initialize a small subset of the N2 mixture components of the observation model. Second, since there is one stimulus condition per trial, the mixture components a... explaining a given trial must be allocated to the same but different epoch states z. Initialization schemes based on simple clustering approaches on the entire trial (such as that outlined for HMMs above) would be agnostic to this structure and fail to provide feasible initial set of parameters. To overcome these issues, we introduce an 'incremental initialization' scheme, which is applied to all the estimated means . before learning each trial.
To introduce the scheme, we first introduce the notion of the putative 2, x. denoted ā. This is to highlight that learning the correct } does not require inferring the real z, in the generative
process but only requires them to be correct up permutation. Let No, No denote the number of
25

---

epochs and task slots in the learning algorithm, respectively. Note that these do not need to be set as the correct Nz, .V. they can simply be large integers. We do assume that the system knows the correct number of a values, Na- The learner keeps track of (...a.
as well two Boolcan-valued tables, Fe.x of size Ne x N₂ and Fæ,: of size N₂ x Ng. The two tables keep track of which combinations of c, 2 and ž, 2 have been encountered.
When the rth trial observed and the learner has access to (qi.r, Cr), the scheme is threefold:
1. We perform K-means clustering on the entire sequence Ti., This gives set of cluster means. Since we assume the epochs to have piecewise constant inputs and mild noise, the means correspond to the different epochs that appeared in this trial. The challenge now is to assign ž, I.. to these means, and to use them to initialize qg,2 accordingly.
2. All cluster means from step (1) should be assigned the same h,. We check the means against {Ā=.s) and decide on ž, according to set of rules and Fe..- The cr, By pair is marked "encountered" in Fe.a-
3. Given ř, from step (2), we treat the cluster means not found (á,,2)==5, as unfamiliar 2, ã. meaning that they represent previously unseen ž, E combinations and should be used to initialize. Each center is assigned a different & that has not been encountered (according to Fa The ž, 2, pairs are marked as 'encountered" in Fa.z.
B.3 Gated updates to parameters
Since each trial contains information about only one task and the few epochs it is associated with, it does not make sense to update parameters related to other tasks and epochs. For A%, we simply gate it such that only the transition matrix corresponding to the current task (the label of which is given) is updated. For (Q..x), we infer which epochs appeared in this trial using the posterior p( Only epochs with sufficiently high chance of appearance have their (á,.n) updated.
B.4 Epoch identifiability in our set of tasks
For the particular set of tasks considered, since F, M epochs have identical qu. for all ar, they are indistinguishable. Thus, our learning algorithm will combine them into single epoch, which we denote as F/M. Instead of learning the 10 epochs in the generative process for our of tasks, the task-model will end learning 9 epochs, including the combined F/M epoch. In terms of the transition matrices, this creates complication for MemoryPro, MemoryAnti tasks where both F and M epochs appear in each trial. The learned transitions will not be deterministic in the sense that the F/M epoch may transition either the stimulus epoch or the response epoch. The ground-trutth parameters used for plotting in Fig. and Fig. [refer to the optimal parameters with merged F/M epoch. In future work, it will be interesting to investigate how epochs that are computationally distinct (e.g. holding still in F vs. holding still while maintaining a memory in M) but map onto the same observations may distinguished, e.g. via feedback from the downstream network implementing the different computations for each epoch.
26

---

C
RNN architecture and hyperparameter settings
C.1
Default RNN architecture and task parameters
Parameter
Value
a
0.1
0.05
Φ
ReLU
number of hidden units
256
rank of Uz. V input noise std
22/a oin where o. = 0.01
minimum duration of epoch
time steps
p(2+11 Zz
0.9
C.2 Default training protocol
We used a batch size of 256 and trained each task for 1000 batches unless otherwise specified.
For the context-modulated RNN, the learning rate for the weights associated with each context z was initially set to 7)s = 0.001. After training each task
was multiplied by a decay factor = 0.5 for any context with p( c) 0.001. During training task c, L2 weight regularization with coefficient of 10 for contexts with p(z c) > 0.001, and set to 0 for all other contexts.
applied
For baseline algorithms using general ("vanilla" ") RNN architectures, the learning rate set 0.01, and the L2 weight regularization coefficient was set to 10 10-5, Parameter choices were determined by coarse grid search.
C.3 Loss function and performance measure
The loss function is a weighted mean square error similar to ||I 14 where is the index of the output units, m,, for response epochs and my,e 0.2 for all other epochs.
L:= (m,,. ")se)²) s.
trial is considered correct if the network maintained fixation for all time steps before the fixation cue turns off, and responded to the correct direction for time steps in the response epoch. the activity of the fixation output exceeds 0.5, the network is considered have broken fixation. The network's s response direction considered correct its angular difference from the target direction is less than T/ 7/10, Average performance and test loss were calculated 200 test trials for each task.
D Supplementary results
D.1 Additional results transfer learning
We provide additional results on transfer learning. Figure S.2 and Figure S.3 supplement Figure 5p and b, respectively, by showing results for all task pairs. Figure S.4 shows the lack of backward transfer when training with the OWP algorithm [I3,, where the test loss of previously learned tasks did decrease when training on subsequent tasks with overlapping epochs.
D.2 Results with other hyperparameter choices
We observed improved performance of our continual learning algorithm we increased the rank of U. and V. (denoted by r) (Figure b. compared with Figure Using r 3, 5, and in ContextR NN results in 34587, 43803, and 66843 trainable parameters, respectively - all fewer than the 69379 parameters of general RNN with the same number of hidden units. With ContextRNN performed worse with tanh than with ReLU activation, but this gap closed 10
27

---

after DelayPro
after DelayAnti
after MemoryPro
after MemoryAnti
after DMPro
after DMAnti
DeloyPro log
DolaAAnii
MemoryPrg
-5
DMiPro MemoryAnii
-5
DMArt
100
100
100
100
100
100
trials (x 2580)
trials (x 2580)
trials (x 2580)
trials (x 2560)
trials (x 2580)
trials (x 2680)
Figure S.2: Additional results on forward transfer. Colored curves show the log test loss of task B (indicated by the row label) when trained after task A (column label). For comparison, black curves in each row show the log test loss of task B when trained from scratch. We plot results across three random seeds for each task order.
(Figure S.5 d). Similarly, performance with and = 0.2 was worse than with 0.1, but this gap was also closed at r 10 (Figure S.5c. f).
E Details of baseline implementations
E.1 Elastic Weight Consolidation
slows down learning rates for network weights deemed important for previous tasks. This is achieved by adding regularization term to the training loss function. For a loss Z(9), the EWC objective is given as
Elastic Weight Consolidation (EWC) 31,
A Σ F. O,) C CEWC ((e) + 2
(25)
Here F; is an importance weight that ties the ith parameter value θ to it's value O at the end of training on the previous task. The important weights are computed as the diagonal of the Fisher Information matrix F evaluated at the parameter values 0 We set À to 10 after a coarse grid search.
E.2 Orthogonal Weight Projection
The original learning rule in Duncker et al. was derived for a parameterization of network dynamics of the form
-= W Win
(26)
given an element-wise activation function s(.). The learning rule was intended to maintain the stimulus/response relationship of previous task by applying of projection matrices the gradient used to update the network weights over learning. The projection matrices were intended to remove
28

---

during
during
during MemoryPro
during MemoryAnti
during DMPro
during DMAnti
DelayPro
DelayAnti
DeleyPro
-1.50
DelayAnii lest
-1.25
MemoryPro
-1.1
-12
-1.3
MomoryArt AomooryArti
-1.1
-12
log
-1.3
loss log test
DMiPr
-1.25
-1.50
DMArnti
-1.50
10
0
10
10
10
0
10
10
trials (x 256)
trials (x 256)
trials (X 256)
triats 256)
braals (X 256)
Mele 256)
Figure S.3: Additional results backward transfer. The log test loss of the previously trained task (indicated by the row label) continue to decrease during subsequent training on another task B (column label). Each task trained with 2560 trials and we plot results across three random seeds for each task order.
DelayPro
DeleyAnti
DeleyPro
MemoryPro
MemoryPro
MemoryAnti
DMPro
DeleyPro
log lost lots
10
20
0
10
20
10
20
10
20
trials 258)
trials 256)
triils (X 256)
triil 296)
bg lest lass
10
20
30
10
trials 258)
trials 256)
trials Ox 256)
trials 256)
Figure S.4: No backward transfer with OWP. a: The log test loss during sequential training of two tasks with the OWP algorithm, each tasks is trained for 2560 trials. The loss of the previous tasks does not decrease after switching to a new task, indicating no backward transfer. Results are shown for three random seeds per task order. b: same but with each task trained for 3840 trials.
29

---

denote the concatenated network state and input state at time of trial on ContexIRNN: = rm 5, S ReL.U
ConteXIRNN: d= 0.1, 10, ReLU
order 4
test perfommacc
order
order
order
100
200
300
500
600
100
200
300
400
500
800
2580)
(x (580)
ContextRNN: a= D.1, r= 3, @ tanh
ContextRNN: a 0.1. 10, @ tanh
order 4
tast perfomance
order
order
order
100
200
300
400
500
600
100
200
300
400
500
600
2580)
2560)
ContextRNN: r ReLU
ContextRNN: 0.2, r- 10, + ReLU
task
order
last performance
lask
order
task
order
task
order
100
200
300
500
600
100
300
400
500
800
trials (x 2560)
trials (560)
Figure S.5: Continual learning performance with different hyperparameter choices. a: Color-coded test performance during sequential training of four different task orders. Each row color-codes the average test performance across five random seeds of specific task over training. RNNs used for this panel have = 0.1, r = 5, á ReLU, b-f: same as but with different hyperparameter choice labeled by the title.
directions from the weight update that would interfere with previous tasks and were defined as follows.
Letting 2,
C the collection of all time points and trials on tasks through and task c, Z i.c 1 .............................. T
W rree Win the concatenated weight matrices, we define the projection matrices
Plee
TI 1:c :1:++)=++++))))))r)))
(27)
pl
W2 TT W. WT +1)
(28)
and the modified learning update as
CL p1:c VWL Pl:
(29)
where Vw C is the derivative of the loss on the new task with respect to the network weights. An analogous set of projection matrices and modified learning update is used for the readout weights.
To facilitate direct comparisons with our approach, we adapted the learning rule a modified setting, where the RNN dynamics are expressed as
h h
Wree ((h) iin
(30)
30

---

given the same element- wise activation function o(-). While the two parameterizations in 26 and are generally considered equivalent [48], the linear intuition used to motivate the approach of should be exact in 30 With I; = ((h,), 1:c is unchanged, but we instead use
33
13
P1 = -1 H HHI
(31)
where H 1sc h ...... h7 While this is very similar the version in Duncker et al. [3, the projections matrix now only depends implicitly on the W of previous tasks. We performed all comparisons using this modified learning rule.
E.3
Neuromodulated RNN
Costacurta et al. [16] designed a neuromodulated RNN (NM-RNN), which a neuromodulatory subnetwork (RNN1) outputs low-dimensional signal that dynamically scales the low-rank recurrent weights of an output-generating RNN (RNN2). RNNI and RNN2 are jointly optimized through gradient descent. In our implementation, RNN2 has 256 neurons and with rank (ank(ree)) 27, equal to the default choice for our ContextRNN. Number of neurons in RNNI is set to 125 such that the whole network has 36658 parameters, comparable the 34587 parameters of our contextRNN, is set to 0.1 for RNN2 and 0.01 for RNN1 following the original paper. We used Adam optimizer, with learning rate to 0.001 and L2 weight regularization coefficient to 10 after a coarse grid search.
E.4 Hypernetwork
Von Oswald al. 23 designed task-conditioned hypernetworks tackle continual learning. A hypernetwork receives distinct learnable embeddings for each task, and outputs the weights for the target network (an output-generating RNN in our case). For the output-generating RNN, we the number of neurons to 256. A full-rank recurrent weight matrix, an input weight matrix and an output weight matrix together require a 67843-d output from the hyperetwork. We used a chunked hypernetwork introduced in the original paper, where a set of learnable chunk embeddings serve as additional input to the hypernetwork. Concatenating the hypernetwork output from distinct chunks gives all parameters for the target network. We set the output dimension of the hypernetwork to 2000 and number of chunks 34. Dimension of task embeddings and chunk embeddings is set 32. The hypernetwork is multi-layer perceptron with two 32-d hidden layers. Together, the number of trainable parameters of the hypernetwork equals 70416, comparable other baselines in our paper using a full-rank RNN. After coarse grid search, we set continual learning regularization strength to 1, orthogonal regularization strength to 0.001 and clip gradient norm to (see Ehret et [32] for details). We used Adam optimizer, with learning rate set 0.001 and L2 weight regularization coefficient set to 10-5
31


Reviews:
[{"review_id": "SNGfFd9ERn", "reviewer": "Reviewer_fkcB", "summary": "The authors focus on continual learning in RNNs.\n\nSpecifically on compositional computations and modular reuse across sequentially learned tasks, which in my opinion is a topic of increasing significance. \n\nThe paper introduces a dual-system architecture consisting of a \"what\" system, which probabilistically infers computational contexts or epochs from input data streams, and a \"how\" system, which dynamically modulates RNN parameters based on the inferred context. \n\nThe paper has a  clean theoretical formulation, good interpretability through structured context inference, and it gives solid empirical evidence to demonstrate superior continual learning and rapid generalization to novel compositional tasks.", "questions": "1. is it possible to include a more realistic task that comes from an NLP real-world application in your evaluation?\n\n\n2. I am not convinced that the proposed approach would be effective in the presence of a large number of tasks or epochs (in the order of 100s). Can you please comment on the scalability of your approach? How would you expect the memory/compute requirements to scale with number of distinct elementary tasks in the \"vocabularly\"? What about the number of compositional tasks?\n\n3. I think there is a highly relevant paper you do not cite: https://proceedings.neurips.cc/paper/2021/hash/fe5e7cb609bdbe6d62449d61849c38b0-Abstract.html\nHow does your probabilistic approach compare to their modular/gating-based technique?\nI do not refer only to empirical results -- but also conceptually and in terms of scalability.", "limitations": "For the most part the authors acknowledge the limitations of their study. One thing that is not clear however is the \"maximum cognitive capacity\" of their system -- in a scenario where the number of tasks keeps increasing with time while the memory/compute resources of the system are finite. How would the model behave in that case? Would it start forgetting in a \"graceful/gradual\" way (similar to an older person that gradually forgets tasks he/she has not practiced for a long time) or would it be something more unpredictable and abrupt?", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Quality:\n+  methodologically rigorous, offers clear theoretical formulation of context-modulated RNN architectures.\n+ demonstrates solid empirical validation across standard cognitive neuroscience tasks.\n- experiments primarily involve synthetic tasks -- questionable whether approach would work in more real-world tasks\n\nClarity\n+ clearly defines two-system approach (\"what\" vs \"how\") \n+ concepts & experiments are well-organized/presented\n\n\nSignificance:\n+ significant advance in CL by modeling computational epochs and compositionality\n+ illustrates backward and forward transfer, compositional generalization, task reuse -- goes further than prior work by Duncker et al.\n- it would be better if the tasks also included something more realistic, from real-world applications that require CL using RNNs\n\nOriginality\n- As expected, some conceptual overlap exists with previous works focusing on modular or context-dependent computations\n+ but the originality comes from interpretable probabilistic model to guide modular reuse of computations in RNNs", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their interest in our work and thoughtful comments. Concerns and questions raised in the review are addressed below.\n1. **More realistic tasks.** At the present stage, we have chosen to focus on synthetic problems common in neuroscience modeling work for interpretability and relevance to understanding continual learning in the brain. The set of cognitive tasks has been of recent interest in the field (e.g., see refs. 11, 13, 14, 16, 35), and we followed this literature to use them as a testbed for our proposed architecture. However, we wholeheartedly agree with the reviewer that it is important to evaluate our approach on realistic time-series tasks and that NLP tasks would be particularly well-suited to the kind of compositionality we are concerned with here.\n2. **Scalability** is indeed an important consideration for continual learning systems. In our case, letting $N_z$ denote the number of epochs (i.e., the ‘elementary tasks’), the probabilistic learning/inference algorithm (Sec. 3.3) used by the 'what' system would have time complexity $O(N_z^2)$ and space complexity $O(N_z)$, as is standard for learning/inference algorithms for HMM-like models. For the 'how' system, the main cost arises from running the RNN, which does not scale with $N_z$ (computing the effective weights in Eqs. 3-5 has cost linear in $N_z$). To improve scalability, a possible next step is to discourage the addition of new epochs by penalizing having a large $N_z$. Alternatively, it may also be possible to group epochs that are statistically distinct, but can share the same solution through another learning mechanism. This would be particularly relevant for the response epochs in our task set. It is also worth noting that, since each component of the network uses low-rank recurrent weights, the memory cost of adding a new component is linear in RNN size, as opposed to the quadratic scaling in vanilla RNNs. \n3. **Related work.** We thank the reviewer for pointing us to the highly relevant Ostapenko et al., 2021. Our scheme differs from theirs in several conceptually important ways, as listed below. We have added reference to this work and incorporated a condensed version of this discussion into the 'Related work and our contributions' section (Sec. 2).\n\t1. Ostapenko et al. requires passing every datum to all modules and running the computation in all modules first to determine which module outputs to keep (their Eq. 5). 'Module selection' in our model (computing the effective weights in Eqs. 3-5) is determined by a separate inference algorithm that is independent of the RNN modules. Our proposed scheme is more efficient in settings where the number of modules is large and the cost of running each module is high.\n\t2. Since Ostapenko et al. deals with static inputs and feedforward networks, different components in the same layer can be run in parallel. In our case, tasks have temporal structures and are solved by recurrent networks. Importantly, the composition of network weights changes during each trial dynamically. Since the hidden state activity $\\mathbf{h}_t$ is shared, different components thus interact with each other over time. Our approach allows for more flexible compositions of modules. For example, the same two modules can be used in the same trial in an alternating fashion.\n\t3. Finally, as the reviewer noted, our model provides an interpretable probabilistic description of the compositional structure of each task, which is more opaque in Ostapenko et al.\n4. **Graceful forgetting.** This is an interesting question and not something our current learning algorithm is adapted to. We assume that the number of contexts across tasks does not exceed the capacity of the network (i.e. number of contexts $\\times$ low-rank components per context $\\le$ total dimensionality of the RNN), such that overwriting is not an issue. However, it would be possible to implement a mechanism that decays the association between context (what) and recurrent weights (how) over time, if a given context is not visited for a long duration. This would effectively free up the low-rank component to be overwritten and associated with a new context, and enable graceful forgetting. While this exceeds the scope of the current paper, it would be an interesting topic for future work."}, {"review_id": "yO6MPlTFmS", "reviewer": "Reviewer_jVxy", "summary": "This manuscript proposed to use the two-stream (what and how) processing idea from the visual system as a way to efficiently solve the continual learning problem. The 'what' component is an RNN trained with a large family of tasks to obtain the compositional latents (vocabulary) of tasks for generalization. The 'how' component is a low-rank RNN conditioned on the context inferred from the 'what' system. The epoch variable z and trial variable x as the latent contexts for six cognitive task. The weights to each epoch are distinct, making the network more flexible for different epochs. The authors also proposed a task model and use EM algorithm to infer the latent contexts, guiding the parameter selection of the downstream low-rank RNN. In the comparison experiment with two traditional methods in continual learning (EWC and OWP), the proposed method shows better continual learning ability and exhibits positive forward/backward transfer effect. The reusability of the epoch context has also been verified. In summary, this article provides new insights for the continual learning of RNNs.", "questions": "Opportunities for improvement\n1. More comprehensive comparison. The strength of the 'epoch' concept could be further demonstrated by adding comparison experiments with reference [16], as well as with the case where the weights are modulated by task context c instead of epoch context z.\n2. More detailed explanation. Figure 5b shows that at the beginning of the training process of the second task, the test loss of the second task will first increase and then decrease. The authors are suggested to provide an explanation about that.\n3. Clarification: The image descriptions can be improved to clarify that the curves of different colors indicate the test performance on the corresponding tasks and the background color indicate the training process of the corresponding tasks. In appendix the stimulus angle in section A and the model parameters in section B are both denoted as theta, which may cause potentially confusion.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strength:\n1. Explicit compositional reuse. This paper separates the 'what' model to infer latent contexts rather than training the whole model in an end-to-end way, making  the compositional reuse of RNN components across contexts more explicit.\n2. Efficient network structure. The recurrent weights of RNN are represented as low-rank matrices, reducing the number of parameters. The proposed context-modulated RNN uses half as many trainable parameters to achieve comparative performance to traditional methods.\n3. Probabilistic perspective. This paper utilizes weights of different contexts from a probabilistic perspective, combining them in a linear manner weighted by the belief instead of picking up weights from one context, making the model more flexible.\n\nWeaknesses:\n1. Contradiction between the precision and reusability of the epoch. In this paper, on the one hand, the response period of the six tasks are all defined with distinct epochs in order to model the neural activities more accurately, sacrificing reusability of epoch. While on the other hand, defining stimulus periods of all tasks as the same epoch makes it difficult for the model to capture the different dynamics of the tasks during stimulus period, such as the evidence accumulation phenomenon.\n2. Discontinuous dynamics in a trial. By assigning different weights for different epoch, the dynamics of the RNN in a trial is separated by epochs, making it impossible to simulate dynamics such as sub-threshold evidence accumulation process. As a consequence the ability of the model is largely limited especially faced with complex tasks.\n3. Limited supported task types. According to the appendix, x determines the stimulus angle, a feature that carries consistent meaning across tasks. As a result, the method is constrained to tasks involving stimuli with angular information, thereby limiting its applicability to a broader range of task types. Additionally, Figure 2c exhibits multiple stimulus-response pairs under a single x within a task, which raises questions about the uniqueness and interpretability of the mapping.\n4. Practicality. The data for the six cognitive tasks is generated under relatively simple and controlled rules, which may differ from the distribution of real biological data. This raises concerns regarding the ability of the proposed method to capture the underlying neural dynamics present in real organisms. \n\nSignificance: The question of continual learning is an unresolved topic, leveraging the concept of compositionally and efficiency is of importance. However, the work itself still need more training dataset, testing and model comparison for evaluation.\nOriginality：The concept to separate into \"what\" and \"how\" models, and the weight composition in a probabilistic perspective are relatively novel.\nClarity：This article is well-written, with a complete introduction to relevant concepts and works and a detailed description about the proposed what and how models.\nQuality: The results are not comprehensive enough, lacking controls and oblation study. The proposed method was only compared with the traditional EWC and OWP methods, without comparing with more similar methods like the reference [16], making the comparison experiment less persuasive. Furthermore, the six cognitive tasks are most distinguishable during their response periods. However, the response periods of these tasks are all treated as distinct epochs, making the reusability of the epoch context seem less meaningful. In addition, the data used for the six tasks are simulated, raising concerns about the real-world applicability of the proposed method.", "was_revised": true, "final_justification": "I have updated the final rating to 4.", "rebuttal": "We appreciate that the reviewer thought that we made novel contributions to addressing an important open problem in the space of continual learning, and that the probabilistic treatment of tasks and context inference is a strength of the paper. The explicit suggestions for how to improve the paper and address the concerns are very helpful. In response, we have made the following changes:\n1. **Additional baselines.** We have now included a comparison against the NM-RNN [16], and (suggested by another reviewer) hypernetworks [23, 31]. Neither method matched the performance of ours in terms of continual learning or compositional generalization. Details and performance metrics are provided at the end of this reply. We have updated Sec. 2 to discuss these works further and also expanded the baseline comparisons (Fig. 4) to include these two.\n2. **Explanation for Fig. 5b.** The degree of loss increase varied across random seeds, and happened after the first step of gradient descent on the second task (our batch size equals 256). It could stem from the discrepancy between the first training batch and the test dataset, or that the model moved away from a local minimum of the second task's loss landscape.\n3. **Figure captions.** cf. 'The image descriptions … the corresponding tasks.' We have incorporated this clarification into our captions.\n\nOther concerns:\n\n4. **Contradiction between the precision and reusability of the epoch.** This point seems to refer to the problem of identifying the correct segmentation of tasks into epochs. As the reviewer noted, it is possible to share stimulus periods, and/or response periods depending on how exactly tasks are segmented. Indeed, for a subset of the tasks we consider here (Memory/Delay Pro/Anti) both segmentations are equivalent. The reviewer’s concern, to our understanding, is that we chose a particular segmentation to better model neural activity, at the sacrifice of re-usability of epochs across tasks. To clarify: we infer the segmentation based only on the statistical structure of the task, not neural activity. In addition to this, sharing of the stimulus epoch only happens for tasks that share the same statistical structure during the stimulus epoch, and therefore does not limit the system’s ability to capture different dynamics when this is necessitated by the demands of the task. For example, the stimulus periods between the DM tasks, and the Delay/Memory tasks are not inferred to be shared. More generally, the tension between oversegmentation and undersegmentation is an issue common to mixture models (e.g., clustering, hidden Markov models), and our approach is no exception. Further clarification on this point by the reviewer would be much appreciated.\n5. **Discontinuous dynamics in a trial.** This point is concerned with dynamics of the RNN and whether the model can capture 'dynamics such as sub-threshold evidence accumulation process'. While the inputs and target outputs are piecewise constant and sometimes discontinuous between epochs, the RNN hidden state (Eq. 1) is not constrained in these ways. The network can only respond correctly during response epochs using information contained in its hidden state. Within each epoch (for example during a stimulus epoch) it may exhibit ramping dynamics that are signatures of evidence accumulation. Our DM tasks are designed to emulate evidence accumulation: the network has to integrate noisy inputs to make a decision relating to the stimulus strength. Our choice of piecewise constant input plus noise follows classic models of evidence accumulation in the literature (e.g., Wang, 2002; Ratcliff et al., 2016). In summary, while the weights of the network switch across epochs, the network state evolves continuously and can exhibit complex activity patterns.\n6. **Limited supported task types.** The reviewer was concerned that our choice of $x$ constrains us to tasks involving stimuli with angular information. Our approach, however, does not require a consistent interpretation of $x$ across tasks. In fact, among the 6 tasks we studied, $x$ determines the stimulus angle in only 4 of them (Delay/Memory Pro/Anti). In the two DM tasks, $x$ controls the stimulus location/strength (appendix lines 94-100). Thus, our task set demonstrates the applicability to a broader range of tasks. Regarding the comment on Fig. 2c ('multiple stimulus-response pairs under a single x within a task'): while $x$ is a single discrete variable taking on a single value per trial of a given task, the stimulus ($\\mathbf{s}_t$) and response variable ($\\mathbf{y}_t$) can be higher-dimensional, as is the case in our example (both are 3D, corresponding to the three lines in the figure). We apologize for the confusion and have updated the figure to clarify this.\n7. **Practicality.** This point is concerned with the generation of the cognitive tasks and whether they can match the distribution of real biological data. While we are excited about connecting our model to neural data in the future, we want to clarify that at the moment this work is not concerned with modeling biological neural activity at all. The synthetic data, represented by $\\mathbf{q}_t$, are meant to model the behavior-level input (stimuli on a screen) and output (motor actions, eye saccades) in neuroscience experiments. The simplified dynamics, where the inputs are piecewise constant, are consistent with typical experimental setups and are common assumptions in related modeling work (see refs. on line 108). \n8. The reviewer suggested comparing to the setting where weights are modulated by task context c instead of epoch context z. This would effectively train a separate network per task, and therefore sidestep both the continual learning problem and any compositional reuse across tasks. This would make forward/backward transfer impossible. We are happy to include these comparisons, but would appreciate further clarification on what the reviewer would like us to show.\n\n**References:** \n\nWang, X. J. (2002). Neuron, 36(5), 955-968.\nRatcliff, R., ..., & McKoon, G. (2016). Trends in cognitive sciences, 20(4), 260-281.\n***\n**Summary of additional baselines.** Costacurta et al., 2024 [16] designed neuromodulated RNN (NM-RNN), in which a full-rank neuromodulatory RNN (RNN1) outputs a low-D signal that dynamically scales the low-rank recurrent weights of an output-generating RNN (RNN2). RNN1 and RNN2 are jointly optimized through gradient descent. Their design showed catastrophic forgetting during continual learning (Table 1). In contrast, our method replaces RNN1 with a probabilistic task model, which rapidly learns a task’s structure within a handful of trials, thus is able to provide useful contextual signals when optimizing RNN2 and facilitate continual learning. Von Oswald et al., 2020 [23] designed task-conditioned hypernetworks to tackle continual learning. Their hypernetwork receives a learnable embedding distinct for each task and outputs the weights for the target network (a RNN in our case). Their design underperformed our method (Table 1). More importantly, these hypernetworks do not capture the time-varying epoch structures of each task, and perform worse on compositional generalization (Table 2).\n\n**Implementation details.** \n\n**NM-RNN**:  RNN2 has 256 neurons and with rank($W^{rec}$) = 27, equal to the default choice for our ContextRNN. RNN1 has 125 neurons such that the whole network has a similar number of parameters as ours.\n\n**Hypernetwork**:  The full-rank RNN has 256 neurons. We used a chunked hypernetwork with 34 chunks and 2000-D output. The hypernetwork is a MLP with two 32-D hidden layers, receiving 32-D task embedding and chunk embeddings as input. After a grid search, we set the continual learning regularization strength to 1, orthogonal regularization strength to 0.001 and clip the gradient norm to 1 (see Ehret et al., 2021 for details). \n\nTable 1:\n\n**A. NM-RNN**\n\n|| task 1| task 2| task 3| task 4| task 5| task 6|\n|-|-|-|-|-|-|-|\n| final perf.| 0.00 (± 0.00)| 0.25 (± 0.06)| 0.00 (± 0.00)| 0.15 (± 0.05)| 0.00 (± 0.00)| 1.00 (± 0.00)|\n\n**B. Hypernetwork**\n\n|| task 1| task 2| task 3| task 4| task 5| task 6|\n|-|-|-|-|-|-|-|\n| final perf.| 0.99 (± 0.01)| 0.89 (± 0.18)| 0.79 (± 0.06)| 0.67 (± 0.11)| 0.80 (± 0.15)| 0.87 (± 0.06)|\n\n**C. ContextRNN (ours)**\n\n|| task 1| task 2| task 3| task 4| task 5| task 6|\n|-|-|-|-|-|-|-|\n| final perf.| 0.99 (± 0.02)| 0.99 (± 0.02)| 0.91 (± 0.04)| 0.93 (± 0.04)| 1.00 (± 0.00)| 1.00 (± 0.00)|\n\nCaption: Continual learning performance (supplementary to Figure 4 a-d). A: Each column shows a task's final performance averaged across 5 random seeds (± std) after training NM-RNN sequentially on task 1 to task 6. B: same as A but for hypernetwork. C: same as A but for contextRNN (ours).\n\nTable 2:\n\n**A. OWP**\n\n|                  | after 0 trials | after 20 trials | after 40 trials | after 512 trials |\n| ---------------- | -------------- | --------------- | --------------- | ---------------- |\n| MemoryAnti perf. | 0.00 (± 0.00)  | 0.00 (± 0.01)   | 0.01 (± 0.02)   | 0.53 (± 0.09)    |\n**B. Hypernetwork**\n\n|                  | after 0 trials | after 20 trials | after 40 trials | after 512 trials |\n| ---------------- | -------------- | --------------- | --------------- | ---------------- |\n| MemoryAnti perf. | 0.00 (± 0.00)  | 0.03 (± 0.02)   | 0.19 (± 0.05)   | 0.64 (± 0.03)    |\n**C. ContextRNN (ours)**\n\n|                  | after 0 trials | after 20 trials | after 40 trials | after 512 trials |\n| - | - | - | - | ---------------- |\n| MemoryAnti perf. | 0.00 (± 0.00)  | 0.56 (± 0.06)   | 0.83 (± 0.02)   | 0.83 (± 0.04)    |\n\nCaption: Compositional generalization (supplementary to Figure 6c). A: Each column shows the test performance on MemoryAnti averaged across five random seeds (± std) after training a pretrained RNN using OWP on MemoryAnti for 0, 20, 40 or 512 trials. B: same as A but for hypernetwork. C: same as A but for contextRNN (ours)."}, {"review_id": "n167zfW5qJ", "reviewer": "Reviewer_N8zU", "summary": "This paper investigates how neural systems can continually learn and reuse compositional skills to solve a variety of tasks, especially when those tasks share common structural elements but differ in specific details. The major challenge is to avoid catastrophic forgetting and to enable generalization by recombining learned components. To address this, the authors introduce a two-system framework that explicitly separates \"what\" (context inference) from \"how\" (computation implementation). The what system employs a probabilistic generative model to segment tasks into discrete epochs and learns to infer these contexts online, while the how system use an RNN whose weights are dynamically modulated via low-rank components according to the inferred context from the what system. Experiments on a set of neuroscience-inspired tasks demonstrate that this approach effectively mitigates catastrophic forgetting and supports compositional generalization.", "questions": "Please refer to earlier sections.\nMinor question: Is it possible to replace the RNN in the “how” component with other types of models? For example, could the proposed framework work equally well if the RNN were substituted with a feedforward neural network, a transformer, or another architecture more suitable for the task at hand? It would be helpful if the authors could clarify whether their approach is generally applicable to other model classes beyond RNNs, and discuss any potential challenges or limitations in doing so.", "limitations": "As mentioned in the weaknesses above, the paper would be strengthened by (1) evaluating its applicability to domains without clear compositional task structure, (2) providing ablation studies on the sensitivity to context inference design and ambiguous context boundaries, (3) more thoroughly distinguishing its approach from related work in meta-learning, hypernetworks, and neuromodulation, (4) including a broader set of baselines, particularly recent modular and task-inference-based continual learning methods, and (5) improving the clarity of Section 3 for readers.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths\n1. Comprehensive experiments: The experiments are thorough and cover catastrophic forgetting, transfer, and compositional generalization.\n2. Originality: While related ideas exist, the explicit and interpretable separation of \"what\" and \"how\" in continual learning RNNs is new and well-justified.\n\nWeaknesses\n1. Assumptions on task structure: The probabilistic model assumes that tasks can be decomposed into discrete, shared epochs. This may limit applicability to domains where such structure is not present or easily defined.\n2. Ablation on module interdependence: There is little analysis of how sensitive the approach is to the precise design of the context inference model, or what happens if context boundaries are ambiguous or misspecified.\n3. Related work overlap: While the two-system approach is novel, there is some overlap with prior work in meta-learning, hypernetworks, and neuromodulation; the paper could do more to clarify distinctions and advantages over these related methods.\n4. Limited baselines: The experiments do not include comparisons to more recent or diverse modular, meta-learning, or task-inference-based continual learning frameworks (e.g. Costacurta et al. Structured flexibility in recurrent neural networks via neuromodulation. Advances in Neural Information Processing Systems, 2024). This limitation makes it difficult to fully assess the relative advantages of the proposed approach.\n5. Clarity: Section 3 is dense and may be challenging for readers unfamiliar with probabilistic modeling.", "was_revised": true, "final_justification": "Most of my concerns have been addressed through the addition of extended baseline comparisons and improved discussion of assumptions and limitations. Therefore, I increased my score.", "rebuttal": "We thank the reviewer for their interest in our work and thoughtful comments, and appreciate that the reviewer thought our submission makes a novel contribution that was thoroughly demonstrated through experiments and evaluations. The explicit suggestions for how to improve the paper were very helpful, and we are making the following changes in response to them:\n1. **Additional baselines & more thorough discussion of related work.** We agree with the reviewer that additional baseline comparisons are necessary and have added two additional baselines in the manuscript. The first one is against the recent Costacurta et al., 2024 [16] neuromodulatory network and a second is against a hypernetwork approach (von Oswald et al., 2020, [23]). In short, we found that neither method matched the performance of ours in terms of continual learning or compositional generalization. Details and performance metrics are provided at the end of this reply. We have updated the 'Related work and our contributions' section to discuss these works further and also expanded the baseline comparisons (Fig. 4) to include these two. We hope that these changes will address Limitations (3) and (4).\n2. **Improved clarity of Sec. 3.** We revised the writing of Sec. 3 to make it more accessible to readers unfamiliar with probabilistic modeling. In particular, we made sure to highlight conceptual features of the algorithm first, before introducing formal notation. We also stressed that the model is given the correct task label $c$ and only performs the epoch inference $z$. If there are parts of Sec. 3 that the reviewer found to be particularly inaccessible to the reader, we would appreciate more detail e.g. line numbers, so we can make sure to improve their clarity. We hope that these changes will address Limitation (5).\n\nWe respond to the other concerns below.\n\n3. **Assumptions on task structure.** This concern is related to the assumptions about task structures that we make in the paper, and raises the question of whether our approach is applicable to domains without clear compositional task structures. At the present stage, we have chosen to focus on synthetic problems common in neuroscience modeling work for interpretability and relevance to understanding continual learning in the brain. As stated in line 107-108, we decided to focus on common experimental paradigms, which are often modeled as having clear boundaries between epochs (see refs. on line 108). In related literature (e.g. [11], [13]-[16]), compositional reuse in multi-tasking RNNs has been of major interest. Our work is intended to contribute to this space, making the compositional task structure, and how it may shape a potential solution in the RNN, an explicit focus. Still, we wholeheartedly agree with the reviewer that it is important to evaluate our approach on more realistic tasks, spanning a wider domain. Analyzing the approach in settings with more ambiguous compositional structures will be particularly interesting. While we argue that this is beyond the scope of the current paper, we would like to highlight that the probabilistic treatment of the 'what' system may be especially well-equipped to handle these settings.\n4. **Ablations and ambiguous context boundaries.** The 'what' system is designed to perform posterior inference of context, conditioned on the observables (inputs $\\mathbf{s}_t$ and targets $\\mathbf{y}_t$) of a given task, and, in principle, any architecture that can do this could be used as the context inference model. In our setting,  ambiguity across epoch boundaries stems only from the temporal variability of each trial and observation noise. We agree that it would be worthwhile to further explore the setting where context boundaries are ambiguous due to other sources. For instance, this could happen in our tasks when stimuli change more slowly (e.g. when they are supplied as a ramping signal rather than turning on/off abruptly). Studying this task variation would test how robust our inference system is to a mismatch between the underlying generative assumptions, and the true task design, which is important. Given the uncertainty about context that would arise at these transition points, the 'how' system would express a mixture of weights rather than a single solution. We would expect that this would impact the task performance, and might encourage (or require) learning some form of normalization across weights. Currently, the system only experiences context-uncertainty during learning. We are planning to update the discussion section to include these points. Regarding the suggested ablation studies, we would appreciate additional detail on what exactly the reviewer had in mind.\n5. **Replacing the RNN with other models.** The general framework of using probabilistic inference to 'decompose' time-series tasks is agnostic to the 'how' module used. Indeed we imagine a similar scheme to work for other architectures that deal with time series, such as transformers. For example, with transformers the equations in Eq. 3-5 would be replaced by analogous equations that determine the query, key, value weights used for input from different epochs. Incorporating the two-systems approach we develop here into more complex architectures represents an interesting direction of future work and may also facilitate further applications to more complex tasks.\n***\n**Summary of additional baselines.** Costacurta et al., 2024 [16] designed neuromodulated RNN (NM-RNN), in which a full-rank neuromodulatory RNN (RNN1) outputs a low-D signal that dynamically scales the low-rank recurrent weights of an output-generating RNN (RNN2). RNN1 and RNN2 are jointly optimized through gradient descent. Their design showed catastrophic forgetting during continual learning (Table 1). In contrast, our method replaces RNN1 with a probabilistic task model, which rapidly learns a task’s structure within a handful of trials, thus is able to provide useful contextual signals when optimizing RNN2 and facilitate continual learning. Von Oswald et al., 2020 [23] designed task-conditioned hypernetworks to tackle continual learning. Their hypernetwork receives a learnable embedding distinct for each task and outputs the weights for the target network (a RNN in our case). Their design underperformed our method (Table 1). More importantly, these hypernetworks do not capture the time-varying epoch structures of each task, and performed worse on compositional generalization (Table 2).\n\n**Implementation details.** \n\n**NM-RNN**:  RNN2 has 256 neurons and with rank($W^{rec}$) = 27, equal to the default choice for our ContextRNN. RNN1 has 125 neurons such that the whole network has a similar number of parameters as ours.\n\n**Hypernetwork**:  The full-rank RNN has 256 neurons. We used a chunked hypernetwork with 34 chunks and 2000-D output. The hypernetwork is a MLP with two 32-D hidden layers, receiving 32-D task embedding and chunk embeddings as input. After a grid search, we set the continual learning regularization strength to 1, orthogonal regularization strength to 0.001 and clip the gradient norm to 1 (see Ehret et al., 2021 for details). \n\nTable 1:\n\n**A. NM-RNN**\n\n|| task 1| task 2| task 3| task 4| task 5| task 6|\n|-|-|-|-|-|-|-|\n| final perf.| 0.00 (± 0.00)| 0.25 (± 0.06)| 0.00 (± 0.00)| 0.15 (± 0.05)| 0.00 (± 0.00)| 1.00 (± 0.00)|\n\n**B. Hypernetwork**\n\n|| task 1| task 2| task 3| task 4| task 5| task 6|\n|-|-|-|-|-|-|-|\n| final perf.| 0.99 (± 0.01)| 0.89 (± 0.18)| 0.79 (± 0.06)| 0.67 (± 0.11)| 0.80 (± 0.15)| 0.87 (± 0.06)|\n\n**C. ContextRNN (ours)**\n\n|| task 1| task 2| task 3| task 4| task 5| task 6|\n|-|-|-|-|-|-|-|\n| final perf.| 0.99 (± 0.02)| 0.99 (± 0.02)| 0.91 (± 0.04)| 0.93 (± 0.04)| 1.00 (± 0.00)| 1.00 (± 0.00)|\n\nCaption: Continual learning performance (supplementary to Figure 4 a-d). A: Each column shows a task's final performance averaged across 5 random seeds (± std) after training NM-RNN sequentially on task 1 to task 6. B: same as A but for hypernetwork. C: same as A but for contextRNN (ours).\n\nTable 2:\n\n**A. OWP**\n\n|                  | after 0 trials | after 20 trials | after 40 trials | after 512 trials |\n| ---------------- | -------------- | --------------- | --------------- | ---------------- |\n| MemoryAnti perf. | 0.00 (± 0.00)  | 0.00 (± 0.01)   | 0.01 (± 0.02)   | 0.53 (± 0.09)    |\n**B. Hypernetwork**\n\n|                  | after 0 trials | after 20 trials | after 40 trials | after 512 trials |\n| ---------------- | -------------- | --------------- | --------------- | ---------------- |\n| MemoryAnti perf. | 0.00 (± 0.00)  | 0.03 (± 0.02)   | 0.19 (± 0.05)   | 0.64 (± 0.03)    |\n**C. ContextRNN (ours)**\n\n|                  | after 0 trials | after 20 trials | after 40 trials | after 512 trials |\n| - | - | - | - | - |\n| MemoryAnti perf. | 0.00 (± 0.00)  | 0.56 (± 0.06)   | 0.83 (± 0.02)   | 0.83 (± 0.04)    |\n\nCaption: Compositional generalization (supplementary to Figure 6c). A: Each column shows the test performance on MemoryAnti averaged across five random seeds (± std) after training a pretrained RNN using OWP on MemoryAnti for 0, 20, 40 or 512 trials. B: same as A but for hypernetwork. C: same as A but for contextRNN (ours)."}, {"review_id": "KpvcDcaiz9", "reviewer": "Reviewer_yJUA", "summary": "The paper suggests a two-system model for continual learning of compositional tasks using RNNs.\nIn that, the what-system identifies the current context based on stimulus-response statistics\nand the how-system learns to perform the individual task.\nThe how-system is implemented as an RNN and is modified by the what-system through re-weighting of its low-rank components.\nBoth systems are trained in parallel and online, such that the tasks are inferred while already learning in task.\nThe authors test their model on 6 related neuroscience tasks which involve matching or mirroring of a stimulus angle, memorising it and making decisions.\nThe tasks are formalised in a way that they consist of several epochs like fixation, stimulus, delay and response that can be inferred and composed.\nEmpirically their ContextRNN method outperforms other multi-task learning techniques (Elastic Weight Consolidation and Orthogonal Weight-Space Projection by Duncker 2020) in that it shows quicker context adaptation and less catastrophic forgetting.\nFurthermore, they demonstrate that their model is able to do forward and backward transfer to related tasks as well as a simple example of compositional generalisation.", "questions": "1. How many low-rank components did you take, just 3? How would this affect the capacity of the RNN, in terms of how many tasks it can solve?\n2. How many latents did the task-recognition network learn? Did that correspond to the number of epochs as defined in the task description?\n3. How would your method compare to hyper-networks as used by Von Oswald et al.? Wouldn't this be a more fair comparison than EWC / OWP, as hypernetworks also separate the what and how?\n3. How could your method be applied to settings with less structured tasks? Here you explicitly use tasks that are composed by clear epochs which are temporally separated. In the discussion you hint at non-parametric methods to find epochs. Could this also be trained in parallel to the how-RNN? (One would probably have to train in an interleaved instead of blocked-task regime then to learn what is similar between tasks.)", "limitations": "Yes", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "- A strength of the paper is definitely the separation of task-recognition (what) and task-performance (how), which is not as clear-cut in many models. Most approaches use a single network to learn, such as EWC and OWP and therefore are less flexible.\n-  Another strength is the method's ability to compose new tasks from similar known tasks. This is a benefit that also comes from the separate what-system, that recognises components of the tasks.\n-  In my opinion, the choice of tasks is limited in that they are quite simple input-output mappings.\nEspecially the Delay and the Memory condition only differ in the duration of the delay but not in the input-output mapping (within Pro/Anti). This makes transfer from the Memory to the Delay domain very simple (as the RNN can use the same dynamics, just shorter). However, as one can see from Fig. S.2, the generalisation works less well for the other direction (again showing that the memory condition is harder).\n- Clarity could be improved regarding the description of the what system. It took me a while to understand that it does not infer the context / task $c$ itself, but instead is given the context label c explicitly and only tries to infer the task component / epoch z. This is confusing in the introduction (e.g. line 40 \"separation into contexts (what)\") and only clarified later. Instead it might be called the \"task decomposition\" network.\n-  Another point that needs explanation is the difference between the Delay conditions and the M' conditions.\n-  They only seem to differ in the final epoch, which is just R in the Delay condition but R_M in the M' condition, although the M' is explicitly constructed without the memory epoch.\n-  A minor weakness is the framing of the approach as brain-like. I agree that the brain is probably using a similiar separation of what and how. However, here the RNN needs to be decomposed into low-rank components, which is not a very neurally plausible operation (or could it be implemented in a more plausible way?) -- this caveat should be added to the discussion part that relates the method to multi-task learning in the brain.", "was_revised": true, "final_justification": "The authors have addressed most of my concerns, clarifying distinctions between tasks, the meaning of R, and R_M, and the conceptual (rather than mechanistic) nature of the biological plausibility claim. They provided helpful explanations of the low-rank decomposition, added relevant baselines, and improved clarity on “inference of tasks.” Though the small and closely related task set limits the scope of the findings, the compositional RNN approach to modeling cognitive tasks remains novel and serves as a proof of concept. Overall, the work is promising and invites future extensions to more diverse task compositions. I therefore decided to keep my score at 5.", "rebuttal": "We thank the reviewer for their interest in our work and thoughtful comments. Below, we attempt to address the raised concerns, answer questions and highlight changes we have made in the revision in response to the reviews.\n1. **Limited task set.** We agree that extending to a richer task set with more complex input-output mappings, with a more thorough assessment of the potential for (and directionality of) transfer learning and compositionality, is important. At the present stage, we focus on synthetic problems common in neuroscience models for interpretability and relevance for understanding Continual Learning and compositional re-use in RNNs. In related work studying similarly simple tasks (e.g. refs. 11, 13-16), compositional reuse in multi-tasking RNNs has been of major interest. Our work is intended to contribute to this space, making the compositional task structure, and how it may shape a potential RNN solution, an explicit focus. We note that our approach for modeling the statistical structure of cognitive tasks in a cohesive generative framework is novel. We expect that this contribution will help to establish a more formal notion of task-similarity and when/how transfer might be simple or more difficult. We fully agree that a richer task set will lead to a variety of interesting questions for further work.\n2. **Writing in Sec. 3.** We thank the reviewers for suggestions about improving the readability of Sec. 3, which was a concern also shared by other reviewers. We revised the writing of Sec. 3 to make it more accessible to readers unfamiliar with probabilistic modeling. In particular, we made sure to highlight conceptual features of the algorithm first, before introducing formal notation. We also stressed that the model is given the correct task label $c$ and only performs the epoch inference $z$. If there are additional parts of Sec. 3 that the reviewer found to be particularly inaccessible to the reader, we would appreciate more detail (e.g. line numbers), so we can make sure to improve their clarity.\n3. **Differentiating Memory and Delay tasks.** We have updated the wording to clarify that there are two differences between these sets of tasks. First, the Memory tasks have a fixation epoch after stimulus presentation. Second, the correct stimulus is present during the response epoch for Delay tasks but not during the response epoch for Memory tasks (see Fig. 1a). We have added details on this to section 4.4.\n4. **Biological plausibility.** We agree that the current version of the architecture lacks biological realism and is only brain-like at a conceptual level. We are very interested in exploring more biologically plausible architectures that could implement the modulation of low-rank components of the RNN. [16] connected a similar selection of low-rank components to neuromodulation. [37] implemented the selection of individually controlled low-rank components through an architecture representing interactions between basal ganglia, thalamus and cortex. This model assumes that the low-rank connectivity of the RNN is the result of cortex-to-thalamus and thalamus-to-cortex connections (= thalamo-cortical ‘loop’). When the number of thalamic neurons participating in each loop is small, the connectivity is low rank. Different loops can be selected via inhibitory input from basal ganglia to thalamus, which effectively turns on and off different low-rank components of the RNN – just like our architecture. Contrary to our work, [37] was not concerned with how this inhibitory selection signal is set.\n5. **Low-rank components (Q1).** The number of low-rank weight matrices (a set of low-rank components) equals the number of unique task epochs. For our current task set, the inference algorithm identifies 9 unique task epochs in total. Each task epoch uses a rank-3 recurrent weight matrix in our main figures (defined by $U_z, V_z$, Appendix C.1), though performance increased for higher ranks (Figure S5). For task-epochs solved by an at most rank-$R$ weight matrix, we may implement $N_z = N/R$ epochs for a fixed network dimensionality $N$. The network can solve tasks that can be composed as sequences of the $N_z$ epochs, which could be a large number. It is worth noting that, since each component of the network uses low-rank recurrent weights, the memory cost of adding a new component is linear in RNN size, as opposed to the quadratic scaling in vanilla RNNs. To improve scalability, a possible next step is to discourage the addition of new epochs by penalizing having a large $N_z$. Alternatively, it may also be possible to group epochs that are statistically distinct, but can share the same solution through another learning mechanism. This would be particularly relevant for the response epochs in our task set.\n6. **Number of learned latents (Q2).** Across the 6 tasks we studied, 10 epochs are defined as part of the generative process of the tasks (listed in line 93 of the appendix). The ‘what’ system learns this generative model based on the inputs and targets of each task, and identifies only 9 epochs. This is expected, as the fixation epoch $F$ and memory epoch $M$ are indistinguishable based on their input-output patterns. \n7. **Additional baselines (Q3).** We agree with the reviewer that a comparison against the hypernetwork approach in von Oswald et al. [23] and Ehret et al. [31] is warranted. We have performed these comparisons and are planning to include them in the revised manuscript. Per the suggestion of other reviewers, we also performed a comparison against the neuromodulatory network in Costacorta et al., 2024 [16]. In short, we found that neither method matched the performance of ours in terms of continual learning or compositional generalization. Details and performance metrics are provided at the end of this reply. \n8. **Less structured tasks (Q4).** Handling tasks with less apparent structures is an important next step. Training the 'what' system together with the 'how' RNN is indeed a promising approach. In the case of ambiguous structures, it may be helpful to use the hidden state of the 'how' RNN as part of the observation $\\mathbf{q}_t$ to infer epochs. This is similar in spirit to the scheme in Ostapenko et al., 2021 for modules in a feedforward network, where the output of each module is used to determine its relevance for a given input.\n***\n**Summary of additional baselines.** Costacurta et al., 2024 [16] designed neuromodulated RNN (NM-RNN), in which a full-rank modulatory RNN (RNN1) outputs a low-D signal that dynamically scales the low-rank recurrent weights of an output-generating RNN (RNN2). RNN1 and RNN2 are jointly optimized through gradient descent. Their design showed catastrophic forgetting during continual learning (Table 1). In contrast, our method replaces RNN1 with a probabilistic task model, which rapidly learns a task’s structure within a handful of trials, thus is able to provide useful contextual signals when optimizing RNN2 and facilitate continual learning. Von Oswald et al., 2020 [23] designed task-conditioned hypernetworks to tackle continual learning. Their hypernetwork receives a learnable embedding distinct for each task and outputs the weights for the target network (a RNN in our case). Their design underperformed our method (Table 1). More importantly, these hypernetworks do not capture the time-varying epoch structures of each task, and perform worse on compositional generalization (Table 2).\n\n**Implementation details.** \n\n**NM-RNN**:  RNN2 has 256 neurons and with rank($W^{rec}$) = 27, equal to the default choice for our ContextRNN. RNN1 has 125 neurons such that the whole network has a similar number of parameters as ours.\n\n**Hypernetwork**:  The full-rank RNN has 256 neurons. We used a chunked hypernetwork with 34 chunks and 2000-D output. The hypernetwork is a MLP with two 32-D hidden layers, receiving 32-D task embedding and chunk embeddings as input. After a grid search, we set the continual learning regularization strength to 1, orthogonal regularization strength to 0.001 and clip the gradient norm to 1 (see Ehret et al., 2021). \n\nTable 1:\n\n**A. NM-RNN**\n\n| | task 1| task 2| task 3| task 4| task 5| task 6|\n|-|-|-|-|-|-|-|\n| final perf.| 0.00 (± 0.00)| 0.25 (± 0.06)| 0.00 (± 0.00)| 0.15 (± 0.05)| 0.00 (± 0.00)| 1.00 (± 0.00)|\n\n**B. Hypernetwork**\n\n| | task 1| task 2| task 3| task 4| task 5| task 6|\n|-|-|-|-|-|-|-|\n| final perf.| 0.99 (± 0.01)| 0.89 (± 0.18)| 0.79 (± 0.06)| 0.67 (± 0.11)| 0.80 (± 0.15)| 0.87 (± 0.06)|\n\n**C. ContextRNN (ours)**\n\n| | task 1| task 2| task 3| task 4| task 5| task 6|\n|-|-|-|-|-|-|-|\n| final perf.| 0.99 (± 0.02)| 0.99 (± 0.02)| 0.91 (± 0.04)| 0.93 (± 0.04)| 1.00 (± 0.00)| 1.00 (± 0.00)|\n\nCaption: Continual learning performance (supplementary to Figure 4 a-d). A: Each column shows a task's final performance averaged across 5 random seeds (± std) after training NM-RNN sequentially on task 1 to 6. B: same as A but for hypernetwork. C: same as A but for contextRNN (ours).\n\nTable 2:\n\n**A. OWP**\n\n| | after 0 trials | after 20 trials | after 40 trials | after 512 trials |\n| - | - | - | - | - |\n| MemoryAnti perf. | 0.00 (± 0.00)  | 0.00 (± 0.01)   | 0.01 (± 0.02)   | 0.53 (± 0.09)    |\n**B. Hypernetwork**\n\n| | after 0 trials | after 20 trials | after 40 trials | after 512 trials |\n| - | - | - | - | - |\n| MemoryAnti perf. | 0.00 (± 0.00)  | 0.03 (± 0.02)   | 0.19 (± 0.05)   | 0.64 (± 0.03)    |\n**C. ContextRNN (ours)**\n\n| | after 0 trials | after 20 trials | after 40 trials | after 512 trials |\n| - | - | - | - | - |\n| MemoryAnti perf. | 0.00 (± 0.00)  | 0.56 (± 0.06)   | 0.83 (± 0.02)   | 0.83 (± 0.04)    |\n\nCaption: Compositional generalization (supplementary to Figure 6c). A: Each column shows the test performance on MemoryAnti averaged across five random seeds (± std) after training a pretrained RNN using OWP on MemoryAnti for 0, 20, 40 or 512 trials. B: same as A but for hypernetwork. C: same as A but for contextRNN (ours)."}]

Author rebuttal:

Meta-review: The paper proposes an interesting method for solving tasks through compositional computation. It is based on building sequentially a vocabulary of tasks.  Then it utilizes a two-component model based on the so called "what-system" that identifies the current task context and the "how-system" that learns to perform the specific task. The whole system is based on RNNs and it is trained online. The authors demonstrate their ContextRNN method on neuroscience tasks and they show that it can outperform other techniques,  such as Elastic Weight Consolidation and Orthogonal Weight-Space Projection. More interestingly, the model can do forward and backward transfer to related tasks as well as compositional generalization.

All reviewers found the separation of "what" and "how" in continual learning of RNNs a very novel and interesting idea.
Overall, this is a strong submission with a clear and nice proposal as well as convincing experimental results.
Decision: Accept (poster)