Disentanglement with Biological Constraints: A Theory of Functional Cell Types
Authors: James C. R. Whittington
Will Dorrell
Surya Ganguli
Timothy Behrens
Venue: ICLR 2023 notable top 25% (2023)
Keywords: Disentangling
neurosciece
representation learning
hippocampus
cortex

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.

Published as a conference paper at ICLR 2023
DISENTANGLEMENT WITH BIOLOGICAL CON- STRAINTS: A THEORY OF FUNCTIONAL CELL TYPES
James C.R. Whittington"
William Dorrell UCL
Surya Ganguli Stanford
Timothy E.J. Behrens UCL & Oxford
Stanford & Oxford
ABSTRACT
Neurons in the brain are often finely tuned for specific task variables. More- over, such disentangled representations are highly sought after in machine learn- ing. 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 become se- lective for single factors of task variation. We demonstrate these constraints lead to disentanglement in 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 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.
INTRODUCTION
Understanding why and how neurons behave is now foundational for both machine learning and neuroscience. Such understanding can lead to better, more interpretable artificial neural networks, as well as provide insights into how biological networks mediate cognition. key to both these pursuits lies in understanding how neurons can best structure their firing patterns to solve tasks.
Neuroscientists have some understanding of how task demands affect both early single neuron re- sponses OOshausenn & Field 1996; Yamins et al., Ock4 Ocko et al. 2018 McIntosh et al. (016) and
population level measures such as dimensionality (Gao et al., 2017] Stringer et al.) 2019). However, there is little understanding of neural population structure in higher brain areas. As example, do not even understand why many different bespoke cellular responses exist for physical space, such as grid cells (Hafting al., 2005), object-vector cells (Hoydal al. 2019), border vector cells (Sol- stad et al. 2008; Lever et al. 2009), band cells (Krupic et al. 2012) or many other cells (0 Keefe & Dostrovsky 1971 Gauthier Tank 2018 Sarel et al. 2017 Deshmukh Knierim 2013). Each cell has a well defined, specific cellular response pattern to space, objects, or borders, as opposed mixed response space, objects, and borders. Similarly, we don't understand why neurons in inferior temporal cortex are aligned to axes of data generative factors (Chang & Tsao 2017 Bao al. 2020 Higgins al,, (02T), why visual cortical neurons are de-correlated (Ecker al. 2010),
why neurons in parietal cortex are selective only for specific tasks (Lee al. 2022). why prefrontal neurons are apparently mixed-selective (Rigotti et al.) 2013), and why grid cells sometimes warp
towards rewarded locations (Boccara al.) 2019) and sometimes don't 't (Butler et al. 2019). In
essence, why are some neural representations entangled and others not?
Machine learning has long endeavoured to build models that disentangle factors of variation (Hinton et al. 2011 Higgins et al. 2017a| Locatello et al. 2019 Bengio et al. 2012). We define disentan- glement single neurons responding to single factors of variation (see Appendix A.T for further details). Such disentangled factors can facilitate compositional generalisation and reasoning (Hig-) gins al. 2018 2017b; Whitiigton al. (021a) (though some work has challenged the idea that disentangled representations generalise better (Schott et al. (022)). as well as lead to more in- terpretable outcomes in which individual neurons represent meaningful quantities. Unfortunately, building models that disentangle is challenging (Locatello et al. 2019))
"Correspondence jerwhi werwhittinge itinttonngeaat coc

---

Published as a conference paper at ICLR 2023
[0.8]
-0.6 Vur(s,)
0
[0.8
-0.0
1.4 1.4
0
(.6)
+.02
0.6
0.8
2n*
Σ
Vor(s))
20
Varls 2n
E||a||
20²
=||=||
20 2.4
E||a| |
20' 26
Neuron
Nouron
Neuron 2
Make non-negative
Minimise activity
Nouron
Neuron
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 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 proof does not require uniformity.
In this work we 1) prove simple biological constraints of nonnegativity and minimising activity energy lead to factorised representations in linear networks; 2) empirically show these constraints lead to disentangled representations in both linear and nonlinear networks; 3) obtain competitive disentanglement scores on standard disentanglement benchmark; 4) provide an understanding why neurons in the brain are characterised into specific cell types due to these same biological constraints; 5) empirically show these constraints lead to specific cell types; 6) suggest when and why neurons in the brain exhibit disentanglement versus mixed-selectivin
Please see appendix A.2 for a comprehensive discussion relating our work to existing literature.
LINEAR DISENTANGLEMENT WITH BIOLOGICAL CONSTRAINTS
We first provide a theorem that suggests why the combined biological constraints of nonnegativity and energy efficiency lead to neural disentanglement (proofs of all theorems are in App. A.3
Theorem 1. Let e R* be random vector whose is independent components denote task factors.
We assume each independent task factor e is drawn from a distribution TT hhat has mean 0, variance a², and maximum and minimum values of min(e;) = -a linear neural representation of the task factors given by
a and max(e,) a. Also let z E R" be
Me b..
where M gnnk are mixing weights and b, E Rn is a bias. We further assume two constraints: (1) the neural representation is nonnegative with 2, 2 for all í 1, I,...,n, and (2) the neural population variance a nonzero constant, Σ, Var(2;) = C, so that the neural representation retains some information about the task variables. Under these two constraints we show that in the space of all possible neural representations (parameterised by M and bz). the representations that achieve minimal activity energy || also exhibit disentanglement, by which we mean every neuron zj is selective for most task parameter: i.e. MM | 0 for k + I, a.k.a. each row of M has at most I non-zero entry. (Proof in App. I.
The intuition underlying the proof of the theorem is shown in Fig where the key idea can be seen with two neurons encoding two factors. In particular, the bias must make every z; nonnegative for all values of e₁ and e2. But since e and 6g are independent, the minimum firing of neuron 1 for example obeys min(ze (21) = min (0.8e; 0.6ez) = min (0.8e1) + min(-0.6es = -a(0.8 + 0.6). Thus for neurons that mix factors, larger bias term must be used to ensure non- negativity, which leads to increased expected energy. Minimising this energy (subject to a constant variance) requires the smallest possible bias for each neuron, which occurs when each neuron is selective for single task factor. We note that this is consistent with neurons having baseline firing
Intuition.
We understand that task factors, or indeed neurons in the hrain, are generally not i.i.d., but we have made this assumption for mathematical convenience.
2

---

Published as a conference paper at ICLR 2023
rate where now activity below baseline corresponds to negative/positive values in the distribution, and activity above baseline corresponds to positive/negative values in the distribution.
The above theorem, while simple, is restricted in two ways: (1) the independent task factors e, directly available to the network; (2) representational collapse (i.e. setting x = 0) under energy minimisation is prevented solely by variance constraint. We thus consider more general setting where a neural circuit receives not the independent task factor vector e, but instead receives the mixed combination the neural representation z as a linear generative model that can predict observed data æ via æ Wz + bz Thus prediction, not variance constraint, now prevents collapsing neural representations (proof in Appendix A.3.2). Furthermore in Appendix [..333 A.3.3 we prove the following:
De, where E R"", D B and m 2 n 2 k. We further model
Theorem 2.
Let æ De be observed entangled data, where the independent task factor vector obeys the same distributional assumptions as in Theorem 1. Let neural representation z exactly predict observed data via x Wz + b₂ with zero error. Then for all such data generation models (with parameters D) and all such neural representations (with parameters W and b..). as long as: (I) the columns
are (scaled) orthonormal: (2) the norm of the read-out weights W is finite; (3) the neural representation is nonnegative (i.e. z > 0). then out of all such neural representations, the minimum energy representations are also disentangled ones. By this we mean that each neuron 2 will be selective for at most one hidden task factor E..
We note that D having (scaled) orthonormal columns may seem like a strong constraint, but holds approximately for any random matrix D with many observations (dimensionality of æ) and few independent task factors (dimensionality of e) (proof in Appendix [A.3.4). Appendix [A.3.5] discusses and provides intuition for when disentanglement occurs as D takes more general forms.
Strikingly, the essential content of Theorem 2 is that any linear, nonnegative, optimally energetically efficient, generative neural representation that accurately predicts entangled observations that linear mixtures of hidden task factors, will possess single neurons that are selective for individual task factors, despite never having direct access to them. In terms of applications, this theorem could apply in supervised or self-supervised machine learning settings in any neural network layer that linearly read-out from, or in neuroscience settings where z reflects a neural population from which one attempts to linearly decode task factors. While Theorem 2 holds in a simple setting, we will show through simulations that its essential content, namely that nonnegativity and energy efficiency together promote disentanglement, holds in practice in much more complex multilayer neural networks.
3 DISENTANGLEMENT IN MACHINES
We now present simulation results demonstrating that nonnegativity and energy efficiency (min- imising either activity or weight energy) lead to single neuron selectivity for single task factors. We show this for supervised and unsupervised learning, both for linear and nonlinear tasks and networks (details of datasets, models, and simulations in Appendix A.4 1.4&lp A.5)
A measure for disentangled subspaces. While our theory describes when single neurons become selective for single independent task factors, it does not limit the number of neurons selective for any given factor. For example in Theorem 2, four copies of the same neuron in z, each with half the activity, along with four copies of projecting weights each with half the values, predicts æ just as well and has exactly the same energy in both z and W More interestingly, an underlying task factor may not be -i-dimmnsional, e.g. spatial location, in which case the subspace that codes for this factor have at least the same dimension. This phenomena cannot be captured by many metrics of disentanglement (e.g. the popular mutual information gap; MIG; Chen et al. (2018) since they score highly if each factor is represented in just one neuron. Thus we define a new metric (mutual information ratio; MIR) that instead scores highly if each neuron only cares about one factor (see Appendix A. for details).
Regularizers as constraints. We impose nonnegativity via ReLU activation function, or softly via explicit regularization Crsanee = Swnnee Σ, max(-a;, 0) where indexes a neuron in the network, and B Bomnag determines the regularization strength. Similarly, we apply regularization to the activity
2I carming dynamics may favour fewer neurons per factor as there are fewer weights align.
3

---

Published as a conference paper at ICLR 2023
A
Thaining
Without brological consiratrt
With bistogical comminaints
X
X
0
W Krus
Wi
W2
W₂
Faoge
Stood
Maad
Mioo
Mioo
Mied
Boonaagged
Minedd
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. Wtrae is the true mapping between z and y. B) A model without biological constraints learns entangled internal rep- resentations. Mutual information matrix (scale 0 to 0.6) shown left, cartoon schematic on right. Wi and W2 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.I). Average and standard error shown for random seeds.
energy and weight energy; Lactivity Sactivity E, ||ai| and Lweight wweight Σ, |Will The role of Eweihtt to promote activity (variance) in the network, otherwise activity could be reduced via Lactivity. and such reduced activity could be compensated for with arbitrarily large weights. The total loss we optimise is
C Loonneg + Lactivity + Lweight
Cprediction
Biological constraints
Functional constraints
Here 'functional constraints', are any prediction losses the network has i.e. error in predicting target labels in supervised learning, or reconstruction error in autoencoders.
Disentanglement in supervised shallow neural networks. First we consider a dataset D {ox y). where æ is orthogonal mixture of six i.i.d. random variables (hidden independent task factors; uniform distribution), and y is linear transform of x (dimension 6; Fig. AA) First we train shallow linear networks to read-out from æ. Networks without biological constraints exhibit mixed internal representations (Fig. 2B). However, with our constraints, networks learn distinct sub-netwe networss for each task factor (Fig- 2C). Removing any one of our constraints leads to entangled representations (Fig. D) Lastly we note sparsity constraints do not induce disentanglement (Appendix [A.6). Thus the disentanglement effect of ReLUs is not from sparsity, but instead from nonnegativity.
Disentanglement in supervised deep neural networks. Training deep nonlinear (ReLU) networks on this data also leads to distinct sub-networks, with all layers learning disentangled representations (Fig. 3).. However with nonlinear data (æ + x², y remaining the same), the early layers are mixed-selectivve whereas the later layers are disentangled (Fig. BB-C). Understanding why the final hidden layer disentangles is easy, since it linearly projects to the target and so our theory directly applies. By extrapolating our theory, we conjecture that biological constraints encourage any layer to be as linearly related to task factors and as disentangled as possible. However, carly layers cannot be linear in hidden task factors since they are required to perform nonlinear computations on the nonlinear data, and thus only once activity becomes lincarly related to independent task factors in later layers does disentanglement set in (as predicted by our linear theory).
Disentanglement in unsupervised neural networks. We now consider unsupervised learning, i.e. D- = {z=) where
a linear mixture of multiple independent task factors in Theorem 2. Training O-hidden layer autoencoders on this data, with our biological constraints, recovers the independent task factors in individual neural subspaces (Fig. JA/B). Moreover, this only occurs when all constraints are present (Fig. [A). Again, even though our theory applies to the linear setting, the same phenomena occur when training deep nonlinear autoencoders on nonlinear data; i.e. when æ is nonlinear mixture of multiple i.i.d. random variables. i.e. D {f(=)} (Fig. (C-D).
Disentanglement on a standard benchmark with VAEs. We now consider a standard disentangle- ment dataset (Fig. [AA Kim & Mnih (2018). To be consistent with, and to compare to, the disentan- glement literature we use a VAE and measure disentanglement with the familiar gap (MIG) metric (Chen et al. 2018). For nonnegativity we ask the mean of the posterior to be
4

---

Published as a conference paper at ICLR 2023
A
02
Stop
Gradton
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 disentangled representation. B) For nonlinear data, only later layers learn disentangled representations. C) Example mutual information matrix from the penultimate hidden layer.
LAKEO
Cauer
Sta
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 high MIR model. All learning curves show mean and standard error from mean from random seeds.
nonnegative (via ReLU) but do not add a norm constraint the VAE loss already includes one in its KL. term between the Gaussian posterior and the Gaussian prior.
While Mate-offthhe-ar results not aim (instead wish to elucidate that simple biological constraints lead to disentanglement), our disentanglement results (Fig. (5B) are competitive and often better than those in the literature (comparing to results of many models shown in Locatello et al. (2019) even though those models explicitly ask for a factorised aggregate posterior. The particular baseline model we show here is 3-VAE (Fig. D). We that (1) our constraints lead to disentangiement (Fig. BBCC); (2) including ReLU improves 3-VAE disentanglement (as predicted by nonnegativity arguments above, Fig. 5D): and (3) our constraints give results in the Goldilocks region of high disentanglement and high reconstruction accuracy (Fig. 5)).
4 DISENTANGLEMENT IN BRAINS: A THEORY OF CELL TYPES
We next turn our attention to neuroscience, which is indeed the inspiration for our biological con- straints. While we hope our general theory of neural representations will be useful for explaining representations across tasks and brain areas, for reasons stated below, choose from spatial processing in the hippocampal formation. We show our biological constraints lead to separate neural populations (modules) coding for separate task variables, but only when task vari- ables correspond to independent factors of variation. Importantly, the modules consist of distinct functional cell types with similar firing properties, resembling grid (Hafting et al.) 2005) and object- vector cells (Høydal et al. 2019) (GCs and OVCs).
first example
We choose to focus on spatial representations for two reasons. Firstly, there is significant puz- zle about why neurons deep in the brain, synaptically far from the sensorimotor periphery, almost miraculously develop single cell representations for human-inteen cuumannitterpretbble factors (e.g. GCs for lo- cation in space; Fig. [6A, and OVCs for relative location to objects Fig. 6B). Such observations are not easily accounted for by standard neural network accounts that argue that representations are unlikely to be human-intens interritterreetable (Richards et al. 2019). Secondly, whilst these bespoke spatial
Using nonnegative posterior mean is odd when the prior is Gaussian, but it allows for casier comparison.
5

---

Published as a conference paper at ICLR 2023
A
2000-
With A. Nick n1. Rugg 201
Oradow Stop
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 8-VAE (panel D). Sweight is the regularisation strength of the weight regularisation. C) Mutual information matrix for high scoring model. D) 3-VAE MIG scores. Adding a ReLU improves MIG scores. E) MIG score against R 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 random seeds. Results from an additional dataset in Fig. 13
representations are commonly observed to factorise into single cells, there are situations in which selectivity spans across multiple task variables (Boccara et al. 2019; Hardcastle et al.) 2017). For example, sometimes spatial firing patterns of GCs are warped by reward (Boccara et al. (2019) and
sometimes they are not (Butler et al. 2019). There is no theory for explaining why and when this
happens.
A factorised task for rodents. We consider task in which rodents must know where they are in space, but must also approach one of multiple objects. If objects appear in different places in dif- ferent contexts, the task is factorised into independent factors (Fig. (C): 'Where am I in allocentric spatial coordinates?' and 'Where am I in object-centric coordinates By contrast, if objects always appear in the same locations, the task not factorised (as spatial location can predict object loca- tion). Formally, our task requires predicting spatial location, æ. whether an object is observed, the optimal action, a. If objects move between tasks, then p(æ, a, o) = p(xr)) ((zr)p(( o), where o and a are not factored since optimal actions are dependent on objects (see Appendix 1.9 details). Our theory says the representation will have two subspaces one for allocentric location (for predicting 2) and one for location relative to objects (for predicting and a) and that these sub-spaces should be represented in separate neural populations when biological constraints are present.
and
The representation in rodent brains indeed has two distinct modules of non-overlapping cell popula- tions: (1) GCs (Hafting et al.] (2005) which represent allocentric space via hexagonal firing patterns (Fig. [A);; and (2) OVCs (Høydal al. 2019) which represent relative location to objects through firing fields at specific relative distances and orientations (Fig. (B..
Model with additional structural constraint. Predicting allocentric spatial locations from ego- centric self-motion cues is known as path integration (Burak & Fiete 2009). and is believed to be fundamental function of entorhinal cortex (where GCs and OVCs are found). GCs naturally emerge from training RNNs to path integrate under several additional biological constraints (Sorscher et al.) 2019 2020; Banino et al., 2018 Cueva & Wei 2018). Hence to model this task (with locations and objects) we could train an RNN, z. that predicts (1) what the spatial location, œ, will be and (2) whether we will encounter an object, after an action, a, from the current location, and (3) what the expected action, a, will be. However, here we adopt a far more general framework that does not limit future applications of our approach simply to sequential integration problems.
In particular, it was recently shown (Gao et al. [222]] that path integration constraints can be applied directly on the representation by adding a new constraint in the loss imposing
(x) fW.,<(> a))-
Here f() is an activation function and W. is weight matrix that depends on the action a. This sur- rogate constraint imposes potential path integration by ensuring that a motion in space æ imposes a lawful change in neural representation z, thereby transforming the sequential path integration prob- lem into the problem of directly estimating neural representations of space (see Appendix A.9 for more details). Thus we minimise:
L = Lavaneg + Lactivit + Lweigh + Liocation + Lactions + Lubjects +
Gatt integration
44)
Biological constraints
Functional constraints
Structural constraints

---

Published as a conference paper at ICLR 2023
ReLu
E
ReLu
F
No Relu
Real grid cells
Task
Task 2
Task
Task 2
Task
Task 2
Task
Task
Cell
#1
#3
Internal
Cell
#4
#2
o o o Space Object
o boo
Action
ReLu
H
ReLu
No Relu
Real OVCS
Cell
Cell
Spotial comreation
Spania
Spar commiction
Crudint Stop
Figure 6: Modules of distinet 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 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 I) spatial location, 2) object location, and correct action every location. D-E) Model ratemaps when trained with ReLU activation function. Task 1 contains no objects, while task 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 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 over many random seeds, we compute the cosine distance between the population's contribution and This is done by taking the absolute value of weights projecting from z and o, then summing over the space/object dimension (to obtain a vector the same dimension as 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.
G) To quantify modules in the ReLu
mode of cells is seen without ReLu activation. J) To quantify module-ness,
These are the same biological constraints as above, but now the functional constraints involve pre- dicting location, object, and action, and additional structural constraint imposes equation In- terestingly, the structural constraint leads to a pattern forming optimisation dynamics (see Appendix A.9 for mathematical details).
Modules of distinct cell types when tasks are factorised and representations are nonnegative. Just theory predicts, when training on tasks where objects and space factorised (i.e. objects can be anywhere in space), under our biological constraints of nonnegativity and energy efficiency, distinct neural modules emerge, each selective for single task factor (Fig. D-B). We see GC-like neurons that consistently represent space independent of object locations, and OVC-like neurons that recenter their representations around the moving objects or are inactive if no objects are present (further cells are shown in Fig. [4]. Whereas without the nonnegativity constraint, all cells look qualitatively similar a single module of multi-peaked amorphous cells (Fig. (F). To quantify whether population really has two distinct cell types (two modules) we analyse the consistency of a cell's representation by taking its average spatial correlation between many different object configurations (Fig. 6G-I). In the ReLu case, there are two modes (Fig. 6G-1) showing double dissociation: the mode containing cells that don't 't change have high grid-score (Barry et al. 2012) and do not have consistent activity around objects. These cells that change between tasks have low grid-score and respond consistently around objects (Fig. 6H.. These cells respond objects and are comparable to OVCs.
GCs. Whereas the mode containing

---

Published as a conference paper at ICLR 2023
Spatial Correiation Around Objects
6.-
83
Inconsistent Object Locations
Locations Object
e1-
Model Data type
Figure 7: Entangled tasks lead to entangled representations and grid cell warping. We show a representative selection of cells from a model with 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, 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 model trained from random seed.
compute the average
Grid cell warping and mixed-selectivity when tasks are entangled. Experimental results show GCs sometimes warp their firing fields towards rewarded locations (Boccara et al. 2019) and some- times don't (Butler et al., 2019). Intriguingly, in the warping situation, the rodents exhibited stereo- typed behaviour; sequentially running between rewards using the same spatial trajectories rather than freely exploring the space (i.e. behaviour is entangled with space). We now explain these neuroscience observations consequence of space becoming entangled with objects/rewards.
Modelling factorised versus entangled tasks (objects changing locations versus always staying in the same locations), produces very different GC behaviours. In the factorised case, grid fields unrelated to objects (Fig. 7AA,, whereas in the entangled task grid fields warp to objects (Fig. ( We quantify this by measuring the average spatial correlation of patches around each object. Only when the task is entangled are fields consistently warped towards objects (Fig. (7). Thus we have an explanation for GC warping: they warp when space can no longer be disentangled from other factors (e.g. objects or behaviour) in behavioural tasks.
5 DISCUSSION
We have proven that simple biological constraints like nonnegativity and energy efficiency lead to disentanglement, and empirically verified this in machine learning and neuroscience tasks, leading new understanding of functional cell types. We now consider more neuroscience implica- tions.
Representing categories. Our theory additionally says that representations of individual categories should be encoded in separate neural populations (disentangled; see Appendix A.7 This pro- vides potential explanation for "ggranmmohhe cells" that selectively represent specific concepts or categories (Quiroga et al., 2005). and have long puzzled proponents of distributed representations further potentially explains situations where animals have been trained on multiple tasks, and different neurons are found to engage in each task (Rainer et al.) 1998 Roy et al. 2010 Asaad et al.) 2000 Lee al., 2022; Flesch al. 2022). We note that while our theory says you need at least one neuron per concept that is read-out (true even without disentanglement), it does not say you need neuron for every possible concept only for ones that are explicitly read-out.
When to disentangle? Our theory speaks situations in which brains networks must generalise to new combinations of learnt factors. In this situation, if the input-output (or input-latent-output)) transformations linear, biological constraints will cause complete disentanglement of the net- works. When the mappings are nonlinear, we show empirically that mixed- selectivvetiit exists, but gradually de-mixes layers approach the output (in supervised), or latent (in unsupervised), layers.
WWe note this phenomena could also be accounted for with a sparsity constraint.
8

---

Published as a conference paper at ICLR 2023
Optimising for low firing contrasts with previous ideas which instead optimise for linear read-our This latter situation is akin to kernel regression, where sixec-selectivity through random expansion increases the dimensionality of neural representations, allowing simple linear read-outs in the high dimensional space to perform arbitrary nonlinear operations on the original low dimensional space.
Mixed-selectivity. Mixed-selectivity exists in the brain. For example, Kenyon cells in the Drosophila mushroom body increase the dimensionality of their inputs by an order of magnitude by close random projections (Aso et al. 2014). This may allow linear read-out behaviour via simple dopamine gating. Similarly rodent hippocampal cells encode conjunctions of spatial and sen- sory variables to allow rapid formation of new memories (Komorowski et al., 2009 More recently it has been suggested that PFC neurons have this same property, for the same reason (Rigotti et al.) 2013). However, it is less clear that this is general property of representations (including PFC), which can separate into different neuronal representations of different interpretable
associative cortex
factors (Hirokawa al. 2019. Bernardi et al., 2020) or tasks (Lee et al. 2022 Flesch al. 2022).
One possibility is that in overtrained situations with only relatively small number of categories or trial-rypes (where mixed-selectiiityy has been observed), the task can effectively be solved by categorising the current trial into one of few previous experiences. By contrast in tasks where combinatorial generalisation is required, the factored solution may be preferred.
A program to understand how brain representations structure themselves. This work is one piece of the puzzle. tells us when neural circuits systems should represent different factors in different neurons. It does not tell us, however, how each factor itself should be represented. For example it does not tell us why GCs and OVCs look the ways they do. We believe that the same principles of nonnegativity, minimising neural activity, and representing structure, will be essential components obtaining this more general understanding. Indeed in a companion paper, we use the same constraints, along with formalising structure/path peruurrr/paah-integratonn using group and representations theory, to mathematically understand why grid cells look like grid cells (Dorrell et al. 2023). Simi- larly, our current understanding is limited to the optimal solution for factorised representations, but we anticipate similar ideas will be applicable to neural dynamics (Driscoll et al. 2022).
6 CONCLUSION
We introduced constraints inspired by biological neurons - nonnegativity and energy efficiency (w.rt. either activity or weights)- and proved these constraints lead to linear factorised codes being dis- entangled. We empirically verified this in simulation, and showed the same constraints lead to disentanglement with both nonlinear data and nonlinear networks. We even achieve competitive dis- entanglement scores on baseline disentanglement task, even though this was not our specific aim. We showed these biological constraints explain why neuroscientists observe bespoke cell types, e.g. GCs (Hafting et al. 2005), OVCs (Høydal et al.) 2019), border vector cells (Solstad et al. 2008 Lever et al.] 2009). since space, boundaries, and objects appear in a factorised form (i.c. occur in any independent combination), and so are optimally represented by different neural populations for each factor. These same principles explain why neurons inferior temporal cortex are axis aligned to underlying factors of variation that generate the data they represent (Chang & Tsao 20177 Bao eet al. [2200) Higgins et al., 2021), why visual cortex neurons are decorrelated (Ecker et al.) 2010), or why neurons parietal cortex only selective for specific tasks (Lee et al.) 2022). Lastly, we also explained the confusing finding of grid fields warping towards rewards (Boccara et al. 2019) as the space and rewards becoming entangled.
This work bridges the gap between single neuron and population responses, and offers an under- standing of properties of neural representations in terms of task structure above and beyond just dimensionality. Additionally it demonstrates the utility of neurobiological considerations in design- ing machine learning algorithms. Overall, we hope this work demonstrates the promise of a unified research program that more deeply connects the neuroscience and machine learning communities to help in their combined quest to both understand and learn neural representations. Such a uni- fied approach spanning brains and machines could help both sides, offering neuroscientists deeper understanding of how cortical representations structure themselves, and offering machine learners novel ways to control and understand the representations their machines learn.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
We thank Emile Mathieu and Andrew Saxe for helpful comments and advice on our manuscript. We thank the following funding sources: Sir Henry Wellcome Post-doctoral Fellowship (222817/Z/21/22) to J.C.R. W.; the Gatsby Charitable Foundation to W.D.; the James S. McDonnell, Simons Foundations, NTT Research, and an NSF CAREER Award S.G.; Wellcome Principal Research Fellowship (219525/2/19/Z) Wellcome Collaborator award 877Z) and Jean- François and Marie-Laure de Sermont-TTonnett Foundation award (JSMF220020372) to T.E.J.B.; the Wellcome Centre for Integrative Neuroimaging and Wellcome Centre for Human Neuroimaging are each supported by core funding from the Wellcome Trust (203139/2/16/2, 203147/7/16/).
10

---

Published as a conference paper at ICLR 2023
REFERENCES
Wael F. Asaad, Gregor Rainer, and Earl K. Miller. Task-Specific Neural Activity in the Primate Prefrontal Cortex. Journal of Neurophyssooogy, 84(1))451-459 July 2000. doi: 10.1 3000.84.1.4511 URL
in. 2000 84 451
Yoshinori Aso, Daisuke Hattori, Yang Yu, Rebecca M Johnston, Nirmala A Iyer, Teri-TB Ngo, Heather Dionne, LF Abbott, Richard Axel, Hiromu Tanimoto, and Gerald M Rubin. The neuronal architecture of the mushroom body provides logic for associative learning eLife, 3:e04577, December 2014. doi: 10.7554/eLife.0457. URL
/ 7554/eLife.
04577
Andrea Banino, Caswell Barry, Benigno Uria, Charles Blundell, Timothy Lillicrap, Piotr Mirowski, Alexander Pritzel, Martin Chadwick, Thomas Degris, Joseph Modayil, Greg Wayne, Hubert Soyer, Fabio Viola, Brian Zhang, Ross Goroshin, Neil Rabinowitz, Razvan Pascanu, Charlic Beattie, Stig Petersen, Amir Sadik, Stephen Gaffney, Helen King, Koray Kavukcuoglu, Demis Hassabis, Raia Hadsell, and Dharshan Kumaran. Vector- based navi- gation using grid-like representations in artificial agents. Nature, 557(7705):429 433, May 2018. doi: 10. 00.108//41186-008-00000 s41586 01 8 0102-6
URL http: /wwww nature www.enaateee.eee..rooooaeeoa.aae wwt./////raa.ae.eeeeeeeeeetcteee
Pinglei Bao, Liang She, Mason McGill, and Doris Y. Tsao. A map of object space in primate inferotemporal cortex. Nature, 583(78844 July 2020. doi: 10. 10000//41186-000-255555555 URL https:
org/ 10. 1038 1038/841586-0220 02000/2451556-0222255----- 2350--
David L. Barack and John W. Krakauer. Two views the cognitive brain. Nature Reviews
Neuroscience, 22(6):359-37L. June 2021. doi: 10.1030441533-000000444444 / //www. www.:/aatur... cwm art -
URL hetps:
Adrien Bardes, Jean Ponce, and Yann LeCun. VICReg: Saaauie--naaianne-CCCoo ization for Self-Supervised Learning. arXiv preprint, January 2022. URL httpp://arxiv.
Regular-
org/ org/abs/ 2105. 04906
H. Barlow. Possible Principles Underlying the Transformations of Sensory Messages. In Walter A. Rosenblith (ed.), Sensory Communication, pp. 216-234. The MIT Press, 1961. doi: 107751/mippess777802255580000000 URL /ttp:::///aaa mit
com/
H B Barlow. Single Units and Sensation: A Neuron Doctrine for Perceptual Psychology? Percep- tion, 1(4)3371-394, December 1972. doi: 10.1068/p010371. URL https:://dii.crg/11 1068/p01033
C. Barry, L. Ginzberg. J. O'eeee, and N. Burgess. Grid cell firing patterns signal environmen- tal novelty by expansion. Proceedings of the National Academy of Sciences, 109 17687- 17692, 2012. doi: 10. 10.3
patpppaas..199110000 1209918109. URL 10.1073/pnal www.0///pa..../// iateas.2209911 1209918109
org ///t
David Bau, Bolei Zhou, Aditya Khosla, Aude Oliva, and Antonio Torralba. Network Dissec- tion: Quantifying Interpretability of Deep Visual Representations, April 2017. URL http:
arrxiv waw.iii.ogg/abbaaaaaaaaaa.aaagaa wtt..///g//bb////// 05796
Timothy J Behrens, Timothy H Muller, James C. R. Whittington, Shirley Mark, Alon B Baram, Kimberly L. Stachenfeld, and Zeb Kurth-nelson. What Is Cognitive Map? Organizing Knowl- edge for Flexible Behavior. Neuron, 100(():490-5))) 509, 2018. doi: 10.1111// mewr... 10.0.0 10.002. URL hnttps://wwa cell. wtt.://wooo.ooo.emm fu text S0896 62 73 (18) 30856-0
Yoshua Bengio, Aaron Courville, and Pascal Vincent. ing: A review and perspectives.
Representation
learn-
arXiv preprint, (1993):1-34, 2012. doi:
MOOBBCE--8A96-4995 B88B B88B- 1552ECCS. -0A.-BB-BBBBB-811EEEAAAAAAAA /
URL cwm/ml http://www...oooooooos/sss//....
5538 (

---

Published as a conference paper at ICLR 2023
Silvia Bernardi, Marcus K. Benna, Mattia Rigotti, Jérôme Munuera, Stefano Fusi, and C. Daniel Salzman, The Geometry of Abstraction the Hippocampus and Prefrontal Cortex. Cell,
9377)),94--67.e21, November 2020. doi: 10.10166f.ce1112000000000 2289
URL https://
Charlotte N. Boccara, Michele Nardin, Federico Stella, Joseph O' and Jozsef Csiesvari. The entorhinal cognitive map is attracted to goals. Science, 363(6434):1443-14447 March 2019. doi: 10. /0.11///ccinnce.av44877. 1126/science.aav//
shaping inductive bias, March 2022. URL /tt.ip////ww..seeeeeeoooooooooor
Blake Bordelon and Cengiz Pehlevan. Population codes enable learning from few examples by
11C1/2021. 03. 30. 30.44473333 437743v3
Bariscan Bozkurt, Cengiz Pehlevan, and Alper T. Erdogan. Biologically-Plausiie Determinant Maximization Neural Networks for Blind Separation of Correlated Sources, September 2022.
http: org/ absaabss22a2a 12894
Yoram Burak and Ila R. Fiete. Accurate path integration in continuous attractor network models of grid cells. PLoS Computational Biology, 502)21000221, February 2009, doi: 10. 10.uual/journal. pcbi.1000291. URL httph:tts://// wlw. org/ 10
1000291
Christopher P. Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in S VAL. arXiv preprint, April 2018. doi: 00448500arXiv.1844..3333
org/ abs 1804 03599
William N. Butler, Kiah Hardcastle, and Lisa M. Giocomo. Remembered reward locations restruc- ture entorhinal spatial maps. Science, 363(6434) March 2019, doi: 10. 10.1 126/sceece. aav5297. URL http: http:://www.. www.:/wwwwwce....o....oo http:///ee.aaaa.ooooooooooo/oo/e aav52.7
Stephen Casper, Shlomi Hod, Daniel Filan, Cody Wild, Andrew Critch, and Stuart Russell. Graph- ical Clusterability and Local Specialization in Deep Neural Networks. April 2022. URL
https://openrevieww http://orrummiieeeeeeeeeeeeeee
Le Chang and Doris Y. Tsao. The Code for Facial Identity in the Primate Brain. Cell, 169(6): 1013-1028.e14. June 2017. doi: 00.1016/..eel1.201..5.0111 URL http:/ttp//ss /xwp/xxxcco dol.oooooooo org/
10. 10.1016/3. 2017 05.05..0.00 011
Ricky T. Q. Chen, Xuechen Li, Roger Grosse, and David K Duvenaud. Isolating Sources of Disentanglement in Variational Autoencoders. Advances in Neural Information Processing
Systems, 31, 2018. URL UtL/htttp:://prcceddnn....o [eee3frecco044233337779722222222
wewp:erroo....e.........r...p.pp cc/ pctpppaper/211//aaass..
Chaitanya Chintaluri and Tim Vogels. Metabolically driven action potentials serve neuronal
energy homeostasis and protect from reactive oxygen species, October 2022. URL https:
7/www. www.bbioxxiv.o org/ www.//t/oonnnnnnioeeeeo wtt.eeeenn....11122222...2
51242vvv
Pierre Comon. Independent component analysis, new concept?
Signal Processing, 36
(3):887-314, April 1994. doi: 10.1016/0166--6644449999-2333393
URL https: :// //www.
sciencedirect. com/ http:///cccccccc.aaaaaaaaaaaaaaa
Christopher J. Cueva and Xue-Xin Wei. Emergence of grid-like representations by training recurrent
neural networks perform spatial localization. International Conference Learning Represen- tations, 0:1-19, March 2018. URL httt:://aaaae wwtp:iii.rrraaaoooo abs//a03.... 07770
Sachin S Deshmukh and James J Knierim. Influence of local objects on hippocampal repre- sentations: Landmark vectors and memory. Hippocampus, 23(4):253-67, April 2013. doi: 10.1002/hipo.22101.. WtRL/htt::////e
ncbi .m.mm.m... m. nib www.nih.oooommm http://go//obmmee//3333. hte.l......ov//eee.
Yedidyah Dordek. Daniel Soudry, Ron Meir, and Dori Derdikman. Extracting grid cell char- acteristics from place cell inputs using non-negative principal component analysis. eLife,
(MARCH2016):1336, 2016. doi: 10.7554/eL 10094.
12

---

Published as a conference paper at ICLR 2023
Will Dorrell, Peter E. Latham, Timothy E. J. Behrens, and James C. R. Whittington. Action- able Neural Representations: Grid Cells from Minimal Constraints. In International Confer- ence Learning Representations, 2023. URL https://openreview wtw.://wwwweeoooouueemmm
x xfqDe22zh41,
Laura Driscoll, Krishna Shenoy, and David Sussillo. Flexible multitask computation recurrent
networks utilizes shared dynamical motifs. preprint, Neuroscience, August 2022. URL /
wttp:///1111110022..00200 //biorxiv.org. http://www..rrg/oooo////////////
503870
Alexander S. Ecker, Philipp Berens, Georgios A. Keliris, Matthias Bethge, Nikos K. Logothetis, and
Andreas S. Tolias. Decorrelated Neuronal Firing in Cortical Microcircuits. Science, 327(5963): 584 587, January 2010. doi: 10.1 166sccc//ce 11t.nne...1999777 1179667. URL htttts://wwett wwwp:wscieeee science.
org http:///..oo..11//cceeeecoccce// 11111////ccccee11116.61. 1179867
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-oddds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy Models of Super- position, September 2022. URL http://axx.....g////////////////
10652
Yuguang Fang, K.A. Loparo, and Xiangbo Feng. Inequalities for the trace of matrix product. IEEE Transactions on Automatic Control, (9(12) December 1994, doi: 10.1109/9.362844
URL https:: /eexpp orre leee wwt.or////ooooo
Timo Flesch, Keno Juechems, Tsvetomira Dumbalska, Andrew Saxe, and Christopher Summer- field. Orthogonal representations for robust context-dependent task performance brains and
neural Neuron, 11077):255-11200e 1270 2022. doi: 1.10166f.curron2222.000 https:://bi..rr//00 http://....1111/.oooooo. www.nnuunnnnommo 2022 005
URL
Peiran Gao, Eric Trautmann, Byron M. Yu, Gopal Santhanam, Stephen Ryu, Krishna V. Shenoy, and Surya Ganguli. A theory of multineuronal dimensionality, dynamics and measurement. URL htttt:: wwt.:wwwwb.......mm
bioRxiv preprint, 0:2 14262, 2017. doi: 10. 10112224442..
Ruiqi Gao. Jianwen Xie, Xue-Xin Wei, Song-Chur Zhu, and Ying Nian Wu. On Path Integration of Grid Cells: Group Representation and Isotropic Scaling. Advances in Neural Information Pro- cessing Systems 35, O(NeurPPS):1 June 2021. URL http:://axx...org/abb... ..000 10259
Jeffrey L. Gauthier and David W. Tank. A Dedicated Population for Reward Coding in the Hippocampus. Newron, 99())119-1133.7, 2018. doi: n0.u0066 201.umon.018.660000 URL
https://doi.org//ee 1016/1 www.on...eerrroommmmmeeee 2018 10118.0 06.006.00 008
Gabriel Goh, Nick Cammarata t. Chelsea Voss t. Shan Carter, Michael Petrov, Ludwig Schubert, Alec Radford, and Chris Olah. Multimodal Neurons in Artificial Neural Networks. Distill, 2021.
doi:
Merse E Gáspár, Sherre-Olivier Polack, Peyman Golshani, Máté Lengyel, and Gergö Orbán. Repre- sentational untangling by the firing rate nonlinearity in V1 simple cells. eLife, 8te43625, Septem- ber 2019, doi: 40.5554/eLife.43625. URL http:::://elfeeseeenesss..o 43625
Torkel Hafting, Marianne Fyhn, Sturla Molden, May-britt Britt Moser, and Edvard I. Moser. Mi- crostructure of spatial map in the entorhinal cortex. Nature, 436(7052):801-806, 2005. doi: 10.10388/ature0332...
Kiah Hardcastle, Niru Maheswaranathan. Surya Ganguli, and Lisa M. Giocomo. A Multiplexed, Heterogeneous, and Adaptive Code for Navigation in Medial Entorhinal Cortex. Neuron, 94 (2):37- 2017. doi: mel6//.neuron22000 03.025
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning Basic Visual Concepts with Constrained Variational Framework. International Conference on Learning Representations, 0,
July 2017a. htl.httss:// http://pwneeeeeeeeeeeeeeeeeeeaap
13

---

Published as a conference paper at ICLR 2023
Irina Higgins, Arka Pal, Andrei A.
Rusu, Loic Matthey, Christopher P Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. DARLA: Improving Shoo Transfer in Reinforcement Learning. arXiv preprint, 2017b. doi: 10331009114422622052222233 URL http:/tpp/aaaxi.orggg////// www.o////////////////////e//// 08475
Irina Higgins, Nicolas Sonnerat, Loic Matthey, Arka Pal, Christopher P Burgess, Matko Bosn- jak, Murray Shanahan, Matthew Botvinick, Demis Hassabis, and Alexander Lerchner. SCAN: Learning Hierarchical Compositional Visual Concepts. arXiv preprint, pp. 1-24, 2018. doi: 10.1 10.1186612884-01--155200 UBL.htpp//arrww..cr///bbbs
03389
Irina Higgins, Le Chang, Victoria Langston, Demis Hassabis, Christopher Summerfield, Doris Tsao, and Matthew Botvinick. Unsupervised deep learning identifies semantic disentanglement in single inferotemporal face patch neurons. Nature Communications, 12(1):1-14, 2021. doi: 10.1038 s41467-021-26751-5
Geoffrey E. Hinton, Alex Krizhevsky, and Sida D. Wang. Transforming Auto-Encoders. International Conference Artificial Neural Networks, 6791:44-51, 2011. doi:
1211077978-33-4222177555
URL
1007/
978 3 21735-213
Junya Hirokawa, Alexander Vaughan, Paul Masset, Torben Ott, and Adam Kepecs. Frontal cor- neuron types categorically encode single decision variables. Nature, 576(7787)446 451, December 2019, doi: 10. 10-10088/41186--1-118666
URL https: / /www nature........ com/
articles/s 5416//4116--11 019-6-0019111--111
Daniella Horan, Eitan Richardson, and Yair Weiss. When Unsupervised Disentan- glement Possible? Advances in Neural Information Processing Systems, 34:5150- 5161, 2021. URL https:://proceedingg ing
neurips wttpaaai..si..pppppeeeeeeep ccc/paperr/222//hsh
A. Hyvärinen and E. Oja. Independent component analysis: algorithms and applications. Neural Networks, 13(4)3411-430. June 2000. doi: 10.1001650093336000000000033333 URL https:
h/www.
Hyvärinen, Erkki Oja, Aapo Hyv, Erkki Oja, A Hyvärinen, and Erkki Oja. Indepen-
dent component analysis: A tutorial. Notes for International Joint Conference on Neu-
ral Networks (IJCNN 99), Washington DC, 1:1-30, 1999. doi: 10. (0.10993/mrra//tv2774
Not WeS.Ghhheeestaallalttoeeeeeeeeen
WashingtonDC1
Øyvind Arne Hoydal, Emilie Ranheim Skytøen, Sebastian Ola Andersson, May-Britt Moser, and Edvard Ingjald Moser. Object-vector vector coding in the medial entorhinal cortex. Nature, 568(7752): 400 404, April 2019. doi: 10.1038/s41586-011 019---099-100 1077-7. URL http::////// //ww///www www. naturatture.e.r com/ articles 1191--996--977----1111111
Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational Autoen- coders and Nonlinear ICA: Unifying Framework. Proceedings of the Twenty Third Inter- national Conference on Artificial Intelligence and Statistics, pp. 2207-2217, June 2020. URL httpss://proceedinng. mlr press/ vvsss///vvv///// www.kkheeaaaaaaa.am html
Hyunjik Kim and Andriy Mnih. Disentangling by Factorising, Proceedings 35th International Conference on Machine Learning, pp. 2649-2658, July 2018. URL httpps::/pprccoeennggg
mlr. wwwppp/wwe/..i.n..ppptpp 8b html
Diederik P. Kingma and Jimmy Lei Ba. Adam: A Method for Stochastic Optimization. arXiv 14511444833 1830553. URL
preprint, 0, 2014. doi: http://bb.aammmoogocpppoooo/l///
http://arre www.oogg//sss.. 6980
http:///rr....rrrabb////12222..
14

---

Published as a conference paper at ICLR 2023
Diederik P. Kingma and Max Welling. Auto- Auto-Encoding Variational Bayes. arXiv preprint. 0(MI): 1-14, 2013. doi:
URL http:// /ttp:://aaxivoorrt org. hbtp//bbb//1/1
6114
Robert W. Komorowski, Joseph R. Manns, and Howard Eichenbaum. Robust Conjunctive Item- Place Coding by Hippocampal Neurons Parallels Learning What Happens Where. Jour-
nal of Neuroscience. 29(31)99918 9929, August 2009. doi: 10.1523/NEUROSSCL
2009. URL https: www. www.ronee.ns..... http://wwwaaaca..ooooooooaoa///a 1378-09.2009
Julia Julija Krupic, Neil Burgess, John O'Keefe, and John Keefe. Neural Representations Loca- tion Composed Spatially Periodic Bands. Science, (37666996)3333333338337 857, August 2012. doi: 10.
htt/science.12224400 URL 1222403
Abhishek Kumar, Prasanna Sattigeri, and Avinash Balakrishnan. Variational Inference of Disen- tangled Latent Concepts from Unlabeled Observations. International Conference on Learning Representations, pp. 16, 2018.
Daniel Kunin, Jonathan M. Bloom, Aleksandrina Goeva, and Cotton Seed. Loss Landscapes of Regularized Linear Autoencoders. arXiv preprint, 2019. URL
1901. 08168
Daniel D Lee and H Sebastian Seung. Algorithms for Non-negative Matrix Factorization, Advances in Neural Information Processing Systems, 0(1), 2000.
Julie J. Lee, Michael Krumin, Kenneth D. Harris, and Mattco Carandini. Task specificity in mouse parictal cortex. Neuron, August 2022. doi: 10.10066/.eeumm02222.....111. URL https:/ //
www.sciencedi rect. wtt./cceeece//aaicccccac..ccccca htt.://c.......cc.../....../////
Colin Lever, Stephen Burton, Ali Jeewajee, John Keefe, and Neil Burgess. Boundary vector cells in the subiculum of the hippocampal formation. Journal of Neuroscience, 29(31);9771-97777 2009. doi: 10. 10.152//NEROOSCL. 131-0999..119999990000 19-09.2009.
David Lipshutz, Cengiz. Pehlevan, and Dmitri B. Chklovskii. Biologically plausible single-layer works for nonnegative independent component analysis, March 2022. URL httt::/pp///iii
org/abs/2010. 12632
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging Common Assumptions the Unsupervised Lean- ing of Disentangled Representations, June 2019. URL
12359
Lane McIntosh, Niru Maheswaranathan, Aran Nayebi, Surya Ganguli, and Stephen Bac- cus. Deep Learning Models of the Retinal Response to Natural Scenes. In Ad- vances in Neural Information Processing Systems, volume 29. Curran Associates. Inc., 2016. URL /ttts::/pprcced.ggs... neurips..... cc/paper/2016/haseee html
Samuel Ocko, Jack Lindsey, Surya Ganguli, and Stephane Deny. The emergence of multiple retinal cell types through efficient coding of natural movies. In Ad- vances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. doi: 10.1101/458737. URL https: 399E2774ccceeaa5552eee222205 /ttppppp://ppprr...///s nips. cc/paperr221B/haa.. Abtrtr3-bbstrttt..... .hmll
Erkki Oja. Simplified neuron model as a principal component analyzer. Journal of mathematical
biology, 15(3):267-23, November 1982. doi: (0.1007/BF00256657
John O' Keefe and J. Dostrovsky. The hippocampus as a spatial map. Preliminary evidence from unit activity in the moving rat. Brain Research, 34(): November 1971. doi: 10.1016/ 0006- 0006-8993/7119003338
com/ wtt.//trriee////i.11 TRL RRLhttp: / http::/www..eeeeeeeeeeoe.oeeerea
00066899371903377
15

---

Published as a conference paper at ICLR 2023
Bruno A. Olshausen and David J. Field. Emergence of simple-cell receptive field properties by learning a sparse code for natural images, 1996.
Cengiz Pehlevan, Tao Hu, and Dmitri B. Chklovskii. A Hebbian/Amti-Het -ebbian Neural Network for Linear Subspace Learning: Derivation from Multidimensional Scaling of Streaming Data.
Neural Computation, (117)))1411-195, July 2015. doi: 2,.162/NECO.a.00722 http:
//axxiv.cggaaa///2. 00659
Cengiz Pehlevan, Sreyas Mohan, and Dmitri B. Chklovskii. Blind nonnegative source separation using biological neural networks. Neural Computation, 29(11):2925-29)) 2954, November 2017. doi:
10.1162/neco a 01007. URL. http:///rxii....or/a//////// 00382
M. Plumbley. Conditions for nonnegative independent component analysis. IEEE Signal Processing Letters, 966):177-10, June 2002. doi: 10.1 10.110991LS22200280000002.
M.D. Plumbley. Algorithms for nonnegative independent component analysis. IEEE Transactions
on Neural Networks, 14(3)553--54,, May 2003. doi: 10. 109/TNN, 11661.
R. Quian Quiroga, L. Reddy, G. Kreiman, C. Koch, and I. Fried. Invariant visual representation by single neurons the human brain. Nature, 433(7045): June 2005, doi: 10.1038/ nature03687. WtLhttps:
www. nature aas.uure.cco//rrrrrrleee/rttrers
Gregor Rainer, Wael F. Asaad, and Earl K. Miller. Selective representation of relevant information by neurons in the primate prefrontal cortex. Nature, 393(66855 June 1998. doi: 10.
1038/31235. URL https / www. natura.t.e....
Blake A. Richards, Timothy P. Lillicrap. Philippe Beaudoin, Yoshua Bengio, Rafal Bogacz, Amelia Christensen, Claudia Clopath, Rui Ponte Costa, Archy de Berker, Surya Ganguli, Colleen Gillon, Danijar Hafner, Adam Kepecs, Nikolaus Kriegeskorte, Peter Latham. Grace W. Lind- say, Kenneth Miller, Richard Naud, Christopher C. Pack, Panayiota Poirazi, Pieter Roelfsema, João Sacramento, Andrew Saxe, Benjamin Scellier, Anna C. Schapiro, Walter Senn, Greg Wayne, Daniel Yamins, Friedemann Zenke, Joel Zylberberg, Denis Therien, and Konrad P. Kording. A deep learning framework for neuroscience. Nature Neuroscience, 22(11):1761-1770, November 2019. doi: 00100888-41599--09--5222000-- 641593 019-0520-2
URL http:: http://wwwww.tt...eeeee.. com/ articles/
Karl Ridgeway and Michael C Mozer. Learning Deep Disentangled Embeddings With the F-Statistic Loss. Advances in Neural Information Processing Systems,
URL https::/pproceeddngg 31, 2018.
wwtparisss.cccppppeeeee cc/ppaper:2018//ash
Mattia Rigotti, Omri Barak, Melissa Warden, Xiao-Jing Wang. Nathaniel Daw, Earl Miller, and Stefano Fusi. The importance of mixed selectivity in complex cognitive tasks. Nature,
497(7451):1-66 2013. doi: 10. 10.1038//auur URL http://tt//// dei wtt.loorg.0000000
nature 21 60]
Michal Rolinek, Dominik Zietlow, and Georg Martius. Variational Autoencoders Pursue PCA Di- rections (by Accident), April 2019. htl/tttpp/arxiv.cr// 1ttp://i....//////11..0.007..
Jefferson E. Roy, Maximilian Riesenhuber, Tomaso Poggio, and Earl K. Miller. Prefrontal Cor- tex Activity during Flexible Categorization. Journal of Neuroscience, 30/25):8519-8528, June 2010. doi: 10. V0.522//NUUROSC11-8870022222 content/30/25/851
URL https::///// //wp::///wwww.. www. httts::/www.jeuurcccc..oo /rw.rosci.oorg
Ayelet Sarel, Arseny Finkelstein, Liora Las, and Nachum Ulanovsky. Vectorial representation of spatial goals in the hippocampus of bats. Science, 355(6321):176-180, 2017. doi: 10.1 10.1126/ Science.aak9589.
Lukas Schott, Julius von Kügelgen, Frederik Träuble, Peter Gehler, Chris Russell, Matthias Bethge, Bernhard Schölkopf, Francesco Locatello, and Wieland Brendel. Visual Representation Learn ing Does Not Generalize Strongly Within the Same Domain, February 2022. URL http:
7/arxiv.co ///rrsiw.org/abb/1107....33
16

---

Published as a conference paper at ICLR 2023
Trygve Solstad, Charlotte N Boccara, Emilio Kropff, May-Britt Moser, and Edvard Moser. Repre- sentation of Geometric Borders in the Entorhinal Cortex. Science, 32265909):1865-18868 De- cember 2008. doi: 10.1 http:/eiiinee..1666666
URL https: www.wwiwncema......m..mmmo org/
lookup/do: 10. 1 w11.///cceee scwence.........o 1 1166466
Ben Sorscher, Gabriel C Mel, Surya Ganguli, and Samuel A Ocko. A unified theory for the origin of grid cells through the lens of pattern formation. Advances in Neural Information Processing 2019.
Systems 32, (NNeurPP)::10033-10000
Ben Sorscher, Gabriel C Mel, Samuel A Ocko, Lisa Giocomo, and Surya Ganguli. A uni- fied theory for the computational and mechanistic origins of grid cells. bioRxiv preprint, pp. 2020.12.29 0.00....9.424883,, 2020. URL http:s:::/p.i.oppppppopp
110111111121001111 2020 12. 29. 424583
Carsen Stringer, Marius Pachitariu, Nicholas Steinmetz, Matteo Carandini, and Kenneth D. Harris. High-dimensional geometry of population responses in visual cortex. Nature, 2019. doi: 10. 10.1038/ 141586-019-1346-33
01
James R. Whittington, Rishabh Kabra, Loic Matthey, Christopher P. Burgess, and Alexander Ler- chner. Constellation: Learning relational abstraccioons over objects for compositional imagination. arXiv preprint, 2021a. URL /
James C. R. Whittington, Joseph Warren, and Tim E. J. Behrens. Relating transformers to mod- els and neural representations of the hippocampal formation. International Conference Learning Representations, September 2021b. URL https: //tttpee///penreviee..
BBDVo981YBQ
Daniel L K Yamins, Ha Hong, Charles Cadieu, Ethan Solomon, Darren Seibert, and James J. DiCarlo. performance-optiiizze hierarchical models predict neural responses in higher visual cortex. Proceedings of the National Academy of Sciences of the United States of America, 111 (23):8619-8624, 2014. doi: 10.1073/pmas. 1403112111.
Guangyu Robert Yang, Madhura R. Joglckar, H. Francis Song, William T. Newsome, and Xiao- Jing Wang. Task representations in neural networks trained to perform many cognitive tasks. Nature Neuroscience, 22(2):297-30 February 2019. doi: 0098//4159--08-0331000 URL
htttp:://ww.nntt /tt.:aaarrr........r.reaea//////
Ilker Yildirim, Mario Belledonne, Winrich Freiwald, and Josh Tenenbaum. Efficient inverse graphics in biological face processing. Science Advances, 6(10)caax5979, March 2020. doi:
11116//ciadv.aax5997 URL http::/wwww.iieeeee...oo//o///// aax5979
Dominik Zietlow, Michal Rolinck, and Georg Martius. Demystifying Inductive Biases for (Beta- JVAE Based Architectures. In Proceedings of the 38th International Conference Ma- chine Learning, pp. 12945-12954, PMLR, July 2021. URL https: http://:///ooceeeengg.g..tt .ml.
press/v139/2ieet wttp://////eeeeeeee html
17

---

Published as a conference paper at ICLR 2023
A APPENDIX
A.1 DEFINITIONS
Disentanglement definition. We define disentanglement when single model neurons care about single ground truth factors. We do not mind if more than one model neuron cares about the same factor. We that other definitions of disentanglement are when single ground truth factors of variation are encoded in at most a single model neuron we do not ask for this level of parsimony,
Disentanglement metric. Our metric (mutual information ratio; MIR) measures the mutual infor- mation between neurons and factors, In,f, then calculates each neuron's 's preference (we exclude inactive neurons) for one factor over the rest via
max (In,f)
Tn Ej In,
We then divide by the number (active) neurons, nas and normalise for the number factors, ny i.e.
En Fu 12
MIR
my
(6)
High MIR means single neurons show high preference for single ground truth factor.
We do not propose this metric as a gold-standard metric for disentanglement, indeed in some situ- ations such as spiral shaped functions, it could exhibit pathological properties. Furthermore since our metric is normalised by the sum mutual information of all factors, rather than the just the gap between the factors with 1 and 2nd highest MI (like the MIG metric), it appears lower even when with near perfect disentanglement. However in our situations it works just fine. Again, we use over other metrics since other metrics penalise situations where more than one neuron codes for ground ground-truth factor.
18

---

Published as a conference paper at ICLR 2023
A.2 RELATED WORK
Here we review the literature on disentanglement and the use of biological constraints, in both machines and in the brain, and where possible relate back the work presented in this paper.
While there are many technique for learning linear features of data, by far the most relevant to this work is independent component analysis (ICA), though we briefly mention principal component analysis (PCA) here too. ICA seeks to learn source signals in data that are maximally independent (Comon 1994). and there exist algorithms (Hyvärinen & Oja 2000) that provably extract these com- ponents when the underlying sources are linearly mixed in the observed data (provided the sources are non- manggaasiinnn. PCA, on the other hand, aims to find the minimally correlated components of the signal that explain the most variance, and simple algorithms exist for PCA too.
A.2.1 LINEAR DISENTANGLEMENT
1982) is simple learning rule that learns PCA can be done in neural hardware too. Oja's rule (Oja
principle components sequentially. More recently, neural algorithms for learning the principal com- ponents non-sequentially have been developed, either using Hebbian/anti-Hor -ebbian learning rules in biological networks (Pehlevan et al. 2015) specific weight regularization in linear autoencoders (Kunin al., 2019).
Neural algorithms for ICA also exist. These draw on ideas from nonnegative ICA (Plumbley 2002 2003) where the sources are assumed to be nonnegative and use a similarity matching objective to derive biologically plausible networks for ICA (Lipshutz et al. 2022 Pehlevan et al. 2017), with this extendable to correlated sources (Bozkurt et al.) 2022). These works relate to ours as they consider nonnegativity, however their focus
biological neural networks implementing linear ICA, whereas we mathematically prove biological constraints including nonnegativity lead to source separation in linear ICA, and empirically show that the same is true from nonlinear ICA.
A.2.2 NONLINEAR DISENTANGLEMENT
There exist no known algorithm for exact source recovery ('identifying' the sources) when the fac- tors are nonlinearly mixed. Indeed it is provably the case that nonlinear ICA is not identifiable (Hyvärinen et al. 1999).
Nevertheless, there have many methods for disentanglement in nonlinear models, with most modern methods use variational autoencoders (VAE; Kingma Welling (2013), with various choices that promote explicit factorisation of the learned latent space. For example 3-VAEs (Higgins et al. 2017a) up-weight (by 3) the term in the VAE loss that encourages the posterior to be a factorised distribution. Other variations of disentanglement VAEs (Burgess et 2018 Kim & Mnih 2018
Ridgeway & Mozer 2018; 8; Kumar et al. 2018 Chen et al. 2018) similarly try to explicitly enforce &fcttoried aggregate posterior.
Importantly, analogously to the nonlinear ICA case, it has been shown that disentanglement is not possible without inductive bias in the data or model (Locatello al. 2019). and it has been shown that generic VAEs are only able to disentangle when there is a particular bias in the data (Rolink) jt al. 2019 Zietlow et al. 2021). However, it has also been shown that numerous inductive biases are sufficient for disentanglement (i.e. conditioning the prior on a third variable Khemakhem et al. 2200 or local isometry and non-Gaussianity Horan et al. (202T))
Though, in our work, we show disentanglement in VAEs with biological constraints, we also show disentanglement in supervised DeepNets and linear and non-linear autoencoders. Empirically, seems out constraints are not limited to the VAE case, though further and more intensive empirical investigation are needed to confirm this.
Our work also relates to modern self-supervised learning (Bardes et al. 2022 which seeks decor- related representations with non-zero variance. Similar to many of the VAEs that disentangle, they introduce loss terms specifically to force decorrelation. Here instead, show that simple biologi- cal constraints of nonnegativity and energy efficiency lead to emergent disentanglement without an explicit decorrelation term in the objective.
19

---

Published as a conference paper at ICLR 2023
There is already evidence of disentanglement in machines emergent phenomenon opposed to being baked in, for example interpretable neurons in vision model (Bau et al. 2017). multi-modal neurons in CLIP (Goh et al.) 202I), and other work by superposition/ privileged axes (Elhage et al. 2022) that also discusses the role of nonnegativity constraints. Similarly, other work is on understanding specialization in DeepNets (Casper et al. 2022).
A.2.3 DISENTANGLEMENT IN BRAINS
Barlow first hypothesised that neurons represent independent components via redundancy reduction
miiimising mutual information between their representation (Barlow 1961 1972 Since Barlow there have been numerous experimental results suggesting that brain neurons disentangled: in inferior temporal cortex (Chang & Tsao
2017 [aoo et . 2020 Higgins et al. 2021] Yildirim et al.,
2020), visual cortex (Fcker et al. [2000] Gáspár et al. 2019), parietal cortex (Lee al. (022), and
entorhinal cortex (Hafting et al. 20055 Solstad et al. 2008 Høydal et al. 2019). Similarly, however,
there have been many reported neurons in various brain regions that are apparently mixed-selective
(Rigotii et al. 2013 Boccara et al. 2019 Butler et al. 2019 Komorowski et al. 2009).
There is no formal understanding of when and why neurons in the brain factorise across task param- eters. Nevertheless, single latent dimensions from disentanglement models do predict single neuron activity, implying biological neurons are disentangled (Higgins al. 202T). One conjecture in neu- roscience is that while representing task factors is important for generalisation, mixed-selectivit is important for efficient read-out (Behrens et al. 2018 Rigotti et al. 2013 Bernardi al., (020). Here show disentangled brain representations are preferred if the task is factorised into independent factors. This work bridges the gap between a single neuron understanding (i.e. like Sherrington) and population based understanding (i.e. like Hopfield) (Barack & Krakauer 2021).
A.2.4 BIOLOGICAL CONSTRAINTS
While these constraints of nonnegativity and energy efficiency are not new to the machine learn- ing or neuroscience literature, we package them together learn new insights into how biological constraints shape the space of neural representation.
Nonnegativity. Recent work on multitask learning recurrent neural networks (RNNs) (Yang al. 2019 Driscoll et al. 20022) demonstrated that neural populations, with nonnegative activa- tion function, partition themselves into task specific modules (Driscoll et al., 2022). Nonnegativity is also important in obtaining hexagonal, not square, grid cells (Dordek et al., 2016; Whittington jet al. [202b]] Sorscher et al. 2019) 2020). Also, nonnegative matrix factorisation empirically yields spatially localised factors for images (Lee & Seung. 2000). Other work discusses potential algo- rithmic functions of nonnegativity in again Hebbian/anti Hebbian networks, relating it to clustering, sparse feature discovery and again ICA (Pehlevan et al. 2015). Our work demonstrates theoretically and empirically why nonnegativity leads to single neurons becoming selective for single factors.
Energy efficiency. Using regularisation is common in machine learning, with or 11 losses on weights primarily used. Less commonly do machine learners ask for energy efficiency on activa- tions, though the VAE loss does include energy efficiency latent activations, and additionally some activation functions can be see as energy minimising, e.g. ReLus are sparsity inducing. In neuro- science, there is much work considering energy efficiency (see Chintaluri Vogels (2022) for fun recent example). There are also results more relevant to us, for example work that test whether ori- entation of neural manifolds from mouse V I minimise energy efficiency (spike count) by randomly rotating and optimally shifting them to fit into nonnegative orthants (Bordelon & Pehlevan 2022).
'ee note that Barlow also considered parsimony as a representational pressure too. While we do not consider it in this paper, we believe it is important constraint.
20

---

where M Rink are mixing weights and b2 E R is bias. We further assume two constraints: (I) the neural representation is nonnegative with 2, 2 for all i I,...," n, and (2) the neural C, so that the neural representation Published as a conference paper at ICLR 2023
A.3 DETAILS OF DERIVATIONS
A.3.1 PROOF OF CONSTRAINTS LEADING TO DISENTANGLEMENT FOR A GIVEN POPULATION VARIANCE.
Theorem. Let e E IR* be a random vector whose k independent components denote k task factors. We assume each independent task factor e; is drawn from distribution that has mean variance a², and maximum and minimum values of min(e:) = -a and max(e;) = Also let z E RR" be linear neural representation of the task factors given by
= Me + b2
(7)
population variance is nonzero constant, E, Var(z,) retains some information about the task variables. Under these two constraints we show that in the space of all possible neural representations (parameterised by M and b²), the representations that achieve minimal activity energy ||22 also exhibit disentanglement, by which we mean every neuron 2j is selective for at most one task parameter: i.e. [M Mad = for + 1.
Proof. We aim to find representation, that minimises activity energy (the expected norm of 2), nonnegative, all for fixed population variance, C. This a constrained optimisation problem, and equates understanding what M and b. must look like in order to satisfy our constraints, and minimise E ||
minimise 2| s.t. 20 Var(z))) C M.b.
(8)
The total activity energy, i.e. the expected norm of z, is
E 1| = Σ B(2z)
Σ Var(z,) + (B2z)²
(9)
Σ (ar((Me)) Σ(bb))))
We want minimise this, under the constraint 2j 0. To satisfy this constraint, (bæ), must account for any negativity. This means that
(b.); 2 - mimmmin(MM)) -min(Me)))) = Σ JM
(10)
Where the last equality is because all e are i.i.d. and so the minimum of sum of random variables is the sum of their minima. The modulus sign is since Mÿr can be positive or negative, so we need to consider the maximum and minimum ek, and an assumption of ours was that the maximum and minimum had the same value a. Thus
(b.);
Σ Miaaa P2
(11)
Where Py 2 0. Intuitively, = O since anything else would increase activity energy. Formally,
|||
V Var((Me),) + ((b)))))
d²
Σ
M3
Σ Σ JMgxla + p2)²
(12)
j.k
Σ M3
Σ J
[M + + + +p3 Σ M
+
1.k
Since both En [Maala and pj are greater than zero, Pj always increases activity energy, no matter what Mss is. Thus set it to zero, i.e. Pj 0. Hence we can simplify the expected activity as
21

---

Published as a conference paper at ICLR 2023
follows
|||=1|| = 2 Σ
M3 +
Maala)}
J.K
3 &
o² Σ M3 +
> MMaxaa)
M2r
J.k
&
0²
M + a Σ (M + Σ
(13)
j.k
(1
Var(zz) + +
Σ
j
jkkllk
d² = =11 9 C + a Σ J.
Ma
Now all the constraints have been incorporated, our only job is to minimise E| This is done when Lym.,ty |Mai|| = and that only happens when Mjx |Mal 0 for all and of k. In words, this means that neuron j in only receives information from element of e. This is disentanglement.
A.3.2 PROOF THAT POPULATION VARIANCE BOUNDED BY READ-OUT WEIGITTS NORM
Theorem. Let a De be observed entangled data, where x E Rm, D E Rmnk and e E R is random vector whose & independent components denote k task factors. We assume each independent task factor &; is drawn from a distribution that has mean O, variance o², and maximum and minimum values of min(e;)
Let neural representation z R exactly predict observed data via x = W x + b with zero error, i.e. a = De Wz + bg. Where W E R mxn k are read-out weights and b, E RM is an offset.
and max(e,)
Then for all such data generation models (with parameters D) and all such neural representa- tions (with parameters W and by), as long as: (1) the smallest singular value ofD is non-zero, minn (D) > 0; (2) the norm of the read-out weights [|W]]3 is finite; then the population variance of is bounded from below by the norm of the read-out weights.
Σ Var(z;) 2 2 min
min
(14)
W
F
Proof. Since x Wz + bz, we have
2 - == x W by = = De W
(15)
Where W+ is the Moore-Penrose pseudoinverse. Additionally since the read-out error is zero, W must contain all the same information as D, and so must be W DF+, where F+ E Rgxx is matrix with rank & (n k). The pseudoinverse of W is thus W FDt. The population variance becomes
Σ Var(zz) Tr(Var De w+b)))
Tr(Var(FD** De))
=Tr((Va(Fee)) =Tr(FVar(e)FF)
(16)
= =o²Tr(FF")
= =
Using the fact that the norm of matrix's pseudoinverse is bounded by the norm of the matrix
22
F+1 | = Tr((( F
F)
(17)
F
22

---

Published as a conference paper at ICLR 2023
Along with using the following trace identity (Fang et al.
1994)
F F Ippo min (D) & W 1
DF
|| S F
pro mome (D)
(18)
Where minn ()) and maa (D) are the smallest and largest singular values of D. Thus
0
min
D
(19)
F
W
P
Combining all together
Σ of
1:022
(20)
|
0²0000 min
D
W
We note that the first inequality becomes an equality if and only if F+ has (scaled) orthonormal rows, and the second inequality becomes and equality if and only if D has (scaled) orthonormal columns.
A.3.3 PROOF OF DISENTANGLEMENT WHEN DATA GENERATIVE MATRIX HAS (SCALED) ORTHONORMAL COLUMNS
Theorem. Let De be observed entangled data, where vector whose è independent components denote k task factors. We assume each independent task factor 6, drawn from a distribution that has mean O, variance 0², and maximum and minimum values of min(e:) observed data via æ W b, with zero error, i.e. æ De Wz z + b,. Where W E mxx m 2 k, are read-out weights and b, E Rm is an offset.
mm and e E R* is a random
-a and max(e;)
Let a neural representation z R exactly predict
Then for all such data generation models (with parameters D) and all such neural representations (with parameters W and b,), as long as: (1) the columns of D are (scaled) orthonormal; (2) the norm of the read-out weights ||| is finite; (3) the neural representation is nonnegative (i.c. > 0), then out of all such neural representations, the minimum energy representations are also disentangled ones. By this we mean that each neuron 2; will be selective for at most one hidden task factor Ej-
Proof. We aim to find a representation, 2, that minimises activity energy (the expected norm of x), is nonnegative, and predicts data æ De via Wz b, with zero error (De Wz z b,). This a constrained optimisation problem, and equates to understanding what W and b (and therefore z) must look like in order satisfy our constraints, and minimise E/| =||?.
minimise W b,
E||z||2
s.t. 2, 0, De Wz z + b Ve. W
(21)
We note that this is the classic matrix factorisation setting. The difference here is that we ask the representation to be nonnegative.
Firstly, we note that since the columns of D are (scaled) orthonormal, then the singular values are all equal: min (D) = mmax (D) = DD). This result is useful as now the inequality in equation
19 becomes an equality
1
O²( |
(22)
|
Now we prove disentanglement. Since the read-out enror is zero, W must contain all the same information as D and must be W DFF. Using this, and repeating similar process to the
23

---

Published as a conference paper at ICLR 2023
first proof (Appendix A.3.1
get
B||=|| = (1 + a²
Σ
Var(z)) + Σ
DWs+llll(W"++ D)ad
= (o + a
FIF. + a E
J.
(23)
ha
2 2 + +
a² Σ
F
k2 (D)
(a² + )
a Σ
j.K
Where the inequality is due to equation| 17 This inequality becomes an equality if and only if F+ has (scaled) orthonormal rows (i.e. has (inverse scaled) orthonormal columns).. This means that for a fixed |W the first term in equation 23 is minimised when F+ has (scaled) orthonormal rows, or equally when has (inverse scaled) orthonormal columns. This is interesting to us, as we can additionally make the second term in equation 23 go to zero when F is a particular type of matrix with (inverse scaled) orthonormal columns. Thus we will have fully optimised equation 23 and this will be our solution. First, rewriting F matrix with (inverse scaled) orthonormal columns
F F-Lo and F+ =+=a0" a
(24)
Where O a matrix with orthonormal columns. The particular that minimises the second term in equation 23 is when O only has, at most, single non-zero element per row as this sets FFx P3 This is disentanglement and is casily seen by
=W* W+b. b,
- De W+b. b,r
(25)
=oOD De aOD b
=00- a0
Since O only has a single non-zero element per row, then each 2, must only contain single element (random variable) from e. We note that nonnegativity is easily achieved by b, learning to take value that satisfies nonnegativity, i.e. a0" D by min min Since O and D are full (k) rank, this is always possible.
We also offer an alternate proof in Appendix A.3.5 when analysing 'Simplification 1'.
A.3.4 PROOF THAT A TALL RANDOM MATRIX WITH FINITE WIDTH HAS (SCALED) APPROXIMATELY ORTHONORMAL COLUMNS
Theorem. Let D E mx be random matrix with elements Dij that are i.d.d. with expectation 0 and variance x²/ x²²mm²). Then as with finite k. D D I.
Proof.
lim DD lim m--00
EDjjDjt
= mm2EDj:Dp.
(26)
R2 i=
00 / k
Thus D has orthonormal columns, scaled by and so its singular values are all identical;
amin (D)
Oman D) 2
A.3.5 WHEN THE DATA GENERATIVE MATRIX DOES NOT HAVE (SCALED) ORTHONORMAL. COLUMNS
While we do not prove the general case, we offer some intuition here that suggest disentangled representations are favoured in many situations. Consider the general setting in which D E Rmak
24

---

Published as a conference paper at ICLR 2023
has the following singular value decomposition (SVD):
D=
(27)
Where Σ R is rectangular diagonal matrix with positive entries, and U Rmx and V E Rixk are orthogonal matrices. As in the above proofs, since is predicted with zero emoor from 2, via
x == D= Wz + b..
(28)
Thus W must be of the form W = DF+, where F+ E Rkxn is rank k matrix, since z E IR", R", and E RR m k. We define the SVD of F
OAR F RA O
(29)
Where E &xxn rectangular diagonal matrix with positive entries, d A-1 E R rectangular diagonal matrix with positive entries, 3. and where O Rx" and R Rx orthogonal matrices. From equation the population variance is
= Σ (
(30)
= Σ X
Where a is the variance of the random variables ej. The norm of the weights W is
W 3
DF+
IUVV RA
'0
-EEV" RA
(31)
(D) X3
Vna Ray Rij
M
While we have been previously interested in finding the minimal activity energy with fixed | | we instead allow W to change and instead minimise W
L = 22| + 3.oll ||| where
Br is the strength of regularization the weights. Thus using equation 23 have
of(D)
L (o² a²)
Σ x
Σ 1A
30D
Σ kJ
Ve, Va
(32)
J.R.EER
This is the overall thing we want to optimise. However, for now we not consider the the middle term the interaction induced by the nonnegativity constraint and instead consider the following optimisation problem
OD X3
minimisc L'= = -022 2 V,R.A
X² + B Σ
Σ W!
ViaVi
(33)
Under the constraints that V and R remain orthogonal and remains rectangular diagonal with positive entries. This is difficult in general, but simplify to gain intuition and insights.
Simplification 1. The first simplification is when E is a scaled identity matrix. This is the same simplification used in Proof 3 (Appendix A.3.3). i.e. the singular values of D are all equal, (D) = (D).
Simplification 2. The second simplification is that V is the identify matrix. This corresponds to data being generated via = UE, ie. the random variables are scaled and then orthogonally projected.
25

---

Published as a conference paper at ICLR 2023
Analysing simplification 1 first, to build up intuition (and offer alternative proof of Theorem 3), L' becomes
oD
C'= (o² + a
+ Sew
Ved
1/
= (0² + 2 Σ
X + Ban 0 (D) Σ
REs Rij E
ViaVie
X
(34)
= (o² + x Σ
X7 + 3m (D) Σ X3
Σ R3
= =o²² + d²) Σ
2 + o (D)
Σ
X3
Where we exploited the fact that R and V are orthogonal matrices. Optimising this is easy to do, we can simply take derivatives to get
(D) 1 = r² +
(35)
This result is independent of R and O, and so we are free to choose whatever orthogonal matrices we like. This is good news for (equation 32) which contains an additional term involving R and O: a² Ljnsya |F = >jl]- Thus if we can set ((ARR)) 3 (OARR }jl 0, then will have fully minimised L. This is easy enough to do (keeping R and O orthogonal) and is achieved when R is a permutation matrix, and O has at most one non-zero element per row (or equally O has at most non-zero element per column). This corresponds to disentangled repre- sentations.
the real game are in is minimising L
a² Dj..1ik (OAR") (OAAA""
Returning simplification 2, where V is the identify matrix, now L' becomes:
L' (o² + a²)
X² + B
V&s VisVi, Rij
kl
(36)
o
= (0² + a²)
x
R
Thus we have the following constrained optimisation problem
R
minimise
a² X2 Bu
Σ
S.L. 1 0, R
(37)
R.A
X?
Here we let intuition take over. Taking inspiration from the simplification 1, ansatz is that 2-40² (where and j are related by the permutation). To justify that R should be permutation matrix in this case, we note that as X gets smaller to keep ||W11 Ei of (D) 1 x as small as possible, the Rag must ensure the low valued X3 are matched with the low valued of (D). Intuitively this done when V is permutation matrix. Once (D) R is chosen as a permutation matrix, showing that X = 224.2 is simple just take derivatives and set to zero. While this is not full proof, one derives the KKT conditions, this solution is at least a consistent solution, i.e. it is a minima but may not be the global minima. We note that for specific (small) values of 3.w. it will be the global minima.
(
R is permutation matrix and X
As before, our actual aim is to minimise L. Since we are free to choose O, and choose it to make the cross terms in equation 32 go to zero, which is when O has at most one non-zero element per column. This corresponds to disentangled representations.
No simplifications. Reminding ourselves of the full objective
L= L= (a²
a²) Σ
X3 + a²
Σ J.R.
|F
Fuul +
3
Σ
VeA Vi
(38)
X2
kl
26

---

Published as a conference paper at ICLR 2023
In this case V is an arbitrary orthogonal matrix, and (D) the (not necessarily equal) singular values of D. One potential ansatz is to assume R is permutation matrix once again, which then reduces the problem to simplification 2. Again this offers a minima, but not necessarily the global minima.
In sum, see that there is always pressure disentangle due to the middle term in the loss. However it will have to trade-off against the weight regularization term (last term). Thus we posit that for small values of Sw disentanglement preferred, even for arbitrary D.
27

---

Published as a conference paper at ICLR 2023
A.4 DETAILS OF GENERATED DATASETS
We note that no domain knowledge, like the exact number of factors, is provided too our models any point. However the dimensionality of the network layer must be least this number to disentangle all factors.
Disentanglement in supervised shallow neural networks.
Starting with 6 independent variables, each sampled from a standard uniform distribution, i.c. E [0, 1] We construct a dataset with æ Oe, where O random orthogonal matrix (random for each model/seed). The target Ie is generated from We where W E IR has elements drawn independently from standard normal distribution (W different from each model/seed). We generate 200000 data points.
Disentanglement in supervised deep neural networks.
This is generated in the same way as above but we additionally set æ + x², where the cubing is done elesent-wise.
Disentanglement in unsupervised shallow neural networks.
Starting with 6 independent variables, each sampled from a standard uniform distribution, i.c. E [0, 1|6. We construct dataset E 1R50 with y = We, where W E g50x6 has elements drawn independently from a standard normal distribution (W different from each model/seed). We generate 200000 data points.
Disentanglement in unsupervised deep neural networks.
Starting with 6 independent variables, each sampled from a standard uniform distribution, i.c. e [0, 1]6. We construct dataset y R:50 with y M (e)), where the MLP is 1-hidden layer network with relu activation function on the hidden layer. The hidden later has dimension 28, and all weights of the network are drawn independently from a standard normal distribution (redrawn from each model/seed), We generate 200000 data points.
28

---

Published as a conference paper at ICLR 2023
A.5 NETWORK ARCHITECTURES AND OPTIMISATION DETAILS
All code will be released on publication. All hyper-parameters are shown in Table 1 We describe any additional details here. We the Adam optimiser (Kingma & Ba 2014). Supervised nets square error loss. Unsupervised nets use sigmoid cross entropy loss.
Supervised networks. Cnonnee Lactivity and Lweight are applied all layers. The shallow Network is a I-layer MLP with hidden dimensions of [188]. The deep networks are 5-layer MLP with hidden dimensions of [188,186,188, 18,,,88,11..1111 182,1888, We use squared error loss for Sprediction
Autoencoders. We apply monney and Laciivity to the latent layer only. We apply weight to all layers.
The shallow autoencoder has hidden layers in the encoder or decoder. The deep autoencoder has and encoder with hidden dimensions of [500,300,100 and a decoder with hidden dimensions of (100,300,500] The output a sigmoid and we use sigmoid cross entropy loss for Lprediction
Variational Autoencoder: We use an architecture described in Higgins et al. (217a)) (Table We use the standard 3- VAES loss with the following additions. We apply Lnonney to the latent layer only.
We apply weigght to the dense layers in the decoder only. We do not apply Lactivity as effective norm constrain term already exists in the 3- VAES loss.
Supervised Net
Autoencoder
VAE
Shallow
Deep
Shallow
Deep
Deep
input dimension
6
6
50 n/a
50
(64, 64, 3)
output dimension
6
n/a
n/a
# parameters
2450
138574
1570
414870
766295
(hidden) layers
5
enc:
enc: 3
enc: 6
dec:
dec 3
dec 6
latent dimension
n/a
10
10
learning rate
3e-3
le-4
le-4
batch size
128
64
64
# gradient updates
300000
300000
500000
Bactivity
1e-3
5e-3
n/a
Bweight
le-4
le-3
0.01 + 1.0
Bnonneg
2.0
0.5
100.0
BYAE
n/a
n/a
1.0
Table 1: The values of various hyper-parameters. enc/dec: encoder/decoder. n/a: not applicable.
Encoder
Decoder
Input: 64 x 64 number of channels
Input: 10
4 x 4 conv, 32 ReLUU, stride
FC, 256 ReLU
x 4 conv, 32 ReLU, stride 2
FC, 4x4x64 ReLU
4 x 4 conv, 64 ReLU, stride
x upconv, 64 ReLU, stride 2
x 4 conv, 64 ReLU, stride
x upconv, 32 ReLU, stride 2
FC 256
x upconv, 32 ReLU, stride 2
F2 2 x 10
x upconv, number of channels, stride
Table 2: Encoder and decoder architecture for the variational autoencoder experiments.
29

---

Published as a conference paper at ICLR 2023
A.6 SPARSITY DOES NOT PROMOTE DISENTANGLEMENT
Readers may wonder if imposing sparsity constraint would encourage disentanglement, as ReLUs promote sparsity well as enforcing nonnegativity. This is not the case, and be understood intuitively, and also with numerical simulation. For intuition, a sparsity constraint creates a diamond shaped iso-contour in neural space (Fig. 8 left, and encourages the firing rate distribution to fall inside that diamond, which in the case our our random variables, means maximal entangling. We confirm this intuition simulation (data and setup the same as Fig. 2 with a variety of sparsity regularization strengths (Fig. right... Thus the reason ReL.Us promote nonnegativity is not because they induce sparsity (which they do), but because they enforce nonnegativity.
Neuron
Neuron
Earars Paxvo 0.0000
Auru + D.COOL
Raaoo 0.003
Agin DCC01
Satisfy sparsity
Peam=000
D.C3
Paxy=0.1 m011
Do.
Neuron
Neuron
50000
100009
15:000
200000
250009 23000
Stop
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.
Similarly, one may wonder if sparsity on weights may encourage disentanglement. We do not ob- serve this empirically (Fig. 9)
1.0
Prpasse 0.0001
Eyctity Papase 0.0001
Appace 0.001
Castaty Papaze = 0.001
Baoas 0.01
Captaiy: Papass m 0.01
Aawe 0.1
Escivlyy Papase - 0.1
50000
100000
150000
200000
250000
300000
Gradient Step
Figure 9: Sparsity weights does not lead to disentanglement.
30

---

Published as a conference paper at ICLR 2023
A.7 REPRESENTING CATEGORICAL DISTRIBUTIONS
We note that this is a fundamentally different problem to representing factors, since only one cate- gory is active at any time; categories anti-correlated. This contrasts factors all factors can be active independently. So even though the proofs are similar, they are about different situa- tions.
Theorem. Given nonnegative representation, 2 E R", that linearly reads-out population, æ e I which is itself linear projection categorical variable (one-hot variable)
x == and a Cc
(39)
Where E R is one-hot representation (with each one-hot vector having equal probability) with each element, C₁, representing category identity (1 when present. otherwise, only one category active at once), W E R= read-out weights with fixed norm, ||W||2 and C' E mm are the data generative weights which is a rank e matrix with all singular values equal and positive; miin (C) Omaz
(CC) 0. Then in order to minimise the expected activity energy, B| 2 must structured so that each neuron only represents at most single category.
Intuition. In order linearly read-out categorical variable from neural population, z, then must take be a linear combination of a one-hot vector (c; each element in the one-hot vector corresponding to each category.
== Mc
(40)
For this representation to be nonnegative, all entries of M must be nonnegative. This means that the population response, zi. for each category, i, must be vector in the positive orthant (Fig, I0] Assuming we keep the expected activity energy constant, i.e.
| = - Σ || 1
|| = constant
(41)
Then the easiest representation to read-out from is the one-hot representation, since in all other situations the vectors are closer together which require larger read-out weights to disambiguate. Alternatively, if there is noise on then the categories will be more often confused if the ze vectors are closer to each other. A one-hot representation is the best, and it is disentangled representation for each category.
Proof. We aim to show when minimising =| when & is nonnegative (2; 2 0), and for a finite norm of the read-out weights, ||W],, then it is optimal for each element, 2; in to represent most a single category. This a constrained optimisation problem
minimise E| 2||
s.t. 0, Cc Wz Vc
(42)
Since we read-out with zero error, Cc = Wz, then W must contain all the information of C, i.e. W CF+, where F+ Rexn is rank matrix. Thus z must be
z W* æ
FC Ce
(43)
Fc
M
Figure 10: Schematic for representing categories. Each axis is neural firing rate. Each colour is population activity for single category, i.c. zz.
31

---

Published as a conference paper at ICLR 2023
For this to be nonnegative, since is one-hot, then Fij 2 0. Additionally, the expected activity energy becomes
|z||
El
Fc 1
1 FIl
1 o²
(44)
c cI
1 Cc)
||
Where the final equality if from equation 22 and where the inequality is from equation 17 This inequality becomes equality if and only F+ has (scaled) orthonormal rows (or F has (scaled) orthonormal columns). Thus to minimise |||2||2 we want F+ to have (scaled) orthonormal rows or equally F have (scaled) orthonormal columns. To satisfy nonnegativity we wanted Fij 0. There is only one type of matrix with (scaled) orthonormal columns that has nonnegative entries, and that is when has rows that contain at most one non-zero element, and so each neuron only "attends" to one category; categories are disentangled.
A.7.1 SIMULATION RESULTS
We train an autoencoder on data from categories. The data from each category is noiseless, so we essentially just have data samples that we train on. Each category vector is sampled from uniform distribution. We train an autoencoder with deep encoder (hidden layers: [500,300,100]) and shallow decoder (O-hidden layers), and with latent dimension of 10.
Along with variants of our constraints, fiparsity = Ssparsity E; lagl. where a, is the activity of neuron the latent layer. All other hyper-parameters are as described in the autoencoder section of Table and we set Asparsity to be same as Bronneg.
also train network with sparsity inducing loss:
We see that only our constraints, and the sparsity inducing loss, achieves disentanglement of cate- gories (Fig.
A
B
0.7
0.6
0.5
Latents
0.4
4
MIL
0.3
ReLu, Exctiriyy Cweight
Exctivity: Lweiihh
0.2
Cnanneg Cactrutty Linight
Csparse
ReLu, Lactivity
Exparzer weiight
0.1
ReLU, Lwoight
0.0
factorg
factor1
factor2
factor3
factord
factors
50000
1000C0
150000
200000
250000
300000
Gradient Step
Factors
Figure 11: Learning data generative factors with autoencoders. A) Training linear autoencoders on línear data. Only models with our constraints, or with sparsity, learn disentangled representations. B) Example mutual information matrix from high MIR model. All learning curves show and standard error from 4 mean from 5 random seeds.
32

---

Published as a conference paper at ICLR 2023
A.8 FURTHER VAE SIMULATIONS RESULTS
Figure 12: Latent space traversals. An image is encoded and then the value of single latent dimensions is changed, and the resulting image is generated. Each row shows this image for when 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 single ground truth factor.
16.0
Eita
Auam Auae
10:00
Figure 13: Same as Fig. but for the dSprites dataset. A) We train on the dSprites dataset, with two example images shown. These images have underlying factors. B) MIG scores are higher with higher weight regularization, and generally higher than any 3-VAE (panel D). C) Mutual information matrix for a high scoring model. We note has dropped a latent dimension. D) 3- MIG scores. We note the high disentanglement 3- VAV models have bad reconstruction as they drop latent dimensions. E) MIG score against R2 shows models with our constraints lie the Goldilocks region of high disentanglement and high reconstruction. All learning curves show mean and standard error from random seeds.
33

---

Published as a conference paper at ICLR 2023
A.9 PATTERN FORMING DETAILS
A.9.1 DETAILS ON TRAINING DATA
A factorised task for rodents. We generate many 16x16 worlds (so 256 locations; 256) where objects can be at random locations. The number of objects in each world is chosen uniformly between 0 and 8. In practice it is not always possible to place every object as we do not permit objects to be closer than units away.
From each world we have the following data: {ar, o(x), a(x)) for all 256 locations æ, where æ is one-hot identifier of location, o(z) is a binary variable of whether there is an object or not at location H, and ((x) is 4-dimensional vector describing the optimal action corresponding to North/Southh/asttWea Optimal actions towards the closest object.
zero in each entry (more than one action could be optimal at each location).
We train on this data, with batch size of 8, i.e. presenting such worlds simultaneously. Each world in the batch is randomly sampled from the above process.
entangled task for rodents. This is the same as the above, except we sample only 1 world throughout training, i.e. every world in the 8 batches are identical and all batches identical.
A.9.2 MODEL DETAILS
We use discrete 16x16 world (so 256 locations; n4 256) and optimise independent represen- tation, (xx) E RR at each location. We now detail each component of the following loss
C Caonneg + Cactiviyy + Cwoight + Liocation + Lactions + Lobjects +
Catt integration
(45)
Biological constraints
Functional constraints
Structural constraints
Biological losses. These are exactly the same as in the main paper, but we must also average over all locations, x
Luunneg
Pmonneg
ΣΣ max(-2; 2x(2),
(46)
74
Eweight
Bweight
Wi
(47)
Puetivity
Lactivity
Σ
|xx)|
(48)
where 2x) is a neuron in representation =(x)). indexes the task (i.e. object, action, location prediction), and the values determines the regularization strength.
Location loss. The representation predicts location (a one-hot encoding describing each of the 256 locations) via linear transformation, which is then fed into a softmax cross-entropy loss. In particular, the logits for each location, æ, are Wiz(æ), where Wi E TR XMe, , If we denote each row of Wi as la noting that this row that 'corresponds' location x in the nne-hot encoding, then the loss is as follows
Llocation
Biocation
Σ In End ==(()
(49)
Object loss. An object is either present or not present at each location, so we use sigmoid cross- entropy loss. In particular, the logits for each location x Wo=(æ), where W, E Rlxn, oobject at returns 1 if an object is present location, æ. and 0 otherwise, then the loss is as follows
Bobject
Cobject
In (W.z(()) Lobject not 1nn1n111110 (W,
(50)
T2
Action loss. More than one action be correct at location, for each of the 4 actions. In particular, the logits at location, æ, are W. ((x), where W, E I2 X",, where nh is the number of actions (4 in our case North, South, East, West). We denote each row
sigmoid cross-entropy loss
34

---

Published as a conference paper at ICLR 2023
of W as t,, noting this row 'corresponds' to action in the action encoding. If 1 la=a(m) returns action, a, is a correct action at location æ, and 0 otherwise, then the loss as follows
Paction
Lection
ΣΣ )a=a(a) Ino =<x)) + (1 ===a()) In(1 ota =(x))) (51)
T1
Structural loss. We use a squared error loss for the structural constraint, which asks for neighbour- ing representations to be related to each other by action matrix Wa for each action, a. This is just like a path integration loss. This loss done for every location, æ, and each of the 4 actions, a.
path integration
path integration
ΣE1==0)
1=0(1) ( f(W,=(1
d,)ll²
(52)
T4
Where W. E Rn weight matrix that depends action, matrices one for each action). du means the displacement in the underlying space (the space of æ ). that the action a corresponds
o, (i.e. there 4 trainable weights
This loss is effectively path integration loss and can be understood by considering the path inte- gration equation. While there more than one potential forms for path integration we choose the form relating to Gao et al. (2021) which sequentially updates representations in time via:
2, = f(W
(53)
If the representation at time-step 1- corresponds to location - da. then for path integration be successful, the representation at time é must correspond to location æ. Then the equation becomes like equation 3 This is exactly what our loss asks for. We note, however, that this loss on its own could just make all representations at all locations the same constant (and learn all Wa to be the identity), in which case location isn't represented at all. However the location loss ensures all loca- tions are represented differently. So really it is the structural loss and the location loss together that are like path integration. Overall, this formulation translates sequential path integration in time into an optimisation over representation at locations without needing to incorporate time. This greatly reduces the difficulty optimisation and makes the problem more general than just sequences in time.
Pattern forming dynamics. The overall loss can be optimised with respect to the weights. However, it can also be optimised directly with respect to This is particularly interesting for us. it allows our representation to be dynamic and change rapidly for a single task, and not just slowly via learning over many tasks. This is necessity for
need to represent objects which may move between tasks. To optimise both z (task particularities) and weights (task generalities), we do so in two stages. First, we optimise with respect we optimise with respect to the weights to learn parameters that are general across tasks.
to 'infer' representation for the current task. Second,
When optimising with respect to we only optimise terms in the loss: Lobjects and Lpath integration- We optimise the first term so the systems has the ability to know where the objects are. We optimise the second term so that information can be propagated around (effectively via path integration).
The dynamics of the Lobjeess are:
&Lohjects dz(x:)
WT C [lobjcct o (W.z(x)))
objects
(54)
This says if you get the object prediction wrong, then update z better predict the object. We restrict this update to only take place where the object is, so it is just an object signal. This update is equivalent rodent observing that it is at an object.
The dynamics of the Cpah integration are:
dapat intogration
Σ -x((xx) S(W.z(r doo)
dz(x))
(55)
W (z(x + d.,) f(W.=(x)) o f' f((,z((z)
The two terms in the above equation can be easily understood. The first says that the representation each location, z(x), should be updated according to what its neighbours think it should (this is
35

---

Published as a conference paper at ICLR 2023
the same update rule as path integration!). The second term says the representation each location, (x),, should be updated did not predict its neighbours correctly. This equation tells representa- tions to update based their neighbours. This is just like cellular automata, but instead of a discrete value being updated on the hasis of its neighbours, it is a whole population vector whose elements continuous. Indeed, just like cellular automata, it is also possible to initialise single 'cell' (location) of the cellular automata, and have that representations propagate throughout the space. In this case, it's just like path integration, but spreading through all space at once. We note, however, that in our simulations we initialise representations at all locations (for each task).
We note that while simulated this on a discrete grid, the same principles apply to continuous cases. In this case the sums over location/actions need to be replaces with integrals.
This is very general approach for understanding representations. The structural loss does have to related to the rules of path integration. It can be anything. It could be the rules of graph. It could be rules of topology. could have one set of rules at some locations and another set of rules other locations. The rules don't have to be neighbouring representations telling each other what to be, it could also be long range rules too. If there are structure or rules in the world or behaviour, our constraints say that representations should understand that structure or rules. In mathematics this is known as a homeomorphism. In sum, understanding representations via constraints very general.
Syper-parameters. Snooncg. Sweight Snctivity. Siocation Sobjact Saction Spath integration chosen le-2, le- 7, 1e- 3, 1e - 2,1e le 2, 40. The learning rate 12e - In the Euler updates for pattern formation, the step size for path integration is 0.1, and the step size for objects is 1. We additionally clip the norm of the update in pattern formation for stability. The batch size is 8, i.e. we present worlds simultaneously.
36

---

Published as a conference paper at ICLR 2023
A
B
Figure 14: All cells. A/B) Ratemaps from ReLU model. A) A task with objects. B) Task with objects. C/D) Ratemaps from model with linear activation function, C) task with no objects. D) Task with objects.
37


Reviews:
[{"review_id": "gPspyhFMb1", "reviewer": "Reviewer_9LCk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "I am very impressed by the authors' thorough replies to all the concerns I raised. I can see that a lot of thought and work has gone into improving the papers' clarity and claims. I am raising my score accordingly.\n\nThe authors show that non-negativity of neural responses together with an energy constraint encourage disentangled representations. They provide a mathematical proof for this conjecture in two linear settings: linear neural responses, and linear generative models. They demonstrate that this work on linear data and, apparently, also on nonlinearly mixed data. Finally, the show how these constraints encourage distinct cell types in a model of spatial navigation.", "summary_of_the_review": "This paper addresses some very interesting questions. They provide an interesting mathematical derivation that could provide a substantial contribution at the intersection of ICA and non-negative factorisation (NMF, sparse coding etc.) literature. However, this would need a much more extensive treatment with more details. I think there are big questions related to the nonlinear setting (see Locatello et al., 2019). The arguments on cell types are interesting. However, the RNN model has a lot of degrees of freedom, the experimental results are unclear and one might wonder if this theory of efficient coding and sources recovery (which needs more links to existing literature) would also apply to, e.g., visual neurons.\n\nI strongly encourage the authors to split this up into multiple papers and do a more thorough scientific investigation of the different component claims: 1) ICA and non-negativity; 2) non-negativity as inductive bias in nonlinear ICA?; 3) cell types and non-negativity leading to separated components. Each of these questions is very important and interesting and it would be great to give them a proper treatment that creates new scientific insights, ideally with higher SNR than the average ML paper.", "strength_and_weaknesses": "Strengths: The paper touches on many interesting and highly relevant topics, including: the relationship between non-negative latent variable models (e.g. NMF) and source separation (ICA); nonlinear source recovery; and the reason for the existence of distinct cell types. The mathematical proof and the linear source recovery simulations appear to be valid.\n\nWeaknesses: It is rather frustrating to see how all these fascinating research questions are dealt with in a short inconclusive manner. The relation between non-negativity and source separation alone would provide more substance than could easily fit into a ICLR conference paper. Ideally, this would then also include a much more thorough literature research and connection to existing source separation and (non-negative) tensor factorisation (like NMF) methods since the early 90’s, as well as a number of additional control studies (changing the number of factors, their distributions, testing violations of different model assumptions, benchmarking against other methods like sparse coding etc.). The other two subjects nonlinear ICA (disentanglement) and cell type emergence are also extremely interesting and deserve a much more thorough treatment.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "gB2ajeaPbF", "reviewer": "Reviewer_GUiS", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper asks the question of why neural systems often develop a factorized or disentangled representation, i.e. a representation which aligns its variables/latent units with a meaningful factorization of the underlying problem structure. The authors use a combination of theory and simulations to show that biological constraints of non-negative and energy-efficiency together gives rise to disentangled representations. Finally taking the specialized cells in hippocampus like grid cells and object vector cells as a biological case in point, the authors simulate spatial processing in ANNs and demonstrate that biological constraints lead to disentangled codes for space and objects when these factors are independent, but the same constraints lead to entangled codes when these factors do not vary independently in a task, providing a potential explanation (specifically) for why grid cells warp their receptive fields and (more generally) for when and why disentangled codes arise in the brain. ", "summary_of_the_review": "I think the question of why neural systems have a factorised code has received little attention from a computational perspective and this paper is an important step towards filling this gap. For the AI community, understanding principles that promote disentanglement is important as they make generalization/transfer to new tasks easier. This paper illustrates important principles for achieving disentanglement in models (although they are already often implements in ANNs albeit without a true understanding of its consequences).  Overall, I think this is a very interesting paper that makes a novel contribution and would appeal to both the neuroscience and AI community. ", "strength_and_weaknesses": "Strengths: \n- I think the paper is exceptionally written and conveys the ideas very well. \n- The paper is a nice mix of theory and experiments along with a biological case simulation. The authors also present an intuition accompanying the theorem which nicely illustrates the logic.\n- Instead of simply comparing representations at the population level which has been the dominant approach in neuroAI, the paper compares representations at the level of axes, i.e. why the brain and ANNs often develop privileged axes. This is an intriguing question that has received little attention in the field so far and I laud the authors for tackling this with clarity. \n\nWeaknesses: \n- I think the justification behind non-negativity (of firing rates) as a biological constraint needs some further discussion in the paper. It is known that the firing activity of single neurons can be suppressed (from baseline rates) and activity is always going to be measured with respect to some reference signal/rate. The constraints the authors frame as ‘biological constraints’ are both routinely employed in modern CNNs through ReLU and l1 regularization on activity/weights. In that context, I found the use of the phrase ‘biological constraints’ to be a bit overblown. \n- The authors claim their results provide a potential explanation for concept cells (i.e. grandmother cells) but if energy considerations favor disentangled codes, we would expect to find selective codes for every single concept. Clearly, allowing specialized cells for all concepts in complex domains would be impossible as we’d run out of neural hardware. It would be useful if the authors could discuss the trade-offs or speculate about how such trade-offs might be handled by the brain. \n- I understand that space might be a limiting factor but it would be helpful to situate their work within the context of a broader literature on disentangling in machines (besides beta-VAE). There is plenty evidence of disentangling in machines as an emergent phenomenon (not baked in) for e.g. the work of David Bau et al. who found interpretable neurons in vision models, multimodal neurons in CLIP (https://distill.pub/2021/multimodal-neurons/), other work by on superposition/ privileged axes (https://transformer-circuits.pub/2022/toy_model/index.html) that also discusses the role of non-negativity constraints etc. There is also other work on understanding specialization in DNNs that could be relevant: https://openreview.net/forum?id=HreeeJvkue9. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "dxrMFNXSczl", "reviewer": "Reviewer_pBZV", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The manuscript introduced a bio-inspired representation learning that guarantees disentanglement in linear network, under some assumptions. The authors showed two neurobiological constraints i.e., nonnegativity and minimum energy efficiency (with respect to activity and weights), lead to linear factorization, under iid assumption of true factors. Additionally, the authors utilized the proposed biological constraints to interpret cell types which are functionally distinct. The authors used multiple simulated scenarios to demonstrate the performance of the proposed approach for linear and non-linear networks, under supervised and unsupervised fashion. Additionally, they reported their model disentanglement performance in comparison with beta-VAE for a few examples of benchmark datasets including Shape3D, dSprites and a rodent navigation task.\n\n**Update after rebuttal:** I have read the author's rebuttal and the latest revision. The authors have satisfactorily addressed most of questions/concerns and improved the manuscript by providing more details and clarification. After reading the rebuttal and the response to the other reviewers, I'm happy to raise my score to help the paper getting accepted.", "summary_of_the_review": "I think this is a good submission and the proposed neural representation learning using biological constraints is persuasive. However, I think the authors need to improve section 4 and discuss how these findings are generalizable to the other datasets.", "strength_and_weaknesses": "Strengths: \n\n- The paper is addressing a challenging and important problem, learning disentangled and interpretable latent factors, from a neurobiological perspective, which in my opinion is quite interesting and novel. The authors provide a simple set of constraints which lead to a disentangled representation which are quite valuable and potentially could expand the findings of the previous disentanglement studies.\n\n- The main claim of the manuscript is theoretically and empirically justified. \n\n- The simulation studies illustrate that under iid the proposed metric reveal the disentanglement decrement under entanglement attacks, while the comparable measures fail to present it.\n\n- The authors provide multiple experimental settings to evalute the proposed method. \n\n\nLimitations: \n\n- Although the paper has attempted to make a bridge between neuroscientific findings and the representation learning field, I think the proposed objective is not well-connected to the other studies in this field, except for beta-VAE. It is not clear how the proposed neurobiological constraints improves the already disentangled representations provided by other approaches like factor VAE, weak supervision method (Locatello, ICML, ICLR 2020), or satisfying local isometry (Horan, NeurIPS 2021 ), all of which do not necessarily rely on decorrelation. It is not clear how these neurobiological constraints will complete those approaches. I think the authors need to elaborate more on this.\n\n- IID Assumptions: I understand that making iid assumption for the latent (task) factors makes the theoretical justification easier, but in practice, especially in the case of neurons, it is not correct to assume that all neurons follow the same distribution.\n\n- Domain knowledge: one of the difficulties of learning a disentangled representation is the absence of domain-specific knowledge, like the number of independent factors. According to the problem definition in Eq. 1, to obtain disentangled factors, the network requires access to the true number of K, as I have understood. If this is correct, is there any way to utilize the proposed approach to learn the true number of latent factors?\n\n- The rodent navigation task is a very interesting experimental study. However, the problem formulation in Eq. 4 is quite unclear and confusing. \n\n- There experimental results for benchmark datasets are very limited.\n\nAdditional Comment/Questions:\n\n- Is the rodent navigation dataset a new dataset used in this paper? What are the modalities of the given data (measurements)? \n- Can you elaborate more on your findings in Figure 7? How should one compare the results in Figures 6 and 7?\n\n- How would this model perform when K is not matched with the true number of factors? In the case of cell type, how can one interpret the results in terms of the function of each neuron?\n\n- How can the proposed method be used to functionally characterized neurons in the other data modalities, e.g. RNA-seq? \n\n- The authors may want to move some of the supplementary materials, e.g. sequential path integration to the main text. It is quit unclear what Eq. 4 is formulating. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ZMwqfwMsdiG", "reviewer": "Reviewer_4iCk", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors investigate how biological brains could achieve disentangled representations. Through a combination of theory and modeling, they provide ample evidence for two key constraints: (1) nonnegative activities and (2) a constraint on energy expenditure in model/real neurons. The work is exceptionally well explained and builds out from these simple principles to explain results in modeling complex artificial data and rodent experiments. ", "summary_of_the_review": "This paper is an absolute accept. Beautifully written, deeply insightful, great lit review, and lays out theory that I believe can inform the development of better models for AI and the design of new neuro experiments. Best paper I've reviewed for a conference in years (ever?). I'm going to be a jerk and give it an 8 for now just because the authors didn't release their code. Assuming they do over the discussion period I'll bump them to 10.", "strength_and_weaknesses": "Strengths:\n\n- Great introduction.\n\n- Love how the authors presented and explained the proof relating non-neg and efficiency to independent task factors.\n\n- The related work brings together the parallel investigations of disentangled factors and the use of nonnegativity to do so (e.g., Lee & Seung, 2000) in Neuroscience and ML.\n\n- Honestly I could go on and on, but this is a brilliant paper. I love how it begins with a simple concept, explains it like I'm five, and then shows that the theory holds in complex situations that are relevant to the large-scale ML and Neuro pursuits of today.\n\n\nWeaknesses:\n\n- Can you add a github link with your code?\n- Make the text bigger in your figures? It's a little hard to make out. That's it.\n\nQuestions:\n\n1. Are the colored circles in Figure 2 cartoons or real data?\n2. In Fig 2, if I'm understanding this correctly, relu and a non-neg activity constraint yield the same performance. Does the regularization also affect the signs of the model weights?\n3. There's no constraint for non-negative weights, right? Just non-negative activities?\n4. What’s the explanation for neurons that show mixed selectivity? Is that just a relaxation on the energy constraint?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%