Original: U NDERSTANDING N EURAL C ODING ON L ATENT M ANIFOLDS BY S HARING F EATURES AND D IVIDING E NSEMBLES
Revised: U NDERSTANDING N EURAL C ODING ON L ATENT M ANIFOLDS BY S HARING F EATURES AND D IVIDING E NSEMBLES

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Systems neuroscience relies on two complementary views of neural data, characterized by single neuron tuning curves and analysis of population activity.
Revised: Systems neuroscience relies on two complementary views of neural data, characterized by single neuron tuning curves and analysis of population activity.

Original: 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.
Revised: 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.

Original: This has recently been demonstrated using Gaussian processes, with applications to realistic and topologically relevant latent manifolds.
Revised: This has recently been demonstrated using Gaussian processes, with applications to realistic and topologically relevant latent manifolds.

Original: Those and previous models, however, missed crucial shared coding properties of neural populations.
Revised: Those and previous models, however, missed crucial shared coding properties of neural populations.

Original: We propose feature sharing across neural tuning curves, which significantly improves performance and leads to better-behaved optimization.
Revised: We propose feature sharing across neural tuning curves, which significantly improves performance and leads to better-behaved optimization.

Original: We also propose a solution to the problem of ensemble detection , whereby different groups of neurons, i.e., ensembles, can be modulated by different latent manifolds.
Revised: We also propose a solution to the problem of ensemble detection , whereby different groups of neurons, i.e., ensembles, can be modulated by different latent manifolds.

Original: This is achieved through a soft clustering of neurons during training, thus allowing for the separation of mixed neural populations in an unsupervised manner.
Revised: This is achieved through a soft clustering of neurons during training, thus allowing for the separation of mixed neural populations in an unsupervised manner.

Original: These innovations lead to more interpretable models of neural population activity that train well and perform better even on mixtures of complex latent manifolds.
Revised: These innovations lead to more interpretable models of neural population activity that train well and perform better even on mixtures of complex latent manifolds.

Original: Finally, we apply our method on a recently published grid cell dataset, recovering distinct ensembles, inferring toroidal latents and predicting neural tuning curves all in a single integrated modeling framework.
Revised: Finally, we apply our method on a recently published grid cell dataset, recovering distinct ensembles, inferring toroidal latents and predicting neural tuning curves all in a single integrated modeling framework.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Neural population activity can appear high-dimensional (Stringer et al., 2019), yet much recent work has reported that neural populations in higher brain areas are often confined to low dimensional subspaces (Yu et al., 2008; Harvey et al., 2012;
Revised: Neural population activity can appear high-dimensional (Stringer et al., 2019), yet much recent work has reported that neural populations in higher brain areas are often confined to low dimensional subspaces (Yu et al., 2008; Harvey et al., 2012;

Original: Mante et al., 2013; Stokes et al., 2013; Shenoy et al., 2013; Kaufman et al., 2014; Sadtler et al., 2014; Gallego et al., 2017; Elsayed & Cunningham, 2017;
Revised: Mante et al., 2013; Stokes et al., 2013; Shenoy et al., 2013; Kaufman et al., 2014; Sadtler et al., 2014; Gallego et al., 2017; Elsayed & Cunningham, 2017;

Original: Gao et al., 2017).
Revised: Gao et al., 2017).

Original: The bread and butter of classic systems neuroscience is linking neural activity to experimentally controlled or observable covariates such as orientation (Hubel & Wiesel, 1979), pitch (Lewicki, 2002), movement (Churchland et al., 2012; Kao et al., 2015), posture (Mimica et al., 2018) and orientation in space (Taube et al., 1990).
Revised: The bread and butter of classic systems neuroscience is linking neural activity to experimentally controlled or observable covariates such as orientation (Hubel & Wiesel, 1979), pitch (Lewicki, 2002), movement (Churchland et al., 2012; Kao et al., 2015), posture (Mimica et al., 2018) and orientation in space (Taube et al., 1990).

Original: These two parallel streams of neuroscientific research might at first seem to be at odds with each other (Kriegeskorte & Wei, 2021); tuning studies of individual neurons give a very different picture of neural coding than distributed representations over high-dimensional neural populations.
Revised: These two parallel streams of neuroscientific research might at first seem to be at odds with each other (Kriegeskorte & Wei, 2021); tuning studies of individual neurons give a very different picture of neural coding than distributed representations over high-dimensional neural populations.

Original: However, they combine elegantly in the form of (neural) latent variable models (LVMs, see Lawrence, 2003; Yu et al., 2008; Pandarinath et al., 2018).
Revised: However, they combine elegantly in the form of (neural) latent variable models (LVMs, see Lawrence, 2003; Yu et al., 2008; Pandarinath et al., 2018).

Original: 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 et al., 2013; Gao et al., 2017).
Revised: 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 et al., 2013; Gao et al., 2017).

Original: One advantage is that these models can help us discover latent variables which may not be tracked as classical covariates in systems neuroscience.
Revised: One advantage is that these models can help us discover latent variables which may not be tracked as classical covariates in systems neuroscience.

Original: This is especially appealing when the LVM is set up as a nonlinear Independent Component Analysis (ICA) model (Hyv¨arinen & Pajunen, 1999; Hyvarinen & Morioka, 2016; 2017; Khemakhem et al., 2020a;b;

Original: Klindt et al., 2021) with (provable) recovery of the latent factors that generated the data (e.g., arm position, movement direction, etc.

Original: Zhou & Wei, 2020; Schneider et al., 2022).

Original: However, when the mapping from latent variables to predicted spike rate (decoding) is fully unconstrained, e.g., by using a multi-layer neural network, we lose the simple biological interpretation of tuning curves.
Revised: However, when the mapping from latent variables to predicted spike rate (decoding) is fully unconstrained, e.g., by using a multi-layer neural network, we lose the simple biological interpretation of tuning curves.

Original: 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 curves.
Revised: 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 curves.

Original: These tuning curves have been parameterized as Gaussian processes in the framework of
Revised: These tuning curves have been parameterized as Gaussian processes in the framework of

Original: Gaussian process latent variable models (GPLVM, Wu et al., 2017; 2018).
Revised: Gaussian process latent variable models (GPLVM, Wu et al., 2017; 2018).

Original: Through the tuning curve approach,we limit ourselves to biologicallyplausible solutions that reveal the actual algorithmic structure of the neural system.
Revised: Through the tuning curve approach, biologically structure of the neural system.

Original: Thus, LVMs with simple tuning curve decoders bring together the view on neural populations as distributed representations of low-dimensional latent variables, along with the biologicallymeaningful perspective on individual neural tuning properties.
Revised: Thus, LVMs with simple tuning curve decoders bring together the view on neural populations as distributed representations of low-dimensional latent variables, along with the biologically on individual neural tuning properties.

Original: Some neural populations exhibit topologically interesting latent manifolds (Singh et al., 2008;
Revised: Some neural populations exhibit topologically interesting latent manifolds (Singh et al., 2008;

Original: Peyrache et al., 2015a; Gardner et al., 2022).
Revised: Peyrache et al., 2015a; Gardner et al., 2022).

Original: For instance, grid cells represent navigational space in toroidal coordinates of spatially repeating two dimensional hexagonal grids (Hafting et al., 2005).
Revised: For instance, grid cells represent navigational space in toroidal coordinates of spatially repeating two dimensional hexagonal grids (Hafting et al., 2005).

Original: They appear in different ensembles, commonly referred to as modules, each coding for space at different resolutions (Fyhn et al., 2007; Stensola et al., 2012).
Revised: They appear in different ensembles, commonly referred to as modules, each coding for space at different resolutions (Fyhn et al., 2007; Stensola et al., 2012).

Original: Thus, the complete population of grid cells is best described as a collection of ensembles of neurons, where neurons in each ensemble have tuning curves of specific shapes on their respective toroidal latent representations of space (Curto, 2017).
Revised: Thus, the complete population of grid cells is best described as a collection of ensembles of neurons, where neurons in each ensemble have tuning curves of specific shapes on their respective toroidal latent representations of space (Curto, 2017).

Original: By contrast, a two-dimensional Euclidean representation might also account for the complete population of grid cells, but would completely obscure their efficient and theoretically interesting coding scheme for representing space (Solstad et al., 2006; Sreenivasan & Fiete, 2011; Mathis et al., 2012; Wei et al., 2015; Klukas et al., 2020).
Revised: By contrast, a two-dimensional Euclidean representation might also account for the complete population of grid cells, but would completely obscure their efficient and theoretically interesting coding scheme for representing space (Solstad et al., 2006; Sreenivasan & Fiete, 2011; Mathis et al., 2012; Wei et al., 2015; Klukas et al., 2020).

Original: A driving motivation behind this work was to model this beautiful neural structure with an LVM that separates the algorithmic and biological parts, while uniting shared tuning properties to be more accurate and trainable than previous approaches.
Revised: A driving motivation behind this work was to model this beautiful neural structure with an LVM that separates the algorithmic and biological parts, while uniting shared tuning properties to be more accurate and trainable than previous approaches.

Original: We propose to train neural LVMs that not only have simple tuning curve decoders, but are also fully differentiable.
Revised: We propose to train neural LVMs that not only have simple tuning curve decoders, but are also fully differentiable.

Original: Thus, we use a flexible encoder, i.e., a neural network as in variational autoencoders (Kingma & Welling, 2014; Rezende et al., 2014), and a simple tuning curve based decoder, akin to GPLVMs.
Revised: Thus, we use a flexible encoder, i.e., a neural network as in variational autoencoders (Kingma & Welling, 2014; Rezende et al., 2014), and a simple tuning curve based decoder, akin to GPLVMs.

Original: The encoder can readily be made convolutional to allow better latent estimation from adjacent time points.
Revised: The encoder can readily be made convolutional to allow better latent estimation from adjacent time points.

Original: Additionally, we implement a feature basis for the tuning curve shapes in the decoder which can be shared across neurons.
Revised: Additionally, we implement a feature basis for the tuning curve shapes in the decoder which can be shared across neurons.

Original: We demonstrate that the neural feature sharing, along with the variational end-to-end training, vastly improves both the training stability as well as the final performance of neural LVMs.
Revised: We demonstrate that the neural feature sharing, along with the variational end-to-end training, vastly improves both the training stability as well as the final performance of neural LVMs.

Original: Moreover, we propose hybrid inference at test-time and show that this, again, brings a considerable improvement in performance.
Revised: Moreover, we propose hybrid inference at test-time and show that this, again, brings a considerable improvement in performance.

Original: 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 topological tuning properties.
Revised: 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 topological tuning properties.

Original: An illustration of our approach is provided in Fig.
Revised: An illustration of our approach is provided in Fig.

Original: To summarize, our full model performs the task of finding latent variables, separating distinct ensembles of neurons and fitting the prototypical tuning curves on each ensemble’s latent space into a single efficient framework.
Revised: To summarize, our full model performs the task of finding latent variables, separating distinct ensembles of neurons and fitting the prototypical tuning curves on each ensemble’s latent space in a single, efficient framework.

Original: In the following, we therefore refer to our model as f eature sharing a nd e nsemble detection L atent V ariable M odel or faeLVM .
Revised: In the following, we therefore refer to our model as f eature sharing a nd e nsemble detection L atent V ariable M odel or faeLVM .

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: Let x i be the instantaneous firing rate of a neuron i .
Revised: Let λ i be the instantaneous firing rate of a neuron i .

Original: To relate this to the spiking activity y i , we assume a Poisson noise model y i ∼ P ( x i ) .
Revised: To relate this to the spiking activity x i , we assume a Poisson noise model x i ∼ P ( λ i ) .

Original: We define latent variables z := { z 1 , . . .
Revised: We define latent variables z := { z 1 , . . .

Original: , z n } (in distinct spaces and possibly with different topologies) that change over time and to which neurons are tuned.
Revised: , z k } (in distinct spaces and possibly with different topologies) that change over time and to which neurons are tuned.

Original: More precisely, we suppose that there exists a deterministic function f i , which for every neuron i relates the latent variables to the firing rate, i.e., x i =
Revised: More precisely, we suppose that there exists a deterministic function f i , which for every neuron i relates the latent variables to the firing rate, i.e., λ i =

Original: f i ( z ) .
Revised: f i ( z ) .

Original: One additional assumption is that each f i is only sensitive to a subset J i of latent variables and invariant to all others.
Revised: One additional assumption is that each f i is only sensitive to a subset J i of latent variables and invariant to all others.

Original: That means, for two values z, ˆ z with z j = ˆ
Revised: That means, for two values z, ˆ z with z j = ˆ

Original: z j ∀ j ∈ J i and z j ̸
Revised: z j ∀ j ∈ J i and z j ̸

Original: = ˆ
Revised: = ˆ

Original: z j ∀ j / ∈
Revised: z j ∀ j / ∈

Original: J
Revised: J

Original: i we have f i ( z ) =
Revised: i we have f i ( z ) =

Original: f i (ˆ z ) .
Revised: f i (ˆ z ) .

Original: As a concrete example, consider a mixed population of head direction neurons I 1 and a grid cell module I 2 .
Revised: As a concrete example, consider a mixed population of head direction neurons I 1 and a grid cell module I 2 .

Original: The former are tuned to a circular latent variable z 1 ∈ S 1 (Peyrache et al., 2015a;
Revised: The former are tuned to a circular latent variable z 1 ∈ S 1 (Peyrache et al., 2015a;

Original: Rybakken et al., 2019;
Revised: Rybakken et al., 2019;

Original: Rubin et al., 2019; Chaudhuri et al., 2019), the latter to a toroidal latent variable z 2 ∈ T 2 (McNaughton et al., 2006; Fuhs & Touretzky, 2006; Gardner et al., 2022).
Revised: Rubin et al., 2019; Chaudhuri et al., 2019), the latter to a toroidal latent variable z 2 ∈ T 2 (McNaughton et al., 2006; Fuhs & Touretzky, 2006; Gardner et al., 2022).

Original: That means, for every neuron in the head direction ensemble we have J i
Revised: That means, for every neuron in the head direction ensemble we have J i

Original: = { 1 } , ∀ i ∈ I 1 ; whereas for every neuron in the grid cell ensemble we have J i
Revised: = { 1 } , ∀ i ∈ I 1 ; whereas for every neuron in the grid cell ensemble we have J i

Original: = { 2 } , ∀ i ∈ I 2 .
Revised: = { 2 } , ∀ i ∈ I 2 .

Original: Conceivably, there exist shared latent variables z k , such as attentional state, running speed or pupil dilation, to which multiple or even all neurons are tuned, i.e., k ∈ (cid:84) i J i .
Revised: Conceivably, there exist shared latent variables, e.g. attentional state, running speed or pupil dilation, to which multiple or even all neurons are tuned.

Original: Within each ensemble, the tuning functions f i , i ∈ I
Revised: Within each ensemble, the tuning functions f i , i ∈ I

Original: j are likely to have similar structure (Albright, 1984; Taube et al., 1990;
Revised: j are likely to have similar structure (Albright, 1984; Taube et al., 1990;

Original: Hafting et al., 2005).
Revised: Hafting et al., 2005).

Original: This is the concept of feature sharing (Klindt et al., 2017), which in visual systems neuroscience has proven to be a useful assumption (Batty et al., 2017;
Revised: This is the concept of feature sharing (Klindt et al., 2017), which in visual systems neuroscience has proven to be a useful assumption (Batty et al., 2017;

Original: Sinz et al., 2018; Ecker et al., 2018; Walker et al., 2018; Cadena et al., 2019; Ustyuzhaninov et al., 2019; Cotton et al., 2020; Zhuang et al., 2021; Burg et al., 2021; Bashiri et al., 2021; Safarani et al., 2021; Franke et al., 2021; Lurz et al., 2021;
Revised: Sinz et al., 2018; Ecker et al., 2018; Walker et al., 2018; Cadena et al., 2019; Ustyuzhaninov et al., 2019; Cotton et al., 2020; Zhuang et al., 2021; Burg et al., 2021; Bashiri et al., 2021; Safarani et al., 2021; Franke et al., 2021; Lurz et al., 2021;

Original: Nayebi et al., 2021; Seeliger et al., 2021; Goldin et al., 2021; Ustyuzhaninov et al., 2022).
Revised: Nayebi et al., 2021; Seeliger et al., 2021; Goldin et al., 2021; Ustyuzhaninov et al., 2022).

Original: 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 ¨ assle & Riemann, 1978).
Revised: 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 ¨ assle & Riemann, 1978).

Original: Precisely, let f µ denote the response function of a visual neuron with receptive field center µ ∈ R 2 , let T : R 2 → R 2 be a spatial translation, and let z : R 2 → R 1 be an image (with slight abuse of notation, z is, here, a function that assigns a grey scale value to every position in two-dimensional space, i.e., an image).
Revised: Precisely, let f µ denote the response function of a visual neuron with receptive field center µ ∈ R 2 , let T : R 2 → R 2 be a spatial translation, and let z : R 2 → R 1 be an image (with slight abuse of notation, z is, here, a function that assigns a grey scale value to every position in two-dimensional space, i.e., an image).

Original: Let us think of f i = f µ and f j = f T ( µ ) as two neurons of the same type (i.e., same tuning curve) but with their receptive field centers at different locations µ and T ( µ ) .
Revised: Let us think of f i = f µ and f i ′ = f T ( µ ) as two neurons of the same type (i.e., same tuning curve shape) but with their receptive field centers at different locations µ and T ( µ ) .

Original: Then we have f i ( z ) =
Revised: Then we have f i ( z ) =

Original: f j ( z ◦ T ) .
Revised: f i ′ ( z ◦ T ) .

Original: Thus, for visual neurons we assume translational equivariance .
Revised: Thus, for visual neurons we assume translational equivariance .

Original: Consequently, we should be able to marginalize over spatial translations to learn the shared structure among all f i as demonstrated in Klindt et al.
Revised: Consequently, we should be able to marginalize over spatial translations to learn the shared structure among all f i as demonstrated in Klindt et al.

Original: The same intuition holds for populations of auditory neurons that are equivariant to pitch, i.e., that have similar tuning properties translated across log-frequency space (Kell et al., 2018; Kell & McDermott, 2019) and possibly other sensory neurons such as somatosensory populations (Lieber & Bensmaia, 2019).
Revised: The same intuition holds for populations of auditory neurons that are equivariant to pitch, i.e., that have similar tuning properties translated across log-frequency space (Kell et al., 2018; Kell & McDermott, 2019) and possibly other sensory neurons such as somatosensory populations (Lieber & Bensmaia, 2019).

Original: 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;
Revised: 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;

Original: Fyhn et al., 2007; Gardner et al., 2022).
Revised: Fyhn et al., 2007; Gardner et al., 2022).

Original: Thus, we propose to introduce feature sharing also in this case.
Revised: Thus, we propose to introduce feature sharing also in this case.

Original: More precisely, we argue that every f i within an ensemble I j , should be modeled by a shared tuning curve function g j ( z, θ i ) =
Revised: More precisely, we argue that every f i within an ensemble I j , should be modeled by a shared tuning curve function g j ( z, θ i ) =

Original: f i ( z ) that applies a simple neuron specific transformation, parameterized by θ i , to obtain the specific f i of each neuron.
Revised: f i ( z ) that applies a simple neuron specific transformation, parameterized by θ i , to obtain the specific f i of each neuron.

Original: For instance, in the case of a neuron-specific spatial translation ( µ i , the receptive field center as above) and scaling ( α i ), we would define θ i = { µ i , α i } and thus g j ( z, θ i ) =
Revised: For instance, in the case of a neuron-specific spatial translation ( µ i , the receptive field center as above) and scaling ( α i ), we would define θ i = { µ i , α i } and thus g j ( z, θ i ) =

Original: α
Revised: α

Original: i g j ( T µ i ( z ) , θ 0 )
Revised: i g j ( T µ i ( z ) , θ 0 )

Original: =
Revised: =

Original: f i ( z ) ( θ 0 yielding simply the centered prototype of the shared tuning curve).
Revised: f i ( z ) ( θ 0 yielding simply the centered prototype of the shared tuning curve).

Original: Evidence for this assumption is provided in Fig.
Revised: Evidence for this assumption is provided in Fig.

Original: 2, where centering and rescaling all tuning curves of head direction selective neurons on S 1 (recorded by Peyrache et al., 2015b) clearly exhibits a shared, Gaussian-like tuning.
Revised: 2, where centering and rescaling all tuning curves of head direction selective neurons on S 1 (recorded by Peyrache et al., 2015b) clearly exhibits a shared, Gaussian-like tuning.

Original: In visual perception there is translational equivariance, however, only among cells of the same types (W¨assle & Riemann, 1978).
Revised: In visual perception there is translational equivariance, however, only among cells of the same types (W¨assle & Riemann, 1978).

Original: Therefore, previous work proposed feature sharing in combination with functional cell type identification (Klindt et al., 2017).
Revised: Therefore, previous work proposed feature sharing in combination with functional cell type identification (Klindt et al., 2017).

Original: Specifically, equating neural ensembles with cell types in the retina, this posits that all neurons in a given ensemble I j be tuned to the same set of latent variables J j , i.e., J i =
Revised: Specifically, equating neural ensembles with cell types in the retina, this posits that all neurons in a given ensemble I j be tuned to the same set of latent variables J j , i.e., J i =

Original: J j , ∀ i ∈ I j .
Revised: J j , ∀ i ∈ I j .

Original: As an example, the toroidal structures z j ∈ T 2 of different grid cell ensembles I j only become apparent after the successful separation of distinct modules with different spatial resolutions (Gardner et al., 2022).
Revised: As an example, the toroidal structures z j ∈ T 2 of different grid cell ensembles I j only become apparent after the successful separation of distinct modules with different spatial resolutions (Gardner et al., 2022).

Original: Thus, we argue that feature sharing hinges on the successful identification and separation of distinct neural ensembles in a mixed population recording.
Revised: Thus, we argue that feature sharing hinges on the successful identification and separation of distinct neural ensembles in a mixed population recording.

Original: The problem of ensemble detection can be formalized as partitioning n recorded neurons I = (cid:83) kj I j into k non-empty subsets I i ∩ I j  = ∅ , ∀ i ̸ = j  .
Revised: The problem of ensemble detection can be formalized as partitioning n recorded neurons I = (cid:83) kj I j into k non-empty subsets I j ∩ I j ′ = ∅ , ∀ j ̸ = j ′ .

Original: Unfortunately, this is a combinatorial problem with k n possible states.
Revised: Unfortunately, this is a combinatorial problem with k n possible states.

Original: Even for a ‘simple’ problem such as clustering neurons into k = 2 groups, this already exceeds the number of atoms in the universe ( 10 82 ) at n ≥ 273 .
Revised: Even for a ‘simple’ problem such as clustering neurons into k = 2 groups, this already exceeds the number of atoms in the universe ( 10 82 ) at n ≥ 273 .

Original: With modern neuroscience routinely yielding thousands of recorded neurons per experiment (Jun et al., 2017), brute-force exhaustive search is clearly not a viable approach.
Revised: With modern neuroscience routinely yielding thousands of recorded neurons per experiment (Jun et al., 2017), brute-force exhaustive search is clearly not a viable approach.

Original: To mitigate this issue, we propose a soft-relaxation of the clustering problem.
Revised: To mitigate this issue, we propose a soft-relaxation of the clustering problem.

Original: Specifically, we define a convex combination of weights over responses derived from tuning curves on each latent variable, and train the neurons to maximize their likelihood given those weights in a differentiable fashion (explained in more detail in the next section).
Revised: Specifically, we define a convex combination of weights over responses derived from tuning curves on each latent variable, and train the neurons to maximize their likelihood given those weights in a differentiable fashion (explained in more detail in the next section).

Original:  M ETHODS
Revised:  M ETHODS

Original: As above, we denote neural activity x
Revised: As above, we denote neural activity x

Original: and the collections of latent variables as z .
Revised: and the collections of latent variables as z .

Original: We want to learn a latent variable model p ( x, z ) =
Revised: We want to learn a latent variable model p ( x, z ) =

Original: p ( x | z ) p ( z ) .
Revised: p ( x | z ) p ( z ) .

Original: We choose a 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.e., a variational (input dependent) posterior with non-zero support across the domain, q ( z | x ) > 0 , ∀ z, x .
Revised: We choose a 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.e., a variational (input dependent) posterior with non-zero support across the domain, q ( z | x ) > 0 , ∀ z, x .

Original: At test-time, we focus on a precise estimate for the latents z .
Revised: At test-time, we focus on a precise estimate for the latents z .

Original: These latents are usually inferred with the encoder that computes the approximate variational posterior q ( z | x ) .
Revised: These latents are usually inferred with the encoder that computes the approximate variational posterior q ( z | x ) .

Original: Here, we aim to further refine this estimate by performing inference based on the decoder p ( x | z ) .
Revised: Here, we aim to further refine this estimate by performing inference based on the decoder p ( x | z ) .

Original: 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)
Revised: 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)

Original: (see Appx.
Revised: (see Appx.

Original: I for further details).
Revised: J for further details).

Original: For Euclidean latent spaces z ∈ R n , we can simply rely on the standard VAE framework, using a
Revised: For Euclidean latent spaces z ∈ R n , we can simply rely on the standard VAE framework, using a

Original: Normal ( N ) posterior and prior for those dimensions.
Revised: Normal ( N ) posterior and prior for those dimensions.

Original: However, for spherical latent spaces S n , we build on the work by Davidson et al. (2018), assuming a von Mises (vMF) distribution.
Revised: For spherical latent spaces S n , we use a wrapped Normal distribution ( w N ), akin to previous work (Falorsi et al., 2019; Jensen et al., 2020), which satisfies the topology of the latent space.

Original: For the relevant case of a toroidal latent space, we can use T 2 = S 1 × S 1 , together with the fact that the variational posterior is, usually, assumed to factorize across dimensions.
Revised: For the relevant case of a toroidal latent space, we can use T 2 = S 1 × S 1 , together with the fact that the variational posterior is, usually, assumed to factorize across dimensions.

Original: Additionally, the factorization  helps keeping the equations the same when working with multiple latents (corresponding to multiple ensembles), i.e., z = { z 1 , ..., z k } .
Revised: This factorization also helps keep the equations the same when working with multiple latents (corresponding to multiple ensembles), i.e., z = { z 1 , ..., z k } .

Original: More precisely, we assume that all latents z j live either in Euclidean space R n j or on a torus T n j =
Revised: More precisely, we assume that each latent z j lives either in Euclidean space R n j or on a torus T n j =

Original: S 11 × . . .
Revised: S 11 × . . .

Original: ×
Revised: ×

Original: S 1 n j (with S 1 = T 1 ).
Revised: S 1 n j (with S 1 = T 1 ).

Original: Further, we let our prior factorize (Kingma & Welling, 2014), i.e., p( z ) = (cid:89) j p (
Revised: Further, we let our prior factorize (Kingma & Welling, 2014), i.e., p

Original: [Equation]
Revised: [Equation]

Original: , p
Revised: giving uniform prior distributions. Analogously, we let our variational posteriors factorize as

Original: [Equation]
Revised: [Equation]

Original: where we set κ = 0 , thus obtaining a uniform prior distribution on the hypersphere (with unspecified µ , see Davidson et al., 2018). Analogously, we let our variational posterior factorize, i.e., q ( z | x ) = (cid:89)
Revised: if

Original: [Equation]
Revised: [Equation]

Original: Here, the  parameters of the Normal and von Mises distribution (mean, variance, and concentration) are  input dependent, and trained with the reparametrization trick.
Revised: Here, the distribution-specific parameters µ ( x ) and σ ( x ) are themselves input dependent, and given by a functional mapping which is learned using the reparametrization trick.

Original: We use a temporal (1D) convolutional neural network to parameterize q .
Revised: We use a temporal (1D) convolutional neural network to parameterize q .

Original: To enforce circular latents, we let the encoder output vectors in R 2 , which we subsequently normalize to avoid discontinuities when estimating angles (Zhou et al., 2019).
Revised: To enforce circular latents, we let the encoder output vectors in R 2 , which we subsequently normalize to avoid discontinuities when estimating angles (Zhou et al., 2019).

Original: We remark that although a convolutional filter was used, temporal smoothness was not explicitly included in the latent dynamics of z in the model used for producing the results in Sec. 4.
Revised: We remark that although a convolutional filter was used, temporal smoothness was not explicitly included in the latent dynamics of z in the model used for producing the results in Sec. 4.

Original: It is, however, a straightforward inclusion to incorporate (see Appx.
Revised: It is, however, a straightforward inclusion to incorporate (see Appx.

Original: D.1).
Revised: D.1).

Original: For the decoder, in clear contrast to standard VAEs, we pick a simple parametric function f to parameterize the reconstruction term p ( x | z ) .
Revised: For the decoder, in clear contrast to standard VAEs, we pick a simple parametric function f to parameterize the reconstruction term p ( x | z ) .

Original: 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.
Revised: 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.

Original: Specifically, we have
Revised: Specifically, we have

Original: [Equation]
Revised: [Equation]

Original: where P ( f i ( z )) is a Poisson distribution with rate f i ( z ) , and the weights for each neuron i the output of a softmax function, such that w ij > 0 and (cid:80) j w ij = 1 .
Revised: where P ( f i ( z )) is a Poisson distribution with rate f i ( z ) , and the weights for each neuron i the output of a softmax function, such that w ij > 0 and (cid:80) j w ij = 1 .

Original: That is, w ij models the participation of neuron i in ensemble j .
Revised: That is, w ij models the participation of neuron i in ensemble j .

Original: The function g j
Revised: The function g j

Original: ( z j , θ i ) ∈
Revised: ( z j , θ i ) ∈

Original: R characterizes the (shared) tuning curves of neurons in the latent space of z j , and takes as input the corresponding latent z j and a (set of) neuron specific parameter(s) θ i , to produce the response of the neuron based on each latent space.
Revised: R characterizes the (shared) tuning curves of neurons in the latent space of z j , and takes as input the corresponding latent z j and a (set of) neuron specific parameter(s) θ i , to produce the response of the neuron based on each latent space.

Original: The function g j ( z j , θ i ) can be specified in a multitude of ways, while still embodying the concept of feature sharing.
Revised: The function g j ( z j , θ i ) can be specified in a multitude of ways, while still embodying the concept of feature sharing.

Original: 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.
Revised: 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.

Original: Since neural activity is non-negative, we model the log of the instantaneous firing rates as
Revised: Since neural activity is non-negative, we model the log of the instantaneous firing rates as

Original: [Equation]
Revised: [Equation]

Original: where µ i denotes (as above) the center of the i -th neuron’s tuning curve and d j ( · ) a distance function in the space of z j (i.e., Euclidean distance in R n and geodesic distance in T n ).
Revised: where µ i denotes (as above) the center of the i -th neuron’s tuning curve and d j ( · ) a distance function in the space of z j (i.e., Euclidean distance in R n and geodesic distance in T n ).

Original: The second, and more flexible, feature space consists of a sum of K weighted ( β k ) Gaussian basis functions such that
Revised: The second, and more flexible, feature space consists of a sum of M weighted ( β m ) Gaussian basis functions (i.e., a spline)

Original: [Equation]
Revised: [Equation]

Original: withµ k , σ k the center  and width  of the k -th basis function.
Revised: with center µ m and width σ m of the m -th basis function.

Original: We consider three variations ofthe faeLVM: sharing of features through Eq.
Revised: We consider three variations of faeLVM: sharing of features by utilizing Eq.

Original: 5 (faeLVM-s), utilizing Eq. 4 for modeling g j as a shared  heat kernel (faeLVM-b),  as well as a model allowing each neuron to learn its own set of the basis weights in Eq.
Revised: 4 to model g j with a shared (bump) heat kernel (faeLVM-b), sharing of features through Eq. 5 using a Gaussian basis (faeLVM-s), as well as a model allowing each neuron to learn its own set of the basis weights in Eq.

Original: 5, i.e. no shared features (faeLVM-n).
Revised: 5, i.e. no shared features (faeLVM-n).

Original:  E XPERIMENTS 4. S IMULATIONS : F EATURE S HARING
Revised:  E XPERIMENTS 4. S IMULATIONS : F EATURE S HARING

Original: Motivated by the evaluation pipeline in Pei et al.
Revised: Motivated by the evaluation pipeline in Pei et al.

Original: (2021), we split our data into four parts (Fig. 3), exploring the benefits of feature sharing in the following way: We generate a 1 D periodic latent variable and simulate spikes according to a Poisson distribution with heat kernel tuning curves (see also Appx.
Revised: (2021), we split our data into four parts (Fig. 3), exploring the benefits of feature sharing in the following way: We generate z ∈ S 1 , a circular 1D latent variable and simulate spikes according to a Poisson distribution with heat kernel tuning curves (see also Appx.

Original: F for a retinal ganglion simulation with calcium dynamics).
Revised: F for a retinal ganglion simulation with calcium dynamics).

Original: We fix T test and N test ( 1000 and 30 ), and investigate model performance based on fixed T train and varying N train , and vice versa.
Revised: We fix T test and N test ( 1000 and 30 ), and investigate model performance based on fixed T train and varying N train , and vice versa.

Original: 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 are better, in both cases).
Revised: 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 are better, in both cases).

Original: We also compare our models with mGPLVM (Jensen et al., 2020), an extension of
Revised: We also compare our models with mGPLVM (Jensen et al., 2020), an extension of

Original: GPLVM (Wu et al., 2018) designed for inference on non-Euclidean spaces, as well as a conventional
Revised: GPLVM (Wu et al., 2018) designed for inference on non-Euclidean spaces, as well as a conventional

Original: VAE using a standard multi-layer perceptron decoder (MLP; with essentially a reverse architecture of the encoder).
Revised: VAE using a standard multi-layer perceptron decoder (MLP; with essentially a reverse architecture of the encoder).

Original: As mGPLVM assumes a Gaussian noise model, the Poisson NLLH is unsuitable as a measure, hence the comparison models are only included when evaluating latent recoveries.
Revised: As mGPLVM assumes a Gaussian noise model, the Poisson NLLH is unsuitable as a measure, hence it is only included when evaluating latent recoveries.

Original: Fig.
Revised: Fig.

Original: 4A-B clearly show the regime where feature sharing excels: datasets with a satisfactory number of neurons, but with short recording times.
Revised: 4A-B clearly show the regime where feature sharing excels: datasets with a satisfactory number of neurons, but with short recording times.

Original: 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 (Stevenson & Kording, 2011).
Revised: 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 (Stevenson & Kording, 2011).

Original: 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. 
Revised: 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. A

Original:  We also remark that whilethe faeLVM-n  performs better on rate prediction in the low neuron regime (Fig. 4A) the performance of LVMs ultimately hinges on how well the models can recover the correct latent variable.
Revised: Mean NLLH averaged 20 seeds varying would also reflect the ity We also remark that while faeLVM-n seems to performs better on rate prediction in the low neuron regime (Fig. 4A), the performance of LVMs ultimately hinges on how well the models can recover the correct latent variable.

Original: In the setting of latent recovery (Fig.
Revised: In the setting of latent recovery (Fig.

Original: 4C-D), we see that feature sharing is evidently beneficial both in low neuron- and low recording-time regimes.
Revised: 4C-D), we see that feature sharing is evidently beneficial both in low neuron- and low recording-time regimes.

Original: faeLVM-b outperforms the other models, perhaps somewhat expected given the shape of the synthetic tuning curves, while allthe faeLVMs outperform both mGPLVM and VAE.
Revised: 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.

Original: We also observe that while mGPLVM, which uses tuning curve decoders, is more interpretable, it is worse than neural network decoders.
Revised: We also observe that while mGPLVM, which uses tuning curve decoders, is more interpretable, it is worse than neural network decoders.

Original: However, 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.
Revised: However, 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.

Original: Note that although the faeLVMs leverage the benefit of a convolutional layer (which mGPLVM does not, being a 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 to infer the latents independently at each time point (see Appx. E).
Revised: Note that although the faeLVMs leverage the benefit of a convolutional layer (which mGPLVM does not, being a 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 to infer the latents independently at each time point (see Appx. E).

Original: 4. S IMULATIONS : E NSEMBLE D ETECTION
Revised: 4. S IMULATIONS : E NSEMBLE D ETECTION

Original: Next, we shift our focus to the challenge of ensemble detection, comparing the performance of faeLVM-b against other common methods.
Revised: Next we shift our focus to the challenge of ensemble detection comparing the performance of faeLVM-b against other common methods.

Original: Specifically, we contrast against clustering methods (see
Revised: Specifically, we contrast against clustering methods (see

Original: Appx.
Revised: Appx.

Original: C for more details) that either i) perform 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., 2021), or ii) perform clustering on the neural covariance matrix (loosely inspired by Carrillo-Reid et al., 2015).
Revised: C for more details) that either i) perform 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., 2021) or ii) perform clustering on the neural covariance matrix (loosely inspired by Carrillo-Reid et al., 2015).

Original: As an additional upper bound to the achievable performance, we include a supervised clustering method that receives both the mutual information score between each neuron’s activity, and the true latent variables, as input.
Revised: As an additional upper bound to the achievable performance we include a supervised clustering method that receives both the mutual information score between each neuron’s activity, and the true latent variables, as input.

Original: We run our model over five different random seeds and pick the run that yields the highest LLH (faeLVM-b5 in Fig. 5), a fair approach as the likelihood is an unsupervised metric (k-means is run over 100 seeds, picking the one with the lowest inertia).
Revised: We run our model over five different random seeds and pick the run that yields the highest LLH (faeLVM-b5 in Fig. 5), a fair approach as the likelihood is an unsupervised metric (k-means is run over 100 seeds, picking the one with the lowest inertia).

Original: The setting we chose for these experiments was the separation of ensembles of neurons that live on a ring.
Revised: The setting we chose for these experiments was the separation of ensembles of neurons that live on a ring, i.e. S 1 .

Original: The comparison methods failed to separate the ensembles when scaling up to neural activity on a torus, in contrast to faeLVM, which further highlights the novelty of an unsupervised method that can accomplish this task (see Sec. 4.4).
Revised: The comparison methods failed to separate the ensembles when scaling up to neural activity on a torus, in contrast to faeLVM, which further highlights the novelty of an unsupervised method that can accomplish this task (see Sec. 4.4 for tori, and also Appx.

Revised: I for an additional ensemble separation ablation).

Original: We tested the accuracy on separating two and three  neural ensembles of equal size on rings, in addition to investigatingensemble detection accuracy for all models as a function of available data points (recording length).
Revised: 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 a function of available data points (recording length).

Original: Fig.
Revised: Fig.

Original: 5A-B shows that our model outperforms both comparison methods and quickly approaches the performance of the supervised upper bound, for both two and three rings.
Revised: 5A-B shows that our model outperforms both comparison methods and quickly approaches the performance of the supervised upper bound, for both two and three rings.

Original: This is even more pronounced for the model selected based on LLH — a viable option when the analysis result warrants additional computing time.
Revised: This is even more pronounced for the model selected based on LLH — a viable option when the analysis result warrants additional computing time.

Original: Fig.
Revised: Fig.

Original: 5C shows that the relationship between ensemble detection accuracy and LLH is monotonically increasing.
Revised: 5C shows that the relationship between ensemble detection accuracy and LLH is monotonically increasing.

Original: Thus, the unsupervised likelihood score is indeed a good proxy to arbitrate between different optimization runs.
Revised: Thus, the unsupervised likelihood score is indeed a good proxy to arbitrate between different optimization runs.

Original: [Figure]
Revised: [Figure]

Original: Transitioning to real datasets, we first apply our models to data from Peyrache et al.
Revised: Transitioning to real datasets, we first apply our models to data from Peyrache et al.

Original: (2015b); more specifically the awake trials from the dataset labeled Mouse 28 , session 140313 , recorded from the anterodorsal thalamic nucleus (ADN).
Revised: (2015b); more specifically the awake trials from the dataset labeled Mouse 28 , session 140313 , recorded from the anterodorsal thalamic nucleus (ADN).

Original: Recordings yield approx.
Revised: Recordings yield approx.

Original: 35 min. of data, which we bin in 100 ms bins, before separating it with a train-test split of 95% − 5% .
Revised: 35 min. of data, which we bin in 100 ms bins, before separating it with a train-test split of 95% − 5% .

Original: 3 neurons are allocated as test neurons (out of 26 total) and we run 20 seeds of each model, selecting the best performing model for visualization (faeLVM-n).
Revised: 3 neurons are allocated as test neurons (out of 26 total) and we run 20 seeds of each model, selecting the best performing model for visualization (faeLVM-n).

Original: The inferred latent (Fig.
Revised: The inferred latent (Fig.

Original: 6A) is observed to match the true head direction to a high degree, and the model is able to predict rates on test neurons (Fig.
Revised: 6A) is observed to match the true head direction to a high degree, and the model is able to predict rates on test neurons (Fig.

Original: 6B), even though the average and peak activity over test bins are quite low ( ∼ 1 and 5 Hz respectively) for the first two.
Revised: 6B), even though the average and peak activity over test bins are quite low ( ∼ 1 and 5 Hz respectively) for the first two.

Original: As we can see from Fig.
Revised: As we can see from Fig.

Original: 6C, the unconstrained model (faeLVM-n) outperforms the other two.
Revised: 6C, the unconstrained model (faeLVM-n) outperforms the other two.

Original: This is not an unsurprising result, given the models higher complexity and the sufficiently large amount of training data.
Revised: This is not an unsurprising result, given the models higher complexity and the sufficiently large amount of training data.

Original: As for  mGPLVM and the traditional VAE (Fig.
Revised: As for the latent prediction, mGPLVM and the traditional VAE (Fig.

Original: 6D), they are close in comparison to the performance of faeLVM-s, although less accurate than faeLVM-n.  Training multiple models and making a selection based on train LLH generally correlates well with solid results on the test set, as also shown in Fig5C.
Revised: 6D) are close in comparison to the performance of faeLVM-s, although less accurate than faeLVM-n. 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 a 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 on train LLH generally correlates well with solid results on the test set, as also shown in Fig

Original: Also note that faeLVM-b achieves much more consistent optimization results over different random initializations, suggesting a typical bias-variance tradeoff.
Revised: 5C. Note that faeLVM-b achieves much more consistent optimization results over different random initializations, suggesting a typical bias-variance tradeoff.

Original: Overall, we recognize that while there are differences in performances all models recover the correct circular manifold.
Revised: Overall, we recognize that while there are differences in performances, all models recover the correct circular manifold.

Original: The faeLVMs however, are noticeably faster ( ∼ 5 min.
Revised: The faeLVMs however, are noticeably faster ( ∼ 5 min.

Original: run time for one seed, on a CPU), compared to mGPLVM ( ∼ 50 min.
Revised: run time for one seed, on a CPU), compared to mGPLVM ( ∼ 50 min.

Original: run time for one seed on a GPU).
Revised: run time for one seed on a GPU).

Original: While not critical in the case of this particular dataset, it is paramount when inferring latents of higher dimensionality.
Revised: While not critical in the case of this particular dataset, it is paramount when inferring latents of higher dimensionality.

Original: We also note that while the circular manifold was in this case known, it is possible to do model selection over different latent spaces (toroidal, circular, planar, etc.)
Revised: We also note that while the circular manifold was in this case known, it is possible to do model selection over different latent spaces (toroidal, circular, planar, etc.)

Original: as was demonstrated in Jensen et al.
Revised: as was demonstrated in Jensen et al.

Original: Thus, it is not strictly necessary to know the latent topology when applying our method, rather it allows for the restriction of possible models to a limited hypothesis class.
Revised: Thus, it is not strictly necessary to know the latent topology when applying our method, rather it allows for the restriction of possible models to a limited hypothesis class.

Original: 4. R EAL D ATASETS : G RID C ELL D ATA
Revised: 4. R EAL D ATASETS : G RID C ELL D ATA

Original: Finally, we apply the complete model pipeline to data recorded from the medial entorhinal cortex (MEC) of a freely moving rat (Gardner et al., 2022).
Revised: Finally, we apply the complete model pipeline to data recorded from the medial entorhinal cortex (MEC) of a freely moving rat (Gardner et al., 2022).

Revised: Neurons in the MEC are selective to a number of features like space (Fyhn et al., 2004), head direction (Sargolini et al., 2006), speed (Kropff et al., 2015), as well as conjunctive representations of the three (Sargolini et al., 2006; Hardcastle et al., 2017), though most demonstrate a strong preference for single features (Kropff et al., 2015; Tukker et al., 2022).

Original: 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.
Revised: 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.

Original: Spike data was binned in 100 ms bins, after removing periods in time when the rat was considered stationary (speed < 2 .
Revised: Spike data was binned in 100 ms bins, after removing periods in time when the rat was considered stationary (speed < 2 .

Original: 5 cm s − 1 ).
Revised: 5 cm s − 1 ).

Original: Using recordings from both the open field and maze experiments yields approx.
Revised: Recordings from both the open field and maze experiments yields approx.

Original: 170 min. of data.
Revised: 170 min. of data.

Original: We train 20 models from different seeds, selecting the one withthe highest in-sample likelihood.
Revised: We train 20 models from different seeds, selecting the one with highest in-sample likelihood.

Original: Compared with the modules discovered in the dataset (Gardner et al., 2022), our model’s results coincides with an accuracy of 0 . 83 .
Revised: Compared with the modules discovered in the dataset (Gardner et al., 2022), our model’s results coincides with an accuracy of 0 . 83 .

Original: The inferred weights (Fig.
Revised: The inferred weights (Fig.

Original: 7A) show a clear separation between the two modules, becoming less distinct as the neurons’ spatial information decreases.
Revised: 7A) show a clear separation between the two modules, becoming less distinct as the neurons’ spatial information decreases.

Original: This is also reflected in Fig 7B, where neurons further away from the diagonal tend to be more accurately classi- fied.
Revised: This is also reflected in Fig 7B, where neurons further away from the diagonal tend to be more accurately classified.

Original: However, we note that the original separation of neurons into ensembles is not an intrinsic truth; a number of the less informative neurons exhibit structurally ambiguous rate maps (see Appx.
Revised: However, we note that the original separation of neurons into ensembles is not an intrinsic truth; a number of the less informative neurons exhibit structurally ambiguous rate maps (see Appx.

Original: G), a possible indication that they may belong to a different ensemble altogether.
Revised: G), a possible indication that they may belong to a different ensemble altogether.

Original: Excluding these neurons while evaluating the model (Fig.
Revised: Excluding these neurons while evaluating the model (Fig.

Original: 7C), we observe a significant increase in performance, reaching 0 .
Revised: 7C), we observe a significant increase in performance, reaching 0 .

Original: 90 accuracy after removing approx.
Revised: 90 accuracy after removing approx.

Original: 50 of the least informative neurons from each ensemble.
Revised: 50 of the least informative neurons from each ensemble.

Original: Fig.
Revised: Fig.

Original: 7D-E show a selection of rate maps for neurons from both ensembles.
Revised: 7D-E show a selection of rate maps for neurons from both ensembles.

Original: 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.
Revised: 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.

Original: The inferred centers and widths also correspond well with the rate maps.
Revised: The inferred centers and widths also correspond well with the rate maps.

Revised: 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 to determine the number of ensembles, the corresponding shapes and allow for some degree of conjunctivity.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: 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 , 1 in the sense that similar observations are grouped while differences are highlighted.
Revised: 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 , 1 in the sense that similar observations are grouped while differences are highlighted.

Original: This principle, instantiated in a group equivariant computation core and a cell-type specific readout, has been successfully deployed in visual neuroscience; 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.
Revised: This principle, instantiated in a group equivariant computation core and a cell-type specific readout, has been successfully deployed in visual neuroscience; 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.

Original: It 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.
Revised: It 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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 rotations SO (3) (vision), two spheres S 2 (Singh et al., 2008), or the Klein bottle shape of oriented and phase shifted wavelets (Carlsson et al., 2008) (relevant to multiple sensory modalities).
Revised: 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 rotations SO (3) (vision), 2-spheres S 2 (Singh et al., 2008), or the Klein bottle shape of oriented and phase shifted wavelets (Carlsson et al., 2008) (relevant to multiple sensory modalities).

Original: 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.
Revised: 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.

Original: In 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).
Revised: In 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).

Original: On the other hand, if we can identify the parameters (e.g., tuning width) that describe the continuous variations across neurons and separate them from the stable features of tuning curves (e.g., Gaussian-like heat kernel shape), we can use that to build a more flexible, parametric (but still shared) feature basis.
Revised: On the other hand, if we can identify the parameters (e.g., tuning width) that describe the continuous variations across neurons and separate them from the stable features of tuning curves (e.g., Gaussian-like heat kernel shape), we can use that to build a more flexible, parametric (but still shared) feature basis.

Original: Examples of such approaches include rotation (Ustyuzhaninov et al., 2019), kinetic (Zhao et al., 2020) or more general nonlinear modifications (Shah et al., 2022) to shared feature spaces.
Revised: Examples of such approaches include rotation (Ustyuzhaninov et al., 2019), kinetic (Zhao et al., 2020) or more general nonlinear modifications (Shah et al., 2022) to shared feature spaces.

Original: Further challenges to the feature sharing idea could arise from recent findings of representational drift (Rule et al., 2019) and location- (Qiu et al., 2020) or context dependent tuning (Kanter et al., 2017), although future directions might consider addressing these in the model.
Revised: Further challenges to the feature sharing idea could arise from recent findings of representational drift (Rule et al., 2019) and location- (Qiu et al., 2020) or context dependent tuning (Kanter et al., 2017), although future directions might consider addressing these in the model.

Original: Other limitations are the fact that the grid cell torus is hexagonal, 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.
Revised: Other limitations are the fact that the grid cell torus is hexagonal, 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.

Original: Finally, although the choice of variational posteriors might at first seem restrictive, we would argue that the corresponding topologies (circles, tori, R n ) are natural choices for latent spaces and likely to keep appearing in the brain, as is suggested in recent work (Kriegeskorte & Wei, 2021).
Revised: Finally, although the choice of variational posteriors might at first seem restrictive, we would argue that the corresponding topologies (circles, tori, R n ) are natural choices for latent spaces and likely to keep appearing in the brain, as is suggested in recent work (Kriegeskorte & Wei, 2021).

Original: Other authors conjectured that, e.g., prefrontal cortex deploys grid cell-like codes to explore physical (Doeller et al., 2010; Jacobs et al., 2013) and conceptual spaces (Constantinescu et al., 2016;
Revised: Other authors conjectured that, e.g., prefrontal cortex deploys grid cell-like codes to explore physical (Doeller et al., 2010; Jacobs et al., 2013) and conceptual spaces (Constantinescu et al., 2016;

Original: Bao et al., 2019), as well as rules in a reinforcement learning task (Baram et al., 2021).
Revised: Bao et al., 2019), as well as rules in a reinforcement learning task (Baram et al., 2021).

Original: The challenge of constructing or observing the relevant space has proven a hurdle for previous supervised methods, being unable to discover cognitive maps without a priori knowledge of the represented covariates.
Revised: The challenge of constructing or observing the relevant space has proven a hurdle for previous supervised methods, being unable to discover cognitive maps without a priori knowledge of the represented covariates.

Original: 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.
Revised: 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.

Original:  E THICS S TATEMENT
Revised:  E THICS S TATEMENT

Original: More generally, reducing high dimensional nonlinear dynamical systems to a set of latent variables can have a broad array of possible applications, for instance, in climate sciences, crop forecasting or flood prediction.
Revised: More generally, reducing high dimensional nonlinear dynamical systems to a set of latent variables can have a broad array of possible applications, for instance, in climate sciences, crop forecasting or flood prediction.

Original: 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.
Revised: 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.

Original: As for any complex and versatile modern machine learning method, there also exists a danger of translating the insights from this work to detrimental purposes and intents.
Revised: As for any complex and versatile modern machine learning method, there also exists a danger of translating the insights from this work to detrimental purposes and intents.

Original: Applications to sensitive data with protected attributes, such as gender or ethnicity, should pay special attention to the en- and decoding of those attributes in the learned latent spaces.
Revised: Applications to sensitive data with protected attributes, such as gender or ethnicity, should pay special attention to the en- and decoding of those attributes in the learned latent spaces.

Original: 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 of the ideas developed in this work.
Revised: 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 of the ideas developed in this work.

Original: Lastly, we tried to minimize the environmental impact of our research by performing the hyper-parameter search on the GPU cluster of a carbon-neutral organization
Revised: Lastly, we tried to minimize the environmental impact of our research by performing the hyper-parameter search on the GPU cluster of a carbon-neutral organization

Original:  R EPRODUCIBILITY S TATEMENT
Revised:  R EPRODUCIBILITY S TATEMENT

Original: Code and implementation are included as supplementary material, and will be made available in the event of acceptance.
Revised: Code and implementation are included as supplementary material and will be made available in the event of acceptance.

Original: Datasets are public, and details regarding model parameters, data generation, etc. are specified under the relevant sections (e g., Sec. 4), as well as more thoroughly in the
Revised: Datasets are public, and details regarding model parameters, data generation, etc. are specified under the relevant sections (e g., Sec. 4), as well as more thoroughly in the

Original: Appendix (e.g., Appx.
Revised: Appendix (e.g., Appx.

Original: A, B, C, I & J).
Revised: A, B, C, J & K)

Original: A PPENDIX A M ODEL DETAILS
Revised: A PPENDIX A M ODEL DETAILS

Original: All models are trained with the Adam optimizer (Kingma & Ba, 2014), a learning rate of 0 .
Revised: All models are trained with the Adam optimizer (Kingma & Ba, 2014), a learning rate of 0 .

Original: 001 , temporal chunks of size 128 ( 64 if data length is smaller than 128 ) and a batch size of 1 .
Revised: 001 , temporal chunks of size 128 ( 64 if data length is smaller than 128 ) and a batch size of 1 .

Original: Moreover, training is concluded after the training objective has not improved for 5 steps ( 10 for grid cell data).
Revised: Moreover, training is concluded after the training objective has not improved for 5 steps ( 10 for grid cell data).

Original: All models are implemented in Pytorch (Paszke et al., 2019).
Revised: All models are implemented in Pytorch (Paszke et al., 2019).

Original: The architecture of the encoder is as follows: a 1D grouped convolution with kernel size 9 ; followed by two layers of 1D convolutions with kernel size 1 (i.e., effectively just MLPs at each time point) and 64 filters ( 128 for grid cell data), followed by a linear projection onto the dimensionality of the mean and variance/concentration parameters required by the variational posterior.
Revised: The architecture of the encoder is as follows: a 1D grouped convolution with kernel size 9 ; followed by two layers of 1D convolutions with kernel size 1 (i.e., effectively just MLPs at each time point) and 64 filters ( 128 for grid cell data), followed by a linear projection onto the dimensionality of the mean and variance/concentration parameters required by the variational posterior.

Original: The decoder follows the feature sharing tuning curve structure as presented in the paper.
Revised: The decoder follows the feature sharing tuning curve structure as presented in the paper.

Original: All experiments where executed with specified random seeds for data generation, model initialisation and training batch selection to maintain complete reproducibilty.
Revised: All experiments where executed with specified random seeds for data generation, model initialisation and training batch selection to maintain complete reproducibilty.

Original: A. A FAST APPROXIMATION TO THE OBJECTIVE FUNCTION
Revised: A. A FAST APPROXIMATION TO THE OBJECTIVE FUNCTION

Original: Next, we present a simplification of the theoretical setup in the paper Sec. 3.
Revised: Next, we present our simplification of the theoretical setup from (Davidson et al., 2018), concerning the use of von Mises distributions when inferring latent variables on circular manifolds.

Original: First, we note that as the variance σ 2 of a Gaussian and the inverse concentration 1 /κ of a von Mises approach 0 , they both become delta functions.
Revised: First, we note that as the variance σ 2 of a Gaussian and the inverse concentration 1 /κ of a von Mises approach 0 , they both become delta functions.

Original: Moreover, they each define the maximum entropy distribution for a fixed variance (concentration) in their respective domains.
Revised: Moreover, they each define the maximum entropy distribution for a fixed variance (concentration) in their respective domains.

Original: Hence, the von Mises is also referred to as the circular normal distribution.
Revised: Hence, the von Mises is also referred to as the circular normal distribution.

Original: We therefore explored the model behavior when replacing the KL term in Eq.
Revised: We therefore explored the model behavior when replacing the KL term in Eq.

Original: (11) with a simpler expression, akin to the standard VAE KL, which simply pushes the variance of the variational encoder towards 1 (i.e., standard normal), but without constraints on the mean (uniform over the latent space, as the full von Mises prior).
Revised: (11) with a simpler expression, akin to the standard VAE KL, which simply pushes the variance of the variational encoder towards 1 (i.e., standard normal), but without constraints on the mean (uniform over the latent space, as the full von Mises prior).

Original: Reparameterisation and sampling was then performed by drawing from that Gaussian and adding the scaled perturbation (same as in standard VAEs) to the angle that the variational posterior returned as mean estimate.
Revised: Reparameterisation and sampling was then performed by drawing from that Gaussian and adding the scaled perturbation (same as in standard VAEs) to the angle that the variational posterior returned as mean estimate.

Original: In general, performing variational inference in non-Euclidean 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; Jensen et al., 2020; Falorsi et al., 2019; Jensen et al., 2022).
Revised: In general, performing variational inference in non-Euclidean 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; Jensen et al., 2020; Falorsi et al., 2019; Jensen et al., 2022).

Original: In this work, we follow the ‘ReLie’ approach outlined by Falorsi et al.
Revised: In this work, we follow the ‘ReLie’ approach outlined by Falorsi et al.

Original: for general Lie groups, which involves defining a parameterizable (in our case Gaussian) distribution on the tangent space of the group and projecting it onto the manifold using the so-called ‘exponential map’ of the group.
Revised: for general Lie groups, which involves defining a parameterizable (in our case Gaussian) distribution on the tangent space of the group and projecting it onto the manifold using the so-called ‘exponential map’ of the group.

Original: This approach is particularly simple for the toroidal manifolds we consider, where the projection step can be achieved by a simple modulo 2 π operation (Jensen et al., 2020).
Revised: This approach is particularly simple for the toroidal manifolds we consider, where the projection step can be achieved by a simple modulo 2 π operation (Jensen et al., 2020).

Original: This leaves the challenge of computing the KL term in Eq.
Revised: This leaves the challenge of computing the KL term in Eq.

Original: 11, which is given by
Revised: 11, which is given by

Original: [Equation]
Revised: [Equation]

Original: As a prior p ( z ) , we use a unit Gaussian projected onto the circle, with a mean matched to the posterior mean.
Revised: As a prior p ( z ) , we use a unit Gaussian projected onto the circle, with a mean matched to the posterior mean.

Original: Since the KL divergence in Eq.
Revised: Since the KL divergence in Eq.

Original: 6 is invariant to a rotation around the circle (Falorsi et al., 2019), we can compute it simply as the KL divergence between two wrapped Gaussian distributions centered at the origin with variances σ 2 q and σ 2 p = 1 respectively.
Revised: 6 is invariant to a rotation around the circle (Falorsi et al., 2019), we can compute it simply as the KL divergence between two wrapped Gaussian distributions centered at the origin with variances σ 2 q and σ 2 p = 1 respectively.

Original: In general, the projected densities needed in Eq.
Revised: In general, the projected densities needed in Eq.

Original: 6 are given by
Revised: 6 are given by

Original: [Equation]
Revised: [Equation]

Original: as derived in Falorsi et al.
Revised: as derived in Falorsi et al.

Original: Here, exp G is the exponential map, J ( x ) is the Jacobian at x in the tangent space, and r θ is the parameterized reference distribution on the tangent space.
Revised: Here, exp G is the exponential map, J ( x ) is the Jacobian at x in the tangent space, and r θ is the parameterized reference distribution on the tangent space.

Original: For the circle, the Jacobian is simply given by the identity J ( x ) = 1 ∀ x .
Revised: For the circle, the Jacobian is simply given by the identity J ( x ) = 1 ∀ x .

Original: With a zero-centered
Revised: With a zero-centered

Original: Gaussian reference distribution, we can therefore rewrite Eq.
Revised: Gaussian reference distribution, we can therefore rewrite Eq.

Original: 7 as
Revised: 7 as

Original: [Equation]
Revised: [Equation]

Original: and similarly for the prior p ( z ) (Jensen et al., 2020).
Revised: and similarly for the prior p ( z ) (Jensen et al., 2020).

Original: To approximate this infinite sum, it is necessary to truncate it after a finite number of elements (Jensen et al., 2020; Falorsi et al., 2019; Jensen et al., 2022).
Revised: To approximate this infinite sum, it is necessary to truncate it after a finite number of elements (Jensen et al., 2020; Falorsi et al., 2019; Jensen et al., 2022).

Original: In the case where σ q ≪ 1 , it is sufficient to use a single term since there is negligible probability mass outside the area of injectivity ( [ π, π ] ).
Revised: In the case where σ q ≪ 1 , it is sufficient to use a single term since there is negligible probability mass outside the area of injectivity ( [ π, π ] ).

Original: This is in fact the regime we are operating in, since less than 0 .
Revised: This is in fact the regime we are operating in, since less than 0 .

Original: 2% of the prior probability mass falls outside the interval [ − π, π ] , and the posterior variance will in general be smaller than the prior variance.
Revised: 2% of the prior probability mass falls outside the interval [ − π, π ] , and the posterior variance will in general be smaller than the prior variance.

Original: Additionally, while the integral in Eq.
Revised: Additionally, while the integral in Eq.

Original: 6 covers the domain [ − π, π ] , we can approximate it with an integral from −∞ to ∞ since the posterior probability mass outside the area of injectivity is negligible, (cid:82) ∞ x = π N ( x ; 0 , σ 2 q ) dx ≈ 0 .
Revised: 6 covers the domain [ − π, π ] , we can approximate it with an integral from −∞ to ∞ since the posterior probability mass outside the area of injectivity is negligible, (cid:82) ∞ x = π N ( x ; 0 , σ 2 q ) dx ≈ 0 .

Original: These approximations together imply that we can simply approximate the KL term in Eq.
Revised: These approximations together imply that we can simply approximate the KL term in Eq.

Original: 6 as the KL between our prior and posterior Gaussians directly in the tangent space , and we verified numerically that this approximation is excellent for σ 2 q ≤ 1 .
Revised: 6 as the KL between our prior and posterior Gaussians directly in the tangent space , and we verified numerically that this approximation is excellent for σ 2 q ≤ 1 .

Original: We thus obtain a simple analytical expression for the KL:
Revised: We thus obtain a simple analytical expression for the KL:

Original: [Equation]
Revised: [Equation]

Original: Note that this is only possible because our prior is localized , in contrast to previous work which often uses a more diffuse or even uniform prior distribution (Jensen et al., 2020; Falorsi et al., 2019), and because we consider circular latent spaces where the Jacobian is 1 .
Revised: Note that this is only possible because our prior is localized , in contrast to previous work which often uses a more diffuse or even uniform prior distribution (Jensen et al., 2020; Falorsi et al., 2019), and because we consider circular latent spaces where the Jacobian is 1 .

Original: Since our prior and posterior both factorize across dimensions, these considerations also generalize to higher dimensional spaces, where the multivariate KL divergence is simply the sum across dimensions of the corresponding univariate KL divergences.
Revised: Since our prior and posterior both factorize across dimensions, these considerations also generalize to higher dimensional spaces, where the multivariate KL divergence is simply the sum across dimensions of the corresponding univariate KL divergences.

Original: Finally, to 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.
Revised: Finally, to 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.

Original: The simplified setup trained about 5 × faster, while behaving nearly identically on all test cases that we evaluated. Hence, all experiments in the paper are performed with the simplified objective.
Revised: The simplified setup trained about 5 × faster, while behaving nearly identically on all test cases that we evaluated, hence, all experiments in the paper are performed with the simplified objective.

Original: 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.
Revised: 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.

Original: B D ATA G ENERATION FOR SIMULATIONS
Revised: B D ATA G ENERATION FOR SIMULATIONS

Original: 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.
Revised: 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.

Original: More specifically, we assume a common shape for each curve (location assigned randomly on S 1 ), with a tuning width of 1 . 2 radians.
Revised: More specifically, we assume a common shape for each curve (location assigned randomly on S 1 ), with a tuning width of 1 . 2 radians.

Original: Peak rates are set to 0 .
Revised: Peak rates are set to 0 .

Original: 5 spikes per bin, the background rate to 0 .
Revised: 5 spikes per bin, the background rate to 0 .

Original: 005 , with a signal-to-noise ratio of 1 for the Poisson noise.
Revised: 005 , with a signal-to-noise ratio of 1 for the Poisson noise.

Original: The latent variable is simulated according to a Gaussian process with kernel standard deviation equal to 5 . 0 , and with kernel scale 50 . 0 .
Revised: The latent variable is simulated according to a Gaussian process with kernel standard deviation equal to 5 . 0 , and with kernel scale 50 . 0 .

Original: It is then projected onto S 1 via the modulo operation.
Revised: It is then projected onto S 1 via the modulo operation.

Original: The ensemble detection experiments in Sec.
Revised: The ensemble detection experiments in Sec.

Original: 4.2 use the same data generation, only with multiple ( 2 or 3 ) ensembles.
Revised: 4.2 use the same data generation, only with multiple ( 2 or 3 ) ensembles, with distinct latent varibles for each ensemble.

Original: C E NSEMBLE DETECTION COMPARISONS
Revised: C E NSEMBLE DETECTION COMPARISONS

Original: For raw-pca , we performed k-means clustering on dimensionally reduced neural activity over time (loosely inspired by Hamm et al., 2021; Lopes-dos Santos et al., 2011; Baden et al., 2016).
Revised: For raw-pca , we performed k-means clustering on dimensionally reduced neural activity over time (loosely inspired by Hamm et al., 2021; Lopes-dos Santos et al., 2011; Baden et al., 2016).

Original: Specifically, we took the projection of the spike matrix onto its first 8 (cross-validated) principal components.
Revised: Specifically, we took the projection of the spike matrix onto its first 8 (cross-validated) principal components.

Original: We then computed k-means clustering on this reduced matrix with the standard scikit-learn implementation, using a fixed number of 2 ( 3 , depending on the setting) clusters, 100 initializations and 1000 maximum iterations.
Revised: We then computed k-means clustering on this reduced matrix with the standard scikit-learn implementation, using a fixed number of 2 ( 3 , depending on the setting) clusters, 100 initializations and 1000 maximum iterations.

Original: For cov-pca , we performed clustering on the neural covariance matrix (loosely inspired by CarrilloReid et al., 2015).
Revised: For cov-pca , we performed clustering on the neural covariance matrix (loosely inspired by CarrilloReid et al., 2015).

Original: Specifically, we computed the covariance across neurons from the spike matrix.
Revised: Specifically, we computed the covariance across neurons from the spike matrix.

Original: Again, we computed the first 8 (cross-validated) principal components of the absolute value of this covariance matrix.
Revised: Again, we computed the first 8 (cross-validated) principal components of the absolute value of this covariance matrix.

Original: The absolute value was taken to preserve any negative or positive dependencies across neurons, as one would presumably observe within an ensemble — removing the modulus operator drastically reduced ensemble detection accuracy.
Revised: The absolute value was taken to preserve any negative or positive dependencies across neurons, as one would presumably observe within an ensemble — removing the modulus operator drastically reduced ensemble detection accuracy.

Original: From the projection onto the principal components we proceeded again as above.
Revised: From the projection onto the principal components we proceeded again as above.

Original: We also performed a clustering of the inverse of the covariance matrix, as a 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.
Revised: We also performed a clustering of the inverse of the covariance matrix, as a 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.

Original: Finally, for the supervised method we performed clustering on the mutual information score (see scikit-learn for implementation).
Revised: Finally, for the supervised method we performed clustering on the mutual information score (see scikit-learn for implementation).

Original: between each neuron’s activity, and the true latent variables, as input.
Revised: between each neuron’s activity, and the true latent variables, as input.

Original: Specifically, the features y that were used as input to k-means (same settings as above) were y ij ∼ I [ x i ; z j ] for every neuron i and every ground truth latent dimension j .
Revised: Specifically, the features y that were used as input to k-means (same settings as above) were y ij ∼ I [ x i ; z j ] for every neuron i and every ground truth latent dimension j .

Original: This presents an upper bound to the achievable performance, since it assumes supervised knowledge of the true latent variable values.
Revised: This presents an upper bound to the achievable performance, since it assumes supervised knowledge of the true latent variable values.

Original: D T HE E VIDENCE L OWER B OUND (ELBO)
Revised: D T HE E VIDENCE L OWER B OUND (ELBO)

Original: Here we show the derivation of the ELBO, which is maximized as part of the objective function during the model training procedure
Revised: Here we show the derivation of the ELBO, which is maximized as part of the objective function during the model training procedure

Original: [Equation]
Revised: [Equation]

Original: D. T EMPORAL TRANSITION PRIORS
Revised: D. T EMPORAL TRANSITION PRIORS

Original: All experiments in the article are conducted using the objective function from Eq.
Revised: All experiments in the article are conducted using the objective function from Eq.

Original: Extending this with additional terms however, e.g. by including a Laplacian temporal transition prior (Klindt et al., 2021) or a Gaussian transition prior (i.e. an L 1 or L 2 penalty across latent time steps, respectively) is relatively simple, and both the L 1 and L 2 penalties are included as optional regularization terms in our model.
Revised: Extending this with additional terms however, e.g. by including a Laplacian temporal transition prior (Klindt et al., 2021) or a Gaussian transition prior (i.e. an L 1 or L 2 penalty across latent time steps, respectively) is relatively simple, and both the L 1 and L 2 penalties are included as optional regularization terms in our model.

Original: 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 in the article, which is in agreement with results in Appendix Section E and
Revised: 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 in the article, which is in agreement with results in Appendix Section E and

Original: J regarding the convolutional kernel size.
Revised: K regarding the convolutional kernel size.

Original: E A BLATION STUDY — CONVOLUTIONAL FILTER
Revised: E A BLATION STUDY — CONVOLUTIONAL FILTER

Original: As mentioned in Sec.
Revised: As mentioned in Sec.

Original: 4.1, we conduct a 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).
Revised: 4.1, we conduct a 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).

Original: This is done to reduce the potential gap in information content available between a convolution and a non-convolutional model.
Revised: This is done to reduce the potential gap in information content available between a convolution and a non-convolutional model.

Original: Note, however, that this applies only to the variational posterior.
Revised: Note, however, that this applies only to the variational posterior.

Original: At inference time, both models infer the best latent at each point in time separately.
Revised: At inference time, both models infer the best latent at each point in time separately.

Original: Experimental settings are precisely the same as those in Sec.
Revised: Experimental settings are precisely the same as those in Sec.

Original: 4.1, other than the size of the convolutional kernel.
Revised: 4.1, other than the size of the convolutional kernel.

Original: One of the main observations from Fig.
Revised: One of the main observations from Fig.

Original: 8C-D is that the faeLVMs still outperform mGPLVM, even without contributions from the convolutional kernel.
Revised: 8C-D is that the faeLVMs still outperform mGPLVM, even without contributions from the convolutional kernel.

Original: In fact, faeLVM-s performs much better now, compared with the results from Fig.
Revised: In fact, faeLVM-s performs much better now, compared with the results from Fig.

Original: 4, achieving results that rivals those of faeLVM-b (Fig.
Revised: 4, achieving results that rivals those of faeLVM-b (Fig.

Original: 8A-B)
Revised: 8A-B)

Original: F M ODEL EXTENSION — CALCIUM RESPONSES AND VISUAL DATA SIMULATION
Revised: F M ODEL EXTENSION — CALCIUM RESPONSES AND VISUAL DATA SIMULATION

Original: In an effort to improve on the applicability of faeLVM, and showcase its relevance regarding applications to different brain regions, different latent topologies and different input data, we include a toy example on visual data, where we modeled a moving dot stimulus (Fig.
Revised: In an effort to improve on the applicability of faeLVM, and showcase its relevance regarding applications to different brain regions, different latent topologies and different input data, we include a toy example on visual data, where we modeled a moving dot stimulus (Fig.

Original: 9A) and a population response of retinal ganglion cells with center surround receptive fields and Poisson spiking (Fig.
Revised: 9A) and a population response of retinal ganglion cells with center surround receptive fields and Poisson spiking (Fig.

Original: 9B-C).
Revised: 9B-C).

Original: Each neuron’s spiking response was then convolved with a double exponential calcium fluorescence kernel (Fig.
Revised: Each neuron’s spiking response was then convolved with a double exponential calcium fluorescence kernel (Fig.

Original: 9D).
Revised: 9D).

Original: We further extended faeLVM by including a Gaussian likelihood as a modeling option, as well as an option to infer latent trajectories not restricted to circular and toroidal latent manifolds.
Revised: We further extended faeLVM by including a Gaussian likelihood as a modeling option, as well as an option to infer latent trajectories not restricted to circular and toroidal latent manifolds.

Original: Results can be seen in Fig.
Revised: Results can be seen in Fig.

Original: 9E-F, where the model is able to infer the correct two-dimensional latent variable, as well as accurately predict calcium responses of test neurons on test data.
Revised: 9E-F, where the model is able to infer the correct two-dimensional latent variable, as well as accurately predict calcium responses of test neurons on test data.

Original: Although this experiment is less comprehensive, being a toy example, it still gives a clear indication that our model is also applicable to both calcium traces and non-grid and head direction cells.
Revised: Although this experiment is less comprehensive, being a toy example, it still gives a clear indication that our model is also applicable to both calcium traces and non-grid and head direction cells.

Original: G A LL GRID CELL RATE MAPS
Revised: G A LL GRID CELL RATE MAPS

Original: In this section, we showcase rate maps for all neurons in each of the two ensembles ( 149 andneurons in the first and second grid cell module respectively).
Revised: In this section, we showcase rate maps for all neurons in each of the two ensembles ( 149 andneurons in the first and second grid cell module respectively).

Original: The data used is the same as discussed in Sec. 4.4, pre-processed in the same way, and we used the same model to produce latent decodings, receptive fields and ensembles.
Revised: The data used is the same as discussed in Sec. 4.4, pre-processed in the same way, and we used the same model to produce latent decodings, receptive fields and ensembles.

Original: Fig.
Revised: Fig.

Original: 10 and Fig.
Revised: 10 and Fig.

Original: 11 show the rate maps (using 50 × 50 spatial bins, smoothed for visualization purposes in the same manner as Gardner et al.
Revised: 11 show the rate maps (using 50 × 50 spatial bins, smoothed for visualization purposes in the same manner as Gardner et al.

Original: (2022), with a Gaussian kernel of smoothing width
Revised: (2022), with a Gaussian kernel of smoothing width

Original: As mentioned in Sec. 4.
Revised: As mentioned in Sec. 4.

Original: we clearly see the more ambiguous rate maps at the bottom of these two figures (in particular, the three-four bottom-most rows).
Revised: we clearly see the more ambiguous rate maps at the bottom of these two figures (in particular, the three-four bottom-most rows).

Original: 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.
Revised: 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.

Original: H A DDITIONAL VISUALIZATIONS FOR H EAD D IRECTION AND G RID C ELL D ATA
Revised: H A DDITIONAL VISUALIZATIONS FOR H EAD D IRECTION AND G RID C ELL D ATA

Original: 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.
Revised: To supplement the results shown in Sec. 4.3 and Sec. 4. we have included additional figures for both the head direction application and the grid cell application.

Original: More precisely, Fig.
Revised: More precisely, Fig.

Original: 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.
Revised: 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.

Original: 6A), while Fig.
Revised: 6A, while Fig.

Original: 12B shows inferred latent dynamics on grid cell data, as function of spatial position.
Revised: 12B shows inferred latent dynamics on grid cell data, as function of spatial position.

Original: As there is no way of recording the ’true’ toroidal coordinates (akin to what is done in the head direction dataset), this figure is included in an effort to 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.
Revised: As there is no way of recording the ’true’ toroidal coordinates (akin to what is done in the head direction dataset), this figure is included in an effort to 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.

Original: I H YBRID INFERENCE I. O VERVIEW
Revised: I E XTENDED ENSEMBLE DETECTION ABLATION

Revised: In this section we present an additional experiment on the ensemble detection task.

Revised: Namely, we ask the model to separate five distinct ring-like ensembles, z j ∈ S 1 , to complement the experiments from Section 4.2, which only considered two and three ensembles.

Revised: From Fig.

Revised: 13, we can see that faeLVM-b is able to separate the five ensembles quite clearly.

Revised: The weights in Fig.

Revised: 13A is approaching the desired one-hot encoding, and we see that the five inferred ensembles correspond very nicely with the true ensembles (Fig. 13C).

Revised: The model is also able to accurately predict spike rates on held out neurons from each ensemble (13B).

Revised: Thus, the ablation 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.

Revised: J H YBRID INFERENCE J. O VERVIEW

Original: Here, we give an overview on how the hybrid inference is performed.
Revised: Here, we give an overview on how the hybrid inference is performed.

Original: In contrast to the traditional variational approach, where one uses a single sample from the variational posterior as prediction for the latent variable, we propose drawing a selection of samples, then performing gradient descent on said latents using the Adam optimizer.
Revised: In contrast to the traditional variational approach, where one uses a single sample from the variational posterior as prediction for the latent variable, we propose drawing a selection of samples, then performing gradient descent on said latents using the Adam optimizer.

Original: Specifically, if z are the latents for the test set (as inferred by the encoder and variational posterior) and x the responses of the training neurons on the test set, we maximize
Revised: Specifically, if z are the latents for the test set (as inferred by the encoder and variational posterior) and x the responses of the training neurons on the test set, we maximize

Original: [Equation]
Revised: [Equation]

Original: (12) for 2000 steps with the Adam optimizer and a learning rate of 0 . 001 .
Revised: (12) for 2000 steps with the Adam optimizer and a learning rate of 0 . 001 .

Original: 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).
Revised: 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).

Original: We emphasize that all model parameters are fixed during hybrid inference, and the optimization only affects the sampled latents.
Revised: We emphasize that all model parameters are fixed during hybrid inference, and the optimization only affects the sampled latents.

Original: 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.
Revised: 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.

Original: This means that we are not just overfitting but actually improving our point estimates of the latents in this inference procedure.
Revised: This means that we are not just overfitting but actually improving our point estimates of the latents in this inference procedure.

Original: We leave optimisation of the full likelihood, including uncertainty estimates on the latents, for future research.[Table]
Revised: We leave optimisation of the full likelihood, including uncertainty estimates on the latents, for future research.

Original: I. H YBRID I NFERENCE SIMULATIONS
Revised: J. H YBRID I NFERENCE SIMULATIONS

Original: For evaluate of the performance gain with hybrid inference at test-time [Table] we consider six variations of the faeLVM; that being the three cases described in Sec. 3, each alternative including a case of using hybrid inference and one using the standard variational posterior approach (fae vs v-fae).
Revised: For evaluate of the performance gain with hybrid inference at test-time, we consider six variations of the faeLVM; that being the three cases described in Sec. 3, each alternative including a case of using hybrid inference and one using the standard variational posterior approach (fae vs v-fae).

Original: As in Sec. 4. [Table] we divide our data into four parts (Fig. 3), generate a 1 D periodic latent variable and simulate spikes according to a Poisson distribution with Gaussian tuning curves.
Revised: As in Sec. 4.1, we divide our data into four parts (Fig. 3), generate a 1 D periodic latent variable and simulate spikes according to a Poisson distribution with Gaussian tuning curves.

Original: We fix T test and N test ( 1000 and 30 ), vary T train and N train over a selection of values, and for each condition, repeat the experiment over 20 seeds (synthetic data is different for each seed, but all of the models are trained on the same dataset, to ensure that comparisons are still valid), reporting mean negative log-likelihood (NLLH) and mean geodesic error (GE) on test data.
Revised: We fix T test and N test ( 1000 and 30 ), vary T train and N train over a selection of values, and for each condition, repeat the experiment over 20 seeds (synthetic data is different for each seed, but all of the models are trained on the same dataset, to ensure that comparisons are still valid), reporting mean negative log-likelihood (NLLH) and mean geodesic error (GE) on test data.

Original: The results presented in Table 1 indicate that performance, in almost every case, increases substantially by including the hybrid inference step.
Revised: The results presented in Table 1 indicate that performance, in almost every case, increases substantially by including the hybrid inference step.

Original: Hence, all other experiments performed in this article are done solely using the hybrid inference where applicable (e.g., when applying the model to grid  cell data, evaluations are done in-training-sample, and without the focus on latent prediction on a test set, hence there is no need for hybrid inference at test-time).
Revised: Hence, all other experiments performed in this article are done solely using the hybrid inference where applicable (e.g., when applying the model to grid [Table] cell data, evaluations are done in-training-sample, and without the focus on latent prediction on a test set, hence there is no need for hybrid inference at test-time) [Table]

Original: J H YPERPARAMETER SEARCH
Revised: K H YPERPARAMETER SEARCH

Original: To investigate hyperparameters and their effect, we selected relevant parameters and estimated their ranges based on related literature.
Revised: To investigate hyperparameters and their effect, we selected relevant parameters and estimated their ranges based on related literature.

Original: The relevant parameters can be seen on the left in Table 2.
Revised: The relevant parameters can be seen on the left in Table 2.

Original: To determine optimal values of hyperparameters, we relied on a random search strategy.
Revised: To determine optimal values of hyperparameters, we relied on a random search strategy.

Original: A grid search is not feasible due to the curse of dimensionality, and a random search can also be more effective (Bergstra & Bengio, 2012).
Revised: A grid search is not feasible due to the curse of dimensionality, and a random search can also be more effective (Bergstra & Bengio, 2012).

Original: Furthermore, the random search can easily be used for discrete and continuous variables, that could be on linear or logarithmic scales.
Revised: Furthermore, the random search can easily be used for discrete and continuous variables, that could be on linear or logarithmic scales.

Original: The selected scales and ranges, can be found in the middle in Table 2.
Revised: The selected scales and ranges, can be found in the middle in Table 2.

Original: In total, we sample 500 hyperparameter configurations.
Revised: In total, we sample 500 hyperparameter configurations.

Original: Next, we train and test each configuration for 3 random seeds.
Revised: Next, we train and test each configuration for 3 random seeds.

Original: We evaluate the models on the test accuracy and mean correlation.
Revised: We evaluate the models on the test accuracy and mean correlation.

Original: The final selected model is shown on the far right in Table 2.
Revised: The final selected model is shown on the far right in Table 2.

Original: Lastly, we also measure whether there is an ordinal correlation (Spearman) between the varied factors and the test LLH.
Revised: Lastly, we also measure whether there is an ordinal correlation (Spearman) between the varied factors and the test LLH.

Original: The results are shown in Fig.
Revised: The results are shown in Fig.

Original: We see that learning the variance (learn var) increases the performance on average.
Revised: We see that learning the variance (learn var) increases the performance on average.

Original: Also, lower learning rates and larger batch-sizes lead to better results.
Revised: Also, lower learning rates and larger batch-sizes lead to better results.

Original: Lastly, also lower weighting for the KL term leads to better results.
Revised: Lastly, also lower weighting for the KL term leads to better results.

Original: Hyperparameter
Revised: Hyperparameter

Original: Ranges
Revised: Ranges

Original: Best kernel size choice([1, 9, 17]) num hidden choice([16, 32, 64, 128, 256, 512]) shared choice([True, False])
Revised: [Table] [Table] choice [Table]

Original: False
Revised: [Table] [Table] [Table] [Table]

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: choice([True, False])
Revised: True [Table] [Table]

Original: True learn mean choice([True, False])
Revised: True [Table] choice([True [Table]

Original: False learn var choice([True, False])
Revised: [Table] choice([True [Table] [Table]

Original: False isotropic choice([True, False])
Revised: isotropic choice([True, [Table] False [Table]

Revised: [Table]

Original: False num basis choice([1, 2, 4, 8]) nonlinearity choice([exp, softplus]) softplus batch size choice([1, 16, 32]) batch length choice([64, 128]) learning rate loguniform(1e-4, 1e-1) 0.
Revised: [Table] choice [Table] [exp,

Revised: [Table]

Revised: [Table]

Revised: choice

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: choice

Revised: [Table]

Revised: [Table]

Revised: 128]) learning

Revised: [Table]

Revised: loguniform

Revised: [Table]

Revised: 1e

Revised: [Table]

Revised: 4, 1e

Revised: [Table]

Original: num worse choice([10, 50, 100]) weightkl choice([0., loguniform(1e-9, 1e0)]) 0.
Revised: [Table] choice [Table] ]) weight

Revised: [Table]

Revised: choice

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: 1e

Original:  weight time choice([0., loguniform(1e-9, 1e0)])0.
Revised: [Table] [Table] [Table] weight time choice([0., loguniform(1e-9, 1e0)])

Revised: [Table]

Original: weight entropy choice([0., loguniform(1e-9, 1e0)]) 0.
Revised: weight entropy choice([0., loguniform(1e-9, 1e0)]) 0.