Title: Separating the 'what' and 'how' of compositional computation to enable reuse and continual learning
Venue: NeurIPS 2025
Accepted: yes
Keywords: continual learning

compositionality

contextual inference

low-rank RNN
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

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-)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. 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 timevarying 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. Contextual inference facilitates the creation, learning, and reuse of 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 fast compositional generalization to unseen tasks.

1 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 [1-9]. An abstract, task-relevant context could guide the selection and composition of different skills, while the maintenance of context-specific

memories could counteract forgetting and aid continual learning [1]. 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 a two-systems approach: one that infers an abstract computational 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 in 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 [11]. Here, we expand upon previous work in RNNs, where continual learning is still a major challenge [12], and compositional reuse of computational building blocks has only been shown to emerge implicitly over simultaneous training on multiple related tasks [13-16], 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 a generative model for cognitive tasks. This allows us to define a computational context within a 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 a 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 a few examples - a property called compositional generalization [17-19]. Finally, we discuss limitations and future directions in section 5.

Our work contributes to a 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 [1, 20-22]. Contextual inference has been proposed as a model explaining many learning-related phenomena in classical conditioning [20] and motor control [1]. Here, a latent discrete-valued random variable is used to describe distributional changes in taskrelevant, continuous observations. Given a set of observations, this variable can be inferred as a familiar or novel context to determine whether a 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 [11, 14]. 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 a different computation on a particular task epoch. A key aspect of our model is that these contexts can reappear in a 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 a response based on information presented at an earlier time in the trial [11, 14].

Task-dependent modulation of neural networks. The general notion of a separate system modulating computation in a neural network in a task-dependent manner has been previously studied. An example of this is the Task-Conditioned Hypernetwork [23], where the weights of a neural network are generated by a second, separate network based on modifiable task embeddings. Other approaches often assume the network has its own learnable weights but that their involvement in computation

pro

anti

response

stimulus

time

stimulus-guided response

Delay Pro/Anti

fixation

pro

anti

response

delay

stimulus

time

memory-guided response

Memory Pro/Anti

fixation

pro

anti

response

decision

time

decision-guided response

DM Pro/Anti

fixation

{

. . . }

,

,

,

,

,

inputs and target

responses, given

task epochs z

and trial variable x

z = 1

all x

z = 2

x = 1

z = 3

all x

z = 4

x = 1

z = 5

x = 3

a

z1

z2

z3

. . .

Λ c

q1

q2

q3

c

x

task id

epoch id

inputs,

targets

trial variable

b

trial variable (x)

determines the

stimulus

x = 1

x = 2

x = 3

. . .

is modulated, for instance via a gating mechanism [16, 24, 25]. These works typically study tasks that can be solved by a single subcomponent (e.g., [25-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 is similar to that of Costacurta et al. [16], where an RNN receives a low-dimensional, time-varying 'neuromodulatory' signal from a second RNN that adjusts how different low-rank components of its recurrent weights are combined. While Costacurta et al. [16] train this architecture end-to-end with backpropagation, we use the interpretable output of a 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 to use and does not make use of an probabilistic inference procedure.

Continual learning in RNNs. Continual learning is relatively understudied in RNNs and sequenceto-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) [31], which penalizes weight changes during future learning according to an estimate of how important each individual weight is for the performance of previous tasks. EWC is representative of a family of weight-regularization-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-35]. 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 it may be inferred as part of RNN training. We do this for a family of compositionally related cognitive tasks which have been popular both in experimental and computational neuroscience [11, 13, 14, 16, 36].

In these tasks, the learner receives a sequence of time-varying inputs s t =1 ,...,T and needs to produce a target response sequence y t =1 ,...,T , following different distributions p ( s 1: T , y 1: T | c ) depending on the task identity c ∈ { 1 , . . . , N c } . These sequences are structured according to particular segments (often informally referred to as 'epochs') such as a fixation period, a stimulus presentation period, or a response period. Each of these epochs defines particular distributions of inputs and target responses

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 task-dependent transitions from one epoch to the next. Fig. 1a 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. 1b). We can describe complex dependencies in the distribution of inputs and target responses p ( s 1: T , y 1: T | c ) for a task by introducing a discrete latent variable z t ∈ Z + denoting the task epoch at time t . We model epochs as evolving over time with task-dependent Markovian transitions p ( z t = i | z t -1 = j, c ) = Λ c ij . Inputs and target responses q t ≡ [ s t , y t ] ∈ R D q are drawn from an epoch-dependent distribution. Each epoch may involve a number of different input values, reflecting e.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 q t ∼ p ( q t | z t , x ) (Fig 1a). For example, x may index a stimulus, which differs from one trial to the next and determines the distribution of inputs and target responses (see Figure 2c for examples). Importantly, sharing the same x 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 is naturally defined as the time-varying sequence of task epochs z t . Thus, performing contextual inference requires learning the statistical structure and dependencies of the underlying generative model, and inferring the latent task epoch from observations q t . 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.1.

3.2 Overview of our task set

We initially focus on six commonly studied neuroscience tasks [11, 13, 14, 16] that can be implemented in our generative framework. The tasks are schematized in Fig. 1a 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 temporal structure involving e.g. delayed responses. We summarize the different tasks in terms of their general goal and epoch structure in Table 1. 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 x , and additional variability in q t = [ s t , y t ] . All tasks involve a 5-dimensional time series of inputs s t and a 3-dimensional time series of target responses y t , which take on noisy values around a piecewise constant mean. The value of the mean is dependent on the task epoch z and trial variable x , chosen to reflect the underlying task goals. We refer to Appendix A.2 for a more detailed description of the tasks, their epoch structures, and distributional choices.

Note that inclusion of the trial variable x has important consequences for how tasks are segmented into epochs. Considering the MemoryPro/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' ( R M,P/A ), the target response is along/opposite the stimulus direction, thus making their dependencies on x different. This necessitates treating them as two different epochs, R M,P and R M,A (see Table 1). 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 it relates to earlier stimuli, e.g. to move an effector in a particular direction. Indeed, empirical results suggest that this solution arises in RNN models trained on multiple tasks [14].

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, { ( q r 1: T r , c r ) } r =1 ,...,N c N trials , where N trials 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 t | c ) and transitions between epochs. It then infers the unobserved latent variables z t and x on each trial.

a

trial index

c=1, x=3

c=1, x=1

c=1, x=6

...

c=2, x=4

task 1

block

c=2, x=2

c=2, x=8

...

task 2

block

c=3, x=5

c=3, x=1

c=3, x=2

...

task 3

block

trial parameters

DMPro

DelayPro

DMAnti

MemoryPro

MemoryAnti

Log likelihood per trial (avg. over x)

c

0

1

-1

0

1

-1

b

prob.

0

1

-1

0

1

-1

DelayAnti, x=4

0

1

-1

0

1

-1

MemoryAnti, x=1

trial index

When training trials are available as batches across all trials and tasks, inference and learning can be solved using the classic Expectation-Maximization (EM) algorithm. EM performs a coordinate ascent on the complete data log-likelihood 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. A pseudocode of the algorithm and details on the inference and parameter update equations are available in Appendix B.1. To overcome convergence to bad local optima in the log-likelihood, we 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. 2b), 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 is robust to changes in task ordering (see examples in Appendix B), 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 z t during each trial, and we considered two types of inference. The more accurate inference (used to produce Fig. 2b) is p ( z t | q 1: T , c ) . This uses both the inputs s 1: T and targets y 1: T for inference and produces a smoothing distribution in that

IPX

R GLYPH<c=30,font=/WHQPCS+HelveticaNeueLTStd-Md>

\T GLYPH<c=13,font=/WHQPCS+HelveticaNeueLTStd-Md> Z ^

U

U

U

UaTL JEGLYPH<c=27,font=/WHQPCS+HelveticaNeueLTStd-Md> c

UaTL NPEFM

cPOUFYUVaM

JOGFSFOcF

UaTL NPEFM

cPOUFYUVaM

JOGFSFOcF

JOQVUTGLYPH<c=27,font=/WHQPCS+HelveticaNeueLTStd-Md> T

U

UaSHFUTGLYPH<c=27,font=/WHQPCS+HelveticaNeueLTStd-Md> Z

U

cPNQGLYPH<c=15,font=/WHQPCS+HelveticaNeueLTStd-Md>

cPOUFYU

Q GLYPH<c=9,font=/WHQPCS+HelveticaNeueLTStd-Md>[GLYPH<c=10,font=/WHQPCS+HelveticaNeueLTStd-Md>

U

a

FGGFcUJWF

OFUXPSL XFJHIUT

GLYPH<c=30,font=/WHQPCS+HelveticaNeueLTStd-Md>

GLYPH<c=12,font=/WHQPCS+HelveticaNeueLTStd-Md>

GLYPH<c=12,font=/WHQPCS+HelveticaNeueLTStd-Md>

MPXGLYPH<c=14,font=/WHQPCS+HelveticaNeueLTStd-Md>SaOL

cPOUFYU GLYPH<c=18,font=/WHQPCS+HelveticaNeueLTStd-Md>

MPXGLYPH<c=14,font=/WHQPCS+HelveticaNeueLTStd-Md>SaOL

cPOUFYU GLYPH<c=19,font=/WHQPCS+HelveticaNeueLTStd-Md>

MPXGLYPH<c=14,font=/WHQPCS+HelveticaNeueLTStd-Md>SaOL

bFMJFG GSPN UaTL NPEFM

b

x

2

x

N

x

1

.

.

.

network

space

context 1

context 3

context 2

c

the inference of z t depends on data from t ′ > t . Since this requires access to the target responses, we call this 'training time inference.' (Fig. 2c, 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, s 1: t . This produces a filtering distribution, p ( z t | s 1: t , c ) , which turns out to be highly accurate after learning (Fig. 2c, bottom row). In particular, it is able to disambiguate situations (e.g., R M,P , R M,A at the end of memory-guided response tasks have identical s t statistics) by using the ground truth task label and the learned epoch transition structure for each task (e.g., knowing that MemoryPro and MemoryAnti tasks respectively end with R M,P , R M,A ). Hereafter, we will refer to 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 a 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 to 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/response transformations. The general framework for how the two systems interact is illustrated in Fig. 3.

4.1 Recurrent neural network architecture

We chose a simple RNN architecture, where the hidden-state activity h t evolves in time according to

with decay rate α , activation function ϕ , external input s t , input bias b in and uncorrelated Gaussian noise ξ t . The network activity is read out via

with output bias b out . Given the external input and initial state, the network is trained to produce a target response time series y t by minimizing the weighted mean squared error between y t and ˆ y t (Appendix C.3).

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 p t ( z ) inferred by the task model (Fig. 3a). More specifically, each computational context in our design corresponds to a set of weights: a low-rank component of recurrent weights U z V T z , input weights W in z , input bias b in z , output weights W out z , and output bias b out z . The effective network weights at time t are a weighted

DelayPro

MemoryPro

MemoryAnti

DMPro

DMAnti

a

g

trials (x 2560)

task

order 4

task

order 3

task

order 2

task

order 1

test performance

0.0

0.5

1.0

b

c

d

h

0

100

200

300

400

500

600

trials (x 2560)

0

100

200

300

400

500

600

e

f

-5

0

log test loss

Adam

NM-RNN

EWC

0

200

400

600

total trials (x 2560)

0.0

0.5

1.0

performance

0

200

400

600

total trials (x 2560)

0

200

400

600

total trials (x 2560)

Hypernetwork

OWP

ContextRNN (ours)

0

200

400

600

total trials (x 2560)

0

200

400

600

total trials (x 2560)

0

200

400

600

total trials (x 2560)

-5

0

log test loss

0.0

0.5

1.0

performance

ContextRNN (ours)

OWP

sum of these components (Fig. 3b):

This architecture allows the network to express distinct dynamics for different computational contexts. (illustrated in Fig. 3c). And computations performed by this downstream how system are selectively gated by the what system. Contrary to previous work, e.g. Duncker et al. [13], 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 t , y t } 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 1) for the task model over a whole batch, then computed contextual inference p t ( z ) := p ( z t | s 1: t , y 1: t , c ) for the batch to instruct the downstream RNN. During testing, the RNN instead received test time inference p t ( z ) := p ( z t | s 1: t , c ) from the task model, which did not make reference to the target response y t .

We compared the performance of our approach to five baselines and summarized the results in Fig. 4, 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 (1) and (2), where W rec , W in , b in , W out , b out are unconstrained and shared across tasks. The general RNNs additionally receive a task identity input W c c , where c is a one-hot encoding of the current task. Under our default choices of hyperparameters (Appendix C.1), 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 to catastrophic forgetting, as is readily visible from the log-loss and the task performance, which quickly degraded after a change in task identity (Fig. 4a). The second baseline is Elastic Weight Consolidation (EWC) from Kirkpatrick et al. [31], 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-weighted approach failed to mitigate catastrophic forgetting (Fig. 4b). For a baseline that specifically addressed continual learning in RNNs, we included the approach from Duncker et al. [13], which selectively slows learning rates in subspaces the network explored on 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 a result, it became harder for OWP to achieve proficient task performance across all tasks (Fig. 4c).

The other two baselines are more similar to our two-system approach in that they have a dedicated component that receives the task identity input and controls the weights used by an RNN that only receives the task input s t . 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 an 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. 4d). The fifth baseline uses hypernetworks for continual learning [23, 32]. A hypernetwork receives distinct learnable embeddings for each task and outputs the task-dependent weights for a target network (an output-generating 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. 4e). By contrast, Fig. 4f shows results from our 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 facilitates forward & backward transfer

We next asked whether, given the compositional structure of the task design, a shared computational context across tasks could facilitate transfer learning. Positive forward transfer refers to the phenomenon where the test loss for a given task decreases more rapidly when it is trained after another task, compared to when it is trained from scratch. We observed positive forward transfer between nearly all pairs of tasks (Fig. 5a and Appendix D). 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 test loss on Memory occasionally converged to a slightly higher final value after a rapid initial

0

10

20

trials (x 256)

-1.5

-1.0

-0.5

log test loss

DelayPro

0

10

20

trials (x 256)

DelayPro

MemoryPro

0

10

20

trials (x 256)

MemoryPro

MemoryAnti

0

10

20

trials (x 256)

DMPro

DelayPro

a

b

0

50

100

trials (x 2560)

-5.0

-2.5

0.0

DelayPro

log test loss

as the first task

after DelayAnti

after MemoryPro

after DMPro

0

50

100

trials (x 2560)

-5.0

-2.5

MemoryPro

log test loss

as the first task

after MemoryAnti

after DelayPro

after DMPro

0

50

100

trials (x 2560)

-5.0

-2.5

0.0

DMPro

log test loss

as the first task

after DMAnti

after DelayPro

after MemoryPro

decrease relative to training from scratch. This phenomenon was also observed in Duncker et al. [13]. 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. 5b 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 facilitates compositional generalization on new tasks

As a final experiment, we 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, 15, 36] and machine learning more generally [17-19]. Specifically, we introduced a variant of the memory-guided response task in which the memory epoch ( M ) is omitted. These variants are denoted as M'Pro ( F → S → R M,P ) and M'Anti ( F → S → R M,A ) (Figure 6a; note that they are different from DelayPro / DelayAnti 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 → S → M → R M,P ), and then assessed whether it can rapidly learn a new task, MemoryAnti ( F → S → M → R M,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 to 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 on average. In contrast, baselines (also pretrained on the M'Pro/Anti and MemoryPro tasks) using Adam, OWP and Hypernetworks respectively achieved only 56%, 53%, 64% accuracy after 512 trials of full model training (Figure 6b, c). These results demonstrate that our model can rapidly generalize to new tasks by recombining previously learned components.

5 Discussion

How does the brain learn to solve tasks through compositional computation? We approached this question by first developing a formalism of compositional tasks that is sufficiently expressive to capture 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 - a time-varying, low-dimensional computational context corresponding to the

pre-training on task epochs

train contextual inference model and low-rank RNN components

compositional generalization

only train contextual inference model

pro

anti

response

stimulus

time

M' Pro/Anti

pro

response

delay

stimulus

time

Memory Pro

anti

response

delay

stimulus

time

Memory Anti

a

b

c

-4

-2

0

MemoryAnti

log test loss

0

100

200

300

400

500

0.0

0.5

1.0

MemoryAnti

test performance

0

100

200

300

400

500

trials

ContextRNN (ours)

OWP

Adam

Hypernetwork

trials

task epoch. This signal is then used to create or (re-)select computational components (one for each context) in a second how system. Numerical experiments suggest that incorporating compositionality and contextual inference in this way allows the network to 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 [9, 36, 41] . While we don'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 pattern-generating dynamics for motor sequences can aid flexible recomposition and learning, and be implemented via thalamocortical loops modulated via selective inhibition from basal ganglia [38]. A similar architecture would be possible for our model to extend these ideas from the motor domain to 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 a large literature on action selection, implicating basal ganglia circuits in the selection and composition of motor sequences [42, 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, 4, 7] 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 a 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 a 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 different task model may result in a different segmentation and opportunities for reuse. In addition, we have assumed the trial variable x to be a 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 a more principled way by using non-parametric 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.

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

James B Heald, Máté Lengyel, and Daniel M Wolpert. Contextual inference underlies the learning of sensorimotor repertoires. Nature , 600(7889):489-493, 2021.

Robert C Wilson, Yuji K Takahashi, Geoffrey Schoenbaum, and Yael Niv. Orbitofrontal cortex as a cognitive map of task space. Neuron , 81(2):267-279, 2014.

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-1326, 2022.

Yael Niv. Learning task-state representations. Nature neuroscience , 22(10):1544-1553, 2019.

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.

Timo Flesch, Andrew Saxe, and Christopher Summerfield. Continual task learning in natural and artificial agents. Trends in neurosciences , 46(3):199-210, 2023.

Jingfeng Zhou, Chunying Jia, Marlian Montesinos-Cartagena, Matthew PH Gardner, Wenhui Zong, and Geoffrey Schoenbaum. Evolving schema representations in orbitofrontal ensembles during learning. Nature , 590(7847):606-611, 2021.

Oded Bein and Yael Niv. Schemas, reinforcement learning and the medial prefrontal cortex. Nature Reviews Neuroscience , pages 1-17, 2025.

Lucas Y Tian, Kedar U Garzón, Adam G Rouse, Mark AG Eldridge, Marc 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.

Xu He, Jakub Sygnowski, Alexandre Galashov, Andrei A Rusu, Yee Whye Teh, and Razvan Pascanu. Task agnostic continual learning via meta learning. arXiv preprint arXiv:1906.05201 , 2019.

Guangyu Robert Yang, Madhura R Joglekar, H Francis Song, William T Newsome, and XiaoJing Wang. Task representations in neural networks trained to perform many cognitive tasks. Nature neuroscience , 22(2):297-306, 2019.

Andrea Cossu, Antonio Carta, Vincenzo Lomonaco, and Davide Bacciu. Continual learning for recurrent neural networks: an empirical evaluation. Neural Networks , 143:607-627, 2021.

Lea Duncker, Laura Driscoll, Krishna V Shenoy, Maneesh Sahani, and David Sussillo. Organizing recurrent network dynamics by task-computation to enable continual learning. Advances in neural information processing systems , 33:14387-14397, 2020.

Laura N Driscoll, Krishna Shenoy, and David Sussillo. Flexible multitask computation in recurrent networks utilizes shared dynamical motifs. Nature Neuroscience , 27(7):1349-1363, 2024.

Reidar Riveland and Alexandre Pouget. Natural language instructions induce compositional generalization in networks of neurons. Nature Neuroscience , 27(5):988-999, 2024.

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.

Thaddäus Wiedemer, Prasanna Mayilvahanan, Matthias Bethge, and Wieland Brendel. Compositional generalization from first principles. Advances in Neural Information Processing Systems , 36:6941-6960, 2023.

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.

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.

Samuel J Gershman, David M Blei, and Yael Niv. Context, learning, and extinction. Psychological review , 117(1):197, 2010.

Samuel J Gershman, Angela Radulescu, Kenneth A Norman, and Yael Niv. Statistical computations underlying the dynamics of memory updating. PLoS computational biology , 10(11): e1003939, 2014.

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.

Johannes Von Oswald, Christian Henning, Benjamin F Grewe, and João Sacramento. Continual learning with hypernetworks. arXiv preprint arXiv:1906.00695 , 2019.

Ben Tsuda, Kay M Tye, Hava T Siegelmann, and Terrence J Sejnowski. A modeling framework for adaptive lifelong learning with transfer and savings through gating in the prefrontal cortex. Proceedings of the National Academy of Sciences , 117(47):29872-29882, 2020.

Nicolas Y Masse, Gregory D Grant, and David J Freedman. Alleviating catastrophic forgetting using context-dependent gating and synaptic stabilization. Proceedings of the National Academy of Sciences , 115(44):E10467-E10475, 2018.

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):e1010808, 2023.

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):8275, 2024.

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-6978, 2024.

Oleksiy Ostapenko, Pau Rodriguez, Massimo Caccia, and Laurent Charlin. Continual learning via local module composition. Advances in Neural Information Processing Systems , 34:3029830312, 2021.

Shagun Sodhani, Sarath Chandar, and Yoshua Bengio. Toward training recurrent neural networks for lifelong learning. Neural computation , 32(1):1-35, 2020.

James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences , 114(13):3521-3526, 2017.

Benjamin Ehret, Christian Henning, Maria R Cervera, Alexander Meulemans, Johannes Von Oswald, and Benjamin F Grewe. Continual learning in recurrent neural networks. arXiv preprint arXiv:2006.12109 , 2020.

Guanxiong Zeng, Yang Chen, Bo Cui, and Shan Yu. Continual learning of context-dependent processing in neural networks. Nature Machine Intelligence , 1(8):364-372, 2019.

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.

Zeyuan Yang, Zonghan Yang, Yichen Liu, Peng Li, and Yang Liu. Restricted orthogonal gradient projection for continual learning. AI Open , 4:98-110, 2023.

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.

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 B (methodological) , 39 (1):1-22, 1977.

Laureline Logiaco, LF Abbott, and Sean Escola. Thalamic control of cortical dynamics in a model of flexible motor sequencing. Cell reports , 35(9), 2021.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Sen Lin, Li Yang, Deliang Fan, and Junshan Zhang. Beyond not-forgetting: Continual learning with backward knowledge transfer. Advances in Neural Information Processing Systems , 35: 16165-16177, 2022.

Weinan Sun, Madhu Advani, Nelson Spruston, Andrew Saxe, and James E Fitzgerald. Organizing memories for generalization in complementary learning systems. Nature neuroscience , 26 (8):1438-1448, 2023.

Melissa J Sharpe, Thomas Stalnaker, Nicolas W Schuck, Simon Killcross, Geoffrey Schoenbaum, 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-76, 2019.

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. Spontaneous behaviour is structured by reinforcement without explicit reward. Nature , 614(7946):108-117, 2023.

James CR Whittington, Timothy H Muller, Shirley Mark, Guifen Chen, Caswell Barry, Neil Burgess, and Timothy EJ Behrens. The tolman-eichenbaum machine: unifying space and relational memory through generalization in the hippocampal formation. Cell , 183(5):12491263, 2020.

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.

CF Jeff Wu. On the convergence properties of the em algorithm. The Annals of statistics , pages 95-103, 1983.

Krishna Nathan, Andrew Senior, and Jayashree Subrahmonia. Initialization of hidden markov models for unconstrained on-line handwriting recognition. In 1996 IEEE International Conference on Acoustics, Speech, and Signal Processing Conference Proceedings , volume 6, pages 3502-3505. IEEE, 1996.

Kenneth D Miller and Francesco Fumarola. Mathematical equivalence of two common forms of firing rate models of neural networks. Neural computation , 24(1):25-31, 2012.

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

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: A discussion of limitation is supplied at 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 in 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 on 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: Our results do not contain proof-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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

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 a 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 it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to 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 a 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or 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 reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.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: //nips.cc/public/guides/CodeSubmissionPolicy ) 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 a 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 in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

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 to a 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, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It 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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: In general, our experiments have very little compute requirement. The configurations used are stated in 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 of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have reviewed the code of ethics and can confirm that our work conforms with it.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Our research aims to better understand task compositionality and how the brain may handle it. These do not have immediate relations to ML 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.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: See previous question.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses 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 not use 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 a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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.

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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was 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

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or 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 on 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 NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

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, or 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, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Separating the what and how of compositional computation to enable reuse and continual learning Appendix

All our codes are available on the public repository.

A Formalization of compositional tasks and task designs

We propose a formalization of shared compositional structures in sequence-to-sequence (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 s t =1 ,...,T and needs to produce a target output sequence y t =1 ,...,T . Different tasks (indexed by c ∈ Z + ) are distinguished from each other by their input distributions as well as the underlying input-to-target rules, together encoded in p ( s 1: T , y 1: T | c ) . For notational convenience let q t ≡ [ s t , y t ] ∈ R D q .

We model this distribution over task-observables by introducing a set of latent variables resulting in the joint distribution

We initially provide intuition about our modeling choices, before formally specifying the distributions for the tasks of the main paper in section A.1 and A.2 below.

To model different tasks with a shared compositional structure, we first observe that in many tasks used in neuroscience experiments, each trial can be temporally segmented into discrete epochs, which we denote as z t ∈ Z + . Each epoch has its distinctive input and target statistics, though they may exhibit complex temporal dependencies across epochs throughout the trial. To illustrate this, we can consider a simple task testing the learner's ability to memorize: During a stimulus epoch, inputs provide information about a trial-specific latent variable (e.g., a dot on the screen indicating an angle θ ). The inputs eventually turn off, and the learner needs to maintain the relevant information in memory until a cue solicits a θ -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 ( q t | z t ) . However, this approach would fail to capture that the y t during the response epoch is coupled to the s t 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 θ . To take this cross-epoch dependency into account, we explicitly model the 'stimulus condition', indexed by the trial variable x . In the specific example above, there may be a list of possible directions and x selects one to use as θ . Thus, in our model of cognitive tasks, the observed inputs and target responses depend not only on the latent epoch z t 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 q t | z t , x to describe them. For a simple example of different tasks sharing compositional structure, we can consider the MemoryPro task described above, and construct a second task, MemoryAnti , which shares the same stimulus epoch but requires a response in the opposite direction of θ . Here, the dependence of the response epoch on θ is different (and MemoryAnti will therefore have a different response epoch relative to MemoryPro ), but the statistics of the stimulus epoch are identical across the two tasks. The differences in what epochs contribute to each task and how one epoch transitions to the next is captured via task-dependent Markovian transitions over epochs p ( z t +1 | z t , c ) .

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 ( x | c ) is assumed to be a uniform distribution

over some finite set of size N x . p ( z 1: T | 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 t | z t , x ) is assumed to be a multivariate Gaussian with a ( z t , x ) -dependent 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 = 1 , . . . , N

This model thus allows us to specify a set of N c tasks sharing a compositional structure as follows. Let N z denote the number of epochs shared among these tasks. The composition of the c -th task is specified by its epoch-transition parameters: the transition probabilities Λ c ∈ R N z × N z : Λ c z,z ′ = p ( z t = z | z t -1 = z ′ , c ) and initial probabilities Π c ∈ R N z : Π c z = p ( z 1 = z | c ) . The observation model of each epoch is specified by the means corresponding to different x values, { ¯ q z,x ∈ R D q } x =1 ,...,N x , where ¯ q z,x is the mean of the multivariate Gaussian p ( q t | z t = z, x ) . Thus, altogether an ensemble of N c tasks composed from N z epochs is specified by the tuple ( { ¯ q z,x } z =1 ,...,N z ; x =1 ,...,N x , { (Λ c , Π c ) } c =1 ,...,N c ) , where the first component specifies the epochs and the second specifies how they are used to compose the tasks.

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 1 of 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 S DM ), the fixation cue is on and the required output is to maintain fixation without producing a response. During the response epochs ( R P , R A , R M,P , R M,A , R DM,P , R DM,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 x = 8 ) and generate x 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 θ , chosen from { 0 , π/ 4 , ..., 7 π/ 4 } depending on x . The stimulus presentation stays on during the ensuing response epoch ( R P for DelayPro or R A for DelayAnti ), where the target output becomes ϕ = θ ( R P ) or ϕ = θ + π ( R A ).

Memory-guided response tasks ( MemoryPro , MemoryAnti ) . After S , stimulus presentation disappears in the memory epoch ( M ). During the ensuing response epoch ( R M,P for MemoryPro , R M,A for MemoryAnti ), there is still no stimulus presentation and the learner must produce a response based on the memorized θ : ϕ = θ ( R M,P ) or ϕ = θ + π ( R M,A ).

Decision making tasks ( DMPro , DMAnti ) . After F , a decision stimuli epoch ( S DM ) presents two stimuli simultaneously. θ = 0 with strength γ in input dims 1, 2 and θ = π with strength γ ′ in input dims 3, 4. The strengths ( γ, γ ′ ) are determined from a set of pairs by x and scale the input channels. During the ensuing response epoch ( R DM,P for DMPro , R DM,A for DMAnti ), the stimuli persist and the required output ϕ is the direction of the stimulus with a higher strength ( R DM,P ) or the one with a lower strength ( R DM,A ).

The epoch structure of the different tasks is encoded via Λ c . Altogether, the 6 tasks can be described as compositions using a shared pool of 10 epochs. Given the epoch identity, the conditional distributions of the inputs and responses are generated as follows:

Where 1 is an indicator that takes on value 1 if the subscript is true and 0 otherwise. This fully specifies the conditional distributions p ( q t | z t , x ) for each task epoch with ¯ q z,x = [¯ s z,x , ¯ y z,x ] . For the DM tasks, we restrict the stimulus values to two locations θ = 0 , θ ′ = π , but pick N x = 8 different combinations of possible ( γ, γ ′ ) pairs. Note that while we generate the input and response distributions using stimulus values θ to follow convention from the literature [11, 14], each value of θ (or ( γ, γ ′ ) in the DM tasks) maps onto a different x value. For σ we used 0.05; for ( γ, γ ′ ), x selects from [0 . 5 , 1] , [1 , 2] , [0 . 5 , 2] , [0 . 2 , 1 . 5] , [1 , 0 . 5] , [2 , 1] , [2 , 0 . 5] , [1 . 5 , 0 . 2] .

Note that while we have followed many conventions from previous work in the task design, these previous approaches tend to implement each task individually but with related distributional assumptions [11, 13, 14, 16]. 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 a 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, 13-15].

B Online learning and inference of compositional task structures

B.1 Online learning and inference

In this section, we provide additional details on the algorithms developed for performing posterior inference over the latent variables of the task model, and online (one-trial-at-a-time) 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 α r t ( z, x, c ) = p ( q r 1: t , z r t = z | x r = x, c r = c ) denote the forward (filtering) message for a given task. During a filtering pass, we compute

Marginalizing the forward message at the final time-step allows us to compute the marginal likelihood over observations

(10)

Let β r t ( z, x, c ) = p ( q r ( t +1): T | z r t = z, x r = x, c r = c ) denote the backwards (smoothing) message. During the smoothing pass, we compute

Given these quantifies, we can compute the joint posterior over z r t , x r and c r as

Note that the filtering and smoothing passes can be done for each value of x and c in parallel. Finally, for later use in the online learning algorithm we also compute the joint posterior over z r t -1 , z r t , x r and c r as

Online learning. Learning aims to recover parameters of the generative model, ( { ¯ q z,x } z =1 ,...,N z ; x =1 ,...,N x , { (Λ c , Π c ) } c =1 ,...,N c ) , up to a permutation over z and x . Performing 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 as the set of sufficient statistics needed to update the set of model parameters Θ . Let Θ ( i,k ) denote the learned parameters after seeing the i -th trial and running k EM iterations. Let X r, ( i,k ) denote the single-trial statistics on trial r , computed using parameters Θ ( i,k ) . When the learning algorithm has access to all trials throughout learning (batch EM), the updates take the form

For example, for the parameter estimates for the transition matrix, this takes the form

with

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 one-trial-at-a-time and cannot be revisited (online EM), the sums of 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

where K denotes the number of iterations per trial. After learning the i -th trial, we only need to store Θ ( i,K ) and S ( i,K ) online ( X ) , taking the form

We summarize this online learning approach in Algorithm 1 and show performance for different training orders (supplementing Figure 2 in the main paper) in Figure S.1.

trial index trial index

B.2 Incremental initialization of model parameters

Even for simple models such as the Gaussian Mixture Model, EM is susceptible to local optima [46]. To avoid convergence to bad local optima, it is important to obtain good initializations for the model parameters. For simple time-series models such as the HMM, this is typically achieved by collapsing the sequences across time, performing clustering (e.g. with K-means [47]), and using the resulting cluster centers as the initial observation-model parameters, given each latent discrete state.

In our model, we need to initialize the underlying cluster centers of the observation model, { ˆ q z,x } z =1 ,..., ˆ N z ,x =1 ,...,N x . Here, ˆ N z reflects the fact that the total number of epochs in the entire dataset is unknown a priori . We overcome this by setting up a large number of 'slots' ( ˆ N z ≥ N z ), exceeding the likely total number of epochs in the task family. Our setting adds two significant challenges relative to simple HMMs: First, each trial contains only a few epochs and a specific stimulus condition ( x value). Therefore, clustering each trial can only initialize a small subset of the N x N z mixture components of the observation model. Second, since there is one stimulus condition x per trial, the mixture components ¯ q z,x explaining a given trial must be allocated to the same x , 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 a feasible initial set of parameters. To overcome these issues, we introduce an 'incremental initialization' scheme, which is applied to all the estimated means { ˆ q z,x } before learning each trial.

To introduce the scheme, we first introduce the notion of the putative z, x , denoted ˜ z, ˜ x . This is to highlight that learning the correct { ¯ q z,x } does not require inferring the real z, x in the generative process but only requires them to be correct up to a permutation. Let ˜ N z , ˜ N c denote the number of

epochs and task slots in the learning algorithm, respectively. Note that these do not need to be set as the correct N z , N c - they can simply be large integers. We do assume that the system knows the correct number of x values, N x . The learner keeps track of { ˆ q ˜ z, ˜ x } ˜ z =1 ,..., ˜ N z , ˜ x =1 ,..., ˜ N x as well as two Boolean-valued tables, F c,x of size ˜ N c × N x and F x,z of size ˜ N x × ˜ N z . The two tables keep track of which combinations of c, ˜ x and ˜ x, ˜ z have been encountered.

When the r th trial is observed and the learner has access to ( q r 1: T r , c r ) , the scheme is threefold:

We perform K-means clustering on the entire sequence q r 1: T r . This gives us a 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 ˜ z, ˜ x to these means, and to use them to initialize ˆ q ˜ z, ˜ x accordingly.

All cluster means from step (1) should be assigned the same ˜ x r . We check the means against { ˆ q ˜ z, ˜ x } and decide on ˜ x r according to a set of rules and F c,x . The c r , ˜ x r pair is marked 'encountered' in F c,x .

Given ˜ x r from step (2), we treat the cluster means not found in { ˆ q ˜ z, ˜ x } ˜ x =˜ x r as unfamiliar ˜ z, ˜ x , meaning that they represent previously unseen ˜ z, ˜ x combinations and should be used to initialize. Each center is assigned a different ˜ z that has not been encountered (according to F x,z ). The ˜ z, ˜ x r pairs are marked as 'encountered' in F x,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 Λ c , 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 z,x } , we infer which epochs appeared in this trial using the posterior p ( z r 1: T | q r 1: T ) . Only epochs with a sufficiently high chance of appearance have their { ˆ q z,x } updated.

B.4 Epoch identifiability in our set of tasks

For the particular set of tasks we considered, since F, M epochs have identical ¯ q z,x for all x , they are indistinguishable. Thus, our learning algorithm will combine them into a single epoch, which we denote as F/M . Instead of learning the 10 epochs in the generative process for our set of tasks, the task-model will end up learning 9 epochs, including the combined F/M epoch. In terms of the transition matrices, this creates a 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 to either the stimulus epoch or the response epoch. The 'ground-truth' parameters used for plotting in Fig. 2 and Fig. S.1 refer to the optimal parameters with a 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 be distinguished, e.g. via feedback from the downstream network implementing the different computations for each epoch.

C RNN architecture and hyperparameter settings

C.1 Default RNN architecture and task parameters

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 η z = 0 . 001 . After training each task c , η z was multiplied by a decay factor γ = 0 . 5 for any context with p ( z | c ) > 0 . 001 . During training on task c , L 2 weight regularization was applied with a coefficient of 10 -5 for contexts with p ( z | c ) > 0 . 001 , and set to 0 for all other contexts.

For baseline algorithms using general ("vanilla") RNN architectures, the learning rate was set to 0.01, and the L 2 weight regularization coefficient was set to 10 -5 . Parameter choices were determined by a coarse grid search.

C.3 Loss function and performance measure

The loss function is a weighted mean square error similar to [11, 14]. L := ⟨ m i,t ( y i,t -ˆ y i,t ) 2 ⟩ i,t , where i is the index of the output units, m i,t = 1 for response epochs and m i,t = 0 . 2 for all other epochs.

A 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. If the activity of the fixation output exceeds 0.5, the network is considered to have broken fixation. The network's response direction is considered correct if its angular difference from the target direction is less than π/ 10 . Average performance and test loss were calculated on 200 test trials for each task.

D Supplementary results

D.1 Additional results on transfer learning

We provide additional results on transfer learning. Figure S.2 and Figure S.3 supplement Figure 5 a 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 [13], where the test loss of previously learned tasks did not 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 as we increased the rank of U z and V z (denoted by r ) (Figure S.5 a , b , compared with Figure 4 f ). Using r = 3 , 5 , and 10 in ContextRNN results in 34587, 43803, and 66843 trainable parameters, respectively - all fewer than the 69379 parameters of a general RNN with the same number of hidden units. With r = 3 , ContextRNN performed worse with tanh than with ReLU activation, but this gap was closed at r = 10

-5

0

DelayPro

log test loss

after

DelayPro

after

after

MemoryPro

after

MemoryAnti

after

DMPro

after

DMAnti

-5

0

log test loss

-5

0

MemoryPro

log test loss

-5

0

MemoryAnti

log test loss

-5

0

DMPro

log test loss

0

100

trials (x 2560)

-5

0

DMAnti

log test loss

0

100

trials (x 2560)

0

100

trials (x 2560)

0

100

trials (x 2560)

0

100

trials (x 2560)

0

100

trials (x 2560)

E Details of baseline implementations

E.1 Elastic Weight Consolidation

Elastic Weight Consolidation (EWC) [31], slows down learning rates for network weights deemed important for previous tasks. This is achieved by adding a regularization term to the training loss function. For a loss L (Θ) , the EWC objective is given as

Here F i is an importance weight that ties the i th parameter value Θ i to it's value Θ ∗ i 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 Θ ∗ i . We set λ to 10 5 after a coarse grid search.

E.2 Orthogonal Weight Projection

The original learning rule in Duncker et al. [13] was derived for a parameterization of network dynamics of the form

given an element-wise activation function ϕ ( · ) . The learning rule was intended to maintain the stimulus/response relationship of previous task by applying a set of projection matrices to the gradient used to update the network weights over learning. The projection matrices were intended to remove

-1.50

-1.25

DelayPro

log test loss

during

DelayPro

during

during

MemoryPro

during

MemoryAnti

during

DMPro

during

DMAnti

-1.50

-1.25

log test loss

-1.3

-1.2

-1.1

MemoryPro

log test loss

-1.3

-1.2

-1.1

MemoryAnti

log test loss

-1.50

-1.25

DMPro

log test loss

0

10

trials (x 256)

-1.50

-1.25

DMAnti

log test loss

0

10

trials (x 256)

0

10

trials (x 256)

0

10

trials (x 256)

0

10

trials (x 256)

0

10

trials (x 256)

0

10

20

30

trials (x 256)

-2

-1

0

log test loss

0

10

20

30

trials (x 256)

0

10

20

30

trials (x 256)

0

10

20

30

trials (x 256)

0

10

20

trials (x 256)

-2

-1

0

log test loss

DelayPro

0

10

20

trials (x 256)

DelayPro

MemoryPro

0

10

20

trials (x 256)

MemoryPro

MemoryAnti

0

10

20

trials (x 256)

DMPro

DelayPro

log test loss

a

b

trials (x 2560)

task

order 4

task

order 3

task

order 2

task

order 1

test performance

0.0

0.5

1.0

0

100

200

300

400

500

600

trials (x 2560)

0

100

200

300

400

500

600

ContextRNN: α = 0.1, r = 5, Φ = ReLU

ContextRNN: α = 0.1, r = 10, Φ = ReLU

trials (x 2560)

task

order 4

task

order 3

task

order 2

task

order 1

test performance

0.0

0.5

1.0

0

100

200

300

400

500

600

trials (x 2560)

0

100

200

300

400

500

600

ContextRNN: α = 0.1, r = 3, Φ = tanh

ContextRNN: α = 0.1, r = 10, Φ = tanh

a

b

c

d

trials (x 2560)

task

order 4

task

order 3

task

order 2

task

order 1

test performance

0.0

0.5

1.0

0

100

200

300

400

500

600

trials (x 2560)

0

100

200

300

400

500

600

ContextRNN: α = 0.2, r = 3, Φ = ReLU

ContextRNN: α = 0.2, r = 10, Φ = ReLU

e

f

directions from the weight update that would interfere with previous tasks and were defined as follows.

Letting z c,r t = [ x c,r t s c,r t ] denote the concatenated network state and input state at time t of trial r on task c , Z 1: c = [ z 1 , 1 1 , . . . z c,r T ] the collection of all time points and trials on tasks 1 through c , and W = [ W rec W in ] the concatenated weight matrices, we define the projection matrices

and the modified learning update as

where ∇ W L 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 to a modified setting, where the RNN dynamics are expressed as

given the same element-wise activation function ϕ ( · ) . While the two parameterizations in (26) and (30) are generally considered equivalent [48], the linear intuition used to motivate the approach of [13] should be exact in (30). With x t = ϕ ( h t ) , Z 1: c is unchanged, but we instead use

where H 1: c = [ h 1 , 1 1 , . . . h c,r T ] . While this is very similar to the version in Duncker et al. [13], 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), in which a neuromodulatory subnetwork (RNN1) outputs a low-dimensional signal that dynamically scales the low-rank recurrent weights of an output-generating RNN (RNN2). RNN1 and RNN2 are jointly optimized through gradient descent. In our implementation, RNN2 has 256 neurons and with rank( W rec ) = 27, equal to the default choice for our ContextRNN. Number of neurons in RNN1 is set to 125 such that the whole network has 36658 parameters, comparable to 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 set to 0.001 and L 2 weight regularization coefficient set to 10 -5 after a coarse grid search.

E.4 Hypernetwork

Von Oswald et al. [23] designed task-conditioned hypernetworks to 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 set 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 hypernetwork. 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 to 34. Dimension of task embeddings and chunk embeddings is set to 32. The hypernetwork is a multi-layer perceptron with two 32-d hidden layers. Together, the number of trainable parameters of the hypernetwork equals 70416, comparable to other baselines in our paper using a full-rank RNN. After a coarse grid search, we set continual learning regularization strength to 1, orthogonal regularization strength to 0.001 and clip gradient norm to 1 (see Ehret et al. [32] for details). We used Adam optimizer, with learning rate set to 0.001 and L 2 weight regularization coefficient set to 10 -5 .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n- 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.\n\nQUESTIONS:\n1. 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.)\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\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.\n\nQUESTIONS:\nPlease 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.\n\nLIMITATIONS:\nAs 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 1, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\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.\n\nQUESTIONS:\nOpportunities 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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nQuality:\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\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nFor 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?", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}