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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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. Wtrue{\bm{W}}_{true} is the true mapping between x{\bm{x}} and y{\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. W1{\bm{W}}_{1} and W2{\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). βweight\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 R2R^{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 x{\bm{x}} and o{\bm{o}} . This is done by taking the absolute value of weights projecting from z{\bm{z}} to x{\bm{x}} and o{\bm{o}} , then summing over the space/object dimension (to obtain a vector the same dimension as z{\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 NetAutoencoderVAE
 ShallowDeepShallowDeepDeep
input dimension665050(64, 64, 3)
output dimension66n/an/an/a
# parameters24501385741570414870766295
# (hidden) layers15enc: 0enc: 3enc: 6
   dec: 0dec 3dec 6
Appendix
Table 1: The values of various hyper-parameters. enc/dec: encoder/decoder. n/a: not applicable.
EncoderDecoder
Input: 64 × 64 × number of channelsInput: 10
4 × 4 conv, 32 ReLU, stride 2FC, 256 ReLU
4 × 4 conv, 32 ReLU, stride 2FC, 4×4×64 ReLU
4 × 4 conv, 64 ReLU, stride 24 × 4 upconv, 64 ReLU, stride 2
4 × 4 conv, 64 ReLU, stride 24 × 4 upconv, 32 ReLU, stride 2
FC 2564 × 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. zc{\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 R2R^{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.
More in this topic

Disentangled Representation Learning

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

See all 231 papers

Explore similar work

CardsList
Sep 30, 2026

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

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.

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 AnalysisNeural Representation GeometryDeep Linear Networks
AskPDF
Oct 9, 2026

Neural networks are more modular than single neurons suggest

Neural networks routinely develop representations in which individual neurons contribute to multiple distinct functions, leading to the widespread view that such networks lack modular organization. We challenge this conclusion by introducing a formally nested hierarchy of modularity: unit-level, which requires distinct computations to occupy disjoint neurons; orthogonal, which asks whether a rotation reveals functionally independent subspaces; and linear, which asks whether a bounded linear map can separate task-relevant computations into independent subspaces. We develop a unified framework for learning decompositions at each level and evaluating them causally through selective ablation. Across vision models, language models and multi-task Recurrent Neural Networks (RNNs), we find that networks which lack unit-level modularity nonetheless often exhibit clear functional modularity at the subspace level. In vision and language models, content versus style and syntax versus semantics respectively dissociate under subspace ablation despite sharing overlapping populations of neurons. These learned subspaces also align differentially with brain data, with content and style subspaces predicting higher and lower visual cortex respectively, and semantic subspaces outperforming syntactic ones in predicting language cortex activity. In RNNs trained on multiple cognitive tasks and previously shown to lack modular structure, our decomposition recovers sharp double dissociations between tasks. The absence of unit-level modularity frequently coexists with clear subspace-level modularity across domains, enabling targeted intervention on specific computations even in networks where no individual neuron is specialized, with broad implications for mechanistic interpretability and for understanding functional organization in both biological and artificial neural systems.

Shreya Saha, Benjamin Bergen, Meenakshi Khosla

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

Modular Neural NetworksNeural Network InterpretabilityDisentangled Representation Learning
AskPDF
Jul 9, 2026

Contravariance Theory: Strong Alignment for Minimal Solutions to Hard Tasks

A series of results from the NeuroAI over the past fifteen years have raised core questions both about how to compare Deep Neural Network (DNN) models to the brain, and about how much convergent evolution to expect between artificial networks and real brain networks. Here, we show that for any two minimal DNN solutions to a sufficiently hard task: (i) "weak" alignment of network representations based on affine mappings guarantees "strong" alignment of privileged axes, and (ii) alignment "zippers" up the network hierarchy, causing the emergence of privileged axes from end-to-end task optimization. These results formalize the notion of contravariance from Cao and Yamins [2024], and illustrate important consequences for the theory of NeuroAI: with sufficiently strong tasks, choice of metric for inter-network comparison is not all that sensitive, and that convergent evolution is probably inevitable.

Dan Yamins, Aran Nayebi

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

Neural Representation Geometry
AskPDF