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

TL;DR

The paper tackles how to separate what to compute from how to compute it to enable continual, compositional learning in neural networks. It introduces a two-system framework: a what system that infers a time-varying computational context over task epochs $z_t$ and a trial variable $x$, and a how system that implements the corresponding computation via a context-modulated, low-rank RNN with weights $W^{rec}= abla_t? $; corrected: $W^{rec} = \sum_z p_t(z) U_z V_z^T$. The authors formalize task compositionality with a probabilistic generative model and develop an online EM-based procedure to infer contexts online, demonstrating its effectiveness across six neuroscience-style tasks with results showing mitigated forgetting, forward and backward transfer, and rapid compositional generalization to unseen tasks. The work links computational architecture to potential neural mechanisms, offering a principled route to reuse and extend learned skills without catastrophic interference and with improved generalization in structured task families.

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

Figures (11)

• Figure 1: A generative model of a family of cognitive tasks used in neuroscience. a: Schematic illustration of a set of cognitive tasks and their compositional nature. b: Directed Acyclic Graphical model (DAG) description of the generative model capturing the distribution over cognitive tasks we consider. The time-varying inputs and target responses of each task can be modeled as a mixture, where the observations for each task are composed of a set of discrete epochs, as illustrated in (a).
• Figure 2: Online continual learning of the compositional structure of tasks. a: Schematics of online continual learning of a set of compositional tasks. b: Colored lines: log likelihood (LL) of tasks (averaged over trials) from the learned model over the course of learning. Background colors: the task being trained at each trial. Dashed horizontal lines represent LLs computed using ground-truth parameters. c Single-trial epoch inference by the learned model. Top row: inputs and target responses of example trials. Dashed vertical lines indicate boundaries between epochs. Trial indices are dropped for brevity. Center row: output of training-time inference ($p(z_{t}|\pmb{q}_{1:T},c)$). Bottom row: output of test-time inference ($p(z_{t}|\pmb{s}_{1:t},c)$), which does not require access to the target responses or information from the future. Note that, since the $F$ and $M$ epochs in memory-guided tasks have indistinguishable observation models, they are merged by the learner.
• Figure 3: Schematic illustration of two interacting learning systems and contextual gating of RNN dynamics. a: The task model (what) processes incoming input and target pairs $\pmb{q}_t = \{\pmb{s}_t, \pmb{y}_t\}$ and infers the time-varying computational context of the current task. A downstream RNN (how) learns the relevant input-to-target transformation and is modulated by the task model's belief $p_t(z)$ over the current context. b: The effective recurrent weights of the downstream RNN are a linear combination of low-rank components, weighted by belief over each corresponding context. c: The gating by the task model causes the network to express different low-dimensional dynamics across different computational contexts, each of which can be re-expressed and reused to execute different tasks.
• Figure 4: Continual learning. a-f: log test loss (colored lines, first row) and test performance (colored lines, second row) throughout sequential training, each plot shows results for five random seeds. Background colors denote the task being trained at each trial. a: Adam optimizer on general RNNs. b: Elastic Weight Consolidation (EWC) on general RNNs. c Orthogonal Weight-Space Projection (OWP) on general RNNs. d: Adam optimizer on neuromodulated RNN. e Adam optimizer on Hypernetwork RNN. f: task model with context modulated RNNs (ours). g-h: Color-coded test performance during sequential training of four different task orders. Each row color-codes the average test performance across five random seeds of a specific task over training. We compare OWP (g) with our method (h).
• Figure 5: Transfer learning. a: The log test loss during training of a task (specified in y-axis label), either as the first task (black) or after training another task (coded by color). We observe faster training after pre-training on tasks that share similar epochs, illustrating forward transfer. b: The log test loss during sequential training of two tasks, when each tasks is trained for fewer numbers of trials than our default choice. The loss of the previous tasks continues to decrease after switching to a new task, illustrating backward transfer. We plot results from three random seeds for each task order.