Building compositional tasks with shared neural subspaces

Abstract

Cognition is remarkably flexible; we are able to rapidly learn and perform many different tasks¹. Theoretical modeling has shown artificial neural networks trained to perform multiple tasks will re-use representations² and computational components³ across tasks. By composing tasks from these sub-components, an agent can flexibly switch between tasks and rapidly learn new tasks⁴. Yet, whether such compositionality is found in the brain is unknown. Here, we show the same subspaces of neural activity represent task-relevant information across multiple tasks, with each task compositionally combining these subspaces in a task-specific manner. We trained monkeys to switch between three compositionally related tasks. Neural recordings found 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. Subspaces were flexibly engaged as monkeys discovered the task in effect; their internal belief about the current task predicted the strength of representations in task-relevant subspaces. In sum, our findings suggest that the brain can flexibly perform multiple tasks by compositionally combining task-relevant neural representations across tasks.

Introduction

Humans, and other animals, can combine simple behaviors to create more complex behaviors^5–9. For example, once we learn to discriminate whether a piece of fruit is ripe, we can reuse this ability as a component of a variety of foraging, cooking, and eating tasks. The ability to compositionally combine behaviors is thought to be central to generalized intelligence in humans¹⁰ and a necessary component for artificial neural networks to achieve human-level intelligence^11–15. When artificial neural networks are trained to perform multiple tasks, they reuse representations and computational components in different tasks^2–4. Whether the brain similarly reuses sensory, cognitive, and/or motor representations across tasks, and how these representations may be flexibly combined, remains unknown. To test this hypothesis, we trained monkeys to perform three different compositionally related tasks.

Monkeys perform three compositionally related tasks

All three tasks followed the same general structure: animals 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 color (Fig. 1b, see Methods for details). Animals performed three categorization tasks. In the Shape-Axis1 (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 animal made a saccade to upper-left (UL) target, and when the shape was more similar to a ‘tee’, the animal saccaded to the lower-right (LR) target (Fig. 1c, top row). The Color-Axis2 (C2) task required the animal to categorize the color of the stimulus and respond on Axis 2: if the color was ‘red’, they saccaded to upper-right (UR), and if ‘green’, they saccaded to lower-left (LL; Fig. 1c, bottom row). Finally, in the Color-Axis1 (C1) task, the animals categorized the color of the stimulus (as in the C2 task) and responded on Axis 1 (as in the S1 task; red=LR, green=UL, Fig. 1c, middle row). In this way, the tasks were compositionally related – C1 can be considered as combining the color categorization sub-task of C2 with the motor response sub-task of S1.

Fig. 1|. Monkeys perform three compositional tasks.

a, Time course of trial. After a brief fixation (500–800ms), a visual stimulus and four response targets appeared on the screen. Animals reported the category of the stimulus with a saccade to one of the four target locations. b, Stimuli were morphs in a two-dimensional feature space that varied in shape (left) and color (right). Stimulus categories are indicated by vertical lines and labels. 50% morph stimuli were randomly rewarded. c, Schematic of the task design: Task S1 required categorizing the stimulus based on its shape and responding on Axis 1. Task C2 required categorizing the stimulus based on its color and responding on Axis 2. Task C1 required categorizing the stimulus based on its color and responding on Axis 1. Colored backgrounds indicate shared sub-tasks: color categorization sub-task between the C1 and C2 tasks (blue) and response sub-task between S1 and C1 tasks (orange). d, Example of task transition sequence. Task switches occurred when performance was equal or greater than 70%. Animals were not instructed as to the identity of the upcoming task, although the axis of response switched between blocks. Number of included blocks for S1/C1/C2 were 94/97/189, respectively. e, Average psychometric curve for both animals for the 102/102/51 trials before the switch for S1/C1/C2 tasks, respectively. Number of included blocks for S1/C1/C2 were 94/97/189, respectively. f, Average performance of both animals after (left) and before (right) a switch (average window of 15 trials, horizontal line for p ≤ 0.001; binomial test uncorrected for multiple comparisons across trials) g, Schematic of locations of neural recordings. h-l, To quantify information about task variables in individual neurons, we computed the proportion of variance in neural response uniquely explained by each task variable using the coefficient of partial determination (CPD, average CPD for permuted dataset was subtracted from observed CPD for each neuron, see Methods). Panels show CPD of h, color category, i, shape category, j, task identity, k, reward and l, response direction. Lines show mean CPDs across all neurons for each region. Regions were omitted if they did not have a significant number of neurons encoding a cognitive variable. m, Time course of average normalized CPD across all recorded neurons, normalized by maximum value to show temporal order.

Overall, both animals performed all three tasks well (Figs. 1e and S1; Monkey Si/Ch: S1:81%/77%, C1: 83%/78%, C2: 92%/92%; all p<0.001, binomial test, see Fig. S1 for individual animal performance). Animals performed the same task for a block of trials (Fig. 1d). When they reached a behavioral criterion (performance ≥70%), the task would change (see Methods for details). The animal was not instructed as to the identity of the new task, although the axis of response always changed between blocks (see Methods). So, animals had to learn which task was in effect on each new block. This was reflected in the animal’s behavior, which improved over the first 75 trials of the S1 and C1 tasks (Fig. 1f, left, performance on S1/C1 increased from 0.47/0.62 at trial 15 to 0.71/0.77 at trial 75, both p<0.001, Chi-squared test). As only the C2 task used Axis 2 responses, the animal identified the task quickly, performing it above chance within 15 trials after the switch (Fig. 1f, performance at 15 trials=0.85, p<0.001, binomial test). Altogether, the animals’ behavior suggests they rapidly identified the change in response axis but slowly integrated feedback to learn whether S1 or C1 was in effect¹⁶.

Task variables were represented in independent subspaces of neural activity

To understand the neural representations used during each task, we simultaneously recorded neural activity from five cortical and subcortical regions (Fig 1g): 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). All five regions represented task-relevant cognitive variables^17,18, including the identity of the current task, the color and shape of the stimulus, the response direction, and whether a reward was received (Fig. 1h–1l). The identity of the task was consistently represented throughout the trial (Figs. 1j and S2). After stimulus onset, information about the color (Fig. 1h) and shape of the stimulus (Fig. 1i) was followed by information about the direction of the animal’s response (Fig. 1l), and then the reward received (Fig. 1k, see Fig. 1m for average time course across all regions).

To understand how each task variable was represented in the neural population, we trained classifiers to decode the stimulus color, shape, and motor response from the pattern of neural activity across all neurons (using a pseudo-population to combine neurons from all recording sessions, see Methods for details). Classifiers trained on LPFC neural activity accurately decoded the category of the stimulus’ color in the C1 and C2 tasks (Fig. 2a; 79ms, 87ms after stimulus onset for C1 and C2, respectively; similar results were seen when controlling for motor response, Fig. S3a, p<0.001 permutation test). In this way, the classifier defined a “subspace” within the high-dimensional space of neural activity that represented the color category of the stimulus input. Similarly, classifiers trained on LPFC activity decoded the category of the stimulus’ shape during the S1 task (Fig. 2b; 100ms after stimulus onset, p<0.001 permutation test).

Fig. 2|. LPFC encodes shared color and response representations.

a-c, Accuracy of classifier trained to decode a, color category, b, shape category, and c, motor response from LPFC neural activity. Lines and shading show mean ± s.e.m. over time. Color of line indicates which task was used to train/test the classifier. Distribution reflects 250 resampled classifiers (see Methods for details). Horizontal bars (top right of each plot) indicate above-chance classification (p ≤ 0.05, 0.01, and 0.001 for thin, medium, and thick lines, respectively; permutation test with cluster mass correction for multiple comparisons). d, Schematic of classifiers used to test whether color category and response location information were shared across tasks (see Methods). e, Cross-temporal cross-task classification accuracy in decoding shared color category in LPFC, trained on C2 task and tested on C1 task. f, Time course of cross-task classification accuracy when trained on LPFC neural activity to decode color during C1 task (blue) or C2 task (green) and then tested on the other task. g, Difference in the onset of color information during C1 task when decoded from classifiers trained on C1 task versus generalized from C2 task (red star and red x in panels a and f, respectively. p ≤ 0.05, 0.01 and 0.001 for *, **, ***, respectively; t-test). To compare timing, classifier accuracy for each area was smoothed with a 50ms boxcar filter. h, Cross-temporal cross-task classification accuracy in decoding shared response direction in LPFC. Trained on S1 task and tested on C1 task. i, Time course of cross-task classification accuracy when trained on LPFC neural activity to decode response during C1 task (orange) or S1 task (green) and then tested on the other task.

Task irrelevant stimulus information was attenuated. In LPFC, classification accuracy about the shape of the stimulus was reduced during the C1 and C2 tasks (Fig. 2b) and information about the color of the stimulus was reduced during the S1 task (Fig. 2a). Similarly, task-relevant information, but not task-irrelevant information, was represented in other regions, including in higher-order visual cortex (Fig. S3b–c).

The direction of the animal’s response could be decoded within each task’s response axis (Fig. 2c, see Methods for details). In LPFC, information about the response occurred after the stimulus’ category (115ms, 133ms after stimulus onset in S1 and C1 tasks, respectively; both p<0.001, permutation test, see Fig. S4c for C2 task). Similar results were seen in other regions (Fig. S4a). Altogether, these results show the stimulus color category, stimulus shape category and response direction are broadly represented at the population level in multiple recorded regions.

Neural representations were shared across tasks

To test whether representational subspaces were reused across tasks, we quantified how well a classifier trained to decode the stimulus’ color category or the animals’ motor response in one task generalized to other tasks (Fig. 2d, Methods). Consistent with a shared representation in LPFC, a classifier trained to decode the color category of a stimulus during the C2 task was able to significantly decode the stimulus’ color category during the C1 task (Figs. 2e and 2f, 65ms after stimulus onset in LPFC, p<0.001, permutation test). The reverse was also true: a classifier trained on C1 could decode stimulus’s color category during the C2 task (Fig. 2f and S5c, 84ms after stimulus onset, p<0.001, permutation test). Importantly, the two tasks require different motor responses and so the shared representation reflected the color category of the stimulus and not the motor response (see Figs. S5a and S6a for further controls for movement). Although color information was reduced in the S1 task, the weak color category representation that did exist also generalized between the C1 and S1 tasks (Fig. S5b).

While color category information was represented in a shared subspace in LPFC, it was represented in task-specific subspaces in other brain regions (Fig. S7). Generalization was weaker in FEF, PAR and IT and was delayed with respect to task-specific sensory information (and was delayed relative to LPFC, Fig. 2g). There was no significant generalization in STR (although stimulus color category could be decoded from those regions, Fig S3b).

Motor response representations also generalized across tasks. A classifier trained to decode response direction in the C1 task generalized to decode response direction in the S1 task (and vice versa; Fig. 2h and 2i, 128/128ms after stimulus onset in LPFC, when training on S1/C1, testing on C1/S1, respectively, p<0.001, permutation test). Unlike stimulus information, the motor representation was shared in all regions (Fig. S4b; perhaps reflecting a widely broadcasted motor signal¹⁹, but see²⁰).

Shared representations were sequentially engaged during a task

So far, our results suggest the representation of both sensory inputs and motor responses were shared across tasks. In this framework, performing a task requires selectively transforming representations from one subspace to another to support each task²¹. Consistent with this, there was a sequential representation of the stimulus color in the shared color subspace followed by the response in the shared motor subspace during the C1 task (Fig. 3a, 63ms difference in onset time of color response using C2 classifier and motor response using S1, p<0.001 t-test).

Figure 3|. Shared representations were sequentially transformed during the task.

a, Sequential processing of shared color category information and shared response direction information in LPFC. Classifier accuracy was normalized to range between maximum accuracy and baseline 200ms before stimulus. b, Schematic showing prediction that shared color representation is transformed to the Axis 1 and Axis 2 response axis during the C1 and C2 tasks, respectively. c, Cross-temporal, trial-by-trial correlation of shared color category encoding (trained on C2 task, tested on C1 task) and shared response location encoding (trained on S1 task, tested on C1 task). Thin and thick red lines indicate p≤0.05 and p≤0.001, respectively, uncorrected t-test. d, Same as panel b, but correlating shared color category encoding (trained on C2 task, tested on C1 task) and response direction encoding on Axis 2 (trained on C2 task, tested on C1 task). e, Average cross-temporal correlation along anti-diagonal axis in panel c (shared color encoding in C1 task predicts response direction in S1 task, blue line) and panel d (shared color encoding in C1 task predicts response direction in C2 task, red line). The shift of the curve’s center towards negative time values reflects the extent to which encoding in the shared color subspace predicts future encoding in the shared response subspace. Dotted lines show mean of gaussian fit to each curve.

To directly test whether this reflected the transformation of information between subspaces, we tested whether information about the stimulus color in the shared color subspace predicted the response in the shared response subspace on a trial-by-trial basis (Fig. 3b). Figure 3c shows the correlation between the representation of the stimulus’ color category in the C1 task, decoded using the classifier from the C2 task, and the representation of the motor response in the C1 task, decoded using the response classifier from the S1 task (see Methods for details). Correlation was measured across all possible pairs of timepoints, quantifying whether color or motor response representations at one time point were correlated with representations at a future time point. The correlation was shifted upward with respect to diagonal line (36ms before saccade start, Figs. 3c and 3e, blue line), indicating that the encoding in the shared color subspace predicts the future encoding in the shared response subspace.

Importantly, this transformation was specific to the task: during the C1 task the shared color representation was not predictive of the associated response direction along Axis 2 (8ms after saccade start, Figs. 3d and 3e, red line). This is consistent with the shared color representation being selectively transformed into a motor response along Axis 1, and not Axis 2, when the animal was performing the C1 task. In contrast, when the animals performed the C2 task, the shared color representation was transformed into a motor response along Axis 2, and not Axis 1 (Fig. S6b–d).

Together, these results suggest tasks sequentially engage shared subspaces, selectively transforming stimulus representations into motor representations in a task-specific manner. Interestingly, there was a negative correlation between the sensory and motor responses early in the trial (Fig. 3c), which may reflect suppression of the motor response during fixation or integration of the stimulus input.

Shared subspaces are dynamically engaged by task belief

Theoretical modeling suggests shared subspaces could facilitate cognitive flexibility by allowing the brain to engage previously learned, task appropriate, representational and computational subspaces^2,3. If true, then this predicts task-appropriate shared subspaces should be engaged as the animal discovers the task in effect.

To begin to test this hypothesis, we first measured the animal’s internal representation of the task (i.e., their ‘belief’ about the task encoded by the neural population). To do so, we trained a classifier to decode the identity of the task using neural activity in LPFC during the fixation period (i.e., before stimulus onset,²²). Training was restricted to the last 75 trials of the S1 and C1 tasks, when behavioral performance was high (Fig. 1f). We then applied the task classifier to the beginning of the C1 task blocks to measure the animals’ internal belief about the task as they learned which task was in effect (Fig. 4b–1, see Methods for details).

Fig. 4|. Shared sub-task subspaces are dynamically engaged during task discovery.

a, Behavioral performance of monkeys during the C1 task depended on the sequence of preceding tasks. Dashed gray line is 25% chance level. Horizontal bars indicate significant difference between performance in two sequences (p ≤ 0.05, 0.01, and 0.001 for thin, medium, and thick lines, respectively; Chi-squared test uncorrected for multiple comparisons across trials). Number of included blocks for C1-C2-C1/S1-C2-C1 sequences were 44/53, respectively. b, Schematic of classifier tests during task discovery.1: Task identity classifier was trained to decode task identity on 75 trials before switch in S1 task and C1 task, tested on sliding window of 50 trials starting from switch into C1 task. 2: Color category classifier was trained to decode color category on 75 trials before switch in C2 task, tested on sliding window of 50 trials starting from switch into C1 task. 3: Shape category classifier was trained to decode shape category on 75 trials before switch in S1 task, tested on sliding window of 50 trials starting from switch into C1 task. 4: Response direction classifier was trained to decode response direction on 75 trials before switch in S1 task, tested on sliding window of 50 trials starting from switch into C1 task. Same test trials were used for 1, 2 and 3 classifiers. Red and black arrows denote train and test trials, respectively. c, Classifier accuracy for task identity during −400ms to 0ms period from sample onset. Darker to lighter color show progression in trial blocks of 50 shifted by 5 trials from switch trial in S1-C2-C1 task transition. Lines show mean classification accuracy after stimulus onset for 250 iterations of classifiers. Horizontal bars indicate above chance trend for classifier accuracy during task discovery (p ≤ 0.05, 0.01, and 0.001 for thin, medium, and thick lines, respectively; modified Mann-Kendall test, Benjamini-Hochberg correction for multiple comparisons across time) d, Comparison of progression of task identity classifier accuracy in C1-C2-C1 task sequence (purple) and S1-C2-C1 task sequence (black). Average accuracy was computed using classifier performance in −400ms to 0ms before stimulus onset period͘ Lines and shading show mean ± s.e.m. classification accuracy after stimulus onset. Distribution reflects 250 iterations of classifiers. e, Correlation of behavioral performance and task belief encoding. f, As in panel c, but showing classifier accuracy for color category classifier in S1-C2-C1 task transition. Grey box shows period with significant color category information (see Methods). g, Comparison of progression of color category classifier accuracy in C1-C2-C1 sequences (purple) and S1-C2-C1 sequences (black). Average accuracy was computed using classifier accuracy in 100ms to 300ms from stimulus onset. h, Correlation between task belief encoding and color category encoding (see Methods for details). p ≤ 0.05, 0.01, and 0.001 for thin, medium, and thick lines, modified Mann-Kendall test, Benjamini-Hochberg correction for multiple comparisons across time. i,j same as f,g but for shape category classifier. Grey box shows period with significant shape category information. k,l same as f,g but for response direction classifier. Classifier accuracy was computed using classifier accuracy in 200ms to 400ms after stimulus onset. Grey box shows period with significant response direction information.

As noted above, monkeys slowly learned whether the S1 or C1 task was in effect¹⁶. Interestingly, animals’ rate of learning depended on the sequence of tasks. When switching into a C1 task, animals learned more quickly when the previous Axis 1 task was C1 compared to S1 (Fig. 4a, purple: C1-C2-C1 task sequence, black: S1-C2-C1; Δ=10.26% change in performance over first 20 trials, p<0.001, Chi-squared test). A similar pattern was seen for S1: performance on task sequence S1-C2-S1 was greater than C1-C2-S1 (Fig. S8a, Δ=4.99%, p=0.032, Chi-squared test). As the C2 task had a unique response axis, it was not affected by the preceding task (Fig. S8b Δ=1.42% between S1-C2 and C1-C2 sequences, p=0.234, Chi-squared test).

As with behavior, the performance of the task classifier increased over trials (Fig. 4c) and depended on the sequence of tasks (Fig. 4d). During S1-C2-C1 sequences, the classifier was slightly below chance initially, suggesting a slight encoding of the S1 task, before increasing as the animals learned (Fig. 4c–d, from 43% to 54% from trial 50 to 125, p<0.001 modified Mann-Kendall trend test). On C1-C2-C1 task sequences, classifier performance was high immediately after the switch to C1 (69% classifier accuracy on trial 50, p<0.001, permutation test) and slightly increased with trials (Fig. 4d, purple line; 8% increase from trial 50 to 125, p<0.001 modified Mann-Kendall trend test). Moreover, the task belief during the C1 task was correlated with behavioral performance on the color categorization task (Fig. 4e, $\tau =0.46$, p=0.0463 modified Mann-Kendall trend test) and anti-correlated with how much their behavior depended on the shape category (Fig. S9a, $\tau =-0.55$, p=0.02 modified Mann-Kendall trend test). Together, these results suggest the monkeys tracked whether the S1 or C1 task was in effect, and that this belief was at least partially maintained between blocks of the Axis 1 tasks, leading to a bias in the neural representation and behavior.

Given this measure of the animal’s internal representation of the task, we next asked whether this representation predicted the strength of representations in the shared color category and shared motor response subspaces. Similar to the task representation, the strength of the representation in the shared color subspace increased as the animals discovered the C1 task (Fig. 4f, Fig. 4b–2, see Methods for details). Like task representations, color representations depended on task sequence (Fig. 4g). Shared color representations increased during the C1 task when it followed a S1-C2-C1 sequence (62% to 68%, Δ=6% from trial 50 to 125, p<0.001 modified Mann-Kendall trend test), but remained stable during C1-C2-C1 sequences (62% to 59%, Δ=3% from trial 50 to 125, p=0.12 modified Mann-Kendall trend test). This suggests that shared color subspace is dynamically engaged as the animal updates its belief about the current task. Indeed, during learning, the strength of internal task representation in LPFC during fixation was correlated with the strength of shared color category representation after onset of the stimulus (Fig. 4h, see Methods for details).

The animal’s behavior and internal representation of the task suggest they initially expected to perform the S1 task in S1-C2-C1 task sequences. If true, and task representations modulate the strength of shared subspace representations, then shape information should initially be strong and then decay as the animal learns the C1 task is in effect. To test this, we trained a classifier to decode the stimulus’ shape category during the last 75 trials of the S1 task, and then tested it as the animal discovered the C1 task (Fig. 4i, Fig. 4b–3, see Methods for details). In contrast to color, shape representation was significant immediately after the switch to the C1 task during S1-C2-C1 sequences (54%, p=0.02, permutation test) and significantly decreased as the animal learned (Fig. 4i–j, 49%, Δ=5%, p=0.049, modified Mann-Kendall trend test). In contrast, shape information was reduced overall and remained stable during C1-C2-C1 task sequences (Fig. 4j, purple, 51% on trial 50, Δ=1%, p=0.203, modified Mann-Kendall trend test). Finally, consistent with task belief modulating engagement of shape representations, belief about the C1 task was inversely correlated with the representation of shape (Fig. S9b).

In contrast to shared color and shape subspaces, the animals’ motor response was stably decoded in a shared subspace (Fig. 4k–l, Fig. 4b–4, see Methods for details, Δ=3%, p=0.106, and Δ=1%, p=0.284, for S1-C2-C1 and C1-C2-C1 sequences, respectively; modified Mann-Kendall trend test; see Methods for details). The representation in the shared response subspace was also not correlated with the strength of internal task representation in LPFC ($\tau =0.38$, p=0.070, modified Mann-Kendall trend test, Fig. S9c).

Altogether, these results suggest the animals use their internal belief to selectively engage the relevant shared color and shape subspaces during the C1 task. The task-relevant stimulus subspace (color) was then dynamically projected onto the shared representation of the response.

Task belief scaled representations within the shared subspaces

Previous work suggests the gain of stimulus features can change, depending on their relevance for the current task^23–26. Indeed, we found task-irrelevant information was attenuated (Fig. 2a, 2b). Furthermore, this attenuation depended on the animals’ internal belief about the task: as they discovered the C1 task was in effect, the representation of color category was magnified (Fig. 5a), while shape representation was attenuated (Fig. 5b).

Fig. 5|. Stimulus representations were compressed during flexible behavior.

a, Distance between stimuli along color category encoding axis (Methods). Darker to lighter color show progression in trial blocks of 50 shifted by 5 trials from switch trial in S1-C2-C1 sequences. Lines show mean value for 250 iterations of classifiers. Horizontal bars indicate significant trend in distance during task discovery (p ≤ 0.05, 0.01, and 0.001 for thin, medium, and thick lines, respectively; modified Mann-Kendall test, Benjamini-Hochberg correction for multiple comparisons across time). Grey box shows period with significant color category information. b, same as a but for distance along shape category encoding axis. Grey box shows period with significant shape category information. c, Compression Index during trial for three tasks. Compression index is taken as the log of the ratio of the separability of stimuli in color and shape subspaces. d, Compression Index during task discovery in S1-C2-C1 sequences. e, Correlation between task belief encoding and average compression index during 100ms-300ms since stimulus onset (Methods). f, Correlation between color category encoding and average compression index during 100ms-300ms since stimulus onset (Methods).

Such amplitude modulation can be thought of as a geometric scaling of stimulus representations in feature space^23,24,27. To quantify this, we defined a ‘Compression Index’ (CPI) as the log of the ratio between the separability of stimuli in (i) the color subspace and (ii) the shape subspace. Color and shape representations were scaled in all three tasks. There was greater separation of color representation during C1 and C2 tasks, and greater separation of shape representations during the S1 task (Fig. 5c). This scaling changed as animals learned which task was in effect (Fig. 5d). In fact, CPI was positively correlated with the strength of the task representation and color category encoding in LPFC (Fig. 5e and 5f, $\tau =0.73$, p<0.001 and $\tau =0.73$, p<0.001, respectively).

Together, these results suggest scaling of stimulus information is a common computation across all the three tasks that can be flexibly engaged based on the monkeys’ belief about the current task, dynamically adjusting the representational geometry of sensory representations^28–30.

Discussion

Our results suggest the brain can perform a task by compositionally combining a set of representational subspaces. We found ‘shared’ subspaces of neural activity in prefrontal cortex that represented sensory inputs and motor actions across multiple tasks. This contrasts with previous work that has found high-dimensional, non-linear representations that ‘mix’ sensory, motor, and task information in a unique way for each task^31,32. Such non-linear, high-dimensional, representations are easy to linearly separate, facilitating learning and performing arbitrary tasks³³. However, because each task representation is essentially unique, learning cannot be generalized across tasks²⁴. Instead, our results are consistent with recent theoretical models that suggest neural representations and computational components are reused across tasks. Shared representational subspaces could speed learning by allowing knowledge to generalize across tasks⁴. For example, the neural representation (and associated computations) that were learned to categorize the color of a stimulus during C2 could be generalized to support behavior during C1. Future work is needed to test this hypothesis as subjects learn new tasks.

While sharing representations across tasks may facilitate generalization between tasks, it could also lead to interference. Given that the animals were trained for months on the task, one might expect the brain to form independent task representations in order to reduce interferences between tasks^34,2. One potential explanation for our observation that representations are shared across tasks could be the uncertainty in which task to perform (as there was no cue for the current task). This may lead to continual learning, which would be facilitated by shared representations³⁵. Instead, we found task-irrelevant dimensions were compressed^28,36 as a function of the animals’ internal belief about the task state (Fig. 5). Compression could reduce interference from non-relevant stimulus inputs (specifically, incongruent stimuli in the task).

Generalization of subspaces across tasks was strongest in prefrontal cortex. This is consistent with previous work that has found projecting multiple stimuli³⁷ or memories³⁸ into a common subspace within prefrontal cortex can allow a single task to generalize across stimuli and/or memories³⁹. Our results show subspace representations within prefrontal cortex can be shared across multiple tasks. Of course, other regions may also be involved – for example, previous work has found generalization of representations in the hippocampus^37,40.

Our results are consistent with prefrontal cortex acting as a ‘global workspace’⁴¹, with different subspaces of neural activity representing different task-relevant information. In this model, performing a task requires transforming information from the relevant sensory subspace to the appropriate motor response subspace. Consistent with this, we found sensory representations in the category subspace predicted responses in the motor subspace on a trial-by-trial basis (Fig. 3). The mechanism by which information is transformed between subspaces remains unknown, although previous work suggests rotation of neural representations may be one possible mechanism^42,43.

Which subspaces were engaged, and how information was transformed, depended on the animals’ internal belief about the current task (Figs. 3 and 4). This is consistent with computational models of cognitive flexibility that suggest task representations can act as a control input to select task-appropriate representations and computations in a neural network (whether feed-forward²⁶ or recurrent^2,3,44). Thus, if the necessary representations already exist, then learning a new task only requires discovering the appropriate control input – i.e., the representation of the task that engages the appropriate representations/computations. This could dramatically speed learning, either by learning the task representations through reward feedback^45,46 or recalling it from long-term memory⁴⁷.