DISENTANGLEMENT WITH BIOLOGICAL CON-STRAINTS: A THEORY OF FUNCTIONAL CELL TYPES

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.

1. 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. A 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 responses (Olshausen & Field, 1996; Yamins et al., 2014; Ocko et al., 2018; McIntosh et al., 2016) 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 an example, we do not even understand why many different bespoke cellular responses exist for physical space, such as grid cells (Hafting et al., 2005) , object-vector cells (Høydal et al., 2019) , border vector cells (Solstad et al., 2008; Lever et al., 2009 ), band cells (Krupic et al., 2012) , or many other cells (O' 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 to a mixed response to 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 et al., 2020; Higgins et al., 2021) , why visual cortical neurons are de-correlated (Ecker et al., 2010) , why neurons in parietal cortex are selective only for specific tasks (Lee et al., 2022) , why prefrontal neurons are apparently mixed-selective (Rigotti et al., 2013) , and why grid cells sometimes warp towards rewarded locations (Boccara et al., 2019) and sometimes don'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 disentanglement as single neurons responding to single factors of variation (see Appendix A.1 for further details). Such disentangled factors can facilitate compositional generalisation and reasoning (Higgins et al., 2018; 2017b; Whittington et al., 2021a) (though some work has challenged the idea that disentangled representations generalise better (Schott et al., 2022) , as well as lead to more interpretable outcomes in which individual neurons represent meaningful quantities. Unfortunately, building models that disentangle is challenging (Locatello et al., 2019) . * Correspondence to: jcrwhittington@gmail.com 1

