Disentanglement with Biological Constraints: A Theory of Functional Cell Types

James C. R. Whittington; Will Dorrell; Surya Ganguli; Timothy E. J. Behrens

q-bio.NCSep 30, 2022

Disentanglement with Biological Constraints: A Theory of Functional Cell Types

Authors: James C. R. Whittington, Will Dorrell, Surya Ganguli, Timothy E. J. Behrens

Organizations: Stanford & Oxford · UCL · Stanford · UCL & Oxford

Abstract

Neurons in the brain are often finely tuned for specific task variables. Moreover, such disentangled representations are highly sought after in machine learning. Here we mathematically prove that simple biological constraints on neurons, namely nonnegativity and energy efficiency in both activity and weights, promote such sought after disentangled representations by enforcing neurons to become selective for single factors of task variation. We demonstrate these constraints lead to disentanglement in a variety of tasks and architectures, including variational autoencoders. We also use this theory to explain why the brain partitions its cells into distinct cell types such as grid and object-vector cells, and also explain when the brain instead entangles representations in response to entangled task factors. Overall, this work provides a mathematical understanding of why single neurons in the brain often represent single human-interpretable factors, and steps towards an understanding task structure shapes the structure of brain representation.

Figures & tables

Figure 1: Proof intuition. Two uniformly distributed independent factors represented with two entangled neurons (left). The representation can be made nonnegative at the expense of activity energy (middle). Activity energy is minimised under a nonnegativity (and variance) constraint when the neurons are axis aligned to task factors (i.e. disentangled, right). Grey boxes denote uniform distributions over neural activity induced by uniform distributions over task factors. Note our proof does not require uniformity.

Figure 2: Shallow linear networks disentangle. We train 1-hidden layer linear networks on linear data. A) Cartoon schematic showing both input and output are entangled linear mixtures of factors (colours). Neurons colours schematically denote which of the factors it codes for. ${\bm{W}}_{true}$ is the true mapping between ${\bm{x}}$ and ${\bm{y}}$ . B) A model without biological constraints learns entangled internal representations. Mutual information matrix (scale 0 to 0.6) shown on left, cartoon schematic on right. ${\bm{W}}_{1}$ and ${\bm{W}}_{2}$ are the learnable weights projecting to and from the hidden layer. C) A model with our constraints learns disentangled representations (MI matrix scale 0 to 2.25). D) Several model variants, in which only those with all our constraints learn disentangled representations (definition of metric MIR in Appendix A.1 ). Average and standard error shown for 5 random seeds.

Figure 3: Deep nonlinear networks disentangle. We train 5-hidden layer nonlinear networks with our constraints on linear and nonlinear data. A) For linear data, all layers in the network learn a disentangled representation. B) For nonlinear data, only later layers learn a disentangled representations. C) Example mutual information matrix from the penultimate hidden layer.

Figure 4: Learning data generative factors with autoencoders. A) Training linear autoencoders on linear data. Only models with our constraints learn disentangled representations. B) Example mutual information matrix from a high MIR model. C) Nonlinear autoencoders trained on nonlinear data. Only models with our constraints learn disentangled representations. D) Example mutual information matrix from a high MIR model. All learning curves show mean and standard error from 4 mean from 5 random seeds.

Figure 5: Learning data generative factors with variational autoencoders. A) We train on the Shapes3D dataset, with two example images shown. These images have 6 underlying factors. B) MIG scores are higher with higher weight regularization, and generally higher than any $\beta$ -VAE (panel D). $\beta_{\text{weight}}$ is the regularisation strength of the weight regularisation . C) Mutual information matrix for a high scoring model. D) $\beta$ -VAE MIG scores. Adding a ReLU improves MIG scores. E) MIG score against $R^{2}$ shows models with our constraints lie in the Goldilocks region of high disentanglement and high reconstruction. All learning curves show mean and standard error from 5 random seeds. Results from an additional dataset in Fig. 13

Figure 6: Modules of distinct cell types form with nonnegativity and factorised tasks. A) When rodents navigate environments with objects, GCs encode location in physical space with firing fields lying on a hexagonal lattice, while B) OVCs encode relative location to objects with firing fields at specific distances and orientations from (white) objects. These plots are ratemaps; the average firing of a given cell at every location. C) Top: To model these cells we use a task environment where objects move location in different contexts - space and objects are factorised. Bottom: We train a representation to predict 1) spatial location, 2) object location, and 3) correct action at every location. D-E) Model ratemaps when trained with a ReLU activation function. Task 1 contains no objects, while task 2 has several objects (white dots). We see two types of cell representation: D) GCs that do not change across tasks and E) OVCs that only appear in the presence of objects. F) Model ratemaps when trained without ReLu activation. All representations are multi-peaked but amorphous. Further cells representations shown in Fig. 14 . G) To quantify modules in the ReLu model, we compute the distribution of individual cell spatial correlations across different tasks. H) The mode with high spatial correlation are cells that have high grid-score ( Barry et al., 2012 ) and are grid cells (GCs), the mode with low spatial correlation are cells that respond similarly around objects (OVCs). I) Only one mode of cells is seen without a ReLu activation. J) To quantify module-ness over many random seeds, we compute the cosine distance between the population’s contribution to ${\bm{x}}$ and ${\bm{o}}$ . This is done by taking the absolute value of weights projecting from ${\bm{z}}$ to ${\bm{x}}$ and ${\bm{o}}$ , then summing over the space/object dimension (to obtain a vector the same dimension as ${\bm{z}}$ ), then computing the cosine distance. We see low cosine distance for the ReLu indicating different cells code for space vs objects - i.e. modules.

Figure 7: Entangled tasks lead to entangled representations and grid cell warping. We show a representative selection of cells from a model with A) a factorised task (as in Fig. 6 ) and B) an entangled task. The phases of the firing fields in the entangled task are locked to object location - they have warped their firing fields. This is not the case for the factorised task (asides from object specific cells). C) To quantify phase locking for each of the task/model variant, we compute the average spatial correlation of patches around objects. Only the entangled task shows high correlation, i.e. the cell representations have warped around objects. Each point is a model trained from a random seed.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

	Supervised Net		Autoencoder		VAE
	Shallow	Deep	Shallow	Deep	Deep
input dimension	6	6	50	50	(64, 64, 3)
output dimension	6	6	n/a	n/a	n/a
# parameters	2450	138574	1570	414870	766295
# (hidden) layers	1	5	enc: 0	enc: 3	enc: 6
			dec: 0	dec 3	dec 6

Appendix

Table 1: The values of various hyper-parameters. enc/dec: encoder/decoder. n/a: not applicable.

Encoder	Decoder
Input: 64 × 64 × number of channels	Input: 10
4 × 4 conv, 32 ReLU, stride 2	FC, 256 ReLU
4 × 4 conv, 32 ReLU, stride 2	FC, 4×4×64 ReLU
4 × 4 conv, 64 ReLU, stride 2	4 × 4 upconv, 64 ReLU, stride 2
4 × 4 conv, 64 ReLU, stride 2	4 × 4 upconv, 32 ReLU, stride 2
FC 256	4 × 4 upconv, 32 ReLU, stride 2

Appendix

Table 2: Encoder and decoder architecture for the variational autoencoder experiments.

Figure 8: Sparsity intuition and simulations. Sparsity does not encourage disentanglement. Here the best situation is maximal entangling. Simulations verify that sparsity constraints do not induce disentanglement.

Figure 9: Sparsity on weights does not lead to disentanglement.

Figure 10: Schematic for representing categories. Each axis is neural firing rate. Each colour is a population activity for a single category, i.e.

{\bm{z}}_{c}

Figure 11: Learning data generative factors with autoencoders. A) Training linear autoencoders on linear data. Only models with our constraints, or with sparsity, learn disentangled representations. B) Example mutual information matrix from a high MIR model. All learning curves show mean and standard error from 4 mean from 5 random seeds.

Figure 12: Latent space traversals. An image is encoded and then the value of a single latent dimensions is changed, and the resulting image is generated. Each row shows this image for when a single latent dimension is traversed. The histogram shows the marginal distribution of that latent variable. From the images, we can see that each latent dimension predominantly cares about a single ground truth factor.

Figure 13: Same as Fig. 5 , but for the dSprites dataset. A) We train on the dSprites dataset, with two example images shown. These images have 5 underlying factors. B) MIG scores are higher with higher weight regularization, and generally higher than any $\beta$ -VAE (panel D). C) Mutual information matrix for a high scoring model. We note it has dropped a latent dimension. D) $\beta$ -VAE MIG scores. We note the high disentanglement $\beta$ -VAE models have bad reconstruction as they drop latent dimensions. E) MIG score against $R^{2}$ shows models with our constraints lie in the Goldilocks region of high disentanglement and high reconstruction. All learning curves show mean and standard error from 5 random seeds.

Figure 14: All cells. A/B) Ratemaps from a ReLU model. A) A task with no objects. B) Task with objects. C/D) Ratemaps from a model with linear activation function. C) A task with no objects. D) Task with objects.

Disentangled Representation Learning

Disentangled representation learning separates distinct factors of variation so models can interpret, control, and generalize their predictions.

See all 231 papers

Not all solutions are created equal: An analytical dissociation of functional and representational similarity in deep linear neural networks
Sep 30, 2026Lukas Braun, Erin Grant, Andrew M. SaxeRepresentational Similarity Analysis Neural Representation Geometry
Department of Experimental Psychology, University of Oxford, Oxford, UK · Gatsby Unit & Sainsbury Wellcome Centre, University College London, London, UK
Neural networks are more modular than single neurons suggest
Oct 9, 2026Shreya Saha, Benjamin Bergen, Meenakshi KhoslaModular Neural Networks Neural Network Interpretability
Electrical and Computer Engineering, University of California, San Diego · Cognitive Science Department, University of California, San Diego · Computer Science and Engineering, University of California, San Diego
Contravariance Theory: Strong Alignment for Minimal Solutions to Hard Tasks
Jul 9, 2026Dan Yamins, Aran NayebiNeural Representation Geometry
Departments of Computer Science and Psychology, and Wu Tsai Neurosciences Institute, Stanford University, Stanford, CA 94305 USA · Machine Learning Department, Neuroscience Institute, and Robotics Institute, Carnegie Mellon University, Pittsburgh, PA 15213 USA

Sep 30, 2026J/K move · Enter open · S save

Not all solutions are created equal: An analytical dissociation of functional and representational similarity in deep linear neural networks

Lukas Braun, Erin Grant, Andrew M. Saxe

Department of Experimental Psychology, University of Oxford, Oxford, UK · Gatsby Unit & Sainsbury Wellcome Centre, University College London, London, UK

Representational Similarity Analysis Neural Representation Geometry Deep Linear Networks

Open paper AskPDF arXiv

A foundational principle of connectionism is that perception, action, and cognition emerge from parallel computations among simple, interconnected units that generate and rely on neural representations. Accordingly, researchers employ multivariate pattern analysis to decode and compare the neural codes of artificial and biological networks, aiming to uncover their functions. However, there is limited analytical understanding of how a network's representation and function relate, despite this being essential to any quantitative notion of underlying function or functional similarity. We address this question using analysable two-layer linear networks and numerical simulations in non-linear networks. We find that function and representation are dissociated, allowing representational similarity without functional similarity and vice versa. Further, we show that neither robustness to input noise nor the level of generalization error constrain representations to the task. In contrast, networks robust to parameter noise have limited representational flexibility and must employ task-specific representations. Our findings suggest that representational alignment reflects computational advantages beyond functional alignment alone, with significant implications for interpreting and comparing the representations of connectionist systems.

Abstract

Disentangled Representation Learning

Explore similar work

Not all solutions are created equal: An analytical dissociation of functional and representational similarity in deep linear neural networks

Neural networks are more modular than single neurons suggest

Contravariance Theory: Strong Alignment for Minimal Solutions to Hard Tasks