Building compositional tasks with shared neural subspaces

Abstract Cognition is highly flexible—we perform many different tasks 1 and continually adapt our behaviour to changing demands 2,3 . Artificial neural networks trained to perform multiple tasks will reuse representations 4 and computational components 5 across tasks. By composing tasks from these subcomponents, an agent can flexibly switch between tasks and rapidly learn new tasks 6,7 . Yet, whether such compositionality is found in the brain is unclear. Here we show the same subspaces of neural activity represent task-relevant information across multiple tasks, with each task flexibly engaging these subspaces in a task-specific manner. We trained monkeys to switch between three compositionally related tasks. In neural recordings, we found that task-relevant information about stimulus features and motor actions were represented in subspaces of neural activity that were shared across tasks. When monkeys performed a task, neural representations in the relevant shared sensory subspace were transformed to the relevant shared motor subspace. Monkeys adapted to changes in the task by iteratively updating their internal belief about the current task and then, based on this belief, flexibly engaging the shared sensory and motor subspaces relevant to the task. In summary, our findings suggest that the brain can flexibly perform multiple tasks by compositionally combining task-relevant neural representations.