Attention to task structure for cognitive flexibility

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.

• The paper reveals that attention-based architectures significantly improve generalization and stability by leveraging compositional task environments.
• The study systematically benchmarks MLPs against attention-gating and concatenation models within multi-n environments varying in richness and connectivity.
• Attention models develop disentangled, cue-selective representations, underscoring the critical role of task structure in driving cognitive flexibility.

Attention to Task Structure as a Driver of Cognitive Flexibility in Multi-Task Neural Networks

Introduction

The ability to simultaneously retain prior knowledge (stability) and generalize to new situations (generalization) is essential for cognitive flexibility in both biological and artificial agents. Despite extensive focus on inductive biases and architectural choices (e.g., MLPs, attention-based models) for optimizing this trade-off, the role of global environmental structure—particularly the richness and connectivity of task regimes—remains underexplored. The paper "Attention to task structure for cognitive flexibility" (2604.13281) systematically addresses the interaction between environmental structure and network architecture in multi-task settings. Specifically, the authors examine how compositional task environments and attention-augmented neural architectures affect generalization to novel tasks and retention of previously learned mappings.

Multi-Task Environment Formalism

The experimental framework is built around a family of Multi-$n$ environments, in which tasks are defined by unique combinations of sensory and motor cues. Each environment (Multi-2, Multi-3, Multi-4) is characterized by the dimensionality of the cue spaces and by the number and organization of regimes (i.e., subsets of tasks presented sequentially). Task regimes are manipulated with respect to richness (the coverage of cue combinations) and connectivity (the overlap of cue components among tasks), which are formalized using graph-theoretical metrics such as average shortest path length (ASPL) and graph diameter (LSPL). This setup systematically exposes the interaction between environmental structure and neural model performance.

Figure 1: Schematic illustrating the Multi-$n$ task structure, regime partitioning, and training/evaluation protocol.

Model Architectures

Three model families are benchmarked:

• MLPs (with modest variation in depth),
• Attention-Gating models (multiplicatively modulating stimulus representations using task cues at two stages),
• Attention-Concatenate models (injecting learned cue representations at multiple locations via concatenation).

Attention-based architectures feature a dedicated cue-processing stream that learns to extract compositional structure for selective feature routing. Both gating and concatenation variants allow cues to condition separate stages of stimulus processing, supporting hierarchical selection and routing.

Figure 2: Mechanistic comparison of MLP, Attention-Gating, and Attention-Concatenate architectures, including their internal cue-routing mechanisms.

Results

Learning and Performance Dynamics

In all environments, networks undergo two-phase sequential learning: train on a first-regime (old tasks), test generalization on novel second-regime tasks without feedback, then train on the second-regime (new tasks), and finally test retention (stability) on the original tasks. All models consistently reach ceiling accuracy within their respective training regimes, indicating effective task acquisition.

Figure 3: Training accuracy curves for all three model families in the Multi-2 environment.

Impact of Richness

Increasing environmental richness (i.e., the diversity and coverage of tasks in the first regime) substantially benefits both generalization and stability, but the effect is especially pronounced in attention-augmented architectures. In the Multi-3 rich environment, Attention-Gating and Attention-Concatenate models achieve generalization rates ranging from 66.3% up to 94.7%, outstripping MLPs (49.2%). This advantage is mirrored in stability, where attention models retain near-ceiling performance on previously learned tasks, exhibiting robust resistance to catastrophic interference, while MLPs display pronounced forgetting.

Figure 4: Summary of generalization and stability metrics as a function of environmental richness for all three models.

Effects of Connectivity

The introduction of a richer Multi-4 environment allows a more granular investigation of connectivity and its impact. Using graph-based definitions, the authors distinguish connected and disconnected regimes and manipulate their connectivity strength (ASPL, LSPL). All models show improved performance in connected vs. disconnected first regimes, but attention architectures are much more sensitive to increasing connectivity, exhibiting monotonic gains in generalization approaching 93%, and stability gains well above 95% even at intermediate richness levels.

Figure 5: Effects of richness and regime connectivity on generalization and stability across architecture types in Multi-4.

Systematic exploration across all 17 unique connected regimes in Multi-4 further clarifies that attention-based models consistently outperform MLPs across all levels of connectivity. Linear regression analyses demonstrate a robust negative correlation between connectivity (lower ASPL/LSPL) and generalization accuracy for attention-based models, but not for MLPs (except a weak effect on stability; see Figures 10–11). This quantifies the unique ability of attention mechanisms to exploit environmental structure for compositional transfer, supporting highly flexible generalization.

Figure 6: Generalization and stability accuracy across all 17 unique connected regime variants in the Multi-4 middle environment.

Figure 7: Generalization and stability as a function of first-regime connectivity (ASPL) for all model types.

Analysis of Internal Representations

Attention-based models develop more disentangled, cue-selective internal representations as richness increases, supporting the hypothesis that hierarchical routing and selective attention promote compositional generalization. Cosine similarity analyses reveal that in richer environments, attention models learn to segregate sensory and motor features early in their processing hierarchy, while MLPs maintain entangled representations across all layers. This cue-separability is invariant to regime connectivity, indicating that it is driven by richness rather than overlap per se.

Figure 8: Layerwise cue sensitivity (cosine similarity) after single-cue perturbations, across environments and model families.

Theoretical and Practical Implications

This study provides robust evidence that the global structure of task environments is a critical determinant of generalization and retention in multi-task learning. While prior work has emphasized the role of computational architecture and local task similarity, these findings demonstrate that regime connectivity—not merely the presence of compositional structure, but its systematic organization—modulates performance, provided that architectures can leverage it (e.g., attention-based routing). The results challenge the sufficiency of approaches focusing exclusively on architectural innovations (e.g., regularization, replay) to solve catastrophic forgetting. Instead, architectural flexibility must be matched to the graph-theoretic structure of experience.

These insights are highly relevant for both cognitive modeling (linking human flexibility to structured environments) and applied neural computation (designing neural agents for continual and compositional learning). For reinforcement learning and curriculum design, graph-based analyses of environment connectivity could guide efficient regime construction and sequence ordering for optimal generalization and stability.

Future Directions

One limitation of the present framework is the controlled, low-complexity task structure. Scalability of these findings to high-dimensional, noisy, or partially observable tasks—including hierarchical and temporal dependencies—remains to be explored. Incorporating richer structure (e.g., causal, temporal, or hierarchical graphs) and investigating how attention mechanisms scale in these domains are natural extensions. Moreover, further probing of internal representations and the potential for synergistic integration with algorithmic approaches (e.g., replay, regularization) would clarify the boundaries of connectivity-driven flexibility.

Conclusion

The systematic dissection provided in "Attention to task structure for cognitive flexibility" (2604.13281) establishes that cognitive flexibility in neural models is co-determined by both model architecture and the connectivity structure of the multi-task environment. Attention-based architectures equipped for selective routing can harness regime connectivity to achieve superior generalization and retention, especially in rich, well-connected environments. These results point toward the necessity of jointly optimizing architectural inductive biases and environmental structure analyses to achieve human-level flexible learning in artificial agents.