Disentangling Representations through Multi-task Learning

Response to the reviewer (1/2)

We thank the reviewer for the thoughtful comments and suggestions, and for acknowledging the contributions of this work. Here we address the key points raised: ## Theoretical Contributions and Linear Disentanglement The reviewer asks whether linear (abstract) and orthogonal (disentangled) representations constitute interesting disentanglement. Linear representations are interesting for a couple of reasons: 1) they tend to emerge in neural networks (e.g. Templeton et al., 2024), and we don't really understand why, and 2) they make it easy to extract meaningful latents one layer downstream, which is a useful property since linear probes are often used. Disentanglement (orthogonality) is also useful for this reason, i.e. it allows linear decoders to tease apart individual latents with minimal interference from others. We show that abstract representations emerge from optimal (or approximately optimal) multi-task classification on decision boundaries that are linear in the “ground truth” latent space (implying $N_{task} \geq D$), despite potentially complex observation maps $f$. We also show that disentangled representations (linear + orthogonal) have to emerge from (potentially approximately) optimal multi-task classifications when $N_{task} \gg D$ span the ground truth latent space. We also foreground that our theory is derived for affine classification boundaries in the latent space. ## Clarification of Multi-Task Learning Theoretically, the tasks fundamentally define the latents the model must represent within our problem formulation. If tasks span only a subset of dimensions in $D$, we guarantee linear representation of that subset if the model is optimal. For sub-optimal models with zero-mean error, the machinery of least squares can be applied to combat prediction error if there if $N_{task} > D$ (equation 19, lines 1507-1520). We foreground this result in the new version of the manuscript in the theory in the main text, and as a corollary of Theorem B.5 (Corollary B.7). Furthermore, we want to clarify that indeed by evidence aggregation we mean evidence accumulation, and we also include experiments more closely mirroring drift diffusion tasks in Appendix A.11. We will also think of an example task that we could best explain this multi-task setting to the neuroscience reader. ## Connection to Predictive Processing Our methods complement predictive processing approaches in the following way. Our key theoretical contribution is establishing a connection between * A system’s **tasks/classification outputs** and * A system’s representation of underlying latent variables. The central question for connecting to the predictive processing view is: **“What set of latent variables is implied by the task of predictive processing?”.** This represents a promising direction for future work, particularly in analyzing the latent representations in pre-trained autoregressive language models. ## Bayesian Filtering We note interesting connections to state space models and particle filters on manifolds, which relate to our framework when the observation map $f$ is a manifold. A current area of active research is extending the theory to include noise both before and after a manifold observation map $f$. Conditions relating to the reach of the manifold $f$ and its Fisher information are of great interest for understanding the conditions under which the observation map $f$ is expected to be invertible, particularly for computationally or energetically bound filters/observers. ## Sparse Coding We appreciate the connection to the sparse coding literature. Of particular interest is the connection between sparse representation learning and the particular distribution of decision boundaries in the latent space, noise properties, and the structure of the observation map $f$. Theoretically, our approach centers on connecting the multi-task predictions $\hat Y_i(t)$ produced by the evidence aggregator and the structure of the underlying latents $x^*$. A key challenge in deriving sparsity results on $Z(t)$ -- as opposed to $\hat Y(t)$. We apply the data processing inequality to show that $Z(t)$ contains the information of $\mu(t)$ -- the optimal estimate of $x^*$ given the classifier’s input data. The specific structure of this representation in $Z(t)$ is challenging to constrain by our approach, except for the final layer L of $Z(t)$ is mapped through element-wise activation function $g$ to $\hat Y(t)$ (i.e., $\hat Y(t) = g(Z_L(t))$). The specific structure of the representations in $Z(t)$ before the final layer are challenging to constrain in terms of sparsity or lack thereof as different regularization schemes on the same RNN that achieves similar performance may lead to different coding schemes for non-final layers in $Z(t)$. Yet, we now rigorously quantify sparsity in our trained RNNs, and show that it consistently emerges across all values of $N_{task}$ and $D$ of our trained networks (Appendix A.12).

Response to reviewer (2/2)

A future direction of interest for understanding the advent of sparsity within our framework is to develop theory on (approximately) optimal multi-task classifiers under energetic constraints on the representations in $Z(t)$. This also relates to neuroscientific theories of efficient coding. ## Information Maximization and Noise Our results relate to the infomax principle but focus specifically on mutual information between ground truth latents $x^*$ and classifier state $Z(t)$, rather than mutual information maximization between the inputs of the network $f(X_1) ... f(X_t)$ and outputs $\hat Y(t)$. Our work particularly resembles nonlinear independent component analysis, with (approximately) optimal multi-task performance guaranteeing inversion of nonlinear observation map $f$ under certain conditions. The parallel with random Gaussian fluctuations around potentially non-Gaussian distributed ground truth latents is especially noteworthy. While the noise distribution in the main theorem is Gaussian, the theorem does not require Gaussian-distributed latent variables $x^*$. Importantly, our framework reveals that the tasks effectively define which latent dimensions of $x^*$ are represented in $Z(t)$ -- if your tasks only span a subset of $x^*$’s dimensions in $\mathbb R^D$, an optimal classifier is only guaranteed to represent those task-specific dimensions. This aligns with information theoretic intuitions: the mutual information between $Z(t)$ and the dimensions of $x^*$ irrelevant to the classification tasks need not be preserved for optimal performance. This provides a simple, principled way to think about how tasks demands shape representation learning. In contrast to infomax and nonlinear ICA, the multiple classification tasks in our framework specify which latent dimensions are relevant to the tasks at hand. This **task-driven perspective on representation learning** aligns with neuroscientific observations of task-specific representations in different brain areas (i.e., various levels of hierarchical organization). Different neural circuits appear to disentangle only those features relevant to their computational role. This suggests that task demands, in addition to general information preservation, may be a useful organizing principle for both biological and artificial representation learning. ## Generative Model Formulation We can indeed formulate our approach as a generative model by $x^*$ by sampling from $\hat X = \mathcal N(\mu(t), t^{-1} \sigma^2 I_D)$, where $\mu(t)$ is a deterministic function of the multi-classifier’s state $Z(t)$. This requires calibration of expected noise $\sigma$ with knowledge of time $t$. For non-Gaussian distributions (covered in B.5), knowledge of the particular shift-and-scale distribution is required. The generative model’s accuracy depends on available data at time $t$ -- fewer samples results in higher variance in the optimal estimator $\mu(t)$ about $x^*$ due to limited evidence aggregation. Since $\mu(t)$ is a deterministic function of classifier state $Z(t)$ and outputs $\hat Y(t)$, we are particularly interested in leveraging the reviewer’s observation to understand how predictive multi-task classification models (e.g., autoregressive language models) model the world differently based on their probability distribution over next tokens (analogous to $\hat Y(t)$). Our current work suggests a path to *conjuring* latent variables/underlying causes $x^*$ from a set of existing $\hat Y(t)$ (discussed above under connections to predictive processing, information maximization and noise). We are especially interested in this as a technique to characterize the “world models” developed by such systems. Finally, we would like to thank the reviewer for the references; we have cited some of that literature when talking about predictive coding (Rao & Ballard (1999), Rao (1999)). ### References: 1. Snoussi, Hichem. "Particle filtering on riemannian manifolds. application to covariance matrices tracking." In Matrix information geometry, pp. 427-449. Berlin, Heidelberg: Springer Berlin Heidelberg, 2012. 2. Tompkins, Frank, and Patrick J. Wolfe. "Bayesian filtering on the Stiefel manifold." In 2007 2nd IEEE International Workshop on Computational Advances in Multi-Sensor Adaptive Processing, pp. 261-264. IEEE, 2007.

Thank you

We would like to sincerely thank the reviewer for their thoughtful review, insightful comments, and support for our work.

Response to reviewer (1/2)

We would like to thank the reviewer for the very thorough and useful review, and for taking the time to read the paper including the proofs very carefully. In regards to the specific points raised: ## Theoretical Framework and Mathematical Precision ### Theoretical Framework Clarifications The reviewer correctly highlights the constraint implicit in our problem formulation that the $N_{task}$ “ground truth” task boundaries $y_i(x^*) \in \\{ 0, 1 \\} , x^* \in \mathbb R^D$ are linear. In other words, our theoretical results rely on the multi-task classifier making optimal (or approximately optimal) predictions on $N_{task}$ linear classifications on ground truth latents $x^*$ based on noisy, non-linearly mapped observations. We now explicitly mention that in the theory section of the main text (paragraph “More general decision boundaries”) along with desiderata and propositions on generalizing our main theoretical results on linear boundaries in latent space to smooth manifold decision boundaries. We also include these considerations here for reviewer convenience: ### Extending theoretical results to nonlinear boundaries on latents Our existing results on linear boundaries $y_i(x^*)$ in latent space function as a solution to locally linearized smooth manifold decision boundaries. Specifically, we can apply our existing results on linear decision boundaries to manifold decision boundaries $y_i(x^*)$ when $\tau_i \gg \sigma$, where $\tau_i$ is the reach of manifold $y_i$. The reach of a manifold is the largest distance from the manifold at which every point in space has a unique closest point on the manifold. Deriving maximally permissive conditions on multi-task manifold decision boundaries $y_i(x^*)$ that guarantee the emergence of disentangled/abstract representations of latents $x^*$ is left for future work. We have also re-worded and clarified the statements that the reviewer mentioned in the abstract. In particular, we now specify that the theoretical results are on “classification tasks affine in the latent space” and we removed the word “general” from the last sentence of the abstract. ### Simplicity We acknowledge that our main theoretical results build on straightforward arguments. We believe this simplicity is a strength, as it provides an accessible foundation for extending our results on the disentanglement of $x^*$ in the state of optimal multi-task classifiers. **Our results aim to clearly and simply show the conceptual insight that distances in latent space are reflected in an optimal (or approximately optimal) evidence aggregating multi-task classifier’s uncertainties.** ### Extension to sub-optimal classifiers While our main results are stated for optimal classifiers, Equation 19 (line 1515) handles sub-optimal classifiers with zero mean error, inheriting from least squares estimation. We formalize this into a corollary of theorem B.5 (Corollary B.7) and give it more prominence in the main text. ### Data processing inequality The application of the data processing inequality (DPI) in Theorem B.5 follows from the markov chain $x^* \to {X(t)} \to Z(t) \to \hat Y(t) \to \mu(t)$, which simplifies to $x^* \to Z(t) \to \mu(t)$. The DPI states that if $X \to Y \to Z$, then $I(X; Y) \geq I(X; Z)$. Applying this to our chain yields $I(x^*; Z(t)) \geq I(x^*; \mu(t))$, which directly supports our theorem’s main claim about information content in the latent state $Z(t)$. We would be thankful for further clarification from the reviewers of errors committed in this line of reasoning. ## Technical Refinements ### Column space and activation functions We acknowledge that our statement about column space representations in line 231 was overstated. We have removed that statement completely, focusing on empirical results instead (e.g. robustness to noise distribution/activation function mismatch in Appendix A.9). ### Singular value analysis Regarding corollary B.10, we incorporate a discussion of the Marchenko-Pastur law while noting important caveats: while the law typically applies to matrices with i.i.d. entries $N(0, \sigma)$, our $C$ matrix consists of random row vectors in $R^D$ of unit norm. This structure, while not strictly meeting the law’s conditions, still supports our conclusions about orthogonalization. We also corrected the statement in line 1683 to read “non-orthogonal” instead of “orthogonal”. ### Statistical precision We improve precision in our treatment of random variables and estimators in lines 1496-1499. Additionally, we clarify that $\hat Y_i(t)$ represents a Bernoulli random variable parameter estimator for assigning probabilities to binary classes, rather than a Bernoulli random variable itself in other places of Appendix B.

Response to reviewer (2/2)

## Cognitive maps We state in the manuscript that our work offers a principle for the formation of cognitive maps (or “world models”). The suggested principle is that cognitive maps (world models) emerge from parallel processing (multi-task learning). In biological systems, this could be implemented at the level of cortical columns which have been proposed to operate independently from each other (Hawkins et al, 2019), and in artificial neural networks this bears close resemblance to self-supervised learning regimes (see l. 1908-16 in the manuscript). Cognitive maps is a concept often invoked to describe brain representations that capture the structure of the real world, in a spatial (e.g. grid cells) or a more high-level semantic level. Abstract, or disentangled, representations do exactly that, as they separate latent factors present in the real world in orthogonal axes. Finally, the link to zero-shot generalization was established in the experiments throughout the paper where we report zero-shot generalization performance $r^2$. We hope this helps clarify how these concepts come together in this work. ## Decreasing dimensionality of $Z$ To examine when the latent encoding fails, we decrease the number of neurons when $D=10$, i.e. when we need the most amount of neurons to disentangle all latents. For $N_{task} = 24$ and $N_{neu}=64$, median $r^2 = 0.9$, therefore representations are abstract. However when decreasing $N_{neu}=32$, median $r^2 = 0.52$, and when $N_{neu}=16$, median $r^2 = 0.19$, i.e. generalization performance is completely hampered. Therefore we conclude that while there is no strictly defined threshold, $N_{neu}$ should be a multiple of $D$ to ensure enough capacity to perform all the functions that the network does (encode the factors, invert nonlinear encoder, denoise etc.). ### In regards to the reviewer’s other questions: * We thank the reviewer for pointing out the vagueness in Fig 2a. We now describe how the right plot is created by adding noise, and how the color-coded boolean variables in $y$ are determined by the classification lines of the corresponding color. * Line 298: This is more of a statement here, rather than a finding, because for $D=2$ the classification lines will always span the space. This assumption is really tested in Figure 5b. * We agree that Fig. S10 (now Fig. S12) might be a bit hard to read, but it can be seen that the mean angle for GPT is always close to 90 degrees (with the exception of $(N_{task}=6, D=10)$ which is below the required $N_{task} \geq D$), while for RNNs it can vary (especially for $D=2$ and $D=4$). * We thank the reviewer for catching the “disentangled” vs. “abstract” statement regarding Fig. S3b (now Fig. S4b). We have corrected it. We have also addressed all minor comments of the reviewer in the new version of the manuscript. Again, we thank the reviewer for the very careful reading of the paper.

Response to reviewer

We thank the reviewer for their clarification of the issues with the current proof. We have revised the proofs for Theorem B5-B6 to remove the inaccurate argument as follows: Since $\mu(t)$ is the optimal estimator (and sufficient statistic) for $\mathbf x^*$ given measurements $\{\mathbf X(1), \dots, \mathbf X(t)\}$, it contains all information about $\mathbf x^*$ contained in the measurements (Cover & Thomas, 1991). In other words, $\mu(t)$ is a deterministic function of $\mathbf Z(t)$, implying that $\mathbf Z(t)$ will contain all information about $\mathbf x^*$ contained in the measurements $\{\mathbf X(t)\}$ An alternative way to answer this question would be to condition mutual information on the noisy samples, which are known to the agent. We chose this way because it is cleaner. Happy to answer any more questions that the reviewer may have, as the discussion period winds down.

Response to reviewer

We thank the reviewer for their continued assistance in improving the clarity of the proof. Here we streamline the logical reasoning for Theorem B.5, showing that $\mathbf Z(t)$ must contain all information about $\mathbf x^*$ given the noisy measurements $\mathbf X(1), ..., \mathbf X(t)$. We show below how that ties in with the notion of sufficient statistics which we introduced in our previous response. Most of the proof is already in the paper, we outline new additions in detail:  ## Proof Sketch: Theorem B.5 **Goal**: Show that optimal multi-task classifier latent state $\mathbf Z(t)$ contains all information about $\mathbf x^*$ given noisy measurements $\{\mathbf X(1), \dots, \mathbf X(t)\}$ if $\mathbf C \in \mathbb R^{N_{task} \times D}$ is full-rank and $N_{task} \geq D$. **Background on sufficient statistics:** "A statistic $\mu(t)$ is called sufficient for $\mathbf x^*$ if it contains all the information in $\{\mathbf X(1), \dots, \mathbf X(t)\}$ about $\mathbf x^*$." (verbatim from Cover and Thomas' Elements of Information Theory, 1999, Section 2.10, substituting variable names). Therefore, if we show that $\mathbf Z(t)$ contains a sufficient statistic $\mu(t)$ of $\mathbf x^*$, we have completed the proof. More formally, "A function $T(\{\mathbf X(1), \dots, \mathbf X(t)\})$ is said to be a sufficient statistic relative to the family \[of probability density functions indexed by $\mathbf x^*$\] $f(\{\mathbf X(1), \dots, \mathbf X(t)\} | \mathbf x^*)$ if $\{\mathbf X(1), \dots, \mathbf X(t)\}$ is independent of $\mathbf x^*$ given $T(\{\mathbf X(1), \dots, \mathbf X(t)\})$, i.d., $\mathbf x^* \to T(\{\mathbf X(1), \dots, \mathbf X(t)\})$ forms a Markov chain. This is the same as the condition for equality in the data processing inequality, $$I(\mathbf x^*; \{\mathbf X(1), \dots, \mathbf X(t)\}) = I(\mathbf x^*; \mu(t))$$ for all distributions on $\mathbf x^*$. Hence sufficient statistics preserve mutual information and conversely." (Cover and Thomas' Elements of Information Theory, 1999, Section 2.10, substituting variable names) $\mu(t) = \text{mean}(\mathbf X(1), \dots, \mathbf X(t))$ **is a sufficient statistic for** $\mathbf x^*$ **given measurements** $\mathbf X(i) \sim \mathbf x^* + \sigma \mathcal N(0, \mathbf I_D)$: For Gaussian noise it's a well known result that the sufficient statistic for the underlying mean given i.i.d. samples is the sample mean of the observations (Cover and Thomas, Elements of Information Theory, 1999, Section 2.10). **Showing that $\mu(t)$ is a deterministic function of $\hat {\mathbf Y}(t)$:** We showed in Theorem B.4 (Trilateration Theorem) that if $\mathbf C \in \mathbb R^{N_{task} \times D}$ is full-rank and $N_{task} \geq D$, then $\hat {\mathbf Y}(t)$, $t$, $\mathbf b$, and $\mathbf C$ are sufficient to reconstruct the exact value of $\mu(t)$, the mean of $\mathbf X(1), \dots, \mathbf X(t)$. Rearranging Equation 16 and applying Equation 15, $$ \mu(t) = (\mathbf C^\top \mathbf C)^{-1} \mathbf C^\top (\frac{\sigma}{\sqrt{t}} \Phi^{-1}(\hat{\mathbf Y}(t)) + \mathbf b) $$ Replacing $\hat{\mathbf Y}(t) = g(\mathbf Z(t))$ from our problem setup reveals that $\mu(t)$ is a deterministic function of $\mathbf Z(t)$. Therefore, optimal multi-task classifier latent state $\mathbf Z(t)$ contains a sufficient statistic $\mu(t)$ of $\mathbf x^*$, which implies that $\mathbf Z(t)$ must also contain all information about $\mathbf x^*$ given noisy measurements $\{\mathbf X(1), \dots, \mathbf X(t)\}$ if $\mathbf C \in \mathbb R^{N_{task} \times D}$ is full-rank and $N_{task} \geq D$. $$\square$$ We hope this clarifies our proof to the reviewer. We cannot make changes to the manuscript now, but we will include these new clarifications in the new version. The same fix applies to theorem B.6.

Response to reviewer (1/4)

We thank the reviewer for their review. In regards to the specific points raised: ## Definition of multi-task learning The reviewer seems to disagree with framing our setting as multi-task learning, citing the lack of diversity in the tasks. They also point out that related work by Huh et al (2024) met this diversity requirement, by training networks on multimodal tasks. However, diversity in the sense of multimodality is not a requirement for multi-task learning. Multi-task learning is defined as the “training of [related] tasks in parallel while using a shared representation” (Caruana (1997)). Therefore, our work falls within the multi-task learning framework. Furthermore, our definition of task diversity is mathematical, not colloquial (“multimodality”). By diversity of task demands we refer to whether the decision boundaries *span* the latent space, involving all latent variables in the decisions, which is a condition for disentangling all factors. So while we do provide results from a mixture of classification boundary geometries in the appendix, we keep the tasks as simple as possible to derive theoretical insight into the nature of generalization and world model construction by means of rigorous proofs, which would be difficult to get a grasp on if we used rich multimodal datasets instead (as **reviewer 2Aef** neatly points out). We agree that it would be interesting to extend the experiments to multi-modal tasks, but this would be a new follow-up study on its own. ## Regarding the feasibility of disentanglement There seem to be two concerns here, that are of different nature and which we therefore address separately. First, from a theoretical standpoint, the reviewer argues that disentangling representations might be difficult, or even impossible, citing previous work on unsupervised learning (Locatello et al (2019), Hyvarinen & Pajunen (1999)), which is not what we do. Meanwhile, we mathematically prove conditions under which disentangling representations is not just possible, but *guaranteed*. In particular, a system being approximately optimal at tasks involving the same latents *is* an inductive bias; therefore our results do not contradict Locatello et al (2019). Furthermore, the “impossibility result” in Locatello et al rests on the narrow view of disentangled representations as axis-aligned, i.e. a single ground truth factor change should lead to a single change in the representation, similar to the single neuron coding assumption in neuroscience. We argue in favor of a mixed representation (Rigotti et al 2013) view, because by looking at single units we might be missing cases where disentanglement exists at the population level; a point also raised by Johnston & Fusi (2023). On a second note, the reviewer raises a concern about the supervised structure of our tasks, claiming that such supervised datasets would be difficult to find in real-world settings. We disagree, as in fact modern self-supervised learning involves converting an unsupervised objective (“scene understanding”) into a well-posed supervised objective (image patch-filling) (see l. 1909-16 in the manuscript). Filling missing patches *is* a multi-task objective, with a common, shared latent space (objects present in an image, and their relations), and solving each one of the objectives acts synergistically in understanding the image latent space as a whole. The same logic can be extended to other modalities, for example in language, when predicting masked words that rely on a common latent sentence “meaning”. Our study is a first effort towards understanding such parallel learning, and providing guarantees for its performance. In the future, we plan to extend the insights from this study to real-world dynamic datasets. ## Regarding “abundance of supervised tasks” One of the main concerns of the reviewer is the apparent use of “dense supervised signals”. Yet, our setting requires the minimum possible number of tasks ($N_{task}$) mathematically required for identifiability in $D$ dimensional space. To uniquely identify a point in $D$ dimensional space, one needs distance information from at least $D$ boundaries, i.e. $N_{task} \geq D$. This is exactly the condition we find for linear decodability in our setting (eqs 3 and 4). Therefore, for linear decodability, no abundance of supervised tasks is required - only the minimum $N_{task}$ that would make identifiability mathematically feasible. Regarding orthogonality, our theoretical bound is less tight here ($N_{task} \gg D$). However, empirically we find that neural networks disentangle much earlier (figure 5b), and in particular, transformers disentangle for the minimum possible number of tasks that identifiability is possible (again, $N_{task} = D$). Compared to alternatives for representation learning in the brain (Mante et al 2013), which requires the construction of task-specific subspaces scaling exponentially with $D$, our work confers enormous algorithmic advantages.

Response to reviewer (2/4)

The reviewer concludes by saying that: > “future work in this area should investigate whether disentanglement can emerge from a smaller number of more diverse tasks, where the need for efficient generalization across tasks drives abstraction” However, this is exactly what our work is showing. Task diversity is defined by the dimensionality of the subspace they span in the latent space, and $N_{task} = D$ tasks that span the entire latent space is the minimal task set that ensures identifiability. Furthermore, one of our main findings is not that the need for generalization drives abstraction; rather generalization is a by-product that comes about by the pressures imposed when learning to solve multiple tasks optimally, given noise and time constraints. On that note, even having less tasks than the dimensionality $D$ is not a problem: in that case, we expect that our networks would learn to disentangle the subspace spanned by the available tasks. We expand more on the dynamic interplay of $N_{task}$ and $D$ implied by our work in the discussion and main rebuttal above. We would also like to note that the reviewer’s assertion that the >“ underlying structure of the data… [is] provided” in our setting is incorrect; the only thing provided to the networks are binary decisions driven by that underlying structure, not the underlying structure itself. The structure is discovered by means of optimally performing (classification) tasks in parallel. Being able to extract latents just from solving multi-task learning objectives is an exciting guarantee that opens up a lot of future research directions. ### Conflation of “competence at N tasks” with “receiving feedback signals” on N tasks The reviewer seems to conflate “competence at N tasks” with “receiving feedback signals on N tasks”. Our theory doesn’t require external signals. It just requires approximate optimality on the multi-task classification objective -- regardless of how and when that was learned. E.g., you could find such reps in a pre-trained LLM if you prompt it right. **Our contribution is that competence on N tasks -> uncovering latent variables.** In other words, the labels provided to the network in Fig. 2b need not come externally. They can originate internally, for tasks that the organism is already competent at. Biological agents are remarkable *because* they achieve high performance on many tasks, and as a result they uncover latent variables that underlie the tasks (as we show here). We expand more on this in the next section. ## Biological plausibility of multi-task learning The reviewer mentions that this multi-task learning setting might be not biologically relevant, as organisms rarely receive such rich feedback from the environment. However, while our theory stems from parallel processing, i.e. multi-task learning, it is not contingent upon the parallel execution of multiple tasks, i.e. multitasking, or the receipt of such feedback in parallel. Behaviorally, the agent need only perform one action, the one most appropriate to its current internal state (e.g. its level of thirst vs. hunger might control the slope of the decision boundary in 2D space of water & food). What we posit is that tasks that have been performed by the agent before and rely on the same input are still resolved somewhere in the brain, by the brain circuits (e.g. cortical columns Hawkins et al. (2019)) previously responsible for them, instead of the entire decision-making brain area focusing only on the current task (Mante et al., 2013). Therefore, the output of these tasks is still placing pressure on the representation, even though they are not actively driving behavior. We feel that this is a more natural way of thinking about how the brain manages different tasks, with older tasks still leaving traces somewhere in the brain ([Losey et al., 2024](https://www.sciencedirect.com/science/article/abs/pii/S0960982224002987)). This is in a similar vein to the quotation from Johnston & Fusi (2023) about prediction of outcomes under different hypotheticals that the reviewer provides (see [Jensen et al, 2024](https://www.nature.com/articles/s41593-024-01675-7)). A link to the predictive coding could also be established here (see **reviewer b5qB’s** insightful comments). The key question for connecting to predictive processing-centric views is: “What set of latent variables is implied by the task of predictive processing?” Furthermore, we also show that our setting generalizes to interleaved learning of tasks with different decision geometry structure; the relationship between interleaved and multi-task learning is an interesting direction for future research, and would more closely mimic neuroscience experiments where new tasks are introduced sequentially. We have added these comments in Appendix A.14 of the revised manuscript.

Response to reviewer (3/4)

## Regarding “incremental progress” claim The reviewer argues that our work is incremental progress over [Johnston & Fusi (2023)](https://www.nature.com/articles/s41467-023-36583-0). While our paper builds on their work, and we did utilize the experimental framework of this study, for its conceptual simplicity and amenability to theory, **our findings go way beyond the findings of Johnston & Fusi (2023) in several ways** (see final paragraph in section 1.1, “Related work” in the original manuscript), which we repeat here for reviewer convenience: * We extend the findings of Johnston & Fusi (2023) to **autoregressive architectures** that can update their representations as further information arrives. Going from feedforward to recurrent/autoregressive architectures already presents technical challenges that make the work non-incremental, and opens up the framework to a family of **dynamic tasks** that feedforward architectures are not suited for. We go beyond that to: * **Prove theorems that guarantee** the emergence of linearly decodable and orthogonal representations in **any** optimal multi-task classifier, and experimentally confirm the findings * Showcase the **importance of noise for disentanglement**, showing that more noise during pretraining results in better out-of-distribution generalization. Johnston & Fusi (2023) did not include noise in their experiments, which is an essential component of biological neural networks * Explore **a range of values for $D$**, providing experimental validation of our theory * Include **correlations between latent factors**, rendering entire areas in state space virtually unseen by the network during training * **Quantify disentanglement**, in addition to abstractness * Show that the substrate for disentangled representations in RNNs is **continuous attractors**, that maintain a short-term memory of the current evidence and updating it over time * Show the advantage of our approach over previously utilized frameworks for representation learning in the brain (“context-dependent computation”, Mante et al 2013) * Demonstrate that **transformers are ideally suited for disentangling representations**, achieving linearly decodable, orthogonal representations at the theoretically minimum possible number of tasks ($N_{task} \geq D$). This is of particular relevance for the machine learning community at large, since there is a lot of interest in what makes transformers so good at world understanding. While judgments about novelty are always subjective, and as acknowledged in the manuscript we build our work on the important contribution of Johnston & Fusi (2023), **it is hard for us to see how these advancements constitute incremental progress, since Johnston & Fusi (2023) focused on feedforward architectures that cannot handle noisy dynamical tasks, and did not provide any theoretical proofs for disentanglement or, as a matter of fact, abstractness.** As a side note, we would like to note that the experimental framework that we utilized was also used in [Maziarka et al. (2023)](https://dl.acm.org/doi/10.1007/978-3-031-26387-3_38) which has a very similar setting to Johnston & Fusi (2023) (generate synthetic objectives based on some sampled ground truth) and arrived at similar conclusions (multi-task setting results in disentanglement in feedforward architectures), although less unequivocally. In addition, the idea that multitask learning can lead to better generalization has been around for decades (Caruana, 1997). Our personal inspiration stemmed from the fruit fly head direction system, where an explicit representation of head direction is enforced presumably due to its downstream usage by many circuits for flight control etc. (Wilson, 2023). **Our work is, to our knowledge, the first time that theoretical guarantees for disentanglement through multi-task learning in autoregressive architectures have been provided and experimentally confirmed.**

Response to reviewer (4/4)

## Sparsity We now rigorously quantify sparsity in our trained networks using the sparseness metric suggested by the reviewer (Vinje and Gallant, 2000). We find that RNNs are consistently sparse across values for number of tasks $N_{task}$ and latent dimensionality $D$ (see Appendix A.12 in updated manuscript). Therefore we conclude that sparsity is a phenomenon that consistently emerges in our trained RNNs, an attribute they share with their biological counterparts. Furthermore, as we now note in Appendix A.12, we do not impose any regularization or other constraints for sparsity; rather it naturally emerges from the optimization objective and network architecture. We would like to note though that **while sparsity is an interesting emergent phenomenon in our networks, our work presents much more far reaching and interesting findings than sparsity, as recognized by all other reviewers.** ## More nonlinear encoding In new experiments, we replace the ReLUs in the encoder with power-law (quadratic) nonlinearities (Appendix A.5 in revised manuscript). We find no impact on generalization performance. Therefore, we conclude that our setting is robust to the choice of encoder nonlinearity, even when the nonlinearity is not injective, going beyond our theoretical proofs. We have also cited the paper suggested by the reviewer ([Alleman, Lindsey & Fusi 2024](https://arxiv.org/abs/2401.13558)), as motivation for these experiments. Finally, we would like to clarify that the observation map is determined by the environment, not the brain, therefore the units of the encoder are not modeling biological neurons. Overall, we feel that these new experiments enhance the overall presentation of our paper, and quantify aspects (e.g. sparsity) that were not quantified before.

Questioning the framing of this work as 'multi-task learning'

I appreciate the authors’ thorough responses, the additional experiments, and the revisions they implemented in their submission. However, my primary concern remains: the **supervised paradigm** fundamentally limits the work's biological plausibility and its relevance to real-world, unsupervised scenarios. --- Below are my detailed responses to specific points made by the authors: ## 1. Framing this work as “multi-task learning” is misleading The paper frames multiple binary classification tasks as multitasking, which I find misleading. > They also point out that related work by Huh et al (2024) met this diversity requirement, by training networks on multimodal tasks
 The multimodal example I mentioned is not central to my critique. Multitasking does not require multimodality: rather, it requires distinct but related tasks that differ in their objectives. > By diversity of task demands we refer to whether the decision boundaries span the latent space … Task diversity is defined by the dimensionality of the subspace they span in the latent space This definition does not align with the standard understanding of task diversity in multitask learning. Tasks that share the same objective—such as binary classification across different decision boundaries—do not constitute diverse tasks in the traditional sense. For example, in a single-modality multitask learning scenario, such as image-based tasks, diversity arises from fundamentally different objectives: identifying the primary object in the image (*Object Recognition*), predicting bounding boxes around objects (*Object Localization*), classifying the type of scene (*Scene Understanding*), estimating pixel depth (*Depth Estimation*), and reconstructing pixel values (*Reconstruction*). Despite using the same modality, each task in this example interprets the input differently and requires specialized processing. In contrast, this paper's setup appears more like **single-task learning with multiple outputs**, as the tasks are essentially variations of the same binary classification task. While decision boundaries may span the latent space, the lack of task-specific diversity (i.e., distinct objectives) makes the classification framework less representative of what is traditionally understood as multitask learning. To address this, I recommend amending the paper's title to reflect its content more accurately. - The current title, "*Disentangling Representations through Multi-task Learning*," may create expectations of true multitask learning involving diverse objectives. - A more precise title could be: "*Disentangling Representations through Multiple Binary Classification Tasks*". While this revised title may sound less broad or ambitious, it more accurately represents the work and its contributions.

The limited relevance of disentanglement through strongly supervised frameworks

## 2. Disentanglement in unsupervised settings > modern self-supervised learning involves converting an unsupervised objective (“scene understanding”) into a well-posed supervised objective (image patch-filling) (see l. 1909-16 in the manuscript) I appreciate the mention of modern self-supervised learning techniques, but I find this tangential to the discussion. The authors' approach is fundamentally supervised, and drawing parallels with self-supervised methods may introduce unnecessary conflation between the two paradigms. While self-supervised learning can incorporate elements of supervision through engineered objectives, it fundamentally differs from the explicitly supervised framework employed here. (Minor: It also appears the cited line references [1909-16] may be incorrect.) > The same logic can be extended to other modalities, for example in language, when predicting masked words that rely on a common latent sentence “meaning”. This example of masked word prediction in language models again seems tangential. My critique is not about the applicability of disentanglement to other modalities, but rather the reliance on a supervised framework in this paper. The focus should remain on the relevance of the current setup to disentanglement in unsupervised or weakly supervised settings, which this paper does not address. > To uniquely identify a point in D dimensional space, one needs distance information from at least D boundaries, i.e. N_task ≥ D. This is exactly the condition we find for linear decodability in our setting (eqs 3 and 4) Thank you for clarifying this. The linear relationship between the number of necessary tasks and input dimensionality is indeed less problematic than I initially thought. However, the reliance on supervision still limits the broader applicability of this work to real-world, unsupervised scenarios. The strength of disentangled representations lies in their emergence from data without requiring extensive supervision, and the current paradigm misses this key point. > the only thing provided to the networks are binary decisions driven by that underlying structure, not the underlying structure itself Indeed, the underlying structure is not directly provided, but the supervision signals effectively encode it through the binary decisions spanning the latent space. Proving this theoretically is neither surprising nor interesting. This approach positions the paper closer to the fully supervised end of the learning spectrum. Within supervised learning, the reliance on many supervision signals spanning the entire latent space limits the paper's relevance to unsupervised learning, which has been an essential focus of disentanglement research for decades. > Conflation of “competence at N tasks” with “receiving feedback signals” on N tasks From the perspective of building AI systems that learn to disentangle representations, your results appear to advocate for providing extensive supervised signals across all tasks (please correct me if not). Is this the most scalable or realistic approach? This reliance on supervision restricts the relevance of the findings for advancing models capable of disentanglement in real-world, minimally supervised scenarios. > They can originate internally, for tasks that the organism is already competent at. This raises an important question: How does the organism become competent in these tasks to begin with? Supervised or unsupervised learning? Without addressing how such competence is acquired, the paper does not provide actionable insights for models that have to learn disentangled representations from scratch. This omission limits the real-world applicability of the results, leaving them speculative at best. > The reviewer mentions that this multi-task learning setting might be not biologically relevant, as organisms rarely receive such rich feedback from the environment. To clarify, my primary concern lies in the reliance on a supervised paradigm. While I question the biological plausibility of this feedback-rich setup, my critique extends beyond biological relevance to the fundamental limitation of this work in advancing disentanglement beyond supervised frameworks.

Incremental contributions within a supervised paradigm

## 3. Clarifying my comment regarding incremental progress I appreciate the authors highlighting how their work extends Johnston & Fusi (2023) and agree that it provides several valuable results beyond that study. However, my central critique remains unchanged: the **supervised paradigm** fundamentally limits the scope and impact of this line of work. Johnston & Fusi (2023) acknowledge in their discussion the inherent limitation of their supervised framework and emphasize the importance of future work moving beyond this artificial setup. By building on Johnston & Fusi (2023) without transitioning to an unsupervised framework, the current submission remains firmly grounded in the same supervised paradigm. While the paper refines and extends Johnston & Fusi (2023), it does not take the critical leap toward unsupervised disentangled representation learning that the field has been striving for. ## Final thoughts Researchers grappling with the complexities and challenges involved in unsupervised disentanglement recognize that advancing the field requires adhering to unsupervised, or at least weakly supervised, frameworks, which have been the focus of this area for decades. From this perspective, I maintain that the present work—while rigorous and well-executed—is incremental in essence, as it does not address the broader challenges or opportunities for unsupervised disentanglement.

Our theory is agnostic to the specific learning paradigm, focus of our work

**We are glad our responses resolved a lot of the reviewer’s previous confusions and contentions**, including the concern about requiring an abundance ($N_{task}>>D$) of tasks, the claim that we directly hint the underlying latent structure to the network, and that our work constitutes incremental progress over Johnston & Fusi (2023), and that **the reviewer acknowledged that our work is rigorous and well-executed.** We will address some contentions that seem to persist or have emerged in their new round of responses, but first it would be useful to establish what our work is and what it is not about. **More generally, it seems that the reviewer has a fundamentally different view of what our paper should be. Even though we understand their particular interests, we think we should be allowed the freedom to set our own goals for our paper.** ## Focus of our work > The focus should remain on the relevance of the current setup to disentanglement in unsupervised or weakly supervised settings, which this paper does not address. **While we appreciate and share the reviewer’s interest in the area of unsupervised disentanglement, our focus is not in this area.** Instead, we take a different perspective, **focusing on providing a formal link between competence at multiple tasks and disentanglement**, and experimentally confirming and extending the theory on tasks canonical in the cognitive neuroscience literature. While we recognize that proposing **new algorithms for unsupervised disentanglement** is a particularly important area of research, it **is not the only important question when it comes to disentangled representations** (see section *“Multiple ways towards understanding disentangled representations”* below). In addition, our theory goes beyond our experiments on supervised disentanglement, because it does not assume how competence at tasks was achieved, leaving the door open for representations that have been learned through supervised, self-supervised, or unsupervised means; our supervised experiments are simply the fastest way to get good performance at multiple tasks in a tractable manner, and test our theory. Therefore, **overly focusing on the supervised nature of the experiments misses a key message of our study.** >The strength of disentangled representations lies in their emergence from data without requiring extensive supervision, and the current paradigm misses this key point. Our theory is agnostic to the specific learning paradigm by which disentangled representations emerge. **It provides a formal argument why disentangled representations may emerge in a system: because the system is competent at multiple tasks.** Key strengths of disentangled representations include the fact that they separate latent factors present in the real world along orthogonal directions, enabling OOD generalization, efficient downstream learning etc, therefore providing any guarantees on their emergence is valuable. >How does the organism become competent in these tasks to begin with? Supervised or unsupervised learning? This is an important question, but given the reviewer’s frequent references to neuroscience, they surely understand that it cannot be fully addressed within a single study. Also, it is not the focus of this study. >the paper does not provide actionable insights for models that have to learn disentangled representations from scratch. We have to respectfully disagree here. Our work shows that if you train a network from scratch on a set of objectives, the network learns a disentangled representation of the latents that are spanned by these tasks. Our experiments are on supervised learning, but as mentioned throughout our responses and in the paper, the obvious next step is scaling that to self-supervised learning in future work. ## Self-supervised learning The connection between our framework and self-supervised learning is not tangential, but deep and promising. As we mentioned in the section *"Regarding the feasibility of disentanglement"* above, the two frameworks share a common structure of an underlying latent truth (e.g. objects in an image and their relationships), and solving each one of the objectives (e.g. filling each one of the missing image patches) acts synergistically in understanding the latent space as a whole. Same logic applies to predicting masked words that rely on a common latent sentence “meaning”. Our study provides guarantees on the disentanglement of representations learned as a function of the number of objectives solved in parallel. As such, self-supervised is a natural extension of our work, and we will be testing how the main findings of our theory generalize in this setting. The same points are mentioned in the new l. 1962-70 of the final version of the manuscript (apologies for the wrong line numbers before, as we added some more things in the meanwhile).

Our title is not misleading, etc

Other points: >From the perspective of building AI systems that learn to disentangle representations, your results appear to advocate for providing extensive supervised signals across all tasks (please correct me if not). We are not trying to advocate in favor of scaling supervised learning here. **Supervised training is our way of getting to our goal of competence at $N_{tasks}$ efficiently, and it is not at all our suggestion of how to scale up disentanglement efforts.** Self-supervised approaches offer much more promise at that, as just mentioned. > While I question the biological plausibility of this feedback-rich setup, my critique extends beyond biological relevance to the fundamental limitation of this work in advancing disentanglement beyond supervised frameworks. As just mentioned, we are not trying to suggest scaling up disentanglement through supervised training efforts. And there’s more than one way to advance the field of disentangled representations, for more see section *“Multiple ways towards understanding disentangled representations”* below. >but the supervision signals effectively encode it through the binary decisions spanning the latent space. Proving this theoretically is neither surprising nor interesting This looks like one of the main findings of our work. Judgements as to whether a certain result is surprising or not are of little practical value, especially in hindsight. ## Framework is multi-task learning by definition The reviewer claims that calling our framework multi-task learning is misleading. At the risk of sounding like a broken record, we will reiterate the definition of multi-task learning from the formative work of [Caruana (1997)](https://link.springer.com/article/10.1023/A:1007379606734), which is widely accepted: > Multi-task learning is defined as the “learning of [related] tasks in parallel while using a shared representation” Therefore, our work neatly falls within multi-task learning. Furthermore, we clearly stated our definition of diversity of tasks, which is mathematical, well-motivated and easily understandable by anyone. Tasks are diverse when they engage a lot of latent factors. The less latent factors that are engaged, the less diverse the set of tasks. There’s different ways to get to task diversity (i.e. engage more latents), including multimodality, different vision objectives that the reviewer references, or the way we sample classification lines here. Our way happens to be clean and amenable to theory. Regarding the title, while we appreciate the reviewer’s suggestion, we will keep it as is. This is not because the current title is more ambitious (why is multi-task learning ambitious?), but because it better reflects the nature of our work. In principle, it doesn’t have to be binary classification tasks, and in practice we also show that our framework extends to integration-to-bound tasks where networks directly report their confidence, not just a binary decision (see Appendix A.11). **The importance of the confidence of a network’s output is a recurring theme in machine learning (e.g. knowledge distillation etc.). We here show that it fundamentally connects to how neural networks construct world models, either directly (integrate-to-bound task) or indirectly (classification task).**

Reviewer seems to hold the view that only unsupervised disentanglement is not incremental

## Regarding incremental progress We are glad that the reviewer now seems to agree that our work is not incremental with regards to Johston & Fusi (2023). The title that the reviewer chose for this section of their response *“Incremental contributions within a supervised paradigm”* can be a bit confusing, but based on context we will assume it means that they think our work is incremental because our experiments are on supervised learning; not because our results are incremental with regards to other work that used supervised learning for their experiments (e.g. Johnston & Fusi (2023)). Just to reiterate, **the supervised paradigm is simply our way to ensure that the networks are competent at multiple tasks which is a requirement of our theory, and secondary to the main message of our work, which is that competence at $N_{task} \geq D$ tasks leads to disentangled representations.** The theoretical insights stand separately from the specific paradigm used, and apply to networks that have achieved competence at multiple tasks through other means, including unsupervised learning. We acknowledge that a limitation and future direction of our work would be to apply our setting to richer tasks (e.g. self-supervised learning) in the paper and throughout the rebuttals. In the limitation section of our paper (now appendix A.15) we also mention other limitations, such as that factorization and feature-based generalization might not be the best strategy always. We cannot make any more modifications to the paper at this point, but we will make sure to connect our mentions of self-supervised learning to the importance of advancing the field of unsupervised learning of disentangled representations. Finally, the discussion in Johnston & Fusi the reviewer is referencing is centered around offline prediction for more biologically plausible versions of the multitasking framework, which we already addressed in the previous round of rebuttals (section *"Biological plausibility of multi-task learning"*). >I maintain that the present work—while rigorous and well-executed—is incremental in essence, as it does not address the broader challenges or opportunities for unsupervised disentanglement. **It seems from the above that while the reviewer no longer considers our work incremental in regards to Johnston & Fusi (2023) as they did before, they now consider it incremental because it doesn’t focus on unsupervised disentanglement.** This brings us to the next point. ## Multiple ways towards understanding disentangled representations Throughout the rebuttals, the reviewer seems to **argue strongly in favor of a particular approach for advancing the field of disentangled representations**: in particular, they seem to argue that if one is to advance the field of disentangled representations, they should focus on answering the question of how neural networks learn disentangled representations through unsupervised learning (also see section “Focus of our work” above). This seems quite restricted to us. While we recognize the importance of and the complexities associated with the unsupervised learning of disentangled representations, **this clearly cannot be the only important research question.** Two other key ways to advance the field of disentangled representations that come to mind are: 1. **Provide theoretical guarantees** for the existence of disentangled representations under some conditions - **our work makes important progress on that front**, irrespective of how these representations were learned and 2. Define new metrics that can **better detect** disentangled representations in existing networks. Our study informs thinking in this space, by arguing in favor of a mixed representation approach to disentanglement (fig. 4), while most metrics to date look for axis-aligned representations. **Our work, along with others** (e.g. Johnston & Fusi, 2023) **suggests that we might be missing disentanglement in networks that perform well in many tasks, due to the metrics used.** We are sure that there will be many other possible ideas and research directions that will be useful as the field of disentangled representations moves forward, and we are excited for the future of the field. Finally, we genuinely hope that our remaining interactions with the reviewer are constructive, and that the reviewer takes the time to carefully consider our responses.

Unfortunately, the reviewer repeats the same objections after we addressed them

We thank the reviewer for their response, and **we are glad that they now agree that testing our theory on self-supervised learning is an actionable next step.** We are not the first to come up with the idea for mixed disentangled representations (e.g. we cite Johnston & Fusi 2023), however a lot of influential work focuses on axis-aligned representations (Higgins et al. 2017, Kim & Mnih 2018, Locatello et al. 2019); we want to advocate in favor of mixed representations before axis-alignment becomes an indispensable part of the definition of disentanglement in the CS community. Regarding multi-task learning, we will be sure to outline our definition of task diversity, which is mathematical in nature. **When it comes to our theoretical guarantees and the origin of competence at multiple tasks however, it seems that the reviewer is still misunderstanding our work, and repeating their previous objections that we just addressed** (e.g. section [*Addressing reviewer Mvb9’s misunderstanding of a key contribution of our work*](https://openreview.net/forum?id=yVGGtsOgc7&noteId=ez9jQvXiOy)). ## Our theoretical guarantees apply to representations learned through unsupervised means > the theoretical framework and most of the experimental results explicitly rely on supervised signals to achieve this competence > In my view, guarantees in unsupervised or minimally supervised settings would be far more impactful. As mentioned several times in our rebuttal, and contrary to the reviewer’s assertion, **our theoretical framework does not rely on supervised signals** to achieve competence at tasks; it just assumes competence. **The conditions for our theory are not related to the learning paradigm** (supervised or unsupervised), **but rather to the geometry of the tasks on which competence is achieved.** Therefore, **our theory also covers cases where competence was achieved by unsupervised or minimally supervised means.** ## Focus of our work > Without addressing how such competence could emerge in real-world, unsupervised scenarios, the broader significance of the results remains unclear. **Addressing how competence at tasks can emerge in real-world settings through unsupervised means is an entirely different question** than our focus here. As we mentioned before, there are more than one ways to make progress towards understanding disentangled representations (see corresponding section [here](https://openreview.net/forum?id=yVGGtsOgc7&noteId=ikR5XGez9y)), and **we focus on providing theoretical guarantees on disentanglement, and rigorously testing them in tractable settings** (see [here](https://openreview.net/forum?id=yVGGtsOgc7&noteId=QjrzaVKFN0)). Testing our theory in self-supervised settings would introduce significantly more overhead, that goes far beyond the scope of a single study (see below). ## Our experimental results already extend past our theoretical guarantees > Moving beyond this paradigm toward weakly supervised or self-supervised approaches, as the authors themselves discuss, would significantly enhance the relevance and impact of this line of work. In our discussion that the reviewer quotes, we were referring to the experiments specifically. While it remains to be seen how the theoretical insights generalize to self-supervised settings, this is not because of the learning paradigm per se, but because of the different **task geometry** implied by the specific application (e.g. image patch-filling, next-token prediction, iterative de-noising). **We have experimentally shown that our findings generalize to other task geometries beyond what is strictly guaranteed from our theory** (e.g. see Appendix A.10, Appendix A.5), **which is promising for the broader application of our findings across task geometries, including those that might be encountered in self-supervised settings** (also mentioned under *“New research directions”* in the [main rebuttal](https://openreview.net/forum?id=yVGGtsOgc7&noteId=0GglZZbzCV)). Extending our results to self-supervised settings would likely require the introduction of new datasets that can take advantage of the temporal dimension, like dynamic versions of dSprites; **thus, this would constitute an entirely new paper on its own.** ## Summary Overall, it appears that the reviewer has a specific vision of what our paper should address, which differs from the actual scope and contributions of our work. **It remains unclear what the reviewer expects in response to their repeated objections regarding “strong reliance on supervised signals”, despite our repeated clarification that supervised signals are not necessary for our theory;** competence on multiple tasks is. We do not agree with the reviewer, and it seems that they remain unwilling to reconsider their stance despite thorough clarifications.

Response to reviewer (1/2)

We would like to thank the reviewer for the in depth review and comments. Regarding the individual comments: ## Dimensionality of latents One thing that we might need to specify here is that $D$ is the dimensionality of the latents (e.g. features like ripeness, fruit identity, etc.), not the dimensionality of the observation (image). Therefore, $D$ is of much lower dimension than the dimensionality of the input, and only corresponds to underlying latent features that are relevant for the tasks. Also, while $N_{task} \gg D$ is the condition for orthogonality, the condition for linear decodability (abstractness), for which equations (4) and (5) hold, is $N_{task} \geq D$. In addition we find empirically that disentanglement (orthogonality) is also achieved for $N_{task} \geq D$, particularly in transformers. Therefore both our theory and experiments point to a linear relationship between the number of tasks required to disentangle factors, and the dimensionality of said factors $D$. We believe that this makes our setting applicable to real life settings, in particular because as noted in the discussion, there exists an interplay between $N_{task}$ and $D$: multi-task learning guarantees that neural networks will learn to disentangle the subspace within the $D$-dimensional latent space that is spanned by the decision boundaries. In other words, we are guaranteed to learn disentangled representations of all latent factors that are involved in the classification tasks, which is a very powerful guarantee. Furthermore, the tasks themselves do not need to be orthogonal; just to collectively span the latent subspace of interest; then the latents can become orthogonal. Finally in regards to optimality, we expect near-optimal performance to be achieved by any efficiently trained network (e.g. one trained with backpropagation), if adequate training samples are provided. Sub-optimal classifiers with zero-mean error are also handled in Corollary B.7 in the Appendix. Regarding dSprites, the latent factors for this dataset are $D=5$: Shape, Scale, Orientation, X and Y Position. Therefore, we would need at least 5 binary classification tasks that collectively depend on all of these variables (example task: $+1$ if we have a shape of scale$>0.7$ centered at $(x,y)$ where $x>y$, $0$ otherwise). ## Connection to neuroscience We thank the reviewer for allowing us to elaborate on this topic. Disentangled (or abstract) representations have been discovered in a multitude of brain areas, serving processes as disparate as memory (Boyle et al, 2022), emotion (Saez et al, 2015) and decision making (Bongioanni et al, 2021). Therefore, it seems natural to ask how such disentangled representations may come about in neural networks, including the brain. A main takeaway from this study is that parallel processing (multi-task learning) enforces the learning of such representations, with mathematical guarantees. Incidentally, it turns out that the cortex itself is built for parallel processing (Hawkins et al, 2019), with individual columns thought to be able to perform computation independently. Therefore, we believe that might be how the brain might build disentangled world models that allow efficient downstream out-of-distribution generalization (as we show). Of course this is still an assumption, and we hope this study motivates experimental neuroscientists to have a closer look at computation in cortical columns. Furthermore, other brain areas, such as the thalamus, could be involved in this parallel processing, as we mention in the manuscript. In regards to the second point, it is hard to quantify the level of noise in the brain. Furthermore, by noise in our setting we refer to noise in the latents (perceptual noise) not noise in the brain (neural noise), which in vision would correspond to viewing an object under different angles/lighting positions etc, which will allow the agent to “localize” the object in latent factor space. This idea of building faithful visual representations by being robust to perturbations like rotation etc. has been around in the literature for a while (e.g. DiCarlo, J. J., Zoccolan, D., & Rust, N. C. (2012). How does the brain solve visual object recognition? Neuron, 73(3), 415–434).

Response to reviewer (2/2)

## Determining the number of tasks for disentanglement The reviewer raises a great point. Indeed most of the time ANNs are not trained with multi-task learning, which might explain their poor out-of-distribution generalization capabilities. Note, however, that our framework neatly falls within the modern machine learning practices, and more specifically self-supervised learning (see l. 1908-16 in the manuscript). In short, self-supervised learning involves converting an unsupervised objective (“scene understanding”) into a well-posed supervised objective (image patch-filling). Filling missing patches *is* a multi-task objective, with a common, shared latent space (objects present in an image, and their relations), and solving each one of the objectives acts synergistically in understanding the image latent space as a whole. Therefore in this case, if a lack of out-of-distribution generalization is detected, our theory proposes a simple remedy: increase the number of masked out patches that have to be predicted from the image, which will result in a construction of a more descriptive latent space of the features underlying images. Of course, that prediction has to be tested, which would be an interesting topic for future research. Another dimension we would like to point out here is that sometimes we don’t necessarily want to uncover the entire latent space. As we note in the discussion, our work reveals an interplay between the number of disentangled latents and the number of tasks. We guarantee that multi-task learning will uncover the latents that are spanned by the tasks we train for, which in many cases is adequate for the tasks under consideration. In other words, there is no need to uncover all possible latents, but only the ones that are relevant for the tasks we are interested in. In that sense, it is the tasks that define the set of useful, disentangled latents, and not the other way around. Regarding to the reviewer’s questions: ## Higher $D$ While a lot of the main text figures focus on the $D=2$ case, a main reason we do this is that for $D=2$ it is easier to visualize the representations. However we also conduct experiments in high-dimensional latent spaces (up to $D=10$) in Figure 5b, experiments that confirm our theoretical findings, showing excellent out-of-distribution generalization (abstractness) for $N_{task} \geq D$ in transformers and RNNs. Extending these findings to more real-world datasets (like CIFAR or dSprites) would be an interesting direction. However these datasets lack the dynamical structure of the tasks we focus on here, where there is some background truth buried in perceptual noise that can be uncovered if multiple independent views of the object are afforded. Therefore, a promising direction would be to develop novel datasets like dynamic versions of dSprites that afford multiple views of the same object under different lighting conditions/angles. However, constructing and testing such datasets with CRNN or ViT architectures would be an independent study of its own, and far beyond the scope of this paper. ## Higher $N_{task}$ In general, learning more classification tasks with low error is not a problem. In particular, increasing the number of tasks while keeping $D$ constant does not negatively affect the performance of these tasks (we tried values of $N_{task}$ up to $192$). Similarly, if the number of training samples is kept constant, increasing the number of tasks does not negatively affect performance. Each sample provides information useful for each task.

Thank you

We thank the reviewer for the kind response. While we share the interest in potential causal connections to neuroscience, we emphasize that **this work focuses on representation learning** rather than computational neuroscience, and that we cannot provide causal evidence about brain function – this would require experimental validation through wet lab studies. **Our primary goal is to advance understanding of how multi-task learning enforces generalizable representations in neural networks, with theoretical guarantees, independently of whether such processes occur in the brain in this exact way or not.** We hope that our work provides motivation and a theoretical foundation that others (experimentalists or theoreticians) may use to test or explore connections to brain function in future research. Once again we would like to thank the reviewer for their thoughtful feedback, and for supporting our work.

Response to reviewer (1/2)

We would like to thank the reviewer for the thoughtful and useful review. ## Main comments: While most of the experimental results are on linear classification boundaries in latent factor space, we would like to point out that the classification boundaries are rendered non-linear by the encoder $f$, therefore the network is de facto still performing non-linear classification tasks. In addition, in the appendix (A.10) we extend our framework to non-linear ($1/x$) classification boundaries, showing that abstract representations are also learned in this setting. Therefore, this generalizes scenarios where agents have to multiply latent factors to make decisions, like for example in situations where one latent corresponds to reward magnitude and the other to reward probability (Bongioanni et al, 2021). Finally, while we would like to extend our framework to more naturalistic tasks, like dynamic versions of dSprites that afford multiple views of the same object under different lighting conditions/angles, which would correspond to noise in the latents as in our setting, we are not currently familiar with any such datasets, and constructing and testing such datasets would be an independent study of its own, beyond the scope of this paper. ## Impact of dimensionality of network and of the encoder We have tested different sizes of networks (dimensionality of $Z$), ranging from $32$ to $512$ neurons. The dimensionality of $Z$ should be greater than $D$, to ensure that the network has enough capacity to capture all the latents, as well as perform other computations (average evidence, invert non-linear mappings, etc.). As a standard machine learning practice, we choose $Z$ to be a few times greater than the largest $D$ we would like to capture. We found that $64$ neurons were enough for $D$ up to $10$. To examine when the latent encoding fails, we decrease the number of neurons when $D=10$, i.e. when we need the most amount of neurons to disentangle all latents. For $N_{task} = 24$ and $N_{neu}=64$, median $r^2 = 0.9$, therefore representations are abstract. However when decreasing $N_{neu}=32$, median $r^2 = 0.52$, and when $N_{neu}=16$, median $r^2 = 0.19$, i.e. generalization performance is completely hampered. Therefore we conclude that while there is no strictly defined threshold, $N_{neu}$ should be a multiple of $D$ to ensure enough capacity to perform all the functions that the network does (encode the factors, invert nonlinear encoder, denoise etc.). Similarly, the output dimensionality of the encoder (let’s call it $E$) should be greater than the largest $D$, to avoid throwing away information at the encoding stage. We have tried all the way down to $E=D$, e.g. by feeding the latents directly to the network. Realistically, $E$ will be much larger than $D$, as an observation (e.g. image) typically has much larger dimensionality than its latents. ## More nonlinear encoding In new experiments, we replace the ReLUs in the encoder with quadratic nonlinearities. We find no reduction in OOD generalization performance (see Appendix A.5 in updated manuscript). ## Transformer architecture Self-attention (or simply, attention) layers are a central piece of the transformer architecture in all situations, and unrelated to whether a transformer received independent temporal information or not. Did the reviewer mean positional encodings here? If we would include positional encodings in the form of learnable parameters, while the noise is sampled IID, we would expect to see no improvement in performance, but no deterioration in terms of disentangled representations either. The transformer would simply not use the positional encodings, since they are not relevant for the execution of the task (the inputs can be treated as a set). Since the transformer has to learn that positional encodings are irrelevant, this might slightly slow down learning, depending on how strong the initialization of the positional encodings is, but the effect will likely not be noticeable. ## Noise correlations We have already shown that our setting works in the case of autocorrelated noise which introduces temporal structure to the inputs in section A.9 in the appendix. If we wanted to extend these findings to transformers, performance would benefit from including positional encodings, since now position matters (the inputs are no longer IID, and adjacent samples have correlated noise).

Response to reviewer (2/2)

## Much greater values of $N_{task}$ and $D$ We have tested values of $N_{task}$ up to $192$, with no negative impact on performance. Testing for much larger values of $D$ poses a problem, because of the increased computational runtime and memory resources associated with increasing $D$. Specifically, in our OOD generalization code, we perform quadrant splits which scale exponentially with $D$, making testing generalization performance for large $D$ difficult. However, we did train networks for $D=30$, $N_{task}=256$, $E=150$ and $N_{neu} = 256$ up to ~ $90$ percent classification accuracy in new experiments, and we also trained and tested networks with $D=20$, $N_{task}=128$, $E=80$ and $N_{neu} = 128$ resulting in great generalization performance (median OOD $r^2 = 0.86$). ### About the more minor points: * Yes indeed, $x_{in}$ is the encoded input $f$ concatenated with the fixation input (see Fig. 2a). We have added this information in the paper * For the experiments with higher $D$, we resorted to randomly sampling classification boundaries. This naturally introduces some randomness in the performance of the networks, which is why RNN performance when $D=10$ can happen to be slightly higher than performance for $D=8$ (the two are not that different when $N_{task}$ is only 3), and also why variance in the fits for $D=8$ is high. In both cases, performance is far below what would be considered as good generalization performance, as expected by the theory.

Any other questions we can answer?

Dear Reviewer, Thank you for taking the time to review our paper and providing valuable feedback. We noticed that there hasn’t been further engagement with our rebuttals during the discussion period. As the discussion is nearing its conclusion, we wanted to kindly ask if there are any additional questions or concerns we could address to help clarify any remaining points. We greatly appreciate your efforts in this review process. Sincerely, The Authors

Core limitation: Reliance on a supervised paradigm

Since I was called out singularly in this segment of the authors’ global rebuttal, I feel it is necessary to respond directly. --- I appreciate the authors' efforts and their response, but I am genuinely struggling to understand the broader significance of this paper—and this line of work in general—where abstract representations are obtained through strong supervised signals. Could the authors clarify how their work contributes to advancing the field? The discussion of parallel processing seems peripheral to my original critique. To ensure clarity, let me restate my concern in concrete terms, highlighting why I find the study of disentangled representation learning impactful and where this paper appears to fall short: - **From a neuroscience perspective**: Abstract representations have been observed across various brain regions. It is an interesting question to explore how and why these representations emerge in evolved biological systems. - **From a machine learning and AI perspective**: Disentangled representations are hypothesized to improve efficiency and facilitate out-of-distribution generalization. Understanding how artificial systems can autonomously develop such representations is a valuable goal. The core issue with this paper is its reliance on supervised signals to achieve abstract representations. In my view, this undermines its relevance to understanding how biological systems develop such representations, as supervision is rare in real-world or naturalistic environments (outside highly controlled laboratory settings). Similarly, the supervised paradigm limits the practical utility of this work in developing artificial systems capable of learning disentangled representations directly from experience. In both neuroscience and AI, the scarcity of reliable supervision in real-world scenarios renders this approach less impactful. To illustrate this point, I'd like to reference a related study from ICLR last year ([Whittington et al., 2023](https://openreview.net/forum?id=9Z_GfhZnGH)), which the authors themselves cite. Whittington and colleagues showed that biologically plausible constraints, such as non-negativity of representations and a metabolic cost term, can lead to disentangled representations in a fully unsupervised framework. Their work is impactful because: - **In neuroscience**: It starts from biologically plausible assumptions, rather than starting with something contrived like multiple binary classification tasks, followed by forcing a post-hoc biological interpretation. - **In AI**: It demonstrates how these constraints lead to disentangled representations that outperform established benchmarks (e.g., beta-VAEs) on certain tasks while providing clear, actionable paths for future work to contribute to improving disentanglement in a way that is biologically grounded. Crucially, Whittington et al. operate in an unsupervised setting, addressing the challenges of both biological plausibility and real-world AI applications. My goal here is not to directly compare the two papers. Instead, I hope to convey, clearly and concretely, why I am struggling to see the significance of the current work. The key questions remain: - How does this paper help us understand the emergence of abstract representations in biological systems, which rarely receive supervision? - How do these results inform the design of models that develop abstract representations in real-world, unsupervised settings? If these questions are outside the scope of the authors' intended contribution, I would appreciate clarification on what specific insights this work provides to the broader community. Is the contribution primarily theoretical, with the aim of laying groundwork for future extensions toward more realistic and impactful scenarios? If so, how does this work pave the way for that future research? My struggle, again, is understanding how providing models with extensive supervised signals to achieve abstract representations could positively influence future research in disentangled representation learning. Given the significant gap between the paper's assumption of abundant supervision, versus real-world scenarios where supervision is rare, I am not yet convinced of its relevance. I would genuinely appreciate it if the authors could articulate how their work is meaningful despite this fundamental limitation, while staying focused on the core critique.

Clarifying a key reviewer misunderstanding on competence at tasks vs. supervision

We thank reviewer Mvb9 for their overview of Whittington et al., 2023, which is excellent work. To briefly address the key concerns related to the usage of supervised learning and its impact on the broader significance of our work: 1. Conceptually and theoretically, our proof on disentanglement in potentially suboptimal multi-task classifiers clarifies the mechanism by which **competence** on $N_{task}$ classifications implicitly represents latent variable estimator $\mu(t)$. **Competence in $N_{task}$ tasks does mean the ability to accurately assign likelihoods to $N_{task}$ classifications necessarily arose from supervised training signals**. A direct counterexample is pre-trained large language models that exhibit low perplexity/accurate likelihood assignment on novel few-shot tasks. We previously mentioned this point regarding **competence vs. supervision** in reviewer Mvb9’s previous round of individual rebuttals (see section *Conflation of “competence at N tasks” with “receiving feedback signals" on N tasks*), but we will further elaborate below. 2. Unsupervised learning objectives (as in Whittington 2023) are not mutually exclusive with applying the theoretical and experimental results of our work. A direct comparison of our work with Whittington et al, which has a different focus, approach and architectures might not be particularly useful here. ## Addressing reviewer Mvb9’s misunderstanding of a key contribution of our work To address reviewer Mvb9’s contention regarding supervision, we briefly outline a main contribution of our work here, to make it extremely easy to parse. Of course it’s impossible to capture the full nuance of our work in a few lines. * **Theoretical contribution**: we prove that competence in $N_{task}\geq D$ enforces abstract/disentangled representations. * **Experimental confirmation**: we confirm the key predictions of our theory with regards to $N_{task}$, effect of noise, etc in autoregressive models that achieve near-optimal (fig. 5a) competence at the tasks * **How we achieve competence at tasks**: multi-task learning with supervised labels. We choose this because it is the most straightforward way to become competent at multiple tasks, however our theory is not limited to this supervised framework. **Our theory establishes a formal link between competence at $N_{task}\geq D$ tasks and disentanglement, irrespectively of how that competence was achieved.** **The focus on the supervised nature of our tasks hence misses a key point of our work, namely that the supervised paradigm is simply our way of guaranteeing the theory requirement of competence at multiple tasks, and as such it is secondary to the key message above.** Scaling up supervised disentanglement efforts is not our suggestion for future research directions. ## Future directions Since our theory is agnostic to the way by which competence at multiple tasks is achieved, our work raises exciting applications in several areas of modern machine learning. As already mentioned in the manuscript (l. 1962-70), and the previous rebuttal to reviewer Mvb9 (section *Regarding the feasibility of disentanglement*), the most straightforward application is self-supervised learning, that relies on the prediction of multiple objectives which share a common latent space, and as such shares commonalities with our experimental approach here. Some actionable next steps spun by this work that we are actively pursuing are: * Extending our experiments beyond the supervised tasks considered here, to real-world datasets to see **how well the main conclusions of our theory hold in self-supervised settings** * **Investigating representations in pretrained language models, which are competent at many tasks**, and hence are likely to display disentanglement according to our theory * **Explaining emergence of disentangled reps in biological organisms naturally competent at many tasks**, and investigating more biologically plausible ways to induce it (e.g. continual learning) Of course we cannot expand too much on these as they are ongoing research. We are sure this work will spark more research directions, as it gets shared with the research community. **Our work also motivates the development of metrics that detect disentanglement without relying on the axis-alignment hypothesis, as now we have good reason to believe that capable models contain disentangled representations, just not necessarily at the individual unit level.** The same shift happened in neuroscience, where people switched from looking for single neuron coding of variables, to population coding. We expand more on reviewer Mvb9’s response at their [individual rebuttal below](https://openreview.net/forum?id=yVGGtsOgc7&noteId=QjrzaVKFN0) to avoid repeating things, including what our paper is about and what it isn’t, which we believe is not accurately reflected in the reviewer’s framing and questions.