Rapid context inference in a thalamocortical model using recurrent neural networks

Context is an abstract collection of goals, memories and external cues that inform appropriate behavioral responses. Humans and animals flexibly adapt actions based on changing contexts, requiring rapid context inference without overwriting past learning. Impaired context inference is linked to reduced flexibility and psychiatric disorders such as schizophrenia. The prefrontal cortex (PFC) is central to cognitive flexibility, and the mediodorsal thalamus (MD), with dense reciprocal connectivity to PFC, has emerged as a key regulator of PFC dynamics and effective connectivity for adaptive behavior. MD amplifies local cortical connectivity, sustains rule representations in PFC, encodes cueing context, and suppresses context-irrelevant PFC activity, enabling rapid context inference observed behaviorally within a few trials. Despite these insights, the neural mechanism for fast temporal context inference remains unclear. Meanwhile, recurrent neural networks (RNNs) have modeled PFC dynamics across tasks but suffer catastrophic forgetting when tasks are learned sequentially. Existing continual learning approaches (replay-, regularization-, and architecture-based) often rely on explicit task identifiers and do not explain how temporal task identification can be rapidly inferred during performance. This study proposes a biologically inspired thalamocortical model that rapidly infers temporal context online and supports continual learning.

Prior work shows PFC’s role in cognitive control and flexibility, with MD thalamus regulating PFC dynamics, amplifying local connectivity, sustaining rule representations, and encoding cueing context. Experimental studies in mice demonstrate selective context coding in MD and mixed selectivity in PFC (cue-selective and cue-invariant neurons), with fast-spiking PFC neurons carrying context information. Behaviorally, context inference occurs within a few trials. In computational neuroscience, RNNs capture context-dependent PFC computations but forget when trained sequentially. Continual learning methods—replay, regularization (e.g., EWC, SI), and architectural/masking strategies—mitigate forgetting but typically need explicit task/context labels and do not address rapid online context inference from neural activity. Previous thalamocortical models proposed MD-like mediation but used limited MD units or required separate change-point detection modules. These gaps motivate a biologically grounded mechanism enabling MD to infer context online and gate PFC to prevent interference while learning.

The authors develop a two-system thalamocortical model comprising a PFC module and an MD module, later extended with a PFC-context (PFC-ctx) module for complex tasks. - Core PFC-MD model for attention-guided task: PFC is a recurrent rate network receiving cue input, recurrent input, and MD gating. MD is feedforward, receiving input from PFC and providing feedback to PFC. A Hebbian synaptic plasticity rule is used between PFC and MD to enable online context inference: synaptic updates are proportional to the product of post-synaptic (MD) and pre-synaptic (PFC) traces relative to an adaptive threshold, combined with winner-take-all (WTA) normalization in MD. Pre-synaptic traces integrate PFC activity over a short history (few trials) to extract context; post-synaptic traces stabilize MD outputs within the same context. WTA activates the top K MD neurons per context for a low-dimensional, discrete representation. Adaptive thresholding balances potentiation and depression to learn new contexts quickly while forgetting old ones slowly. - MD-to-PFC effects: MD exerts multiplicative (enhancing recurrent input to context-relevant PFC neurons) and additive (suppressing context-irrelevant PFC neurons) influences, implementing gating that maintains current-context PFC representations and inhibits others. - Extended architecture for complex tasks: To handle dynamic inputs and simultaneous task learning, task optimization and context inference are separated via an added PFC-ctx layer (feedforward, fast-spiking-like), which receives (task-disjoint) inputs and projects to MD. PFC handles task computation with learned recurrence via supervised training. MD receives PFC-ctx input and, via Hebbian plasticity with traces and WTA, infers temporal context and gates PFC activity (disinhibition) additively. PFC recurrence (and input-to-PFC, output weights) are trained with Adam on mean square error; PFC-ctx→MD weights are updated by a modified Hebbian rule with adaptive thresholds and bounded weights. Noise and probabilistic synaptic transmission are used in PFC-ctx to test robustness. For forward transfer studies, MD-to-PFC gating is manipulated to be disjoint (non-overlapping PFC subsets per task) or overlapping (some shared PFC neurons spared from inhibition) based on measured task similarity. - Tasks and training: (1) Attention-guided behavior task with two contexts (two cues each), trained in alternating blocks with a return to the first context for continual learning assessment. (2) NeuroGym/Yang19 tasks comprising working memory, decision making, categorization, inhibitory control families; networks learn two or three tasks sequentially in blocks. Evaluation includes decoding analyses (context vs rule information in PFC, MD, PFC-ctx), robustness to noise and hyperparameters (number of MD neurons, number of contexts), trajectory analyses (PCA), and comparisons against EWC and SI baselines. Continual learning (CL) is measured as performance on task 1 after learning task 2; forward transfer (FT) as performance on task 2 after learning task 1. Task similarity is quantified by inner product of task variance vectors from an RNN trained with interleaved tasks.

- Selective MD context encoding and rapid inference: MD neurons selectively encode temporal context, not rules, while PFC contains both rule and context information. MD infers temporal contexts within about six trials (three cycles), demonstrating few-trial, online inference. t-SNE shows clearer context separation in MD than in PFC. - Mechanistic validation: Pre-synaptic traces accumulate context-relevant PFC activity across trials; WTA yields stable, low-dimensional MD representations; adaptive thresholding speeds learning of current context while slowing forgetting of past contexts, allowing sequential encoding of multiple contexts. - Robustness: MD context decoding is more robust than PFC under intermediate noise levels in inputs and PFC units, especially during delay periods when cues are absent. MD selectivity persists across variations in the number of MD neurons and number of contexts, provided K (top active MD units) is appropriate. - Complex tasks and separated pathways: With the added PFC-ctx pathway, MD outperforms PFC and PFC-ctx in context encoding under noise and synaptic instability, integrating PFC-ctx activity across trials into robust context codes. MD regulation promotes more modular, non-overlapping task representations in PFC compared to PFC-only models (neural trajectories do not overlap in PFC-MD). - Continual learning benefits: In three-block switching (context1→context2→context1), PFC-MD shows substantially lower errors upon returning to a previous context and aligns with mouse data where bilateral MD suppression impairs performance in the third block (behavioral comparisons show significant effects; e.g., P<0.001). MD-mediated gating protects synapses of context-irrelevant PFC neurons (smaller weight changes), reducing interference and catastrophic forgetting. - Benchmarking vs baselines: On NeuroGym tasks learned sequentially, PFC-MD achieves near-perfect preservation when switching tasks with disjoint gating and outperforms EWC and SI on both Task 1 retention and Task 2 acquisition. Regularization methods provide soft protection and show reduced performance compared to MD-gated modularization. - Forward transfer trade-off: Disjoint MD-to-PFC effects yield perfect CL but no FT. Introducing overlapping MD effects (shared PFC neurons across tasks) enables forward transfer at the cost of some forgetting. FT improves with higher task similarity; CL moderately decreases but remains better than in baselines for similar tasks. All models show FT increasing with task similarity, but PFC-MD (overlapping) balances CL and FT more favorably.

The study addresses how thalamocortical circuitry could support rapid temporal context inference and continual learning. A Hebbian plasticity mechanism with pre/post traces, WTA, and adaptive thresholding enables MD to integrate PFC activity over a few trials to extract low-dimensional, discrete context representations. MD then gates PFC, enhancing context-relevant populations and suppressing irrelevant ones, producing modular task representations that limit interference and sustain previously learned contexts. This mechanism explains empirical observations that MD encodes context and regulates PFC dynamics, and reproduces mouse behavioral effects of MD suppression during context switching. Extending to complex tasks, separating task learning (PFC) and context inference (PFC-ctx→MD) preserves rapid context inference despite evolving PFC representations. The model demonstrates a fundamental trade-off between stability (continual learning) and flexibility (forward transfer), which can be tuned via MD-to-PFC overlap guided by task similarity. The framework is biologically plausible in its use of Hebbian rules, feedforward MD, traces, and competition, and provides a unifying account of context inference, gating, and weight protection in PFC during sequential learning.

The authors propose a biologically inspired PFC-MD model that performs rapid, online temporal context inference via Hebbian plasticity with traces, adaptive thresholding, and MD winner-take-all, and uses MD gating to support continual learning. The model: (1) infers context within a few trials and encodes it robustly in low dimensions; (2) gates PFC to maintain context-relevant activity and suppress irrelevant populations, reducing interference; (3) scales to complex tasks with a separated PFC-ctx pathway; (4) outperforms EWC and SI on sequential NeuroGym tasks; and (5) enables forward transfer by overlapping MD-to-PFC effects according to task similarity, revealing a tunable trade-off between forgetting and transfer. Future work should develop more flexible, learnable gating mechanisms, online task similarity estimation modules, and test the framework in more naturalistic, continuous contexts and additional definitions of context.

- Context operationalization: Experiments use blocked designs with discrete cue sets or task IDs; real-world contexts are continuous and uncertain. The model’s context reflects stimulus pairs/task identification, limiting generalizability to other context types (e.g., reward contingencies). - Representation constraint: MD’s context code approximates a discretized first principal component of time-averaged PFC activity; context may not always align with the largest variance component. - Forward transfer vs stability: Disjoint MD gating prevents FT; overlapping gating introduces forgetting. Online measurement and control of task similarity were not implemented during training. - Architectural assumptions: PFC-ctx receives task-disjoint inputs and directly informs MD; while motivated by fast-spiking/context signals, this is an assumption. Gating is hand-crafted in part (e.g., WTA K, inhibition levels), and multiplicative gating on recurrence was simplified for complex tasks. - Training paradigm: Blocked training facilitates linking PFC populations to MD; performance under interleaved or un-cued context switches in highly dynamic environments needs further evaluation. - Biological mapping: While components are plausible (Hebbian rules, traces, feedforward MD, competition), precise physiological implementations and broader thalamic computations require additional empirical validation.