Disentangling Representations through Multi-task Learning
Authors: Pantelis Vafidis
Aman Bhargava
Antonio Rangel
Venue: ICLR 2025 Poster (2025)
Keywords: zero-shot generalization
disentanglement
interpretability
world models
multi-task learning
computational neuroscience
neuroAI
evidence accumulation
cognitive maps
continuous attractors
RNNs
transformers
Primary area: interpretability and explainable AI
TL;DR: We theoretically prove multi-task learning is guaranteed to lead to disentangled, generalizable representations in autoregressive models, and validate our theory on RNNs and transformers performing cognitive neuroscience evidence accumulation tasks.

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.

Published as a conference paper at ICLR 2025
DISENTANGLING REPRESENTATIONS THROUGH MULTI-TASK LEARNING
Pantelis Vafidis', Aman Bhargava" Computation and Neural Systems
Antonio Rangel
Humanities and Social Sciences California Institute of Technology a arangel 1 cal wolltech.eee edu
California Institute of Technology (pvafeidi: abhargav @eteeaavvvallleee /tc.ha/aaaaaaaaaaaaecccccccc
ABSTRACT
Intelligent perception and interaction with the world hinges on internal represen- tations that capture its underlying structure ("disentangled" or "abstract" repre- sentations). 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 agents that optimally solve multi-task evidence accumulation classification tasks, canonical in the neuroscience literature. The key conceptual finding is that, by producing accu- rate multi-task classification estimates, system implicitly represents ordinates specifying disentangled representation of the underlying latent state of the data it receives. The theory provides conditions for the emergence of these rep- resentations 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 rep- resentation from the model's latent state (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 bound- ary 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 sys- tems, and helps explain why ANNs often arrive at human-interpretable concepts, and how they both may acquire exceptional zero-shot generalization capabilities.
of co-
INTRODUCTION
The ability to construct representations that capture the underlying structure of the world from data, a hallmark of intelligence. Humans and animals leverage their experiences to construct such faith- ful representations of the world ("world models". " "cognitive maps"), resulting in near-effortless ability to generalize to new settings (Lake et al. 2015) 2016). Modern foundation models also dis- play emergent out-of-distriben di-tiiiiitttion (OOD) generalization abilities, in the form of zero- or few-shot learning (Brown et al., (020 Pham et al. 2021 Jia et al. 2021 Oquab et al.) (002)); however whether artificial systems learn world models remains unclear. Understanding the conditions under which that occurs is bound to lead to better generalizable systems, and explain why artificial systems often converge to human interpretable, aligned representations of the world (Templeton al. 2024).
A promising direction towards understanding the construction of world models is abstract, or disen- tangled representations (Higgins et | 2017 Kim & Mnih 2018 Johnston & Fusi (023).. These two concepts are interrelated yet somewhat distinct (see definitions adapted from Ostojic & Fusi (2024) in Appendix A.I Shortly, an abstract representation of 21...., Zn represents each x; linearly and
"Equal contribution.

---

Published as a conference paper at ICLR 2025
approximately mutually orthogonally. Disentangled representations encode each 2; orthogonally, without the necessity of linearity. Both representations preserve the latent structure present in the world in their geometry by isolating factors of variation in the data, which facilitates downstream generalization. When a representation is abstract, a linear decoder (i.e. downstream neuron) trained to discriminate between two categories can readily generalize stimuli not observed in training, due to the structure of the representation. Furthermore, the more disentangled the representation is, the lower the interference from other variables and hence the better the performance. This corresponds to decomposing novel stimulus into its familiar features, and performing feature- based generaliza- tion. For instance, imagine you are at grocery store, deciding whether a fruit is ripe or not. If the brain's internal representation of food attributes (ripeness, caloric content, etc.) is disentangled, then learning to perform this task for bananas would lead to zero-shot generalization to other fruit (e.g. mangos, Figure D Crucially, the visual representation of a mango is high-dimensional, non-linear and noisy, making it particularly challenging to extract a low dimensional latent like "ripeness".
Several brain areas including the amygdala, prefrontal cortex and hippocampus encode variables of interest in an abstract format (Saez et al. 2015 Bernardi et al. 2020, Boyle et al. 2022 Nogueira et al. 2023 Courellis al. 2024). This raises the question of under which conditions do such rep- resentations emerge in biological and artificial agents alike. Previous work showed that feedforward neural networks develop abstract representations when trained to multitask (Johnston & Fusi 2023). However, real-world decisions typically rely on imperfect, noisy information, evolving dynamically over time (Britten et al. 1992, Krajbich et al.) 2010). To account for this important feature of the world, we train autoregressive models (RNNs, LSTMs, transformers) to multitask canonical neuro- science tasks involving the accumulation of evidence over noisy streams. The tasks tie closely to Bayesian filtering theory, and should be solved by any agent that deals with a noisy world.
Contributions. Our main contributions are the following:
We prove that any optimal multi-task classifier is guaranteed to learn an abstract represen- tation of the ground truth contained in the noisy measurements in its latent state, if the classification boundary normal vectors span the input space (Appendix BB. Furthermore, the representations are guaranteed to be disentangled as the number of tasks Nusk greatly exceeds the input dimensionality D. Intriguingly, noise in the observations necessary to guarantee the latent state would compute a disentangled representation of the ground truth.
We confirm that RNNs trained to multitask develop abstract representations that zero-shot generalize OOD, when Nusk D, and orthogonal, disentangled representations for greater Nusk. The computational substrate of these representations is a 2D continuous attractor (Amari 1977) storing a ground truth estimate in a product space of the latent factors. In addition, the representations are sparse and mixed, attributes of biological neural networks.
We reproduce these findings in GPT-2 transformers, which generalize better due to them learning disentangled representations already from Nusk 2 D, confirming their appropri- ateness for constructing disentangled world models.
We demonstrate that our setting is robust inputs, interleaved learning of tasks and free reaction time tasks canonical in the cognitive neuroscience literature (Britten et al. 1992 Krajbich et al. 2010).
number of manipulations, including correlated
Finally, we discuss implications for generalizable representation learning in biological and artificial systems, and demonstrate the strong advantage of multi-task learning over previ- ously proposed mechanisms of representation learning in the brain Maante et al. 2013).
Although framed in the context of canonical neuroscience tasks, our results are general; they apply to any system aggregating noisy evidence. While our experiments focus on supervised multi-task learning for tractability, the theory only assumes competence at multiple tasks, thus enabling alter- native methods of acquiring such competenee, such as self-supervised unsupervised pre- pre-training.
1.1 RELATED WORK
Disentanglement has long been recognized as a promising strategy for generalization (Bengio et al.
2012) (although note Locatello et al. (2019); Montero et al. (2020) for contrarian view), yet most
classic work focuses on feedforward architectures (Higgins et al. 2017: Kim & Mnih 2018: Whit-
frigton al. 2022 Maziarka al. 2023). In autoregressive models, Hsu et al. (2017): Li & Mandt
2

---

Published as a conference paper at ICLR 2025
2018) showed that variational LSTMs disentagle representations of underlying factors in sequen- tial data allowing style transfer; however the underlying representational geometry was not char- acterised. Other work focuses on fitting RNNs to behavioral data while enforcing disentanglement for interpretability (Dezfouli et al. 2019 Miller et al. 2023). Work on context-dependent decision making has shown that RNNS re- ee-prpooe learned representations in a compositional manner when trained in related tasks (Yang et al. 2019 Driscoll et al., 2022); however, the abstractness of the resulting representations was not established. Finally John al. (2018) show that multitasking re- sults in disentanglement, however unlike us they directly enforce latent factor separation through their adversarial optimization objectives. Our approach is most closely related weakly supervised disentanglement, without comparing across samples (Shu et al. 2019).
Our work relates to previous work on linear identifiability. [Roeder et al. (0021) show that repre- sentations of models trained on the same distribution must be linear transformation. of each other; yet we go beyond their results to show that abstract representations are guaranteed emerge un- der moderate conditions, irrespectively of the dimensionality of the input and model architecture. Lachapelle et al. (2023) proved that disentangled representations emerge in feedforward architec- tures from multitask learning in sparse tasks when sparsity regularization constraint is placed on the predictors; we place no such constraints and still uncover disentangled representations.
Previous ceroosciencc-insppie work showed that multitasking feedforward networks learn abstract representations, as quantified by regression generalization (Johnston & Fusi 2023). We expand upon these findings in several ways. First, extend the framework autoregressive architectures (RNNs, LSTMs, transformers) that can update their representations as further information arrives. Second, we prove theorems that guarantee the emergence of abstract representations in any opti- mal multitask classifier if the number of tasks exceeds the input dimensionality D, and showcase disentanglement in our trained networks. Third, we rigorously analyze the role of noise in forming disentangled representations, extending the noise-free regime studied in Johnston & Fusi (2023) Finally, we explore a range of values for D, providing experimental validation of our theory.
PROBLEM FORMULATION
Multi-Task Classification with Evidence-Aggregattinn We study the evidence accumulation multi-task classification paradigm shown in Figure [p. An agent with latent state Z(t) receives noisy, non-linearly mapped observations (f( J where each X(t) = x - o (0, Ip) is a iniform(-0.55 0.5)). /V being
noisy mcasurement of unknown ground truth vector x" E RR Gaussian noise. The noisy measurements are transformed by an injective observation map f. which can be non-linear and high dimensional, representing the wide range of sensory transformations found in real-world scenarios. The agent is tasked with simultaneously solving Nuak classifica- tion problems by accumulating information over time (a canonical neuroscience task (Britten al. 1, each defined by a random linear decision boundary in the ground truth space R² i.e.
v((x') = {
if if. x² by otherwise
where yxx") E {0, Now represents the Nuas classifications of x", {(c, bi Nam are the classi- fication boundary normal vectors and offsets, and let (t) = g(Z(t)) M) represent the agent's predicted likelihood of (x") = over each of the binary classifications at time t. The classification lines reflect criteria based which decisions will be made. Imagine for example that I corresponds to food and D2 to water reward. Depending on the agent's 's internal state, one takes precedence over the other, and the degree of preference reflected in the slope of the line.
Criterion for Disentangled Representation Learning: We investigate how solving the multi- task classification problem (Figure [p) leads to agents learning disentangled representations of the latent ground truth x" in its internal state Z(t). Specifically, we ask whether there exists linear- affine transformation (A, b) such that x* At + b. Such mapping would imply Ztt) linearly represents x*. If the rows of A are approximately orthogonal, the representation is disentangled.
'Du to the observation map f. the tasks may appear non-linear from the perspective of the multi-task classification agent.
3

---

Published as a conference paper at ICLR 2025
b
Ground truth
latent
Classification
observations
Estimates
(X(I), X(O) Z(O)
Y())
X2
y
Riperos
classificationss
Prcn
Prott
-1
Figure Disentangled representations and a framework to learn them. (a) A disentangled representation directly lends itself to OOD generalization: downstream linear decoder that can dif- ferentiate ripe from unripe bananas can readily generalize mangos, even though has never been trained on mangos. (b) Overview of our multi-task classification framework. A ground truth x" is sampled and Gaussian noise is added to arrive at observations (X(1), ([)).... X(t) These observations processed by the filter-based model illustrated graphically in Figure S7] maintaining latent state Z(t). The latent state Z(t) is then used to produce classification outputs Y; (t). (2(.). Theorem B.6 proves that (t) must encode the optimal estimator of x" given the noisy observations, p(t).
3 THEORETICAL RESULTS
Here we provide conditions and guarantees for the emergence of disentangled representations in op- timal multi-task classifiers with latent state Ztt) the paradigm described in Section and Figure
By "optimal multi-task classifier", we refer to any agent or system whose outputs Y (I) correspond to the correct posterior classification probabilities given the noisy, non-linearly transformed obser- vations; that is, for each task I,..., Neask
(1(1) = Pr (xx(x)) =
(2)
The notion of optimality allows us to make precise statements about the informational content of the agent's internal state since (1) = (Z(()), Let E Nu be a matrix where each row is
decision boundary normal vector. Then
Theorem 3.1 (Disentangled Representation Theorem). If Nax is a full-rank matrix and Nusk 2 D and noise > 0, then
1. Any optimal estimator yx()) must encode a finite-sample, maximum likelihood esti- mate jutt of the ground truth evidence variable x" in its latent state Z(t).
2. If the activation function is sigmoid-like, u(t) will be linearly decodable from Z(t), thus implying that Ztt contains an abstract representation of p(t) (Ostojic & Fusi 2024).
3. The representation guaranteed to be disentangled (orthogonal) as Nuask > D for random decision boundaries.
Specifically, uu(t) is the maximum likelihood estimate (MLE) of x" given observations ((X)),,,,,,(((())) closed-form expression for extracting u(t) from Z(t) if Nusk 2Dz is:
μ(t) = (C (OZ()) +b)
(3)
where the CDF of the normal distribution, the noise magnitude and the trial duration. Furthermore, if the activation function g is of the sigmoid family of functions (tanh, sigmoid), then the term + -1 (g(-)) approximately cancels out, leading to:
p(1) ~ Og o (CT (CT 2(1) CC
(4)
Affine Term
Linear
Zi(t)
where we have approximated Oranh Gaussian IID noise, p(t) is the sample mean of (x()) (i)eee i.e. with non-linearity removed.
2x3 for g tanh and 0.5886 for sigmoid. For
4

---

Published as a conference paper at ICLR 2025
Proof. Point and Equation proven in Appendix B in Theorem B.6 Point 2 and Equation are proven in Corollary f.8 tanh and Corollary B.9 for sigmoid. Point is proven in Corollary B. 10
The key conceptual insight in the proof of Theorem 3.1 is that each of the multi-task classification probability estimates (t) represents an estimated projection distance between the MLE fut) and the given classification boundary bc). Once distances to classification boundaries are recovered, (() can be inferred if the Nusk classification boundaries span the D-dimensional space of x"
Robustness of results While Theorem 3.1 applies to optimal multi-task classifiers, Corollary B.7 shows that a sub-optimal multi-classifier with zero-mean independent errors will represent (t) in state Z(t) (Equation 3) with residual errors w.r.t. optimal jut) expected to decrease at a rate of approximately O (Ntak). See Appendix B.4 B.9 for extensions of the theory to anisotropic and non-Gaussian noise distributions (Elliptical, I-distribution, Laplace distributions). The linear approximation for decoding μu(() from Z(4) in Equation is enabled by the remarkable similarity between sigmoid functions and the Gaussian CDF & (Corollary B.8). The sigmoid- like structure of d suggests many similar activations (e.g., softmax) would exhibit approximate linear decodability.
More general decision boundaries Decision boundaries yi on latents x" may appear non- non-linear in the image of observation map f. but Theorem |. applies to linear boundaries y; on latent space (Eq. D. Our results extend naturally to smooth manifold decision boundaries through local linearization when the manifold yi 's reach 7 is much larger than the noise scale a. Intriguingly, classification boundary distances are only guaranteed to be recoverable when there is non-zero noise >> 0 such that (t) does not saturate to or O, and thus still carries useful decision boundary distance informa- Nisk tion (see Lemma B.3 An intriguing open question is what conditions manifolds {ır } J==1
necessary and sufficient to preserve the decodability of (l)). We leave a complete characterization of representation learning with multiple manifold decision boundaries for future work.
4 METHODS
1ife"x 500
X,
And
Ciacnfication lincs
We trained autoregressive models (RNNs, LSTMs, GPT-2 transformers) with latent state ztt), to output multi-task classifications y(x*) 10,1)MM M given noisy and non- ((())))...f((((() (Figure - We subsequently trained linear probes A on z(t) to estimate x", denoted A(t) = A z(t.. We here focus on leaky RNNs, representing a brain area making decisions; for more details on GPT-2 experiments see Ap- pendix A.6 The networks contain Nnen neu- rons, and their activations zt) obey:
-02
linearly mapped inputs
-0.
0.
0.
1/
Encoder
T4- -2 [ me + Win Xtn b]+
(5)
Figure 2: Data generation and architecture. (a) For each trial, we sample a ground truth vector x", and add IID noise to arrive at X(t). The task is to report whether x" lies above (1) or below (0) each of the classification lines (color matches corresponding boolean variable in y), given the noisy and non-linearly transformed sam- ples f (X(1)). (X(()). (b) Models (RNN depicted) are trained report the outcome of all the binary classifications in at the end of the trial (indicated by the fixation input turning 0).
where W ree is the recurrent weight matrix, W
is the matrix carrying the input vector Xic, is unit-specific bias vector, T is the neuronal time constant, [.]4 is the ReLU applied element-wise and time dependencies have been dropped for brevity. We discretize Equation 5 the for- ward Fuler method for T 20 timesteps of duration At = 100 ms, which we find to
maximum distance at which each point on a manifold has a unique closest point on the manifold II fact, Equations 3 and do not hold for - 0, as they were derived hy means of Bayesian estimation which assumes the presence of noise.
2"Raahh":
5

---

Published as a conference paper at ICLR 2025
be stable. The RNN's output yt) R is given by y(t 9( our z(t)), where W our is read- out matrix and g = sigmoid the output activation function applied elementwise. The encoder ff is a 3-layer MLP with hidden dimensions 100, 100, 40 and ReLU non- non-linearities, and it is randomly ini- tialized and kept fixed during training as it represents a static mapping from latents to observations (observation map). An additional fixation input is directly passed to the hidden layer. It is during the trial and turns O at the end of the trial, indicating that the network should report its decisions (Figure 2p. The fixation input is concatenated with f (()( to form X and it precludes the RNN from learning a specific timing in its response. We refer to this kind of tasks as fixed reaction- time (RT). The network is trained with a cross-entropy loss and Adam default settings, except learning rate 1/0 10-3, to produce the target outputs yx").. By minimizing loss across trials, the network is incentivized to estimate (t)) (x() 1]. Table S1 summarizes all hyperparameters and their values, which are shared across all architectures.
5 EXPERIMENTS
5.1 MULTI-TASK LEARNING LEADS TO DISENTANGLED REPRESENTATIONS
We train RNNs to do simultaneous classifications for Nusk linear partitions of the latent space for D (Figure partitions shown). To quantify the disentanglement of the representations after learning, we evaluate regression generalization by training a linear decoder to predict the ground truth x* while network weights are frozen. We perform out-of-distribution 4-fold crossvalidation, i.e. train the decoder on 3 out of 4 quadrants and test in the remaining quadrant (Appendix A.. for details). We also evaluate in-distribution (ID) performance by training the decoder in all quadrants. An example of train and test losses is shown in Figure SITE. We find that the network's OOD and ID generalization performance are excellent (median 2 0.96, 0.97 respectively across example net- works); therefore the network has learned abstract representation that zero-shot generalizes OOD. In addition, ID performance increases with the number of tasks Nusts, and the O0D) generalization gap decreases (Figure 3)). Performance identical when choosing more nonlinear, power-law nonlinearity for the encoder (Appendix (A.5). Therefore we conclude that multi-task learning leads to abstract representations in the RNN's hidden layer, when tasks span the latent space.
Since x* can be decoded by this representation in unseen (by the decoder) parts of the state space, it follows that the representation can be used to solve any task involving the same latent variables, without requiring further pretraining. In other words, to solve any other task we do not need to deal with the denoising and unmixing of the latent factors 31, Z2; we would just need to learn the (potentially non-linear) mapping from I1, 32 to task output. Furthermore, the representation scales linearly with input dimensionality D (see Figure 5b). This marks a significant improvement from previously proposed models for representation learning in the brain where one task is executed at a time (Mante et al. 2013, Yang et al. 2019). which scale linearly with Nass and exponentially with D (see Appendix[A.8 for details). Crucially, these findings are cehiitettureagnnostee for non-leaky ("vanilla") RNNs, which outperform leaky ones for small Nuau. LSTMs which per- form the best, and GPT-2 transformers (details in Appendix A.6) which have excellent performance already from Nuask = 2 (Figure Bp). Note that state-space models have superior asymptotic perfor- mance, which expected due to the nature of the task. We focus on leaky RNNs because of their closer correspondence to biological neurons, which have membrane voltage that decays over time.
they hold
So far we showcased abstractness, but not disentanglement. For disentanglement, it is crucial that the latents lie in orthogonal subspaces. Looking at the angles between the decoders of the latents, we find that they become orthogonal as N Yusk& >> for RNNs (Figure as predicted by our theory. Intriguingly, this already occurs from Niss 2 D in transformers, showcasing their superior ability to separate latent factors. Furthermore, orthogonality strongly correlates with OOD generalization performance, which emphasizes the close link between disentanglement and abstractness: the more orthogonal the representations are, the cleaner the readout of the latent factors by linear decoders.
We further demonstrate the robustness of our setting by showing that abstract representations emerge for different noise distributions and correlated inputs (Appendix A.9). non-linear bound- aries (Appendix 10) and for cognitive neuroscience integrate- bound tasks where the agent can make their decision whenever confident enough, not at a fixed time (Krajbich et al., 2010) (Ap- pendix

---

Published as a conference paper at ICLR 2025
1.0
A-gio anans loogs
100
inout dimensionelity 0-2
d
Time
PCA
0.0
0.8
82-
0.7
distitutioon diatroution
0.6-
70
Component
0.5
10
50
10
bascs
of tesks
Egorndiuns onay beal part Lower
Qun of b.dee
1.0
Dot
09-
0--
X
0.7-
PC3
PC2
10
-18
1031 00
Figure 3: Learning disentangled representations. (a) ID and OOD generalization performance for networks trained different number of tasks Nustt- We report the 25, 50 and 75 percentile of r2 for each network size (see Appendix (A.2). ID and OOD performance increase with Nunk and the generalization gap decreases, indicating that the networks have learned abstract representations. (b) The results hold for other autoregressive architectures, including LSTMs and GPT-2 transformers. (c) Angles between latent factor decoders (see Appendix [.3] for how they were estimated). The angles approach 90 degrees Nask > D for RNNs. but already fror Nusk D for transformers. Remaining errors around 90 degrees are attributed to variability in the linear decoder fits. (d) Top PCs of RNN activity (Nins 24, = 2), capturing 85% of variance (see inset). Each line a trial, while color saturation indicates time. All trials start from the center and move outwards, towards the location of x" state space. We color the last timepoint in each trial (squares) according to the quadrant this trial was drawn from. Red x's correspond to attractors (see Appendix A.7). Here we move input noise so that trajectories can be visualized casier. The network learns a -iv-diiensional continuous attractor that provides disentangled representation of the 2D state space. (e) Spectral plot resulting from linearizing RNN dynamics around every fixed point (Appendix [A.7]. First two eigenvalues of the difference system near O, while the rest decay much faster, indicating marginal stability across two dimensions for every fixed point, a signature 2D continuous attractor.
5.2 REPRESENTATIONAL STRUCTURE IN RNNS AND TRANSFORMERS
In this section, we open the black box and investigate the representations learned by the networks, starting with RNNs. Figure 31 shows the top PCs (capturing 85 % of the variance) of network activity after training (final accuracy 95%) for multiple trials, along with the fixed points of net- work dynamics. To find the fixed points, we follow a standard procedure outlined in [Sussllo & Barak (2013) (see Appendix A.7 for details). Looking at Figure 31 the fixed points span the entire two-dimensional manifold that the trials evolve in, which corresponds to a continuous attractor with stable states across 2D "sheet". Linearizing the dynamies around each fixed point and computing the eigenvalues of the linearized system (Appendix A.. for details), reveals marginal stability across two eigenvectors, i.e. near-0 eigenvalues which correspond slow, integration dimensions in net- work dynamics, therefore confirming the continuousness of the attractor (Figure Be). This implies that the network can store a short-term memory (Wang] 200T) of the current amount of accumulated evidence in product space of the latent variables, and update it as further evidence arrives.
Furthermore, compared to the representations after the encoder which are non non-linearly mixed, high- dimensional and overlapping (Figure [SITa), the representation in Figure 3 looks disentangled as we would expect from the theory and metries above. Individual trials with noise show how the rep- resentation maintains a sense of metric distances in the RNN representation space (Figure SST1)) Figure SIT;: demonstrates how this representation comes about during learning, and Figure S that the short-term memory persists when delay period is included before the decision. Therefore, multi-task learning has led to disentangled, persistent representations of the latent variables. Impor- tantly, and in line with our theory, this only happens when noise is present in the input, which forces the network to learn notion of distance from classification boundaries (Lemma (.3). Indeed, when the network trained without input noise, it does not learn 2D continuous attractor [Figure(STTp)

---

Published as a conference paper at ICLR 2025
a
RNN
GPT
Combaton combrtnssoaaarr conPotet
RNN
Finz oen (eckooa)
GPT
Weunons
Figure 4: RNN and GPT representations and relation to latent variables. (a) Hidden layer activations of RNN in Figure P (left) and GPT-2 transformer (right), while systematically varying the latent factors 1 and 22 from -0.5 to 0.5. Activations are plotted in 8*8 grids, one for each value of 21 and z2. Each grid contains firing rates for a total of 64 neurons for the RNN, and activations for 8 units for each of the 8 heads from the final embedding of the sequence for GPT-2. (b) Correlation coefficient of activations for both models with 21 and ø2, respectively.
Finally, we examined RNN and GPT unit activations, and their relation to the latent variables. In Figure 4a we plot activations for all 64 units for both networks, while regularly sampling æι and I2. RNNs representations are sparse, with only ~ 10 % of neurons active at any time, which is in line with sparse coding in the brain (see Appendix AA.1 for quantification of sparsity as a function of Nust. D and RNN architecture). In addition, the average firing rate is ~ spike/s, which sur- prisingly close cortical values. Transformers on the other hand, do not have these features, shared by RNNs and their biological counterparts. Furthermore, we find that both networks display mixed selectivity, i.c. neurons are tuned to both variables, which is a known property of cortical neurons (Rigotti et al. 2013) (Figure 4b). This suggests that metrics of disentanglement that assume that individual neurons encode distinct factors of variation (Higgins et al.) 2017, Kim & Mnih 2018] Chen et al.] 2018 Eastwood et al.) 2022] Hsu et al. 2023) might be insufficient in detecting disen- tanglement in networks that generalize well. While recent work incorporates such axis-alignment in the definition of disentaglement, our work along with others (Johnston & Fusi, 2023) showcases the advantages of approaching disentanglement from mixed representations perspective. Importantly, these properties were not imposed during training, was there any parameter fine tuning involved; they emerged from task and optimization objectives.
5.3 EXPERIMENTS CONFIRM AND EXTEND THEORETICAL PREDICTIONS
Here we expore the relation between the theory in Section 3 and Appendix B and experiments in Section in more depth. First, we wondered why performance saturates in our networks to a high yet non-I 2 The central limit theorem predicts that the estimate of the ground truth x" in any optimal multi-task classifier becomes more accurate with Vt, providing theoretical maximum ,2 given trial duration T (Appendix A.4). Since the RNNs trained on the free reaction-time (free RT) output their decision confidence at any time in the trial,
task in Appendix [I.I are required we can compute OOD 22 of free RT network predictions at any timepoint I, and compare that to the theoretical prediction. Figure 5p shows that indeed the highest RNN , falls in the vicinity of just short of the theoretical maximum. This indicates that RNNs trained with BPTT on these tasks behave like near-optimal multi-task classifiers that create increasingly accurate predictions with time, tightening the relation between theoretical and experimental results.
An important prediction of our theory is that to learn abstract representations, Nua should exceed D. To test this, increase D (adding more inputs Figure 2p), while varying Naatt Sampling classification hyperplanes homogeneously (similar to Figure 2a. center) in high- high-dimensional spaces is non-trivial; therefore resort to randomly sampling them. Figure 5 shows OOD generalization
8

---

Published as a conference paper at ICLR 2025
b
RNN
GPT
tasks
Out Out deoo--dtiibttonn
1.0
1.0
1.0
buttor-derruutom
0.8
0.8
0.0
0.0
0.6
0.6
0.8-
Experiments Theory
0.4
0.4
0.0
0.7
0.0
0.0
0.7
0.5
1
1 2.
6
B
10
4
6 3
10
0.1
0.2
0.3
0.5
Thal duration (s)
Ingut dimerskonaltty
D
Input dimersionnlity
Treining nobe sid
Figure 5: Experiments confirm theoretical predictions. (a) OOD ,2 for free RT RNN required to report its estimate of x' at different times (see Appendix AA.1 for details, Nuust 24, D 2).
Maximum network y2 matches optimal multi-task classifier theory predictions (Equation 7 in Ap- pendix A.4). (b) OOD 2² function of input dimensionality D and number of tasks Nuast- Good
values of are obtained when Nust 2 D, especially for GPT models, confirming our theoretical
results. (c) Increasing amounts of noise in pretraining results in better OOD generalization (D 2).
performance for various combinations of D and Nusk- We observe that performance is bad when the Nuak < D, but it increases when Nuatt 2 D. For RNNs, this increase is abrupt for smaller D and more gradual for higher, which is in line with remarks by Johnston & Fusi (203)) that it is easier to learn abstract representations when D is high. Transformers on the other hand display higher generalization performance than RNNs, and always perform almost perfectly when Nank > D, demonstrating their superior performance in learning abstract representations. Looking at the angles between latents for higher D (Figure(ST2), we find that transformers have excellent disentanglemen as long as Nusk 2 D, which might explain their superior generalization performance RNNs for lower Nast These results, together with Figure 3 demonstrate the superior ability of transformers in disentangling latent factors. Overall, our findings confirm our theory that abstract representations emerge when Nust 2 D, and even go beyond to suggest that disentangled representations emerge earlier than the theoretical condition Naak > D, as long as the architecture is appropriate. These results remarkable, especially for high D, because they go against our intuition that Nun should scale exponentially with D to fill up the space adequately; instead it need only scale linearly.
Importance of noise for generalization Our theory and experiments provide insight on the im- portance of noise for developing efficient, abstract representations (Figure (ITe).. The closer to classification boundary the ground truth x" is, the more likely noise will cross over the boundary. Since, our theory shows, any optimal multi-task classifier has to estimate Pr(y, said probability directly relates to the actual distance from the boundary, it follows that noise allows the model to learn distances from boundaries (Lemma B.3 ) leading to efficient localization. We reasoned that additional noise might be even more beneficial, as it would allow more accurate esti- mation of (x)) = 1}. especially when x* far from the boundary. To test this, we increase noise strength when pretraining RNNs, while testing with the same = 0.2. Indeed, increasing amounts of noise consistently result in better OOD generalization (Figure 55).. This benefit comes for smaller numbers of tasks, allowing us to consider less supervised tasks (e.g. 3). train on them with more noise, and achieve the same performance as more tasks (e.g. 12). So even though net- works with more noise perform worse in pretraining (low 90%s classification accuracy), they learn more abstract representations. These findings are highly non-trivial, and have informed our thinking about generalization and inherent variability of the underlying latent factors.
1}. and
DISCUSSION
In this work, we proved that disentangled, generalizable representations must emerge in agents optimally solving multi-task evidence accumulation tasks canonical in the neuroscience literature. We also conducted experiments in suite of autoregressive models (RNNs, LSTMs, transformers) which confirmed all of the main theoretical predictions. A key takeaway is that transformers more readily disentangle representations, which may explain their unique world understanding abilities. Here we discuss the broader impact of this work for representation learning and neuroscience alike. I Limitations of this study and how can be extended in the future are discussed in Appendix A.15

---

Published as a conference paper at ICLR 2025
6.1 IMPLICATIONS FOR REPRESENTATION LEARNING
Topology-prees preservving representation learning Our work has profound implications for learning representations that inherit the topological structure of the world. We prove this naturally happens as long as there are enough tasks to uniquely identify the location of x". Crucially, the constraints from different tasks should be placed simultaneously on the representation, which explains why represen- tations from context-dependen computation (Mante et al. 2013) are typically not disentangled.
Representational alignment across individuals Our results provide a new perspective on the Platonic representation hypothesis (Huh al. 2024), which suggests that the convergence in deep neural network representations is driven by a shared statistical model of reality, like Plato's 's concept of an ideal "Platonic" reality. Theorem 3.] suggests that the key factor driving convergence is the diversity and comprchensiveness of the tasks being learned. As long as individuals are faced with similar day-to-day tasks that collectively span the space of the underlying data representation, con- vergence to a shared, ceality-aligned representation can occur. This could explain why for example modern LLMs come to encode high-level, interpeattereee concepts (Templeton et al., 2024).
Manifold hypothesis While our problem is framed in terms of arbitrary injective observation map f. the formulation encompasses many scenarios relevant to the manifold hypothesis (Fefferman cet al., 2013). The function can represent a smooth manifold embedded in a high dimensional space, directly modelling the manifold hypothesis of deep learning. In neuroscience, f could be non-linear encoding of stimuli in a neural population response, connecting our work to neural manifold research (Langdon et al., (023). By developing and testing theoretical guarantees for the emergence of disentangled representations in this multi-task problem formulation, we provide insight on how neural networks can inherently discover and linearize low-dimensional manifolds within high-dimensional, non-linear observations, enhancing our understanding of how complex data structures are captured and represented in deep learning models and biological systems alike.
Interplay between number of tasks and fine- theegrainness of representations Finally, the theorem and experimental results presented here are not a ste-rettt from dimensionality 1 of the la- tents to how many tasks Nus are required to uncover them. Instead, there is a fundamental interplay between richness of tasks performed and detail of the representation learned. a high-dimensional world, the richness of the tasks at hand directly affects the dimensionality D of the latents that can be extracted, allowing for "ground truths" x* at different levels of granularity to be explored. The richer the label information available, the more fine-graine grained the resulting world model will be.
6.2 IMPLICATIONS FOR NEUROSCIENCE
The brain encodes variables of interest in disentangled format, in processes disparate mem- ory (Boyle et al.) 2022), emotion (Saez et al. 2015), and decision making (Bongioanni et al. 202T). Furthermore, performance in tasks has been shown to degrade once abstract representations col- lapse (Sacz et al.] 2015),, supporting their role in guiding generalizable behavior. Our findings put forth parallel processing as unifying mechanism for generalization in brains. The cortex, with its massively parallel architecture (Markram et al., 20155) Hawkins et al., (219)) is a prime candi- date area for the construction of disentangled, generalizable world models. Another candidate area is the thalamus; it is posited that thalamocortical loops operate in parallel, and combined with in- ternal state-dependent mechanisms lead to state-dependent action selection (e.g. prioritizing water when thirsty), while evidence integration occurs in corticostriatal circuits (Rubin et al. 2020). The representations discovered here (continuous attractors, CANs) have been widely found in the brain when solving similar tasks, highlighting their role as a general computational substrate for cognitive functions in the brain (for relations of our work to the neuroscience literature, see Appendix A3.1)). Notably, the receipt of rich supervisory signals from the environment is not a reguirement for our setting, as it can leverage the output of previously learned tasks (see Appendix A. 14 on the biological plausibility of multi-tase learning).
The algorithmic efficiency of multi- task learning compared to alternatives ("context-dependent com- putation" "( Mantee et al. (2013). Appendix AA ) makes think that it is coincidence that the cor- tex can support parallel processing; all the pieces are there, and we feel that the brain has to leverage this feature to construct faithful models of the world, it does.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
All code used to generate the results can be found in experiments are seeded, ensuring exact reproducibility results. Five networks have been trained for every configuration shown, to provide sufficient statistics to support our conclusions. To ensure full clarity of the thcoretical results in the main text, a full proof provided in Appendix
The
AUTHOR CONTRIBUTIONS
P.V. conceived the project, wrote the code, performed the experiments and wrote the manuscript. A.B. developed the theory, wrote the theory sections and parts of the manuscript. A.R. supervised research and provided funding.
ACKNOWLEDGMENTS
P.V. would like to thank the Onassis Foundation and A.R. the NOMIS Foundation for funding. A.B. thanks the NIH PTQN program for funding. No competing interests to declare. We would like to thank Yisong Yue for early discussions, Aiden Rosebush for discussions proof methods and Stefano Fusi for early feedback.
REFERENCES
Eddie Aamari, Jisu Kim, Frédéric Chazal, Bertrand Michel, Alessandro Rinaldo, and Larry Wasser- man. Estimating the reach of manifold, 2019. URL
04565
Abdullahi Ali, Nasir Ahmad, Elgar de Groot, Marcel Antonius Johannes van Gerven, and Tim Chris- tian Kietzmann. Predictive coding is a consequence of energy efficiency in recurrent neural net- works. Patterss, 3(12):000639, December 2022. ISSN 2666-3899 3899, doi: 10.1016/j 10.1016/j.paterr.2000 2022
100639. /
dx doi org/ 1016 patter www.aaaeee.e.........aeeeeee.e 100639
Matteo Alleman, Jack Lindsey, and Stefano Fusi. Task structure and nonlinearity jointly determine learned representational geometry. In The Twelfth International Conference Learning Repre- sentations, 2024. URL
S Amari. Dynamics of pattern formation in lateral-inhibition type neural fields. Biological Cyber- netics, 27(2)777-87, 1977. doi: 10.1007/bt03332259. URL
bf00331259
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives, 2012. htthst/:::// /ttp/vaiiiv/../////i//eae/aaaaev argg/bs// /aabbs////..... 1206. 5538
Silvia Bernardi, Marcus Benna, Mattia Rigotti, Jérôme Munuera, Stefano Fusi, and C. Daniel Salzman. The geometry of abstraction in the hippocampus and prefrontal cortex. Cell, 183(4): e21, November 2020. doi: 10. cell..2.e....020000..
URL httph:tt///.m :tttss://oo:. org/
10.1016/3.c0 www.eeee...0....... 2020. 09.09..00 031
Aman Bhargava, Cameron Witkowski, Alexander Detkov, and Matt Thomson. Prompt baking, 2024.
URL
220.:/ 13697
A. Bongioanni, D. Folloni, L. Verhagen, J. Sallet, M. C. Klein-Fllügge, and M. F. S. Rushworth. Activation and disruption of a neural mechanism for novel choice in monkeys. Nature, 591(7849): 270-20 274, January 2021. doi: 1038/ 0 4 586 020 03115-5
URL httpt:: /tttps://oi..o /wi. //tp:///o.o..0//..ooo org/10.
George EP Box and George C Tiao. Bayesian inference statistical analysis. John Wiley & Sons, 2011.
Lara M. Boyle, Lorenzo Posani, Sarah Irfan, Steven A. Siegelbaum, and Stefano Fusi. Tuned ge- ometries of hippocampal representations meet the demands of social memory, January 2022. URL
477361

---

Published as a conference paper at ICLR 2025
KH Britten, MN Shadlen, WT Newsome, and JA Movshon. The analysis of visual motion: comparison of neuronal and psychophysical performance. The Journal of Neuroscience, 12(12):
4745-4765, December 1992. doi: 10. 10.1553/jmeurosci..222 wtt.://urrocc...2200004........ 12--14445.91992
org/10. 10.10.1.223 1tte/10.552/jjeeeeee / 12 12 04 0474 1992
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Vosse Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess. Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL
https: //tpxi://axxvvv org/ abs wwwpa//s..0. 2005. 14165
Bingni W. Brumton, Matthew M. Botvinick, and Carlos D. Brody. Rats and humans can optimally accumulate evidence for decision-making Science, 340(6128):95- April 2013. doi: 10.1126/ science.
htttttts: org. wwwp///ce//e....ceccccc ence. 1233912
Yoram Burak and Ila R. Fiete. Accurate path integration in continuous attractor network models of grid cells. PLoS Computational Biology, 5(2)ee000029, February 2009. doi: 10.1371/joural. pebi. 1000.10 URL httlh:tts:///o/ http:///ooorrro.////// 1001110.3311/orrrrrr .pt.rnaall..........m 1000291
Tian Qi Chen, Xuechen Li, Roger B. Grosse, and David Duvenaud. Isolating sources of disentangle- ment in variational autoencoders. CoRR, abs/1 802. 04942, 2018. (RL/ttt:://arxii org/
abs/1880 04942
Amit Choudhury. A simple approximation to the area under standard normal curve. Mathematics and Statistics, 2(3):177-149, 2014.
Hristos S. Courellis, Juri Minxha, Araceli R. Cardenas, Daniel Kimmel, Chrystal Reed, Tau- fik A. Valiante, C. Daniel Salzman, Adam N. Mamelak, Stefano Fusi, and Ueli Rutishauser.
Abstract representations emerge in human hippocampal neurons during inference. Nature, 632 URL
(8026):841-849, August 2024. ISSN 1476-4687. doi:
http://dx. doi. wwtp//g/......... 1038/ 024-07799
Thomas M Cover and Joy A Thomas. Information theory and the stock market. Elements of Infor- mation Theory. Wiley Inc, New York, pp. 543-556, 1991.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. CoRR, abs/1810.04805, 2018. URL
http://arxi:: org. abs 1810 04805
Amir Dezfouli, Hassan Ashtiani, Omar Ghattas, Richard Nock, Peter Dayan, and Cheng Soon Ong. Disentangled behavioural representations. Advances in neural information processing systems, 32, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition scale, 2020. URL httpttts///axii..oooo http://xxvvc..g./eee.// /ab/ 201...../2000.99.
Laura Driscoll, Krishna Shenoy, and David Sussillo. Flexible multitask computation in recurrent networks utilizes shared dynamical motifs. bioltxi, pp. 2022-08, 2022.
Cian Eastwood, Andrei Liviu Nicolicioiu, Julius von Kügelgen, Armin Kekić, Frederik Träuble, Andrea Dittadi, and Bernhard Schölkopf. Dei-es: An extended disentanglement framework with connections to identifiability. 2022. URL htt.:ttts:///aat
abs. 2210. 2ttpp///2210000224.223033
Kai Wang Fang. Symmetric multivariate and related distributions. Chapman and Hall/CRC, 2018.
Charles Fefferman, Sanjoy Mitter, and Hariharan Narayanan. Testing the manifold hypothesis, 2013. URL
abspgg/bbb/11100 0425
Charles Fefferman, Sanjoy Mitter, and Hariharan Narayanan. Testing the manifold hypothesis. Journal of the American Mathematical Society, 2944)9933-1049 2016.
12

---

Published as a conference paper at ICLR 2025
Timo Flesch, David G. Nagy, Andrew Saxe, and Christopher Summerfield. Modelling continual learning in humans with hebbian context gating and exponentially decaying task signals, 2022.
URL htttt::/aaxiiiii org/ abt//203...2. 11560
David J. Foster and Matthew A. Wilson. Reverse replay of behavioural sequenees in hippocampal place cells during the awake state. Nature, 40007004))0000000-00080 683, February 2006. ISSN 1476-4687.
doi: 10. 10.1038/nature04557..
Richard J. Gardner, Erik Hermansen, Marius Pachitariu, Yoram Burak, Nils A. Baas, Benjamin A. Dunn, May-Britt Moser, and Edvard Moser. Toroidal topology of population activity in grid cells. Nature, 602(7895):123 128, January 2022. doi: 00.1088841156--011-4420000 h -20.--10000011111-0-----------11
URL
James Gornet and Matt Thomson. Automated construction cognitive maps with visual predictive coding. Nature Machine Intelligence, 6(7):820-833, July 2024. ISSN 2522-5839. doi: 10.1038/
>42256-024-0086--L. URL http: //tp:////xx/// doi http://xddooooggg//...oooo htt.::////........oo/11111118/1/
Jeff Hawkins, Marcus Lewis, Mirko Klukas, Scott Purdy, and Subutai Ahmad. framework for intelligence and cortical function based on grid cells in the neocortex. Frontiers in Neural Circuits, 12. January 2019. ISSN 1662-5110 10. doi: 10.3389/fncir.211800000
URL //xt doi.
33t9// fncir 2018. 00121
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners, 2021.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAF: Learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations, 2017.
net http://rrreeeeeeeoooorrmaaaaaaaa
Irina Higgins, David Amos, David Pfau, Sebastien Racaniere, Loic Matthey, Danilo Rezende, and Alexander Lerchner. Towards a definition of disentangled representations, 2018. URL https:
7/ www.iv.org//ab... 02230
Kyle Hsu, Will Dorrell, James C. R. Whittington, Jiajun Wu, and Chelsea Finn. Disentanglement via latent quantization, 2023. URL https://pp/at:: org/ abs 2200//11358
Wei-Ning Hsu, Yu Zhang. and James Glass. Unsupervised learning of disentangled and in- terpretable representations from sequential data, 2017. URL https://arxiv: org/abs/
7709907900
Minyoung Huh, Brian Cheung, Tongzhou Wang, and Phillip Isola. The platonic representation hypothesis. arXiv preprint arXiv:2405.009987 2024.
Kristopher T. Jensen, Guillaume Hennequin, and Marcelo G. Mattar. A recurrent network model of planning explains hippocampal replay and human behavior. Nature Neuroscience, 27(7): 1340-1348, June 2024. ISSN 1546-1726. doi: 10.
URL http:
7/dx.
10. 1tt../11111310022222000
Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoe V. Le, Yunhsuan Sung, Zhen Li, and Tom Duerig. Scaling up visual and lasion-l representation learning with noisy text supervision, 2021. URLpttts:://axxt
orgooog///a......ooo wttgg//bs//......c 05918
Vineet John, Lili Mou, Hareesh Bahuleyan, and Olga Vechtomova. Disentangled representation learning for non non-paralle con-parllll text style transfer, 2018. URL Utt/tttt.s//avvvv.o///b/////0
04339
W. Jeffrey Johnston and Stefano Fusi. Abstract representations emerge naturally in neural networks trained to perform multiple tasks. Nature Communications, 14(1), February 2023. doi: 10.1038/ 441467-023-36588 /tt.:httpss///o..orr//////
13

---

Published as a conference paper at ICLR 2025
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2649 2658. PMLR, 10-15 Jul 2018. URL 8b. html
Samuel Kotz and Saralees Nadarajah. Multivariate i-distributions and their applications. Cambridge university press, 2004.
Krajbich, C. Armel, and A. Rangel. Visual fixations computation and comparison value in simple choice. Nature Neuroscience, ((110):1222-1999 2010. doi: 10.1038/mn.2635. URL https: orgp::///do.og/////o/ooo wtw.....111111118811111111 26355
Sebastien Lachapelle, Tristan Deleu, Divyat Mahajan, Ioannis Mitliagkas, Yoshua Bengio, Simon Lacoste-Julien, and Quentin Bertrand. Synergies between disentanglement and sparsity: General- ization and identifiability in multi-task learning. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th In- ternational Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 18171-18206. PMLR, 23-29 Jul 2023. URL https: httpss//prcceedding...m press/ /2133/22222 http://1chhpppeeeecc Te23a html
B. M. Lake, R. Salakhutdinov, and B. Tenenbaum. Human-level concept learning through prob- abilistic program induction. Science, 35066266) December 2015. doi: 10.1 science.aab3050.
Utt/ttss:://aii.
012.0....6665880 aab3050
Brenden M. Lake. Tomer D. Ullman, Joshua Tenenbaum, and Samuel J. Gershman. Build- ing machines that learn and think like people. Behavioral and Brain Sciences, 40, Novem- ber 2016. doi: 10. 00.1017/s01402222 40525x
URI. https://dci http:///..oorroo....o. orgg/1..111111
0115525xx60118..
Christopher Langdon, Mikhail Genkin, and Tatiana A. Engel. A unifying perspective neural man- ifolds and circuits for cognition. Nature Reviews Neuroscience, 24(6))363-377 April 2023, ISSN 1471-0048. doi: 11238/441583-022-00066-9-9 s41583- 02 023-00033------
URL http: http://xxxccii..
Kenneth Levenberg. A method for the solution of certain non-linear problems in least squares. Quarterly of Applied Mahematics, 2(2):164-168, 1944. ISSN 1552-4485. doi: 10. 1090/a// 10666. URL htttp://dx doi.doi.ooooooooo http://xw......l////00000///////
Yingzhen Li and Stephan Mandt. Disentangled sequential autoencoder, 2018. URL https:/
Francesco Locatello, Stefan Bauer, Mario Lučić, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Frederic Bachem. Challenging common assumptions the unsuper- vised learning of disentangled representations. In International Conference on Machine Learn- mlr. pttip.pres/////occtellll
ing, 2019. Paper Award.
html Best
Darby M. Losey, Jay A. Hennig, Emily R. Oby, Matthew D. Golub, Patrick T. Sadtler, Kristin M. Quick, Stephen I. Ryu, Elizabeth C. Hyler-Kabara Aaron P. Batista, Byron M. Yu, and Steven Chase. Learning leaves memory trace in motor cortex. Current Biology, 34(7):519-11311 April 2024. ISSN 0960-9822. doi: 10. 10.1016/j.cu..2244.0..
URL /ot.
10.1016 cub. 2024. 003
Yi Ma, Doris Tsao, and Heung- Yeung Shum. On the principles of parsimony and self-consistencyy for the emergence of intelligence, 2022.
V. Mante, D. Sussillo, K. V. Shenoy, and W. T. Newsome. Context-dependent computation by recurrent dynamics in prefrontal cortex. Nature, 503(7474):78 84, November 2013. doi: 10. 1 URL CRI.https::/// sattps::/do::. wtw.:oorgr/1... 1.200.000330 www.:/10t////1e112/00
14

---

Published as a conference paper at ICLR 2025
Henry Markram, Eilif Muller, Srikanth Ramaswamy, Michael W. Reimann, Marwan Abdellah, Car- los Aguado Sanchez, Anastasia Ailamaki, Lidia Alonso-Nanclares, Nicolas Antille, Selim Arse- ver, Guy Antoine Atenekeng Kahou, Thomas K. Berger, Ahmet Bilgili,, Nenad Buncic, Athanas- sia Chalimourda, Giuseppe Chindemi, Jean-Denis Courcol, Fabien Delalondre, Vincent Delattre, Shaul Druckmann, Raphael Dumuse, James Dynes, Stefan Eilemann, Eyal Gal, Michael Emiel Gevaert, Jean-Pierre Ghobril, Albert Gidon, Joe W. Graham, Anirudh Gupta, Valentin Haenel, Etay Hay, Thomas Heinis, Juan Hernando, Michael Hines, Lida Kanari, Daniel Keller, John Kenyon, Georges Khazen, Yihwa Kim, James G. King, Zoltan Kisvarday, Pramod Kumbhar, Sébastien Lasserre,
Le Bé, Bruno R.C. Magalhães, Angel Merchán-Pérez Julie Meystre, Benjamin Roy Morrice, Jeffrey Muller, Alberto Muñoz-Céspedess, Shruti Muralidhar, Keerthan Muthurasa, Daniel Nachbaur, Taylor H. Newton, Max Nolte, Aleksandr Oveharenko, Juan Palacios, Luis Pastor, Rodrigo Perin, Rajnish Ranjan, Imad Riachi, José- Rodríguez, Juan Luis Riquelme, Christian Rössert, Konstantinos Sfyrakis, Ying Shi, Julian C. Shillcock, Gilad Silberberg, Ricardo Silva, Farhan Tauheed, Martin Telefont, Maria Doledo-Rodriguez Thomas Tränkler, Werner Van Geit, Jafet Villafranca Díaz, Richard Walker, Yun Wang, Ste- fano M. Zaninetta, Javier DeFelipe, Sean L. Hill, Idan Segev, and Felix Schürmann. Recon- struction and simulation of neocortical microcireuitry. Cell, 163(2)4556-99,, October 2015. ISSN 0092-8674. doi: T.ce11.2015.0 02.029
URL http: doi www.ooggg//.... http:///x.......1111..1...1.1... 1016/
Donald W. Marquardt. An algorithm for least-squarees estimation of nonlinear parameters. Journal of the Society for Industrial and Applied Mathemattics, 11(2):431-441 June 1963, ISSN 2168-3484.
doi: 10.1137/01 1030. URL htt.:t/pp/bbk.co....////00/11//1
Paul Masset, Torben Ott, Armin Lak, Junya Hirokawa, and Adam Kepecs. Behavior- and modality- general representation of confidence in orbitofironta cortex. Cell, 182()):11--126.1188 July 2020. doi: 10.1006ff.cel1.22200...00 http:///00/0.../////////.//.....
022
Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner. dsprites: Disentanglement testing sprites dataset.
2017.
Lukasz Maziarka, Aleksandra Nowak, Maciej Wolczyk, and Andrzej Bedychaj. On the Rela- tionship Between Disentanglement and Multi-task Learning, pp. 625-641. Springer Interna- tional Publishing, 2023. ISBN 9783031263873, doi: 10. 0000000977778300--66337777777770 http: /ttp://dxx doi. -30
URL
Kevin J. Miller, Maria Eckstein, Matthew M. Botvinick, and Zeb Kurth-Nelson. Cognitive model discovery via disentangled mns, June 2023. URL http:h/t:// doi /ttp:///oo....r.....1110//22.220
06. 54.250
M. L. Mittelstaedt and H. Mittelstaedt. Homing by path integration in mammal. Naturwis- senschaften, 67(11):566 567, November 1980. doi: 10.1007/b/00450672. URL https //doi.
Milton Llera Montero, Casimir JH Ludwig, Rui Ponte Costa, Gaurav Malhotra, and Jeffrey Bow- ers. The role of disentanglement in generalisation. In International Conference on Learning Representations, 2020.
Ramon Nogueira, Chris C. Rodgers, Randy M. Bruno, and Stefano Fusi. The geometry of cortical representations of touch in rodents. Nature Neuroscience, 26(2):239-25 Jan- uary 2023. doi: 10.
URL Battps://cci.org http://dooo...//....o
s41593 022-0022-2222 01237-9
OS Oguejiofor, AN Aniedu, HC Ejiofor, and AU Okolibe. Trilateration based localization algorithm for wireless sensor network. International Journal of Science and Modern Engineering (IJISME), ((10):319-6338)) 2013.
Chris Olah. Neural networks, manifolds, and topology colah's blog, April 2014. URL https: 7/colah giubb. htt.:/.oooootttttoooooo /201:tt//////4-// -N-MMan fol Topolloooooosso
15

---

Published as a conference paper at ICLR 2025
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Mare Szafraanie, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaackdin El-Nouby, Mahmoud Assran, Nico- las Ballas, Wojciech Galuba, Russell Howes, Po- Yao Huang, Shang- Wen Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Hervé Jegou, Julien Mairal, Patrick Labatut, Ar- mand Joulin, and Piotr Bojanowski. Dinov2: Learning robust visual features without supervision, 2023. http:s///a......o//.a/a/////////
07193
Srdjan Ostojic and Stefano Fusi. Computational role of structure in neural activity and connectivity, Trends in Cognitive Sciences, 28(7))677-690, July 2024. ISSN 1364-6613, doi: 10. 1016/j.tiss. 2024.03.003. URL httL:t/:://xx doi wtt.:d.................. /tt.p//t.////p.o................ tics. 2024. 03. 003
E Page. Approximations the cumulative normal function and its inverse for use pocket calculator. Journal of the Royal Statistical Society Series C: Applied Statistics, 26(1):75 76, 1977.
Hieu Pham. Zihang Dai, Golnaz Ghiasi, Kenji Kawaguchi, Hanxiao Liu, Adams Wei Yu, Jiahui Yu, Yi-Ting Chen, Minh-Thang Luong, Yonghui Wu, Mingxing Tan, and Quoc V. Le. Combined scal- ing for zero-shot transfer learning, 2021.
10000
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Rajesh P. Rao and Dana H. Ballard, Predictive coding in the visual cortex: a functional interpre-
tation of some cxtra-classical receptive-field effects. Nature Neuroscience, 2(1):79-87, January 1999. ISSN 1546-1726. doi: 10. 10.1038/4580. URL | http:
doi, wtt/10.103//4/11116..0 4580
Rajesh P.N. Rao. An optimal estimation approach visual perception and learning. Vision Research, 19(1)):19331198, June 1999. ISSN 0042-6989.
doi: 10. 10.1016500022-699999900000111171
URL
http: doi. org 101..1016611 1ww11101600000 6989 (98) (0299)00C29999 0C279-X
Mattia Rigotti, Omri Barak, Melissa R. Warden, Xiao-Jing Wang, Nathaniel D. Daw, Earl K. Miller, and Stefano Fusi. The importance of mixed selectivity in complex cognitive tasks. Nature,
497(7451):585-590 May 2013. doi: 10. 10.1038/ WtLhhttp:://// htttps:://oi.org::: http///i...g.....//// I038/nature12160
Geoffrey Roeder, Luke Metz, and Durk Kingma. On linear identifiability of learned representa- tions. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Con- ference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 9030-9039. PMLR, 18-24 Jul 2021. URL Foeder21a html
ngg. mlr. wwwps//esssr..e.e 1t.es/..esssr..s//1111333
Jonathan E. Rubin, Catalina Vich, Matthew Clapp, Kendra Noneman, and Timothy Verstynen. The credit assignment problem in cortico-basal tanglia-thalamic networks: review, problem and possible solution. European Journal of Neuroscience, 53(7):2234 May 2020. ISSN 1460-9568. doi: 10.1111/ejn.1477
doi htt.:///o/./.at/////////ooooooo. /tt 14745
Ueli Rutishauser, Tyson Aflalo, Emily R. Rosario, Nader Pouratian, and Richard A. Andersen. Single-neuron representation of memory strength and recognition confidence in left human poste- January 2018. doi: 10. net.o...nro.0.0...1111 rior parietal contex. 000 neuron. wwt.uuronn.01..11. 2017.11.0100 029
URL
A. Saez, M. Rigotti, S. Ostojic, S. Fusi, and C.D. Salzman. Abstract context representations in primate amygdala and prefrontal cortex. Neuron, 87 August 2015. doi: 10. 1016/. ccuron.2015.07.o URL httts:: 101000...011111
Rui Shu, Yining Chen, Abhishek Kumar, Stefano Ermon, and Ben Poole. Weakly supervised disen- tanglement with guarantees, 2019. URL
Wayne W.M. Soo, Vishwa Goudar, and Xiao-Jing Wang. Training biologically plausible recurrent neural networks on cognitive tasks with long-tem dependencies. October 2023. doi: 1011 2023. 10. 10.561588. URL http:://xxx doi .org/ 1100..111112000 10. 561588
16

---

Published as a conference paper at ICLR 2025
Ben Sorscher, Gabriel C. Mel, Samuel A. Ocko, Lisa M. Giocomo, and Surya Ganguli. A unified theory for the computational and mechanistic origins of grid cells. Neuron, 111(0):111137..13.. January 2023. doi:
URL /ttp://p./......///1/..
neuron. 2022. 10.003
David Sussillo and Omri Barak. Opening the black box: Low-dimensional dynamics in high- dimensional recurrent neural networks. Neural Computation, 25(3):626 649, March 2013. doi: 10.1162/neco a 00409.
1ttgg//...162/eecc.. a_ 00409
Adly Templeton, Tom Conerly, Jonathan Marcus, Jack Lindsey, Trenton Bricken, Brian Chen, Adam Pearce, Craig Citro, Emmanuel Ameisen, Andy Jones, Hoagy Cunningham, Nicholas Turner, Callum McDougall, Monte MacDiarmid, C. Daniel Freeman, Theodore R. Sumers, Edward Rees, Joshua Batson, Adam Jermyn, Shan Carter, Chris Olah, and Tom Henighan. Scaling monosemanticity: Extracting interpretable features from claude sonnet. Trans- former Circuits Thread, 2024. htttp::/tttaattt reu
scaling monosenanticiiy//iniee htmml html
Loring W Tu. Differential geometry: connections, curvature, and characteristic classes, volume 275. Springer, 2017.
Pantelis Vafidis, David Owald, Tiziano D' A'Albs,, and Richard Kempter. Learning accurate path integration in ring attractor models of the head direction system. eLife, 11:c69841, jun 2022. ISSN 2050-084X. doi: 10.7554/eLiffe.6841..
69841
William E. Vinje and Jack L. Gallant. Sparse coding and decorrelation in primary visual cortex during natural vision. Science, 28765456) February 2000. ISSN 1095-9203. doi: 2..126lsienne22275555.....
URL http: doi. 11t.g///.....///e... htt.://ccincce..oo 287. 287.
5456.
Gary R Waissi and Donald Rossin. A sigmoid approximation of standard normal integral. Applied Mathematics and Computation, 77(1)911-95, 1996.
Xiao-Jing Wang- Synaptic reverberation underlying mnemonic persistent activity. Trends in Neurosciences, 24(8):455 463, August 2001. doi: 10.1016/s0166- 2236(00))1868)
URL
http::://doo.oooo
00116//01666 2236 (00) (00)01868-3
James R. Whittington, Will Dorrell, Surya Ganguli, and Timothy J. Behrens. Disentanglement with biological constraints: A theory of functional cell types, 2022. URL https://arxiv:
oxg/abs/2210. 01768
G. R. Yang and X.-J. Wang. Artificial neural networks for neuroscientists: primer. Neuron, 107 (6):1048-1070, September 2020. doi:
URL https://doi.
bxq/10.1016/1 neuron. 2020 00.
Guangyu Robert Yang, Madhura R. Joglekar, H. Francis Song, William T. Newsome, and Xiao-Jing Wang. Task representations in neural networks trained perform many cognitive tasks. Nature
Neuroscience, 22(2):297-306, Feb 2019. ISSN 1546-1726 1726. doi: 10. 10.1088844159--08--0010
URL PRLLhttts:://o...oo /t/.////////////oooo///////////0
Yu Zhang and Qiang Yang. overview of multi-task learning. National Science Review, 5(1):
30-43, 09 2017. ISSN 2095-5138, 1093 wt/33//r/nwxxll
doi: 10. 10.10933ns//wxx105.
htth:t::: http:::/dooo.o//////
17

---

Published as a conference paper at ICLR 2025
A SUPPLEMENTARY MATERIALS
A. DEFINITIONS
In the main text we used the terms "abstractness" and "disentanglement"" Since there exists some ambivalence about their meaning in the literature, we would like to strictly define them here, We will be using definitions adapted from Ostojic Fusi (024):
An abstract representation of latent factors ,1,...,37 represents each æi linearly and ap- proximately mutually orthogonally. Thus, abstractness ensures a simple linear map can decode each z; regardless of variation in 2zzs.
Disentangled representations of ; encode each æ; orthogonally, without the ne- cessity of linearity.
Note that under this definition, axis-alignment not requirement for disentanglement (also see Higgins et al. (2018)) Our work suggests that the computer science and neuroscience communities should adopt this broader definition of disentanglement, because otherwise we might be missing cases where the factors are not axis-aligned, but they still orthogonal and can still be isolated by linear decoder. Our argument is that there is nothing special about individual factors being encoded by individual neurons. Rather, think that allowing for mixed representations within the definition of disentanglement leads to more holistic view of disentanglement. A contribution of this work, along with others (Johnston & Fusi 2023). is bring this argument the forefront.
A.2 QUANTIFICATION GENERALIZATION PERFORMANCE
To assess OOD generalization performance, we keep the trained networks fixed and train linear decoder A predict the ground truth x" from network activity the end the trial. We train the decoder in 3 out of 4 quadrants and test OOD in the remaining quadrant, repeating this process 5 times for each quadrant, which results total of 20 OOD ,2 values for each network. To account for randomness in initialization and sythetic generation of datasets, combination of number of tasks Nusk and dimensionality D, resulting in a total of 100 OOD ,' values for each pair of (Nask, D). We report the 25, 50 (median) and 75 percentiles of those values in Figure b and throughout the text. For ID) generalization performance, we train on all quadrants and test in quadrant time. For input dimensionality D we keep the same logic by choosing every 4-th quadrant to be sampled only in testing, repeating the process for every mod 4 group of quadrants.
train 5 networks for each
A.3 ESTIMATION OF ANGLES BETWEEN LATENT FACTORS
To estimate the angles between latent factors in the representation, we obtain the normal vectors of the decoders A for each of the latents, and compute pairwise angles for all of them. To account for variability in the decoder fits, we repeat the decoder fit 5 times for each out-of-distribution region (see Appendix A.2 for details). We also repeat this process across 5 trained networks for each combination of task D), and report the 25, 50 (median) and 75 percentiles of all values for each (Nusk, D) combination in Figure 3D and Figure S12
A.4 DERIVATION OF THEORETICAL ? FOR OPTIMAL MULTI-TASE CLASSIFIERS
Here derive the theoretical 222 for the estimation of ground truth x" from noisy data for discrete time optimal multi-task classifier at time t. r² is defined as:
2 MSE(x*, ,1)
Var(x*)
where the mean of X(1). X(t), is the prediction of the multi-task classifier (see Ap- pendix B The optimal estimator of x" given observations X(1). X(t) is denoted X(t) N(u(t), ID) where a is the noise strength. Note that put) + x" as t - oo by the central limit theorem, and pll) is the optimal estimator of x" given Gaussian-noised observations. Since
18

---

Published as a conference paper at ICLR 2025
the dimensions in both noise and ground truth are independent, can focus dimension at time i.e.:
2 MSE(2, (4(() =1 [x o +N(00 =1
o²
Var(z;)
Var(x;)
(Vr((r;)
Remembering that x Uhiforrm--05.. 0.5) it follows that Var(x;) 30.5 and replacing or for given trial duration 7 which we compare RNN
0.2 from Table S1 we arrive to x²
OOD generalization performance in Figure
A.5 MORE NONLINEAR ENCODING
In the main text we used ReLU nonlinearities for the encoder f. Here extend our findings more nonlinear observation maps
stribuiot
0.9
which are likely to be encountered in the real world, e.g. ones with
0.8
power-law nonlinearities. Recent work showed that the choice of activation function can influence the geometry of representations Alllman al., 2024). Interestingly, find that replacing ReLUs in the encoder with a quadratic nonlinearity results in virtually iden- tical OOD generalization performance compared Figure (Fig- ure (ST. Therefore we conclude that our setting is robust to the choice of encoder nonlinearity, even when the nonlinearity not injective, going beyond our theoretical proofs (Appendix B.3
0.7
Nonlinearity
ReLU
Quadratic
of tasks
Figure S1: OOD generaliza- tion is robust to choice of en- coder nonlinearity.
A.6 GPT-2 EXPERIMENTS
We train GPT-2 causal transformers with dmodel = Nee 64, Naver = 1, Nhedd = in the multi-task classification task of the main text. The networks receive continuous, noisy and and should output multi-task classifications
non-linearly mapped inputs yxx)) E Nm... The output of the network is Y(t) := (Z((t), where Z(t) is the last em- bedding of the sequence in the last layer and g sigmoid. Since the imput is continuous, we omit the tokenization and embedding steps, and project the input directly to the hidden state with linear map. Furthermore, since the inputs are IID, we do not include positional encodings. The networks trained with binary cross-entropy loss for Noaroh 2 10 batches, while the rest of the param- eters are identical to the fixed RT networks of the main text (Table
A.7 FINDING FIXED POINTS AND LINEARIZATION OF DYNAMICS
To find approximate fixed points of RNN dynamics after training, we follow a standard procedure outlined in [ussillo & Barak 2013). Specifically, we keep network weights fixed, provide no inputs the network. and instead optimize over hidden activity. Specifically, we penalize any changes in the hidden activity, motivating the network to find stable states of the dynamics in the absence of input, i.e. attractors of the dynamics. This process finds all states of accumulated evidence that can be stored in this network as short-term memory. Network dynamics could then leverage these states to maintain and update the internal representation of the ground truth x on a single trial level, and drive downstream decisions.
Then for every approximate fixed point z/, we linearize RNN dynamies around and estimate the eigenmodes which describe how the system behaves in a small region óz around z', Specifically, following Sussillo & Barak (2013); Mante et al. (2013) we take the difference system öz(t + to) = z(t to) z (ft) and linearize it, i.c.
Oz = Fzz))6 8z
where à = F(z) is the function describing the RNN dynamics and M = Fz'z is its Jacobian at zz. To estimate F'(z²), we let network dynamics run in the absence of inputs for one time step At starting from zf, i.c. ōz(At) = z z(At) z! and autodifferentiate óz( (t).... We then perform eigendecompositionn of M and report the eigenvalues around each approximate fixed point.
19

---

Published as a conference paper at ICLR 2025
Eigenvalues near 0 indicate that the difference system öz(t) z(t) z changes slowly time, ie. they correspond to "slow" dimensions in network dynamics which can integrate inputs and maintain them over time (continuous attractors) (Amari 1977: Mante et al. 2013).
A.8 STATE-SPACE EFFICIENCY OF CONTEXI-DEPENDENT COMPUTATION
The workhorse model for computational neuroscience has traditionally been context-dependent computation, where tasks are carried out one at a time and task identity is cued the RNN by a one-hot vector Mante et al. (2013). However, this approach can be algorithmically inefficient, because as we show here scales linearly with the number of tasks Nasks and exponentially with input dimensionality D. This is because context- context-dependent computation utilizes different parts of the state space for different tasks, and the resulting representations collapse to what is minimally required for each task (also see (Mante et al. 2013 Yang & Wang 2020 This can be detrimental computation within a large yet limited neural substrate. In for brains, which need to pack a general, compact (Ma et al. 2022). can be used for any down- contrast, abstract representations stream task involving the same variables, scale linearly with D, and as we show readily emerge from relatively simple tasks.
To compare context- context-dependee dependent decision making, where one task is performed at a time (Mante et al,, 2013), to multitasking, in terms of state-space usage efficiency, dependent decisions on the same tasks encountered in the main text. Compared to the network the main text (Figure(b) the RNN now also receives one-hot task rule vector indicating the current task, and it outputs the decision for that task only (iggure[S2p,, We have also omitted the non-linear encoder, making the tasks easier. We train the RNN for two tasks, one task a time, in interleaved
train RNNs perform context-
batches (Figure S2h). In one task, the RNN is required to decide which stream has more evidence, and the other whether the sum of evidence across streams exceeds a certain decision threshold (here 0).
a
Tasks
Taek
0
-0.5
0.0
b
Heebon
Xi
Exidnce Extoares
6.0-
1
PG3
PG1
Time
Figure S2: Representations in RNN trained in context-dependent decision making. (a) We trained RNNs for two classification tasks: two-alternative forced choice (where the decision boundary is the .2 line) and evidence integration (corresponding the .1 -22 line). Fach task corresponds to different one-hot task rule vector. (b) Network architecture. In addition to the inputs Figure Zp. the network also one-hot vector indicating the current task. (c). Top PCs of RNN activity example trials (40 in total). The task rule biases the network towards learning separate solutions in different parts of the state space for different tasks, in the form of separate line attractors; red x's for wwo-alternative forced choice and orange x's for evidence integration.
We find that in this setting the network is not incentivized learn abstract representations. Instead a separate line attractor is present in the dynamics for each task (red and orange x's in Figure S22;; one of them is presumably tracking the difference of evidence (similar Yang & Wang (0200 but
20

---

Published as a conference paper at ICLR 2025
for independent evidence streams) and the other the sum of evidence. That is to say, the task rule biases the network to learn different computations in separate regions of state space, as in Mante] et al. 2013) As result, the 2D latent space has collapsed and cannot be decoded from network activity; therefore generalization to any task that involves these two variables is not possible.
It follows that the network can be inefficient in terms of state space usage, because instead of com- pressing all of its activity around the same region, it spreads it across multiple regions, one for each task, which scales badly (linearly with the number of tasks Nua and exponentially with input di- mensionality D). To demonstrate the latter, imagine a family of tasks with classífication boundaries of the form + 2 ① 22 @ .. æD = 0, where @ {+1. 11,1.-1.00 }} is an operator indicating contribution with a positive sign, negative sign or absence of contribution for factor to a specific task, respec- tively. As just shown, each one of this tasks will require its own line attractor, resulting in total of 3D line attractors lying in separate regions of the state space, just for this simple family of tasks. As mentioned in the main text, such inefficiency can be detrimental for brains, which need to pack a lot of computation within a large yet limited neural substrate. Compare that to multitasking, which builds representations that can serve any task that involves the same latent variables, scaling linearly with D (as saw that we only need Nusk D to learn them). Note that context-depeneen compu- tation can still be efficient, if tasks have a compositional structure where the solution for one task is part of the solution for another (Yang et al. 2019 Driscoll et al.] 2022); in this case, representations developed for the former can act as a scaffold for representations for the latter.
Overall, we believe that multitasking may present a paradigm swift for generalizable representation learning in biological and artificial systems alike. That is not to say that context-depended rep- resentations are not useful; they are great leveraging the compositional structure of tasks (Yang) et al. 2019; Driscoll et al., 2022), but tend to overfit to the specifics of the task, while multitask representations serve as world models applicable to various scenarios. Both types of representations are likely to be found in the brain. One possibility is that context coneext-dependent representations may emerge as a first quick solution to task, while disentangled representations come about with more experience or when more tasks are performed over time to support better generalization.
A.9 ROBUSTNESS TO OTHER NOISE DISTRIBUTIONS AND CORRELATED INPUTS
We here show that our setting is robust to Gaussian anisotropic and autocorrelated noise, and other asymmetric distributions of noise (Gumbel) whose CDF no longer matches the sigmoid functions in shape, with almost no drop in performance, and correlated inputs. This demonstrates that abstract representations are also learned outside of the specific assumptions made by our theory.
Starting from anisotropic noise, we observe that doubling the standard deviation of noise across one dimension (a = 0.4) does not result in reduction in OOD generalization performance (median r² = 0.96 for No 24, D- = 2). This is in line with our theory that can be extended to cover anisotropic noise (Lemma (. IT.. Non-IID noise should not be a problem either, since we are training network for many samples and the effects of correlations will cancel out over long ensembles. Indeed, we find that including autocorrelated AR(1) noise with an AR coefficient of 0.7 results in only minor reduction in performance (median 2 0.95 for Namt 24, D 2).
We were also curious to see the impact of correlated inputs. A problem with high correlations is that they render parts of the state space virtually invisible to the network (Figure(S3p) Surprisingly, OOD generalization performance is very weakly affected by input correlations, even though the state space sampled uniformly in test (Figure S3p). The behavior is highly non-linear: performance is great until 0.97, but for perfectly correlated inputs (p- 1), the performance drop is sharp.
Finally, our theory pointed out sigmoid functions as a choice for activation function because of their close resemblance to the Gaussian CDF, resulting in the best OOD r². Still we find that for an asymmetric noise distribution (Gumbel) whose CDF does not match sigmoid functions well, there is only a slight drop in performance (median ? = 0.95 from 0.96). Therefore the conditions for the activation function/CDF should be quite lax; any monotonic bijective function should work with small performance drop. This drop in performance is because the representation would be "stretched out" and "compressed" in a non- non-linear manner in regions where there is discrepancy between the noise CDF and the activation function. But this nonlinear squishing (determined by the term & -1 (g(( (t))) would be geometrically inoffensive - no cutting or gluing together would be required to map from (t) to linear representation of (()). As result, the representations would
21

---

Published as a conference paper at ICLR 2025
b
Correlation coefficient 0.97
0.5
Out--f-distrrbutton
0.9
0.7
0.6
-0.5
-0.5
0.97 Factor correlation
0.7
Figure S3: Disentanglement and factor correlations. (a) We introduce strong correlations in the latent factors, rendering parts of the state space virtually invisible the network during pre- training (trained for total of 24 classification tasks). (b) Despite that, generalization performance excellent for correlations very close to Once the factors perfectly correlated, performance drops significantly. This implies that the network can learn an abstract representation from correlated inputs, long as there is some signal about the factors independently. This finding goes beyond (Johnston & Fusi, 2023) to show that the multi-task learning setting allows OOD generalization when the distribution during training the RNN itself vastly different that the one during testing.
remain approximately linearly decodable. Monotonicity and bijectivity are quite mild assumptions for the activation function used by neurons in the brain.
A.10 NONLINEAR CLASSIFICATION BOUNDARIES AND INTERLEAVED LEARNING
In the main text trained networks linear classification boundaries. The tasks still non- linear, since the encoder renders these boundaries non-linear to the network. However, there are cases where the latent factors themselves might need to be combined non-linearly, to make decisions. For instance, if the two factors represent the amount and probability of reward respectively, an agent needs to multiply the two and decide whether the expected value exceeds certain (metabolic) cost of performing an action to obtain said reward. Figure S44 shows the classification lines for the multiplicative task, where the network should decide whether the ground truth lies above below the curve because they violate the injectivity condition. Hence, we wondered how the representation would look like if the network was trained on both the linear and multiplicative boundaries, as animals do.
= % for multiple values of This family of tasks is not covered by Theorem B.6
b
Classification lines
0.5
0.0
-0.5
-0.5
0.0
0.
PC3
X1
PC1
Figure S4: Interleaved learning of linear and non-linear boundaries. (a) Classification lines for the multiplicative task. There a total of 48 classification lines, 12 per quadrant. (b) The network learns an abstract representation when trained for the linear and multiplicative boundaries in interleaved batches.
22

---

Published as a conference paper at ICLR 2025
For that we perform interleaved training of both tasks (i.e. train in batches sampled from one of the tasks at time), a setting where neural networks excel at, compared to humans who excel at blocked training, where tasks are learned sequentially (but see|Flesch et al. (2022)) ). Figure (S4 shows that the network still learns an abstract, two-dimensional continuous attractor. OOD generalization for this network is excellent, and almost identical to ID performance (median 2 0.94, 0.97 respectively). Overall, we conclude that our framework extends to interleaved learning of mixture of linear and nonlinear boundaries, which better reflects the challenges encountered by agents in the real world. Note that during interleaved training. linear and non-linear tasks are not performed simultaneously; yet they are in immediate succession which can also place pressure to the network to gradually learn representations that satisfy all tasks. The relation between multi-task and interleaved learning is promising topie for future research.
A.11 ABSTRACT REPRESENTATIONS ARE LEARNED FOR A FREE REACTION TIME, INTEGRATE TO BOUND TASK
In the main text we trained networks to produce a response at the end of the trial. However, in many situations agents are free to make a decision whenever they are certain enough. Therefore, we here seek to extend our framework free reaction time (RT) decisions. A canonical model accounting for choices and reaction times in humans and animals is the drift-diffusion model (Krajbich et al., 2010 Brunton et al.) 2013). is composed of an accumulator that integrates noisy evidence over time, until certain amount of certainty, represented by bound, is reached, triggering decision. In the linear classification task setting, the accumulated amount of evidence at time / for line with slope Aa (t) is given by:
() X2(t)
(9)
a
A.1t) X:00 oXxi)
Excener 1
Rood
Free
Timm
stral
Figure S5: Free reaction time task. (a) Data generating process. Every classification line from Figure 2 now corresponds to an accumulator (see corresponding colors), and the desired output for the RNN is the accumulator values for the entire trial. The accumulator is quantized to integer values between +5. (b) Representation for RNN trained on free reaction time task. The network learns a two-dimensional continuous attractor, similar to Figure 31. A 3D rotating figure to better visualize this representation is provided in the Supplementary Material. (c) OOD generalization performance for the free reaction time (RT) task. Free RT outperforms fixed RT for a small number of tasks.
Intuitively, A. (t) reflects the amount of confidence at time & that the ground truth x* lies above or below the classification line with slope a. Essentially, the network has to explicitly report distance from the classification lines. not just in which side of the line x" lies for that trial. We set the decision bound to t5, and plot the accumulators Aq for all lines in Figure 2.. Note that once the bound is reached a decision is effectively made and A, is kept constant. Also, instead of using continuous values, we quantize Aa, because going to used as target signal to train the network, and we do not want to introduce strong inductive bias towards integrating the evidence streams.
We then train the RNN to reproduce confidence estimates from Figure for the entire trial. Com- pared to previous experiments, the fixation input is no longer available to determine when to produce a decision. Instead, decisions evolve dynamically throughout the trial. We also use a MSE loss, change the activation function to g 5 tanh, and the Adam learning rate nu = * 10 but all other parameters remain the same as in the main text. To have a closer correspondence to the free RT experiments here, also train the fixed RT task from the main text with MSE loss, symmet-
23

---

Published as a conference paper at ICLR 2025
ric labels yxx) {-1,+1) Now and output non-linearity tanh. We find that the change of objective and loss only has minor effects on generalization performance.
Figure SSp shows that in this setting the network still learns two-dimensional continuous attrac- tor of the latent space. Furthermore, the free RT outperforms the fixed RT network from the main text (Figure (5p) for small number of tasks, since it is explicitly required to report distance from the classification lines. However, as our theory shows (Lemma [B3)) the fixed RT network is also implicitly reporting distance from the boundaries, when behaving like an optimal multi-task classi- fier, which explains the similar performance for a larger number of tasks. Overall, we showed that our setting accounts for naturalistic free RT decisions, and provides theoretical justification for the importance of confidence signals in the brain (Rutishauser et al.) 2018 Masset et al. 2020).
The importance of the confidence (i.e., calibrated likelihoods) of a network's output, is a recurring theme in machine learning too (e.g. knowledge distillation (Bhargava et al. 2024)). We here show that confidence fundamentally connects how neural networks construct world models, either directly Lintegrate-t--bumn task here) or indirectly (classification tasks in main text). Under this framework, knowledge distillation can be cast as smaller models directly copying the world models (logits) of larger ones.
A.12 QUANTIFICATION OF SPARSITY
In the main text, we observed that RNN representations are sparse. We here seek to more precisely quantify the sparsity in these networks, and investigate how it is affected by the number of tasks
Nusts latent dimensionality D, and specific recurrent architecture. To do so, we sample n 1000 ground truth vectors x" randomly for every network, and compute the sparseness (Vinje & Gallant 2000) of a neuron in the hidden layer as:
Eleffn) ({)/m))
100 %
(10)
1
where is the steady-state response of the neuron ground-truth stimulus i. Sparseness ranges from 0 to 100 %, with greater sparseness indicating greater selectivity of the neuron to stimuli. Then, the sparsity of a network is given as the average of the sparseness of all its neurons.
a
b
100
100
Sparsit (%)
80
Spersity (%)
80
60
60
40
40
20
Non-leaky RNN
20
LSTM
0
10
30
10
Inpat dimensi onality
# of taste
Figure S6: Quantification of sparsity choice. (a) Sparsity of a recurrett network as a function of number tasks and network archi- tecture. Five networks trained for each network configuration. Greater levels of sparsity indicate that the network activations are more sparse. (b) Sparsity of RNNs as a function of number of tasks and latent dimensionality D. Five network are trained for each combination of (N D).
function of Nusk, D and recurrent architecture
Figure S6 shows that RNNs and non-leaky RNNs are very sparse, with sparsity values around 90 for different values of Nust supporting the claim the main text. LSTMs on the other hand, which are less brain-likd have lower sparsity values, although interestingly sparsity increases with
'LSTMs architecturally enforce intricate, high- cageccap multiplicative gating mechanisms, while in biologi- cal neural networks gating has to he learned. For other aspects of biological implausibility of LSSS compared RNNs, see Appendix B in (Soo al. 2023)
24

---

Published as a conference paper at ICLR 2025
Notably, we did not do anything to promote sparsity (e.g. regularization) in these networks. Therefore we conclude that sparsity naturally emerges from the optimization objective of multitask learning, particularly in architectures that are more brain-like.
Nust-
Next we wondered how latent space dimensionality D would affect sparsity in our trained RNNs. Figure shows that networks remain very sparse for the whole range of dimensionality D tested in the main text, with sparsity values above 75 %. Greater dimensionality results in less sparsity on average, which is expected since Nnew 64 in our networks, therefore a significant amount of their capacity must be used as D increases. This effect plays in only as Nust increases, as networks will only learn disentangle the input dimensions that are spanned by the tasks, as our theory predicts. Overall, there seems to be a proportional relationship between the number of active neurons and dimensionality D, as long as there are enough tasks to uncover the D latents.
A.13 RELATION TO NEUROSCIENCE LITERATURE
An ongoing debate in the brain sciences is whether to solve tasks the brain learns abstracts represen- tations, or simple input-output mappings. Here we show that training RNNs to multitask results in shared, disentangled representations of the latent variables, in the form of continuous attractors. In this multitask setting, one task acts as a regularizer for the others, by not letting the representation collapse, or overfit, to specific tasks (Zhang & Yang 2017).
Our findings directly link to two important neuroscientific findings: spatial cognition and value- based decision-making. First, the tasks here bear close resemblance to path-integration, i.e. the ability of animals to navigate space only relying on their proprioceptive sense of linear and angular velocity (Mittelstaedt & Mittelstaedi 1980 Burak & Ficte| 2009 Vafidis et al.) 2022 Sorscher ett al. 2023). In path-integration animals integrate velocity signals to get location, while here we integrate noisy evidence to get rid of the noise. In path-integration, networks have to explicitly report
distances, while in our setting distances are estimated implicitly (Lemma B.3) ). We learn abstract representations in the form of 2D "sheet" continuous attractor, while the computational substrate for path integration is 2D toroidal attractor (Gardner et al. 2022; Sorscher et al. 2023) not an abstract representation. The conditions under which a 2D sheet vs. toroidal continuous attractor is learned is a potential area of future research. Second, decision making experiments in monkeys result in a 2D abstract representation decisions (Bongioanni et al. 2021). Likewise, context-dependent decision-making experiments in humans also resulted in orthogonal, abstract representations (Flesch et al. 2022).
the medial frontal cortex, which supports novel inferential
A. 14 BIOLOGICAL PLAUSIBILITY OF MULTI-TASK LEARNING
While our theory stems from parallel processing, i.e. multi-task learning, it is not contingent upon the parallel execution of multiple tasks, i.c. multitasking, or the receipt of rich supervisory feedback from the environment 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 the 2D latent 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. In other words, our theory assumee competence at Ni tasks, independently of when and how that competence was achieved. We feel that this is 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); after all, biological agents are remarkable because they achieve high performance on many tasks. This theory is also closely related to the widely observed phenomenon of memory replay (Foster& & Wilson, 2006), or mental simulation of counterfactuals (Jensen et al. 2024). A future direction to further enhance the biological relevance of our work would be to investigate the relation between multi-task learning and slow, interleaved learning (see Appendix A. in a continual learning setting.
25

---

Published as a conference paper at ICLR 2025
A.15 LIMITATIONS AND FUTURE DIRECTIONS
A limitation of the present work is that factorization is assumed. Yet not all problems are factoriz- able, or should be factorized. For instance, more coarse-grained understanding of the world, that doesn't disentangle all factors, might be more suitable in many cases, and that might be reflected in the nature of the tasks. Furthermore, we focus on canonical cognitive neuroscience tasks which are somewhat removed from standard ML benchmarks. Normally, disentanglement methods would be tested against a benchmark such as dSprites (Matthey et al. 2017); however to the best of our knowledge no such benchmark exists for sequential tasks where evidence has to be aggregated over time. Future work could endeavor to apply our setting to richer tasks, like extracting latent item attributes from item embeddings when sequential decisions are made in online retailer settings.
Our theory agnostic to the way by which competence at multiple tasks is achieved. Thus, a natu- ral next step is investigate whether disentangled representations exist in wider range of models capable solving multiple tasks. A prime example is large language models that display excellent zero- and few-shot generalization capabilities, with progress already made in that direction (Tem- pleton et al. 2024). Moreover, the pre-training objective for LLMs (cross entropy loss/likelihood maximization) fits well within theoretical framing (approximately) optimal multi-classifiers. Another application area, as already mentioned, is neuroscience; animals are naturally competent at multiple tasks, thus our work provides theoretical justification for why disentangled representations have been found in many brain areas, and motivates looking for more.
Our experiments showed parsimony of our theoretical results under conditions not covered by our theory, including non-injective observation maps (Appendix (A.10) and decision boundaries (Ap- pendix A.1T which is encouraging for testing our findings on settings beyond what is strictly cov- ered by the theory. It would be interesting to see how the theoretical insights generalize to different task geometries, for example those implied by supervupe learning applications (e.g. image patch-filling, next-token prediction, iterative de-noising). The connection between our framework and self-supervised learning deep and promising. Both frameworks share a common structure, where underlying latent truth (e.g., objects in an image and their relationships) is inferred. Each objective (c.g., filling in missing image patches) contributes synergistically to understanding the la- tent space whole. A similar logic applies predicting words, where the latent "meaning" of sentence is shared, whether a causal (e.g.. LLMs) or masked setting (e.g., BERT). Our study is first effort towards understanding such parallel learning, and providing guarantees for its perfor- mance.
26

---

Published as a conference paper at ICLR 2025
B THEORETICAL DERIVATIONS
Here we prove our main theoretical result outlined Section
High-level summary of proof We prove that competence at Nask tasks guarantees linear decod- ability when Nusk 2 D for non-degenerate tasks (Theorem B.5). and orthogonality when Nusk >> D and task boundaries are sampled randomly (Corollary [B.T0p. T0..... To that end, we first show that optimal evidence aggregation in multi-task classification framework enforces the multi-task classifier to encode a notion of distances from classification lines (Lemma B.3). Given a suitable set distances from classification lines, we show that can uniquely identify optimal estimate of the latents given noisy data in closed form (Trilateration Theorem, Theorem B.4). In addition, if the readout function of the multi-task classifier is sigmoid-like, the optimal estimate will be approximately lin- early decodable from the representation (Corrolary B.8 B.9). We then prove by contradiction that all the above results hold when arbitrary injective observation map f is applied to the input after noising (Theorem B.6 The theory generalizes to sub-optimal classifiers via least-mean squares approximation and the Moore-Penrose Pseudoinverse (Theorem B.7]. and different noise distri- butions (Section B.4). Finally, we discuss implications of the theorem for representation learning, manifold learning and the Platonic representation hypothesis, and future directions (Section [B6).
Notation:
lower case variables denote scalars (e.g., x), upper case variables denote random variables (e.g.. X), and boldfaced variables denote vector quantities (e.g, x, X). We denote the D x D identity matrix as D.
Multi-task
Y(1)
{()
f(3)
Y(4)
Y(1)
classification
g]
8
g
8
8
estimates
Z(1)
Z(2)
Z(3)
2(4)
Z(t)
E.g. RNN latent state
fl
fl
f
f
fl
X(1)
X(2)
X(3)
X(4)
...
X(t)
+ i.i.d. isotropic Gaussian noise
Ground truth evidence
Ground truth classifications
Figure S7: Bayesian graphical model framework representing our theoretical framework for multi- task classification. The agent with latent state Z(t) estimates the ground truth decision output y(x*) E {0, N from noisy observations X(t) transformed by injective observation map f. We prove that latent state Z(t) must encode an optimal, linearly decodable estimate of the de-noised environment state x* when the decision boundary normal vectors {c; Nus span R..
Variable Glossary:
x" E RRD: Ground truth (un-noised) input variable of dimension D.
x(t) + x² ooN(0, Ip) are i.i.d. noisy measurements of x*, where
- - is the amount of equivariant Gaussian noise, and
- is the discrete time index within a trial.
/ JRD - Z : An injective observation map that transforms the noisy measurements X(I)
before they reach the latent state Z(I) of the optimal estimator. The map is injective, meaning that it preserves the uniqueness of the input, ie. if f(x1) ) = f(x2), then X1 = X2. The codomain Z can be any suitable space, such as CM 1000 or other spaces.
Nusk is the number of classification tasks,
27

---

Published as a conference paper at ICLR 2025
{(o,, i are the classification boundary normal vectors and offsets respectively, with E RD and by E IR. We assume each ||c,|| = 1.
(C, b) are matrix and vector representing each of the Nust classification tasks where Ce RMaaa x
y(x") E {0, 13 Ground truth classification outputs, where each ground truth classifi- cation y: (xx)) is given by
(x(x) = f c x>
(11)
otherwise
Z() Latent variable of a multi-task classification model, conditional on X(t).
g Map from latent state (t) to multi-task classification estimates (t).. For most of our experiments, readout map g = sigmoid, for instance.
(() (t))) E [0, 1] Nout Output vector of the multi-task classification model time t, where each (t) is a Bernoulli random variable estimator, estimating the conditional probability Pr (x")
}} given the noisy observations (via latent variable Z(l) see
Equation 12). 12).
((t) = (tttt), E() Optimal estimate of x" given measurements X(), , X(t), derived in Lemma
Problem Statement:
We consider optimal estimators of y(x*) in the multi-classification shown graphically Figure SS7
paradigm Equation
x + X1)..... X(t) 4 Z(t) 4 (t)
(12)
Contribution: We prove in Theorem B.6 ("Optimal Representation Theorem") that any optimal estimator of (x") described above will represent an optimal estimate of x in latent state Z(t). We begin by proving results optimal estimators in Sections B. B.2 with identity observation map f. developing the linear case of the optimal representation theorem (Theorem B.5) showing that the latent state Z(t) must encode an estimate of x* (visualized in Figure (8)). We generalize this result any injective observation map ff in Section [n33and derive closed-form solutions for extracting the estimate of x" from (t)). We derive approximation results for g taah in Corollary [B.8)
and g sigmoid in Corollary B.9 that show the representation of x" in Ztt) will be linear-affine decodable if g the sigmoid Tamily of functions.
B. SINGLE DECISION BOUNDARY
First, will derive (t) for single decision boundary with parameters (c. b). We focus
(()) X()),, X()), reintroducing the latent variable Z(t) later on.
Since v(x*) is a deterministic function of non-random variable x', we will derive the probability distribution over Pix* (X(),..... .())) denoted - determine y )(((()))).
Lemma B.1. Assuming no prior on x", the conditional probability distribution (()
P(x*|X(11 (X(1, X(I)) is given by
..()..
X(t) NN((tt.
(13)
where jut)= = mean ((1),.... X(t) and E(t) - o0o
Proof. Since X(1), X(t) are i.i.d. from a Gaussian distribution with mean x" and identity covariance, the sample mean is known to be distributed normally centered at the ground truth x". We apply the known standard deviation of the underlying distribution (identity covariance scaled by σ) arrive at (t) -² as the variance on the sample mean (derived from the central limit theorem).
'Noee that the intermediate computation x() does not imply that system must compute this value predict Y, the full computation of x(() may not be necessary to determine Y()).
28

---

Published as a conference paper at ICLR 2025
Ground truth evidence
Gaussian noisy
RNN latent state
Classification Estimates
observations
(X(1, X()) 1 (1)
Y(1)
()())
y(x** Ground truth classifications
M
X(O
X(O)
P200
Figure S8: An overview of the classification process using RNN with Gaussian noisy observa- tions. The ground truth x" generates the noisy observations (X(1), X(t))- These observations processed by the filter-based model illustrated graphically Figure 57 maintaining latent state Z(t). The latent state (t) is then used to produce classification outputs (t) and (t). The- orem B.6 proves that Z(t) must encode an estimate of x*, visualized in this figure, shown as X including its mean u(t), which is the optimal estimator for x* given the noisy observations.
We can use estimator X(t) to construct Y (t) by expanding Y = v(X(t)) via Equation
In essence, we are interested in the amount of the probability density of X that lies on each side of the decision boundary. Deriving this probability is simplified by the fact that X is isotropic i.e., it inherits the spherical covariance of the underlying data generation process (Lemma B.2).
Lemma B.2. X(t) = (m((t), E{t) with isotropic covariance E(t) oo and mean p(t) E
RDD The probability density of (() the positive side of the decision boundary (x x > b)
can be expressed as
((I) A Pr(cT x b)b
(14)
where t is the CDF of the normal distribution and k = c. cp(() - b is the signed projection distance between the decision boundary and the mean p(t) of X(t).
Proof. Since the X(t) is isotropic, the variance every axis is equal and independent. We may rotate our coordinate system such that the projection line between the plane and the mean of (() aligns with an axis we denote as "axis 0". The rest of the axes must be orthogonal to the plane. Since each component of an isotropic Gaussian is independent, the marginal distribution X on axis 0 is univariate Gaussian with variance 1-1 mean at distance l from the boundary. Equation 14 applies the normal distribution CDF 4 to determine the probability mass on the positive side of the boundary.
Observe that (1) in Equation 14 monotonically scales with the signed distance between hyperplane and (t) (CDFs are monotonic).
Lemma B.3. Knowledge of time and optimal classification estimate (t) is sufficient to determine the projection distance between μ(t) mean mee.(X()..... X(t)) and the decision boundary (c,b), b).
Proof. Recall Equation 14 from Lemma B.2 We may solve for projection distance k separating the decision boundary and the mean p(t) of observations X(1)... X(t) as
σ 4 (}()))
k
(15)
29

---

Published as a conference paper at ICLR 2025
Since & is the CDF of the normal distribution, and the normal distribution is not zero except at =oo, the inverse 4 is well-defined.
Note that non-zero noise is required for Lemma B.3 to hold, as zero noise would yield zero proba- bility mass on one side of each decision boundary, meaning that no distance information would be recoverable from Y (t) (and Equation (15) would lead 0 indeterminacy).
B.2 TRILATERATION VIA MULTIPLE DECISTON BOUNDARIES
To recap Section B.1 : We derived an optimal estimator of x" (denoted X(t) based on noisy
i..d. measurements (I)..... X(t) ~ Nxxx,, ,o"Ip) in Lemma B.1 In Lemma B.2 we derived the equation for Bernoulli variable estimator Y (1) to estimate a single classification output y(x*) based on the same noisy measurements via ((). Finally, we showed in Lemma B.3 that the uncertainty in (t) and the time t is sufficient to determine the projection distance between the decision boundary and p(t) = mean mean(X(1)).... (()) via Equation
We now have the tools to prove our final result via trilateration. Much like distance information from cell towers can be used to trilaterate one's 's position, we will leverage Lemma B.3 and use distances from decision boundaries {(c,, b.) Fie[NNa] to constrain the positions.
Let (t) denote the vector of classification estimates (() from Equation 15
Theorem B.4 (Trilateration Theorem). C E RN is full-rank and Naask 2 D, then (t), t, b, and C are sufficient to reconstruct the exact value of u(t), the mean of (1))........ X(),, which also the optimal estimator for x".
Proof. We may prove this claim by providing an algorithm to reconstruct μ(tt (1)),,.
from (t), C, and t. Invoke Lemma B.3 to compute the signed projec- tion distance between jut) and each decision plane (c,, b,). Let k [k....., KNeaee where each k; corresponds to decision boundary c, Then the mean (u(t) must satisfy
Cp(t)= k + b
(16)
Thus, for full rank C and Ntask 2 D, we will have uniquely determined (u() value.
Sufficient statistics and optimal estimators: "A statistic p(t) is called sufficient for x" it con- tains all the information in (I)....,, X(t) about x"." (from Cover and Thomas' Elements of Infor- mation Theory, 1999, Section 2. 10, substituting variable names).
More formally, "A function T(X(1).. X(t)) is said to be sufficient statistic relative to the family [of probability density functions indexed by x"] f(X(1), X(t) if
X(t)
independent of x" given T (X),),, X()), i.e. x* + TIX(1), ...., ... X(!)) X(I), X(1)
forms Markov chain. This is the same as the condition for equality in the data processing inequality,
[(x x(),, X(t) 1xx" m(t))
for all distributions on x*. Hence sufficient statisties preserve mutual information and conversely." (Cover and Thomas' Elements of Information Theory, 1999, Section 2.10, substituting variable names)
u(t)) = mean mean(X(1)..).... (X(1).... X(t)) is sufficient statistic for x given measurements X(i) ~ x" + (0 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. 20).
Trilateration differs from triangulation, and it is more frequently used practice. Triangulation is when one has angle information w.r.t the cell towers. Usually, this is not available so one trilaterates their position Oguejiofor et al. (013). This more closely matches our setting, where we just have distances information w.r.t. the decision boundaries and must determine the position.
30

---

Published as a conference paper at ICLR 2025
Theorem B.5 (Optimal Representation Theorem, Linear Case), Any system that optimally estimates classification probabilities Y (t) based on noisy measurements (X(1)..... X(t)) must implicitly
code a representation of pu(t) = mean mean(X(1)).... (X(1)....,,X)) in its latent state Z(t) if decision boundary
matrix C is full rank and N Ntask D.
Proof. We showed in Theorem B.4 (Trilateration Theorem) that C E & full-rank and Niask 2 D. then Y()), b, and are sufficient reconstruct the exact value of (t), the mean of X(1)...... .I))...... Rearranging Equation 16 and applying Equation 15,
σ p(t) = (CTC)-
-- bb Vt
Replacing (t) g(((t)) from our problem setup reveals that put) is deterministic function of Z(t).
Therefore, optimal multi-task classifier latent state Z(t) contains a sufficient statistic p(t) of x", which implies that (tt) must also contain all information about x' given noisy measurements ((1)....,(( if X full-rank and Ntas D.
Theorem B.5 boils down the observation that the confidence associated with each Y; in Y()) measures of distance between an implied estimate of x" (denoted put)) and classification boundary i (denoted (c., b))- Y specifies the position of X via "coordinates" defined by decision boundary normal vectors C
For sub-optimal estimators of Y, we may still obtain an understanding of the implied estimate X using the same methods. In fact, the machinery of least-squares estimation for Ax = b provides readily accessible formula for in sub-optimal estimators of Y (Equation 16) in the form of the Moore-Penrose pseudoinverse:
-c (k+
(17)
Conveniently, the estimation errors in sub-optimal Y have mean of zero, additional decision boundaries in C (e.g.. beyond the minimum D linearly independent boundaries) result in improved estimation of x by the central limit thcorem, thus generalizing our results to sub-optimal estimators (see Corollary B.7
B.3 OPTIMAL REPRESENTATION THEOREM (GENERAL CASE)
We extend the results from the linear case (Theorem B.5 to the general case where observations transformed by an injective observation map Theorem B.6
Theorem B.6 (Optimal Representation Theorem). Letx" E R be latent representation for linear binary classification task yx*) E {0, N
Nuse and X(t) (O,, In)) be noisy observations transformed by an injective observation map f.
Nauk X D is full-rank matrix representing the decision boundary normal vectors in RDD and Nask Z D. then any optimal estimator of y(x") must encode an optimal estimator μ(t) of the latent variable x" in its latent state Z(t). Furthermore, j(t) is sufficient statistic of x", ensuring that all the information about x" contained in (X(t) } also contained in Z(t). Consequently, j(t) the optimal estimate of x" based on (X(1).....((((( can be written deterministic
If C
function (Equation 18 of latent state (tt).
=(t)= = (C ( (o(((t) +b))) +b)
(18)
Proof. We use proof by contradiction to extend the linear case of the general representation the- orem account for injective observation maps f that map X(t) before they input to Z(t).
31

---

Published as a conference paper at ICLR 2025
Assume toward contradiction that there exists superior way of computing Y based injectively mapped (t) other than learning f and following the same procedure as when X(t) was fed
in directly (which derived the optimal estimator for in Lemma [.] and Lemma B.2). This sumption implies there is some additional information in (()(1)) that is not in ((),, violating the data processing inequality.
Formally, consider the following Markov chain:
x(t) 4 Z(t) + Y(t) μ(t).
x
(19)
Since injective, f-1 exists, making /(X(1)) X() an equivalent transformation in terms of
information content. Hence, any optimal estimator that processes f(X(()) can only perform as well as if had directly processed X(t).
To complete the proof, we show that uut) can be reconstructed from Z(t). Given the full-rank matrix C, we can use the same trilateration process in the linear case. The optimal estimate p(t) can be written as:
p(t) = (C
+b)
(20)
where g(((t) represents the transformation from the latent state to the classification probabilities.
Since (u(t) the optimal estimator (and sufficient statistic) for x* given mcasurements
{X()), ((t)). it contains all information about x" contained in the measurements (Cover &
Thomas (1911). In other words, p(t) a deterministic function of Z(t), implying that Z(t) will
contain all information about x' contained in the measurements {X(t)).
Corollary B.7 (Recovery of p(t) for Sub-Optimal Classifiers). Let (t) E [0, 1] represent the output of a sub-optimal classifier with zero-mean independent errors, i.e., Y;(l) = Pr(mx(x)) =
1} + 6G, where Ble,] and Varle,] = o for all E ;l,,...,NN Nask).
IfC E N is a full-rank and well-conditoneed matrix of decision boundary normal vectors with Niak 2 D, the estimated mean ju(t) xf can be recovered using the Moore-Penrond pseudoinverse:
ji(1) CC b))
where k = -/ (V((1((( and 0-1 is the inverse CDF of the standard normal distribution.
For sub-optimal classifiers, as the number of tasks Nusk increases:
The redundancy in C reduces sensitivity to classification errors.
Under the assumption of independent. zero-mean errors in Yt), the residual error in ju(t) is expected decrease rate of approximately O(1/ ((1///oa)) (Nua). driven by the averaging effect of least-squares estimation.
Motivated by the similarity between ((2) and sigmoid-like activation functions g(z), we show that the two be approximately canceled in Equation 18 implying that jut) be reconstructed with high accuracy with a linear-affine transformation (e.g.. linear decoding) when g = tanh or g sigmoid. This implies that Z(t) contains abstract representation of put) (Ostojic & Fusi 2024).
Corollary B.8. If the readout function is tanh, then the reconstruction equation for v(l) from Z(I) can be simplifted using the approximation ((2) x 1/2 tanh( =)
=) 3. Consequently, p(t) can be expressed directly in terms of Z(t) without the need for the inverse CDE.
2/30 TVE
m(1)
+ (CTC) CTb
(21)
Proof. Consider the readout function given by g(((t)) Yt) = tanh tanh this function allows for linear decoding of x" from Z(t),, we need to leverage the similarity between tanh and +.
2- To show that
32

---

Published as a conference paper at ICLR 2025
The normal distribution CDF (=) and the function taah(2) + known to be very similar, as both functions are sigmoid-like centered at zero, and asymptotically approach and (Choudhury)
2014).
Page 1977) proposed a simple approximation of & via tanh. Eliminating higher order terms, their / 22) + z.2 We found the following approximation yiclded a superior mean squared error:
approximation is (xx) = 2/ tanh(
1
((2) tanh
2/3
2
Using this approximation, express + in terms of tanh:
Φ- tanh(2) + 2 = 2v3
Substituting this approximation into the reconstruction equation for p(t) from Theorem B.6
p(t) (CCC)--C'' ( ( ()tanh(Z(t))+ 1) +)+b)) +b)
= <CTC)-''"
(vr (2/3 + +b)
TVi 2v30 (C (CC ((t) (C C)
Therefore, we have shown that ju(t) can be expressed as a linear transformation of Z(t) when the readout function g is sigmoid-like. This confirms the corollary.
Corollary B.9. For 9 = sigmoid, linear scaling by a = 0.5886 yields mean absolute error of 0.0038699 from (tt) in the range [-10, 10). enabling the following accurate linear-affine approxi- mation of j(t) from Ztt) given g = sigmoid:
Og Vi
m(()
(22)
Proof. We found a 0.5886 by computationally minimizing the mean squared eтror between 2(az2) and (2) (see s s wwt.pppx..pgguussii..ooaaasasss. ing a, on successively larger optimization bounds T = 10, 100, we found that d, converged to 0.5886. We note that sigmoid approximations to the Gaussian distribution CDF have existed in the literature for some time (Waissi & Rossin 1996).
.m in supporting code). Upon comput-
Corollary B.8 and B.9 are visualized in Figure S9 showing the close approximations to the Gaussian CDF.
Corollary B.10. Nuas >>>>> D implies orthogonal representations in latent (t).
Proof. Recall Equation from the disentangled representation theorem:
m(t) = (CT ( x (0'(((Z(t)) + +b)
For sigmoid-like g, we can approximate
mlt)
ag Vi
The orthogonality of the representations in Z(t) is therefore governed by the orthogonality of the in the matrix A (CT CCT R D If AA DxD is diagonal, then the rows of A are orthogonal.
AAT ((CTC)- ((C C)-
33

---

Published as a conference paper at ICLR 2025
Sigmoid and Tanh Approximations of Gaussian CDF
aa)
ttanh/4(2x)
0.8
0.6
0.4
0.2
0.0
-4
-2
Figure S9: Sigmoid (a(-) and tanh approximations of the normal distribution CDF & via horizontal scaling.
= =(CT))=((TC(((CC))
=((CTC)-1)
B B is symmetric matrix for any matrix Recall that the inverse of symmetric matrix is also symmetric. So (CTC) C)-1 is also symmetric. Therefore
AA T CC
As the columns of C are high-dimensional randomly sampled vectors, their probability of being non-orthogonal vanishes the dimensionality Nuask increases. We can also state the condition in terms of the singular value decomposition (SVD) of C = UEVT where U E RD Σ (d1,..., V E DxD, and o; is the ith singular value of C and U, V are orthonormal. Then
AAT VE P
the singular values are approximately uniform a ap R a then
AAT ~ V ID
AA
VVT
In
Therefore, uniform singular values C is a sufficient condition guarantee an orthogonal, disen- tangled representation of ju()) in Z(t)
D
As sidenote, would like point out that the setting here relates to the Marchenko-Pasee law while noting important caveats: while the law typically applies to matrices with i.i.d. entries N (0, .), C matrix consists of random vectors in RD of unit strictly meeting the law's conditions, still supports our conclusions about orthogonalization.
This structure, while
"This uniformity condition in the singular value decomposition is analogous to the outcome of the LM damping technique (Levenberg 1944 Marquardi 1963), used for least squares inversion problems in various applications.
34

---

Published as a conference paper at ICLR 2025
B.4 SUITABLE NOISE DISTRIBUTIONS
While the original proof leverages Gaussian noise due to its mathematical convenience, the key property required for the proof is more general. Specifically, the essential requirement that the marginal posterior distributions along the decision boundary normals C₁ have invertible cu- mulative distribution functions (CDFs), allowing us to recover the distances from the observed classification probabilities.
We will now provide precise mathematical description of the class of noise distributions where this key property holds, generalizing the disentangled representation theorem beyond Gaussian noise.
Invertibility of the Marginal Posterior CDFs Along Decision Boundary Normals. For each decision boundary normal vector ct the marginal posterior distribution of x* projected onto c₁ must have an invertible CDF. This allows to map the observed classification probabilities to unique distances between the estimated mean put) and the decision boundaries.
Key Noise Property Required for Proof:
Mathematical Description of Suitable Noise Distributions Let define class of noise dis- tributions <(t) where the key property holds and is straight forward to solve analytically.
Definition:
The proof is immediately generalizable to noise distribution c(t) if it satisfies the fol- lowing conditions:
1. Additive noise model:
X(t) = + ctt)
where (tt) i.i.d. random vectors.
2. Posterior Distribution Tractability: The posterior distribution ((x))[X(s)} must be analytically tractable or well-approximated, allowing us compute the posterior mean or maximum posteriori estimate ju(t) of x" and understand its properties.
3. Existence of Invertible Marginal Posterior CDFs: For each decision boundary normal vector C₁, the marginal posterior distribution of c, x" has continuous and strictly increas- ing CDF (k), which is invertible.
4. Support over X:: Let X' C R be the connected subset of allowable x* values. The noise distribution must have full support over X, ensuring that any real-valued x" is possible to trilaterate. For our proof, we assume support over the maximally permissible R is used.
Implications: expressed as
Any suitable noise distribution allows classification task probability {,()) to be
Yi(t) Pr(c. x" b([((s))
Pilb: u())
where F; is the marginal distribution of x* Since F; invertible, we can solve for distance k; c. (((t b; as
h; F (1 (t)).
This equation allows us to reconstruct decision boundary distances k; from optimal classification probabilities (t).) Thus the proof via trilateration for the disentangled representation theorem feasible for any suitable noise distribution E(t) as described above.
B.5 EXAMPLES OF SUITABLE NOISE DISTRIBUTIONS
B.5.1 ELLIPTICAL DISTRIBUTIONS
Definition: multivariate distribution Fang 2018 is elliptical its density function f(x can
be expressed
f(x) = 1 1 ((x 11) E - 1x 11))
35

---

Published as a conference paper at ICLR 2025
where g 10, oo) + [0, oo) is non-negative function, f R is the location parameter, and hxnn the scale matrix.
[0,
Properties:
Symmetric and unimodal around
Projections onto any direction c, yield univariate elliptical distributions.
Marginal distributions along c, have invertible CDFs if leads to such marginals.
Examples of Suitable Elliptical Distributions:
Multi-variate Gaussians with Full-Rank Covariance Matrix (Lemma B.11
Multi-variate T-distributions with Full-Rank Scale Matrix: Heavy-tailed alternative to the Gaussian.
Multivariate Laplace Distribution With Full-Rank Scale Matrix: Has exponential tails.
B.5.2 EXPONENTIAL POWER DISTRIBUTIONS
Definition: Also known as the generalized Gaussian distribution, defined by the density:
x
f(x) = (1/(1 exp
where 3 0 controls the kurtosis [ox & Tiao (201I)). Properties:
For 3 2, reduces the Gaussian distribution.
For 3 1, it becomes the Laplace distribution.
Symmetric and unimodal.
Marginal distributions are the same form and have invertible CDFs.
Generalizing the Proof
Given the above, the proof can be generalized to any noise distribution c(t) satisfying the conditions stated. The key steps are as follows:
1. Compute the Posterior Distribution:
Since e(t) is i.i.d., the likelihood function is:
P((X(s)) x²) II ((((ss) -x).
Without prior (uniform prior), the posterior is proportional to the likelihood.
The posterior distribution P(x* (X(s))==) can be found (or approximated) based
on the noise distribution.
2. Marginalize Along Decision Boundary Normals:
For each c, compute the marginal posterior distribution of c, x"
Due to the symmetry and unimodality of the noise distribution, this marginal will also be symmetric and unimodal.
3. Compute Classification Probabilities:
The classification probability is: F. F.b.-c.TM())
Y1(1) Pr(c. x" > by (X(+))
where F; is the marginal posterior CDF of c.'x" x".
4. Invert Marginal CDFs to Find Distances: Since F, is invertible, we can solve for k;:
k; F1 (1-()--)
36

---

Published as a conference paper at ICLR 2025
5. Set Up Linear System to Recover (tt):
The distances k; relate to (t() via:
c. p(t) k; + bj-
Collecting all Nusk equations:
Cu(t) = k b.
6. Solve for p(t): If C is full rank, we can solve for (u(t):
(u() = (CTC)- (k b).
Implications
Optimal Estimator Must Encode p(t):
The latent state (()) must contain sufficient information recover (at), it is essential for optimal classification across all tasks.
The theorem holds for any noise distribution satisfying the stated conditions, not just Gaus- sian noise.
B.5.3 EXAMPLE WITH MULTIVARIATE T-DISTRIBUTION
Suppose e(f) follows multivariate t-distribution dom v 2:
Kotz & Nadarajah (2004)) with degrees of free-
1. Posterior Distribution: The posterior P(x*
(X(s))==1 is also multivariate
distribution.
2. Marginal Posterior Distributions: Projections onto c; yield univariate t-distributions.
3. Invertible Marginal CDFs: The CDF t-distribution is known invertible.
4. Recover Distances: Use the inverse t-CDF find ky:
k; 8: 1 (()),
where 8, is the scale parameter, and T is the inverse CDF of the t-distribution with degrees of freedom.
5. Proceed with the proof: Follow the same steps as before reconstruct μ(t).
B.5.4 EXAMPLE WITH ANISOTROPIC GAUSSIAN NOISE
Lemma B.11. Suppose e(l) follows an anisotropic multi-variate Gaussian distribution with full- rank covariance matrix E and zero mean. Then we can update Equation from Lemma B.2 as
(() A Pr(c Pr(c b)
(23)
Kvi VeT Ea
(24)
Proof. Anisotropic noise results in the quadratic form Ee in the denominator representing the variance of the marginalized anisotropic noise distribution along decision boundary normal vector c, As long as is non- -ingula,, the remainder of the disentangled representation proof may proceed substituting Equation 14 with Equation 24
B.5.5 CONCLUSION
The key property enabling to recover distances from classification probabilities is the invertibility of the marginal posterior CDFs along the decision boundary normals. This property is not exclusive to Gaussian noise but is shared by broader class of noise distributions, including but not limited to:
Elliptical distributions (e.g., multivariate t-distributions, Laplace distributions).
37

---

Published as a conference paper at ICLR 2025
Exponential power distributions.
Other symmetric and unimodal distributions with invertible marginals.
Therefore, the proof of the Disentangled Representation Theorem generalizes to any noise distribu- tion satisfying the conditions outlined above. The essential requirement is that we can uniquely map the observed classification probabilities to distances along the normals, allowing us to reconstruct the posterior mean ((t) and establish that any optimal estimator must encode this information in its latent state (tt).
Correspondence the structure of noise distribution CDF along marginals F, and point-wise acti- vation functions (the activation function (1) = (((())))
B.6 DISCUSSION
The theoretical results presented in this appendix, particularly the Optimal Representation Theorem (Theorem B.6), provide insights into the factors driving representational convergence and alignment in neural networks and, more generally, any optimal multi-task classifier the setup shown in Figure S7 This theorem establishes clear connection between the latent representations learned by optimal multi-task classifiers and the true underlying data representation, offering a principled explanation for the emergence of disentangled representations aligned with the intrinsic structure of the data.
Connection to Manifold Hypothesis: Our theoretical results have important implications for the manifold hypothesis, which posits that real-world nigh-dimensional data tend to lie on or near low-dimensional manifolds embedded in the high-dimensional space (Fefferman al. 2016 Olah 2014). The key insight is that our proofs show an optimal multi-task classifier must encode an esti- mate of the disentangled coordinates of the true underlying environment state in its latent represen- tation. Consider the disentangled space in which x" resides, denoted X The injective observation map f X + X, where decision boundaries yi X - {0, 1} are linear. Our results imply that an optimal classifier's latent state Z()) must encode disentangled coordinates in X rather than ambient coordinates X.
The injective observation map f aligns closely to the typical conception of a data manifold (c.g.. if f E C or f e C'", as described in Tu (2017). The disentangled space X can be seen as the intrinsic coordinate system of the manifold, while f maps these coordinates to the high-dimensional observation space X. Our findings suggest that an optimal classifier will implicitly learn to invert this mapping to recover the disentangled coordinates. Moreover, for natural data where the manifold hypothesis holds, the learned latent representation would plausibly capture the manifold structure, as this is essential for disambiguating noisy observations and estimating the true underlying state. The low- -ww-iiensioonl manifold structure is key prior that an optimal classifier can (and in our case must) exploit improve its performance.
Relation to Autoregressive Language & Multi-Modal Transformers: Consider analogy with masked autoencoder vision foundation models, where x" is the "ground truth" of scene (ob- jects, positions. states, and relationships). the measurement variable X is an image with missing patches (Dosovitskiy et al.) 2020; He et al. 2021), and the model predicts the missing patch data v(x"). The model's 's latent variable Z exhibits some "understanding" of x" in the form of abstract representations useful for downstream tasks. This analogy extends to masked language models (De- vlin al. 2018) and autoregressive language models (Radford et al., (2019). where x* is "meaning" in a semantic space, X(!) are words, and y is the next word. Localizing x* from Z(t) relates to constructing a world model, showing that Z represents x* abstractly and with high fidelity.
Ordering of Noise and Observation Map: The ordering of the noise and the non-linear ob- servation map matters for the latent space representation. When the noise is applied before the observation map, the noisy observations are constrained to lie on manifold with the same intrinsic dimension as the true latent space X*. In contrast, when the noise is applied after the observation map, the noisy observations can deviate from the low-dimensional manifold, potentially introducing degeneracy where two noised observations arising from different x" may appear identical (i.e., non- injective). Imagine x as 2D piece of paper. An injective, smooth, continuous observation map
38

---

Published as a conference paper at ICLR 2025
A
B
x E X(O) x(==x++o((01)
Ex I (((()) E I
x 2
f(x)) E a X()) x(0)=1x*)
Figure S10: (A) x" is noised before being transformed by injective observation map f. resulting in observations f(X(t)) lying on the image of f (here f is a 2D folded surface). (B) x" is first transformed by injective observation map f and noise is added afterward, resulting in observations f(x*) + aN /) that not on the image of. ʃ.
f: X" - X' where X' is a 3-dimensional space "crumples" the sheet of paper X into a crumpled ball in X. If you add noise after the mapping, a point on one corner of the paper could get "popped out" of the 2D manifold by the noise and end up very far away on the crumpled surface if you were to examine flattened out (illustrated in Figure
The curvature of the observation map f and the level of noise a are fundamental factors influencing the extent of the degeneracy introduced by the noise after the observation map. High curvature in f can make the intrinsic geometry of the data more challenging to identify (e.g., more tightly crumpled paper). Large noise levels can push observations further from the underlying manifold, similarly worsening the potential degeneracy in the observations. The reach of the manifold (Aamari et al. (2019)) can be used as an immediate loose bound for post-observation map noise e2(() to ensure that the derived theorems still hold. Interestingly, this observation map degeneracy is a key concern work on visual predictive coding for map learning by Gornet & Thomson (2024). They demonstrate that a predictive coding network can learn to map Minecraft environment with visually degenerate states by integrating information to perform predictive coding along trajectories within the environment.
Connection to the Platonic Representation Hypothesis: Our results provide a new perspective on the Platonic representation hypothesis (Huh al. 2024). The Platonic representation hypothesis suggests that the convergence in deep neural network representations is driven by shared statistical model of reality, like Plato's 's concept of an ideal reality. Convergence of representations is analyzed in terms of similarity of distances between embedded datapoints among AI models trained on var- ious modalities. While the authors of the hypothesis argue that energy constraints might lead to divergence from shared representation for specialized tasks, Optimal Representation Theorem suggests that the key factor driving convergence is the diversity and comprehensiveness of the tasks being learned. As long the tasks collectively span the space of the underlying data representation. convergence to a shared, reality-aligned representation can occur, even in the presence of energy or computational limitations. Our theoretical results amount to necessary condition for optimal multi-task classifiers to represent a disentangled representation of the data within their latent state. With energy constraints, extraneous network activity may be regularized out of the model, resulting in greater alignment between disentangled representations in energy constrained models that "under- stand" the Platonic nature of reality. The very energy constraints Huh et al. (2024) suggest may lead to divergence could actually facilitate convergence of the platonic representations, as they may en- courage the learning of simple, generalizable features that capture the essential structure the data. This insight opens up interesting avenues for future research on the interplay between task diversity, energy constraints, and the emergence of shared representations. Finally, energy constraints have been shown naturally lead to predictive coding Rao & Ballard (1999)dd(999)).... Ali al. (2022), tightening
39

---

Published as a conference paper at ICLR 2025
the relationship between energy efficiency, prediction, and cognitive map learning. A relationship between predictive coding and optimal Bayesian estimation has also been established (Rao 1999).
Implications and Future Directions: The theoretical analysis presented in this appendix sheds light on the factors driving the emergence of disentangled representations in neural networks and their alignment with the intrinsic structure of the data. By formalizing the conditions under which learned representations recover the true underlying data manifold. our work provides foundation for understanding the remarkable success of representation learning across diverse domains. Av- enues for future research include exploring the sample complexity of learning under different obser- vation maps and noise levels, and empirically validating the convergence of representations across models and modalities in the context of task diversity and energy constraints.
40

---

Published as a conference paper at ICLR 2025
C SUPPLEMENTARY FIGURES & TABLES
Table S1: Hyperparameter values for RNN training. These values apply to all simulations, unless otherwise stated. 100ms chosen as conservative estimate of membrane time constant. was varied in some simulations (c.g. Figure We found that free RT tasks benefited from a higher learning rate. Other hyperparameters worked out of the box.
PARAMETER
VALUE
EXPLANATION
At
100 MS
BULER INTEGRATION STEP SIZE NEURONAL TIME CONSTANT
100 MS
N
64
NUMBER
NEURONS
a
0.2
INPUT NOISE STANDARD DEVIATION
T
20
TRIAL DURATION (IN Ats)
0.001/0.003
ADAM LEARNING RATE FIXED/FREE RT BATCH SIZE
16
batch
10
NUMBER OF TRAINING BATCHES
D
2
DIMENSIONALITY OF LATENT SPACE
Nuver
1
RNN/LSTM NUMBER OF LAYERS
41

---

Published as a conference paper at ICLR 2025
a
b
PC3
PC2
C
d
0.06
Tran
0.04
Loss
0.02
1000
Figure S11: Details of learned representations and of learning. (a) Representation after the decoder. out of 40 features were chosen randomly representations wrap around non-linearly, and the quadrants are overlapping. The RNN needs to invert the non-linear mapping and remove the noise to arrive at disentangled representations. (b) Individual trial examples from network in Figure 2b. Plotting conventions same as in Figure St. except here every trial has been mapped to a separate color (lines and final dot). The ground truth x for these example trials is shown in the bottom right, and the attractors have been made transparent for better visibility. Note that these trials include noise, like the ones the network has been trained on. As can be seen, the network maintains state space are also close in representation space. (c) Representation early in learning, for network trained with 1/4 of the examples compared Figure 31. The representation is not disentangled yet, however it is visible how the quadrants start separating and the attractors start spreading in the 2D manifold. (d) Same as in b, but for network with a delay period of 500 ms darker dots at the end of trajectories). Activity remains localized after the removal of the evidence streams, maintaining short-term memory the joint evidence with only minor leaks. (e) Representation learned network trained without input noise (a 0). Trajectories separate from the beginning, and there is no pressure to learn a 2D continuous attractor anymore. (f) Train and test errors for linear decoder for the OOD generalization task. Transparent lines correspond to different quadrants while opaque lines to the average across quadrants for one network.
displayed. Compared Figure 33 the
of metric distances in the 2D space: examples close in
42

---

Published as a conference paper at ICLR 2025
RNN
GPT
120
120
105
105
lotents
Mastt
Input dimensionality D
Input immessonntiyy D
Figure S12: Angles between latent factor decoders higher dimensions. Angles converge to 90 degrees Nusk >> D for RNNs, and early Nusk 2 D for transformers (see Appendix A.3 for angle estimation details). This confirms that multi-task learning leads to orthogonal, disentangled representations, in some cases even earlier than our theoretical guarantees.
43


Reviews:
[{"review_id": "RdfEfx123h", "reviewer": "Reviewer_2Aef", "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. \nThe 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.", "questions": "- In Equation 6, could you clarify if $x_{in}$ is the encoded input, i.e., $f(x)$?\n\n- 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?\n\n- Have you conducted experiments with other encoding functions, such as a quadratic mapping? \n\n- In Figure 5b, the blue curve for GPT appears to increase at the end. Could you clarify the reason behind this behavior?\n\n- 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?\n\n- On the other hand, if the inputs contained some temporal structure, how might this affect the nature of the representations?\n\n- Have you tested the framework with much larger values of $N_{task}$ and $D$ (e.g., over 100)?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ypesXno42U", "reviewer": "Reviewer_w8Z5", "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.", "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? \n2. 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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Originality: the study of the emergence of disentangled representation in temporal tasks and models is relatively new. \n2. Quality: the empirical validation is comprehensive. \n3. Clarity: the theory is well explained. The background section is also very nicely written and thorough. \n4. 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?\n2. 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?\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "teliXaCtby", "reviewer": "Reviewer_Mvb9", "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.", "questions": "**Sparsity:**\n\nThe 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?\n\nAdditionally, 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.\n\n**More nonlinear encoding:**\n\nThe 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.\n\nCould 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)).", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 1, "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.\n\nThe 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.\n\nThe 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)).\n\nIt 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.\n\nOverall, 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.\n\nInstead 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "588umFVt9a", "reviewer": "Reviewer_xvfc", "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.", "questions": "Comments:\n\nLine 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?\n\nCorollary 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?\n\nFigure 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.\n\nLine 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?\n\nFigure 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.\n\nLine 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.\n\nLine 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”.\n\nMinor comments:\n\nLine 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].\n\nLine 1408: Similar comment here to Line 1321 on the use of Bernoulli variable for \\hat{Y}(t).\n\nLine 1426: In Equation 16, boldfaced t should be scalar.\n\nLine 1592: “Proof: Recall equation (3)…” should be equation (4).\n\nLine 1608: “B^TB is a symmetric for any matrix B.” Some LaTeX formatting errors here and possibly missing the word matrix following “symmetric”.\n\nLine 309: “disentaglement” misspelled\n\nFig. 3d Caption: May be good to include a short statement about the inset on PCA components accounting for variance.\n\nLine 1070: “Once the factors are perfectly, performance…” missing the word correlated after “perfectly”\n\nLine 1162: “…for all lines in Figure S4a.” Did you mean all lines from Fig. 2a, shown in Fig. S4a?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "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.  \n\nIn 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\". \n\nThe 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. \n\nIt 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). \n\nThe 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fyUS1zRJ4Q", "reviewer": "Reviewer_b5qB", "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.", "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?\n\nBy evidence aggregation, did you mean evidence accumulation in the spirit of drift diffusion and race to bound models?\n\nIs 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?\n\nCan 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)?\n\nOn 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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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?\n\nThe 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. \n\nRelated 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).\n\nFicco, L., et al., 2021. Disentangling predictive processing in the brain: a meta-analytic study in favour of a predictive network. Scientific Reports. 11, 16258.\nGeorge, D., et al., 2021. Clone-structured graph representations enable flexible learning and vicarious evaluation of cognitive maps. Nature Communications. 12, 2392.\nGros, C., 2009. Cognitive Computation with Autonomously Active Neural Networks: An Emerging Field. Cognitive Computation. 1, 77-90.\nHu, 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.\nLinsker, R., 1990. Perceptual Neural Organization - Some Approaches Based on Network Models and Information-Theory. Annual Review of Neuroscience. 13, 257-281.\nMaass, 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.\nOlshausen, 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.\nRao, R.P., 1999. An optimal estimation approach to visual perception and learning. Vision Res. 39, 1963-89.\nRao, 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.\nSimoncelli, E.P., Olshausen, B.A., 2001. Natural image statistics and neural representation. Annu Rev Neurosci. 24, 1193-216.\nSuh, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)