Attention to task structure for cognitive flexibility

TL;DR

This paper investigates how environmental structure influences cognitive flexibility in multi-task learning, highlighting the interaction between environment connectivity and attention-based models in improving generalization and stability.

Contribution

It introduces a graph-theory based multi-task environment and compares attention models to MLPs, revealing environment connectivity's critical role in cognitive flexibility.

Findings

Rich environments enhance generalization and stability.

Connectivity between tasks significantly affects model performance.

Attention models benefit more from environment connectivity than MLPs.

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.