Understanding Neural Coding on Latent Manifolds by Sharing Features and Dividing Ensembles
Authors: Martin Bjerke
Lukas Schott
Kristopher T Jensen
Claudia Battistin
David A. Klindt
Benjamin Adric Dunn
Venue: ICLR 2023 poster (2023)
Keywords: neuroscience
neural activity
tuning curves
neural ensemble detection
grid cells
latent variable models

Abstract: Systems neuroscience relies on two complementary views of neural data, characterized by single neuron tuning curves and analysis of population activity. These two perspectives combine elegantly in neural latent variable models that constrain the relationship between latent variables and neural activity, modeled by simple tuning curve functions. This has recently been demonstrated using Gaussian processes, with applications to realistic and topologically relevant latent manifolds. Those and previous models, however, missed crucial shared coding properties of neural populations. We propose $\textit{feature sharing}$ across neural tuning curves which significantly improves performance and helps optimization. We also propose a solution to the $\textit{ensemble detection}$ problem, where different groups of neurons, i.e., ensembles, can be modulated by different latent manifolds. Achieved through a soft clustering of neurons during training, this allows for the separation of mixed neural populations in an unsupervised manner. These innovations lead to more interpretable models of neural population activity that train well and perform better even on mixtures of complex latent manifolds. Finally, we apply our method on a recently published grid cell dataset, and recover distinct ensembles, infer toroidal latents and predict neural tuning curves in a single integrated modeling framework.

Published as a conference paper at ICLR 2023
UNDERSTANDING NEURAL CODING ON LATENT MAN- IFOLDS BY SHARING FEATURES AND DIVIDING EN- SEMBLES
Martin Bjerke', Lukas Schott?, Kristopher Jensen',
Claudia Battistin', David A. Klindt Benjamin A. Dunn
"Nowwegian University of Science and Technology. Bosch Center for Artificial Intelligence.
University of Cambridge, 'Equal contribution jaiin jenjamin.
imartin.bjezke, ben
no, klindt. daviddgmaiilllddd.d com.
ABSTRACT
Systems neuroscience relies on two complementary views of neural data, charac- terized by single neuron tuning curves and analysis of population activity. These two perspectives combine elegantly in neural latent variable models that constrain the relationship between latent variables and neural activity, modeled by simple tuning curve functions. This has recently been demonstrated using Gaussian pro- cesses, with applications realistic and topologically relevant latent manifolds. Those and previous models, however, missed crucial shared coding properties of neural populations. We propose feature sharing across neural tuning curves which significantly improves performance and helps optimization. We also propose so- lution the ensemble detection problem, where different groups of neurons, i.e., ensembles, can be modulated by different latent manifolds. Achieved through a soft clustering of neurons during training, this allows for the separation mixed neural populations in unsupervised manner. These innovations lead to more in- terpretable models of neural population activity that train well and perform better even mixtures of complex latent manifolds. Finally, we apply our method a recently published grid cell dataset, and recover distinct ensembles, infer toroidal latents and predict neural tuning curves a single integrated modeling framework.
INTRODUCTION
Neural population activity can appear high-dimensional (Stringer et al., 2019), yet much recent work
has reported that neural populations higher brain areas are often confined to low dimensional
subspaces (Yu et al. 2008; Harvey al. 2012)] Mante al.] 2013 Stokes al.] [203]] [Shenoy al.
2013 Kaufman et al. 20144 Sadtler et al. [014 Gallego et al. 2017 Elsayed & Cunningham 2017
Gao et al.] (017). The bread and butter of classic systems neuroscience linking neural activity
to experimentally controlled or observable covariates such as orientation (Hubel & Wiesel 1979)
pitch (Lewicki 2002), movement (Churchland et al. 2012 Kao al. 2015). posture (Mimica al. 2018) and orientation in space (Taube et al.) 1990). These two parallel streams of neuroscientific research might at first seem be at odds with each other (Kriegeskorte & Weil 202T): tuning studies of individual neurons give a very different picture of neural coding than distributed representations over high-dimensional neural populations. However, they combine elegantly in the form of (neural)
latent variable models (LVMs, see Lawrence] 2003 Yu et al.) 2008 Pandarinath et al. 2018).
In their basic form, neural LVMs find the low-dimensional structure of neural population activity, for instance, when a large network of neurons is coding mostly along few linear subspaces (Mante) cet al. 2013 Gao et al. 2017). One advantage is that these models can help us discover latent variables which may not be tracked as classical covariates in systems neuroscience. However, when the mapping from latent variables to predicted spike rate (decoding) is fully unconstrained, e.g., by using a multi-layer neural network, lose the simple biological interpretation of tuning curves.
In an effort to maintain a biologically interpretable relationship between the latent variables and the neural activity, recent work has proposed more constrained decoders approximating simple tuning

---

Published as a conference paper at ICLR 2023
population responses
encoder
separate latent spaces
shared tuning curves
ensemble selection
spike predictions
111
II II
111
II
Figure I: Model outline. Our main contributions are outlined in purple. The input is a matrix of neural spiking activity. The encoder (1) is a multilayered neural network. The latent spaces separate with, potentially, different topologies (e.g. IR. and T9). The decoder a parametric tuning curve model with feature sharing
(2). The ensemble detection is weighted (ideally one-hot) selec- tion of latent spaces for each neuron (3). For decoding of the activity, we assume Poisson spiking.
curves. These tuning curves have been parameterized as Gaussian processes in the framework of Gaussian process latent variable models (GPLVM. We et al., a 2018)2018). Through the tuning curve approach, we limit ourselves to biologically plausible solutions that reveal the actual algorithmic structure of the neural system. Thus, LVMs with simple tuning curve decoders bring together the view on neural populations as distributed representations of ow-dimensional latent variables, along with the biologically meaningful perspective on individual neural tuning properties.
Some neural populations exhibit topologically interesting latent manifolds (Singh et al. 2008 Peyrache et 2015a Gardner et al., 2022). For instance, grid cells represent navigational space in toroidal coordinates of spatially repeating two dimensional hexagonal grids (Hafting et al.) 2005)
They appear in different ensembles, commonly referred as modules, each coding for space at different resolutions (Fyhn et al. [007 Sttnsolaa et al. (012). Thus, the complete population of grid cells is best described as a collection of ensembles neurons, where neurons in each ensemble have tuning curves of specific shapes their respective toroidal latent representations of space (Curto 2017). By contrast, a two-dimensionaal Euclidean representation might also account for the complete population of grid cells, but would also completely obscure their efficient and theoretically interest- ing coding scheme for representing space (Solstad et al.) 2006, Sreenivasan & Fiete 2011 Mathis eet al. 2012 Wei et al.) [2155] Klukas al., (020)). A driving motivation behind this work to model this beautiful neural structure with an LVM that separates the algorithmic and biological parts, while uniting shared tuning properties be more accurate and trainable than previous approaches.
We propose to train neural LVMs that not only have simple tuning curve decoders, but are also fully differentiable. Thus, we use flexible encoder, i.c., a neural network as in variational autoencoders (Kingma& & Welling 2014 Rezende et al. (014). and simple tuning curve based decoder, akin to GPLVM. The encoder can readily be made convolutional to allow for better latent estimation from adjacent time points. Additionally, we implement a feature basis for the tuning curve shapes in the decoder which can shared across neurons. We demonstrate that the neural feature sharing, along with the variational end-to-end training, vastly improves both the training stability well as the final performance of neural LVMs. Moreover, we propose hybrid inference at test-time and show that this, again, brings considerable improvement in performance. Finally, we integrate the problem of separating distinct ensembles of neurons into our approach - a crucial task for the discovery of different biological structures and the precise mathematical understanding of their topologiaal tuning properties. An illustration of our approach is provided Fig.
To summarize, our full model performs the task of finding latent variables, separating distinct en- sembles of neurons and fitting the prototypical tuning curves on each ensemble's latent space in a single efficient framework. In the following, therefore refer to our model as feature sharing and ensemble detection Latent Variable Model or faeLVM.
2 BACKGROUND
Let A be the instantaneous firing rate of a neuron i. To relate this to the spiking activity T;, we assume Poisson noise model x; P():). We define latent variables == (-1,---, =x) } (in distinct
2

---

Published as a conference paper at ICLR 2023
A 140 120
B L. 0.8
Figure 2: The feature sharing sumption. A, Tuning curves for 26 head direction tuned neurons. B, The same 26 tuning curves, shifted and scaled to visualize structural similarity, along with a standard Gaussian shape (dotted line).
Normalized Hz
so
0.8
0.4
0.2
Head direction
Head direction
spaces and possibly with different topologies) that change over time and to which neurons are tuned. More precisely, we suppose that there exists a deterministic function fi, which for every neuron i relates the latent variables to the firing rate, i.e., A = f(()). One additional assumption is that each Si is only sensitive to subset J; of latent variables and invariant to all others. That means, for two values 2, å with 2j- 2 Vjj E Jr and zj + Vj & J; we have f1(z) = f((3).
As a concrete example, consider a mixed population of head direction neurons 1, and a grid cell module 1z. The former are tuned to a circular latent variable 21 S (Peyrache et al., 2015a Ry- bakken et al. 22119] Rubin et al. [2019 Chaudhuri et al. 2019), the latter to a toroidal latent variable 22 T McNauggton et al. 2006 Fuhs & Touretzky 2006 Gardner et al. 2022). That means, for every neuron in the head direction ensemble we have J. = (1,, Vi 11; whereas for every neuron in the grid cell ensemble we have Jj = {2], Vi E Iz. Conceivably, there exist shared latent variables, e.g. attentional state, running speed or pupil dilation, to which multiple or even all neurons are tuned.
Within each ensemble, the tuning functions f;; E 1 are likely to have similar structure (Albright 1984 Taube et al.) 1990, |Hafting et al.. 2005). This is the concept of feature sharing (Klindt et al.
2017). J which in visual systems neuroscience has proven to be a useful assumption (Batty et al. 2017
Sinz et al., 2018 Ecker et al. 22018] Walker et al.] 2018 Cadena et al., 2019 Ustyuzhaninov et al.
2019
Cotton et al. 2020 Zhuang et al. [202]] J Burg et al. 2021
Bashiri et al., 2021 Safarani et al,
2021
Franke et al. 2021 Lurz et al. [222]] |ayebiy et al. 2021 Seeliger al. 202 [| Goldin al.,
2021
Ustyuzbaninov et al. (022). In vision, this can be motivated by the fact that visual (especially
early retinal) neurons form mosaics, that tile the input space and compute the same response function
(i.e., tuning curve) across space (Wässle & Riemann 1978). Precisely, let fn denote the response function of a visual neuron with receptive field center l E TR²,, let T : R²² + R² be a spatial translation, and let R -> IR be an image (with slight abuse of notation, is, here, a function that assigns a grey scale value to every position in two-dimensional space, i.e., an image). Let us think of fi - f and fr- = fIiM) as two neurons of the same type (i.e., same tuning curve shape) but with their receptive field centers at different locations μ and T(x). Then we have f((2) (2 Thus, for visual neurons we assume Iranslational equivariance.
Consequently, should be able marginalize over spatial translations learn the shared structure among all fi as demonstrated in Klindt et al. (2017). The same intuition holds for populations of au- ditory neurons that are equivariant pitch, i.e., that have similar tuning properties translated across log-frequency space (Kell et al., 2018 Kell & McDermoti 2019) and possibly other sensory neu- rons such somatosensory populations (Lieber & Bensmaia 2019). But it also holds true for higher cognitive neurons, such as the ones discussed above, where, for instance, each grid cell is tuned to a different location on the torus and the receptive fields (tuning curves) across grid cells are remarkably similar in shape (Hafting et al. 2005: Fyhn et al. 2007 Gardner et al. (022). Thus, we propose to introduce feature sharing also in this case. More precisely, we argue that every fa within an ensem- ble 1,. should be modeled by shared tuning curve function (z, 0,) fi(z) that applies simple neuron specific transformation, parameterized by 0;, to obtain the specific fi of each neuron. For instance, in the case of neuron-specific spatial translation (f4, the receptive field center above) and scaling (ax), we would define 8, = {A4. ax} and thus (2, 82) = a,g, (T2, (z), fu) = f(() (9, yielding simply the centered prototype of the shared tuning curve). Evidence for this assump- tion is provided in Fig. 2 where centering and rescaling all tuning curves of head direction selective neurons on 81 (recorded by Peyrache al. 2015b) clearly exhibits shared, Gaussian-like tuning.
In visual perception there is translational equivariance, however, only among cells of the same types (Wässle & Riemann| 1978). Therefore, previous work proposed feature sharing in combination with functional cell type identification (Klindt et al. 2017). Specifically, equating neural ensembles with cell types in the retina, this posits that all neurons in given ensemble I, be tuned to the same set of latent variables Jj, i.e., Jj Jj, Vi E Ty. As an example, the toroiaal structures zj E T of different
3

---

Published as a conference paper at ICLR 2023
grid cell ensembles I only become apparent after the successful separation of distinct modules with different spatial resolutions (Gardner et al. 2022). Thus, we argue that feature sharing hinges on the successful identification and separation of distinct neural ensembles in a mixed population recording.
The problem of ensemble detection can be formalized as partitioning n recorded neurons / U 1 into k non-empty subsets T5 = 0. vj J Unfortunately, this a combinatorial problem with kn possible states. Even for 'simple' problem such as clustering neurons into k = 2 groups, this already exceeds the number of atoms in the universe (10*2) at n 2 273. With modern neuroscience routinely yielding thousands of recorded neurons per experiment (Jun et al., 2017). brute-force ex- haustive search is clearly not a viable approach. To mitigate this issue, we propose a soft-relaxation of the clustering problem. Specifically, we define convex combination of weights over responses derived from tuning curves on each latent variable, and train the neurons to maximize their likeli- hood given those weights in a differentiable fashion (explained in more detail in the next section).
3 METHODS
As above, we denote neural activity x and the collections of latent variables as 2. We want to learn latent variable model p(a, =) p(=))=)((2). We choose variational autoencoder (VAE, Kingma & Welling 2014), which maximizes a lower bound to the data likelihood called the evidence lower bound, which necessitates the introduction of an approximation to the posterior distribution, i.c., a variational (input dependent) posterior with non-zero support across the domain, ((2xz) > 0, V2, x.
At test-time, we focus on precise estimate for the latents z. These latents are usually inferred with the encoder that computes the approximate variational posterior (<<[2). Here, we aim to further refine this estimate by performing inference based on the decoder p(z]z). Technically, this can be achieved by directly performing gradient descent on the latents, with the decoder starting from the
encoder estimate (Schott et al. 2018 Ghosh et al., 2019) (see Appx. [ for further details).
For Euclidean latent spaces = E RR, we can simply rely on the standard VAE framework, using Normal (N) posterior and prior for those dimensions. For spherical latent spaces S", we use a wrapped Normal distribution (wN), akin to previous work (Falorsi et al. 2019 Jensen et al. 2020). which satisfies the topology of the latent space. For the relevant case of a toroiaal latent space, we use T S' x S', together with the fact that the variational posterior is, usually, assumed to factorize across dimensions. This factorization also helps keep the equations the same when working with multiple latents (corresponding multiple ensembles), i.e., (21,-..,
More precisely, we assume that each latent 2j lives either in Euclidean space R.y or on a torus T = 8t X x... x Sh, (with St = TT), Further, we let our prior factorize (Kingma & Welling 2014), iee.
p(=) II p(=)), p(zz)) j=1
II II 1=1
N(0. 1), if E R"y wN(0, 1), if E T
giving uniform prior distributions. Analogously, we let our variational posteriors factorize as b ((2|x)= = II d(>z]æ)) j=1
a(z,)=)= = II a(< (l)))lz)) (z),
ny
=50)
N(u(t), if E R2 wN (())) if E Th,
(2)
Here, the distributionnsseccfie parameters ((x) and (xx are themselves input dependent, and given by functional mapping which is learned using the reparametrization trick. We use temporal (1D)
convolutional neural network to parameterize q. To enforce circular latents, we the encoder output vectors in R23 which we subsequently normalize to avoid discontinuities when estimating angles (Zhou et al.) 2019). We remark that although convolutional filter was used, temporal smoothness was not explicitly included in the latent dynamics of the model used for producing the results in Sec. 4 It is, however, a straightforward inclusion to incorporate (see Appx. D.I).
For the decoder, in clear contrast to standard VAEs, we pick a simple parametric function f to parameterize the reconstruction term p(r|z). As argued above, this is a crucial choice in the interest of defining simple biologically meaningful link functions (i.e., tuning curves) between (potentially complex) latent variables and neural activity. Specifically, we have
p(a.12) P(f(()),,
f((z) = Σ wigh (zj, 0t),
(3)
4

---

Published as a conference paper at ICLR 2023
train time bine
best bins
Figure 3: Schematic outline of model evaluation. Evalua- tion pipeline as suggested in[Pei et al. (202T). During training, the models are all trained on spikes from "train' neurons (1), learning their tuning curves, and the latent variable represen- tation which is used to infer tuning curves on "test" neurons at train time (2). For testing, we fix all tuning curve shapes and infer the latents from test data (3), before performing spike prediction on held-out data from held-out neurons (4). Thus, our encoder only receives train neurons (1, 3) to infer latents.
trai meurooe
Mosan X Vecke
Noss x Feo
test nesrone
Nep Toan
Nes TTa
where P(f((z)) is a Poisson distribution with rate fi(z), and the weights for each neuron the output That is, Waj models the participation of neuron in ensemble j. The function 9A (zj, 0:) E R characterizes the (shared) tuning curves of neurons in the latent space of 23, and takes as input the corresponding latent 2, and a (set of) neuron specific parameter(s) 8, produce the response of the neuron based on each latent space.
of softmax function, such that wij 0 and Σ, Whj
The function gj (zz, 02) can be specified in multitude of ways, while still embodying the concept of feature sharing. Here we present two possible options: the first simply assumes a single heat kernel shape (i.e., Gaussian-like bump) with shared tuning width for all neurons. Since neural activity is non-negative, we model the log of the instantaneous firing rates as
log (2). u:) d, (z)
where l4 denotes (as above) the center of the i-th neuron's tuning curve and (. a distance function in the space of 2, (i.c., Euclidean distance in IR" and geodesic distance in T²). The second, and more flexible, feature space consists of a sum of M weighted (8m) Gaussian basis functions (i.e., a spline)
dz Um 22)2
log (5() = Σ Bm exp
(5)
o²
with center am, and width o. of the m-th basis function. We consider three variations of faeLVM: sharing of features by utilizing Eq. 44to to model with a shared (bump) heat kernel (faeLVM-b), sharing of features through Eq. S using Gaussian basis (faeLVM-s), as well model allowing each neuron to learn its own set of the basis weights in Eq. 5 i.e. no shared features (faeLVM-n).
4 EXPERIMENTS
4.1 SIMULATIONS: FEATURE SHARING
Motivated by the evaluation pipeline in Pei et al. (02T), we split our data into four parts (Fig. 3) exploring the benefits of feature sharing in the following way: We generate ID latent variable and simulate spikes according to a Poisson distribution with heat kernel tuning curves (see also Appx. for retinal ganglion simulation with calcium dynamics). We fix Tiss and Niea (1000 and 30). and investigate model performance based on fixed Timin and varying Nirain. and vice versa. For each condition, we repeat the experiment over 20 seeds (while data is different for each repeat, all models are trained on the same seed, ensuring valid comparisons), reporting mean negative log-likelihood (NLLH) and geodesic error (GE) on test data (lower numbers better, in both cases). We also compare our models with mGPLVM (Jensen et al. 2020), an extension of GPLVM (Wu et al., 2018) designed for inference on non- non- on-Eullide spaces, as well as conventional VAE using a standard multi-layer perceptron decoder (MLP; with essentially reverse architecture of the encoder). As MGPLVM assumes Gaussian noise model, the Poisson NLLH is unsuitable as measure, hence it is only included when evaluating latent recoveries.
S'. circular
Fig. A clearly show the regime where feature sharing excels: datasets with a satisfactory number of neurons, but with short recording times. This coincides well with the current state-of-affairs for neural recording techniques, where experiments are usually limited by the possible experimental length, not by the amount of recordable neurons (teevenson & Kording 2011). Moreover, even with further technical innovations, we may eventually be able to record from all neurons, but we will likely never be able to record during all possible natural inputs or behaviors.
5

---

Published as a conference paper at ICLR 2023
A
D
Raie pred. (NLLH)]
Latent pred [G]
06
0.0
1.0
Newn
1.4
Trah
0.8
0.4
1.2
0.4
0.9
0.2
1.0
10
20
40
50
21
1/
2
2
0.0
30
0.0
2
250
2
Meain
Nooin
Tora
Figure 4: Latent and Rate Prediction. The models facLVM-n,s.b are described in the main text, mGP is the model from Jensen et al. (2020), VAE traditional approach using an MLP decoder. A, Mean NLLH of predicted test neuron rate, as function of amount of training neurons. As the performance is averaged over 20 seeds of varying data, error bars would also reflect the variability NLLH across data, resulting in inappropriate visualization. Hence, we report the mean rank of 20 seeds for each model at each condition, in an effort to more appropriately capture the variability of the models. B, Same as A, now as a function of training time points. C, Mean GE between true latent and inferred latent, test data, with corresponding SEM error bars, as function of number of training neurons. D, Same as C, now as a function of training time points.
We also remark that while faeLVM-n seems to performs better on rate prediction in the low neuron regime (Fig. 44, the performance LVMs ultimately hinges on how well the models can recover the correct latent variable. In the setting of latent recovery (Fig. FC-D), we see that feature sharing evidently beneficial both in low neuron- and low recording-tines regimes. faeLVM-b outperforms the other models, perhaps somewhat expected given the shape of the synthetic tuning curves, while all faeLVMs outperform both mGPLVM and VAE. We also observe that while mGPLVM, which uses tuming curve decoders, is more interpretable, it is worse than neural network decoders. How- ever, our approach with feature sharing closes that gap and even improves upon the existing VAE method, thus highlighting the importance of sharing features; more accurate inference, while also being interpretable. Note that although the faeLVMs leverage the benefit of a convolutional layer (which mGPLVM does not, being non-convolutional model), we emphasize that the faeLVMs still outperforms mGPLVM, even after reducing the convolutional filter size down to 1, which forces the model infer the latents independently each time point (see Appx.
4.2
SIMULATIONS: ENSEMBLE DETECTION
1
B
1.0
0.9
level
24c
Accuracy
raw_pca
Accuracy
0.8
Amoount chtta
0.8
cav_pca
0.7
feeLVM-D5 supervised
0.6
0.
2 20
24
27
2 Amount of data
29
24
-1.00
00.75
20
Amount of data
Figure 5: Ensemble Detection. Ensemble detection accuracy, with corresponding SEM error bars, on two (A) and three (B) equally sized ensembles of neurons tuned separate S² latent variables. The models are: k-means on PCA reduced data ('raw-pea'), k-means on PCA reduced covariance matrix of the data ('cov-pea') ) supervised clustering based on the mutual information scores between the cell responses and the latent variables, fael.VM with shared heat kernel tuning curves ('fael b'), and the same model optimized over seeds, selecting the one that produced the highest L.HH ("faeLVM-55'). Chance levels computed from 100, 000 random labels with optimal permuta- tions. C, Relation between likelihood and accuracy for different data amounts.
Next, we shift our focus to the challenge of ensemble detection, comparing the performance of faeLVM-b against other common methods. Specifically, we contrast against clustering methods (see Appx. C for more details) that either i) perform k-means clustering on dimensionally reduced neural activity time (loosely inspired by Lopes-dos Santos al., 2011] Baden et al. 2016; 2 Hamm et al.
202T), or ii) perform clustering on the neural covariance matrix (loosely inspired by Carrillo Reid
al. 2015). As additional upper bound to the achievable performance, include supervised

---

Published as a conference paper at ICLR 2023
clustering method that receives both the mutual information score between each neuron's activity, and the true latent variables, as input. We run our model over five different random seeds and pick the run that yields the highest LLH (faeLVM-b5 in Fig. [5) fair approach as the likelihood is an unsupervised metric (k-means is run over 100 seeds, picking the one with the lowest inertia).
The setting chose for these experiments ring, i.e. S The comparison methods failed to separate the ensembles when scaling up to neural activity on a torus, in contrast faeLVM, which further highlights the novelty of an unsupervised method that can accomplish this task (see Sec. 4.4 for tori, and also Appx. [ for an additional ensemble separation ablation). We tested the accuracy on separating two and three distinct (no shared latents across ensembles) neural ensembles of equal size on rings, in addition to investigating accuracy for all models as function of available data points (recording length).
the separation of ensembles of neurons that live
Fig. 5B shows that our model outperforms both comparison methods and quickly approaches the performance of the supervised upper bound, for both two and three rings. This even more pronounced for the model selected based on LLH - viable option when the analysis result war- rants additional computing time. Fig. [C shows that the relationship between ensemble detection accuracy and LLH is monotonically increasing. Thus, the unsupervised likelihood score is indeed good proxy to arbitrate between different optimization runs.
4.3
REAL DATASETS: HEAD DIRECTION DATA
2n
B Neur
Rate pred, [NLLH]
12
Latent pred. [GE]
Interned latent
2.0
3.0
Neur 15
3.5
0.8
I
Nour 20
1.0
0.4
True head direction 2n
1000 0.5
0.2
favt/VM-n
VaiLVV-bb
mGP
Figure 6: Application to head direction data. A, Variational mean of fa-L plotted against recorded mouse head direction (for additional visualizations, see Appx. H.. We trained 20 models, selecting the one with highest training likelihood on training neurons, before performing inference at test-time. B, True test neuron spikes (grey), and predicted test neuron rates (black) from the model used in A, per bin, for three randomly selected neurons. C, NLLH of test neuron rate, over 20 seeds, as violin plot. Performance of the best model, based train LLH, is indicated by the black dot. D, Same as C, but comparing GE between recorded head direction and inferred latent, on test data.
Transitioning to real datasets, we first apply our models to data from Peyrache et al. (2015b); more specifically the awake trials from the dataset labeled Mouse 28, session 140313, recorded from the anterodorsal thalamic nucleus (ADN). Recordings yield approx. 35 min. of data, which we bin in 100 ms bins, before separating it with a train-test split of 95% - 5%. 3 neurons are allocated as test neurons (out of 26 total) and run 20 seeds of each model, selecting the best performing model for visualization (facLVM-nn). The inferred latent (Fig. ( is observed to match the true head direction high degree, and the model is able predict rates on test neurons (Fig, 6B). even though the average and peak activity over test bins are quite low (~~ 1 and 5 Hz respectively) for the first two.
As we can see from Fig. 6C, the unconstrained model (facLVM-n) outperforms the other two. This is not an unsurprising result, given the model's higher complexity and the sufficiently large amount of training data. As for the latent prediction, mGPLVM and the traditional VAE (Fig. (P)) are close in comparison to the performance of fael VV-s, although less accurate than fael.VM-. We also note that the VAE does perform better on the rate prediction, which, conceivably, might be explained by the fact that it is not constrained by particular tuning curve model and thus might learn interactions that incorporate non-head direction variables in an effort to improve its rate prediction. Training multiple models and making a selection based train LLH generally correlates well with solid results on the test set, as also shown in 5C. Note that faeLVM-b achieves much more consistent optimization results over different random inttializations, suggesting a typical vaaiinei tradeoff.
Overall, we recognize that while there are differences in performances, all models recover the correct circular manifold. The fael VMS however, noticeably faster (~ 5 min. time for one seed, on
7

---

Published as a conference paper at ICLR 2023
1:25
Weights
0.75
ens
1.00
Accuracy
0.95
0.30
0.50
WW
0.85
0.00
50
100
150
200
250
204
50
125
150
Neurons
NLLH
Neurons removed
Open Field
Toroidal
Open Field
Toroidal
Open Field
Toroidal
Open Field
Toroidal
Figure 7: Application grid cell data. A, Learned ensemble weights. Neurons are ordered by de- creasing spatial information, dashed line separates first and second reported module (Gardner et al.) 2022 B, Greedy test for each neuron, checking which latent provides higher likelihood (i.e., one- hot ensemble weights), colored by reported module. C. Classification accuracy when successively removing neurons with low spatial selectivity. The same number of neurons were removed from each ensemble, starting with least spatially informative ones. Accuracy is measured against super- vised ensemble detection from Gardner et al. (2022). D. Rate map of every 23th neuron, first grid cell module (smoothed for visualization purposes). Maps shown function of rat's position in arena, as well as inferred coordinates on latent torus. Red circles indicate inferred receptive field mean and width (SD) for the heat kernel (i.e., Gaussian-like bump). E, Same D for second module. For full set rate maps, as well as latent dynamics, see Appx.
CPU), compared mGPLVM (~ 50 min. time for one seed GPU). While critical in the case of this particular dataset, it is paramount when inferring latents of higher dimensionality. possible do model
We also note that while the circular manifold selection over different latent spaces (toroidal, circular, planar, etc.) as was demonstrated in Jensen al. (2020). Thus, is not strictly necessary know the latent topology when applying our method, rather it allows for the restriction of possible models to limited hypothesis class.
in this case known,
4.4 REAL DATASETS: GRID CELL DATA
Finally, we apply the complete model pipeline to data recorded from the medial entorhinal cortex (MEC) of freely moving rat (Gardner et al. 2022). Neurons in the MEC selective number of features like space (Fyhn et al. (2004), head direction (Sargolini et al.) 20066), speed (Kropff et al., 2015), as well conjunctive representations of the three (Sargolini et al. 2006 Hardcastle al.
2017 though most demonstrate a strong preference for single features (Krofff et al. 2015 Tukker al. 2022). We consider data from the rat R, day 1-session, using recordings from neurons labeled module 2 and 3, considering 149 and 145 non-conjunctive neurons respectively. Spike data was binned in 100 bins, after removing periods in time when the rat was considered stationary (speed < 2.5 cm s 1)). Recordings from both the open field and maze experiments yields approx. 170 min. of data. We train 20 models from different seeds, selecting the one with highest in-sample likelihood.
Compared with the modules discovered in the dataset (Gardner et al.) 2022), our model's results coincides with an accuracy of 0.83. The inferred weights (Fig. 7 show a clear separation between the two modules, becoming less distinct the neurons' spatial information decreases. This is also reflected Fig BB where neurons further away from the diagonal tend to be more accurately classi- fied. However, note that the original separation of neurons into ensembles is not intrinsic truth; a number of the less informative neurons exhibit structurally ambiguous rate maps (see Appx. G). possible indication that they may belong to different ensemble altogether. Excluding these neurons
8

---

Published as a conference paper at ICLR 2023
while evaluating the model (Fig. 7c, we observe a significant increase in performance, reaching 0.90 accuracy after removing approx. 50 of the least informative neurons from each ensemble.
Fig. 7D-E show selection of rate maps for neurons from both ensembles. We see that the model infers toroidal decodings that keep the structural integrity of the receptive fields, akin to the findings in Gardner et al. (0022). The inferred centers and widths also correspond well with the rate maps. We note that, while promising, the dataset included here contains just 294 of the 2460 neurons originally recorded and that the full dataset (currently not available) would require additional steps determine the number of ensembles, the corresponding shapes and allow for some degree of conjunctivity.
5 CONCLUSION
In this paper we provide a proof-of-concept that the paradigm of share and divide is a successful application of the old idea of carving nature at the joints in the sense that similar observations are grouped while differences are highlighted. This principle, instantiated in a group equivariant computation core and a cell-type specific readout, has been successfully deployed in visual neuro- science; we show that it also provides an excellent inductive bias (or prior) in the case of higher cognitive variables such as (but not necessarily limited to) head direction or position in space. is advantageous from a statistical perspective to leverage the similarity in tuning functions across neurons, as well as from a biological point of view, as finding protoypical tuning curves can help us unravel the common tuning properties across individual neurons.
The idea of separating ensembles is at the core of the scientific endeavor where the task is to group our observations of nature into distinct categories. In the case of grid cell modules, it is absolutely crucial, as the topological properties that define this class of neurons are completely hidden when the separate ensembles of different spatial resolution are not properly distinguished. One might wonder whether other brain areas such as primary visual cortex, for which a power law in the population activation space has been stipulated (Stringer et al. 2019), may not resolve into simpler underlying topological ensembles of neurons, such as the special orthogonal group of three-dimensional rota- tions S0(3) (vision), 2-spheres S (Singh et al., 2008), or the Klein bottle shape of oriented and phase shifted wavelets (Carlsson et al. 20088 (relevant to multiple sensory modalities).
Furthermore, the assumption of feature sharing relies on an equivariant neural code, such as the approximate translation equivariance in the visual system, which certainly is an abstraction from the actual inter-neural variability that exists in the real world. some cases, it might be that neurons are better described along some continuous space of variations, rather than by a fixed set of discrete prototypical tuning curves (Ustyuzhaninov et al., 2022). On the other hand, we can identify the parameters (c.g., tuning width) that describe the continuous variations across neurons and separate them from the stable features of tuning curves (c.g.. Gaussian-lide like heat kernel shape), we can use that to build more flexible parametric (but still shared) feature basis. Examples of such approaches include rotation (Ustyuzhaninov et al. 2019), kinetic (Zhao et al. 2020) or more general nonlinear modifications (Shah et al.) (202)) to shared feature spaces. Further challenges to the feature shar- ing idea could arise from recent findings of representational drift (Rule et al./ (019)) and location- (Qiu et al. 2020) or context dependent tuning (Kanter al., 2017), although future directions might consider addressing these in the model. Other limitations are the fact that the grid cell torus is hexag- onal, while we use an orthogonal basis; further analyses might reveal if the encoder compensates for that, e.g., by adjusting the distribution of speeds on the torus per direction.
Finally, although the choice of variational posteriors might at first seem restrictive, we would ar- gue that the corresponding topologies (circles, tori, I²") are natural choices for latent spaces and likely to keep appearing in the brain, as is suggested in recent work (Kriegeskorte & Wei [021). Other authors conjectured that, e.g., prefrontal cortex deploys grid cell-like codes to explore phys- ical (Doeller et al. 2010 Jacobs et al., 2013) and conceptual spaces (Constantinescu et al.) 2016 Bao et al. 2019). well rules in reinforcement learning task (Baram et al.) 1002 The chal- lenge of constructing or observing the relevant space has proven a hurdle for previous supervised methods, being unable discover cognitive maps without a priori knowledge of the represented covariates. Our model is particularly suited in this setting, offering a targeted method to identify neural ensembles and summarizing as well as leveraging their shared tuning properties.
Plato (428-348 BC), Phaedrus 265e; also, Zhuangzi (369-286 BC), Nourishing the Lord of I Life.

---

Published as a conference paper at ICLR 2023
ETHICS STATEMENT
More generally, reducing high dimensional nonlinear dynamical systems set of latent variables can have a broad array of possible applications, for instance, in climate sciences, crop forecasting or flood prediction. Restraining the decoders to interpretable mechanisms and groups of observations could, also in those cases, help us further our understanding of the underlying processes or relations between low dimensional summaries and data observations. As for any complex and versatile mod- machine learning method, there also cxists danger of translating the insights from this work to detrimental purposes and intents. Applications to sensitive data with protected attributes, such as gender ethnicity, should pay special attention the en- and decoding of those attributes in the learned latent spaces. Therefore, while we do not see any obvious misuse, nor want to explicitly name any possible malicious purposes, we still strongly discourage any nefarious applications the ideas developed in this work. Lastly, we tried to minimize the environmental impact of our research by performing the hyper- parameeeeeetrr search on the GPU cluster of
organization.
REPRODUCIBILITY STATEMENT
Code and implementation are included as supplementary material, and is also made available at: https://gthbb..mm/aaaddkiiieiee
Datasets are public, and details regarding model pa- rameters, data generation, etc. are specified under the relevant sections (e.g., Sec. 4), as well as more thoroughly in the Appendix Appx. & K
comm/avidl-lindd/NeualllVa
ACKNOWLEDGMENTS
We thank Erik Hermansen for valuable discussions and feedback related the grid cell data, as well as Ta-Chu Kao for fruitful discussions. This work was supported by a Norwegian Research Council Large-scale Interdisciplinary Researcher Project Grant iMOD (NFR grant no. 14).
REFERENCES
Thomas D Albright. Direction and orientation selectivity of neurons in visual area mt of the macaque. Journal of neurophysiology, 52(6): 106-1130, 1984.
Tom Baden, Philipp Berens, Katrin Franke, Miroslav Román Rosón, Matthias Bethge, and Thomas Euler. The functional diversity of retinal ganglion cells in the mouse. Nature, 529(7586):345-3333 2016.
Xiaojun Bao, Eva Gjorgieva, Laura K Shanahan, James D Howard, Thorsten Kahnt, and Jay A Gottfried. Grid-like neural representations support olfactory navigation of a two-dimensional odor space. Neuron, 10265):1066-175, 2019.
Alon Boaz Baram, Timothy Howard Muller, Hamed Nili, Mona Maria Garvert, and Timothy Ed- ward John Behrens. Entorhinal and ventromedial prefrontal cortices abstract and generalize the structure of reinforcement learning problems. Neuron, 109(4):713-723, 2021.
Mohammad Bashiri, Edgar Walker, Klneses Klemens Lurz, Akshay Jagadish, Taliah Muham- mad, Zhiwei Ding, Zhuokun Ding, Andreas Tolias, and Fabian Sinz. A flow-based latent state generative model of neural population responses to natural images. Advances in Neural Informa- tion Processing Systems, 34, 2021.
Eleanor Batty, Josh Merel, Nora Brackbill, Alexander Heitman, Alexander Sher, Alan Litke, EJ Chichilnisky, and Liam Paninski. Multilayer recurrent network models of primate retinal ganglion cell responses. In 5th International Conference on Learning Representations, 2017.
James Bergstra and Yoshua Bengio. Random search for hyper-parameter optimization. Journal of machine learning research, 13(2), 2012.
Max F Burg, Santiago A Cadena, George H Denfield, Edgar Y Walker, Andreas S Tolias, Matthias Bethge, and Alexander S Ecker. Learning divisive normalization in primary visual cortex. PLOS Computational Biology, 1706)21009028, 2021.
10

---

Published as a conference paper at ICLR 2023
Santiago A Cadena, Fabian H Sinz, Taliah Muhammad, Emmanouil Froudarakis, Erick Cobos, Edgar Walker, Jake Reimer, Matthias Bethge, Andreas Tolias, and Alexander Ecker. How well do deep neural networks trained on object recognition characterize the mouse visual system? 2019.
Gunnar Carlsson, Tigran Ishkhanov, Vin De Silva, and Afra Zomorodian. On the local behavior of spaces of natural images. International journaal of computer vision. 76(1):1-12 2008.
Luis Carrillo-Reid, Jae-eun Kang Miller, Jordan P Hamm, Jesse Jackson, and Rafael Yuste. En- dogenous sequential cortical activity evoked by visual stimuli. Journal of neuroscience, 35(23): 8813-8828, 2015.
Rishidev Chaudhuri, Berk Gerçek, Biraj Pandey, Adrien Peyrache, and lla Ficte. The intrinsic attractor manifold and population dynamics of canonical cognitive circuit across waking and sleep. Nature neuroscience, 22(9):1512-1520) 2019.
Mark M Churchland, John Cunningham, Matthew Kaufman, Justin Foster, Paul Nuyujukian, Stephen Ryu, and Krishna V Shenoy. Neural population dynamies during reaching. Nature, 487 (7405):51-56, 2012.
Alexandra O Constantinescu, Jill X O'Reilly, and Timothy EJ Behrens. Organizing conceptual knowledge in humans with a gridlike code. Science, 352(6292):14444111 2016.
Ronald James Cotton, Fabian Sinz, and Andreas Tolias. Factorized neural processes for neural processes: K-shot prediction of neural responses. Advances in Neural Information Processing Systems, 33:11368:11- 1368-11379, 2020.
Carina Curto. What can topology tell us about the neural code? Bulletin of the American Mathe- matical Society, 54(1):63-78, 2017.
Tim Davidson, Luca Falorsi, Nicola De Cao, Thomas Kipf, and Jakub M Tomezak. Hyperspheri- cal variational auto-encoders. arXiv preprint TrXiv:/804.0099,, 2018.
Christian F Doeller, Caswell Barry, and Neil Burgess. Evidence for grid cells in human memory network. Nature, (663(728)):677-66 2010.
Alexander S Ecker, Fabian Sinz, Emmanouil Froudarakis, Paul G Fahey, Santiago Cadena, Edgar Y Walker, Erick Cobos, Jacob Reimer, Andreas S Tolias, and Matthias Bethge. A rotation-equivariant convolutional neural network model primary visual cortex. arXiv preprint arXiv: 10504, 2018.
Gamaleldin Elsayed and John Cunningham. Structure in neural population recordings: an expected byproduct of simpler phenomena? Nature neuroscience, (009)1310-13111 2017.
Luca Falorsi, Pim de Haan, Tim R Davidson, and Patrick Forré. Reparameterizing distributions on lie groups. In The 22nd International Conference on Artificial Intelligence and Statistics, pp- 3244-3253 PMLR, 2019.
Katrin Franke, Konstantin F Willeke, Kayla Ponder, Mario Galdamez, Taliah Muhammad, Saumil Patel, Emmanouil Froudarakis, Jacob Reimer, Fabian Sinz, and Andreas Tolias. Behavioral state tunes mouse vision to ethological features through pupil dilation. bioRxiv, 2021.
Mark C Fuhs and David S Touretzky, A spin glass model of path integration in medial entorhinal cortex. Journal of Neuroscience, 26( 16):6666644 2006.
Marianne Fyhn, Sturla Molden, Menno P Witter, Edvard I Moser, and May-Britt Moser. Spatial representation in the entorhinal cortex. Science, 305(5689) 2004.
Marianne Fyhn, Torkel Hafting, Alessandro Treves, May-Britt Moser, and Edvard I Moser. Hip- pocampal remapping and grid realignment in entorhinal cortex. Nature, 446(7132);190-199)) 2007.
Juan A Gallego, Matthew G Perich, Lee E Miller, and Sara A Solla. Neural manifolds for the control of movement. Neuron, 94(5):978-98) 2017.
11

---

Published as a conference paper at ICLR 2023
Peiran Gao, Eric Trautmann, Byron Yu, Gopal Santhanam, Stephen Ryu, Krishna Shenoy, and Surya Ganguli. A theory of multineuronal dimensionality, dynamics and measurement. BioRxiv, pp- 214262, 2017.
Richard Gardner, Erik Hermansen, Marius Pachitariu, Yoram Burak, Nils A Baas, Benjamin A Dunn, May-Britt Moser, and Edvard Moser. Toroidal topology of population activity in grid cells. Nature, pp. 1-6, 2022.
Partha Ghosh, Arpan Losalka, and Michael J Black. Resisting adversarial attacks using gaussian mixture variational autoencoders. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 33, pp. 2019..
2019.
Matías A Goldin, Baptiste Lefebvre, Samuele Virgili, Alexander Ecker, Thierry Mora, Ulisse Fer- rari, and Olivier Marre. Context-dependent selectivity natural scenes in the retina. bioRxiv, 2021.
Torkel Hafting, Marianne Fyhn, Sturla Molden, May-Britt Moser, and Edvard I Moser. Microstruc- ture of spatial map in the entorhinal cortex. Nature, 436(7052):001-866, 2005.
Jordan P Hamm, Yuriy Shymkiv, Shuting Han, Weijian Yang, and Rafael Yuste. Cortical ensembles selective for context. Proceedings of the National Academy of Sciences, 1 118(14), 2021.
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): 375-387, 2017.
Christopher D Harvey, Philip Coen, and David W Tank. Choice-specific sequences in parietal cortex during virtual-navigation decision task. Nature, 484(7392):62-68, 2012.
D. H. Hubel and T. N. Wiesel. Brain mechanisms of vision. Scientific American, 241:130, 1979.
Joshua Jacobs, Christoph T Weidemann, Jonathan F Miller, Alec Solway, John Burke, Xue-Xin Wei, Nanthia Suthana, Michael R Sperling, Ashwini D Sharan, Itzhak Fried, et al. Direct record- ings of grid-liks like neuronal activity human spatial navigation. Nature neuroscience, 16(9):1188- 1190, 2013.
Kristopher Jensen, Ta-Chu Kao, Marco Tripodi, and Guillaume Hennequin. Manifold gplvms for discovering non-euclidean latent structure in neural data. Advances in Neural Information Pro- cessing Systems, 133:22500-22599 2020.
Kristopher T Jensen, David Liu, Ta-Chu Kao, Máté Lengyel, and Guillaume Hennequin. Beyond the euclidean brain: inferring non-euclidean latent trajectories from spike trains. bioRxiv, 2022.
James Jun, Nicholas A Steinmetz, Joshua H Siegle, Daniel Denman, Marius Bauza, Brian Barbar- its, Albert Lee, Costas A Anastassiou, Alexandru Andrei, Çağatay Aydın, et al. Fully integrated silicon probes for high-density recording of neural activity. Nature, 551(7679):232-236, 2017.
Benjamin R Kanter, Christine M Lykken, Daniel Avesar, Aldis Weible, Jasmine Dickinson, Ben- jamin Dunn, Nils Z Borgesius, Yasser Roudi, and Clifford G Kentros. A novel mechanism for the grid-to-place cell transformation revealed by transgenic depolarization of medial entorhinal cortex layer ii. Newron, (366):1400-442. 2017.
Jonathan C Kao, Paul Nuyujukian, Stephen Ryu, Mark M Churchland, John Cunningham, and Krishna V Shenoy. Single-trial dynamics of motor cortex and their applications to Train-machine interfaces. Nature communications, 6(1):1-12, 2015.
Matthew T Kaufman, Mark M Churchland, Stephen Ryu, and Krishna V Shenoy. Cortical activity in the null space: permitting preparation without movement. Nature neuroscience, 17(3):440- 448, 2014.
Alexander JE Kell and Josh McDermot.. Deep neural network models of sensory systems: win- dows onto the role of task constraints. Current opinion in neurobiology, 55:121-132, 2019.
12

---

Published as a conference paper at ICLR 2023
Alexander JE Kell, Daniel LK Yamins, Erica N Shook, Sam V Norman-Haignere, Dermott. task-optimized neural network replicates human auditory behavior, predicts brain responses, and reveals cortical processing hierarchy. Neuron, 98(3):630-644 2018.
Haignere, and Josh H Mc-
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.6980, 2014.
Diederik P Kingma and Max Welling, Auto-encoding variational bayes. In Proceedings of the International Conference Learning Representations (ICIR), 2014.
David Klindt, Alexander S Ecker, Thomas Euler, and Matthias Bethge. Neural system identification for large populations separating "what" and "where". Advances in Neural Information Processing Systems, 30, 2017.
David A. Klindt, Lukas Schott. Yash Sharma, Ivan Ustyuzhaninov. Wieland Brendel, Matthias Bethge, and Dylan Paiton. Towards nonlinear disentanglement in natural data with tempo- ral sparse coding. International Conference on Learning Representations, 2021. URL
https::/cpenreev/ net/
Mirko Klukas, Mareus Lewis, and lla Fiete. Efficient and flexible representation of higher- dimensional cognitive variables with grid cells. PLoS computational biology, 16(4):e1007766 2020.
Nikolaus Kriegeskorte and Kue-Xin Wei. Neural tuning and representational geometry. Nature Reviews Neuroscience, 22(11):703-718, 2021.
Emilio Kropff,, James E Carmichael, May-Britt Moser, and Edvard Moser. Speed cells in the medial entorhinal cortex. Nature, 523(7561):419 424, 2015.
Neil Lawrence. Gaussian process latent variable models for visualisation of high dimensional data. Advances in neural information processing systems, 16, 2003.
Michael S Lewicki. Efficient coding of natural sounds. Nature neuroscience, 5(4):356-33) 2002.
Justin D Lieber and Sliman Bensmaia. High-dimensional representation texture somatosen- sory cortex of primates. Proceedings of the National Academy Sciences, 116(8):3268-35 -277,, 2019.
Vítor Lopes-dos Santos, Sergio Conde-Ocazionez, Miguel AL Nicolelis, Sidarta T Ribeiro, and Adriano BL Tort. Neuronal assembly detection and cell membership specification by principal component analysis. PloS one, 6(6):e0096,, 201
Lurz, Mohammad Bashiri, Konstantin Willeke, Akshay K Jagadish, Eric Wang, Edgar Walker, Santiago A Cadena, Taliah Muhammad, Erick Cobos, Andreas Tolias, et al. Generalization in data-driven models of primary visual cortex. bioRxiv, pp. 2020-10, 10, 2021.
Valerio Mante, David Sussillo, Krishna Shenoy, and William T Newsome, Context- dependetd computation by recurrent dynamics in prefrontal cortex. Nature, 503(7474):78-84, 2013.
Alexander Mathis, Andreas VM Herz, and Martin Stemmler. Optimal population codes for space: grid cells outperform place cells. Neural computation, 24(9):2280-233 2012.
Bruce L McNaughton, Francesco P Battaglia, Ole Jensen, Edvard Moser, and May-Britt Moser: Path integration and the neural basis of the' cooniiive map'. Nature Reviews Neuroscience, 7(8): 663-678, 2006.
Bartul Mimica, Benjamin A Dunn. Tuce Tombaz, VPTNC Srikanth Bojja, and Jonathan R Whitlock. Efficient cortical coding 3d posture in freely behaving rats. Science, 362(6414)) (1)66441))588888) 2018.
Aran Nayebi, Javier Sagastuy-Bress Brena, Daniel M Bear, Kohitij Kar, Jonas Kubilius, Surya Ganguli, David Sussillo, James DiCarlo, and Daniel LK Yamins. Goal-driven recurrent neural network models of the ventral visual stream. bioRxiv, 2021.
13

---

Published as a conference paper at ICLR 2023
Chethan Pandarinath, Daniel J O' O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D Stavisky, Jonathan C Kao, Eric Trautmann, Matthew Kaufman, Stephen Ryu, Leigh R Hochberg, et al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10)8805-8) 805.815,-88808880 2018.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. http::/pppers..sepppme/// 8024-8035, Curran Associates, Inc., 2019. URL http: /ptpppaapssssssssseeeeeeeaeeeeee hwwppppprrsuuurpppppeepepppepeee nwwrips. cc/paper/ 9015 patlppyytrccchahe an -n-immaatiive style high hetth-eerfrrrmaam deceeicee-eeeppee -eee-dddeeeppppeeeeeeeeee learning libraryy..... pdf
Felix Pei, Joel Ye, David M. Zoltowski, Anqi Wu, Raeed H. Chowdhury, Hansem Sohn, Joseph E. O'Doherty, Krishna V. Shenoy, Matthew T. Kaufman, Mark Churchland, Mehrdad Jazayeri, Lee E. Miller, Jonathan Pillow, II Memming Park, Eva L. Dyer, and Chethan Pandarinath. Neu- ral latents benchmark '21: Evaluating latent variable models of neural population activity. In Advances in Neural Information Processing Systems (NeurIPS), Track on Datasets and Bench- marks, 2021. URL https: /awpiv.ivvvvva////ia//a///aaavaa org/ abaaaass/1100 04463
Adrien Peyrache, Marie M Lacroix, Peter Petersen, and György Buzsáki. Internally organized mechanisms of the head direction sense. Nature neuroscience, 18(4):56--575 2015a.
Adrien Peyrache, Peter Petersen, and György Buzsáki. Extracellular recordings from multi-site silicon probes in the anterior thalamus and subicular formation of freely moving mice. CR- CNS.org. http: //tx /wwp/xxxxxii..oo////oo/////... org 60800600100000000000 2015b.
Yongrong Qiu, Zhijian Zhao, David Klindt, Magdalena Kautzky, Klaudia P Szatko, Frank Schaeffel, Katharina Rifai, Katrin Franke, Laura Busse, and Thomas Euler. Mouse retinal specializations reflect knowledge of natural environment statistics. bioRxiv, 2020.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and ap- proximate inference in deep generative models. In International conference on machine learning, pp. 1278-1286. PMIR, 2014.
Alon Rubin, Liron Sheintuch, Noa Brande-Filat, Or Pinchasof, Yoav Rechavi, Nitzan Geva, and Yaniv Ziv. Revealing neural correlates of behavior without behavioral measurements. Nature communications, 10(1):1-14 2019.
Michael E Rule, Timothy O'Leary, and Christopher D Harvey. Causes and consequences of repre- sentational drift. Current opinion in newrobiology, 58: 181-141 2019.
Erik Rybakken, Nils Baas, and Benjamin Dunn. Decoding of neural data using cohomological feature extraction. Neural computation, 31(():68-9)) (01]))66-33,100 2019.
Patrick T Sadtler, Kristin M Quick, Matthew D Golub, Steven M Chase, Stephen Ryu, Elizabeth C Tybar-Kabara, Byron M Yu, and Aaron P Batista. Neural constraints on learning. Nature, 512 (7515):423 426. 2014.
Shahd Safarani, Aree Nix, Konstantin Willeke, Santiago Cadena, Kelli Restivo, George Denfield, Andreas Tolias, and Fabian Sinz. Towards robust vision by multi-task learning on monkey visual cortex. Advances in Neural Information Processing Systems, 34, 2021.
Francesca Sargolini, Marianne Fyhn, Torkel Hafting, Bruce L McNaughton, Menno P Witter, May- Britt Moser, and Edvard Moser. Conjunctive representation of position, direction, and velocity in entorhinal cortex. Science, 312(5714)758--722 2006.
Elad Schneidman, Michael J Berry, Ronen Segev, and William Bialek. Weak pairwise correlations imply strongly correlated network states in a neural population. Nature, 44077087) 2006.
14

---

Published as a conference paper at ICLR 2023
Lukas Schott, Jonas Rauber, Matthias Bethge, and Wieland Brendel. Towards the first adversarially robust neural network model on mnist. arXiv preprint arXiv:/805.091900 2018.
Katja Seeliger, Luca Ambrogioni, Yagmur Güglütürk, Leonieke M van den Bulk, Umut Güçlü,, and MAJ van Gerven. End-to-end neural system identification with neural information flow. PLOS Computational Biology, 17(2):e 1008558, 2021.
Nishal P Shah, Nora Brackbill, Ryan Samarakoon, Colleen Rhoades, Alexandra Kling, Alexander Sher, Alan Litke, Yoram Singer, Jonathon Shlens, and EJ Chichilnisky. Individual variability of neural computations in the primate retina. Neuron, 110(4):698-708 2022.
Krishna V Shenoy, Maneesh Sahani, and Mark M Churchland. Cortical control of arm movements: dynamical systems perspective. Anmual review of neuroscience, 36:337-359, 2013.
Gurjeet Singh, Facundo Memoli, Tigran Ishkhanov, Guillermo Sapiro, Gunnar Carlsson, and Dario I. Ringach. Topological analysis population activity in visual cortex. Journal of vi- sion, 888):11-11, 2008.
Fabian Sinz, Alexander Ecker, Paul Fahey, Edgar Walker, Erick Cobos, Emmanouil Froudarakis, Dimitri Yatsenko, Zachary Pitkow, Jacob Reimer, and Andreas Tolias. Stimulus domain transfer in recurrent models for large scale cortical population prediction on video. Advances in neural information processing systems, 31, 2018.
William Skaggs, Bruce Mcnaughton, and Katalin Gothard. An information theorrtti approach to deciphering the hippocampal code. Advances in neural information processing systems, 5, 1992.
Trygve Solstad, Edvard Moser, and Gaute T Einevoll. From grid cells to place cells: mathematical model. Hippocampus, 16(12): 1026-1031, 2006.
Sameet Sreenivasan and Ila Fiete. Grid cells generate an analog error-correcting code for singularly precise neural computation. Nature neuroscience, 14(10):1330-1337 2011.
Hanne Stensola, Tor Stensola, Trygve Solstad, Kristian Froland, May-Britt Moser, and Edvard I Moser. The entorhinal grid map is discretized. Nature, 492(7427):72-78, 2012.
Ian H Stevenson and Konrad P Kording. How advances in neural recording affect data analysis. Nature neuroscience, 14(2):139-1) 2011.
Mark G Stokes, Makoto Kusunoki, Natasha Sigala, Hamed Nili, David Gaffan, and John Duncan. Dynamic coding for cognitive control in prefrontal cortex. Neuron, 78(2):364-37 2013.
Carsen Stringer, Marius Pachitariu, Nicholas Steinmetz, Matteo Carandini, and Kenneth D Harris. High-dimensional geometry of population responses in visual cortex. Nature, 571(7765):361- 365, 2019.
Jeffrey S Taube, Robert Muller, and James B Ranck. Head-direction cells recorded from the post- subiculum in freely moving rats. i. description and quantitative analysis. Journal of Neuroscience, 10(2):400435, 1990.
John J Tukker, Prateep Beed, Michael Brecht, Richard Kempter, Edvard I Moser, and Dietmar Schmitz. Microcircuits for spatial coding in the medial entorhinal cortex. Physiological reviews, 102(2):653-688, 2022.
Ivan Ustyuzhaninov, Santiago A Cadena, Emmanouil Froudarakis, Paul Fahey, Edgar Walker, Erick Cobos, Jacob Reimer, Fabian H Sinz, Andreas S Tolias, Matthias Bethge, et al. Rotation- invariant clustering of neuronal responses in primary visual cortex. International Conference Learning Representations, 2019.
Ivan Ustyuzhaminov, Max F Burg, Santiago Cadena, Jiakun Fu, Taliah Muhammad, Kayla Ponder, Emmanouil Froudarakis, Zhiwei Ding. Matthias Bethge, Andreas Tolias, et al. Digital twin veals combinatorial code of non-line non-linear computations in the mouse primary visual cortex. bioRxiv, 2022.
15

---

Published as a conference paper at ICLR 2023
Edgar Y Walker, Fabian H Sinz, Emmanouil Froudarakis, Paul G Fahey. Taliah Muhammad, Alexan- der Ecker, Erick Cobos, Jacob Reimer, Xaq Pitkow, and Andreas Tolias. Inception visual cortex: in vivo-silico loops reveal most exciting images. bioRxiv, pp. 506956, 2018.
Wässle and HJ Riemann. The mosaic of nerve cells in the mammalian retina. Proceedings of the Royal Society of London. Series B. Biological Sciences, 200( 141):441-461, 1978.
Xue-Xin Wei, Jason Prentice, and Vijay Balasubramanian. A principle of economy predicts the functional architecture of grid cells. Elife, 4:e08362, 2015.
Anqi Wu, Nicholas A Roy, Stephen Keeley, and Jonathan W Pillow. Gaussian process based non- linear latent structure discovery in multivariate spike train data. Advances in neural information processing systems, 30, 2017.
Anqi Wu, Stan Pashkovski, Sandeep R Datta, and Jonathan W Pillow. Learning a latent manifold of odor representations from neural responses in piriform cortex. Advances in Neural Information Processing Systems, 31, 2018.
Byron M Yu, John P Cunningham, Gopal Santhanam, Stephen Ryu, Krishna V Shenoy, and Ma- neesh Sahani. Gaussian-procest factor analysis for low-dimensional single-trial analysis of neural population activity. Advances in neural information processing systems, 21, 2008
Zhijian Zhao, David A Klindt, André Maia Chagas, Klaudia P Szatko, Luke Rogerson, Dario A Protti, Christian Behrens, Deniz Dalkara, Timm Schubert, Matthias Bethge, et al. The temporal structure of the inner retina single glance. Scientific reports, 10(1):1-17, 2020.
Yi Zhou, Connelly Barnes, Jingwan Lu, Jimei Yang, and Hao Li. On the continuity of rotation representations in neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 5745-5753, 2019.
Chengxu Zhuang, Siming Yan, Aran Nayebi, Martin Schrimpf, Michael Frank, James DiCarlo, and Daniel LK Yamins. Unsupervised neural network models of the ventral visual stream. Pro- ceedings of the National Academy of Sciences, 118(3), 2021.
16

---

Published as a conference paper at ICLR 2023
APPENDIX
A MODEL DETAILS
All models are trained with the Adam optimizer (Kingma & Ba 2014), a learning rate of 0.001, temporal chunks of size 128 (64 if data length is smaller than 128) and a batch size of 1. Moreover, training is concluded after the training objective has not improved for 5 steps (10 for grid cell data). All models are implemented in Pytorch (Paszke al. 2019). The architecture of the encoder is as follows: a ID grouped convolution with kernel size 9; followed by two layers of ID convolutions with kernel size (i.e., effectively just MLPs each time point) and 64 filters (128 for grid cell data), followed by linear projection onto the dimensionality of the mean and variance parameters required by the variational posterior. The decoder follows the feature sharing tuning curve structure as presented in the paper. All experiments where executed with specified random seeds for data generation, model initialisation and training batch selection to maintain complete reproducibilty.
A.1 A FAST APPROXIMATION TO THE OBJECTIVE FUNCTION
Next, we present our simplification of the theoretical setup from DDavidson et al. 2018 concerning the use of von Mises distributions when inferring latent variables circular manifolds. First, we note that as the variance x² of a Gaussian and the inverse concentration 1/K of a von Mises approach 0, they both become delta functions. Moreover, they each define the maximum entropy distribution for fixed variance (concentration) in their respective domains. Hence, the von Mises is also referred to as the circular normal distribution.
We therefore explored the model behavior when replacing the KL term in Eq. with a simpler expression, akin to the standard VAE KL, which simply pushes the variance the variational en- coder towards 1 (i.e., standard normal), but without constraints on the mean (uniform over the latent space, as the full von Mises prior). Reparameterisation and sampling was then performed by draw- ing from that Gaussian and adding the scaled perturbation (same as in standard VAFs) to the angle that the variational posterior returned as mean estimate.
In general. performing variational inference in -uclldelid spaces is challenging Falorsi et al. 2019), which has warranted the use of specialized inference procedures in previous latent variable models on such manifolds (Davidson et al. 2018 Falorsi et al. 2019 Jensen al. 2020 2022). In this work, we follow the 'ReLie' approach outlined by Falorsi et al. (019) for general Lie groups, which involves defining a parameterizable (in our case Gaussian) distribution the tangent space of the group and projecting it onto the manifold using the so-called 'exponential map' of the group. This approach is particularly simple for the toroidal manifolds we consider, where the projection step can be achieved by simple modulo 2T operation (Jensen et al. 2020). This leaves the challenge of computing the KL term in Eq. II which is given by
(l(((())lo())) = Fyy()) (() - logp(=)).
As prior p(=), we use a unit Gaussian projected onto the circle, with a mean matched to the poste- rior mean. Since the KL divergence in Eq. Glis invariant to rotation around the circle (Falorsi al.) 2019), we can compute it simply as the KL divergence between two wrapped Gaussian distributions centered the origin with variances of and of respectively. In general, the projected densities needed in Eq. 0 are given by
Lo (z) =
Σ ZER eepp (x)==
[o(+))
as derived in Falorsi et al. (2019). Here, expc is the exponential map. (xx) is the Jacobian at æ in the tangent space, and ra is the parameterized reference distribution on the tangent space.
For the circle, the Jacobian is simply given by the identity J(x)) 1 1. With a zero-centered
Gaussian reference distribution, we can therefore rewrite Eq. 7 as
(o(2)
Σ A(22 + 2kT; p 0, oz))
(8)
kez
17

---

Published as a conference paper at ICLR 2023
and similarly for the prior p(2) (Jensen et al. 2020). To approximate this infinite sum, it necessary to truncate it after a finite number of elements (Falorsi et al.) 2019; Jensen et al. 2020 2022). In the case where og < 1, it is sufficient to outside the area of injectivity (IT, x)) This is in fact the regime we are operating in, since less than 0.2% of the prior probability mass falls outside the interval -,, =] and the posterior variance will in general be smaller than the prior variance.
single term since there negligible probability mass
Additionally, while the integral in Eq. covers the domain [- m], we can approximate it with an integral from -oo to Do since the posterior probability mass outside the area of injectivity is These approximations together imply that we can simply approximate the KL term in Eq. aas the KL between our prior and posterior Gaussians directly in the tangent space, and we verified numerically that this approximation excellent for o² S 1, We thus obtain a simple analytical expression for the KI.:
negligible, Ja -== (x; 0, og )dz
KL(g(=)lp(z)) ~ KL [M 0, o,)1
(9)
2 - -1- og)
(10)
Note that this is only possible because prior is localized, in contrast to previous work which often uses a more diffuse or even uniform prior distribution (Falorsi et al. 2019 Jensen et al. 2020).
and because consider circular latent spaces where the Jacobian is 1. Since our prior and posterior both factorize across dimensions, these considerations also generalize to higher dimensional spaces, where the multivariate KL divergence simply the sum across dimensions of the corresponding univariate KL divergences.
Finally, verify that these approximations do not affect our results or conclusions, we repeated several key analyses with the variational inference framework described in Davidson et al. (2018). The simplified setup trained about 5x faster, while behaving nearly identically on all test cases that we evaluated, hence, all experiments in the paper are performed with the simplified objective. We note that this simplified procedure may be of interest to those who might need to iterate over models without access to a large compute cluster.
B DATA GENERATION FOR SIMULATIONS
When conducting the experiments that showcase the scaling of various models with respect to the amount of data (Sec. 4.1). we generate the spikes according to Gaussian-like tuning curves. More specifically, we assume a common shape for each curve (location assigned randomly on S²), with a tuning width of 1.2 radians. Peak rates are set to 0.5 spikes per bin, the background rate to 0.005, with a signal-to-nn -o-iee ratio of for the Poisson noise. The latent variable is simulated according to a Gaussian process with kernel standard deviation equal to 5.0, and with kernel scale 50.0. is then projected onto g via the modulo operation. The ensemble detection experiments in Sec. 4.2 use the same data generation, only with multiple (2 or 3) ensembles, with distinct latent varibles for each ensemble.
C ENSEMBLE DETECTION COMPARISONS
For raw-pca, we performed k-means clustering on dimensionally reduced neural activity over time (loosely inspired by Lopes-dos Santos et al. 2011 Baden et al. 2016) Hamm et al.] (02TT. Specif- ically, we took the projection of the spike matrix onto its first 8 (cross-validaade validated) principal compo- nents. We then computed k-means clustering on this reduced matrix with the standard sciikitllean implementation, using fixed number of 2 (3, depending on the setting) clusters, 100 initializations and 1000 maximum iterations.
For cov-pca, we performed clustering on the neural covariance matrix (loosely inspired by Carrillo- Reid et al., 2015). Specifically, we computed the covariance across neurons from the spike matrix. Again, we computed the first 8 (cross-validated) principal components of the absolute value of this covariance matrix. The absolute value was taken to preserve any negative or positive dependencies across neurons, one would presumably observe within ensemble removing the modulus operator drastically reduced ensemble detection accuracy. From the projection onto the principal components we proceeded again as above. We also performed clustering of the inverse of the
18

---

Published as a conference paper at ICLR 2023
covariance matrix, first order approximation to an Ising model (Schneidman et al. 2006), but this approach did not yield results above chance level, hence it was not included in the figures.
Finally, for the supervised method performed clustering the mutual information score (see scikit-learn for implementation) between each neuron's s activity, and the true latent variables, as imput. Specifically, the features y that used input to k-means (same settings as above) were Vij ~ [ 2,] for every neuron i and every ground truth latent dimension j. This presents an upper bound to the achievable performance, since assumes supervised knowledge of the true latent variable values.
D THE EVIDENCE LOWER BOUND (ELBO)
Here we show the derivation of the ELBO, which maximized as part the objective function during the model training procedure
logp(x) log / p(zl))p(z)dz= = log p(xrl))p( ((((2)
((2(2x)
log Esng(ele) p(x])z) ((2)
Jensen's inequality
E Esng(sla) logp((]=) p(z)
(11)
((=(0r) )
((=[x)
Esoe(=(x) [log (ogp(zr(z))
E Ezng(s)a) (llgg(=())))))))))))))))())())) log [og)())] ELBO(g).
-KLIqi
D.1 TEMPORAL TRANSITION PRIORS
All experiments in the article are conducted using the objective function from Eq. Extending this with additional terms however, e.g. by including a Laplacian temporal transition prior (Klindi et al., 2021) or a Gaussian transition prior (i.e. an L1 or L2 penalty across latent time steps, respectively) relatively simple, and both the L1 and L2 penalties are included as optional regularization terms our model. While exhaustive experiments have not been performed regarding these priors, results seem to indicate that the inclusion of these priors do little to improve the model performance in the settings we have studied the article, which is in agreement with results in Appendix Section and K regarding the convolutional kernel size.
E ABLATION STUDY CONVOLUTIONAL FILTER
As mentioned in Sec. we conduct similar model performance comparison between the faeLVMs and mGPLVM, while reducing the kernel size of the convolutional layer for the faeLVMs to 1 (effectively allowing the encoder network to let the variational posterior only depend on the instantaneous neural population activity). This is done to reduce the potential gap in information content available between convolutional and non-convolutional model. Note, however, that this applics only to the variational posterior. At inference time, both models infer the best latent at cach point in time separately. Experimental settings are precisely the same those in Sec. 4.1 other than the size of the convolutional kernel.
One of the main observations from Fig. BC-D is that the faeLVMs still outperform mGPLVM, without contributions from the convolutional kernel. In fact, faeLVM-s performs much better now, compared with the results from Fig. achieving results that rivals those of MacL.V-b (Fig. 8A-B)
F MODEL EXTENSION CALCIUM RESPONSES AND VISUAL DATA SIMULATION
In an effort to improve on the applicability of faeLVM, and showcase its relevance regarding appli- cations to different brain regions, different latent topologies and different input data, include toy example on visual data, where we modeled a moving dot stimulus (Fig. (AA) and a population response of retinal ganglion cells with center surround receptive fields and Poisson spiking (Fig.
19

---

Published as a conference paper at ICLR 2023
A
D
Latent gred, [GE]
Rate [NLL)
1.00
8.6
0.95
1.4-
Teat
0.5
10
20
90
40
50
2
2
50
40
50
Myrio
Moraic
Tome
TooLUW5nn
Taal996s
TceLU/M-
MGP
Figure 8: Model performance comparison, with convolutional filter size 1. The models faeLVM- n.s.b are described in the main text, mGP the model from Jensen al. (2020). A. Mean NLLH (negative log-likelihood, lower is hetter) of predicted test neuron rate, function of amount of training neurons. As the performance is averaged over 20 seeds of varying data, error bars would also reflect the variability in NLLH across data, resulting in inappropriate visualization. Hence, report the mean rank of 20 seeds for each model at each condition, in an effort to more appropriately capture the variability of the models. B, Same in A, now function of training time points. C, Mean GE (geodesic error, lower is better) between true latent and inferred latent, on test data, with corresponding SEM error bars, function of number of training neurons. D, Same in C, now as function of amount of training time points.
A
B
C
D
E
F
40
Figure 9: Visual data, toy example. A. Scatter plot of trajectory of simulated visual stimulus, colored according its temporal evolution. B. Receptive fields for three training neurons. C. Same three training neurons, showing firing rate in shaded dark blue and the corresponding Poisson spikes in orange. D, Neural responses after passing the neural activity from C through a calcium kernel. E, Scatter plot of true latent trajectory against inferred latent trajectory (variational mean). F True Ca2+ neural responses (grey) plotted against predicted neural responses for three test neurons (black).
BCC.). Each neuron's ' spiking response was then convolved with a double exponential calcium flu- orescence kernel (Fig. (). We further extended fa by including Gaussian likelihood modeling option, as well as an option to infer latent trajectories not restricted to circular and toroidal latent manifolds.
Results can be seen in Fig. 9 B-F, where the model is able to infer the correct two-dimensional latent variable, as well accurately predict calcium responses test neurons on data. Although this experiment is less comprehensive, being toy example, it still gives a clear indication that our model also applicable both calcium traces and non-grid and head direction cells.
G ALL GRID CELL RATE MAPS
In this section, we showcase rate maps for all neurons in each of the two ensembles (149 and 145 neurons in the first and second grid cell module respectively). The data used is the same as discussed in Sec. 4.4, pre-proee proceeccd in the same way, and we used the same model to produce latent decodings, receptive fields and ensembles.
Fig. 10 and Fig. II show the rate maps (using 50 x 50 spatial bins, smoothed for visualization purposes in the manner Gardner al. (2022,, with Gaussian kernel of smoothing width
20

---

Published as a conference paper at ICLR 2023
Open Feld
Tereical
Open Field
Torokcal
Open
Torokal
Open Fieid
Fidid
Toraidal
Open Field
Toroidal
Open Fiedi
dalo
Open Field
Taroidal
O
O
O
Figure 10: All rate maps, first grid cell module.
21

---

Published as a conference paper at ICLR 2023
Open
Tornical
Open
Toroidal
Open
Open
Down
Open Field
O
O
C
Figure 11: All rate maps, second grid cell module.
22

---

Published as a conference paper at ICLR 2023
A
Figure 12: Supplementary visu- alizations for analyses on real datasets. A, Recorded mouse head direction ("True HD') and inferred latent trajectories
(6)
data. Re- sults from the experiment on head direction data, Sec. 4.3 B. Dy- namics of the inferred toroidal co- ordinates plotted as function of the (x.,) v) spatial position of the animal (two tori, each with two circular co- ordinates, hence four figures). Re- sults from the experiment on grid cell data, Sec.
600
1000
B
2.75), both as function of spatial position and inferred toroidal coordinates. For the rate maps in torodial coordinates, we also include the inferred receptive field center and width. Neural rate maps are shown in decreasing spatial information content (Skaggs et al. 1992 Gardner et al. 2022 going from left to right, top to bottom.
As mentioned in Sec. 4.4 we clearly see the more ambiguous rate maps at the bottom of these two figures (in particular, the three-four bottom- rows). Here, the spatial resolution is not as distinct as for the more informative neurons, and the rate maps less clear, for which one might surmise that these neurons need not necessarily belong to the module originally assigned to.
H ADDITIONAL VISUALIZATIONS FOR HEAD DIRECTION AND GRID CELL DATA
To supplement the results shown in Sec. 4.3 and Sec. 4.4 we have included additional figures for both the head direction application and the grid cell application. More precisely, Fig. 12A shows the inferred variational mean of faeLVM-n plotted against the recorded mouse head direction (now as trajectories, in contrast to the scatterplot in Fig. o while Fig. [2 shows inferred latent dynamics on grid cell data, function of spatial position. As there coordinates (akin to what done in the head direction dataset), this figure is included in an effort demonstrate the clear correspondence between each of the inferred toroidal axes and the spatial position, with the results being in agreement with what was found in Gardner et al. (2022).
way of recording the 'true' toroidal
I EXTENDED ENSEMBLE DETECTION SIMULATION
In this section we present an additional experiment related to the ensemble detection task. Namely, ask the model separate five distinct ring-like ensembles, 2j E S1, complement the experi- ments from Section 4.2 which only considered two and three ensembles.
From Fig. we can see that facLVM-b able separate the five ensembles quite clearly. The weights in Fig. 13 is approaching the desired one-hot encoding, and we see that the five inferred ensembles correspond very nicely with the true ensembles (Fig. 13C). The model is also able to accurately predict spike rates on held out neurons from each ensemble (). Thus, this simulation provides evidence that our model is also able to work on more challenging ensemble separation tasks, and is not only limited to lower ensemble numbers.
J HYBRID INFERENCE
J.1 OVERVIEW
Here, we give an overview on how the hybrid inference is performed. contrast to the traditional variational approach, where uses single sample from the variational posterior prediction for the latent variable, we propose drawing a selection of samples, then performing gradient descent on said latents using the Adam optimizer. Specifically, if are the latents for the test set (as inferred by
23

---

Published as a conference paper at ICLR 2023
Ens.1 1
Ens.2
Ens.3
Ens.4
Ens.5
A
C
6.
Weights
Eno.
90
100
150
2.0
Neurons
B
MM " poreereyy
la.u.]
showe Seaesslll Alas الار M Mmma M
Figure 13: Ensemble detection of 5 circles. A, Learned ensemble weights for the neurons belong- ing five unique circular ensembles. B, True spike rate (grey) predicted spike rate (black) for for five randomly selected test neurons, one from each ensemble. C, Scatter plot of true latent trajectory against inferred latent trajectory.
the encoder and variational posterior) and x the responses of the training neurons on the test set, maximize
max log p(x|2)
(12)
for 2000 steps with the Adam optimizer and learning rate of 0.001. We perform this in parallel on M = 10 distinct samples from the variational posterior and pick the one with the highest final likelihood (note that all samples usually converge to the same value). We emphasize that all model parameters are fixed during hybrid inference, and the optimization only affects the sampled latents.
Empirically, we observe that this not only improves the likelihood for the training neurons (i.e., the objective we are maximizing), but crucially also the likelihood of the test neurons. This means that we are not just overfitting but actually improving our point estimates the latents in this inference procedure. We leave optimization of the full likelihood, including uncertainty estimates on the latents, for future research.
J.2 HYBRID INFERENCE SIMULATIONS
For evaluate of the performance gain with hybrid inference test-time, we consider six variations of the faeLVM; that being the three cases described in Sec. 3] each alternative including case of using hybrid inference and one using the standard variational posterior approach (fae vs v-fae).
As in Sec. 4.1 we divide our data into four parts (Fig. generate 1D periodic latent variable and simulate spikes according to a Poisson distribution with Gaussian tuning curves. We Tiest and Niett (1000 and 30), vary Tiran and Nurin over selection of values, and for each condition, repeat the experiment over seeds (synthetic data is different for each seed, but all of the models trained on the same dataset, ensure that comparisons still valid), reporting mean negative ligelikeldoood (NLLH) and mean geodesic error (GE) on test data.
The results presented in Table indicate that performance. in almost every case, increases substan- tially by including the hybrid inference step. Hence, all other experiments performed in this article done solely using the hybrid inference where applicable (e.g., when applying the model to grid
24

---

Published as a conference paper at ICLR 2023
Table 1: Hybrid inference comparison. Table showing negative log-likelihood (NLLH) and mean geodesic error (GE) for the six models specified (lower is better for both criterions), under varying amounts of synthetically generated training data with different numbers of training neurons (N) and different training dataset lengths (7).
NIT
30 100
30
500
30
1000
15 / 1000
45
1000
LVM
NLLH
GE
NLLH
GE
NLLH
GE
NLLH
GE
NLLH
GE
v-fae-n
12745
0.57
8971
0.28
8781
0.27
9110
0.45
8642
0.27
v-fae-s
11798
0.64
9167
0.34
8949
0.31
9718
0.40
8820
0.31
v-fae-b
10912
0.50
9152
0.36
8830
0.30
9492
0.33
8633
0.26
fac-n 11975
0.58
8775
0.20
8608
0.17
9018
0.44
8415
0.09
fae-s 10963
0.42
8694
0.13
8657
0.14
9460
0.29
8491
0.13
fae-b
10040
0.25
8636
0.12
8501
0.12
9227
0.23
8267
0.05
cell data, evaluations are done tn-ttriinng-samplle and without the focus on latent prediction on test set, hence there is no need for hybrid inference at test-time).
K HYPERPARAMETER SEARCH
To investigate hyperparameters and their effect, we selected relevant parameters and estimated their ranges based on related literature. The relevant parameters can be seen on the left in Table 2
To determine optimal values of hyperparameters, we relied on a random search strategy. A grid search is not feasible due the curse of dimensionality, and random search can also be more effective (Bergstra & Bengio 2012). Furthermore, the random search can easily be used for discrete and continuous variables, that could be on linear or logarithmic scales. The selected scales and ranges, can be found in the middle in Table
In total, we sample 500 hyperparameter configurations. Next, we train and test each configuration for 3 random seeds. We evaluate the models the test accuracy and mean correlation. The final selected model shown on the far right in Table
Lastly, we also measure whether there is an ordinal correlation (Spearman) between the varied fac- tors and the test LLH. The results are shown in Fig. 14 We see that learning the variance (learnvvr) increases the performance average. Also, lower learning rates and larger batch-sizes lead to better results. Lastly, also lower weighting for the KI. term leads to better results.
25

---

Published as a conference paper at ICLR 2023
Table 2: Hyperparameter search. We show selected hyperparameters and their corresponding
ranges.
Hyperparameter
Ranges
Best
kerneLsiiee num_hidden
choice([1, 9, 17])
choice([16, 32, 64, 128, 256, 512])
128
shared
choice((True, Falsel)
False
learn coeff
choice((True, Falsel)
True
learn_mean
choice((True, Falsel)
False
learn_var
(oooice(True,FFlsee))
False
isotropic
choice((True, False])
False
num_basis
choice([1, 2, 81)
4
nonlinearity
choice([exp. softplus])
softplus
batch size
choice([1, 16, 32])
32
hatch length learning_rat num_worse weight kl weight.time time weight,entropy
choice([64, 128])
64
loguniform(1e-4. le-1)
0.001003
choice([10, 50, 100])
10
choice([O., loguniform (1e-9, 1c0)))
0.0
choice([O., loguniform( (1e-9, 1e0)))
0.0
choice([O., loguniform (1-99 Ie0)))
0.0
A
B
pueluee beiler)
te-
Spearmen
23
lernll sire
hoder
shared
leam couff
learn mean
leam, var iottopic
batch stre
basig
batch length
learingg rate num worse
weight kd mndight time weight, entroyy
karel size num madden shared
lean couff
learn mean
laargnic run taais knstoh size
batch, lengrh leaming rare num worse
weight, weight, trmee
weight, antropy
lean
Figure 14: Correlations with test LLH. A, Spearman correlation. B, P-values.
26


Reviews:
[{"review_id": "zKnWwqNM03b", "reviewer": "Reviewer_xqoS", "summary": "", "questions": "", "limitations": "", "rating": 6, "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 authors propose a new latent variable model for neural data (faeLVM). The faeLVM is a generative model that combines the insights provided by tuning curves with the interpretability of low-dimensional latent variables. Specifically, faeLVM models a population of neurons are a collection of ensembles, where neurons in an ensemble are sensitive to a subset of the latent variables and share the functional form of the tuning curve, which the authors call feature sharing. For training, the authors use the VAE framework and utilize a variational posterior and optimize the ELBO. The main point of this paper, in my mind, is to make VAE-style models interpretable by making the decoder portion very simple, while still allowing for an arbitrarily complex encoder.", "summary_of_the_review": "Overall, I like the basic idea of this paper. VAE decoders are too often treated as black boxes in scientific applications where interpretability is centrally important. I think this paper has some flaws and the empirical demonstrations are in idealized datasets. Nonetheless I am giving it the benefit of the doubt and recommending it as a borderline accept for now, pending the authors willingness to address some of the comments above.", "strength_and_weaknesses": "The motivation behind this paper is very nice. The idea is similar to Manifold GPLVMs (Jensen et al. 2020), but in many ways the approach taken in this paper is simpler to implement and conceptualize. I really like the idea of having flexible and interpretable \"deep\" latent variable models for neuroscience.\n\nThe main weakness of this paper is that it uses what I would consider \"easy\" experimental datasets to make the case for simple tuning-curve-like decoders. For example, the grid cell data used pre-selected, non-conjunctive grid cells from MEC. This is a very special subset of the full MEC population, and it would be much more compelling if the method worked well on the raw data. At a minimum, the authors should discuss this limitation and cite [Hardcastle et al. 2017](https://doi.org/10.1016/j.neuron.2017.03.025), which characterizes the messy, conjunctive coding in the full MEC population.\n\nIn short, I think the feature sharing idea proposed by the authors can be useful on certain neural datasets, but it will be a far from universal component you'd like to build into the model. I have my doubts about whether this will become a commonly used model in the neuroscience literature (outside of Neuro-ML overlapping conference venues).\n\nI would also like to see the paper edited for clarity, as discussed below.", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "prmtzNu_Oww", "reviewer": "Reviewer_3QYX", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "This paper introduces a latent variable model for both generating and clustering neural activity. In simulated datasets, they study the performance of the model and its ability to recover underlying ensembles in a synthetic and real-world toroid task.\nOverall, the paper provides an interesting model and a strong set of different evaluations of the method, across different grid cell datasets and a synthetic toroid task where they set up an ensemble detection task. However, there are some concerns about the generalization of the approach to different tasks and neural datasets. \n\n\n", "summary_of_the_review": "This paper provides a new latent variable model for neural population activity that can leverage shared features across different neurons to learn population-level representations and also detect ensembles (neurons that use shared features). The model provides an interesting blend between population-level analysis and more interpretable neuron-level measures of tuning. \n\nIn empirical evaluations, the authors show that their model can provide good predictions of held out neural activity and thus serves as a good generative model, and also can recover underlying ensembles in the data. However, it is unclear how general the model is, how much they rely on simplified tasks where the subgroups are orthogonal or well separated, and how much the geometry of the underlying task matters. The relatively simple conditions in which the model is tested makes it hard to assess the generalization of the approach to different tasks or conditions, or even larger numbers of clusters.", "strength_and_weaknesses": "Strengths:\n+ The method seems to provide some good advances in both its generative capability as well as the ability to predict ensembles.\n\n+ They introduce a new toroid task and measures of ensemble detection in this case.\n\n\nWeaknesses:\n- The authors compare three variants of their general approach - a shared, no-shared feature model, and a heat kernel model. While there is some discussion of how the different models provide varying benefits in different regimes (for synthetic data), the differences between the models and underlying assumptions is not entirely clear. In the experiments it seems that the shared features model doesn't give improvements across the board -- but isn't this is the core motivation of the work? \n\nE.g., What is the difference between the non-shared model (faeLVM-n) and the geometric LVM model (mGP)? It would be helpful to have a more clear description of the ablations and other models included in the comparisons.\n\n- In the experiments, the number of clusters is relatively low (2-3) and there isn’t much discussion about how the number of clusters is selected (or the regularization is performed to encourage soft clustering of neurons). How does the method work when there are more clusters or larger amounts of overlap between ensembles? How do you select the number of ensembles or clusters and/or regularization in the soft clustering assignment?\n\n- The applications and datasets that are examined are somewhat limited and it's unclear how well it will work in different tasks that don’t have strong geometry and when the feature subspaces are non-overlapping. In particular, is it possible to use their method in conditions where ensembles are not as separable or more entangled?\n\n- The set of baselines are limited. There are a number of other methods for latent variable modeling in neural data analysis, including many cited in the work and in the Neural Latents Benchmark cited in the evaluation section. It would be interesting to understand how their model compares to some of these other methods and in other datasets that don’t have clear toroid structure or where neurons are separated as cleanly as in the datasets tested here.\n\n- For the ensemble detection method, what is the performance for the other variants of your approach? What about feature sharing? It seems this would be the natural model used for these experiments (and not just the -b model).", "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": "cGmf-8d5DH2", "reviewer": "Reviewer_qWBb", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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 propose a latent variable model that extracts meaningful low-dimensional representations from large scale neural recordings. The proposed method assumes Gaussian-like tuning curves to the latent variable on the neurons which share an ensemble-wise basis. The spherical variational posterior captures the toroidal manifold of variables such as head direction that is encoded in the neural activity.", "summary_of_the_review": "The authors propose a latent variable model that extracts meaningful low-dimensional representations from large scale neural recordings. The proposed method assumes Gaussian-like tuning curves to the latent variable on the neurons which share an ensemble-wise basis. The spherical variational posterior captures the toroidal manifold of variables such as head direction that is encoded in the neural activity. Strong assumptions such as tuning curve and vMF posterior do help the inference to get latent variable desirably matches the behavior. They also limit the generalization. What more do we know from it about the data than an encoding model with the same tuning curves?", "strength_and_weaknesses": "Strength:\nUsing biologically plausible tuning curves\nSpherical variational posterior captures toroidal manifold\nIdentifying neuron ensembles\n\nWeaknesses:\nStrong assumptions such as tuning curve and vMF posterior limit the scope of the proposed method. Most LVMs for spike trains aim at explaining most of the shared variability of the data using a low-dimensional representation. On the contrary, the demonstration of this work only shows it perfectly fits to coding toroidal behavioral variables. The vMF posterior forces the latent variable to be toroidal and mimicking Gaussian-shape tuning curve to head direction of the target neurons also forces the latent variable to be close the head direction. It is fine as   we know the toroidal head direction and the shape of tuning curves beforehand, but this demonstration makes the method seem less nontrivial. \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": {}}]

Author rebuttal:

Decision: Accept: poster