Disentangling Representations through Multi-task Learning

Authors: Pantelis Vafidis, Aman Bhargava, Antonio Rangel
Primary area: interpretability and explainable AI
Keywords: zero-shot generalization, disentanglement, interpretability, world models, multi-task learning, computational neuroscience, neuroAI, evidence accumulation, cognitive maps, continuous attractors, RNNs, transformers

Abstract:
Intelligent perception and interaction with the world hinges on internal representations that capture its underlying structure ("disentangled" or "abstract" representations). Disentangled representations serve as world models, isolating latent factors of variation in the world along approximately orthogonal directions, thus facilitating feature-based generalization. We provide experimental and theoretical results guaranteeing the emergence of disentangled representations in agents that optimally solve multi-task evidence accumulation classification tasks, canonical in the neuroscience literature. The key conceptual finding is that, by producing accurate multi-task classification estimates, a system implicitly represents a set of coordinates specifying a disentangled representation of the underlying latent state of the data it receives. The theory provides conditions for the emergence of these representations in terms of noise, number of tasks, and evidence accumulation time, when the classification boundaries are affine in the latent space. Surprisingly, the theory also produces closed-form expressions for extracting the disentangled representation from the model's latent state $\mathbf Z(t)$. We experimentally validate these predictions in RNNs trained on multi-task classification, which learn disentangled representations in the form of continuous attractors, leading to zero-shot out-of-distribution (OOD) generalization in predicting latent factors. We demonstrate the robustness of our framework across autoregressive architectures, decision boundary geometries and in tasks requiring classification confidence estimation. We find that transformers are particularly suited for disentangling representations, which might explain their unique world understanding abilities. Overall, our framework establishes a formal link between competence at multiple tasks and the formation of disentangled, interpretable world models in both biological and artificial systems, and helps explain why ANNs often arrive at human-interpretable concepts, and how they both may acquire exceptional zero-shot generalization capabilities.

Decision: accept (poster)
Reviewer scores: 6, 6, 3, 6, 8

Reviews and discussion:
<thread>
creation date:
1729359513562

thread_type:
official_review

<thread_object>
creator: reviewer_b5qb

<content>
summary:
This paper demonstrates the emergence of disentangled representations in recurrent neural networks and transformer architectures in classification tasks based on evidence aggregation or accumulation over time. The conditions under which disentangled (latent) representations emerge are studied in relation to the number of classification criteria, relative to the dimensionality of the input. The paper motivates its focus with reference to the neuroscience literature on evidence accumulation and offers a number of theorems predicated on mutual information and the functional forms of nonlinearities in various mappings, implicit in the network architectures that could be adopted.

strengths:
The strengths of this paper rest upon a detailed treatment of the functional forms of mappings — implicit in neural network architectures — and how they can be unpacked in relation to maximising the mutual information between inputs and (classification) outputs. Another nice feature of this paper is its reference to neuroscience; in the sense of potential implications for neuronal processing in the brain.

weaknesses:
The specific contributions of the analysis are not foregrounded sufficiently. For example, one could read the setup considered by the authors as characterising a mapping between inputs x and classification estimates (Y^) using a bottleneck architecture (i.e., mapping through a low dimensional latent space or manifold). If any such mapping were optimised with respect to cross entropy or mutual information, would the latent representations not be disentangled? In other words, is it a linear (orthogonal) disentanglement that qualifies as interesting disentanglement? And, if so, does this inherit from the assumptions and functional form of the mappings in question?

The second weakness is the rather colloquial appeal to the neuroscience literature. For example, evidence aggregation or accumulation would not be considered canonical in the cognitive neurosciences. Canonical paradigms will be things like oddball paradigms, psychophysical paradigms, working memory paradigms, attentional paradigms, and so on. Furthermore, most of neuroscience nowadays cast disentanglement in terms of predictive processing under hierarchical or deep generative world models (i.e., disentanglement is a way of describing the inversion of a generative model of how independent causes become entangled). This inversion is usually articulated in terms of predictive coding or belief propagation. Given that the authors mention Bayesian filtering — and predictive coding can be read as extended Kalman (Bayesian) filtering — there might have been a missed opportunity to join the dots between the authors work and current formulations of perceptual synthesis and evidence accumulation in the brain. 

Related to this, there is a long history of sparse coding models in neuroscience that speak to the current questions; for example, sparse coding models, variants of independent component analysis, liquid computation, echo state machines, and so on (Gros, 2009; Hu et al., 2020; Maass et al., 2002; Olshausen and Field, 1996; Simoncelli and Olshausen, 2001; Suh et al., 2016). More recently, people have been looking at latent representations at various hierarchical levels to address compositionality and disentanglement; especially in relation to spatiotemporal receptive fields that can be assessed empirically in the brain e.g., (Ficco et al., 2021; George et al., 2021; Rao et al., 2023).

Ficco, L., et al., 2021. Disentangling predictive processing in the brain: a meta-analytic study in favour of a predictive network. Scientific Reports. 11, 16258.
George, D., et al., 2021. Clone-structured graph representations enable flexible learning and vicarious evaluation of cognitive maps. Nature Communications. 12, 2392.
Gros, C., 2009. Cognitive Computation with Autonomously Active Neural Networks: An Emerging Field. Cognitive Computation. 1, 77-90.
Hu, H.-y., et al., 2020. RG-Flow: a hierarchical and explainable flow model based on renormalization group and sparse prior. Machine Learning: Science and Technology. 3.
Linsker, R., 1990. Perceptual Neural Organization - Some Approaches Based on Network Models and Information-Theory. Annual Review of Neuroscience. 13, 257-281.
Maass, W., Natschlager, T., Markram, H., 2002. Real-time computing without stable states: a new framework for neural computation based on perturbations. Neural Comput. 14, 2531-60.
Olshausen, B.A., Field, D.J., 1996. Emergence of simple-cell receptive field properties by learning a sparse code for natural images. Nature. 381, 607-9.
Rao, R.P., 1999. An optimal estimation approach to visual perception and learning. Vision Res. 39, 1963-89.
Rao, R.P.N., Gklezakos, D.C., Sathish, V., 2023. Active Predictive Coding: A Unifying Neural Model for Active Perception, Compositional Learning, and Hierarchical Planning. Neural Computation. 36, 1-32.
Simoncelli, E.P., Olshausen, B.A., 2001. Natural image statistics and neural representation. Annu Rev Neurosci. 24, 1193-216.
Suh, S., et al., 2016. Echo-State Conditional Variational Autoencoder for Anomaly Detection. In: 2016 International Joint Conference on Neural Networks. IEEE International Joint Conference on Neural Networks (IJCNN), Vol., ed.^eds., pp. 1015-1022.

questions:
For the neuroscience reader, it would be good to clarify what multi-task means in this setting. Does it mean that there are many different classification criteria or decision boundaries that have to be met? In other words, is the multi-task aspect a way of saying that the number of classification criteria have to exceed the dimensionality of the underlying (ground truth) evidence? Perhaps it would be useful to have a simple example that people can imagine applying to a perceptual classification task in a psychophysics laboratory?

By evidence aggregation, did you mean evidence accumulation in the spirit of drift diffusion and race to bound models?

Is there any relationship between the functional form of your formulation and schemes based upon the principle of maximum mutual information (a.k.a. infomax principle) (Linsker, 1990). Canonical schemes here would include canonical variates analysis and non-linear versions, such as independent component analysis that take us into the realm of sparse coding models. One interesting aspect of independent component analysis is the identification of latent variables that have a non-Gaussian structure under the assumption that the random fluctuations are Gaussian. Is this related to your interesting observation that noise is required for the kind of disentanglement you are considering?

Can you formulate your evidence accumulation scheme as a Bayesian (e.g., extended Kalman Bucy) filter? If so, can you relate this to predictive coding in the brain (Rao, 1999)?

On a related note, can you formulate your description in terms of inverting a generative model? In other words, can you turn your scheme on its head and generate ground truth evidence from the ground truth classifications; i.e. causes. If so, is this one perspective on the requisite invertibility of f?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730574914464

thread_type:
official_review

<thread_object>
creator: reviewer_xvfc

<content>
summary:
This paper studies the problem of learning disentangled representations using a multi-task evidence aggregation approach. In contrast to previous works which have studied a similar problem of abstract and disentangled representations using contextual information, the authors here adopt an approach which simultaneously collects multiple streams of evidence to enable linear classification in multiple tasks at each step. They prove the emergence of an abstract representation of ground truth data given an optimal multi-task classifier, which is shown to be learnable using RNN, LSTM, and GPT-2 based transformer architectures.

strengths:
The key contribution is an extension of the Johnston and Fusi (2023) feedforward framework to recurrent networks. They show that abstract variables are indeed represented in the latent states of RNNs, LSTMs and GPT-2 style transformer architectures when these models are trained using the multi-task paradigm. The results are interesting and point to a marked advantage in adopting the multi-task paradigm for understanding canonical neuroscience experiments in an AI framework. The discussion of these results in the context of contemporary neuroscience literature is also well framed and places this paper in a good position to appeal to both the broader neuroscience and task-learning communities.

weaknesses:
While the experimental results are solid contributions, the theoretical contributions are less appealing. The key theoretical results follow from fairly straightforward arguments, but appear a bit overstated. This can potentially be remedied with some re-wording and clarifying statements.  

In my opinion, the following two sentences in the abstract are overstated for reasons that I expand further below:  "The key conceptual finding is that, by producing accurate multi-task classification estimates, a system implicitly represents a set of coordinates specifying a disentangled representation of the underlying latent state of the data it receives.....Overall, our framework puts forth parallel processing as a general principle for the formation of cognitive maps that capture the structure of the world in both biological and artificial systems, and helps explain why ANNs often arrive at human-interpretable concepts, and how they both may acquire exceptional zero-shot generalization capabilities". 

The conditions under which the authors prove the theoretical result are rather strong. Specifically, the authors require 1) that the network is an *optimal* classifier in the sense that the network correctly computes the N_task logits associated with each of the N_task tasks, 2) that decision boundaries in the world are linear in the space of abstract variables, and 3) a highly simplified evidence accumulation problem. 

It follows from 1) and 2) that the output of an optimal network encodes the distance to the N_task decision boundaries. The optimal network will thus represent a linearly transformed version of the input if N_task is larger than the dimensionality of the input (provided the decision boundaries are not degenerate). The mathematical proofs presented in this paper formalize this argument. I believe it will help clarify the results if the authors provide this intuition upfront. The proofs have various conceptual issues, for example, in the use of the data processing inequality in Theorem B.5 and the mix-up of random variables and their estimators (lines 1442-1445). 

The authors state that their framework offers a "general principle for the formation of cognitive maps". I believe it will be helpful if the authors can clarify what this general principle is -- is it the proposal that decision boundaries in the world are linear in the space of the abstract variables or that biological networks are solving multiple tasks at the same time? How is the current work related to the formation of cognitive maps and zero-shot generalization?

questions:
Comments:

Line 231: “…in the column space of (C^TC)^{-1}C^T potential with some element-wise non-linearity.” This statement seems too strong given that the restriction on the activation function g has been significantly relaxed. Are there any preliminary results to qualify this statement?

Corollary B9: While I agree with the results of this corollary, it would be good to explicitly connect the uniformity of the singular values to the corollary statement using the Marchenko-Pastur law. Also in this corollary, at line 1613, it is stated that the probability of being orthogonal vanishes as the dimensionality N_task increases. Was this supposed to be “non-orthogonal” instead?

Figure 2: This figure is particularly unclear, and it is perhaps one of the important foci in explaining all the experiments to follow. Upon first viewing with the caption, it is not clear what the arrows (two grey and one red) in Fig. 2a are supposed to imply.

Line 294: “…in the RNN’s hidden layer, when tasks span the latent space.” Does this really imply anything about the tasks formally spanning the latent space?

Figure 4a: I agree that the RNN representations are sparse, and this looks like a good spot to test if there is any lower threshold on the size of the hidden units before the latent encoding fails.

Line 450: “…which might explain their superior generalization performance to RNNs for lower N_task.” It is not really clear from Fig. S10 that the RNNs are indeed worse at orthogonalization than the GPT.

Line 1119/1120: “Figure S3b shows that the network still learns a disentangled, two dimensional continuous attractor.” Not convinced that this actually shows disentanglement since the metric for that qualification from before was the orthogonalization of different tasks. Did you mean to say abstract here? Also, small typo in the spelling of “disentangled”.

Minor comments:

Line 1321: \hat{Y}_i(t) is noted as a Bernoulli random variable. This implies the support of the estimator is {0,1}, but it should really be [0,1].

Line 1408: Similar comment here to Line 1321 on the use of Bernoulli variable for \hat{Y}(t).

Line 1426: In Equation 16, boldfaced t should be scalar.

Line 1592: “Proof: Recall equation (3)…” should be equation (4).

Line 1608: “B^TB is a symmetric for any matrix B.” Some LaTeX formatting errors here and possibly missing the word matrix following “symmetric”.

Line 309: “disentaglement” misspelled

Fig. 3d Caption: May be good to include a short statement about the inset on PCA components accounting for variance.

Line 1070: “Once the factors are perfectly, performance…” missing the word correlated after “perfectly”

Line 1162: “…for all lines in Figure S4a.” Did you mean all lines from Fig. 2a, shown in Fig. S4a?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730588400474

thread_type:
official_review

<thread_object>
creator: reviewer_mvb9

<content>
summary:
This paper extends the work of [Johnston and Fusi (2023)](https://www.nature.com/articles/s41467-023-36583-0)  by investigating how multiple supervised binary classification tasks can lead to abstract or disentangled representations in recurrent neural networks that accumulate evidence from noisy inputs. The authors present theoretical and experimental evidence suggesting that optimal multi-task classifiers can learn abstract representations of underlying ground truth factors.

strengths:
The paper introduces a theoretical framework connecting optimal multi-task classification to disentangled representations, building upon the empirical observations in [Johnston and Fusi (2023)](https://www.nature.com/articles/s41467-023-36583-0) . To support these theoretical claims, the paper provides a range of experiments exploring different architectures, task structures, and decision boundary geometries.

weaknesses:
A significant weakness of this paper is its heavy reliance on dense supervised signals to achieve abstract representations. Similar to [Johnston and Fusi (2023)](https://www.nature.com/articles/s41467-023-36583-0), this work utilizes multiple binary classification tasks, framed as "multi-task" learning. This framing implies a diversity of tasks that is not truly reflected in this setup. The dependence on numerous, highly similar, supervised tasks raises questions about the biological plausibility and real-world applicability of the proposed mechanism.

The requirement of a large number of tasks (N_task >> D) to achieve disentanglement adds to concerns regarding the model's biological plausibility and practical relevance. Biological and artificial agents operating in real-world environments rarely encounter such an abundance of supervised signals.

The authors assert (lines 486-493) that their Theorem 3.1 implies "the key factor driving convergence [to a Platonic representation of reality] is the diversity and comprehensiveness of the tasks being learned." This claim is not supported by the results presented. The tasks in this paper are essentially multiple variations of the same supervised classification task, a far cry from the genuinely diverse tasks (e.g., multimodal vision and language tasks) considered in the Platonic representation paper ([Huh et al., 2024](https://arxiv.org/abs/2405.07987)).

It is well-established that achieving disentangled representations in real-world, unsupervised settings is incredibly difficult, if not impossible, without incorporating appropriate inductive biases ([Locatello et al., 2019](https://arxiv.org/abs/1811.12359); which the authors cite). A substantial body of research, spanning decades, has demonstrated that source signals become non-identifiable when mixed non-linearly ([Hyvarinen and Pajunen, 1999](https://www.sciencedirect.com/science/article/pii/S0893608098001403)). Considering this extensive history of work in machine learning dedicated to tackling this complex problem, this paper's approach—attaining abstract or disentangled representations by providing the model with an abundance of supervised signals in the form of binary classification tasks—appears contrived and ultimately of limited practical value. The reliance on such a heavily supervised setting makes it difficult to see how these findings could generalize to more realistic scenarios where supervision is scarce and the underlying structure of the data must be inferred, rather than provided.

Overall, this paper makes incremental progress on the work of [Johnston and Fusi (2023)](https://www.nature.com/articles/s41467-023-36583-0)  but fails to address the fundamental concern of over-reliance on dense supervised signals. As suggested by [Johnston and Fusi (2023)](https://www.nature.com/articles/s41467-023-36583-0), exploring unsupervised mechanisms such as "predicting the sensory consequences of our actions," might offer a more promising path forward.

Instead of relying on a multitude of supervised signals to enforce disentanglement, 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. This shift in focus would significantly enhance the biological plausibility, practical relevance, and overall impact of this line of research to understanding disentangled representation learning.

questions:
**Sparsity:**

The emergence of sparsity in Fig. 4 is interesting. Could the authors specify if they applied any particular architectural constraints, such as regularization techniques or sparsity-promoting mechanisms, to encourage this behavior?

Additionally, quantifying the sparsity (e.g., lifetime sparsity, [Vinje and Gallant, 2000](https://www.science.org/doi/10.1126/science.287.5456.1273)) and examining how this metric varies with N_task, the number of latent dimensions, and specific RNN architecture choices would provide further insights into the effects of task density on network representation.

**More nonlinear encoding:**

The authors' use of a piecewise linear MLP+ReLU encoding of ground truth signals raises questions about the robustness of their empirical results under more complex nonlinear encodings.

Could the authors explore how the N_task vs. out-of-distribution r2 relationship (Fig. 3) might change if the encoding network incorporated more nonlinear functions such as exponential or power-law activations, to simulate more biologically plausible ([Priebe et al., 2004](https://www.nature.com/articles/nn1310)) nonlinear encoding? This line of inquiry is particularly relevant given that the choice of activation function significantly influences the geometry of representations ([Alleman et al., 2024](https://openreview.net/forum?id=k9t8dQ30kU)).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730604549949

thread_type:
official_review

<thread_object>
creator: reviewer_w8z5

<content>
summary:
The paper presents theoretical and empirical results showing that, in multi-task evidence aggregation classification tasks, representations become disentangled as the number of tasks N_task greatly exceeds the input dimensionality 𝐷 when the agent solves the tasks optimally.

strengths:
1. Originality: the study of the emergence of disentangled representation in temporal tasks and models is relatively new. 
2. Quality: the empirical validation is comprehensive. 
3. Clarity: the theory is well explained. The background section is also very nicely written and thorough. 
4. Significance: The paper addresses the key problem of learning the world model in representation learning and also discusses in detail the biological relevance. The results offer meaningful insights for future research on disentangled representation in both artificial systems and the brain.

weaknesses:
1. The theory assumes that agents are optimal multi-task classifiers, which may not be achievable in realistic settings where the input dimension D is already large and the number of tasks N_task >>D. This raises questions about the practical relevance of the regime considered in the paper. Additionally, it’s difficult to imagine a large number of orthogonal tasks in real-world settings. For instance, in the dSprite dataset, as the authors noted, how could a meaningful, larger set of tasks be constructed, and what implications would this have for the applicability of the framework?
2. The connection to neuroscience also appears somewhat speculative. The authors could clarify how their findings might directly relate to brain function. For example, can the theory predict the relationship between the disentanglement of neural representations and the level of noise or variability in different brain regions?
3. While explaining zero-shot generalization capabilities in artificial neural networks (ANNs) is a central claim, this remains a substantial challenge for most networks. The theory should address whether prior models fail in zero-shot generalization due to insufficient N_task and how practitioners could estimate an adequate N_task, especially in settings where underlying latent factors are unknown, making it impractical to sample random classification tasks as in the paper.

questions:
1. Most examples and empirical validations focus on low-dimensional cases (e.g., D=2). Could you provide examples using moderately higher-dimensional datasets, such as CIFAR or dSprites, to illustrate how N_task scales with D in these settings? 
2. For a given D, what is the maximum number of tasks N_task can be learned with low error? Additionally, how do the test and generalization errors scale with N_task if we keep the number of training data in at a realistic scale?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730694176009

thread_type:
official_review

<thread_object>
creator: reviewer_2aef

<content>
summary:
This paper explores the important question of how abstract (linear and approximately orthogonal) and disentangled (orthogonal without the necessity of linearity) representations can emerge in biological and artificial agents. 
The authors present both theoretical and experimental results demonstrating that multi-task learning, specifically within the framework of evidence aggregation classification tasks, can lead to the development of such representations.

strengths:
The paper presents theoretical results that establish specific conditions—relating to the number of tasks, input dimensionality, input noise, and more—that lead to the emergence of abstract and disentangled representations in agents solving multi-task evidence aggregation classification tasks. The authors conduct thorough experiments across several architectures (RNNs, LSTMs, and transformers) to validate their theoretical results, showing that even architectures like GPT-2 can exhibit these properties. The work bridges AI and neuroscience, with potential explanations for how abstract, human-aligned representations might arise in artificial neural networks, making it valuable for both fields.

weaknesses:
One notable limitation of this work is the assumption of factorization, as acknowledged by the authors. Additionally, the theoretical framework is tailored to a specific type of multi-task learning problem—evidence aggregation classification with linear decision boundaries—which may not capture the full diversity of tasks and decision boundaries that agents encounter in dynamic environments. It would be valuable to explore how these ideas generalize to other multi-task learning scenarios. Furthermore, the experiments focus on synthetic data with simple latent structures. While these experiments effectively support the theoretical results, testing the framework on slightly more complex data could provide further insights.

questions:
- In Equation 6, could you clarify if $x_{in}$ is the encoded input, i.e., $f(x)$?

- Does the dimensionality of the latent state Z influence the results? Also, does the dimensionality of the encoded input play a significant role? Have these aspects been examined?

- Have you conducted experiments with other encoding functions, such as a quadratic mapping? 

- In Figure 5b, the blue curve for GPT appears to increase at the end. Could you clarify the reason behind this behavior?

- Since the noise at each time step is independent, temporal information isn’t modeled in this setup. How might the results differ if a transformer with self-attention layers was used? Would such a setup still yield disentangled representations?

- On the other hand, if the inputs contained some temporal structure, how might this affect the nature of the representations?

- Have you tested the framework with much larger values of $N_{task}$ and $D$ (e.g., over 100)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732265777862

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Main response (2/2)

comment:
## Biological plausibility of multi-task learning
**Reviewer Mvb9** mentioned that our framework may lack biological relevance because of its “reliance on supervisory signals”. **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). In other words, **our theory assumes competence at $N_{task}$ tasks, independently of when and how that competence was achieved**. 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/pii/S0960982224002987)), and biological agents are remarkable *because* they achieve high performance on many tasks. This theory is also related to the widely observed phenomenon of memory replay ([Foster & Wilson, 2006](https://pubmed.ncbi.nlm.nih.gov/16474382/)), or mental simulation of counterfactuals ([Jensen et al, 2024](https://www.nature.com/articles/s41593-024-01675-7)).

At the same time, we recognize that parallel processing is still an assumption, and the study already presents fundamental insights in the nature of generalization and world model construction in neural networks as is, independently from whether this exact process materializes in biological neural networks or not. Future work will explore the relationship between multi-task learning and interleaved training (see Appendix A.10 for preliminary results), which more closely mimics neuroscience experiments where new tasks may be introduced to the animal.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732265838961

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Main response (1/2)

comment:
We thank the reviewers for the time and effort they put in reading the manuscript. The vast majority of reviewers recognized the contribution of the work, and their comments and suggestions have proven very useful in improving its presentation.

Here we summarize additions made in the revised manuscript, and clarify some aspects of our paper that were confusing for more than one reviewer.

## Summary of changes

In the updated version of our manuscript we:

* Rigorously quantify sparsity of our trained networks, showing that RNNs are consistently sparse across hyperparameter choices ($N_{task}$, $D$), similar to the brain (see **new Appendix A.12**). We also show that the number of active neurons is proportional to $D$, as long as there are enough tasks to uncover all latents ($N_{task} \geq D$) (**reviewers Mvb9** and **b5qB**).

* Demonstrate that our framework is robust to the choice of encoder nonlinearity, even handling non-injective observation maps (**new Appendix A.5**) (**reviewers 2Aef** and **Mvb9**)
* Expand upon the biological plausibility of multi-task learning (**Appendix A.14**) (**reviewer Mvb9**)
* Include a reproducibility statement
* Specify and clarify a lot of the pain points identified by the reviewers (e.g. description of Fig. 2).
* Summarize robustness of the theory under different noise distributions, readout functions and suboptimal classifiers in the new paragraph “Robustness” in section 3 of the main text
* Include another paragraph on extending the theory to non-linear boundaries in latent space in the theory section (section 3) (**reviewer xvfc**)
* Include a corollary in Appendix B (**Corollary B.7**) discussing extensions of the theory to sub-optimal multi-task classifiers (**reviewers b5qB** and **xvfc**)
* Move limitations to Appendix A.15 to make space in the main text.
* Increase the accuracy of our treatment of mutual information of the point estimator $\mu(t)$ and latent state $Z(t)$ with $x^*$ (l. 1550-3, 1617-9) (**reviewer  xvfc**)

All additions are marked with green font in the manuscript, for reviewer convenience. We also provide more experiments to address specific reviewer comments in the individual rebuttals. **Overall, we feel that we have addressed most reviewer comments, making this a compelling contribution for how autoregressive architectures can disentangle their representations, with provable guarantees.**

## Theoretical results and number of tasks

Our theory provides separate conditions for disentanglement and abstractness, which created some confusion (e.g. **reviewers w8Z5** and **Mvb9**). We show that linear identifiability is guaranteed when $N_{task} \geq D$ and orthogonality when $N_{task} \gg D$. Equations 3 and 4 hold under linear identifiability, i.e. when $N_{task} \geq D$ (condition 3 in Theorem 3.1 doesn’t have to hold). Empirically, we find that our networks, in particular transformers, readily disentangle when $N_{task} \geq D$, going beyond our theoretical guarantees. This might have caused some confusion to **reviewer Mvb9**, who mentions that an abundance of supervised signals is required, which is clearly not the case. Our networks require the minimum number of tasks that theoretically guarantee linear identifiability.


## Interplay between $N_{task}$ and $D$

Extending the above, having less tasks than latents is not a problem either. Our theory guarantees that the networks would disentangle the subspace in latent space that is spanned by the classification objectives. Therefore, as we note in the discussion, our work reveals an interplay between  $N_{task}$ and $D$; in a rich world where a lot of latents may exist, we are guaranteed to uncover the latents that are relevant for the tasks we are concerned with, as long as the subspace of these latents is spanned. If we would like to uncover more latents (i.e. understand the world in more detail), we would have to incorporate more tasks.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733202252630

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Main contributions & implications

comment:
We would like to thank reviewers for their feedback, and for finding our work interesting. For reviewer convenience, we summarize here some of the key contributions of our work, their implications, and potential future steps. This section acts synergistically with the more technical *Contributions* section in the main text.
## Main contribution
A main contribution of our work is that it **establishes a formal link between competence at many tasks, and the emergence of disentangled representations.** We first establish this argument theoretically, and then confirm and extend it experimentally. While our experiments are on supervised multi-task learning, **our theory is agnostic to the way by which competence was achieved**, and leaves the door open for other possibilities (e.g. self-supervised or unsupervised pre-training). 
## Role of experiments
**Our supervised learning experiments are the most efficient way to confirm and extend our theoretical findings, and allow for greater tractability of results**, compared to unsupervised approaches. Note that scaling up supervised disentanglement efforts is **not** our suggestion for future research, due to scalability concerns. Instead, future work will endeavor to extend our results to **self-supervised learning**, which shares a lot of commonalities with our framework (i.e. making several predictions in parallel that all rely on a common latent ground truth; also see l. 1962-70 in the manuscript), and test how well our theoretical findings hold up in this setting.
## New research directions
Since our theory is agnostic to the way by which competence at multiple tasks is achieved, **a natural next step is to investigate whether disentangled representations exist in a wider range of models capable of solving multiple tasks**. A prime example is large language models that display excellent zero- and few-shot generalization capabilities, with some progress already made in that space. Moreover, the pre-training objective for LLMs (cross entropy loss/likelihood maximization) fits well within our theoretical framing on (approximately) optimal multi-classifiers.

**Our experiments showed parsimony of our theoretical results under conditions not covered by our theory**, including non-injective observation maps (Appendix A.5) and decision boundaries (Appendix A.10) which is encouraging for testing our findings on settings beyond what is strictly covered by our theory. Another direction we are actively pursuing is extending the theory to handle these more general cases (see “Discussion” on pages 36-38 of the manuscript).
## Need for better metrics
Our work also motivates new metrics that can **better detect** disentangled representations in existing networks. We argue in favor of a mixed representation approach to disentanglement (fig. 4), while most metrics to date look for axis-aligned representations. **Our work suggests that we might be missing disentanglement in networks that perform well in many tasks, due to the metrics used.**
## Impact for neuroscience
Our findings are also relevant for neuroscience, since the evidence integration tasks we use are canonical decision-making neuroscience tasks. Naturally, animals are capable of solving multiple tasks, thus our work provides theoretical justification for why disentangled representations have been found in many brain areas, and motivates looking for more. Note that we do not make strong specific claims about how competence at multiple tasks is achieved by organisms here.

**One hypothesis put forth by this work is that parallel computation may be happening**, where an animal may be executing only one task at a time but its predictions for previously encountered, similar tasks are also computed. Expanding upon this idea and other alternatives ways to achieve competence at multiple tasks in biological organisms (e.g. continual learning) will be a topic of future research. The potential benefits of this line of research are clear: **multi-task learning is extremely more efficient than previously proposed mechanisms for representation learning in the brain** (Mante et al, 2013), which is particularly relevant for organisms that support diverse computations within a limited neural substrate.
## Confidence and world models
Finally, our work emphasizes the importance of the confidence (i.e., calibrated likelihoods) of a network’s output, which is a recurring theme in machine learning (e.g. knowledge distillation etc.). **We here show that confidence fundamentally connects to how neural networks construct world models, either directly** (integrate-to-bound task in Appendix A.11) **or indirectly** (classification tasks). Under this framework, knowledge distillation can be cast as smaller models directly copying the world models (logits) of larger ones.

For more general implications of our work, we would like to point the reviewers to the discussion in the main text. Overall, we are very excited for the future of the field.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper explores the question of how abstract and disentangled representations can emerge in biological and artificial agents. The authors present both theoretical and experimental results demonstrating that multi-task learning, specifically within the framework of evidence aggregation classification tasks, can lead to the development of such representations.

Strengths: The paper introduces a theoretical framework connecting optimal multi-task classification to disentangled representations, building upon the empirical observations in Johnston and Fusi (2023). The paper provides a range of experiments exploring different architectures, task structures, and decision boundary geometries to validate these claims. The discussion of these results in the context of contemporary neuroscience literature is also well framed and places this paper in a good position to appeal to both the broader neuroscience and task-learning communities.
Weaknesses: The conditions for the theoretical results are rather strong, which may limit their practical relevance.  The experiments focus on synthetic data with simple latent structures. While these experiments effectively support the theoretical results, testing the framework on more complex data could provide further insights.

Overall, while the theory and experiments have limitations,  the paper still contributes meaningfully  to the understanding of disentangled representations and how they might arise. I think it will be valuable for the community and recommend accepting it.