Thalamus: a brain-inspired algorithm for biologically-plausible continual learning and disentangled representations
Authors: Ali Hummos
Venue: ICLR 2023 poster (2023)
Keywords: brain-inspired learning
neuroscience
recurrent neural networks
context inference
bayesian brain

Abstract: Animals thrive in a constantly changing environment and leverage the temporal structure to learn well-factorized causal representations. In contrast, traditional neural networks suffer from forgetting in changing environments and many methods have been proposed to limit forgetting with different trade-offs. Inspired by the brain thalamocortical circuit, we introduce a simple algorithm that uses optimization at inference time to generate internal representations of the current task dynamically. The algorithm alternates between updating the model weights and a latent task embedding, allowing the agent to parse the stream of temporal experience into discrete events and organize learning about them. On a continual learning benchmark, it achieves competitive end average accuracy by mitigating forgetting, but importantly, the interaction between the weights dynamics and the latent dynamics organizes knowledge into flexible structures with a cognitive interface to control them. Tasks later in the sequence can be solved through knowledge transfer as they become reachable within the well-factorized latent space. The algorithm meets many of the desiderata of an ideal continually learning agent in open-ended environments, and its simplicity suggests fundamental computations in circuits with abundant feedback control loops such as the thalamocortical circuits in the brain

Published as a conference paper at ICLR 2023
THALAMUS: A BRAIN-INSPIRED ALGORITHM FOR BIOLOGICALLY-PLAUSIBIE CONTINUAL LEARNING AND DISENTANGLED REPRESENTATIONS
Ali Hummos
Brain and Cognitive Sciences Department Massachusetts Institute of Technology Cambridge, MA 02139 ahummos @MIT.
ABSTRACT
Animals thrive in constantly changing environment and leverage the temporal structure to learn well-factorized causal representations. In contrast, traditional neural networks suffer from forgetting in changing environments and many meth- ods have been proposed to limit forgetting with different trade-offs. Inspired by the brain thalamocortical circuit, we introduce a simple algorithm that uses opti- mization at inference time to generate internal representations of the current task dynamically. The algorithm alternates between updating the model weights and a latent task embedding, allowing the agent to parse the stream of temporal ex- perience into discrete events and organize learning about them. On a continual learning benchmark, it achieves competitive end average accuracy by mitigating forgetting, but importantly, by requiring the model to adapt through latent updates, it organizes knowledge into flexible structures with cognitive interface to con- trol them. Tasks later in the sequence can be solved through knowledge transfer they become reachable within the well-factorized latent space. The algorithm meets many the desiderata an ideal continually learning agent in open-ended environments, and its simplicity suggests fundamental computations in circuits with abundant feedback control loops such as the thalamocortical circuits in the brain
INTRODUCTION
Animals thrive in constantly changing environmental demands at many time scales. Biologi- cal brains seem capable of using these changes advantageously and leverage the temporal struc- ture to learn causal and well-factorized representations (Collins & Koechlin, 2012, Yu et al.] [02T]] Herce Castañón et al.) 202I). In contrast, traditional neural networks suffer in such settings with sequential experience and display prominent interference between old and new learning limiting most training paradigms to using shuffled data (McCloskey & Cohen 1989). Many recent meth- ods advanced the flexibility of neural networks (for recent reviews, see Parisi et al. (2019) y Hadsell
ett al. (2020); Veniat et al. (202T)). However, in addition to mitigating forgetting, several desirable properties continually learning agent have been recently suggested (Hadsell al. 2020 Veniat et al,, (021) including: accuracy on many tasks at the end of a learning episode or at least fast adaptation and recovery of accuracy with minimal additional training. The ideal agent would also display knowledge transfer forward, to future tasks and backwards to previously learned tasks, but also transfer to tasks with slightly different computation and or slightly different input or output distributions (Veniat et al.] (021). The algorithm should scale favorably with the number of tasks and maintain plasticity, or the capacity for further learning, Finally, the agent should ideally able to function unsupervised and not rely on access to task labels and task boundaries (Hadsell et al.) 2020 Rao et al. 2019). We argue for another critical feature: contextual behavioral, where the same inputs may require different responses at different times, feature that might constrain the solution space to be more relevance to brain function and to the full complexity the world. A learning agent might struggle to identify reliable contextual signals in high dimensional input space,

---

Published as a conference paper at ICLR 2023
Output
Weights updales
Learn
Input
Open-ended discovery tasks and disentangied task learning
Thalamus-inspen innppieed Attenatiin Latent and Weight Updalees for Sequential learning
Conticethalamic
projections.
Thalamocortial
Feadback error
projections
Gradients using backprop. but
Multiplicative
Latent embedding updates
Infer
MD (2)
Model schematic. Controlled triggering of latent updates and weight updates and alternat- ing them allows for unsupervised task discovery. Latent updates retrieve previously learned tasks or choose a new embedding for new ones (PFC: prefrontal cortex, MD: mediodorsal thalamus)
Figure 1:
they are knowable at all, and with many contextual modifiers it might not be feasible experience all combinations sufficiently to develop associative responses.
Neuroscience experiments have revealed thalamic role in cognitive flexibility and switching be- tween behavioral policies (Schmitt et al. 2017 Mukherjee et al. 2021). The prefrontal cortex (PFC), linked advanced cogmition, shows representations of task variables and input to output transfor- mations (Johnston et al. 2007;; Mansouri et al. 2006 RRuugier et al. 2005 Rikhye et al. 2018 while the medio-dorsal thalamus shows representations of the task being pertormed (Schmitt al,,
2017 Rikhye al. 2018), and uncertainty about the current task (Mukherjee et al. (021). The mediodorasal thalamus, devoid recurrent excitatory connections, and with extensive reciprocal connections to PFC is thought gate its computations by selecting task relevant representations to enable flexible behavioral switching (Wang & Halassa [021] Hummos et al.) 2022). The connec- tions from cortex thalamus engage in error feedback control to contextualize perceptual attention (Briggs) 2020), motor planning (Kao et al. 2021). and cognitive control (Halassa & Kastner] [017]]
Wang & Halassa 202T) and representations of error feedback can be observed in the thalamus (Ide)
& Li 2011 Jakob et al. [2021; Wang et al. 2020). Moreover, these thalamic representations can be composed produce complex behaviors (Logiaco al., 202T).
In this paper, we take inspiration from the thalamocortical circuit and develop simple algorithm that uses optimization at inference time to produce internally generated contextual signals allowing the agent to parse its temporal experience into discrete events and organize learning about them (Fig []. Our contributions are as follows. We show that a network trained on tasks sequentially using traditional weight updates, with task identifiers provided, can be used to identify tasks dynamically by taking gradient steps in the latent space (latent updates). We then consider unlabeled tasks and simply alternate weight updates and latent updates to arrive at Thalamus. parsing sequential experience into events (tasks) and contextualizing its response through the simple dynamics of gradient descent. The algorithm shows generalization to novel tasks and discover temporal events at any arbitrary time-scale and does not require a pre-specified number of events or clusters. Additionally, does not require distinction between training phase or testing phase and accordingly suitable for open-ended learning.
algorithm capable of
MODEL
We consider the setting where tasks arrive sequentially and each task k is described by dataset D of inputs, outputs, and task identifiers (x* y² z²). We examine both settings where the task identifier is and is not available to the learning agent input.
As learning agent, we take function fa with parameters that takes input x and latent embedding
2

---

Published as a conference paper at ICLR 2023
=o(x""
(I)
We define some loss L as a function of predicted output y and ground truth y:
(2)
To optimize the model we define two updates of interest, updating the model parameters 0 in the traditional sense according to their gradient 20 (weight update), and updating the latent embedding
z along the direction of the the enror gradients Az (latent update), by backpropagating gradients from C through fe. but keeping unchanged.
dC d0
dC
A
Az ~
dz
We showcase our algorithm by parametrizing So with vanilla RNN model, as an abstraction of PFC (although note applicability to any neural network architecture). Pre-activation of the RNN neurons vtt) evolve in time as follows (bias terms excluded for clarity):
To dt v(t)= = -00)-vvpp v(t) + W x(t) + W ((v(((((())))) O W
Where is the pre-nonlinearity activations, @ is a relu activation function, W is the recurrent weight matrix, x(t) is the input the network at time W are the input weights, is the latent embedding vector (representing the thalamus), and W are the weights from the latent embedding initialized from Bernoulli distribution Wa.j (pp 0.5) (or rectified Gaussian). These projections had multiplicative effect, consistent with amplifying thalamocortical projections (Schmitt et al. 2017).
Network output was projection layer from the RNN neurons
@(Wout o(Wonts(v(t))
(5)
When task identifiers j were available they were input directly to the latent embedding (2 it) vector as one-hot encoded vectors. When task identifiers were not available, was set to uniform value. The mean squared error loss was used for both latent and weight updates.
Biological-plausibiiic While we rely on backpropagation to obtain gradients with respect to the latent embedding z, these gradients can be learned using a second neural network (Marino et al.) 2018 Greff et al. 2019), which would compare the output to feedback from environment and pro- duce credit assignment signals to the latent embedding. The model would be implemented as two neural networks. One does the task computations and is gated by latent embedding, and second that compares the output to feedback from environment and produces gradients for the latent em- bedding. These learned credit assignment projections then would correspond the corticothalamic projections while the weights from the latent embedding to the RNN would correspond to the thala- mocortical projections. Once the credit assigning corticothalamie projects learned, the structure can now show ability to adapt its behavior with no need for backpropagation or any parameters updates. The circuit now adapts by updating the latent embedding, rather than re-learning parame- ters, which matches modern formulations of adaptive behavior in neuroscience (Heald et al. [021I). We assume learning relies on the credit assignment mechanisms the brain, though these are yet to be fully understood (Lillicrap et al. 2020). For this work, we use backpropagation to maintain generality, as learning the credit assignment network might impart level of domain specialization.
EXPERIMENTS
We first describe continual learning dynamics an RNN model trained with usual weight update paradigm, then we demonstrate the flexibility of taking gradient descent steps in the latent space (latent update) in the resulting pretrained model. Second, we consider unlabeled tasks stream and combine both updates in simple algorithm capable of meeting most desiderata of an ideal continual
3

---

Published as a conference paper at ICLR 2023
Model pretrained with Task IDs
Infer tasks with latent updates
Task ID
Apws
Trak
raat
Z
Output
dosss (2)
z
Output
RNN
RNN
Input
Input
B
clxdm2anti
go
rtanti
Curren Acruncy
100
cbxcm11
uml plyanti
dnmc
M
Latentz
26
m
10032
Accurecy 330
30
40 120 150 160
210 240
270
Figure 2: (A) Schematic of the experiment, model pretrained with task IDs and transitioning iden- tifying tasks through gradients backpropagated from the loss function to update the task embedding z (latent updates). (B) Current accuracy as the latent is updated repeatedly using only one batch from the current task. We switch to the next task once accuracy criterion is reached. (C) current values of the latent task embedding vector during pretraining.
Circular markers indicate the one-hot task used (D) Gradients of loss with respect to z. (E) Average accuracy on all tasks if we were to use the original task IDs used during pretraining, shows no forgetting of previous tasks.
learning agent, including learning disentangled representations that factorize along the structure of the environment (Fig D. We test the model on a set of cognitive neuroscience tasks, and the standardized split MNIST dataset.
The algorithm was implemented in PyTorch and optimized with the Adam optimizer (Kingma & Ba (017)) (learning rate 0.001 for weights updates and latent updates). Simulations were run the MIT OpenMind computing cluster and estimate in the order of 10a GPU hours were used in developing the model. Code available at GitHub. com/ huwmosa Tww.aaaaaaaa
3.1 PRETRAINING RNN WITH TASK IDENTIFIERS, REHEARSAL, AND TRAINING TO CRITERION
We by pretraining our vanilla RNN with task identifier input a collection of 15 cognitive tasks commonly used in neuroscience and cognitive science experiments (Yang et al. 2019), al- though we also present results on split MNIST dataset in a later section. The tasks are composed of trials with varying sequence lengths and input with brief noisy stimuli representing angles. The input-output computations require cognitive abilities such as working memory, perceptual decision making, contextual decision making amongst others, and importantly the same input may require opposite outputs depending on the task (example tasks shown in Fig A[7). We train tasks only to criterion prevent over-training and use rehearsal, ubiquitous feature in animal learning, acquire expertise in new domains. The network achieves high end average accuracy on all tasks successfully (details of pretraining in Appendix A.2 Figs. A7 Ag This RNN trained with one-hot encoded task identifiers became proficient at all tasks at the end of training, and we refer to it as the pretrained network.
they

---

Published as a conference paper at ICLR 2023
3.2 PRETRAINED NETWORK CAN INFER TASKS THROUGH LATENT UPDATES
To handle unlabeled tasks, one can possibly transition this pretrained network through weight up- dates and due to its accelerated ability at learning new tasks and rehearsing previous tasks (Fig A7), it can adapt its weights rapidly to transition between tasks with small number of batches (Fig Ag)
This is analogous to fine-tuning pretrained vision and language models to new data distributions. However, this assumes there is enough data to describe the new distribution fully (which is the case in our simple dataset), and the network would have to constantly update weights as it transitions between distributions, making it prone to forgetting. Most critically, the network is not inferring change in distribution and has no representation of it, but merely adapting the weights to it.
Here, consider an alternative solution using latent updates to adapt to changing data distribution. By consuming only one batch, we take gradient steps in the latent space until accuracy reaches criterion again, and then allow the model to process incoming batches of the new task. Thereby the switch consumes only one batch (although other datasets and environments might require more data to accurately disambiguate context). We note several interesting aspects of this algorithm:
The algorithm correctly switches the RNN to solve incoming tasks without forgetting any of its previous learning as evident by no decrease on its accuracy to perform all 15 tasks given the task identifiers / again (Fig (3).
At moments displays serial search behavior through several possible task labels before converging on the correct one, as by-product of Hamiltonian dynamics (Fig 2C, e.g. in the solution of cxtdm1 from latent update 60 to 130).
The algorithm might not solve the tasks using the embedding vectors it was pretrained on but rather using any latent vectors that perform well given its recent trajectory in the latent space (Fig 2 Interestingly, even though it was pretrained on uninformative latent space (equally distributed one- hot vectors), the latent updates retrieve a rich latent space that is functionally organized and clusters tasks according to similarity and compatibility (Fig 3)
Pretraining the network however assumes one has access to data where data distribution changes are labelled across the correct dimensions. With this motivation, we next describe an algorithm that drops the requirement for any task labels.
3.3 ALTERNATING WEIGHT AND LATENT UPDATES TO DISCOVER TASK REPRESENTATIONS IN A STREAM OF UNLABELED TASKS
We combine weight updates and latent updates in a simple algorithm that can recover the la- tent variables in the environment as it follows the dynamics of the data distribution in time. We start with naive RNN attempting to solve a stream of incoming unlabeled tasks. This net- work was freshly initialized and had no pre- vious training on any tasks. It learns the first task through weight updates, and as the average accuracy begins to increase, any sudden drops from running average by more than a specified amount 10.1) would trigger switching to latent updates. We take a specified number of gradi- ent descent steps in the latent space. If average accuracy has not recovered, the algorithm then goes back to learning through weight updates (see Algorithm AT). More intuitively, once the current latent embedding no longer solves the current task, use latent updates to reach at an- other one. If the task was seen before likely this will solve the current task and you move on. If not, then it is likely new task and gradient de-
15
dmmc
10
multiom
ctxdm2
diygo
dm2
dnmcc
itanti
110
ddlannt
-15
-0
annt
-20
-10
0
10
20
Figure 3: Dimensionality reduction with TSNE of the latent vectors after latent updates each task. Data point colored by the ground truth task ID)
5

---

Published as a conference paper at ICLR 2023
Weight updates to criterion
A
Latent updates to criterion
A Accuracy due to latent updates
100
M
RNN
Thalamus
400
0.4
50
200
0.2
Thalamus
0.0
50
100
50
100
50
100
Task ordinal position
Task ordinal position
Task ordinal position
Figure 4: Thalamus performance on an unlabeled sequence of tasks. A) Comparing number of weight updates needed to solve tasks along the training sequence for Thalamus and baseline RNN. Thalamus reliance on weight updates begins to vanish. B) Latent updates needed for each task in the sequence. C) The change in accuracy after the latent update loop. Latent updates contribute increasingly more to recovering accuracy at task transitions. Results averaged over 20 random seeds.
scent arrived at an untrained embedding. The algorithm then goes back to weight updates and learns the task with the loaded embedding.
These dynamics allow the algorithm to successfully uncover the latent structure of the environment (discrete stream of tasks) without supervision, while remaining invariant to the time scale and not requiring a predefined number of clusters or events. Successful convergence leads the algorithm to require fewer weight updates and eventually, handle the stream of unlabelled tasks through latent updates only 4 example run in Fig ATO).
For quantitative comparison with other models, refer the reader to the next section for exper- iments on split MNIST. For the cognitive task, we noted that continual learning models were not designed to handle contextual behavior. For example, methods for task-agnostic continual learning (Rao et al. 2019 Zeno et al., 2021] AAhille et al. 2018) have a task inference network that maps inputs task identity, which means that transitioning between tasks requiring different output for the same inputs would require training and weight updates. To quantify results on cognitive tasks, we compared Thalamus to an RNN that used weight updates to transition between tasks, as a proxy. We run both Thalamus and the RNN on a stream of unlabeled tasks drawn from 5 cognitive tasks. Thalamus reliance on weight updates to transition between tasks drops rather rapidly after the first few times it sees a task, while reliance on latent updates begins to dominate (Fig(4). While Thalamus required more model runs due to the latent updates, had higher average accuracy due to recovering promptly from dips in accuracy task transitions (Table []. More importantly, we evaluated on the unseen 10 cognitive tasks, with the weights frozen, and Thalamus through latent updates only was able to solve 27% of unseen tasks (see quantification on more standardized benchmark next section) (Table
A vexing finding in brains is that despite internal consistency and intact function, recordings show that, to an outside observer, brain representations seem to drift and change over days and weeks (for e.g. see Marks & Goard )202T), causing difficulties, for example, in designing brain controller interface (BCI) (Mridha et al. 202T). By training classifier to predict task identifier from latent z vector, we show that the internal-gent -eeerrtegeee representations in our model reveal similar dynamics (Fig
3.4 COMPARISON ON EXTERNAL BENCHMARK: SPLIT MNIST
We compare our model to other continual learning methods on a popular continual learning bench- mark that splits MNIST digits into 5 pairs, making each pair into binary classification task (0 or 1, chosen arbitrarily) (LeCun & Cortes, 2010), (referred to as 'domain incremental learning' Ven & Tolias 2019)). Models are shown an image at time from each task and are required to classify it as 0 or 1. Importantly, models have no access to task identifiers and are trained on the 5 tasks sequentially. As Thalamus embodies slightly different paradigm meant operate in dynamic environments, testing its performance on a task requires prompting with a batch from the training- set of the same task, it can adapt its latent embedding z, while keeping the weights frozen (i.e.

---

Published as a conference paper at ICLR 2023
Thalamus vs. RNN, from the last 20 blocks in 5-task sequence learning over 20 random seeds. Solved defined as > 0.8 accuracy
Table 1: Model comparison on NeuroGym tasks:
Model
Weight updates
Avg accuracy
Total updates (compute)
Ratio novel tasks solved
RNN
23.73 + 1.21
0.91 + 0.00
23.73 1 1.21
0.21
Thalamus
3.64 + 0.50
0.96 + 0.00
271.56 55.28
0.27#00.14
Table 2: End average accuracy (tested on all 5 tasks, mean (# SEM) on the split MNIST binary classification tasks. Performance results from Ven & Tolias (2019)
Method
Accuracy
Prompt with a training batch
Upper bound (iid)
98.42 (1: 0.06)
No
EWC
63.95 (#: 1.90)
No
SI
65.36 (#: 1.57)
No
LwF
71.50 (:: 1.63)
No
DGR
95.72 (#: 0.25)
No
RUF
97.31 (= 0.11)
No
Thalamus (Ours)
97.71 (# 0.13)
Yes
it requires feedback from the environment whether as supervised signal or a reinforcement re- ward). Thalamus shows competitive end average accuracy (Table 2 when compared the popular regularization-base bbased methods synaptic intelligence (SI) (Znkee et al., 2017). elastic weight consoli- dation (EWC) (Kirkpatrick al. 2017) although these methods are at a great disadvantage here as they generally require access to task labels and task boundaries. We also compare to three genera- tive rehearsal methods that are of relevance to neuroscience theories: deep generative replay (DGR) (Shin et al. 2017). learning without forgetting (LwF) (Li & Hoiem, 2018), and the model with the current state of the art, replay through feedback (RtF) (Ven & Tolias 2019) (Table 2).. Of note, other more recent works do not report results on MNIST and moved to more complex datasets requiring specific optimizers, data augmentations, and loss functions, making comparison to their core architecture difficult (Pham al. [021] He & Zhu (02T), and as we argue here, not all desider- ata of continual learning have been met for the simpler MNiST-based tasks. Further experimental details can be found in Appendix Al.
3.5 KNOWLEDGE TRANSFER AND CONTEXTUAL BEHAVIOR
To assess knowledge transfer we use the split MNIST task but limit the model to learning only the first 4 tasks and then test on the last unseen task. The model requires one batch from the training set of the last task to update its latent z (while weights frozen), and achieves 91.1 91.14 accuracy on the last task (avg accuracy on all tasks 93.62 (# 0,30)) (example run in Fig ATT]. Through learning the first four tasks the MLP and the latent Z become function parametrizable along relevant dimen- sions to solve related tasks, while the weights are frozen. To highlight this difference compared to other methods we plot the accuracy task 5 models learn the tasks sequentially (Fig (). This generalization relies on balance between adapting via weight updates or latent updates, and conflict between the weight updates required for each task leads to richer semanties offloaded to late up- dates. To showcase these effects, we consider that continual learning literature recommends using stochastic gradient descent (SGD) with momentum and low learning rate limit interference in weight updates thereby limiting forgetting of earlier tasks (Hsu et al., 2019 Mirzadeh et al.) 2020). Thalamus shows the opposite effect, overall accuracy drops (to 83%) but we see disproportionately higher drop in accuracy of the unseen fifth task (drops to 67%) when using SGD with 10-1 learning rate.
Contextual split MNIST: Finally, we introduce a new challenge by taking the first 3 tasks from split MNIST but then adding task 4 and 5 by using images from tasks 1 and 2 but flipping the labels. It is trivial to show that other continual learning methods, which rely on input images to infer context, achieve accuracy around 59%, but Thalamus maintains similar accuracy (96.09 +-083).

---

Published as a conference paper at ICLR 2023
A
Early classifier
B
Representational Drift Index (RDI),
0.35
Late classifier
example neuron
Task Task prodiction accuracy
0.30
RDI
PDG
MOV
0.15
mmy
0.2
0.10
0
Chance level
0.05
4
7
10000
20000
30000
Week
Trial
Figure 5: (A) Representational drift in the network trained on 10 tasks, measured by training logistic regression classifier to predict task identity from the latent z embeddings, either from early (first 1000 batches) or late batches and tested accuracy on all. (B) Drift measured in the responses of a visual cortex neuron in response to natural movie (MOV) or passive drifting gratings (PDG). Adapted from Marks & Goard (02TT)
RELATED WORK
Our work adds to growing continual learn- ing literature, we refer the reader to recent re- views (Hadsell et al., 2020 Parisi et al.) 2019), and only discuss papers of specific relevance to our algorithm. Continual learning
1.0
EWC
unlabeled tasks: several methods have been proposed for continual learning when task boundaries and task identities are not known and most meth- ods use variational inference approach (Rao 2019 Zeno et al., 2021; Achille et al. 2018) to infer tasks and map current data points categorical variable representing task iden- tity. Most methods require learning task in- ference network to a latent space, running in- ference the latent space and decoding to out- put. To handle new tasks, they keep a buffer of unexplained data points and instantiate a categorical task once the buffer fills. In con- trast, Thalamus relies on simply training net- work to solve a given task, and task inference is done by taking gradients through its computa- tions. The shape of the latent space inferred by latent updates fully specified by the forward network weights. Additionally, by using a con- tinuous vector to represent tasks, can handle open ended number of tasks naturally. A critical difference is that these methods unable to handle contextual behavior which maps the same input to different outputs different contexts.
0.9
RP
RtF
task accuraty
0.8
AGEM
ER
Thalamus (ours)
0.7
0.6
0.4
5
Task being leamed
Figure 6: Mcasuring accuracy of the last task only as various models learned the sequence of tasks from to 5. Simulations run as specified in Ven Tolias (2019).
Thalamus is intimately related to the theory of predictive coding in neuroscience (Rao & Ballard 1999) and our gradient descent updates to the thalamus vector are conceptually similar to propaga- tion of prediction errors. The algorithm also connects to adaptive learning from behavioral science Yu et al.] 202I), and to perspectives in neuroscience that consider learning as interaction between learning in the "weight space' and 'activity space' (Sohn al. 2021 Heald et al. 2021). Our work
8

---

Published as a conference paper at ICLR 2023
instantiates these concepts by learning by using gradient descent for latent updates. Similar use of gradient descent to update latent variables was proposed as model for schizophrenia by Yamashital
& Tani (2012) and more recently to retrieve language instructions for cognitive tasks (Riveland & Pouget 2022 Our work uses these methods as a component in a continually learning agent to learn disentangled representations.
The concept of fast weights also prescribes changes in the model parameters at inference time adapt computations recent context (Ba al., 2016), however our model extends this fast changes in the neural space and uses credit assignment and gradient descent at inference time. RL² algorithm uses meta-learning to instantiate an RNN that can eventually adapt itself changing en- vironment (Duan et al.) 2016). In common with Thalamus, it also instantiates a learning system that operates in neural activity space to adapt computations to ongoing task demands. Hypernetworks utilize one network that generates parameters of another network to similarly adapt its computations (Ha et al., 2016). Both algorithms use gradient flow through the task loss to determine the optimal neural activity, however they gradients heavily during the extensive training of the meta-learner (hypernetwork) and use only forward propagation during inference. Our algorithm combines use of gradient updates with forward propagation inference time resulting in much simpler algo- rithm and continued adaptability with no specified training phase. Deep feedback control operates on similar concepts by calculating an inverse of neural network computation and allowing deep feedback controller to dynamically guide the network at inference time (Meulemans et al.) (022). Our model uses proeppropaattion through the neural network and operates on gradients instead of estimating a function inverse. Finally, the algorithm is also related to energy-based models that also rely on optimization- based inference (Mordatch 2018) but our algorithm does not require the lengthy training with negative samples to train the energy- based model.
5 CONCLUSIONS
Thalamus is relatively simple algorithm that can label ongoing temporal experience and then adapt future similar events by retrieving those internally generated labels. A notable limitation is since we use simple latent space, the form of knowledge transfer between tasks is quite limited con- ceptually to adapting the same computations if they apply to a new or old task. richer form of knowledge transfer would emerge with latent space that allows for compositionality over compu- tational primitives. Future work will consider latent spaces with hicrarchical or recurrent structures. Another possible extension is once the gradients of the latent embedding approximated with second neural network, the algorithm can be deployed in hardware, allowing it to show adaptive behavior and 'learn' without requiring backpropagation or weight updates. The algorithm shows several conceptually interesting properties, and the representations it might learn from more com- plex and naturalistic datasets will be of interest to compare to brain representations.
ACKNOWLEDGMENTS
Thanks to Matthew Nassar, Brabeeba Wang, Ingmar Kanitscheider, Robert Yang, and Michael Ha- lassa for insightful discussions and valuable comments on earlier drafts.
REFERENCES
Alessandro Achille, Tom Eccles, Loic Matthey, Christopher P Burgess, Nick Watters, Alexander Lerchner, and Irina Higgins. Life-long disentangled representation learning with cross-domain latent homologies. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, 18, pp. 9895-9905, Red Hook, NY, USA, December 2018. Curran Associates Inc.
Jimmy Ba, Geoffrey Hinton, Volodymyr Mnih. Joel Z. Leibo, and Catalin Ionescu. Using Fast Weights to Attend to the Recent Past. arXiv: 1610.06258 [cs, star), December 2016. URL http: arXiv: 1610.06258.
// //arxiv.com www.or///aao.. abs 610. 0. 06258
Farran Briggs. Role of Feedback Connections in Central Visual Processing. Annual Re- view of Vision Science, 6:313-334, September 2020. ISSN 2374-4650. doi: 10.1
mmnnee-visinn1122199-87777

---

Published as a conference paper at ICLR 2023
Anne Collins and Etienne Koechlin. Reasoning, Learning. and Creativity: Frontal Lobe Func-
tion and Human Decision-Making. PLOS Biology, 106)se1001293, March 2012. ISSN
1545-7885. doi: http:///oooaa....o1000000 plosbi wwwgo//attlllee/e/oooooo Library of Science.
URL hKttt:///oouuraalsss /
Publisher: Public
Guy Davidson and Michael C. Mozer. Sequential Mastery of Multiple Visual Tasks: Net-
works Naturally Learn to Learn and Forget to Forget. pp. 9282-9293, 2020. URL
VisaaVisulllaaeeeeeeeeee Tew.kas_eeewrrssaaaarrreeee Learn to CVPR 2020 paper.
Sequential Mastery_offfaaoaaa
Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and Pieter Abbeel. RL$
Fast Reinforcement Learning via Slow Reinforcement Learning. M.Kv:1661..02777 [cs, stat), November 2016. URL http: http::::/axii..og///a.. htrg/abs/1611.. 02779 arXiv: 161 1.02779.
Vishwa Goudar, Barbara Peysakhovich, David Freedman, Elizabeth A. Buffalo, and Xiao-Jing Wang. Elucidating the neural mechanisms of -eaminngtt--Laam bioRxiv, 202.0.004..0 January 2021. doi: 10.11101/20211000024555777
httpt::://biiii..ooo////ooo////e early/2021 /9/ 02.00//222222 02..99.....57..... 455707 www. abstract]
Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Christopher Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi Multi--Object Representation
Learning with Iterative Variational Inference. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2424 2433. PMLR, June 2019. URL https: html
//proceedings.. m
press /v97/v97/// httg//geffffl// prr
David Ha, Andrew Dai, and Quoc Le. HyperNetworks. arXiv:1609.09100 [cs], December 2016.
URL
09106) arXiv: 1609.09106.
Raia Hadsell, Dushyant Rao, Andrei A. Rusu, and Razvan Pascanu. Embracing Change: Con- tinual Learning in Deep Neural Networks. Trends in Cognitive Sciences, 24(12):1028-144 doi:
December 2020. ISSN 1364-6613.
2020.09..04 URL hhtpps:
sciencedirect com/ secssciennn
Michael M. Halassa and Sabine Kastner. Thalamic functions in distributed cognitive control. December 2017. ISSN 1546-1726. doi: 10.1038/
Nature Neuroscience, 20(12):1669-167
141593-017-0002-11
Jiangeng He and Fengqing Zhu.
Unsupervised Continual Leanning Via Pseudo Labels. 2ttp:///aab/111011.. arXiv:
CoRR, abs/2104.07164, 2021. URL 2104.07164.
James B. Heald, Máté Lengyel, and Daniel M. Wolpert. Contextual inference underlies the leamn- ing of sensorimotor repertoires. Nature, 600(7889):4894499 493, December 2021. ISSN 1476- 4687. doi: 10. 158-002-41---- URL https: http://wwwwwww...tt//tttt/t/ s41586-
nature. com/articles/ Number: 7889 Publisher: Nature Publishing Group.
Santiago Herce Castañón, Pedro Cardoso-Leite, Irene Altarelli, C. Shawn Green, Paul Schrater, and Daphne Bavelier. A mixture of generative models strategy helps humans generalize across tasks. bioRxiv, pp. 2021.02..6.4331500 January 2021. doi: 10. 10.1201/222 16. 431506.
431506 abstract
Yen-Chang Hsu, Yen-Cheng Liu, Anita Ramasamy, and Zsolt Kira. Re-evaluating Continual Learn- ing Scenarios: A Categorization and Case for Strong Baselines. arXiv:1810.12446 [cs], January 2019. URL
1810. 1810.12488
arXiv: 1810.12488.
Ali Hummos, Bin A. Wang, Sabrina Drammis, Michael M. Halassa, and Burkhard Pleger. Tha- lamic regulation of frontal interactions in human cognitive flexibility. PLOS Computational Biology, 18(9):e101050) September 2022. ISSN 1553-7358. doi: 10. /0.
1010500. URL / plos. org 1371/journal wpt.maaalpppppppp...m 1010500
Publisher: Public Library of Science.
10

---

Published as a conference paper at ICLR 2023
Jaime S. Ide and Chiang-shan Ray Li. A cerebellar thalamic cortical circuit for enror-related cogni- tive control. Neurolmage, 54(1): January 2011. doi: 1010006/..cuuoiieeee000000077100 URL https: http://wwwwwpw///ppttppppppptttt nchi. http:///ww...taee///a/aaatttte/a Publisher: NIH Public Access.
Anthony M. V. Jakob, John G. Mikhael, Allison E. Hamilos, John A. Assad, and Samuel J. Gersh- man. Dopamine mediates the bidirectional update of interval timing. Technical report, Harvard University, November 2021. URL https: www. www.:.booxxiv.. org/ pontentt/00 1101/ 202111111 02. 4 68 0 3v11 Company: Cold Spring Harbor Laboratory Distributor: Cold Spring
Harbor Laboratory Label: Cold Spring Harbor Laboratory Section: New Results Type: article.
Kevin Johnston, Helen M. Levin, Michael J. Koval, and Stefan Everling. Top-Down Control- Signal Dynamics in Anterior Cingulate and Prefrontal Cortex Neurons following Task Switch-
ing. Neuron, 53(3):553445 February 2007. ISSN 0896-6273. doi: 10.1016/j.neuron.
12.023, URL httpps//wwet wwwl. htt.:/e.w.......eeoo///aaattta/a
6273 (06)
01029-4 4. Publisher: Elsevier.
Ta-Chu Kao, Mahdieh S. Sadabadi, and Guillaume Hennequin. Optimal anticipatory control as theory of motor preparation: A thalamo-cortical circuit model. Neuron, 109(9):1567-- 1581.e12, May 2021. ISSN 0896-6273. doi: 10. me0o://.eeerrn222200..00000 URL https:
Diederik Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. 142.6994 [cs], January 2017. URL htt.a///ce....ccc.cc..e.ooooe//n
7 //www. www.sscieeeeeiieeoooo http://eeee.ceeeeeeeeeeeee/eeeee
6980
arXiv: 1412.6980.
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A.
Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Baabska-Barwinska Demis Hass- abis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell. Overcoming catastrophic forget- ting in neural networks. Proceedings of the National Academy of Sciences, 1 114(13):321-3352, March 2017. doi: 10.1073/pnas.16 161 18351 14. URL pttpps: /ttps://wwwr /www. pnas , wtt////////// Publisher: Proceedings of the National Academy of Sciences.
1073/pnas. 61 835114
Quoe V. Le, Navdeep Jaitly, and Geoffrey E. Hinton. A Simple Way Initialize Recurrent Networks of Rectified Linear Units. arXiv: 1rXxvv:504.000 554.00991 [cs], April 2015. URL http://arxiv.org
abs/1504. 00941
arXiv: 1504.00941.
Yann LeCun and Corinna Cortes. MNIST handwritten digit database. 2010. URL WRRhhttp///yan
lecun.
Zhizhong Li and Derek Hoiem. Learning without Forgetting... IEEE Transactions on Pattern Analysis and Machine Intelligence, (0(12):2935-2944 December 2018. ISSN 1939-3539. doi: 10. 1 109/ Intelligence.
TPAMI.2017.2773300 Conference Name: IEEE Transactions on Pattern Analysis and Machine
Timothy Lillicrap, Adam Santoro, Luke Marris, Colin J. Akerman, and Geoffrey Hinton. Back- propagation and the brain. Nature Reviews Neuroscience, 21(6):335-344 June 2020. ISSN 1471- 0048. doi: 10. 10010008//15333-0000027777 Number: 6 Publisher: Nature Publishing Group.
URL (otthhtps:/ www. nature. www/ aottieetllee//////
02
Laureline Logiaco, L. F. Abbott, and Sean Escola. Thalamic control of cortical dynamies in a model of flexible motor sequencing. Cell Reports, 3509): June 2021. ISSN 2211- 109090. 10900. URL /htpps:///// /ttp:://wwww... www.ww.....c...................e rect. com/
1247. doi: 10. 10.10066).eelrep.2221.... 2021 sci
Farshad A. Mansouri, Kenji Matsumoto, and Keiji Tanaka. Prefrontal Cell Activities Related to Monkeys' Success and Failure in Adapting to Rule Changes in a Wisconsin Card Sorting Test Analog. Journal of Neuroscience, 26(10):745-2755 March 2006. ISSN 0270-6474, 1529. 2401. doi: 10. 10.152/INEROSSCI.2338 5228055..060 URL https: www. ineurosci.ooo org/ content/26/133 / 17/5 Publisher: Society for Neuroscience Section: Articles.
Joseph Marino, Yisong Yue, and Stephan Mandt. Iterative Amortized Inference. Proceedings
of Machine Learning Research, 80:3403-3412, July 2018. ISSN 1938-7228. URL https:
11

---

Published as a conference paper at ICLR 2023
//resolver: wtt../ccc..cccooaaaaaaaooooooo..
Conference
Name: 35th International Conference on Machine Learning Meeting Name: 35th International Conference on Machine Learning Place: Stockholm, Sweden Publisher: PMLR.
Tyler D. Marks and Michael J. Goard. Stimulus-dependent representational drift in pri- mary visual cortex. Nature Communications, 12(1):5169, August 2021. ISSN 2041- 1723. doi: 10. 641467-021-25458
URL (tttt::: /www nwwuraataee.... wtt.sarr..co.m.oommo cww.artiiarss////////ooom. Number: Publisher: Nature Publishing Group.
Michael McCloskey and Neal Cohen. Catastrophic Interference Connectionist Net- works: The Sequential Learning Problem. In Gordon H. Bower (ed.), Psychology of Learn- ing and Motivation, volume 24, pp. 109-165. Academic Press, January 1989. doi:
1016550099-7222118880333332 URL htttps:/wwet www.s/eeeeeeeee..coeeeeeeeeeeeee atticcle/ss1/50797742116666
Alexander Meulemans, Matilde Tristany Farinha, Javier Gareía Ordóñez,, Pau Vilimelis Aceituno, João Sacramento, and Benjamin F. Grewe. Credit Assignment in Neural Networks through Deep Feedback Control. Tow:2106.07885 [cs], January 2022. URL
org. abs/
2106. 07887
XrXiv: 2106.07887.
Seyed Iman Mirzadeh, Mehrdad Farajtabar, Razvan Pascanu, and Hassan Ghasemzadeh. Un- derstanding the Role of Training Regimes in Continual Learning. Technical Report arXiv:2006.06958. arXiv, June 2020. URL http: /tap:v///rivv.... org worg/absse abtp//bbb//00../ 06958 arXiv:2006.06958 [cs, stat] version: type: article.
ggor Mordatch. Concept Learning with Energy-Based Models. arXiv: /81102280 [es), November 1811. 02486
2018. httt:://arxiv.org/e
arXiv: 181 1.02486.
M. F. Mridha, Sujoy Chandra Das, Muhammad Mohsin Kabir, Aklima Akter Lima, Md. Rashedul Islam, and Yutaka Watanobe, Brain- Commp-Computer Interface: Advancement and Challenges. Sensors (Basel, Switzerland), 21(17):5746, August 2021. ISSN 1424-8220. doi: 10.3390/s21175466 URL https:/ ncbi nlm. nih. hwtp//me//arriicle///MM55/5/
Arghya Mukherjee, Norman H. Lam, Ralf D. Wimmer, and Michael M. Halassa. Thalamic circuits for independent control of prefrontal signal and noise. Nature, pp. 1-8, October 2021. ISSN 1476- 4687. doi: 10. 1038//4 1586-021--1111111 021002200--0333 s4158 241586-222 001-6660211005---
URL https: www. nat ure. commeariilee//////
German I. Parisi, Ronald Kemker, Jose L. Part, Christopher Kanan, and Stefan Wermter, Continual lifelong learning with neural networks: review. Neural Networks, 113:54-71, May 2019. ISSN 0893-6080. doi: 10. n00.00161.ccmmet2211..0000000
URL https: /ttp://wwwwww///o/ www sciencedirectt.
comm http://ccccceeeerr.aaaa//////1.0
Quang Pham, Chenghao Liu, and Steven Hoi. DualNet: Continual Learning, Fast and Slow. In
Advances in Neural Information Processing Systems, volume 34, pp. 16131-16144, Curran As- sociates, Inc., 2021. URL htttpss::/pppcceedaasss....
neurips. wtw.:eeripss.cpppeeeeeeeeeeeeeee hctppaaprr//22///aa...
html
Dushyant Rao, Francesco Visin, Andrei A, Rusu, Yee Whye Teh, Razvan Pascanu, and Raia Hadsell. Continual Unsupervised Representation Learning. 191iv: 1448/ [cs, stat), October 2019.
URL http::::/wx....or///.ao.aooo wttp:lvv.rrg/ab//111900 14481
14481
arXiv: 1910.14481.
Rajesh N. Rao and Dana H. Ballard. Predictive coding in the visual cortex: a functional in- terpretation of some extra-classical receptive-field effects. Nature Neuroscience, 2(1):79-87, January 1999. ISSN 1546-1726. doi: 10.1038/4580. URL hittps://www..r
com/
articles/nn019g
Rajeev V. Rikhye, Aditya Gilra, and Michael M. Halassa. Thalamic regulation of switching between cortical representations enables cognitive flexibility. Nature Neuroscience, 21(12):1753-1765 December 2018. ISSN 1546-1726. doi: 10. 10.1038/<41533--08--22666
12

---

Published as a conference paper at ICLR 2023
Reidar Riveland and Alexandre Pouget. A neural model of task compositionality with natural lan- guage instructions. bioRxiv, pp. 2022.02.2..481299 January 2022. doi: 10. 10.1101/2022.02.22 481293. http:::/ilox....ccr//ootteetetee 481293 abst ract.
Nicolas P. Rougier, David C. Noelle, Todd S. Braver, Jonathan D. Cohen, and Randall C. O O'Reilly. Prefrontal cortex and flexible cognitive control: Rules without symbols. Proceedings of the Na- tional Academy of Sciences, 102(20):7338-7777 70117733337733,, May 2005. doi: (0.1073/pnas.0502/ 0...05552252222
URL https: //wwww wwwwwppaa..........mm..m/....... pww.ppaas.a.. org/ /1. 1073/pnas 0502455102 ceedings of the National Academy of Sciences.
Publisher: Pro-
L. lan Schmitt, Ralf D. Wimmer, Miho Nakajima, Michael Happ, Sima Mofakham, and Michael M. Halassa. Thalamic amplification of cortical connectivity sustains attentional control. Nature, 545
(7653):219-22) May 2017. ISSN 1476-4687. doi: 10. 1038/nature22073.
Hanul Shin, Jung Kwon Lee, Jachong Kim, and Jiwon Kim. Continual Learning with Deep Gen- erative Replay. In Advances in Neural Information Processing Systems, volume 30. Curran As- sociates, Inc., 2017. URL https: 3dba
Defbe
html
Hansem Sohn, Nicolas Meirhaeghe, Rishi Rajalingham, and Mehrdad Jazayeri. A Network Per- spective on Sensorimotor Learning. Trends in Neurosciences, 1703);170-181, March 2021. ISSN 1878-108X. doi: 10.1016/j.tims 1...0166.tins.222.111111 1020.11.007.
Gido M. de Ven and A. Tolias. Three scenarios for continual learning. ArXiv, 2019.
Tom Veniat, Ludovic Denoyer, and Mare' Aurelio Ranzato. Efficient Continual Learning with Modular Networks and Task-Driven Priors. arXiv:20/2.12633 [cs]. February 2021. URL
httt::/arxiv.oort
2012 12631
arXiv: 2012.12631
version: 2.
Jing Wang, Eghbal Hosseini, Nicolas Meirhaeghe, Adam Akkad, and Mchrdad Jazayeri. Reinforce- ment regulates timing variability in thalamus. eLife, 9xe55872, December 2020. ISSN 2050-084X. doi: 10.7554/eL.ff.55555 eLife Sciences Publications, Ltd.
htl.ittts://doo..rrttt 7t54//1175554111111.1 55872, Publisher:
Mien Brabeeba Wang and Michael M. Halassa. Thalamocortical contribution solving credit as- signment in neural systems. urXiv:21044 004444 [q-bio], April 2021. URL htttp://arxiv org/ ht5/abs/21000 2004. 01474
arXiv: 2104.01474.
Yuichi Yamashita and Jun Tani. Spontaneous Prediction Error Generation in Schizophre- nia. PLOS ONE, 7(5)xe37443, May 2012. ISSN 1932-6203, doi: 10.137
0037843. URL https: /tttp:://jouznals journal. pone. 0037843
journals ..s.s.pli. plos. http://poossmeerrttteeeeeeeeeeee
Publisher: Public Library of Science.
Guangyu Robert Yang. Madhura R. Joglekar, H. Francis Song, William T. Newsome, and Xiao-Jing Wang. Task representations in neural networks trained to perform many cognitive tasks. Nature Neuroscience, 22(2):297-306. February 2019. ISSN 1546- 1726. doi: 10.1038//4139--01--03300 641593-018-0310-2
URL https: /wwwwww... www..nature... com/ articles/
Linda Q. Yu, Robert C. Wilson, and M. Nassar. Adaptive learning is structure learning time. Neuroscience & Biobehavioral Reviews, 2021. doi: 10. 2...10166.eebiirre2000......
Friedemann Zenke, Ben Poole, and Surya Ganguli. Continual Learning Through Synaptic Intelli- gence. Technical Report arXiv: 703.00420, arXiv, June 2017. URL
abs/ 04200
Xrviv: 703. 04200 [cs, q-bio, stat] type: article.
Chen Zeno, Itay Golan, Elad Hoffer, and Daniel Soudry. Agnostgi Continual Learning Using Online Variational Bayes With Fixed-Point Updates. Neural Computation, 33(11)33139-3177,
October 2021. ISSN 1530-888X, doi: 10. 10.1162/neco a 01430.
13

---

Published as a conference paper at ICLR 2023
A APPENDIX
A.1 TRAINING DETAILS AND HYPERPARAMTERS FOR THIALAMUS TRAINING WITH NO TASK IDENTIFIERS
For the cognitive tasks neurogym dataset, we used 356-dimensional RNN and a dt of 100ms and tau of 200ms. Recurrent weight matrix initialized as an identity matrix (Le et al. 2013 Yang et al. 2019). Multiplicative binary projections from z to the RNN had sparsity of 0.5. The training sequence used 5 randomly sampled cognitive neurogym task. We initially present the tasks to be learned in the same order twice for 200 batches cach and allow the model to adapt with weight updates only. The model then sees the same sequence repeatedly. The first 3 repetitions are with block size of 200 and then we lower it to 100 for efficiency and easier visualization.
For the split MNIST dataset and comparison with other models, employ the same architecture and hyperparameters reported in Ven & Tolias (2019) [Hsu et al. (2019). The classifer network had two fully-connected layers with 400 ReLU units and 2-unit output layer. We used 10 dimensional latent z embedding space that projected multiplicative binary gates with 0.8 sparsity. The same gat- ing weights were used for both layers for simplicity. For this dataset we follow a more incremental sequencing of the tasks to align better with other methods on the same dataset. Tasks are incremen- tally added and the tasks seen so far are rehearsed twice after each new task. Only latent updates were used when evaluating on unseen tasks. See Fig ATO for an example run.
A.2 PRETRAINING THE NETWORK FOR SECTIONS 3.1 AND 3.2 WITH TASK IDENTIFIERS
We provide an example learning session of sequence of 15 tasks with rehearsal and training to criterion (Fig A8). During training, accuracy on all other tasks was tested every 10 batches using the task IDs as imput indicate the testing task. Ablations show that task identifiers, training to criterion and rehearsal are the minimal components for this network and this dataset. As the network learns more tasks we see decrease in the batches needed to acquire new tasks or rehearse early ones consistent with previous literature (Goudar al. 2021 Davidson & Mozer 2020).
A.3 RAPID ADAPTATION OF PRETRAINED NETWORK
We provide an example of training the network using weight updates. After pretraining with task IDs it can now transition its behavior without task IDs, by adapting its parameters rapidly to handle oncoming tasks and requiring few updates to achieve criterion performance on each task (Fig AD). But see main text for discussion of the limitations of this approach to adaptable behavior.
A.4 THALAMUS ALGORITHM
For completeness we provide the algorithm details below in Algorithm 1
Clustering the latent embedding values Thalamus assigned to tasks shown a somewhat difficult to interpret distribution, with tasks having two embedding clusters, sometimes shared with another task. Some functional organization along task computations but also tasks seem to have limited overlap (Fig AT2). Further study of the topography of the latent surface under different conditions and creating more compositional structures will be the subject of future studies.
To clarify the need for the latent updates to be triggered judiciously and only when an environment change is suspected, we lesion the algorithm by removing the accuracy criterion and simply updating the weights and then the latents every trial (Fig AT3). The two updates seem to disrupt each other, with weight updates need to transition between tasks never approaching zero (Fig A13A), and on average, the latent updates seem to decrease accuracy (Fig AT3C).
A.5 FURTHER EXAMINATION OF LATENT EMBEDDING TO RNN WEIGHITS W
For the results reported in the paper we draw W from Bernoulli distribution with p=0.5 for neu- rogym tasks and 0.2 for split MNIST. To examine relaxing such a requirement we ran Thalamus on sequence of 5 neurogym tasks using Wz drawn from Gaussian distribution with mean of 1,
14

---

Published as a conference paper at ICLR 2023
Table 3: hyperparameters used
Parameter
Neurogym RNN
split MNIST MLP
Weight update learning rate
0,001
0.001
Latent Update learning rates
0.01
0.001
Latent Update L2 regularization
0,0001
0.1
Latent size
15
10
Max latent updates
min(1000, batchno)
W sparsity
0,5
0.8
Loss function
MSE
MSE
prevent multiplying the RNN units inputs by values each time step leading exponential de- cay. Additionally, since having negative gating values made little sense numerically or biologically, rectified the values and any negative values to We found that having W drawn from rectified Gaussian distribution with variance had similar performance to the baseline model with Bernoulli W (Fig AT4). We show that sparsity matters and multiplying the Gaussian values with Bernoulli improves performance further (Fig A[5), presumably so a proportion of neurons and their weights are dedicated some tasks. Additionally, the results reported used a fixed W and trained the other RNN parameters {W", W Woot), }. While this initially seemed to stabilize the learning problem to prevent the latent embedding meaning from changing rapidly, it the last version of the model, allowing all parameters to be trainable showed reasonably similar performance, and fixing some of the weights is not a necessary assumption (Fig AT6).
Algorithm Thalamus algorithm
Init z to uniform value
repeat
Sample x= y from D* in sequential blocks g* (x*
Accuracy L(g", ,y*)
Running mean accuracy Accuracy 2 0.1 then
while remaining latent updates AND Accuracy < Criterion do latent updates ramp up
to max of 1000
Az < a
dc + L2reg de
b a learning rate was 0.01
yhh fx z)
Accuracy = L(sk )*)
end while
end if
Running mean accuracy Accuracy 2 0.1 then AW + W de
>> Update weights only if accuracy not recovered
TW
end if
Update Running mean accuracy
until Converged: not requiring any weight updates for tasks tested
15

---

Published as a conference paper at ICLR 2023
A
B
Tasks being trained -
dm2
dnmo
dm1
RNN
drne dm2 rtanii
Fixed
(PFC)
Leamed
Accuracy on other tasks
and
Januarr FAL
dmm
foomection
dmn1
Iisine
Latent Task embedding (Thalamic representation)
500
1000
1500
Trials (x100)
C 1.0
E
+
1.0
5000
End avg accuracy
0.8
OHH
End avg accurany
0.8
4000
T
0.6
Total training
3000
0.2
1000
0.0
0.0
TTC
+TTC
+TTC
-TC
+TTC
Shuffled
Sequential training
Shuffed training
Sequential
+Reh
+Reh
+Reh
training
training
+IDs
+IDs
+IDs
+IDs
IDs
F
140
300
learning
120
reheersal
70
rehearsal
1st rehearsal
learning
Trials to criterion
Trials to criterian
60
criterion
250
2nd rehearsal
100
200
80
50
200
rehearsal
rehearsal
40
150
60
100
40
20
50
20
Task ordinnal position
Task ordinal position
Task ordinal position
Task ordinal position
Figure 7: (A) Network schematic. (B) Accuracy all tasks the network learns them sequentially. (C) Accuracy at the end of the training session averaged across all 15 tasks, measured over 10 ran- dom orderings the tasks. The axis labels indicate which components were used in the training paradigm. TTC indicates training each task to criterion, Reh, indicates rehearsal of previous tasks after each new task, and IDs indicates whether task identity was input the network. The combi- nation of the three components was necessary for optimal accuracy. (D) For comparison, we train network using shuffled data until achieves similar accuracy and (E) compare the number of batches needed using shuffle training and sequential training. (F) Decreased batches to criterion as a function the order of which task was learned in the initial learning and in subsequent rehearsals. The right subplot is the same lines overlaid to show progression.
16

---

Published as a conference paper at ICLR 2023
chodm
COUI STTE 0993
1030
1500
2000
1000
Trinkc
Figure 8: Training all 15 tasks with rehearsal and training criterion.
cixdm1
dms
dmc
current accurasy
1.00
00
2oddm2
tlm
0.75-
0.50-
0.25
0.00
Latent embedding
142288 6 442 2
0.085
Accuracy task DDs
1.00
0.75
0.50
0.25
0.00
120
150
180
210
Trials (100)
Figure 9: Pretrained network adapts rapidly to changing tasks. (A) solving stream of tasks by rapidly fine-tuning the weights. (B) Latent task embedding z kept uninformative and uniform with network no longer receiving task IDs as input. (C) the network begins to no longer respond to the pretraining task IDs, and accuracy all tasks drops with progressive adaptation of the weights.
17

---

Published as a conference paper at ICLR 2023
dms
dnms
rtgo
ctxdm1 dlyanti
current task accuracy
0.5
Latent 2 vector
12 6
0.75
0.50
0.25
Weight updates to cntterion
100
XXX XB0
XXX XXX
1000
Latent updates
Weight updates
100
50
updates
500
2000
4000
6000
8000
10000
12000
14000
Trials
Figure 10: Newly initialized network learning unlabelled tasks using the Thalamus algorithm. Train- ing begins with exposure to each task for 200 batches, and then for 100 for visualization purposes. Note that as opposed to Fig A here we only show batches new data points and omit latent update current task. Vertical lines label transition between tasks, color coded to task identity. (B) latent vector representation. (C) Weight updates (left axis) and Latent updates (right) to reach accuracy criterion.
loops (A) Accuracy
18

---

Published as a conference paper at ICLR 2023
A
Task
Task 2
Task
Task 3 Task
1000
2000
3000
4000
5000
B
Task 2
Task
Task
Task
current task accuracy
1.00
0.75
0.50
Latent Z vector
0.15
0.10
8
O
Weight upd ates to criterion
M00 x.
XXXXXX
X 300 x
Xxxx
XX XXXX X
1000
Latent updates
XXc X
Weight updates
20
x
Latent updates
500
1000
2000
3000
4000
5000
Batches
Figure 11: Example run of Thalamus on split MNIST tasks. We provide an example where the 5th
task "task 3" was not seen through out the training sequence but performance it increases after model forms latent representations of the other tasks.
80
'dlygo
mms
40
dmc
20
dlyanti
ebxdm
.mml %
Itanti
20
90
anti
multidm
40
drmm
dnms
-60
-40
-20
40
60
Figure 12: Diagram of T-SNE reduced latent embeddings learned from running Thalamus 10 cognitive tasks. We froze the weights after 30 rehearsals to prevent semantic drift, and gathered data with weight updates.
19

---

Published as a conference paper at ICLR 2023
Weight updates to criterton
A
Latent updates to criterion
C
Accuracy due to latent lstent updates
200
0.0
100
-0.2
-0.3
100
100
100
Task ordinal position
Task ordinal position
Task ordinal position
Figure 13: Showing the effects of updating weights and latent every batch instead of the accuracy criterion. A) The model continues to rely weight updates and never approaches zero weight updates. B) Latent updates in each block simply reflect the length of the task block the latent are updated every batch. C) The change in accuracy after the latent update indicates mainly drop in accuracy after the latent updates. Results averaged over 20 random seeds.
Weight updates to criterion
A
Latent updates to criterion
A Accuracy due to latent updates
75
Bern(0.5)
N(1.1.0)
50
500
25
250
0.2
Bem(0.55
N(1,1.0)
0
100
0
100
100
Task ordinal position
ordinal position
Task ordinal position
Figure 14: Thalamus performance metrics on 5 neurogym tasks. Results from 20 seeds. Comparing
model with W drawn from Bernoulli with (p 0.5) vs. rectified normal distribution with mean and variance 1.
updates to criterion
A
Latent updates to criterion
C
& Accuracy due to latent updates
75
#(0.5)
Weight updates
25
0.2
(I0.5)
(10.5)A71.1.0)
0.0
0
100
100
100
Task ordinal position
Task ordinal position
Task ordinal position
Figure 15: Same as Fig A[4 but the rectified normal distribution with mean and variance further sparsified by multiplying with Bernoulli (p=0.5).
Weight updates to criterion
A
Latent updates to criterion
C
5 Accuracy updates
75
fixed learned
750
0.4
50
250
fixed
learned
0
0
due
0.0
100
Task ordinal position
Task ordinal position
Task ordinal position
Figure 16: Thalamus performance metrics on 5 neurogym tasks. Results from 20 seeds. Comparing model 'Wz fixed', which is the baseline model used in the paper, to the model with W allowed to be trainable.
20

---

Published as a conference paper at ICLR 2023
response
rtgo response
dlygo
dlygo response
No02
Impur Input Koot
21
16:41210
18
16141230
221
12
121518
dm1
dm1 response
ctxdm1
ctxdm1 response
dms
dms response
Moo2
Inpur Nos1
21
16143210
1614
dmc
dmc response
dmz
dm2 response
ctxdm2
ctxdm2 response
No52
Inpur Most
24
24
R
multidm
multidm response
rtanti response
anti
response
Mo32
Input Nodt
163412200 10
30
24
24
diyanti
dlyanti response
response
response
InpuR Mods
16:412210 2230
24
24
27
20
discretized time
discretzzed time
Figure 17: Visualizing the inputs and the ground truth responses for neurogym tasks. The first input dimension has a fixation signal, where agents are to withhold responses and attend to two sensory modalities (mod1, and 2). Inputs come pulse of various lengths at angle in one modality. In 'Go" tasks, agent is to simply report the same angle as output in the same modality. "Anti" tasks require reporting the opposite angle. Delay tasks have the input removed and response is delayed. In decision making tasks ('dm') Inputs are noisy across two modalities and the correct response to infer the peak of the noisy precept. Which modality report depends context, and the type of task (for e.g. in ctxdm2, input in mod2 should be ignored). Delayed non-match to sample ('dnms') ) and similar tasks require comparing to inputs held in memory and responding accordingly. See Yang et al. (2019) for further descriptions
21


Reviews:
[{"review_id": "hT3w6KBuVAv", "reviewer": "Reviewer_W9vo", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The work couples an RNN with a latent space representation that can infer the current task in a continual learning setup. The latent space is inspired by the role of thalamus, and enables the RNN to adapt to new tasks by changing the latent representation rather than the internal dynamics of the RNN. ", "summary_of_the_review": "Good paper, but several aspects (presentation, discussion on plausibility, maybe better comparison with other methods, and adding code) could be improved. I put the current score to 6 but I can raise it if my concerns are addressed.\n\n**UPDATE**: updated the score from 6 to 8 as all of my concerns have been addressed.", "strength_and_weaknesses": "### Strengths\n\nA clear algorithm for continual learning that involves a thalamus-like part and offloads some of the weight updates in the main network to that part.\n\nGood experimental performance.\n\nA clear way to learn from an unlabeled stream of tasks (by triggering learning of z with large accuracy deviations) that also enforces stability in the RNN weights (as z is updated first in those cases).\n\n### Weaknesses\n\nAs the goal of the paper is biological plausibility, it needs a discussion on how Thalamus could be implemented in the brain. As I can see it, the main issue would be credit assignment to the latents, requiring backpropagation through the RNN.\n\nPresentation is lacking in parts (mostly the plots) but can be easily improved.\n\nNo code attached.\n\nComparison with other methods could be improved (see in bold below).\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "_6n4YGKyRXF", "reviewer": "Reviewer_5VKC", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents a model for continual learning that is inspired by thalamocortical networks in the brain. Neuroscience studies have revealed that (i) the prefrontal cortex (PFC) shows representation of relevant task variables, and (ii) the thalamus shows representations of the task being performed. Further, is thought that the connections from the thalamus to the PFC gate the computations performed by the PFC by selecting the relevant representations. \n\nIn the proposed model, a RNN and a latent embedding vector serve as abstractions of the PFC and the thalamus, respectively. The projections from the latent embedding to the neurons in the RNN correspond to thalamocortical projections. The model is trained on a sequence of tasks with alternating RNN weight and latent embedding updates, resulting in a network capable of parsing a sequence of inputs into the appropriate contexts. \n\nFinally, the performance of the model was evaluated on cognitive tasks from (Yang et al., 2019) and on the split MNIST task. ", "summary_of_the_review": "The paper is well written and the main concepts are clearly conveyed. The experimental results also highlight the potential of this approach. Overall, the proposed model is promising and could potentially lead to more interesting avenues of future work. \n", "strength_and_weaknesses": "The proposed model is a simple RNN with projections from a latent embedding vector. But by training this model using a simple procedure of alternating RNN weight and latent embedding updates, we obtain a network that is capable of learning tasks in the absence of any explicit task labels. These results are quite interesting, and as the authors mention in the conclusion there are a lot of interesting future directions (e.g., using a latent space that allows for compositionality over computational primitives). \n\nThe following are some comments/questions I had:\n-  Is there a specific reason why the weights from the latent embedding are drawn from a Bernoulli distribution? \n- Figure 3 shows that the latent embeddings cluster according to task id with the pretrained network. I suppose the same could be expected in the thalamus network trained from scratch? \n- What is the loss function used to train the networks? \n- I wonder if there is a better criterion for switching to latent updates than just a drop in accuracy. This requires the network to always know how well it is performing at the task? \n- Would simultaneously optimizing both the RNN weights and the latent embedding vector do worse than the alternating updates? ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Yj-z-gF4wCM", "reviewer": "Reviewer_ugKD", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors proposed a RNN with latent embedding that uses optimization at inference time to generate internal contextual signals allowing the agent to parse its temporal experience into discrete events and organize learning about them. They showed that the model trained on tasks sequentially using weight updates with task identifiers provided, can be used to identify tasks dynamically by taking gradient steps in the latent space. Moreover, the model shows generalization to novel tasks and can discover temporal events at arbitary time-scale and does not require a pre-specified number of events or clusters. ", "summary_of_the_review": "The paper borrow some idea from how the brain carries out continual learning (especially relating to the Thalamus). But how the simple RNN model is related to the PFC and the Thalamus is not that clear. Even though the authors showed that the proposed model works well on the cognitive benchmark, it is still not clear to what extent the model can be applied to more realistic task (those tasks considered in the ML community). ", "strength_and_weaknesses": "The proposed model is simple and effective on the benchmark. The author claimed that the model shows generalization to novel tasks. However, with such simple network, I am wondering if this is a real generalization, i.e., how different the novel tasks is from those tasks the model has seen during training? Do they share the same underlying task distribution?\n\nThere are several minor questions here: \n\n1, Why is the multiplicative effect consistent with the dis-inh. thalamocortical projections?\n\n2, What is r(t) in Eq.5?\n\n3, During training, why do you keep $W^z$ fixed?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "LKho4VcDyP-", "reviewer": "Reviewer_JbgB", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a brain-inspired algorithm for continual learning and disentangled representations. However, the contributions are generally not clear. ", "summary_of_the_review": "Weakness: \n1、\tThe connection between brain mechanisms such as thalamic-cortex interaction is weak. The background description is very confusing. I could not get the clear message from the description. I feel that they purposely link their algorithm with a brain mechanism. Unfortunately, authors don't tell a complete and feasible story.\n2、\tThe description for method is simple (e.g., the Section 2 contains very short paragraphs) but I can not see any connections between the model and the motivation mentioned above. \n", "strength_and_weaknesses": "The paper is somewhat well-written. The results may show the effectiveness of proposed method. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster