Building Compositional Tasks with Shared Neural Subspaces

Abstract

Main

Humans and other animals can combine simple behaviours to create more complex behaviours. For example, learning to discriminate whether a piece of fruit is ripe can be used as a component of a variety of foraging, cooking and eating tasks. The ability to compositionally combine behaviours is thought to be central to generalized intelligence in humans and a necessary component for artificial neural networks to achieve human-level intelligence.

When artificial neural networks are trained to perform multiple tasks, they reuse representations and computational components in different tasks. Whether the brain similarly reuses sensory, cognitive and/or motor representations across tasks remains unclear. Furthermore, we do not yet understand how the brain flexibly engages these representations to continually adapt to the changing demands of the environment. To address these questions, we trained two monkeys to flexibly switch between three different compositionally related tasks.

Composing Tasks from Subtasks

All three tasks followed the same general structure: the monkeys were presented with a visual stimulus and had to indicate its category with an eye movement (Fig. 1a). The stimuli were parametric morphs, independently varying in both shape and colour (Fig. 1b).

The monkeys performed three categorization tasks. In the shape–axis 1 (S1) task, they categorized the shape of the stimulus and then responded on axis 1: when the shape was more similar to a ‘bunny’, the monkey made a saccade to the upper-left (UL) target, and when more similar to a ‘tee’, to the lower-right (LR) target.

The colour–axis 2 (C2) task required the monkey to categorize the colour of the stimulus and respond on axis 2: if ‘red’, saccade to upper-right (UR), and if ‘green’, saccade to lower-left (LL).

Finally, in the colour–axis 1 (C1) task, the monkeys categorized the colour (as in C2) and responded on axis 1 (as in S1; red = LR, green = UL). In this way, the tasks were compositionally related—C1 can be considered as combining the colour categorization subtask of C2 with the motor response subtask of S1.

Both monkeys performed all three tasks well (monkey Si/Ch: S1, 81%/77%; C1, 83%/78%; C2, 92%/92%; all P < 0.001, binomial test). The monkeys were not instructed as to the identity of the new task. Thus, they had to learn which task was in effect on each new block.

Representations Were Shared Across Tasks

To understand the neural representations used during each task, we simultaneously recorded neural activity from five cortical and subcortical regions: the lateral prefrontal cortex (LPFC; 480 neurons), frontal eye fields (FEF; 149 neurons), parietal cortex (PAR; 64 neurons), anterior inferior temporal cortex (aIT; 239 neurons) and striatum (caudate nucleus; 149 neurons).

To test whether representational subspaces were reused across tasks, we quantified how well a classifier trained to decode the stimulus’ colour category or the animals’ motor response in one task generalized to the other tasks. Consistent with a shared representation in the LPFC, a classifier trained to decode the colour category during the C2 task was able to significantly decode colour category during the C1 task (65 ms after stimulus onset; P < 0.001, permutation test).

Motor response representations also generalized across tasks. A classifier trained to decode response direction in the S1 task generalized to decode response direction in the C1 task (and vice versa; 128 ms after stimulus onset; P < 0.001).

Tasks Sequentially Used Shared Subspaces

If the representation of both the category of the stimulus and the motor response were shared across tasks then performing a task requires selectively transforming representations from one subspace to another. Consistent with this, there was a sequential representation of the stimulus colour in the shared colour subspace followed by the motor response in the shared motor subspace during the C1 task (63 ms difference in onset time; P < 0.001, t-test).

To directly test whether this reflected the transformation of information between subspaces, we tested whether information about the stimulus colour in the shared colour subspace predicted the response in the shared response subspace on a trial-by-trial basis. The correlation was shifted upward with respect to the diagonal line (36 ms before saccade start), indicating that the encoding in the shared colour subspace predicted the future encoding in the shared response subspace.

Importantly, the transformation of stimulus information into response was specific to the task: during the C1 task, the shared colour representation did not predict the associated response direction along axis 2. This is consistent with the shared colour representation being selectively transformed into a motor response along axis 1, and not axis 2, when the monkey was performing the C1 task.

Task Belief Engaged Shared Subspaces

Theoretical modelling suggests that shared subspaces could facilitate cognitive flexibility by allowing the brain to engage previously learned, task-appropriate representations. If true, then this predicts that task-appropriate shared subspaces should be engaged as the animal discovers the task in effect.

We trained a classifier to decode the identity of the S1 and C1 tasks using neural activity in the LPFC during the fixation period. As the animals discovered the current task, the performance of the task classifier increased. During the S1–C2–C1 sequences, the classifier increased from 42% to 54% accuracy (P = 0.0079, permutation test).

The strength of the representation in the shared colour subspace increased as the monkeys discovered the C1 task (64% to 69%, Δ = 5%; P = 0.0478). This increase was predicted by the strength of internal task representation in the LPFC during fixation (P = 0.0478).

In contrast to shared colour and shape subspaces, the animals’ motor response was stably decoded in a shared subspace throughout the block. This makes sense, as both S1 and C1 tasks used the same motor response.

Task Belief Scaled Shared Subspaces

Previous research suggests tasks may modulate the gain of stimulus features, depending on their relevance for the current task. Similarly, we found stimulus information was attenuated when irrelevant. Furthermore, this attenuation depended on the animals’ internal belief about the task: as they discovered the C1 task was in effect, the representation of colour category was magnified, while shape representation was attenuated and delayed in time.

We defined a compression index (CPI) as the log of the ratio between the separability of stimuli in the colour subspace versus the shape subspace. Colour and shape representations were scaled in all three tasks. CPI was positively correlated with the strength of the task representation and colour category encoding in the LPFC (z-scored Kendall’s τ = 2.328, P = 0.008).

In addition to selecting sensory information, the three tasks also require selecting the appropriate motor response. Neurons selective for each motor axis were suppressed when the animal performed a task requiring a response on the other axis. In contrast to the gradual suppression of sensory representations during task discovery, the incorrect response axis representation was quickly suppressed after a change in task (within 3–6 trials).

Discussion

Our results suggest the brain can perform a task by compositionally engaging a series of representational subspaces. Subspaces of neural activity within prefrontal cortex represented task-relevant information. These subspaces were shared across multiple tasks, suggesting they act as task components. Subspaces were sequentially engaged, such that information from the relevant sensory subspace was transformed into the appropriate motor response subspace.

Although our study is limited to three tasks, the underlying mechanism has the ability to be highly expressive—flexibly sequencing task components together could implement a wide variety of behaviours (a form of sequential compositionality). In this framework, the representation of the task acts as a control input that selects the appropriate representations and computations.

If the brain can reuse representations and computations across tasks, then this could allow one to rapidly adapt to changes in the environment, either by learning the appropriate task representation through reward feedback or by recalling it from long-term memory.

Several previous studies have found different tasks use representations specific to that task. Such task-unique representations have the advantage of minimizing interference, but limit generalization. By contrast, shared subspaces increase interference but could speed learning by allowing knowledge to generalize across tasks.

One reason shared representations were found (rather than task-unique ones) may be because the task required continuous learning. The scaling of neural representations may facilitate continual learning by constraining learning to task-relevant representations—neural learning rules are activity dependent and gated by reward, so suppressing irrelevant representations may limit learning to task-relevant ones (addressing the credit assignment problem).

Methods Summary

Full methods including surgical procedures, signal preprocessing, GLM analysis, classifier construction, pseudopopulation methods, cross-validation, CPI calculation, and TDR analysis are available in the full article.

Supplementary & Extended Data Figures

The paper includes 5 main figures and 10 extended data figures covering additional analyses:

Key References

View all 69 references on Nature →