Attention to task structure for cognitive flexibility

Xiaoyu K. Zhang; Mehdi Senoussi; Tom Verguts

cs.NEApr 14, 2026

Attention to task structure for cognitive flexibility

Authors: Xiaoyu K. Zhang, Mehdi Senoussi, Tom Verguts

Organizations: Department of Experimental Psychology, Ghent University, Henri Dunantlaan 2, 9000 Gent, Belgium · Univ Toulouse, Université Toulouse Jean Jaurès, CNRS, CLLE, Toulouse, France

Abstract

Humans and artificial agents must often learn and switch between multiple tasks in dynamic environments. Success in such settings requires cognitive flexibility: the ability to retain prior knowledge (cognitive stability) while also transferring it to novel tasks (cognitive generalization). Cognitive flexibility research has largely focused on the role of model architecture to achieve these complementary goals. However, it is less well understood how the structure of the environment itself influences cognitive flexibility, and how it interacts with model architecture. To address this gap, we design a multi-task learning environment in which tasks are defined by a combination of two cue dimensions, allowing us to characterize the environment with graph-theory methods. We also introduce gating-based (multiplicative) and concatenation-based attention models that can decompose tasks into components and can sequentially allocate attention to them. We compare the attention-based models' performance in the multi-task learning environment to multilayer perceptrons. Generalization and stability are systematically evaluated across environments that vary in richness and task connectivity. We observe that richer environments improve both generalization and stability. In addition, a critical novel observation is that (graph theory based) connectivity between the tasks in the environment strongly modulates both stability and generalization, with especially pronounced benefits for attention-based models. These findings underscore the importance of considering not only cognitive architectures but also environmental structure and their interaction in shaping multi-task learning, generalization, and stability.

Figures & tables

Fig 1: Overview of the multidimensional task structure. (a) Multi-2: A $2\times 2$ dimensional task structure with two sensory cues $(S_{1},S_{2})$ and two motor cues $(M_{1},M_{2})$ . There are four combinations of cues (4 tasks) in the task structure. The first regime includes two tasks, $(S_{1},M_{1})$ and $(S_{2},M_{2})$ , while the second regime contains two novel tasks, $(S_{1},M_{2})$ and $(S_{2},M_{1})$ . (b) Multi-3: A $3\times 3$ dimensional task structure with three sensory cues and three motor cues. The first regime includes three tasks in a poor environment and six tasks in a rich environment. The second regime is identical for both, featuring three novel tasks. (c) Procedure: First regime tasks are learned with feedback, followed by a generalization test on second regime tasks without feedback. After second regime tasks with feedback, a stability test is conducted using the original first regime tasks without feedback.

Fig 2: Learning curves in Multi-2 for three model types. Curves show mean training accuracy over 50 runs for (a) MLP_2, (b) Gate_2, and (c) Concat_2, representing the MLP, Attention-Gating, and Attention-Concatenate architectures, respectively. Orange traces correspond to learning step 1 (training on the first-regime tasks), and blue traces correspond to learning step 2 (training on the second-regime tasks). Shaded regions indicate ±1 standard deviation across runs (cropped at 100%). The vertical dashed line marks the transition between learning regimes.

Fig 3: Performance of models in generalization and stability across different environments (Multi-2, Multi-3 Poor, Multi-3 Rich). Mean accuracy (averaged over 50 runs) for MLP_2, Gate_2, and Concat_2. (a) Generalization results show higher accuracy for attention-based models compared to the MLP model, with near-perfect performance in the rich environment. (b) Stability results highlight the robustness of attention-based models in retaining prior knowledge, while the MLP model exhibits catastrophic forgetting. All models perform better in the rich than poor environment in both generalization and stability. The chance level for each environment is shown on the color bar.

Fig 4: Experimental design for Multi-4 and connectivity graphs. Mean accuracy (averaged over 50 runs) for MLP_2, Gate_2, and Concat_2. (a) Environment richness. In Multi-4, the first regime varies by environmental richness: poor (4 orange tasks), middle (8 orange tasks), and rich (12 orange tasks). The second regime tasks (4 blue tasks) consist of tasks that are different from all first regime tasks. (b) Environment connectivity. Different first regimes are shown for the middle environment (8 orange tasks), that vary in their connectivity. The second regime tasks for each first regime consists of the remaining eight tasks (shown in white) that are not included in the first regime. (c) Cue space connectivity. Connectivity of first regimes from (b) in cue space. Each cue, $M$ or $S$ , is represented as a vertex and each edge between a pair of vertices represents a task. The two diagrams on the left illustrate connected regimes (all vertices are connected through a sequence of edges), while the two diagrams on the right illustrate disconnected regimes (not all vertices are connected).

Fig 5: Learning curves in the Multi-4 Middle environment for 3 model types. Mean training accuracy (averaged over 50 runs) across two learning steps for (a) MLP_2, (b) Gate_2, and (c) Concat_2. Orange curves show accuracy during learning step 1 (training on first-regime tasks), and blue curves show accuracy during learning step 2 (training on second-regime tasks). Shaded regions indicate ±1 standard deviation across runs. The vertical dashed line marks the transition between learning steps.

Fig 6: Effects of environmental richness and first-regime connectivity on generalization and stability in Multi-4. Mean accuracy (averaged over 50 runs) for MLP_2, Gate_2, and Concat_2. (a) Generalization and (b) stability in Multi-4. Richness and connectivity improve both metrics for all models (connected $>$ disconnected), but attention-based models consistently outperform the MLP model. In the middle environment, connected regimes allow attention-based models to approach rich-environment performance, whereas the MLP model remains below rich levels and shows greater forgetting. The chance level for each environment is indicated on the color bar. (Ctd = Connected; Dtd = Disconnected; middle environment.)

Fig 7: Task-wise performance in the Multi-4 Middle environment. (a) MLPs; (b) Attention–Gating models; (c) Attention–Concatenation models. Orange heatmaps show performance on the different tasks in the first regime, and blue heatmaps show performance on the different tasks in the second regime. Within each regime, accuracy varies only marginally across tasks, with no task on which any model is exceptionally strong or weak. Chance level (12.5%) is indicated on the color legend.

Fig 8: Accuracy across the 17 unique connected first-regime variants in the middle environment. Mean accuracy (averaged over 50 runs) for MLP_2, Gate_2, and Concat_2. (a) generalization accuracy; (b) stability accuracy. Attention-based models outperform the MLP model on generalization and stability under all connected first regimes. Regimes are ordered by increasing ASPL (higher to lower connectivity). Black boxes on the y-axis indicate identical ASPL (i.e., matched connectivity).

Fig 9: Violin plots of performance across 17 unique connected first-regime variants in the middle environment. Distributions of generalization and stability across 50 runs for (a) MLP_2, (b) Gate_2, and (c) Concat_2. Regimes are ordered by increasing ASPL (lower connectivity), and black boxes group regimes with identical ASPL values. Each violin plot summarizes the distribution across runs, with the horizontal bar indicating the mean. Chance level is indicated by a dashed line in each plot.

Fig 10: Linear regression of model performance vs. first-regime connectivity (ASPL) across 17 unique connected regimes. (a) MLPs show no significant relationship between ASPL and generalization accuracy, but exhibit a positive association between ASPL and stability. (b–c) For attention-based models, Gate_1 and Attention–Concatenate models show negative correlations between ASPL and generalization accuracy. By contrast, Gate_2 shows no significant correlation with ASPL for generalization, most plausibly due to a ceiling effect; the same ceiling effect likely explains the absence of a significant relationship for stability.

Fig 11: Linear regression of model performance vs. first-regime connectivity (LSPL) across 17 unique connected regimes. (a) MLPs show no significant relationship between LSPL and generalization accuracy, but exhibit a positive association between LSPL and stability. (b–c) For attention-based models, Gate_1 and Attention–Concatenate models show negative correlations between LSPL and generalization accuracy. By contrast, Gate_2 shows no significant correlation with LSPL for generalization, most plausibly due to a ceiling effect; the same ceiling effect likely explains the absence of a significant relationship for stability.

Fig 12: Cue sensitivity across poor, middle, and rich environments. Plots show mean cosine similarity of hidden representations under a single-cue change (averaged over 50 runs) for (a) MLP_2, (b) Gate_2, and (c) Concat_2. Cue1 (blue) denotes the sensory cue and Cue2 (orange) the motor cue. Lower cosine similarity indicates greater sensitivity to that cue change.

Fig 13: Cue sensitivity in the middle environment under connected (Ctd) and disconnected (Dtd) first regimes. Plots show mean cosine similarity of hidden representations under a single-cue change (averaged over 50 runs) for (a) MLP_2, (b) Gate_2, and (c) Concat_2. Cue1 (blue) denotes the sensory cue and Cue2 (orange) the motor cue. Lower cosine similarity indicates greater sensitivity to that cue change. Ctd/Dtd correspond to the regimes in Figure 4; Dtd_1 includes only sensory-cue changes.

ASPL	LSPL	Number
2	3	1
2.07	4	1
2.14	4	5
2.21	5	1
2.29	4	1
2.29	5	2

Table 1: Average shortest path length (ASPL) and longest shortest path length (LSPL) of all unique regimes in the middle environment.

Fig 14: Model architectures (for Multi-4). (a) Dense layer in MLP. (b) Mechanism in Attention-Gating layer. (c) Mechanism in Attention-Concatenate layer. (d) MLP model processes cues and stimuli together with backpropagation. (e) Attention-Gating Model uses Attention-Gating layers (Gate1 and Gate2) to gate stimulus features guided by cue features. In the upper plot, filled circle means active, open circle inactive. (f) Attention-Concatenate Model concatenates cue and stimulus features to guide learning (Concatnate1 and Concatenate2). Dotted lines indicate attention-gating pathways. Full lines indicate attention-concatenate pathways.

	MLP_1		MLP_2		Gate_1		Gate_2		Concat_1		Concat_2
	$N_{w}$	$N_{l}$	$N_{w}$	$N_{l}$	$N_{w}$	$N_{l}$	$N_{w}$	$N_{l}$	$N_{w}$	$N_{l}$	$N_{w}$	$N_{l}$
Multi-2	14	2	18	3	10	2	12	3	10	2	12	3
Multi-3	36	3	42	4	12	3	18	3	12	3	18	3
Multi-4	48	3	56	4	26	3	36	3	26	3	36	3

Table 2: Number of learnable parameters in the hidden layers of each model.

Multi-Task Learning

Multi-task learning trains models to solve related tasks together by sharing representations, parameters, or supervision.

See all 391 papers

Balancing Plasticity and Stability with Fast and Slow Successor Features
May 25, 2026Raymond Chua, Doina Precup, Blake RichardsSequential Decision Making Non-Stationary RL
School of Computer Science, McGill University, Montr´eal, Canada · Mila - Quebec Artificial Intelligence Institute, Montr´eal, QC, Canada · CIFAR Learning in Machines and Brains +3
Beyond Simply Environment Scaling: Designing Effective Environment Distributions for Multimodal Agent Learning
Aug 4, 2026Kejian Zhu, Zhuoran Jin, Dongqi Huang +4Curriculum Learning Multimodal Agents
The Key Laboratory of Cognition and Decision Intelligence for Complex Systems, Institute of Automation, Chinese Academy of Sciences, Beijing, China · School of Artificial Intelligence, University of Chinese Academy of Sciences, Beijing, China
Studying the Effects of Task Diversity Using Small Mathematical Worlds
Oct 9, 2026Emily Xia Meng, Xuanyu Yang, Junaid Hasan +2Multi-Task Learning LLM Training
University of Washington, Seattle · University of Washington Pacific Northwest National Laboratory

May 25, 2026J/K move · Enter open · S save

Balancing Plasticity and Stability with Fast and Slow Successor Features

Raymond Chua, Doina Precup, Blake Richards

School of Computer Science, McGill University, Montr´eal, Canada · Mila - Quebec Artificial Intelligence Institute, Montr´eal, QC, Canada · CIFAR Learning in Machines and Brains+3

Sequential Decision Making Non-Stationary RL Reinforcement Learning Stability-Plasticity Dilemma Continual RL Memory Consolidation+1

Open paper AskPDF arXiv

A hallmark of intelligence is the ability to adapt in non-stationary environments, yet deep Reinforcement Learning (RL) agents often struggle in such settings. Prior studies introduce non-stationarity through abrupt shifts in features or dynamics, whereas real-world environments often evolve gradually through continual drift. This distinction has important implications for the "stability-plasticity dilemma" in RL, as abrupt task changes may demand more plasticity than naturalistic settings. To address this, we modify existing 3D Miniworld and MuJoCo environments to incorporate naturalistic, continual non-stationarity, and use them to examine how stability and adaptation affect performance under continuous environmental change. We find that methods favoring stability, such as synaptic consolidation, outperform approaches focused on plasticity, such as parameters resetting. Motivated by this result, and prior evidence that Successor Features (SFs) reduce interference, we investigate whether SFs are better consolidation targets than Q-values. Across both environments, applying neuro-inspired synaptic consolidation to SFs yields superior performance on continually changing settings. Moreover, consolidation is most effective when SFs are stabilized across multiple timescales, which capture complementary aspects of gradual environmental change. Together, these results suggest that stability is more critical in continual learning when changes are gradual, and that multi-timescale consolidation of predictive representations is an effective approach.

Abstract

Multi-Task Learning

Explore similar work

Balancing Plasticity and Stability with Fast and Slow Successor Features

Beyond Simply Environment Scaling: Designing Effective Environment Distributions for Multimodal Agent Learning

Studying the Effects of Task Diversity Using Small Mathematical Worlds