Cognitive flexibility, the ability to adapt behavior to different contexts, is a fundamental cognitive ability. Context inference, the process of identifying the relevant context, is crucial for this flexibility. The prefrontal cortex (PFC) is known to play a critical role in cognitive flexibility, and recent research highlights the importance of the mediodorsal thalamus (MD) in regulating PFC dynamics and connectivity to support adaptive behaviors. The MD is densely connected to the PFC and appears to encode cueing contexts, sustaining context-relevant PFC representations while suppressing irrelevant ones. Studies in mice suggest that context inference can happen rapidly, within a few trials. However, the neural mechanisms underlying this rapid context inference remain poorly understood. Current computational models, such as recurrent neural networks (RNNs), often suffer from catastrophic forgetting when learning tasks sequentially. Continual learning, or the ability to learn new tasks without forgetting old ones, is a significant challenge in both artificial and biological systems. This study addresses these challenges by proposing a PFC-MD neural circuit model that incorporates biologically plausible mechanisms to achieve rapid context inference and continual learning.

Existing research emphasizes the PFC's critical role in cognitive flexibility and executive functions. Studies have demonstrated the importance of thalamocortical interactions, particularly the PFC-MD pathway, in regulating PFC activity and connectivity. The MD's role in amplifying context-relevant cortical connections and suppressing context-irrelevant ones has been highlighted. RNNs have been successfully used to model PFC dynamics, but they often struggle with catastrophic forgetting when learning tasks sequentially. Various continual learning approaches exist, broadly categorized as replay-based, regularization-based, and architecture-based methods. However, most methods rely on explicit task identification during training, which is not how context inference works biologically. This paper aims to bridge the gap by integrating biologically realistic properties of thalamocortical circuits into an RNN model to achieve rapid context inference and continual learning.

The authors developed a two-system recurrent neural network model consisting of an MD thalamus module and a PFC module. The model incorporates a Hebbian learning-based synaptic plasticity rule between the PFC and MD. This unsupervised learning enables the MD to infer temporal contexts by integrating context-relevant activities over trials. The PFC representations are gated by MD projections to prevent interference between different task representations. The model was initially trained on a simplified context-dependent classification task, analogous to attention-guided behavioral tasks used in mice studies. The network's performance was evaluated based on the accuracy of rule prediction given cue inputs. Further experiments involved training the model on more complex, sequentially learned cognitive tasks from the Neurogym platform, which tested various cognitive functions, including working memory, decision-making, and inhibitory control. To address challenges of continual learning in complex tasks, the model was modified to include a separate PFC-context (PFC-ctx) module, mirroring experimental observations of distinct PFC neuron types involved in cue-selective and cue-invariant responses. This improved architecture separated task learning and context inference into two pathways, with supervised learning in the task learning pathway and gating in the context inference pathway. The model was then compared to other biologically plausible continual learning methods such as Elastic Weight Consolidation (EWC) and Synaptic Intelligence (SI). Task similarity was quantified based on task representations in a recurrent neural network using task variance analysis. The model was further tested for forward transfer, the ability to transfer knowledge from previously learned tasks to similar future tasks, by manipulating the MD-to-PFC projections to create either disjoint or overlapping PFC neuron activations across tasks.

The study's key findings demonstrate that the proposed PFC-MD model effectively addresses the challenges of rapid context inference and continual learning. The MD thalamus accurately infers temporal contexts within a few trials, using Hebbian learning-based synaptic plasticity. This plasticity, incorporating pre-synaptic traces, adaptive thresholding, and winner-take-all normalization, allows for the robust and fast encoding of context in the MD, while gating prevents interference between different task representations in the PFC. The model exhibits significantly improved performance compared to a PFC-only model in flexibly switching between temporal contexts. The MD's context selectivity proves robust against noise in both inputs and PFC units, particularly during delay periods. The improved architecture with the PFC-ctx module successfully balances rapid context inference with complex task learning. The MD outperforms both the PFC and PFC-ctx in context encoding, even under high noise and low synaptic stability conditions. The inclusion of the MD promotes disjoint and modular task representations in the PFC. The MD-to-PFC connections enable continual learning in the PFC by enhancing context-relevant neuronal connectivity and suppressing context-irrelevant activities. The model exhibits lower error rates when switching back to previously learned contexts, contrasting with the catastrophic forgetting observed in PFC-only models and mirroring experimental findings in mice with MD suppression. The PFC-MD model significantly outperforms EWC and SI methods in continual learning performance. Finally, by manipulating the MD-to-PFC projections to create overlapping PFC activations, the model demonstrates forward transfer, particularly for similar task pairs, showcasing a trade-off between preventing forgetting and facilitating knowledge transfer. The model performance of the PFC-MD model and the PFC-only model were consistent with the mice experiment data under the manipulation in the three-block framework. MD suppression significantly degraded the task performance when the model switched back to the previous context.

The findings of this study provide valuable insights into the neural mechanisms underlying rapid context inference and continual learning. The proposed PFC-MD model demonstrates how biologically plausible properties of thalamocortical circuits can be leveraged to achieve these capabilities. The model's success in overcoming catastrophic forgetting and facilitating knowledge transfer has significant implications for the development of more robust and flexible artificial intelligence systems. The model's performance on sequential tasks highlights the potential for integrating context inference into artificial continual learning frameworks. The results align well with experimental observations in mice, supporting the biological plausibility of the proposed mechanisms. Future research could explore more complex and naturalistic scenarios, incorporating uncertainty and dynamic context changes. The model's limitations, such as the assumption of disjoint PFC activations under certain conditions, warrant further investigation.

This research presents a novel PFC-MD neural network model that effectively addresses the challenges of rapid context inference and continual learning. The model, incorporating biologically plausible mechanisms, demonstrates superior performance compared to existing methods. Future work should explore the model's scalability to a larger number of tasks and investigate more sophisticated gating mechanisms. This study makes a significant contribution to our understanding of how the brain achieves cognitive flexibility and offers valuable insights for developing more robust and adaptable AI systems.

The model makes certain simplifying assumptions, such as the assumption of largely disjoint PFC neuron activations for different contexts in the initial model. While this assumption is partially relaxed in later experiments, more sophisticated representations of context and task similarity could be explored. The study focuses on specific cognitive tasks; testing the model's generalizability to a wider range of tasks is important for future research. The computation of task similarity could be further refined to capture more nuanced relationships between tasks, and the model’s performance under dynamic, uncertain contexts could also be investigated.