Learning shapes neural geometry in the prefrontal cortex

The relationship between the geometry of neural representations and the task being performed is a central question in neuroscience1–6. The primate prefrontal cortex (PFC) is a primary focus of inquiry in this regard, as under different conditions, PFC can encode information with geometries that either rely on past experience7–13 or are experience agnostic3,14–16. One hypothesis is that PFC representations should evolve with learning4,17,18, from a format that supports exploration of all possible task rules to a format that minimises metabolic cost4,17,18 and supports generalisation7,8. Here we test this idea by recording neural activity from PFC when learning a new rule (‘XOR rule’) from scratch. We show that PFC representations progress from being high dimensional and randomly mixed to low dimensional and rule selective, consistent with predictions from metabolically constrained optimised neural networks. We also find that this low-dimensional representation facilitates generalisation of the XOR rule to a new stimulus set. These results show that previously conflicting accounts of PFC representations can be reconciled by considering the adaptation of these representations across learning in the service of metabolic efficiency and generalisation.

Here, we tested this idea in two macaque monkeys which learnt an exclusive-or (XOR) rule -a problem that can be solved by a range of representations, from low-to highdimensional (Fig. 1a) 19 .Importantly, we tracked how the dimensionality and geometry of PFC representations changed across multiple training sessions of an XOR rule that was entirely new to the animals at the start of recording (experiment 1) and during subsequent generalisation of this rule to a new stimulus set (experiment 2).We used a classical conditioning paradigm in which the nonlinear combination of the features of two objects presented in succession (XOR) predicted the outcome of the trial.Importantly, the animals were only required to fixate through both experiments.This classical conditioning design carries advantages relative to operant tasks in that changes in representation with learning are unlikely to be confounded with changes in behaviour [28][29][30][31] .Later, we also show that our results hold in a previously collected delayed match-to-sample task [32][33][34] .
We found that initially PFC selectivity was high-dimensional and randomly mixed with strongly linear representations.As learning progressed, PFC activity became lowdimensional, manifesting a non-linear but highly structured selectivity pattern across the neural population.We compared these results against predictions from both normatively defined generative models and optimised artificial neural networks.Collectively, these models provide a theoretical foundation for why representations should change with learning -namely that PFC representations optimise for task performance while also contending with metabolic costs.Finally, we tested whether the learnt low-dimensional representation was reused in conditions where new unseen combinations of stimuli were presented.We found that during this generalisation of the XOR rule, the representation of these new stimuli became aligned with the lowdimensional representation of the pre-existing stimulus set, allowing the same neural code to be used across different stimulus sets.
Figure 1.Potential effects of learning on neural geometry in the prefrontal cortex.Learning can reduce or expand neural dimensionality, changing how many linear decoding axes can be implemented on neural firing rates (discriminability).a, High-dimensional representations enable high discriminability.A high-dimensional regime allows the separation of all task features using three possible readout axes (left), whereas a low-dimensional representation only allows task-relevant features to be separated (right).b, Low-dimensional representations enable high generalisability.When the task representation is high-dimensional (left), XOR discriminant from one task (blue and green shapes) fails to correctly differentiate the XOR feature in the second task (pink and khaki); i.e., a rotation that aligns both XOR discriminant axes is not possible.When the neural code is low-dimensional, both tasks can be immediately aligned to a common axis enabling a shared XOR discriminant that generalises across tasks.c, In Experiment 1, animals were incentivised to combine two passively viewed task features (colour and shape) in a non-linear fashion (XOR).For example, blue+square and green+diamond combinations were rewarded, whereas blue+diamond and green+square were not.In experiment 2 an additional set of colour-shape contingencies was introduced.Monkeys could generalise the learned rule from experiment 1 to the new set of colours in experiment 2. d, Timeline of task events in a single trial.e, Neural data was collected from the lateral surface of the prefrontal cortex in two macaque monkeys (see Methods and Supplementary Figure 5 for further details).

Results
In experiment 1, the animals were trained to combine a colour stimulus (either blue or green) with a subsequently presented shape (either square or diamond) in a nonlinear fashion to predict the reward (the XOR between colour and shape) outcome of the trial (Fig. 1c-d).We recorded 376 neurons from the lateral PFC across both macaques (Fig. 1e; see Methods, data acquisition and pre-processing).Importantly, to capture learning dynamics, we started recording from the first session in which the animals were exposed to the stimuli.Experimental sessions were then pooled into four learning stages.All selectivity and decoding analyses were run in time windows before the animals received feedback about the outcome of the trial (i.e., reward; for details see Methods, data acquisition and pre-processing).

Generative models of random and minimal selectivity
We first wanted to understand how the geometry of the neural representations changed over the course of learning.We compared neural population activity at each stage of learning to the geometries produced by two generative models with different discriminability-generalisability trade-offs 4,17 : (i) a high-dimensional random geometry 3,[14][15][16]35 (high discriminability, low generalisability); and (ii) a lowdimensional, structured geometry (low discriminability, high generalisability) 10,12,[36][37][38][39][40] . Th former is a well-established geometry (e.g., inherent in reservoir computing models 15,16 ) whereas the latter was motivated by recent theoretical work which suggests that complex (nonlinear) tasks may require low-dimensional, highly structured population activity 10,41 .These models make distinct predictions about the distribution of selectivity to the task-relevant variables in a 3-dimensional selectivity space (colour, shape, and their interaction, i.e., XOR).
In this selectivity space, each axis represents a neuron's response to one stimulus variable (e.g., 'shape selectivity' = higher firing rate for square than diamond).Hypothetically, one could imagine different distributions of variable encoding within this selectivity space.The properties of these distributions are determined by a covariance matrix: the diagonal elements describe the strength of coding of each variable (variance) whereas the off-diagonal entries determine the strength of the relationship between variables (covariance).In our generative models neural firing rates were simply constructed as a linear combination of task variables (colour, shape and XOR inputs represented using one-hot encodings; see Methods, generative models for details) with a specified covariance matrix of selectivities to the task variables.In line with previous studies, the high-dimensional random mixed model was constructed by allowing selectivity to be distributed randomly according to a spherical Gaussian distribution (Fig. 2a).In contrast, in the minimally structured XOR selectivity model, only the interaction (i.e., the XOR) term carried variance (Fig. 2b).We derived this model mathematically and demonstrated that it minimises total firing rate activity (i.e., a metabolic cost) while maximising task performance (see Supp.Materials, section 1).Consistent with this, we found that feedforward networks trained using backpropagation to perform the task while also minimising a metabolic cost term converge to the minimal XOR selectivity model (see Supp.Fig. 1a-f; Methods, optimised feedforward networks).Additionally, in line with prior accounts 17,42 , XOR decoding in the minimal model was more robust to noise (Fig. 2e) and required fewer neurons to implement a stable readout (see Supp.Fig. 1g).
We established a metric to measure whether neural activity is better described by the random or minimal model.We first fitted a linear model to surrogate data generated by both our generative models, in which task variables (colour, shape, and colour x shape (XOR)) were used as predictors of each unit's firing rate (see Methods, model section; for similar analysis see 41,43 ).Then we measured the average within-model distance between two random models (Fig. 2c, red line; Methods, eq. 4) and the average between-model distance between the random and the minimal model (Fig. 2c, blue line, Methods, eq. 5).To test whether our measure captures learning dynamics, we constructed four artificial populations with varying proportions of random and minimal selectivity.As the proportion of the minimal model in this mixed population increased, it became more dissimilar to the average random model and more similar to the average minimal model (Fig. 2c, black line).A reflected version of these results held true when the minimal model was used as reference (Fig. 2d).
Subsequently, to gain insight into the predictions of these models, we employed an established technique 3,14 and trained linear decoders to decode all three task variables in both models.As previously suggested 3,14 , a randomly mixed selectivity model yielded a high-dimensional task representation, allowing for all variables, including the nonlinear XOR, to be decoded (Fig. 2f, red; cf.Fig. 1a, far left).In contrast, for the minimal model, only the XOR combination of shape and colour, and not shape or colour independently, could be decoded (Fig. 2f, blue; cf.Fig. 1a, far right).
Figure 2. Predictions from random mixed and minimal selectivity models.a, Each neuron can be represented as a point in the 3-dimensional selectivity space spanned by colour, shape, and XOR (their interaction).In the random mixed model, selectivity is distributed according to a spherical Gaussian distribution in this space (Methods, generative models); the covariance matrix is computed between the selectivity coefficients.b, Analogous to a but for the minimal model; neurons are selective only for the XOR (interaction between colour and shape), as this is the only feature that is necessary to solve the task.c, Relative distance between the covariance matrices of either the random (red), minimal (blue) or the model with varying proportions of minimal selectivity (black) to the random selectivity model.d, Same as panel c but the distance is calculated relative to the minimal selectivity model.Red and blue error bars show standard deviations (±1 ..over 1000 randomly drawn models; see Methods, measuring similarity between selectivity distributions) of the relative Euclidean distance between the covariance matrix of the random, minimal model (with matched total variance to the data) and the covariance matrix expected from random selectivity; black error bars show the standard deviation of the relative distance between the surrogate covariance (with varying proportions of minimal selectivity) and random covariance (±1 ..over 1000 random models).e, XOR decoding as a function of noise () in the neural activities for the random and minimal selectivity models.Dashed grey line shows chance-level decoding.f, Mean (over 100 models) decoding of task variables for the random (red) and minimal (blue) models.Dashed grey line shows chance-level decoding.g, Mean (over 100 models) cross-generalised decoding of task variables for the random (red) and minimal (blue) models.Dashed grey line shows chance-level decoding.
While both models can perform the task, we expected their representation of the XOR variable to differ fundamentally.On the one hand, the minimal model by design should represent the XOR in a format that generalises over all other task variables.On the other hand, this is not guaranteed in the random model.We verified this intuition using cross-generalised decoding, a method in which a linear decoder is trained to decode a given task variable (e.g., XOR = True vs. XOR = False) on a given set of task conditions (e.g., blue colour) and tested on a different set of task conditions 7 (e.g., green colour; see Methods section, cross-generalised decoding).We found that for the random model, cross-generalised decoding was at chance-level for all task variables (Fig. 2g, red).This is because the random model, by design, exhibits no reliable structure in its representation of variables and therefore these dimensions are represented randomly in relation to each other.In contrast, the minimal model displayed maximal cross-generalised decoding for the XOR variable (Fig. 2g, blue), indicating that it can be decoded regardless of which set of task conditions the decoder is trained and tested on.This suggests that the minimal model is able to represent the XOR in a highly cross-generalisable format (Fig. 1a, far right).Consequently, the minimal model also exhibits below-chance cross-generalised decoding for colour and shape.

A low-dimensional task relevant geometry emerges over learning
We applied these metrics to our neural data (Fig. 3).Inspecting the neural representation of task variables on the population level suggested that the geometry in the PFC changed over the course of learning (Fig. 3a) and resulted in a lowerdimensional, task-specific (XOR) representation by the end of learning (Fig. 3a, far right).To quantify this, we fitted a linear regression to our data, just as we did for our generative models (Fig. 2c), in which task variables (colour, shape, and colour x shape (XOR)) were used as predictors of each neuron's firing rate.We then examined how selectivity coefficients changed over learning (Fig. 3b, learning stage 1 and Fig. 3c, learning stage 4).We compared the covariance structure of these selectivity coefficients (Fig. 3b,c, bottom right) to the covariances obtained from the random selectivity model (Fig. 3b, red contours) and minimal selectivity model (Fig. 3c, blue line).At the beginning of learning (stage 1), PFC cells were randomly distributed in selectivity space resembling the random model ( = .219,one-sided, Fig. 3d and  = .858,one-sided, Fig. 3e).However, in late learning (stage 4), selectivity diverged away from randomly mixed selectivity ( = .000,one-sided, Fig. 3d, compare black and red lines) and converged towards the minimal model ( = .000,one-sided, Fig. 3e, compare black and blue lines).
We next applied our linear decoding analyses (Fig. 2f, g) to the neural recordings to establish whether these changes to selectivity distributions led to changes in the geometry of neural representations (Fig. 3f,g; Supp.Fig. 3e-l).At the beginning of training (learning stage 1), the animals exhibited a high-dimensional geometry that allowed for every variable to be decoded, reminiscent of the random mixed model (compare Fig. 2f, red to Fig. 3f, stage 1; see Supp.Fig. 3a-d for temporally resolved decoding).In line with the random model, the task variables also exhibited poor crossgeneralised decoding during early learning (compare Fig. 2g, red to Fig. 3g, stage 1).In late learning (stage 4), however, the representation became low-dimensional and task-specific.As predicted by our minimal model, the coding of the XOR strengthened ( = .008,one-sided) whereas the decoding of shape decreased ( = .000,one-sided; compare Fig. 2f, blue to Fig. 3f, stage 4).Moreover, the XOR variable became progressively represented in a format that generalised across the remaining task variables ( = .000,one-sided,) and cross-generalised decoding for shape decreased over learning ( = .02,one-sided; compare Fig. 2g, blue to Fig. 3g, stage 4).Unexpectedly, colour decoding remained high even in late learning (decoding:  = .518,one-sided; cross-decoding:  = .17,one-sided; Fig. 2f; Supp.Fig. a, e).Finally, when we examined an additional task feature that was not predictive of reward (the width of the shape, see Methods section, data and task), we found a strong reduction in its representation in PFC over learning (decoding:  = .014,one-sided; cross.decoding:  = .002,one-sided; see Supp.Fig. 3c, g, k).
We next explored whether these learning-related changes were accompanied by a decrease in neural dimensionality (as measured with shattering dimensionality; see 14 for details).A neural representation described by three binary input dimensions (colour, shape and width) results in 35 dichotomies (division into two sets of four stimuli) that can be theoretically decoded.We found that the mean decoding accuracy of all dimensions (excluding colour, shape and XOR) decreased significantly over learning ( = .046,one-sided; Fig. 3h).Additionally, a principal component analysis computed on condition averages revealed that the proportion of variance explained by the first principal component increased as a function of learning ( = .01,one-sided; Supp.Fig. 3o; for details see Methods, Principal component analysis).Both the shattering dimensionality and the PCA analyses indicate that neural dimensionality decreased as a function of learning.Furthermore, as colour, shape, and width form a cube in the input space (when one-hot encoding is used) these dichotomies can be split into linear (a hyperplane can be used to divide all vertices into two sets) and nonlinear dimensions (a nonlinear function needs to be used to differentiate two sets of 4 vertices).We found that the decoding of linear dimensions decreased significantly over learning ( = .012,one-sided), while the decoding of non-linear dimensions did not ( = .22,two-sided; Fig. 3i).The progressive change over learning, not only in the neural dimensionality but also in the type of computations that are represented (linear vs non-linear) conflicts with a static model of the PFC as a high-dimensional nonlinear kernel (e.g., such as that utilised in reservoir computing 15,16 ).d, Relative Euclidean distance between the covariance matrix of selectivity coefficients from the data and the covariance matrix expected from random selectivity (with matched total variance) as a function of learning (Methods, measuring similarity between selectivity distributions).Red and blue, respectively, error bars show mean (±1 ..over 1000 randomly drawn models) of the relative Euclidean distance between the covariance matrix of the random, respectively minimal, model (with matched total variance to the data) and the covariance matrix expected from random selectivity; black error bars show standard deviation of relative distance between the observed covariance and random covariance (±1 ..over 1000 random models).e, Same as panel d but we show the relative distance from the covariance matrix expected from minimal selectivity (with matched total variance).f, Linear decoding of task variables for learning stages 1 (grey) and 4 (black).g, Cross-generalised linear decoding of task variables for learning stages 1 (grey) and 4 (black).h, The learning dynamic of the decoding of all possible task dichotomies excluding colour, shape and XOR (shattering dimensionality); solid black line indicates the mean over these dichotomies.i, Mean ±1 ..linear decoding of all linear dichotomies (dashed line), non-linear dichotomies (solid line), and the XOR dichotomy (dotted black) as a function of learning.All p-values were calculated from permutation tests (***,  < 0.01; **,  < 0.01; *,  < 0.05; †,  < 0.1; .., not significant).Grey shading in panels f-g shows the mean ±1 .. of chance-level decoding obtained by shuffling trial labels (for details see Methods, statistical testing).
To replicate our findings, we also performed a re-analysis of an existing dataset [32][33][34] in which recordings were taken from primate ventral and dorsolateral PFC before and after learning a delayed match-to-sample task that was similar in structure to ours (for details see Methods, existing lPFC dataset).This task required active behaviour and the reward signal was disentangled from the XOR (i.e., match/no match) signal.We found that, as in our data, learning pushed neural activity in the PFC towards a minimal XOR regime (Supp.Fig. 4).
Our findings indicate that neural activity in the PFC shifts between two distinct selectivity regimes throughout the learning process.Initially, the PFC maximally expanded the representational space by encoding all available variables.Subsequently, after a combination of task variables that predicted the trial's outcome was identified, the PFC reduced its dimensionality and maintained a minimal model of the task structure.Previous studies on cross-generalisation have suggested that such low-dimensional representations should improve learning of similar problems 7 .This is based on the notion that the classification of new stimuli can be rapidly improved by leveraging already learnt representations, thus resulting in a shared (abstract) representation of task variables between the old and new problems.We tested this prediction in experiment 2.

Generalisation: abstract coding dominates task representations over learning
In experiment 2, we introduced a new colour pair (stimulus set 2) that followed the same shape-outcome associations as the previous colour pair (stimulus set 1) (Fig. 1c).Randomly interleaving stimulus set 1 trials and stimulus set 2 trials allowed us to test whether a shared neural representation would be used for both stimulus sets (Fig. 1b).Similar to experiment 1, we recorded neural activity from the very first session in which the animals were exposed to the new stimulus set and divided the experimental sessions into four distinct learning stages.To explore learning-induced changes to neural representations, we again employed the decoding and cross-generalised decoding metrics.We next compared these metrics to changes in the structure of neural selectivities.Specifically, for each neuron, we computed its selectivity profiles for each variable (colour, shape, XOR), separately for each stimulus set, and measured their correlation between different sets.Notably, given that the colour was presented before the shape in the experimental setup (Fig. 1c), these analyses were conducted for each learning stage in both the colour-locked and shape-locked time windows (Fig 4 and Fig. 5, respectively).
Hypothetically, PFC could represent the new colours separately from the already learned pair (high-dimensional representation), or it could generate an abstract representation (shared across sets) of the colours' contextual meaning (lowdimensional representation).Through an exploration of neural activity in the colourlocked period (Fig. 1d), we observed that the PFC collapsed the stimulus set dimension and progressively aligned corresponding colours from both sets according to their meaning (Fig. 4a).Although context was equally well decodable in early as well as in late learning ( = .56,two-sided, Fig. 4b, black, decoding; see Supp.Fig. 5a for temporally resolved decoding), its format changed from an item-specific representation into an abstract code.We tested this explicitly by training a linear SVM classifier to decode the context in stimulus set 1 and tested this model on differentiating context in stimulus set 2, and vice-versa (see Methods, cross-stimulus set generalisation).We found that the representations of new and old colours became aligned as a function of learning, to allow for a single context readout ( = .000,onesided, Fig. 4b, purple, 'cross.decoding'; see Supp.Fig. 5b for temporally resolved cross.decoding).This change in geometry was reflected in a change in neural selectivity, as with the progression of learning, PFC cells exhibited an increasingly positive correlation between selectivity for colour pair 1 and colour pair 2 ( = .004,one-sided; Supp.Fig. 5d; Fig. 4c-d).Additionally, no learning-induced changes were detected when colour decoding was analysed separately for set 1 and set 2 ( = .94,two-sided and  = .16,two-sided, respectively; Fig. 4e).We also investigated the representation of the stimulus set (the old set 1 vs the new set 2) and found that the animals were able to clearly distinguish between the old and new colours during the early stages of learning (Fig. 4f, black, 'decoding').As learning progressed, however, we observed a reduction in stimulus-set decoding ( = .000,one-sided) and corresponding cross-context decoding of stimulus set ( = .02,one-sided, Fig. 4f, purple, 'cross.decoding'), as well a reduction in the shared selectivity ( = .04,onesided, Supp.Fig. 5c; Fig. 4g, h).This indicates that the animals collapsed the set 1 vs set 2 dimension.The emergence of an abstract representation of context before the appearance of the shape suggests that the network occupies a preparatory activity state that permits the immediate, context-appropriate transformation of shape information to reward prediction (XOR representation). 6,41,43,44e next focused on the neural activity in the shape-locked period, which reflected the state of the PFC when all necessary information for outcome prediction was available (colour and shape).Inspecting the neural representation of task variables on the population level suggested that set 1 and set 2 geometries aligned as a function of learning to allow a single XOR readout (Fig. 5a).We found that in both early learning (stage 1) and late learning (stage 4), XOR decoding was high and no learning-induced changes to the XOR signal were detected ( = .572,two-sided, Fig. 5b, black, decoding).This suggests that the animals were able to rapidly cross-generalise the XOR representation between stimulus sets as early as in stage 1.We tested this explicitly by training a linear SVM classifier to decode the XOR in stimulus set 1 and tested this model on differentiating the XOR in stimulus set 2, and vice-versa.We observed high scores of cross-generalised XOR decoding in both stage 1 and stage 4, as well as a learning-induced increase in this measure ( = .032,one-sided, Fig. 5b, purple, 'cross.decoding').We also observed a learning-induced increase in the shared selectivity for XOR between set 1 and set 2 ( = .042,one-sided, Supp.Fig. 5h; Fig. 5c-d;).These results suggest that, although the PFC constructed an abstract (shared across sets) representation of the XOR early in learning, it continued to align the set-specific XOR boundaries until this abstract format fully dominated the XOR representation.A trend-level difference between shape decoding was detected across the learning stages ( = .084,two-sided, Fig. 5f, black, decoding).Additionally, both cross-generalised decoding ( = .000,one-sided, Fig. 5f, purple, 'cross.decoding') and selectivity analyses ( = .000,one-sided, Supp.Fig. 5g; Fig. 5g,h;) showed that the PFC progressively aligned the shape representation between stimulus sets.It is important to note that the same shapes were used in both stimulus set 1 and stimulus set 2. These results suggest that, at the start of the learning process, the PFC encodes shape differently (using a high-dimensional representation) depending on the stimulus set, and that only during late learning (stage 4), the representations converge to a single shared axis (become low-dimensional).We next explored the representation of context (which we already introduced in Fig. 4) during the shape-locked period.We found no learning-induced changes to the geometry (decoding:  = .86,two-sided; cross.decoding:  = .374,one-sided) or selectivity ( = .4,one-sided, Supp.Fig. 5f; Fig. 5j-l).We speculate that the PFC already transformed the context signal in the colour-locked period, so that after this preparatory state was used to guide early shape processing, context information did not serve any relevant function in the late shapelocked period.Finally, we found that the decoding of colour and XOR in set 2 approached the levels of decoding for these variables in set 1 with learning ( = .044,one-sided, and  = .044,one-sided, respectively; Fig. 5e, m), whilst decoding of shape remained low in both sets ( = .25;one-sided, Fig. 5i).These changes to neural geometry were accompanied by slight changes in the dimensionality of representations (Supp.Fig. 5i-j).Specifically, as compared to the beginning of experiment 1, the mean shattering dimensionality of set 2 decreased slightly (but not significantly) over learning in experiment 2 (stage 1:  = .11,one-sided; stage 4:  = .06,one-sided), whereas the dimensionality of set 1 in experiment 2 remained consistently lower (stage 1:  = .04,one-sided; stage 4:  = .042,one-sided).

Discussion
The prefrontal cortex has the capacity to generate both low-dimensional and highdimensional representations, each of which presents a unique trade-off between generalisability and discriminability.However, the conditions under which each regime is employed currently remain unclear.Our study investigated how the dimensionality and geometry of neural activity changed over learning.We observed that, as learning progressed, neural activity in the PFC transitioned from being high-dimensional with high-discriminability to being low-dimensional and abstract.This transition in the representational strategy was accompanied by a change in population structure, from random and mixed, to minimal and structured.The structured representations that emerged during learning then supported the generalisation of the learned rule to a novel stimulus set.
We found that the transition between high-dimensional and low-dimensional representations can be reproduced by a model that solves an XOR problem under metabolic cost constraints.One recent theoretical proposal 45 used alternative constraints to argue that the emergence of structured selectivity is largely contingent on the complexity of the task.This was particularly evident in tasks that necessitate flexible input-output mappings, such as XOR rules.While our model explicitly defines the XOR nonlinearity as an input, this alternative model instead leveraged the temporal dimension to express the nonlinearity through recurrent dynamics 45 .However, the two models are not mutually exclusive, as both demonstrate that shortly before a decision is made, the system only retains relevant information (i.e., the XOR).This is accomplished by employing different forms of regularisation, such as an L2 cost on neural activities in our model, and a rank constraint on the recurrent connectivity matrix in the alternative model 45 .Another prominent theoretical proposal 42 describes neural geometries similar to those found in our optimised networks: output variables are represented in an abstract, cross-generalisable fashion.However, one difference with this proposal is that the metabolic cost constraints in our model cause task-irrelevant variables to no longer be represented after learning (Supp.Fig. 1h-i); and this is consistent with the findings in our data (Fig. 3h).
The acquisition of a low-dimensional representation following the learning of a single XOR rule raises the question of which regime the PFC might adopt when confronted with more complex tasks.Although XOR operations necessitate nonlinear integration and abstraction from sensory input, they can be reduced to simple stimulus-response pairings once the rule has been learnt or when a memory-based strategy is utilised.However, some tasks are more difficult to decompose, as they require switching between multiple orthogonal or conflicting subtasks.It has been suggested 18,20 that in conditions requiring the performing of multiple tasks in series, a high-dimensional representation could be employed by the PFC in order to maximise flexibility and prevent interference.
We found that the PFC acquired a low-dimensional representation to solve a complex cognitive task.This corroborates the findings of Hirokawa et al., 12 who found that neural activity covaried with behaviourally relevant variables on a single cell level, thus occupying a low-dimensional manifold.On the other hand, some studies have suggested that an increase in neural dimensionality is predictive of performance 14 .It is possible that the structure of the task and training provides an explanation for these contrasting findings 18 .In our study, recordings were initiated from the onset of task training, spanning a period of five weeks, with animals experiencing the entire task structure from the first session.In contrast, many other investigations into the primate PFC's involvement in complex cognitive tasks train animals in a fashion that decomposes the task into multiple subcomponents and either builds up task knowledge across training 14 or presents them in a serial (block-wise) manner 7,46 .Additionally, whereas in our study, one source of information (width) was always irrelevant, in other studies, information becomes periodically relevant and irrelevant across multiple blocks, which may promote encoding of currently irrelevant information 43 .These training differences may promote information encoding in a highdimensional manner that favours discriminability over metabolic efficiency.
Although we operationalised the metabolic constraint using an L2 cost on neural firing rates, several other related formulations are possible.The energy expenditure of a population can be reduced not only through a global (population-level) reduction of firing rates but also through the reduction of the number of actively firing neurons.This would naturally correspond to an L1 (rather than L2) cost on neural firing rates, which encourages sparseness of representations.Indeed, different levels of sparseness of neural firing rates could affect whether the representation is low-or high-dimensional and would also indirectly reduce an L2 metabolic cost on neural firing rates.However, our results indicate that the PFC likely does not rely on sparseness as a mechanism for metabolic cost reduction, as we did not observe any learning-induced changes in this measure (Supp.Fig. 3p).Additionally, imposing costs on the strength of synaptic connections rather than on neural firing rates would also be an interesting avenue for future work, especially when considering more realistic models with recurrent connectivities.
Many experimental efforts are guided by the implicit assumption that the animals have no prior knowledge nor received training related to the task being studied, thus constituting a tabula rasa system.Nevertheless, it is unlikely that the entirety of the animal's experience is unrelated to the experimental task, and that new representations do not interact with information already stored in the neural system.Our second experiment allowed us to address this issue and explicitly explore the interactions between past and new knowledge.In line with previous predictions 7, 17 , we found that when a new task instance is added, both the new and old instances were rapidly aligned to common axes and sensory differences between them were collapsed.
An alternative strategy for generalising previously learnt knowledge is exemplified by neural networks that develop compositional representations 47,48 .After being trained on multiple cognitive tasks, these models have the ability to decompose a variety of complex operations into simple computational motifs.Furthermore, compositional models are able to recombine these basic operations to rapidly solve tasks that were not previously seen 47,48 .Neural dimensionality could thus depend on the similarity of the new problem to the already learnt operations and the number of computational motifs required to effectively solve it.
It is perhaps surprising that, given the key role of the PFC in the development and acquisition of structured knowledge, only a few studies have investigated how the structure of PFC representations changes during several training days of an entirely novel task 49,50 .By tracking changes in neural activity across learning, it is possible to identify the biological principles that are required to produce representations supporting higher cognitive functions 21 .In this study, we have demonstrated that the constraint of metabolic cost reduction parsimoniously describes the changes in neural selectivity and geometry in non-human primates learning a complex cognitive rule.This provides a foundation for explaining how PFC representations should change with learning.Future experiments should extend this paradigm, to track changes in learning even more complex and naturalistic tasks 51 ; those that have a compositional structure 47,52 the influence of different learning curricula 53 ; and how these representations change within the same individual neurons as opposed to pseudo populations [54][55][56] .

Data and task
Animals and task.Two adult male rhesus macaques, monkey 1 and monkey 2, were trained in this study.The experiments were conducted in line with the Animals (Scientific Procedures) Act 1986 of the UK and licensed by a Home Office Project License obtained after review by Oxford University's Animal Care and Ethical Review committee.The procedures followed the standards set out in the European Community for the care and use of laboratory animals (EUVD, European Union directive 86/609/EEC).The animals were seated in a sound-and lighting-attenuated experimental booth.Their heads were restrained and faced a 19-inch screen.The centre of the screen was aligned with a neutral eye position.The animals performed a passive object-association task (Fig. 1c,d).Importantly, the animals were accustomed to an experimental setting but had no previous exposure to the task or stimuli introduced in this protocol.Neural recordings were collected from the first session as one of the main aims of the study was to capture learning dynamics.In the first experiment, the animals were presented with a colour and a shape, a nonlinear combination of which predicted reward (Fig. 1c,d).In experiment 2, a second set of stimuli was additionally introduced to test whether the rule learnt in the first experiment cross-generalised to the new sensory domain (Fig. 1c).The colours used in the coloured circles were designed in the CIELab colour space 57 .The L parameter (luminance) was kept constant for every colour which ensured that the stimuli where approximately isoluminant; parameters a and b varied with regard to valence but not value which resulted in a circular colour representation 57 .As colours were randomly assigned to conditions for each animal, this circular representation ensured that regardless of which colour pair was assigned to which XOR mapping, the initial colour similarity/dissimilarity within colour pair was kept constant.Additionally, in both experiments, the second object had two features: one relevant for reward prediction (shape) and one irrelevant (width) (for the duration and sequence in which stimuli were presented see Fig. 1d) The trial sequence was randomised.All trials with fixation errors were excluded.The dataset contained on average 237.9 ( = 23.9) and 104.8 ( = 2.3) trials for each of the 8 conditions in experiment 1, and 101.0 ( = 18.6) and 54 ( = 1.1) trials per each of the 16 conditions in experiment 2, for monkey 1 and monkey 2, respectively.
Data acquisition and pre-processing.Before the start of the experimental protocol, a titanium head holder with two recording chambers was placed and fixed with stainless steel screws in each animal.The frontal recording chambers were implanted over the lateral prefrontal cortex (lPFC) of the right hemisphere in both animals.Data from a second chamber targeting inferotemporal cortex in the right hemisphere are not considered here.A craniotomy was made beneath each chamber to enable electrophysiological recording.Recording locations for each animal are shown in Supplementary Figure 5. Surgical procedures were carried out under general anaesthesia and were aseptic.A semi-chronic micro-drive system (SC-96, Gray Matter Research) with 1.5 mm interelectrode spacing, interfaced to a multichannel data acquisition system (Cerebus System, Blackrock Microsystems) was used for frontal recordings.Data were recorded over a total of 25 daily sessions in each monkey (monkey 1: 17 sessions in experiment 1 and 8 sessions in experiment 2; monkey 2: 10 sessions in experiment 1 and 15 sessions in experiment 2).The switch to experiment 2 was made after the animal showed a robust reward prediction signal.Notably, electrodes were manually advanced by a minimum of 62.5  before every session to ensure that activity from new cells was recorded.Neural activity was amplified, filtered (300  − 10 ), and stored for offline pre-processing and analysis.Cluster separation was applied (valley seeking algorithm), and the binary spike train was smoothed using a Gaussian window ( = 50).We collected spiking activity from 146 and 230 neurons in experiment 1 and from 205 and 151 neurons in experiment 2, for monkey 1 and monkey 2, respectively.Only cells sampled from the ventral and dorsal lateral frontal cortex were included in the data (Supp.Fig. 6).No neurons were excluded based on their selectivity profiles.Importantly, as the focus of this study was to track how learning influenced neural geometry and not the magnitude of firing rate (e.g., repetition suppression effects), we z-scored firing rates of each neuron across the whole session.During this procedure a small constant was added to the denominator to prevent small firing rates having large impact on neural representations (for similar normalisation see 58 ).The obtained firing rate data were then epoched from 200  before to 1200  after the colour onset.Next, all sessions were divided into four learning stages separately for each animal and experiment; the sessions within each stage were pooled across animals.All subsequent analyses were implemented in Python using custom-written code and run on combined data (monkey 1 and 2).Two types of analyses were used in this study: (1) timepoint-resolved, where a specific method was applied to every time point in the epoch to track how representations evolved in trial time, and (2) time-averaged, where a method was run on time-averaged data (e.g., [t !""#$ , t %""#$ ], colour-locked or shape-locked) in the time window preceding the shape display or trial outcome.In the former time window, we examined the neural geometry when only the colour information is known (i.e., context representation), whereas just before outcome onset (latter time window), we examined whether neural geometry reflected the colour and the shape and their combination (XOR) before the animals received feedback about the value of the trial.

Constantinidis et al. 2016 dataset.
We also used an existing dataset of electrophysiological recordings 32 which have been described in detail previously 33,34 .In brief, neural activity was recorded from the ventral and dorsal lateral PFC (identical to the areas targeted in this study) in four rhesus monkeys who performed a feature match-to-sample task.More specifically, the animals were required to report after a delay period whether the shape of the first stimulus was the same as the shape of the second stimulus.Note that a match/no-match rule is equivalent to an XOR rule.Importantly, neural activity was recorded before the animals were exposed to the task rule (passive viewing) and after they had learned the rule.As both correct match and correct no-match trials were rewarded, the match/no-match signal was not confounded with a reward prediction signal.To test whether neural activity was pushed towards a minimal regime in such experimental conditions we employed the same decoding and selectivity measures as used in the analysis of our dataset (see Fig 2).We examined neural data averaged across the presentation of the second stimulus and the subsequent delay period ([t "#$ , t &"""#$ ]; stimulus 2-locked).Furthermore, neural activity was analysed for all stimulus pairs combined.For the 8 stimuli, we paired them into 4 sets of pairs and performed our analyses separately on each pair of stimuli (and averaged results over all 4 pairs) so that chance decoding was the same as in our dataset (i.e., 50%).

Models
Multiple linear regression.We can model the firing rate  of a neuron (either from our generative models or our data) at a given time point as a linear combination of the three main task variables: colour, shape and the interaction between colour and shape (the XOR term): where  is a vector of 1x K dimensionality containing the time-averaged firing rates for K trials;  is the design matrix of dimensionality K x D where rows correspond to the K trials and columns correspond to the value of the D task variables such as colour, shape and XOR (D = 3) in each trial. is a D-by-1 vector populated with the coefficients for each of the task variable estimated for the th neuron.Finally,  contains K residuals.The  vector specifies the coordinates of the th neuron in the selectivity space spanned by D task variables (Fig. 2a,b and Fig. 3b,c).That is, every neuron can be represented as a point in a space where each axis corresponds to the cell's selectivity for a task variable.
Generative models.Neural selectivity can be defined by the matrix  '()( = (S *' ) +,*,-,+,',/ , where the th row corresponds to the th neuron, and columns correspond to the regression coefficients for the three task variables colour, shape, and colour x shape (XOR) that form the axes of the considered space. 5This cloud of points is then centred by removing the mean (∑  *' = 0 ' , for each of the  task variables).Here, we explored two types of selectivity distributions and their representational properties.Firstly, we examined a random mixed selectivity model in which selectivities are captured by a spherical multivariate Gaussian distribution  0(*'1# ~  ' (,  &   ).In such a model, all variables can be decoded equally well from the population resulting in a high-dimensional representation and there is poor crossgeneralisation between variables.The second selectivity model we examine results from a system performing the XOR task while being constrained to exhibit low overall firing rates (i.e., a form of metabolic cost).We derived analytically that maximising XOR decodability while minimising such a metabolic cost results in neurons being selective only to the interaction term (colour x shape; XOR) and having no selectivity to the linear terms (colour or shape; Supp.Materials Section 1).A matrix describing the selectivity of such a population can be thus formulated as  #3*3#(4 ~  ' Q,  & diag(0,0,1)V.We call this the minimal XOR model.Importantly, to allow comparisons between the observed selectivity and model selectivity (minimal or random), the generative models were constructed using parameters derived from the data.Specifically, we used the mean value of diagonal entries of the covariance matrix  X '()( estimated from  '()( to set the value of the variance parameter  & in both  0(*'1# and  #3*3#(4 (note that this ensures that both models have the same total variance).Furthermore, we showed that multiple linear regression was able to recover the underlying minimal and random models from artificially generated firing rates under various levels of noise (Supp.Fig. 2).
Optimised feedforward networks.We used  = 400 neurons in these networks and their firing rates were described by eq. 1 with  = 2.The output  of these networks was given by a softmax readout where  567 are the two sets of readout weights and  is the readout bias.We optimized these networks with back-propagation using a canonical cross-entropy cost function (3) where the first part of eq. 3 denotes the cross-entropy loss ℋ(, ) between the true probabilities of reward  (which were equal to 0 or 1, depending upon the stimuli for that trial) and the model's readout probabilities  and the second term corresponds to a metabolic cost on all firing rates.Before training, the values of the 's were drawn randomly from a Gaussian distribution with 0 mean and variance 0.05, and elements of  567 and  were set to 0. We trained two types of models: (1) with no regularisation ( = 0) and ( 2) with high regularisation ( = 20; Supp.Fig. 1a and 1b, respectively).In line with our predictions, selectivity to task variables in models trained with high regularisation converged on the minimal selectivity regime whereas models trained with no regularisation produced randomly mixed selectivity (Supp.Fig. 1).

Decoding.
To test what information was represented by the observed neural population as a function of learning, we employed linear SVM decoding 7,14 .In contrast to the regression analysis, where the regression coefficients could be estimated for every neuron separately, a decoding analysis required the construction of pseudo neural populations.To achieve this, were divided into four learning stages.As each session had a different number of trials, the sessions were aligned to the last trial (to maximise learning effects) and the same number of trials per condition was used to construct the new pseudo-population (e.g. the last 99 and 49 trials per condition were selected from each of the sessions in experiment 1 and 2, respectively).This resulted in a matrix  = (X 9*) ) +,9,:,+,*,-,+,),; for each of the four learning stages, where the first dimension corresponds to K trials, the second dimension corresponds to N neurons (combined from two animals), and the last dimension corresponds to T time points.Then binary SVM classifiers were used to decode the task variables (colour, shape, width and XOR in experiment 1 and context, stimulus set, shape, width and XOR in experiment 2) at each time point for temporally resolved decoding (Supp.Fig. 3a-d and Supp.Fig. 5a,b) and from time-averaged firing rates ([t <+""#$ , t "#$ ]; outcome-locked, Fig. 3f, Fig. 5b,f,j; colour-locked, Fig. 4b,f).An equivalent decoding procedure was used when analysing generative models (Fig. 2f).Decoding was performed in a cross-validated way where K trials were split randomly into set 1 and set 2, with each containing 50% of trials.The decoder was fitted using the set 1 and tested on set 2. The procedure was then repeated using set 2 as the training set and set 1 as the test set.Both decoding scores were then averaged.This procedure was repeated 10 times for different random splits of trials in sets 1 and 2, and these 10 resulting scores were then averaged.To estimate shattering dimensionality 7,14 in experiment 1(Fig.3h) and 2 (Supp.Fig. 5i,j), we used the same decoding approach except that we averaged decoding scores over all 35 possible dichotomies that could be theoretically represented given the task structure (i.e., three linear variables form a cube in state space that can be dissected into two sets of 4 vertices in 35 possible ways).We estimated the decodability of each of these dichotomies and tracked their mean decoding accuracy as a function of learning.To assess learning-induced changes to the linear and nonlinear components of neural dimensionality separately, we split the dichotomies into linear (n = 7) and nonlinear decoding axes (n = 28) and tracked their mean decoding accuracy as a function of learning (Fig. 3i).Please note that exploring the theoretically maximal dimensionality of the neural representations in experiment 2 using shattered dimensionality, where four linear variables (colour, stimulus set, shape, width) were used, would result in the identification of 6435 theoretically possible binary decoding problems and thus be computationally expensive.Instead, we tracked shattering dimensionality separately for stimulus set 1 and stimulus set 2.
Cross-generalised decoding.To examine the neural geometry of the task representation we used cross-generalised SVM decoding 42 .Trials were divided according to the to-be-decoded target variable labels as well as the splitting variable labels.More specifically, the splitting variable provides two instances of data which can be used to train and test the target variable classifier, respectively.In contrast to a typical cross-validation procedure, the testing happens not only on trials that were previously not seen but also on trials that correspond to different conditions.To achieve a high cross-generalisation score, it is therefore not sufficient to generalise across trial-wise noise but also to generalise across conditions 7 .As the task examined in experiment 1 has three linear variables (colour, shape, width), 36 possible crossgeneralisation axes can be identified (i.e., combinations of decoded variable and variable over which decoding was cross-generalised) per each of the 35 target variables.That is, the labels for 8 conditions (2   2 ℎ  2 ℎ) are split into two sets of four labels (by colour, shape, or width).To identify training and testing exemplars within each of the sets, every possible binary combination of condition labels was determined (4 ℎ 2).This yielded 6 combinations for each of the sets.Thirty-six unique train-test splits (cross-generalisation axes) can be achieved by combining these two sets.Decoding accuracy scores estimated by running linear SVM decoders for each of these splits were then averaged to obtain a general crossgeneralisation score of the analysed target variable.
Cross-stimulus set generalisation (decoding and selectivity analysis).Crossgeneralised decoding performed for experiment 2 data (both temporally resolved and run on time-averaged firing rates; Fig. 4b,f, Fig. 5b,f,j and Supp.Fig. 3a-d, Supp.Fig. 5a,b, respectively) differed in one aspect from the algorithm described in the Cross-generalised decoding section.As the aim of the analysis described here was to identify the neural format of the main task variables used across stimulus sets, only one splitting variable was used (i.e., stimulus set) to obtain cross-generalisation scores for the task-relevant variables (context, shape, and XOR).This reduced the possible cross-generalisation decoding axes to four possible binary decoding problems (e.g., when performing cross-generalised decoding for the colour variable we can: (1) train on differentiating colour 1 from colour 2 in stimulus set 1 and test on differentiating colour 3 from colour 4 in stimulus set 2, (2) train on differentiating colour 3 from colour 4 in stimulus set 2 and test on differentiating colour 1 from colour 2 in stimulus set 1, (3) train on differentiating colour 1 from colour 3 and test on differentiating colour 2 from colour 4, and (4) train on differentiating colour 2 from colour 4 and test on differentiating colour 1 from colour 3; these four decoding scores were then averaged).Using this procedure, we explored the cross-stimulus set generalisation potential of the colour, shape, width and XOR variables.Additionally, to test how selectivity of PFC cells changed as a function of learning in experiment 2 we employed a Pearson correlation metric.Specifically, we compared how similar the colour, shape and XOR coefficients in stimulus set 1 (Fig. 4c,g, and Fig. 5c,g,k) are to coefficients for the same variables in stimulus set 2 (Fig. 4d,h and Fig. 5d,h,l), which yielded three correlation scores for each of the main task variables (Supp.Fig. 5c-h).This was done for each of the four learning stages to explore whether selectivity for stimulus set 1 aligns with selectivity for stimulus set 2 as a function of learning, consistent with a shared abstract representation.
(5) where the subscript  denotes a random draw and the expectations were computed over 1000 draws.From both the denominators and numerators, the distance within each of the models was subtracted to centre the measure around 0. More specifically, (x X  0(*'1# −  X  0(*'1# x  )(the expected difference between two different randomly drawn selectivity distributions from the random model) was, for example, subtracted from the denominator and numerator of  =01# 0(*'1# to account for within model distance.Additionally, both  =01# 0(*'1# and  =01# #3*3#(4 were normalised by the distance between selectivities generated using both generative models ((x X 3 0(*'1# −  X 3 #3*3#(4 x

𝟐
)) which resulted in the metrics being bounded between 0 and 1 (when x X '()( −  X #3*3#(4 x  is equal or smaller than x X 0(*'1# −  X #3*3#(4 x  ).This was done to allow for a comparison of similarity estimates across learning stages.The Euclidean distance metric was chosen as the main analysis tool in this study based on simulations in which we generated different proportions of random and minimal selectivity across a single population (from 0% minimal and 100% of random to 100% minimal and 0% random) and compared the precision with which multiple metrics recovered the true proportions.We compared the Euclidean distance metric to the PAIRs metric, which has been used previously in the literature 12,35 , and to the symmetric Kullback-Leibler divergence estimate (KL divergence) which benefits from a strong theoretical basis and is assumption-agnostic.We found that, compared to the KL divergence and PAIRs metrics, the Euclidean distance measure yielded the highest precision of tracking learning-induced changes to neural selectivity (Supp.Fig. 7a,b).Specifically, our simulations showed that both the KL divergence and PAIRs can be used to precisely identify extreme selectivity regimes (either strong random selectivity or strong minimal selectivity) but fail at identifying intermediate selectivity regimes showing a strong bias towards random mixed selectivity (Supp.Fig. 7a,b).the focus of this study was to track learning dynamics, a metric that allows to identify a broad range of selectivity regimes was chosen for the final analysis.Nonetheless, the results from experiment 1 (Fig. 3d,e) were broadly replicated using the symmetric KL divergence estimate (Supp.Fig. 7c,d) and PAIRs (Supp.Fig. 7e,f).
Principal Component Analysis.PCA was used to visualise the population firing rate for different conditions in experiment 1 (Fig. 3a) and experiment 2 (Fig. 4a and 5a) and as a measure of neural dimensionality (Supp.Fig. 3o).Firstly, pseudo populations were constructed for each learning stage using the same procedure as described in the Decoding section.Then, firing rates were averaged in the time window preceding the shape presentation ([t !""#$ , t %""#$ ] colour-locked; Fig. 4a) and outcome presentation ([t !""#$ , t %""#$ ], shape-locked; Fig. 3a and 5a).Next, principal components were computed on condition averages and the population firing rates we projected onto the top three principal components.This was done separately for each learning stage.Note that width 1 and width 2 trials were pooled together.To compute how the variance explained (ratio) by the first PC changed as a function of learning (Supp.Fig. 3o), trials were randomly split into test and train 10 times; PCA was fitted then on train trials and the test trial firing rates were projected onto them to compute variance explained.The results from 10 random splits were then averaged.The null distribution for the permutation test was computed by randomly shuffling neurons between stage 1 and stage 4, and repeating the described PCA procedure ( = 500).

Statistical testing
Decoding and cross.gen.decoding.Throughout the study, we employed nonparametric permutation tests to test statistical significance within each learning stage and between learning stages (learning-induced effects).Two types of null distributions were thus constructed: (1) for statistical testing in time-resolved and time-specific decoding analyses the labels describing the trial dimension () of the pseudopopulation matrix  = (X 9*) ) +,9,:,+,*,-,+,),; were randomly permuted 500 times; (2) to test for learning-induced effects, the matrices  $)(GH + and  $)(GH ! were concatenated along the neuron dimension () and then 500 new ′ $)(GH + and ′ $)(GH !matrices were generated by randomly assigning neurons to either ′ $)(GH + or ′ $)(GH ! . One-sided tests were used when testing the predictions of the minimal model and twosided tests were used when no differences were expected.Additionally, to control for time-related family-wise errors a cluster-based permutation correction was added to the time-resolved decoding 59 .Selectivity measures.To test whether observed selectivity was dissimilar to the random selectivity regime and similar to minimal selectivity regime 1000 random and minimal models were generated using data-derived parameters for each learning stage.Next, the  0(*'.=01#0(*'.and  0(*'.=01##3*.distances were computed for 1000 randomly generated models according to eq. 4 and eq. 5 (with  X 0(*'1# as input) to serve as null distributions for both comparisons.Note that the observed selectivity was compared to random model selectivity when analysing the data's similarity to random (Fig. 3d) as well as minimal selectivity (Fig. 3e).Furthermore, as in experiment 2 we tested whether selectivity for task variables was similar in stimulus set 1 to variables in stimulus set 2 and whether this selectivity alignment changed over learning, two null distributions were thus constructed: (1) statistically significant selectivity alignment was assessed by comparing the observed correlation to a distribution ( = 500) of correlations obtained after shuffling one of the selectivity vectors; (2) learning-induced effects in selectivity alignment were assessed by comparing the observed difference in alignment between stage 1 and stage 4 to a distribution of differences computed after randomly shuffling neurons between stage 1 and 4.

Section 1.2: An energy cost on unnecessary neural activity
If we also consider minimising the squared norm of neural activity for each condition (i.e., an energy cost), we have & ‚ = ƒ J,$ " N ƒ J,$ " + Tr( + ) where the subscripts  and  correspond to colour and shape indices, respectively.Therefore, Therefore, the total mean energy cost  is given by =  " N  " +  K N  K +  L N  L +  M N  M + Tr( + ) To minimise  while keeping  M > 0, which we need for performance, we can set  " =  K =  L =  which gives

Figure 3 .
Figure 3. Analysis of neural representations in macaque PFC over learning.a, Trial-averaged population firing rate during the late shape period for each task condition projected onto the top three principal components as a function of learning (Methods, principal component analysis); the same scale was used for all plots.b,c, Neural selectivities in learning stages 1 and 4, respectively.Each point represents the selectivity of one neuron for each of the task variables (colour, shape and XOR).We show all 3 possible pairs of the 3 axes.In panel b, the 3 red rings show spherical Gaussian distributions corresponding to 1, 2, and 3 standard deviations of the data.In panel c, the 3 blue rings show ellipses from the minimal model that correspond to 1, 2, and 3 standard deviations of the data (note that they all lie on top of one another along the XOR selectivity axis).The bottom right panel shows the covariance matrix of the selectivities computed from the data for stage 1 of learning (b) and stage 4 of learning (c).d, Relative Euclidean distance between the covariance matrix of selectivity coefficients from the data

Figure 4 .
Figure 4.The PFC aligns new and old colours as a function of learning to allow a single readout of context.a, Trial-averaged population firing rate during the late colour-locked period for each colour projected onto the top three principal components as a function of learning (Methods, principal component analysis); the same scale was used for all plots.b, Decoding (black) and cross-generalised decoding (purple) of context as a function of learning.The purple and grey shaded areas indicate chance-level cross.gen.decoding and chance-level decoding, respectively.c, d, Selectivity of PFC neurons in stage 1 and stage 4 for context in stimulus set 1 and stimulus set 2 in the colour-locked period.The line of best fit is shown in black.e, Decoding of colour separately for stimulus set 1 (grey, 'blue vs green') and stimulus set 2 (black, 'pink vs khaki').f-h, Analogous to b-d but for set coding (f) and set selectivity (g, h).All p-values were calculated from permutation tests (***, p < 0.01; **, p <0.01; *, p < 0.05; n.s., not significant).Shading in panels b and f shows the mean ±1 s.d. of chance-level decoding.

Figure 5 .
Figure 5.The PFC utilised abstract stimulus representations when previously unseen stimulus combinations were introduced.a, Trial-averaged population firing rate during the late shape-locked period for each condition projected onto the top three principal components as a function of learning (Methods, principal component analysis); the same scale was used for all plots.b, Decoding (black) and cross-generalised decoding (purple) of XOR as a function of learning.The purple and dark grey shaded areas indicate chance-level cross.gen.decoding and chance-level decoding, respectively.c,d, Selectivity of PFC neurons for XOR in stimulus set 1 and stimulus set 2 in the shape-locked period in stage 1 and stage 4, respectively.The line of best fit is shown in black.e, Decoding of XOR shown separately for stimulus set 1 (grey) and stimulus set 2 (black).The grey shaded area indicates chancelevel decoding.f-m, Analogous to b-e but for shape (f-i) and context (j-m).All p-values were calculated from permutation tests (***,  < 0.01; **,  < 0.01; *,  < 0.05; †,  < 0.01; .., not significant).Shading in panels b, f, and j shows the mean ±1 s.d. of chance-level decoding and cross-set decoding.