Learning Time-Invariant Representations for Individual Neurons from Population Dynamics
Authors: Lu Mi
Trung Le
Tianxing He
Eli Shlizerman
Uygar Sümbül
Venue: NeurIPS 2023 poster (2023)
Keywords: population dynamics
neuronal representation
calcium imaging
cell types
TL;DR: We fit dynamical models to neuronal activity to learn an invariant representation by considering the activity of both the individual and the neighboring population.

Abstract: Neurons can display highly variable dynamics. While such variability presumably supports the wide range of behaviors generated by the organism, their gene expressions are relatively stable in the adult brain. This suggests that neuronal activity is a combination of its time-invariant identity and the inputs the neuron receives from the rest of the circuit. Here, we propose a self-supervised learning based method to assign time-invariant representations to individual neurons based on permutation-, and population size-invariant summary of population recordings. We fit dynamical models to neuronal activity to learn a representation by considering the activity of both the individual and the neighboring population. Our self-supervised approach and use of implicit representations enable robust inference against imperfections such as partial overlap of neurons across sessions, trial-to-trial variability, and limited availability of molecular (transcriptomic) labels for downstream supervised tasks. We demonstrate our method on a public multimodal dataset of mouse cortical neuronal activity and transcriptomic labels. We report >35\% improvement in predicting the transcriptomic subclass identity and >20\% improvement in predicting class identity with respect to the state-of-the-art.

Learning Time-Invariant Representations for Individual Neurons from Population Dynamics
Lu Mi M Trung Le Tianxing He², Eli Shlizerman', Uygar Sümbül!
1 Allen Institute for Brain Science
University of Washington org
{lu mi, ayggarsaalleniistittue
(tle45, shlizee)0uw edu goosehe@cs washington edu
Abstract
Neurons can display highly variable dynamics. While such variability presumably supports the wide range of behaviors generated by the organism, their gene expres- sions are relatively stable the adult brain. This suggests that neuronal activity a combination of its time-invariant identity and the inputs the neuron receives from the rest of the circuit. Here, we propose self-supervised learning based method to assign time-invariant representations to individual neurons based permutation-, and population size-invariant summary population recordings. We dynamical models to neuronal activity to learn representation by considering the activity of both the individual and the neighboring population. Our self-supervised approach and use of implicit representations enable robust inference against imperfections such as partial overlap of neurons across sessions, trial-to-trial variability, and limited availability of molecular (transcriptomic) labels for downstream supervised tasks. We demonstrate our method on public multimodal dataset mouse cortical neuronal activity and transeriptomic labels. We report 35% improvement in pre- dicting the transcriptomic subclass identity and > 20% improvement in predicting class identity with respect the otate-of-the-art.
Introduction
Population recordings of neuronal activity enable relating behaviorally-relevan dynamics to the summary activity of the recorded population. While this has produced numerous insights into how the brain works []], the activity and identity of individual neurons should be analyzed to achieve mechanistic understanding at the implementation level (2), which may hold the key to new cologicalllyiinsiree algorithms
B 4 Moreover, emerging experimental evidence suggests that neurons have diverse yet stable molecular identities, which can dictate their computational roles 6
Joint (i.e., multimodal) profiling and alignment of electrophyssioogicall features and gene expression of individual neurons suggest good correspondence between these two modalities in slice experi- ments (7H11. Recently, population recordings of calcium activity followed by spatially registered single-cell transcriptomic recordings enabled similar joint profiling of in-vivo activity and molecular identity [6]. Importantly, such activity depends on both the intrinsic physiological properties of neurons and the exogenous inputs (synaptic and modulatory) to those neurons, which are themselves a product of both sensory inputs to the organism and the recurrent activity in the brain.
While recording from molecularly defined neuron populations has been popular method, these experiments do not allow for studying the concurrent responses different neuron types to stimuli.
Equual contribution
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

On the flip side, joint profiling of panneuronal population activity and transcriptomics is slow, expensive, and not available to many research labs. Learning the association between these two observation modalities can minimize the need for joint profiling and provide neurobiological insights.
The constancy of neuronal identity in adults in the face of a potentially rapidly changing environment represents key challenge: the inferred identity should be invariant to time and the task that the organism engages with, suggesting that the inference method should ideally be invariant to those variables. In the absence of a priori information on identity, it is also desirable that the invariance extends to the number and the (arbitrary) ordering of experimental population. Moreover, technical challenge common to many multimodal datasets is that only a relatively small fraction of the observations tend to be jointly characterized (or otherwise labeled), limiting the applicability of supervised approaches.
To address these problems, here, we develop self-supervised approach Neuronal Time- Invariant Representations (NeuPRINT), to infer neuronal identity from population recordings by forming a model of activity dynamics that depends only on past activity of the neuron itself and statistics of past population activity that are invariant to the ordering of the individuals and asymptotically invariant to the size of the population. We demonstrate the utility of the inferred identities by reporting the performance of a simple classifier of transcriptomic identity on those representations and other baselines. We also study the impact of providing similarly invariant yet more detailed information on the population by partitioning it into center vs surround subsets, reflecting well-known yet simple connectional and functional property of neuronal circuits 14-
2 Related work
Latent dynamical models: Latent modeling of time series of population activity aims to gain insight on the dynamics of the totality of the activity of a single experimental trial [5 When interpretabillty of the latent dynamics is not a priority, powerful recurrent neural networks uncover high-fidelity representations (21 22]. Recently, low-dimensional time series representations were obtained with high reproducibility. albeit without model of the dynamics, using a contrastive loss objective [23 Another recent study obtained parametric representations of individual neurons in population recordings within a linear dynamics framework, where the activity of the neuron is informed by a latent representation of the population so that the results are not transferable across studies that do not have the same population with the same ordering of the neurons (24). The idea of informing neuronal activity prediction of neighboring activity and/or extrinsic covariates such as stimuli or behavior was also explored in previous classical studies 25 26]
Liu et al. 27 proposed the use of the transformer architecture 288 to generate representations of individual neurons in a supervised paradigm, which do not interact with the rest of the population. Instead, those representations are combined to accomplish the supervised task.
Multi-modal neural data: Characterization of intrinsic electrophysiology of neurons typically uses 10 [1I or mechanistic model parameters (30), instcad of time series models, align gene expression and physiology in multimodal slice recordings. Such experiments study only the neuron in isolation and the recordings may not reflect in-vivo activity in the context of a behavior.
set of features 7
9 29
A large-scale study that focuses the in-vivo activity of the individual neurons only, and not their interaction, uses a neural network with attention layers in a supervised setting to predict the neuronal subclass, based on transgenic mouse lines 31
Implementing nonlinear autoregressive models with exogenous inputs (NARX models) using recurrently connected neural networks has a long history [2) Modern imple- mentations typically use gated units to address the vanishing gradients problem [33]. Recently, the transformer architecture (28), has produced impressive results in sequence-to-sequeee
Other predictive models:
tasks. Despite interpretability issues, such models can be significantly more powerful than classical dynamical systems models (e.g., linear dynamical system) [34 35). In this work, we use transformer dynamical model to integrate with both individual and population dynamics, and the representation of individual neuron is assigned with time-invarian embedding.
equation: For point neuron models, which ignore the spatial distribution of the different ion channels, the celebrated Hodgkin-Huxley (HH) equations provide a phenomenologica description of the neuronal membrane potential via coupled differential equa- tions [36]] 37.. The membrane potential controls the generation of spiking outputs and calcium
Phenomenological Hodgkin-Husee

---

imaging provides an indirect measurement of this voltage over time [8). For in vivo measurements, the synaptic input current injected by other neurons (pre-synaptic partners of the neuron of interest) is represented as exogenous input, which can be written sum over the synapses. The HH equations also include multiple neuron-specifie, time-invariant parameters. Thus, the HH equations consider neuronal activity as a function of both the activities of other neurons and the intrinsic parameters of its electrophysiology.
3 Methods
Implicit dynamics models for neuronal activity: Neuronal dynamics are significantly more complex than the HH equations because the physical distribution of the various channels on the neuronal arbor, nonlinear computing abilities of the dendrites, modulatory communication between neurons,et ettc. can (i) modulate the neuron-specific parameter set, (ii) add more parameters that set, and (iii) change the functional form of HH equations. Moreover, the relationship between the membrane voltage and the observable of most neuronal population activity experiments, the calcium dynamics, is itself complex. Finally, it appears impossible to perform the detailed measurements needed to fit the parameters of the HH equations based on in vivo experiments with current technology. Thus, we pursue an implicit modeling approach while trying to capture the fundamental dependencies with the help of flexible deep neural network parametrization. Let x.() denote the calcium activity of neuron i at time and consider the following equation for dynamics:
dxx(() = f(x )D())
(1)
dt
where p(-1) denotes the activity of all the neurons that provide (synaptic or extra-synaptic) input to representation for neuron i, which implicitly captures the intrinsic parameters of HH equations and other such iineiinvariant aspects of neuronal identity.
neuron at time t. 4(()) denotes
Permutation-invaariat summary of population activity: Even if neurons were identifiable in population imaging experiments, the connectivity of neurons and our observations of it can be considered as stochastic events [39]. Therefore, to enable the transfer of knowledge across sessions, experiments, and individuals, it is highly desirable to approximate the dependency of neuronal dynamics on P. (Eq. 1 in a way that is invariant to permutations, number, and detailed identity of the neurons contributing it. To achieve this, we propose to replace P. with multiple (asymptotically) invariant statistics of the activity of a neighboring population of neurons, such as average activity 40]
Behavioral observations, such pupil diameter, can serve indirect readouts the activity of unobserved neurons. Concatenating these observations with the aforemennioned statistics will enrich the exogenous input observed by the dynamical model. We call this new concatenated variable P..
Center-surround partition: Synaptic connection probability between neurons depends on dis- tance 12 [3]. Similarly, neuronal co- variability correlates with spatial distance [14. To include this neurobiological insight while maintaining pormutation-invarian extension: we partition the population activity into two groups, center and surround, and compute the relevant invariant statistics for each group separately. Such partitioning is easy to obtain in calcium imaging experiments since the distances between the somata of neurons are readily available. Beyond its simplicity, this choice is also motivated by surround suppression being a connectivity motif in the brain [41 42).
of our model, we propose simple
Discrete-time dynamical model of neuronal activity: We re- re-irite neuronal dynamics in discrete time as
( = S(x
(c(-i) (C St +41 B B (1)),
5(-1)
where X E NxT represents the activity data for the whole recorded population with N neurons and T time steps. C,S. and B represent D, D, and D' -dimensional permutation- and size-invariant (i.c., N) surrogates for brain activity at each time point. C' and S compute identical statistics of population activity (here, mean and standard deviation) except that the statistics in are calculated over the center partition (neurons whose distance neuron is at most A) and the statistics in S are calculated over the surround partition (neurons whose distance to neuron i is larger than A), see Appendix E Here, A is a hyperparameter. B denotes the contribution of temolvedllved behavioral observations. Hence, the triplet (C,, B.) corresponds to P.. ((()) e æK denotes K-dimensional
3

---

back propagation
behavior
inputs
information
gradient
(peddiction loss
outputs
session
running speed pupil size
session
X-
neuron indek
session
newron Index
session
Transformer
suouaau
X
implicit
dynamical
primary cortex (V1)
model
time
NxK
time
W
masked calcium fluorescence trace inputs
learned time invariant representation +
predicted calcium fluorescence trace outputs
overview of dynamical model f
Figure 1: Overview of self-supervised representation learning framework NeuPRINT. Activities of N neurons (recorded by 2-photon calcium imaging of the mouse primary visual cortex) and behavior information (pupil size, running speed, etc.) across multiple sessions are used as inputs to fit an implicit dynamical model f and learn time-invariant x representation Φ. The learned representations are later evaluated supervised downstream tasks to predict transcriptomic class and subelass identities. In the optimization framework, neuron-specific representation Ø; is repeated at every time step, then concatenated with masked past neuronal activity x W- and permutation- invariant population inputs Pt-w+1d +1:t to form the input. The transformer model is trained to predict neural activity X at the masked step with a causal attention mask over the W- -ssepp context window.
time- invariant representation for neuron i, and W denotes the width of the available temporal context. The subscripts denote the limits of the time interval within which the corresponding variable is available to f.
Self-supervised Representation Learning Framework with Transformer: We thus propose to solve the following self-supervised optimization problem to infer both the function f and c(0);
arg min 1,000))
Σ x
|| 1+1 - -(X f(x (C/ W B. B Il),, (1)),,
(3)
where ||- denotes a norm. worth pointing out that f depends on the neuron of interest or other neurons in the population only through its explicit parameters. The reason for this choice is to maintain the transferability and ubiquity of the learned model f, the first argument the optimization, while summarizing neuronal variability with (1), the second argument.
The idea inferring invariant representations for neurons directly from in vivo recordings by fitting a dynamical model (i.e., predicting activities in the next time step) reminiscent of, and motivated by, the recent spectacular successes of "foundation models" in natural language modeling [43]. where capturing the dynamics of a complicated system produces an implicit understanding of the dynamics and identity of its components. This procedure has been shown useful for multiple downstream tasks (43]. Therefore. we use transformer model 28 to parametrize the function f. Unlike previous uses of the transformers for neural data [34] 35], the input tokens in our model do not come from countable set because raw calcium recordings are best represented by real valued signals.
To train the transformer and the time-invariant representation, we first generate masked activity inputs (i)
X 1+ where neuronal activity time + 1 is zero-out. This masked activity and the past activities
G-1+1:t are concatenated with the time-invariant representations ( and permutation-invant summary of population dynamics Pr-w+1:t to form the inputs to the transformer. We ask the transformer to predict the activity at the masked step, and compute the loss from the predicted activity (i) t+1 and ground truth activity ( The dynamical model f and time-invariant representation ((+) are jointly learned during the optimization. To perform this task, the transformer has to learn the temporal progression of neuronal activity conditioned on the neuronal identity and the causal temporal

---

context of individual and population statistics. An overview of NeuPRINT learning framework is depicted in Figure
4
Experiments
4.1 A Multimodal Dataset
We use recent, public multimodal dataset to train and demonstrate our model: Bugeon et al. [6] obtained population activity recordings from the mouse primary visual cortex (VI) via calcium imaging, followed by single-cell spatial transcriptomics of the tissue and registration of the two image sets to each other to identify the cells across the two experiments. 2-photon calcium imaging recordings were obtained with a temporal sampling frequency of 4.3Hz. And the spatial coordinates of recorded neurons are also provided. We first evaluate our approach on one animal (SB025) across 6 sessions. The recordings from this animal include 2481 neurons in total. We then extend our analysis on functional recordings from 4 mice (SB025, SB026, SB028, SB030) across 17 sessions. They contain 9728 neurons in total. Each session lasts about 20 minutes and records about 500 neurons. A small subset of neurons overlap across sessions. The subsequent transcriptomic experiment profiles mRNA expression for 72 selected genes in ex vivo tissue. These genes were used to identify the excitatory vs inhibitory class labels of neurons. In addition, 51% of the neurons in the inhibitory class of SB025 also have identified subclass labels (Lamp5, Pvalb, Vip, Sneg, Sst). Finally, the dataset includes behavioral information for the mice (running speed and pupil size) during the in-vivo recording as well as an assignment for each image frame (i.c., time point) that we call frame state from the set (running, stationary desynchronized, stationary synchronized).
4.2 Benchmark Evaluation for Transcriptomic Identity Prediction
We use the aforementioned public dataset to introduce new two-step benchmark: (i) self-supervised learning of time- and permutation-invarian representations for individual neurons from population activity, (ii) prediction of labels for each individual neuron based on those representations.
We introduce a downstream classification task to predict the subclass label with supervised learning, where the neurons with subclass labels from all sessions are randomly split into train, validation and test neurons with a proportion of 80%: 10% : 10%. We further introduce another supervised downstream classification task to predict the class identity only (i.e., excitatory vs inhibitory). In this task, the validation and test neurons in the subclass prediction task are used as validation and test neurons for inhibitory neurons, and the same fraction of excitatory neurons are randomly selected as validation and test neurons. The rest of the recorded population from all sessions is used for training.
We first optimize the dynamical model (f in Eq. 2 and the time-invariant representation on the training set. We use past activities of the training neurons and cnnnuutationinnvarian summary of population dynamics including pupil size, running speed, frame state, the mean and standard deviation of population activity, and the mean and standard deviation of center-surroung activity to predict the individual neurons' activity in the next time step. After training, we fix the dynamical model f and only optimize the lime-invariant representations & for training, validation, and test neurons under the same supervieee learning framework.
Following self-supervised optimization, in the second step, we evaluate the learned representation with two supervised downstream tasks (class and subclass prediction) and three simple classifiers including k-nearest neighbor (KNN), linear model, multi-layer perceptron (MLP) with one hidden layer. The training neurons' representations & are used to train the classifier, and validation neurons are used to tune the hyperparameters (learning rate, hidden dimensionality, number of epochs, etc.), and the top-1 accuracies of all models are reported on the test neurons (Table I). See Appendix for an analysis of sensitivity.
4.3 Implementation of a spectrum of implicit dynamical models and downstream classifiers
We explore four different implicit (not mechanistic) dynamical models: linear, nonlinear, gated- recurrent network (GRU), and transformer with self-attention. We optimize the parameters of the dynamics f and the neuronal representation & using gradient descent for all models.
We use linear dynamical system where the activity at the last step is predicted from a linear combination of the activities and statisties of the population activities, behavioral information from previous steps inside temporal window and the iine-invariant representation (which is repeated
Linear model:
5

---

at each step). This corresponds to an autoregressive model with exogenous inputs, where statistics of the activities of other neurons and behavioral information constitute the exogenous input to the dynamics of the neuron of interest.
Nonlinear model: In addition to the linear model, nonlinear activation applied to the weighted activity, behavioral information, repeated neuronal representation at each step before the linear combination (i.e., nonlinear autoregressive model with exogenous inputs).
Recurrent network with gated units: The activity at the next step is predicted from the hidden state in addition to the activity at the current step and the repeated neuronal representation the inputs.
Transformer: We implement a W step causal attention mask such that the transformer predicts the activity of the neuron at the current time step based on the hidden states in the W previous time steps. The hidden state tasks the neuron activities, statistics of the population activities, behavioral information at that time step, and time-invariant representation repeated at each time step as inputs. We use the transformer encoder-only implementation from PyTorch 44 with 2 attention heads.
Training details:
For the objective function to predict the activity, we explore both mean squared error (MSE) and negative log likelihood (NLL) with a Gaussian distribution. To train the dynamical model and representation of neurons, we use 64-dimensional embedding for the time-invariant representation. The temporal trial window size is 200 steps for the linear. nonlinear models, recurrent network and transformer. The batch size is 1024. We use the Adam optimizer 45 with a learning rate of 10-3
Downstream supervised classification: For the linear and MLP classifiers, we use the cross-entropy loss train the model. For KNN, the scikit-leam implementation 46 with the number of nearest neighbors & = 5.
4.4 Baselines
31 is a supervised framework for predicting cell types LOLCAT and its variants: LOLCAT from individual neuronal activities using multi-head attention network. Since the attention in LOLCAT is a simple weighted sum operation, we implement two additional supervised variants Transformer+ISI and Transformer+Raw using the self-attention mechanism employed in NeuPRINT for a fair comparison. These two variants follow the attention design (28 and use special classification token represent the classification of the entire neuronal activity 47 Transformer+ISI operates on the inter-spike interval (ISI) distributions input summarized from non-overlapping sub- windows of continuous 2-photon calcium recordings. We use suite2p package from raw calcium traces and compute the ISI distributions. On the other hand, Transformer+Raw operates directly variants train both the attention network and classifier (linear or MLP) in an end-to-end fashion, using neuronal class and subclass labels during learning (See Appendix for details). As consequence of the supervised training scheme, LOLCAT does not extract time- invariant representations of neuronal identity and its performance is constrained by the number of labels available in the dataset.
[48) to infer spikes
the calcium traces. Unlike NeuPRINT, LOLCAT and the two supervised
Principal component analysis: We project the raw calcium activities to a low-dimensional represen- tation using Principal Component Analysis (PCA) and evaluate the effectiveness of this representation for downstream tasks. The projection is performed on a randomly selected sub-window in the raw recordings of each neuron, therefore its projected representation is also not time-invariant.
Uniform manifold approximation and projection: Similar PCA, we project the raw calcium activities to low- oow-dimesiiona representation using Uniform Manifold Approximation and Projection (UMAP) 49), and evaluate the resulted time-vrrian representation by downstream classifiers.
Random: We further generate random representations for individual neurons, and train a supervised classifier on the random representations measure the chance level of prediction.
5 Results
5.1 Self-supervised Learning Demonstrates Superior Generalization Capabilities in Data-Limited Scenarios
We evaluate our proposed method NeuPRINT and other baselines under three categories as shown in Table [ (i) iime-invariant vs. time-variant; (ii) self-supervised representation learning vs. end-to-end

---

Input
Individual without Population
Individual with Population
Supervision
Lower Bound
Data-Limited
Unsupervised
Self
Data-Limited
Self
Supervised
supervised
Supervised
-supervised
Trans Trans
Trans Trans
Task
Model
Random
LOLCAT
PCA UMAP
NeuPRINT
LOLCAT
NeuPRINT
+IST +Raw
+IST +Raw
KNN
0.260
0.263 0.281
0.415
0.610
Subclass
Linear
0.256
0.404
0.474 0,474
0.316 0.404
0.537
0.474
0.491 0,386
0.683
ML.P
0.302
0.561
0.439 0.330 0.340
0.512
0.526 0.386
0.756
KNN
0.488
0,536
0.584
0.652
0.711
Class
Linear
0.526
0.600
0.664 0.669 0.544 0.576
0.697
0.608
0.680 0.608
0.793
MLP
0.523
0.640 0.664 0,565 0.520
0.752
0.632 0.616
0.807
Table Top-1 accuracy of transcriptomic label prediction based on representations learned by (i) our
proposed self-supervised representation learning from neural dynamics framework NcuPRINT, (ii) the supervised learning method LOLCAT 31 and its variants Transformer++IS and Transformer+RRaw (iii) unsupervised baselines PCA and UMAP, (iv) random representations (to determine the chance- level). Note that this experiment corresponds to a data-limited regime due to limited labeled data. We performed classification using three classifiers (KNN, Linear, MLP) and two tasks: predicting the subclass from the set [Lamp5, Pvalb, Vip, Sneg, Sst}, and predicting the cell class from the set {excitatory, inhibitory J. We study the performance of different models using only individual neuronal activity vs adding population statistics as input.
NeuPRINT
NcuPRINT subciass
LOLCAT
True label
ECC
0.81
0.19
Lamp5
0.80
0:13
0.00
Lamp5
0.85
0.00
0:15
0:00
811
IN
0:29
0.71
True label
0.08
0.75
0.17
0.80
ltue laber
Vip
0.62
0:15
0.23
0.00
0.6
3.0
Predicted libel
Lamp5kdp Wpp Pvolo
LOLCAT EI
Pvalb
0.20
0.10
0.70
0.00
0.4
Pvalb
0.47
0.06
0.47
0.00
0.4
0.4-
True Maw
EC
0.90
0.30
0.-
Set
0.67
0.00
0.00
0.33
0.2
5et
0.50
0.17
0.33
0:00
0.2
IN
0.74
0.26
0.0
Lamp5
Predicted label
Lamps
Vip Predicted
Pvalb
Set
0.0
Prodicted
EC
IN
Figure 2: Left: Relative abundances of subclass and class labels. Right: Confusion matrices of self-supervised representation learning framework NeuPRINT and supervised learning method LOLCAT [11 based on predicting the cell class and subclass labels. While both of the sulf-supervised and supervised steps are learned with all available subclasses, we excluded the Sneg population from the confusion matrices because it represents a negligible fraction of the test set with the 80%: 10% 10% split, so that quantification for this population would not be reliable.
supervised learning vs. unsupervised learning; (ii) activity of the neuron of interest ("Individual") as the only input the dynamics model f vs. permutation-invarite dynamics provided as exogenous input to f. Under the data-limited scenario (i.c., a small amount of labeled samples), which describes vast majority of neuroscience datasets, we find that our self- supervised representation learning model NeuPRINT with a supervised downstream MLP classifier outperforms the current state-of tfte-of--hhe-rr approach LOLCAT by 35% and its variants by > 19% in the subclass prediction task and outperforms LOLCAT by > 20% and its variants by > 13% in the class prediction task. Since LOLCAT is optimized using features extracted from non-overlapping sub-windows in an end-to-end supervised learning approach, it does not generate time-invariant representations that are critical to generalize across trials. Moreover, as shown in the confusion matrices in Figure 2 when the data in the subclass and class prediction tasks has imbalanced distribution under the data-limited regime, our method can generate more balanced predictions across labels than LOLCAT. Both of these methods outperform two other standard unsupervised representation learning baselines, PCA and UMAP, which also extract time-varying representations across non-overlapping sub-windows. All of the evaluated models perform above the chance level, and we find that one-hidden layer MLP classifier improves classification accuracy over the linear classifier or KNN.
representation of population
Ablation studies (Table Figure 3) show that using a permutation-invanee summary of population dynamics is critical to improving NeuPRINT's performance on the downstream subclass prediction.
7

---

1.0
without population
1.0
Top-1 Subclass Accuracy
0.8
with population
0.756
Top-1 Class Accuracy
0.8
0,7910.807
0.664 0.680
0,569
0.600.6088
0.616
0.6
0.526
0.537
0.6
0.4
0.586
0.4
0.2
0.2
0.0
0.0
LOLCAT
Transformer Transformer NeuPRINT
LOLCAT
Transformer Transformer NeuPRINT
ISI
Raw
Raw
Figure 3: Accuracy of transcriptomic subclass and class prediction of NeuPRINT and baselines on single-mouse spontaneous activity recordings, with and without inputs from population statistics.
1.4
1.c
Ablation Studies
Moon Squored Enor
1.0
Top-1 Subciass accuracy
Likellhood
1.2
C.8
Top-1 Subriaes accurary
Top Subclass accuracy
3.8
0.8
0.6
0.6
0.0
0.4
Negafive
0.1
0.6
MSE
NLL
subcless
02
class ace
Linear Nonlinear Recurrent Transformer
Linear Nonlinear Recurrent Transformer
0.0
00
populatton
krunnnng speed
apapli
+tram
apspulatio
centee
surround
we
state
actuity
activty
Figure 4: Left: Top-1 accuracy (or loss) of learned representations with different dynamical models (linear, nonlincar, recurrent, transformer) in the subclass prediction task. Right: Ablation studies to dissect the impact of the different components of the crrmutation-ivvariann summary of population dynamics including running speed, pupil size, frame state, population activity, center-surround activity in improving the accuracy, in Table One component is added time from left right.
On the other hand, this effect not significant in the class prediction task, suggesting that the intrinsic electrophysiology of neurons is significantly different between the excitatory vs inhibitory classes.
5.2 Learning Representations Across a Spectrum of Implicit Dynamical Models
We next investigate learning representations using our NeuPRINT framework over a spectrum of implicit dynamical models ranging from simple linear and nonlinear dynamical models to more advanced deep learning architectures such as gated recurrent networks and transformers. We also evaluate the performance of the models under two different objective functions (mean squared error vs Gaussian negative log likelihood). The results are shown in Figure 4 We find that the transformer, which leverages the powerful attention mechanism (28 to preserve the information
a large temporal context, achieves the best performance in predicting the masked (future) neural activities, as quantified by NLL and MSE loss. This ability to capture the dynamics more faithfully explains the transformer's superior performance inferring neuron identity since neuronal physiology correlates with molecular expression (10, 29]. We note that, for this relatively small dataset, the MSE loss performs better than the NLL. loss based on the downstream classification accuracy.
5.3 crrmutaaton-Invariann Summary of Population Dynamics Enhances the Time-invariant Representation of Individual Neurons
To further investigate the role of each component in the cermutationiivvaiiat summary of population dynamics, we perform series of ablation studies as shown in Figure 4 We add each input (running speed, pupil size, frame state, population activity, and surrer- activity) to NeuPRINT one at time. We find that all of the proposed components contribute to the success of the time-invariant representations as evaluated by the downstream transcriptomic classification task. These results support the perspective put forth in Section 3 how the neuron reacts external inputs (hence summary variables proposed here) forms a part of the neuron identity. Overall, we find using all of the available cormutation-invaiii components the summary of the population recording improves the accuracy by 22% in subclass prediction.

---

Top Subelas Accuracy
1,0
Spontaneous
Drifting Gratings
0.566 9400
Natural Scenes
0.6
Natural Scanes
0.451 0:991
0.579 asst
0.502
0.544
0.4
0.2
0.0
Transformer
Transformer Raw
LOLCAT
NeuPRINT
1.0
0.807 0.809 8.836 0.836
Tap Class Accurary
0.000
0648
0.648
0.640
0.640
0.608
0.eon.
0.6
0.0
Transformer
Transformer
LOLCAT
NeuPRINT
Rew
Figure 5: Accuracy of transcriptomic subclass and class prediction of NeuPRINT and baselines on single-mouse recordings during spontaneous activity and visual stimuli-driven (drifting gratings, natural scenes) activity.
Learned
Time- Variant
Time- Invariant
Representation
Lower
Data Limited
Unsupervised
Self-supervised
Bound
Supervised
Individual
Inputs
Individual
Population
Trans Trans +ISI +Raw
NeuPRINT without population
NeuPRINT
Task
Classifier
Random
LOLCAT
PCA
UMAP
(ours)
KNN
0.174
0.333 0.348
0.419
0.552
Subclass
Linear
0.340
0.457
0.449 0.493 0.304
0.384
0.552
0.590
MLP
0.362
0.442 0.423 0.384
0.406
0.552
0,685
KNN
0.581
0.670
0.613
0.659
0.700
Class
Linear
0.667
0.700
0.710 0.675 0.645
0.660
0.746
0.746
MLP
0.660
0.702 0.707 0.682 0.667
0.770
0.800
Table 2: Extensions to multiple animals: Top-1 accuracy of transcriptomic label prediction based on (i) the representations learned by our proposed sulf-super representation learning from neural dynamics framework NeuPRINT, (ii) the supervised learning method LOLCAT and its variants Transformer+1SI and Transformer+ Raw, (iii) unsupervised baselines PCA and UMAP, (iv) random representations (to determine the chance-level). Note that this experiment corresponds limited regime due to limited labeled data. We performed classification using three classifiers (KNN, Linear, MLP) and tasks: predicting the cell class from the {excitatory, inhibitory) and predicting the subclass from the set (Lamp5, Pvalb, Vip, Sneg, Sst).
data-
5.4 Cell Type Identifiability Tends to Increase with Stimulus Relevance and Complexity
We further test our model NeuPRINT and other baselines on recordings with sets of visual stimuli (drifting gratings and two different natural scene image sets 6 The results summarized in Figure suggest an increase cell type identification accuracy of NeuPRINT from in-vivo activity as stimulus relevance and complexity increases, e.g., higher accuracy (~ 5%) in subclass prediction based on Natural Scenes compared to spontaneous activity recording.
5.5 Extensions to Multiple Animals
We further extend our evaluations from one animal to multiple animals. We report the evaluations for all animals on the extended dataset in Table While the performance of LOLCAT increases on class prediction task with this larger dataset, still remains less accurate than our model across all tasks and classifiers. For subclass prediction, where the distribution of cells across subclasses is highly imbalanced, our method NeuPRINT outperforms LOLCAT by ~ 23%, and also by ~ 10% in class prediction (See discussion the preliminary nature of this study below).

---

6 Discussion and Conclusion
In this work, we studied the problem of inferring neuronal identity from in vivo recordings. Motivated by the fact that in vivo physiology a neuron has two distinct components (the neuron itself, and the synaptic and modulatory imputs it receives), vectors for neurons based on a model of neuronal dynamics with exogenous inputs. While letting the activities of neighboring neurons represent the exogenous input directly could be a natural choice, this would prohibit the transferability of the model because a biologically meaningful ordering of neighboring neurons is not available in these recordings. This key observation led us to propose calculating commutatiio-invariaan statistics of the activity of the neighhoring neurons.
presented self-supervised approach to infer identity
In the field of natural language processing (NLP), learning the dynamics of the underlying system has been shown to enable success multiple downstream tasks [43]. Inspired by this, we proposed to use the identity vectors inferred by our dynamics model in predicting the molecular labels of the neurons. This was enabled by a recent, public dataset that profiles both in vivo calcium activity and subsequent transcriptomic expression in individual neurons [6]. Learning the dynamics is cuff-supervieed task. We trained simple classifiers with limited ground-trun truth labels to subsequently predict the molecular label, based the invariant representations inferred by this self-supervised step. We believe this two-step approach could constitute a benchmark for future studies on this topic.
We experimented with different dynamics models. We reported that, consistent with the impressive findings in NLP 28 433, the transformer architecture with its self-attention mechanism provides the best results among the architectures we considered. Due its feedforward structure, this architecture also enables fast inference, which could be useful for experimental setups with real-time feedback.
Ablation experiments demonstrated the merit of providing permutation- invarriant inputs and infor- mation on the overall state of the nervous system (e.g. pupil size). This set of experiments also revealed that more detailed input representations could further improve performance while retaining the desired coomuuation-invaiiann
property. particular, motivated by insights on cireuit motifs in neuronal networks 14) we demonstrated a center-surround setup, which partitions the available population into two concentric sets and calculates identical statistics for each set separately.
Our approach can be improved when applied to multi-region, cross-session and cross-animal record- ings by learning extra embeddings that represent different brain regions, sessions animals. In this sense, our multi-animal experiments represent a limited, preliminary study. It may also be possible to study recordings from healthy vs diseased brains or from individuals belonging to different species using a single model.
Our study represents initial attempt. Ever-growing neuroscience datasets represent straight- forward way to improve accuracy. On the technical side, engineering of the invariant features can extend the list of meaningful and invariant statistics of population activity. Contrastive learning of neuronal identity can improve classification accuracy by encouraging the model to depend more on the identity input. Our center-surround setup should also be considered an initial attempt at capturing existing neurobiological insights in the kind of dynamics models studied in this paper. Studying the transferability of the proposed model across individuals is high-priority direction for future research.
7 Reproducibility
All optimizations are performed on one NVIDIA Tesla V100 GPU. Details are included in Appendix E We released our software http:/tttw...commmmmmmmmaaaaaaat for reproducibility.
8 Acknowledgement
LM supported by the Shanahan Foundation Fellowship. LM and US thank the Allen Institute founder, Paul G Allen, for his vision, encouragement, and support. ES and TL acknowledge the support in part by A3D3 National Science Foundation grant OAC-21 17997 and the Department of Electrical Computer Engineering at the University of Washington. ES also acknowledges the support in part by the Washington Research Foundation Fund and the Department of Applied Mathematics at the University of Washington.
10

---

Broader Impact Since we studied basic neuroscience problem using recordings from mice, we do not expect this paper to have an immediate societal impact. We hope that, the long term. variants of our method will be helpful to better understand how the brain works and the physiological consequences of the various diseases of the brain. is nevertheless important keep in mind that technical progress, such as reported in our study, could potentially be adapted without much difficulty to malicious use.
References
[1] Shreya Saxena and John P Cunningham. Towards the neural population doctrine. opinion in neurobiology, 55:103-111, 2019.
Current
[2] David Marr. Vision: A computational investigation into the human representation and processing of visual information. MIT press, 2010.
[3] Hannah Bos, Anne-Marie Oswald, and Brent Doiron. Untangling stability and gain modulation in cortical circuits with multiple internurron classes. bioRxiv, pages 2020-06, 2020.
[4] Yuhan Helena Liu, Stephen Smith, Stefan Mihalas, Eric Shea-Brown, and Uygar Sümbül. Cell- type-specific neuromodulation guides synaptic credit assignment in a spiking neural network. Proceedings of the National Academy of Sciences, 1 118(51):e211 111221111111111111111 1821118, 2021.
[5] Anirban Paul, Megan Crow, Ricardo Raudales, Miao He. Jesse Gillis, and Z Josh Huang. Transcriptional architecture of synaptic communication delineates gabaergic neuron identity. Cell, 171 539, 2017.
[6] Stephane Bugeon, Joshua Duffield, Mario Dipoppa, Anne Ritoux, Isabelle Prankerd, Dimitris Nicoloutsopoulos. David Orme, Maxwell Shinn, Han Peng, Hamish Forrest, et al. tran- scriptomic axis predicts state modulation of cortical interneurons. Nature, 607(7918):330-333, 2022.
[7] Nathan W Gouwens, Staci A Sorensen, Fahimeh Baftizadeh, Agata Budzillo, Brian R Lee, Tim Jarsky, Lauren Alfiler, Katherine Baker, Eliza Barkan, Kyla Berry, et al. Integrated morphoelectric and transcriptomic classification of cortical gabaergic cells. Cell, 183(4):935- 953, 2020.
[8] Federico Scala, Dmitry Kobak, Matteo Bernabucci, Yves Bernaerts, Cathryn René Cadwell, Jesus Ramon Castro, Leonard Hartmanis, Xiaolong Jiang, Sophie Laturnus, Flanine Miranda, al. Phenotypic variation of transcriptomic cell types in mouse motor cortex. Nature, 598(7879):144-150, 2021.
[9] Rohan Gala, Nathan Gouwens, Zizhen Yao, Agata Budzillo, Osnat Penn, Bosiljka Tasic, Gabe Murphy, Hongkui Zeng, and Uygar Sümbül. A coupled autoencoder approach for multi-modal analysis of cell types. Advances in Neural Information Processing Systems, 32, 2019.
[10] Dmitry Kobak, Yves Bernaerts, Marissa A Weis, Federico Scala, Andreas S Tolias, and Philipp Berens. Sparse reduced-rank regression for exploratory visualisation of paired multivariate data. Journal of the Royal Statistical Society Seriey C: Applied Statistics, 70(4):980-1000 2021.
[11] Yeganeh Marghi, Rohan Gala, Fahimeh Baftizadeh, and Uygar Sumbul. Joint inference of discrete cell types and continuous type-specific variability in single-cell datasets with mmidas. bioRxiv, pages 2023-10, 2023.
[12] Jennifer S Lund, Alessandra Angelucci, and Paul C Bressloff. Anatomical substrates for functional columns in macaque monkey primary visual cortex. Cerebral cortex, 13(1), 2003.
[13] Robert Levy and Alex D Reyes. Spatial profile of excitatory and inhibitory synaptic con- nectivity in mouse primary auditory cortex. Journal of Neuroscience, 32(16):5609-55 5619, 2012.
[14] Robert Rosenbaum, Matthew Smith, Adam Kohn, Jonathan E Rubin, and Brent Doiron. The spatial structure of correlated neuronal variability. Nature neuroscience, 20(1):107-114, 2017. 2017.
11

---

[15] Byron M Yu. John P Cunningham, Gopal Santhanam, Stephen Ryu, Krishna V Shenoy, and Maneesh Sahani. coussian-process, factor analysis for low-dimensional single-trial analysis of neural population activity. Advances in neural information processing systems, 21, 2008.
[16] Yuanjum Gao, Evan W Archer, Liam Paninski, and John P Cunningham. Linear dynamical neural population models through nonlinear embeddings. Advances in neural information processing systems, 29, 2016.
[17] Yuan Zhao and Il Memming Park. Variational latent gaussian process for recovering single-trial dynamics from population spike trains. Neural computation, 2965):1293-131,, 2017.
[18] Scott Linderman, Matthew Johnson, Andrew Miller, Ryan Adams, David Blei, and Liam Paninski. Bayesian learning and inference in recurrent switching linear dynamical systems. In Artificial Intelligence and Statistics, pages 914 914-922. PMLR, 2017.
[19] Daniel Durstewitz. A state space approach for piecewise-linear recurrent neural networks for identifying computational dynamics from neural measurements. PLoS computational biology, 1306))e 2017.
[20] Felix Pei, Joel Ye, David Zoltowski, Anqi Wu, Raeed Chowdhury, Hansem Sohn, Joseph O' O'Doherty, Krishna V Shenoy, Matthew T Kaufman, Mark Churchland, et al. Neural latents henchmark 21: evaluating latent variable models of neural population activity. arXiv preprint wwXiv:210900446. 2021.
[21] Chethan Pandarinath, Daniel J O' O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey 1) Stavisky, Jonathan C Kao, Eric M Trautmann, Matthew T Kaufman, Stephen Ryu, Leigh R Hochberg, et al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(00):05-811, 2018.
[22] Mohammad Reza Keshtkaran, Andrew R Sedler, Raeed Chowdhury, Raghav Tandon, Diya Basrai, Sarah L Nguyen, Hansem Sohn, Mehrdad Jazayeri, Lee E Miller, and Chethan Pandari- nath. A large-scale neural network training framework for generalized estimation of single-trial population dynamics. Nature Methods, pages 1-6, 2022.
[23] Steffen Schneider, Jin Hwa Lee, and Mackenzie Weygandt Mathis. Learnable latent embeddings for joint behavioral and neural analysis. arXiv preprint 2tt::2200400677. 2022.
[24] Tze Hui Koh, William E Bishop, Takashi Kawashima, Brian B Jeon, Ranjani Srinivasan, Yu Mu, Ziqiang Wei, Sandra Kuhlman, Misha B Ahrens, Steven M Chase, et al. Dimensionality reduction of -mmggdd- neuronal population activity. Nature Computational Science, 3(1):71 15,,71-85,222 2023.
[25] Wilson Truccolo, Uri T Eden, Matthew R Fellows, John P Donoghue, and Emery N Brown. A point process framework for relating neural spiking activity to spiking history, neural ensemble, and extrinsic covariate effects. Journal of neurphyysoloogy. 93(2):1074-1085 2005.
[26] Jonathan W Pillow, Jonathon Shlens, Liam Paninski, Alexander Sher, Alan M Litke, EJ Chichilnisky, and Eero P Simoncelli. Spatio-temporal correlations and visual signalling in a complete neuronal population. Nature, 454(7207):995-99) 999, 2008
[27] Ran Liu, Mehdi Azabou, Max Dabagia, Jingyun Xiao, and Eva Dyer. Seeing the forest and the tree: Building representations of both individual and collective dynamics with transformers. Advances in Neural Information Processing Systems, 35:2377-236 2022.
[28) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neurul information processing systems, 30, 2017.
[29] Rohan Gala, Agata Budzillo, Fahimeh Baftizadeh, Jeremy Miller, Nathan Gouwens, Anton Arkhipov, Gabe Murphy, Bosiljka Tasic, Hongkui Zeng, Michael Hawrylycz, et al. Consistent cross-modal identification of cortical neurons with coupled autoencoders. Nature computational science, 1(2): 122,002222200 2021.
12

---

[30] Yves Bernaerts, Michael Deistler, Pedro J Goncalves, Jonas Beck, Marcel Stimberg. Federico Scala, Andreas Tolias, Jakob Macke, Dmitry Kobak, and Philipp Berens. Combined statistical-mechanstiii modeling links ion channel genes to physiology of cortical neuron types. bioRviv, pages 2023-03, 2023.
[31] Aidan Schneider, Mehdi Azabou, Louis McDougall Vigier, David F Parks, Sahara Ensley, Kiran Bhaskaran-Nair, Tomasz. Nowakowski, Eva Dyer, and Keith Hengen. Transcriptomic cell type structures in vivo neuronal activity across multiple timescales. Cell Reports, 42(4), 2023.
[32] Tsungnan Lin, Bill Horne, Peter Tino, and Lee Giles. Learning long-term dependencies in narx recurrent neural networks. IEEE Transactions Neural Networks, 7(6):1329-1333) 1996.
[33] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 1735:1735-770,, 1997.
[34] Joel Ye and Chethan Pandarinath. Representation learning for neural population activity with neural data transformers, Neurons, Behavior, Data analysis, and Theory, 5(3): 2021.
[35] Trung Le and Eli Shlizerman. Stndt: Modeling neural population activity with spatiotemporal transformers. In Advances in Neural Information Processing Systems, 2022.
[36] Alan L Hodgkin and Andrew F Huxley. A quantitative description of membrane current and its application conduction and excitation in nerve. The Journal of physiology, 117(4):500, 1952.
[37] Wulfram Gerstner, Werner M Kistler, Richard Naud, and Liam Paninski. Neuronal dynamics: From single neurons networks and models of cognition. Cambridge University Press, 2014.
[38] Christine Grienberger and Arthur Konnerth. Imaging calcium in neurons. Neuron, 73(5):862- 885, 2012.
[39] Eric Jonas and Konrad Kording. Automatic discovery of cell types and microcircuitry from neural connectomics. Elife, 4:e04250, 2015.
[40] Jakub M Tomezak, Maximilian llse, and Max Welling. Deep learning with permutation-invarian operator for multi-instance histopathology classification. arXiv preprint 17 2017.
[41] Stephen W Kuffler. Discharge patterns and functional organization of mammalian retina. Journal of neurophysiology. 16(1):37-68, 1953.
[42] Laurent Itti, Christof Koch, and Ernst Niebur. A model of saliency-based visual attention for rapid scene analysis. IEEE Transactions on pattern analysis and machine intelligence, 20(11):1254-1255 1998.
[43] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 33:1877-1901, 2020.
[44] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances Neural Information Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019.
[45] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 6982,690,, 2014.
[46] Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and Duchesnay. Scikit-learn: Machine learning Python, Journal of Machine Learning Research, 12:2825-2834 201
13

---

[47] Jacob Devlin, Ming- Ling-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 04805 2018.
[48] Marius Pachitariu, Carsen Stringer, Sylvia Schröder, Mario Dipoppa, L Federico Rossi, Matteo Carandini, and Kenneth D Harris. Suite2p: beyond 10,000 neurons with standard two-photon microscopy. BioRxiv, page 061507, 2016.
[49] Leland Mclnnes, John Healy, Nathaniel Saul, and Lukas Großberger. Umap: Uniform manifold approximation and projection. Journal of Open Source Software, 3(29):861, 2018.
14

---

Learned
Time Variant
Time Inve-anvaiiant
Representation Lower
Data-Limited Supervised
Unsupervised
superviee
Bound
Individual
Inputs
Individual
Population
Trans +SI
Trans +Raw
NeuPRINT without population
NeuPRINT
Task
Classifier
Random LOLCAT
PCA UMAP
(ours)
0.196
0.343
0.371
0.424
0.546
KNN
0.041
1 0.0.47 1 0.070
0.051
0.056
0.326
0.449
0.414
0,436
0.316
0.386
0.595
0.683
Subclass
Linear
+ 0.055
+ 0.027
= 0.036
0.021
0.046 0.046
= 0.041
0.042
0.382
0.456
0.422
0.392
0.496
0.606
0.722
MLP
0.062
= 0.077
0.054
0.022
0.509
0.569
0.547
0.568
0.705
KNN
Class
+ 0.061
100015
- 0.057
0.017
0.536
0.616
0.632
0.652
0.584 0.583
0.741
0.766
Linear
0.035
+0.017
0.026
0.035
0.031
0.565
0.624
0.651
0.596
0.584
0.782
0.803
MLP
+ 0.030
0.010 0.007 + 0.015 0.013
+ 0.018
0.037
Table 3: Sensitivity analysis across runs with different random seeds: Top-1 accuracy (mean-tstandard deviation) of transcriptomic label prediction based on (i) the representations learned by our proposed self-supervised representation learning from neural dynamics framework NeuPRINT, (ii) the supervised learning method LOLCAT and its variants Transformer+ISI and Transformer+Raw (abbreviated Trans+ISI and Trans+Raw, respectively) (Note that this experiment corresponds to data-limited regime due the size of the dataset), (iii) unsupervised baselines PCA and UMAP, (iv) random representations (to determine the chance-level). We performed classification using three different simple classifiers (KNN, Linear, MLP) and two tasks: predict- ing the cell class from the set {excitatory, inhibitory] } and predicting the subelass from the set (Lamp5, Pvalb, Vip, Sneg, Sst).
A Sensitivity Analysis
To quantify the sensitivity of our methods and baselines the transcriptomic identity (neuron class and subclass) prediction benchmark, we run each model with 5 different random seeds (train/val/test random split, random initializations, etc.). We report their averaged Top-1 accuracies and the corresponding standard deviations across 5 runs for all models in Table
Consistent with the main text, we find NeuPRINT still significantly outperforms other methods including supervised baselines LOLCAT [31 and its variants Transformer+ISI and Transformer+Raw (Note the limited availability of labeled data). The standard deviations for all methods are relatively small, indicating the robustness our evaluation framework and the significance the performance gaps.
B A Spectrum of Implicit Dynamical Models
explore spectrum of implicit dynamical models Linear, Nonlinear, Recurrent neural network, Transformer, and two objective functions mean squared error (MSE) when the model only predicts the mean of the output distribution; negative log likelihood (NI.L) when the model predicts both mean and standard deviation of output distribution to train the models. The results are summarized in Table
C Roles of Permutation-Invarian Summary of Population Dynamics
To investigate the role of each component in the commutation-invariinn summary of population dynamics, we perform a series of ablation studies as shown in Table We add each input (running speed, pupil size, frame state, population activity, and center-surround activity) to NeuPRINT one at a time. We find that all the proposed components contribute to the success of the time-invariant representations as evaluated by the downstream transcriptomic classification task.
15

---

Task
Linear
Nonlinear
Recurrent
Transformer
Objective
NLL
MSE
NLL
MSE
NLL
MSE
NLL
MSE
Loss
1,422
0.984
1,412
1.022
1.231
0.887
0.308
0.221
Subclass
0.488
0.512
0.337
0.512
0.610
0.610
0.707
0.756
Class
0.706
0.716
0.724
0.679
0.743
0.743
0.752
0.807
Table 4: Comparison of reconstruction loss (first row, the smaller the better) and top-1 accuracy of class and subclass prediction using MLP downstream classifier (second and third rows, the larger the better) achieved by different dynamics models f (linear, nonlinear, RNN, and transformer) in NeuPRINT. Two objectives used to train each model (mean squared error (MSE) negative log likelihood (NLL)).
w/o
+
Subelass
population inputs
running speed
pupil
frame state
population activity
center-surround activity
size
MLP
0.512
0.609
0.658
0.683
0.732
0.756
Table Ablation studies of permutation-invarian inputs representing population activity, including running speed, pupil size, frame state, permutaaionninvariii population representation, permutation- invariant center-surround representation. Results are reported on the subclass prediction task with an MLP downstream classifier.
D Extensions to Other Visual Stimulus Conditions
We further test our model NeuPRINT and other baselines recordings with 3 sets of visual stimuli (drifting gratings and two different natural scenes). The results Table o suggest an increase cell type identification accuracy from iivivo activity stimulus relevance and complexity increases.
E
Implementation Details
E.1 NeuPRINT
Permutation-invarian summary of population dynamics: We use C, S to represent D, D dimensional permutation- and size-invariant (i.e., N) surrogates for brain activity each time point. C' and S compute identical statistics (mean and standard deviation) of population activity except that the statistics C' are calculated over the center partition (neurons whose distance to neuron is at most A) and the statistics in S are calculated over the surround partition (neurons whose distance to neuron is larger than A). Here, A hyperparameter:
Po(-o = mean (x(0))x00<d
(4)
c(-- std(X,(t)) dw <))
(5)
Ms(-n mean [x(t)]j:00000 A dd))
(6)
s(--1) = std(X,()))::0
< dj:)
Batch sampling: For each batch of data we randomly sample 512 labeled and unlabeled neurons to be included in the batch. For each neuron further randomly sample 2 sub-windde wab-windows of 512 timesteps from its continuous calcium flourescene traces. Each resulting sample for the 7th neuron is denoted (*).
We use the transformer encoder architecture 28 and the masking strategy as originally proposed in [47]. Random timesteps in X (4) are masked (zero-out) with probability population summary P and time-invariant
Multihead attention:
0.25 and concatenated with the permutationninnaaiit representation (6) along the feature dimension form input ( Note that the same learnable x60) repeated every timestep, enforcing
Similar to [47), we embed input X (4) and
16

---

Task
Classifier linear
LOLCAT
Transforer+ISS
Transformer+Raw
NeuPRINT
0.474
0.491
0.474
0.683
Subclass
Spontaneous Activity
mlp
0.561
0,439
0.756
linear
0.608
0,680
0,669
0.793
Class
mlp
0,640
0.664
0.807
lincar
0.421
0.509
0.404
0.756
Subclass
Drifting Gratings
mlp
0.474
0.509
0.780
lincar
0.600
0,640
0.648
0.809
Class
mlp
0,632
0.640
0.809
linear
0.491
0.579
0.421
0.780
Natural Scenes
Subclass
mlp
0.544
0.439
0.805
linear
0.648
0.672
0.696
0.809
Class
mlp
0,664
0,640
0.836
linear
0.491
0.509
0.544
0.732
Natural
Subclass
mlp
0.561
0.526
0.782
Scenes
linear
0.648
0.664
0.656
0.809
2
Class
mlp
0.664
0.648
0.836
Table 6: Accuracy of transcriptomic subelass and class prediction of NeuPRINT and baselines on
single-mouse recordings during spontaneous activity and visual stimuli-driven (drifting gratings, natural scenes) activity.
employ sinusoidal positional embedding to encode the temporal order in the input sequence, resulting in (+) = Emb(X ()
For each input {(), a set of weights Wº R2'xd,, W E Rxxx, W R''d, are learned to transform input (0) to a set of query, key, and value (Q, K, V), where Q x K
W WVV, Attention between temporal tokens for one attention head is computed as:
QK
Attention(Q (Q, V) softmax
8
ddz
Each head will find a different pattern in the data and produce an output of size dg The final attention output will be concatenation of these single-head outputs. We use 2 heads in our model.
Feedforward layers and residual connections are subsequently applied attention output:
7 (0) MSA(X (0) FF + ())(X(())
where MSA represents the multihead attention operation, FF represents the feedforward layer with ReLU activation, and ( represents the reconstructed calcium trace with masked timesteps recov- ered.
Computational cost: NcuPRINT has 833K parameters and takes 2 hours in total for training and inference on a single NVIDIA Tesla V100 GPU, Details on the size of the datasets are mentioned in Section 4. of the main paper.
E.2 LOLCAT and Its Variants
We use the publicly available implementation of LOLCAT from 311. We implement two additional variants of LOLCAT (Transformer+ISI and Transformer+Raw) following the philosophy in [31 but with the self-attention mechanism as used by NeuPRINT 28). LOLCAT and Transformer+ISI continuous time series of calcium traces are divided into non-overlapping sub-windows of size 64, within which the Spike Interval (ISI) distribution is computed. We use suite2p Python package 48 to infer spikes and compute the ISI distribution with 16 bins, using a spike threshold of 0.2. To perform classification using the self-attention mechanism, we use a learnable special classification token [CLS] appended to the beginning of the input sequence represent the classification output 47]. No positional embedding added to the input sequence. Therefore the transformer output at the CLS token position will represent the pooling operation in [3I] and all sub- aub-windows treated as if they are independent trials, which enables Transformer+ISI to apply to any number of observed trials important design choice for LOLCAT.
further implement TransformrrRRww variant of LOLCAT where the transformer operates directly the raw calcium traces rather than the ISI distribution of calcium traces. Transformer+Rav
17

---

Transformer +ISI
Transformer +Raw
random/
NeuPRINT
LOLCAT
PCA/UMAP
Representation dim
64
64
Encoder hidden dim
70
[32, 16, 16]
[128] x 2048
[28] 4 2048
MIP classifier hidden dim
2048
2048
Number of attention heads
2
4
4
4
Number of attention layers
1
1
1 512
Window size
200
2048
512
2048
Context window size
2
8
1
Batch size
1024
varies
varies 500
varies
Number of epochs
400
500
500
Number of downstream epochs
5000
1000
Learning rate
10
Downstream learning rate
10-4
10-4
10-1
10-4
Dropout
0.1
0.1
0.1
KNN neighbors
5
5
5
5
ISI sub-window size Number of ISI bins
64
64
16
16
Table Hyperparameters of our self-supervised representation framework NeuPRINT and other baselines including end-to-end supervised learning model LOLCAT, its variants Transformer+1S1 and Transformer+ Raw, unsupervised representation learning models PCA and UMAP, and chance-level prediction based on random features.
follows the same architecture as Transformer+IS1, except that now the positional embedding is added the input sequence to denote the temporal relationship between timesteps in the trial window.
E.3 Random/ PCA UMAP
To probe the chance-level classification performance, train and evaluate downstream classifiers using random vectors size 64 as representations for individual neurons, equivalent to the 64-dim (0)), To compare NeuPRINT with unsupervised methods PCA and UMAP, we first project the data to lower dimensional space using 64 components, then train and evaluate downstream classifiers on the low-dimensional represcntation..
E.4 Downstream Classifiers
We 5 nearest neighbors for KNN downstream classifier. For the downstream MLP classifier, use a multi-layer perceptron network with a single hidden layer of size 2048 and ReL.U activation.
E.5 Hyperparameters
We include all the important hyperparameters (representation dim, window size, number of epochs, learning rate, batch size, etc.) for our NeuPRINT and other models (supervised-leannin) baselines LOLCAT, Transformer+IST and Transformer+ Raw, unsupervised representation learning baselines UMAP and PCA, chance-level prediction based random features) in Table 7
F
Pseudo Code
Our NeuPRINT framework includes three main components: implicit dynamical system that uses the state-of-the-art transformer architecture to model neural dynamics; an optimization framework that fits the dynamical model and learns time-invariant representations for neurons; supervised learning framework to train the downstream classifiers for subclass and class prediction, taking the learned invarinvar representations as inputs. The pseudo code for these three components is listed as follows:
18

---

ceneformmrccaccuun fluoroscence trace..
permutation iitaaiiniiaaaaaaaaaa invariant_populltioon _amooon__smmrr..
input_embedder = linear (input_dim, hidden_dim)
transformer_encoder multihead attention
window_side
layer_dim,
num hum_heads)
output_decoder linear (hidden dim. output_dim)
f mask_inputs (calliiu hacpppttccccccimmmoooooseeeeenn (lr
masked_calcium
input concatenate(
masked_calcium calcium htw.oa/acoooooeeeea_______ ttreeeetrr..____e_-eeeeaeee_eeee
masked
permutation gemnuaaiooiinaiiaapppppuuiisssss
tine_avariancccppesnttttoo))
input positional_encoding (input)
input inpu_eebeedder(nput)))
context mask
_onptettecceeetttteeeeee mask (
window_skt size,
context_window_size
output transformer_encoderr (input, context_maski
output cutput_decoder(autt
return output
Time_Invariante Self Selp-inrtdd Stt.:sSuppriiieeeeeereennaatoooo
calcium calcium_fluorosee
time invariant cenerraattreprreentaaiiooo zeros (neuron dim, embedding _im)
recon model occon_model
Transformer(
input_dim. hidden_dim.
window_sizz
hontextwwindow.cm
jayer_dim.
num_heads)
optimizer = Adam
ane__ivvrrintt_epreeee
recon_ model.
lcarning_rate)
recon_model (
calcium Sacccimmffuuoooscenect trace
permutation http:attoonniiarriaaapopolllllia summarry
iine_iivariantceppresnttttooo
recon_loss mse(
seddicted_caciuumm.mm ffaorcicmm eetrrceeeee_reeeeeeeee trace ]
Aalcium_fluorosceenee trace masked_steps
recon_loss backward ()
19

---

optimizer. step
return time iine_iiverrantrrpreeenntaaooo
Downstream Dowasteea..Caassirrreppprrrsssss
rre_invariant__ppreeennte
classifier dip dim. hidden_dim. output_dim)
optimizer Adam( classifier. learning_rate)
predicted_class_labee classificr ( (fine_ivvarinnnrgrreetttttoonn
classification__oss predicted class _la_ssslaasslll___________a_ label
cross_entropy(
ground_truti toant_truth___ cllas fatthhhllsss_aaall_lll
classification_loss backward ()
optimizer.ste ()
return predicted_class label
20


Reviews:
[{"review_id": "z31Q2xdYmB", "reviewer": "Reviewer_QXoV", "summary": "The authors evaluate several approaches for identifying neural types and subtypes (e.g., GABAergic cells are commonly divided into several subclasses Lamp5, Pvalb, Vip, Sncg, Sst). They used a dataset from a recently (2022) published neuroscience paper and compared their approaches to a recent study that had the same objective. To identify neural types, they use several 2-photon calcium imaging recordings (that provide a proxy for neural activity) from mice V1 and behavioral information (running speed and pupil size). They used mRNA expression for 72 genes to label the neurons according to the class (this was then used to compute the accuracy). While performance for class accuracy was relatively similar for all methods, the authors concluded that an approach based on the \"Attention is all you need\" transformer led to the best results classification results for subclass prediction. They evaluated the contribution of different factors of the model (e.g., running speed and pupil size) and concluded that each contributed to the classification accuracy. ", "questions": "While the authors compared several approaches and concluded that the one based on transformers led to the best results, I wonder if they conducted hyperparameter sweeps to evaluate that the difference in the results between different methods was not because of different choices of a learning rate. \n\nWhile table S3 was very useful for examining the hyperparameters that were used, it would be helpful to provide explicitly the number of trainable parameters used in each approach. \n\nThe authors used transformer architecture from \"Attention is all you need\" paper. Would more recent transformed architectures lead to better performance?\n\nWas there a reason to use only 2 attention heads? \n\nWhat was the robustness of the approach? For example, if you run the code twice with different seeds, how many cells would change the class/subclass? \n\nDid different methods provide somewhat complementary results (e.g. were some neurons correctly labeled with recurrent model and incorrectly with transformer)? If yes, would it be possible to use some form of ensemble learning to get an even better performance?\n\n", "limitations": "The paper addresses limitations (although that could be expanded) and has a broader impact section. ", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The authors compared several approaches and provided detailed information about parameters in the supplemental material. They conducted an ablation study by evaluating the contribution of different pieces of data. The results are presented in a clear form, illustrating the overall accuracy as well as confusion matrices. The proposed approach led to better accuracy than in recent (2023) previous work that used the same dataset. ", "weaknesses": "The evaluation was limited to a single dataset.\n\nWhile, as the authors state in the discussion, this is an initial attempt to tackle this problem, it remains unclear whether the results would substantially improve if the authors did a hyperparameter sweep or used a more powerful transformer architecture. \n\n\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive review.\n\n**\"The evaluation was limited to a single dataset.\"**\n\nWe believe our general response (“Uniqueness of the dataset and study” and “Robustness of the method”) addresses the key points identified as weaknesses by this reviewer. Please refer to that response for clarifications and results of new experiments. Briefly:\nIn that response, we argue that the Bugeon et al dataset is the only publicly available dataset that is suitable for the problem posed in this paper. To address the broader point, we will edit the manuscript to summarize the discussion above.\n\nTo address the reviewer’s concern on limited evaluation, we have performed multiple new experiments. Some of these use parts of the Bugeon et al dataset that we have not used in the initial submission, thus representing new evaluations, including testing on a previously unseen animal, testing on recordings with 3 sets of stimuli (drifting gratings and 2 different natural scenes). Please see the tables in our overall response and the attached pdf file. All of these new experiments support our original findings, and will be added to the revised manuscript.\n\n**“It remains unclear whether the results would substantially improve if the authors did a hyperparameter sweep or used a more powerful transformer architecture.”**\n\nWe have followed the reviewer’s suggestions to tune a set of hyperparameters that are critical for improving the models’ performance, please refer to new results Fig R1 and Table R1 in the attached pdf. Using a new set of hyperparameters (i.e. number of attention heads: 8,  embedding dim: 128, learning rate: $10^{-3}$) for our model NeuPRINT improved the accuracy substantially (subclass accuracy from 0.732 to 0.756, class accuracy from 0.752 to 0.807). At the same time, we also tried to improve our baseline LOLCAT by replacing its original simple attention structure with a more powerful self-attention mechanism from transformer, changing the linear classifier to multi-layer perceptron (MLP) classifier, and testing different versions of inputs: spike interval distribution (ISI) (refer to ‘Transformer+ISI’) v.s. raw fluorescence trace (refer to ‘Transformer+Raw’). We also performed additional hyperparameter tuning for these baselines. Subsequently, the subclass accuracy improved from 0.474 to 0.561, and class accuracy from 0.608 to 0.680, which are still worse than the performance of NeuPRINT. \n\n**“Number of trainable parameters used in each approach”**\n\nThe numbers of trainable parameters used in each approach are as follows: LOLCAT: 8.5K, Transformer+ISI: 1.24M, Transformer+Raw: 1.19M, NeuPRINT: 834K.\n\n**“What was the robustness of the approach?”**\n\nWe have conducted a sensitivity analysis in the supplement (Table S1), where the variance of using different random seeds seems relatively small. And we found our approach could also be generalized from a single animal to multiple animals as shown in Supp. Table S2. As mentioned in our general response, we also reported our methods with multiple new comparisons (see attached pdf), and established the robustness of NeuPRINT under different (and practically relevant) experimental setups, including testing on recordings with 3 sets of stimuli (drifting gratings and 2 different natural scenes). Furthermore, we also did a new cross-animal test where NeuPRINT and the downstream classifiers were trained with data from 3 animals. Any data from one animal was held out to be used for testing. At test time, we kept the dynamical model and downstream classifier fixed and only inferred the time-invariant representation for the unseen animal, and directly tested the accuracy for downstream tasks. Interestingly, NeuPRINT generalizes well in this setting, achieving subclass prediction accuracy of 0.586 and class prediction accuracy of 0.804. These new experiments support our original findings. We will incorporate them into the revision.\n\n**“If you run the code twice with different seeds, how many cells would change the class/subclass?”**\n\nTo address the reviewer’s concern, we add the following set of results on the average proportion of cells changing their assignments across runs to the revision: for two NeuPRINT transformers using the same hyperparameters but different random seeds (with class accuracies 0.814 and 0.826), 16.7% of their predictions are not matched: model 1 is incorrect, model 2 is correct (8.9%), model 1 is incorrect, model 2 is correct (7.7%). For subclass prediction, their accuracies are 0.732 and 0.746, and there is 27.8% mismatch in their prediction: model 1 is incorrect, model 2 is correct (14.1%), vice versa (12.7%). \n\n**“Did different methods provide somewhat complementary results?”**\n\nThe reviewer raises an interesting question on the potential to use ensemble learning. We agree that this is potentially a fruitful approach. As a starting point, we performed a mismatch analysis for the results from the recurrent network and the transformer. Particularly for subclass prediction, when the transformer has an accuracy of 0.731, and the RNN has an accuracy of 0.571, there is a large mismatch (54.6%), suggesting there might be benefits to performing ensemble learning. We will add an ensemble study to the revision, where we will train a classifier based on concatenating NeuPRINT representations from both transformer and recurrent network.\n"}, {"review_id": "hSj1CFSblI", "reviewer": "Reviewer_5dAW", "summary": "Understanding neural dynamics is crucial for studying brain function, but most existing methods assume that all neurons are functionally identical, which is not the case at the biological level. The biological variability among neurons can influence their dynamics, and since this variability is genetically encoded, the dynamics associated with neuron types would be time-invariant. This work introduces a method that uses transformers to learn a time-invariant representation of neural dynamics. The learned representation can reconstruct neural activity more accurately than existing methods when applied to neural recordings in mice. The authors also demonstrate that the genetic identity of the cell can be captured more robustly from this representation compared to alternative methods and lesion versions of the model.", "questions": "Can the approach be directly used on other kinds of neural recordings like electrophysiological data?", "limitations": "The authors should include a section or clear description of the limitations and assumptions of the method. This section would also benefit from including information on the computational cost, training time, and data demands associated with the method.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is adequately written and technically good. The method was tested and shown to work well when applied to calcium neural recordings in mice. The introduced method utilizes transformers to identify time-invariant dynamics associated with neural responses and relates them to the genetic heterogeneity observed in neural populations. Understanding the interplay between genetics and other factors, such as learning or external stimuli, in their relationship to neural dynamics is critical for studying neural function, and this work successfully extracts specific dynamics that can help predict neural type.", "weaknesses": "While the authors demonstrate how their method can extract neural dynamics related to neuron type, it is important to note that the dataset used only contains spontaneous neural activity. Neural dynamics are typically associated with specific computations and not just baseline states. Therefore, extracting time-invariant or baseline dynamics is relatively straightforward, and the reconstruction analysis may have limited significance. A more interesting application of these methods would be to disentangle dynamics associated with time-invariant or genetic factors from those associated with task variables when animals are engaged in a task or exposed to stimuli.\nIn terms of identifying neural type in the supervised learning task, the authors compare the classification performance of their method to other methods using neural dynamics. However, they overlook alternative approaches based on multiple neural properties such as firing rates, autocorrelograms, or morphology. If the goal of the introduced method is to label subpopulations or neurons, a direct comparison with commonly used approaches to putatively label neurons would be necessary to fully assess the significance of the method. Moreover, in the context of calcium imaging, the expression of the calcium indicator is typically given by genetic markers, so the type or types of neurons being monitored are known. Incorporating this explicit knowledge of neuron type when extracting dynamics could improve the accuracy of the method.\nThe accuracy results demonstrate that the proposed method can identify neuron types better than alternative methods. However, the overall accuracy is relatively low and fails to identify less prevalent neuron types. This suggests that the practical utility of the method for neuroscience research may be limited. Providing information on the computational and data demands of the method, as well as its applicability to other datasets, would be important to assess its relevance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive review.\n\n**“the dataset used only contains spontaneous neural activity”**\n\nWe believe this is not a limitation of the proposed method. Our method could be directly applied to neural recordings under different conditions, stimuli, and tasks. To address the reviewer’s concern on limited evaluation on spontaneous activity recordings only, we have performed multiple new experiments. We test our model based on parts of the Bugeon et al dataset that we have not used in the initial submission, thus representing completely new evaluations, including testing on recordings with 3 sets of visual stimuli (drifting gratings and two different natural scenes). Please see the tables in our overall response, Table R2, and Figure R2 in the attached pdf file. All of these new experiments still support our original findings that our model outperforms other baselines. Interestingly, we also found the predictive performance of NeuPRINT improved when more complex stimuli were applied (i.e. subclass accuracy improved from 0.756 to 0.805, and class accuracy improved from 0.807 to 0.836 when switching from spontaneous activity to natural scenes), which indicates that intricate, identity-specific dynamics can become more prominent when the underlying computation becomes more complicated. We will add the new results and discussions to the revised manuscript. \n\n**“overlook alternative approaches based on multiple neural properties such as firing rates, autocorrelograms, or morphology”**\n\nAlternative approaches based on activity (e.g., firing rates) are studied in detail in the LOLCAT paper and LOLCAT was found to be superior to those approaches. Since NeuPRINT outperforms LOLCAT. Therefore, it should outperform those alternative approaches. We thank the reviewer for this suggestion. We will revise the manuscript to mention such alternative approaches and cite the LOLCAT paper to suggest that those alternative approaches are not as powerful in in-vivo settings.\n\nOn the other hand, cell type identification based on morphology is a very different problem. The reason is that there are no methods that record both in-vivo activity and neuronal morphology. Nascent studies combining calcium imaging with electron microscopy are upcoming, but these are limited by the much smaller field of view of electron microscopy. Therefore, while being the classical approach to identifying neuronal types, morphological analysis is not (yet) relevant for the purposes of studying in-vivo population activity in the context of neuronal identities. To address the reviewer’s concern, we will edit the Introduction to provide a morphological and historical perspective on neuronal cell types.\n\n**“the expression of the calcium indicator is typically given by genetic markers, so the type or types of neurons being monitored are known”**\n\nWe believe this represents a key confusion and we hope that our overall response addresses this. Briefly, as the reviewer correctly mentions, the typical setup includes markers for, say, Pvalb neurons. Therefore, *only* Pvalb neurons are visible under the microscope. The experimenter already knows the cell type information, so there is nothing to infer on the recorded cells. But, what is the relationship between the activities of these Pvalb neurons and the other neurons? Those other neurons were never visible in such a typical setup and therefore cannot be studied.\n\nIf, instead, a pan-neuronal marker is used, then potentially neurons of all types are visible under the microscope. Now, however, the experimenter does not know anything about cellular identity. This is the problem we address with NeuPRINT, with the help of emerging datasets that use pan-neuronal markers. Once inference accuracy becomes acceptable, post-hoc spatial transcriptomics will no longer be necessary (e.g., this will be recovered computationally by NeuPRINT).\n\t\nIf the researcher is interested in, for instance, only the interactions between excitatory neurons in the mouse cortex, then the same argument can be repeated for pan-Glutamatergic markers. Conceptually, this will be identical to the problem posed here and NeuPRINT should provide a compelling solution as is.\n\n**“However, the overall accuracy is relatively low and fails to identify less prevalent neuron types”**\n\nFirst, our updated results show higher accuracy with more hyperparameter tuning, and under more complex stimuli (natural scenes) i.e. from 0.732 to 0.805 for subclass prediction, from 0.752 to 0.836 for class prediction. \n\nWe agree that the existing accuracy values might not be high enough for NeuPRINT assignments to be trusted for individual cells. However, modern population recording experiments typically profile many cells. In these settings, NeuPRINT enables statistical inference and hypothesis testing. (e.g., is an event more likely to happen for excitatory or inhibitory neurons?)\n\n**“Information on the computational and data demands”**\n\nTo address the reviewer’s concern, we will revise the manuscript to discuss computational and data demands of NeuPRINT. NeuPRINT has 833K parameters and takes 2h in total for training and inference on a single Nvidia Tesla V100 GPU. We have included details on the size of the dataset we used in Section 3.1 of the main paper. \n\n**“Can the approach be directly used on other kinds of neural recordings like electrophysiological data?”**\n\nIf the reviewer means inference in, for instance, high-density electrode recordings using a pre-trained NeuPRINT model (based on calcium data), then, we do not anticipate this to work because the model takes advantage of the specifics of the recording modality as well as intrinsic relations, as discussed under Sec. 2. When such recordings with post-hoc spatial transcriptomic profiling become available, we expect NeuPRINT to be applicable and useful following training with such a dataset. We will add this perspective to the Discussion of the revised manuscript.\n"}, {"review_id": "QLuAhTn1iP", "reviewer": "Reviewer_MNFF", "summary": "In this manuscript the authors describe an approach to predicting the transcriptomic identity of a neuron from its activity dynamics that they call NeuPRINT. Specifically, the authors take a dataset with calcium imaging and transcriptomic identification of the imaged neurons. First, they train up a model – either linear, nonlinear, recurrent, or a transformer – to predict the next timestep of the calcium trace using both past activity and what they term \"time-invariant\" parameters. Next, a downstream classifier is trained to use the learned parameters of the model to predict the genetic identity of each neuron. The authors demonstrate that their model outperforms the current state-of-the-art along with a series of control baselines.", "questions": "1. In Table 1 does the w/o population column mean NeuPRINT *w/o population or behavioral covariates*? In other words, does NeuPRINT have access to information that LOLCAT does not?\n1. Can the authors interrogate the learned parameters of their model to learn anything new about the dataset they are modeling?", "limitations": "The limitations of NeuPRINT could receive a more thorough treatment (see points above).", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Strengths:\n\n1. Their model includes some interesting ideas. Namely, the use of population activity and time-invariant parameters seems to confer substantial predictability. \n1. The performance of the model appears to be impressive on this multi-modal dataset.\n1. Per the ablation study shown in Figure 3, the model appears to make use of all features that were fed to it.\n\n", "weaknesses": "Weakness.\n\nMajor:\n\n1. A major claim made by the authors is that their method outperforms both LOLCAT and simple, effectively random baselines. While their model does appear to perform better, this does not look like an apples to apples comparison.\n    1. First, it would be appear that LOLCAT uses *only* a neuron's activity in order to classify its transcriptomic identity. On the other hand, NeuPRINT includes activity from nearby neurons in addition to behavioral covariates such as pupil size and running speed. Unless I am mistaken, performance comparisons between NeuPRINT and LOLCAT with similar inputs (single neuron activity *only*) are not shown until the ablation study in Figure 3. Figure 3 right shows performance that appears to be similar to LOLCAT. Is the boost in performance with NeuPRINT solely coming from the population activity and behavioral covariates? If so, then authors would need to scale back their claims about superior performance. Alternatively, the authors could add this information to LOLCAT for a more fair comparison.\n    1. Second, I am a little confused about the PCA/UMAP baselines. Here the authors selected a *random* window of time to project the data down to 64 components. Why 64? Did this account for a specific amount of variance? Why a random window of time? How was the window size chosen, and why not use a sliding window approach or use the full time-series?\n1. Whether or not NeuPRINT outperforms LOLCAT, the LOLCAT paper systematically interrogates their model for insight into the underlying biology of neural systems. I could not find any comparable biological insight here. \n\nMinor:\n\n1. The discussion of time-variant and time-invariant is fairly confusing. For example, in line 233 in the description of LOLCAT, it is asserted that \"As a consequence of the supervised training scheme, LOLCAT does not extract time-invariant representations of neuronal identity and its perfomance is constrained by the number of labels available in the dataset.\" Are the authors saying that ISI distributions are time-variant? Also, why wouldn't the performance of *any* classifier be limited to the available labels? I see no reason why this wouldn't also apply to NeuPRINT.\n1. I found a number of typos, the manuscript needs a copy-edit.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**“Is the boost in performance with NeuPRINT solely coming from the population activity and behavioral covariates?”**\n\nFirst, to address the reviewer’s concern on the role of population activity in the improved accuracy, we performed multiple new experiments to also add population activity to LOLCAT and its variants. (Please see the results in Table R1 in the overall response and Fig R1 in the attached pdf.) We reported the performance of the original implementation of LOLCAT (which became publicly available only after the initial submission) as ‘LOLCAT’. We also tuned LOLCAT with different architectures and formats of inputs. We changed its original attention module with a standard transformer architecture, and tested linear and multi-layer perceptron (MLP) classifiers, and different versions of inputs: spike interval distribution (ISI) (reported as ‘transformer+ISI’) v.s. raw fluorescence trace (reported as ‘transformer+raw’). To add population activity to these models, within the windows where ISI is computed, we calculated the mean and standard deviation of population activity in the center and surround partitions of each neuron, as well as the mean and standard deviation of behavior variables, and concatenated them with the ISI features. We also performed additional hyperparameter tuning for our NeuPRINT, LOLCAT and its variants to improve their performance. With these modifications, we evaluated both conditions (with and without population activity). The max subclass accuracy of the LOLCAT variants has improved from 0.474 to 0.561, and class accuracy from 0.608 to 0.680. However, these are still significantly worse than the performance of NeuPRINT, with 0.756 in subclass accuracy and 0.807 in class accuracy.  Interestingly, adding population activity to LOLCAT or its variants does not necessarily improve its accuracy. On the subclass prediction task, the best performance (0.561) is achieved from transformer + ISI without using population activity. For the class prediction task, adding population activity is slightly improving the performance from 0.664 to 0.680. While for NeuPRINT, adding population activity significantly improves the accuracy from 0.537 to 0.756 on subclass prediction. \n\nOn the other hand, we emphasize that using population activity and behavior information to predict neurons’ transcriptomic identity instead of only using individual neuron’s activity alone is one of the major contributions of our work, regardless of whether this novel concept could be easily adapted to existing methods and boost their performance or not.\n\n**“Second, I am a little confused about the PCA/UMAP baselines. Here the authors selected a random window of time to project the data down to 64 components. Why 64? Did this account for a specific amount of variance? ...”**\n\nThe PCA/UMAP baselines are provided merely to orient the reader. We do not expect those to ever outperform LOLCAT. Since the reported experimental sessions last for 5000 time steps, and we evaluate across 6 sessions, simply applying these baselines to the full time-series would create scalability issues. We chose the same setups and parameters for these baselines as those of NeuPRINT, so they would serve as naive unsupervised representation learning baselines, with a similar reconstruction framework, yet without a dynamical model and time-invariance constraint.\n\n**“Whether or not NeuPRINT outperforms LOLCAT, the LOLCAT paper systematically interrogates their model for insight into the underlying biology of neural systems.”**\n\nFirst, while a systematic interrogation is outside of scope (The LOLCAT paper was published in a life sciences journal. Our submission is oriented towards attendees of NeurIPS, primarily a machine learning venue.), our manuscript does reveal a novel neurobiological insight: Our results demonstrate that there is conditional information on the type of the cell of interest in the in-vivo population activities of cells other than the cell of interest. (e.g., Fig. 3: information in the population activity (or pupil size, etc) improves the accuracy of identifying the type of the cell, given the activity of that cell.)\n\nIn addition, shown in Figure R2 and Table R2 in the attached pdf, we tested our model and LOLCAT and its variants on new experiments. We include recordings with 3 sets of stimuli (drifting gratings and 2 different natural scenes) in addition to the spontaneous activity recordings. We find that the predictive performance of NeuPRINT improves when more complex stimuli are used (i.e. subclass accuracy improved from 0.756 to 0.805, and class accuracy improved from 0.807 to 0.836 when switching the condition from spontaneous recordings to natural scenes), which indicates intricate, identity-specific dynamics can become more prominent when the underlying computation becomes more complicated. This is in line with and extends the findings of Bugeon et al. We thank the reviewer for this comment, and we will edit the manuscript to highlight this finding.\n\n**“The discussion of time-variant and time-invariant is fairly confusing. Are the authors saying that ISI distributions are time-variant?”**\n\nFor LOLCAT, the ISI distribution is based on a 3-second trial sample. Therefore, the ISI distributions of neurons vary across trials, creating a time-varying descriptor. Compared to that, NeuPRINT only learns a fixed set of parameters (time-invariant) across all trials and all sessions. Importantly, this calculation scales well to long recordings because NeuPRINT has a dynamics model. We hope this explanation clarifies the confusion and we are happy to further elaborate and correct if needed.\n"}, {"review_id": "U2Ni3wI9YO", "reviewer": "Reviewer_4tWo", "summary": "This paper introduces NeuPRINT, an unsupervised approach that simultaneously learn (i )a dynamical model of recorded neural activity and (ii) a time invariant embedding of individual neurons using transformers. \n\nUsing a downstream transcriptomic classification task, the authors show that the learnt embedding space contains relevant information about the neurons identity, and demonstrated improved performances against a range of baseline methods.", "questions": "\n\n1) How are $C$, $S$ precisely defined ? (at least in the experiment)\n\n2)  The author only evaluate their model using the discovered embedding space. What about the dynamics $f$ ? How do the calcium traces reconstructed by the model compare with linear dynamical systems for example ? \n\n3) Why is the dependence on $i$ dropped from equation (3) to equation (4) ? Although $P$ must be permutation invariant, it might be relevant to model slightly different input $P^{-i}$ for each neuron. Since it does not seem much more computationally expendive to do so, can the author motivate this modeling choice ? \n\n4) How is the embedding dimension chosen ? Can it be interpreted like other method like t-SNE or UMAP ? \n\n5) Would it be possible to jointly learn $f$, $\\Phi$ *and* $P_t$ ? In doing so, the model would yield interpretable embedding and neural trajectories !\n\n", "limitations": "The authors adequately addressed the ethical limitations of their paper. ", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The proposed method not only allows to model the dynamical systems governing individual neurons recorded activity, but it also links stereotyped temporal motifs to time invariant location in an embedding space. The latter can therefore be interpreted as a representation of the intrinsic properties of the recorded neurons, like their conductance or capacitance. \n\nThe approach is conceptually interesting and could have interesting application in characterizing neural population from temporal traces. The authors designed an original model that uses summary of population activity as input to each neurons.\n\nThe approach is adapted on a range of architecture, and demonstrates that one can leverage the flexibility of transformers to find embedding more informative about neural identity than with traditional model like linear dynamical system.\n", "weaknesses": "The method description is not detailed enough and does not allow to reproduce the results presented in the paper. Overall, it is not very well written.\n\n1) The statistics used to build $C_t$ and $S_t$ are never defined. Do they correspond to 1st and 2nd moment of calcium traces or to more complicated statistics ?\n\n2) Only pseudo-code is provided. \n\n3) Almost half a page is dedicated to Hodgkin-Huxley equations, which have no technical value in the present paper. Surely it would be possible to link the time invariant neural embedding with intrinsic properties of the neurons (like its shape, conductance, capacitance etc.) without introducing many variables that are never used subsequently. In contrast, very little is said about the new method itself. \n\n4) Section 4.2:  without further explication, one cannot understand the difference between the two losses used. \"Negative Log Likelihood of Gaussian\"  (NLE) and Mean Squared Error (MSE) are exactly the same thing if the variance of the former model is 1. If the variance is not modeled/learned, then they only differ from multiplicative and additive constant. The fact that the authors did not find a monotonic relationship between MSE and NLE in Table 2 could mean two things: that (i) they indeed used different loses (in which case the method description should be more detailed), or (ii) that the same loss was used in both columns and that the difference in performances simply results from the use of different random seeds. \n\n5) Section 1.1 \"Related work\" needs rewriting. NeuPRINT is interesting in that it combines several feature of well known neural data analysis techniques. For example, the use of summary statistics of population activity $P_t$ echoes many latent models (see Saxena 2019). Yet, contrary to standard techniques like GPFA, the influence of the population on individual is captured by equation (3), which relates the NeuPRINT to coupled model like (Pillow 2008 or Truccolo 2004). In its current state, section 1.1 simply list existing approaches without really relating them to proposed approach.\n\nMinor Points: \n\n1) The vocabulary is often vague (*e.g.* \"numerous insights into how the brain works\", line 20) and many abreviations like \"w/\" or \"w/o\" are unecessarily used in the main text.\n\n2) I agree that the term \"representation\" is conveniently overloaded but the paper would gain in clarity by attempting to better explicit it. In the present study, from my understanding, the representation of a neuron is an \"embedding of its time-invariant properties that have an effect on its dynamics\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive comments on improving the writing and presentation of the paper. Meanwhile, we encourage the reviewer to also check the results in the overall response and attached pdf.\n\n**“The statistics used to build C  and S are never defined.”**\n \nWe provide the precise definitions of $C_t$ and $S_t$ as below. The reviewer is right that $C^{(-i)}_t$ and $S^{(-i)}_t$ correspond to the 1st moment (mean) and 2nd moment (standard deviation) of calcium traces of the center and surrounding partitions of every neuron $i$. We will explicitly include this definition in our revision. \n\n$$\\mu\\_{C^{(-i)}\\_t} = \\text{mean}(X\\_t^j | j: 0<d\\_{ji} < \\Delta)$$\n$$\\sigma\\_{C^{(-i)}\\_t} = \\text{std} (X\\_t^j | j: 0<d\\_{ji} < \\Delta)$$\n$$\\mu\\_{S^{(-i)}\\_t} = \\text{mean}(X\\_t^j | j: d\\_{ji} \\geq \\Delta)$$\n$$\\sigma\\_{S^{(-i)}_t} = \\text{std} (X\\_t^j | j: d\\_{ji} \\geq \\Delta)$$\n\nwhere $\\text{mean}$ and $\\text{std}$ calculate empirical mean and standard deviation of their argument, respectively, $d_{ji}$ denotes the distance between cell bodies of neurons $i$ and $j$, and $\\Delta$ denotes the distance hyperparameter partitioning center vs surround.\n\n**“Only pseudo-code is provided.”**\n\nWe will definitely provide the source code for reproducibility after the paper is published.\n\n**“Almost half a page is dedicated to Hodgkin-Huxley equations, which have no technical value in the present paper.”**\n\nWe presented Hodgkin-Huxley equations to provide biological motivation on the need to introduce the time-invariant components that contribute to neuronal dynamics. As a phenomenological model, Hodgkin-Huxley equations give an explicit format of time-invariant components and the dynamics, while in our methods, both components are defined implicitly. We understand the reviewer’s concern that this introduces extra variables which might be confused as belonging to our main method. We will reframe and move the text surrounding the Hodgkin-Huxley equations to the Related Works section in our revised paper. \n\n**“The difference between the two losses used”**\n\nWe would like to clarify that the two losses are in fact not the same. In our implementation, when trained with Gaussian Negative Log Likelihood loss (NLL), our model predicts a time-invariant Gaussian distribution (both mean and variance) for the trace at the predicted time step. Therefore the variance is learnable and is not simply set to 1. When trained with Mean Square Error loss (MSE), our model only predicts one output for the predicted time step, upon which the mean square error with the ground truth is calculated. \n\n**“Relating existing approaches to the proposed method”**\n\nWe appreciate the reviewer’s suggestions on introducing the links between our paper and previous works. We will cite these references as suggested in the revised version.\n\n**“Definition of the term “representation”**\n\nWe will explicitly define the term \"representation\" in our revised paper as \"embedding of a neuron’s time-invariant properties that have an effect on its dynamics\" as the reviewer correctly suggested.\n\n**“The authors only evaluate their model using the discovered embedding space. What about the dynamics? How do the calcium traces reconstructed by the model compare with linear dynamical systems for example?”**\n\nAs shown in Table 2 and Figure 3 of the main paper how our method compares to other dynamical modeling approaches in terms of reconstructing the calcium traces. We evaluated the reconstruction loss via two metrics: NLE and MSE, which are shown numerically in Table 2 as “loss” row, and graphically in Figure 3 as red curves with left axis readings on both plots. We showed that the reconstruction loss for our transformer-based NeuPRINT method is significantly smaller (MSE: 0.221, NLE: 0.308) compared to other dynamical models, e.g. Linear (MSE: 0.984, NLE: 1.422), Recurrent (MSE:0.887, NLE:1.231), showcasing our transformer-based model’s ability to reconstruct neuronal dynamics with significantly better performance.\n\n**“Why is the dependence on i dropped from equation (3) to equation (4)? Although P must be permutation invariant, it might be relevant to model slightly different input $P^{-i}$ for each neuron?”**\n\nWe thank the reviewer for catching the mismatch between equation 3 and equation 4. We will update our dynamical equation 4 in the main paper as follows:\n\n$$X^{(i)}\\_{t+1} =  f(X^{(i)}\\_{t-W+1:t}, (C^{(-i)}\\_{t-W+1:t}, S^{(-i)}\\_{t-W+1:t}, B\\_{t-W+1:t}), \\Phi^{(i)})$$\n\nwhere for each neuron i, the 1st and 2nd moment of activities of neighboring neurons in the “center” partition $C^{(-i)}$ and “surround” partition $S^{(-i)}$, as well as the global behavior statistics $B$, together contribute to neuron $I$’s dynamics as external input components. Note that B does not depend on $I$ as it represents global, behavioral statistics.\n\n**“How is the embedding dimension chosen? Can it be interpreted like other method like t-SNE or UMAP?”**\n\nThe embedding dimension for our model is a hyperparameter that needs to be tuned to improve the reconstruction tasks and downstream tasks, it is similar to t-SNE or UMAP in the sense that it's for dimension reduction. \n\n**“Would it be possible to jointly learn $f$, $\\Phi$ and $P_t$?”**\n\nWe thank the reviewer for the constructive comments. Currently only time-invariant representation $\\Phi$ and dynamic model $f$ are jointly learned, while $P$ is directly summarized by permutation-invariant operations with no learnable parameters, i.e. 1st and 2nd moments of population activity, since we would like the model to be invariant to varied size of neuronal population across different sessions. We agree that learning $P$ represents a valuable extension. We believe straightforward approaches will have difficulty in ensuring permutation invariance of $P$. Therefore, we did not immediately attempt to incorporate this idea. Nevertheless, we decided to mention this extension in the revised Discussion. \n"}]

Author rebuttal:
We thank the reviewers for their detailed comments. We would like clarify two key points that we thought are at the heart of much confusion:

**Uniqueness of the dataset and study**: Population recordings with calcium imaging are typically performed on mouse lines that are selective for a subset of neurons (e.g., Pvalb neurons). The large datasets produced by the Allen Institute are prominent examples of this approach. While the advantage is that the researcher knows that the recorded neurons are Pvalb neurons, the obvious disadvantage is that only Pvalb neurons are visible in this setting. Hence, interactions between different neuron types, concurrent responses of different neurons to stimuli, etc cannot be studied. Similarly, the information that other neurons carry on the type of a neuron of interest cannot be properly studied with these recordings because only one kind of neuron is available. Finally, classifying the neurons in these datasets is not a problem of practical interest. (e.g., There’s no utility to predicting that the neuron is of type Pvalb when we already know that all the neurons identified in that experiment are of type Pvalb.)

A recent method which pairs calcium imaging with post hoc spatial transcriptomics opens new avenues by allowing experiments with pan-neuronal virus infections (recording from all kinds of neurons). This is because the identity of the neurons recorded by the calcium imaging session will be subsequently identified by the spatial transcriptomics experiment. As of today, the dataset by Bugeon et al is the only such publicly available dataset. This underlines a challenge: spatial transcriptomics is slow, expensive, and not available to many research labs. This is the problem that our study aims to address: by using the available calcium imaging+spatial transcriptomics data as a dictionary, NeuPRINT aims to predict the cell types in calcium imaging-only experiments.

**Robustness of the method**: We would first like to apologize for not referring to the supplement more prominently in the initial submission. Table S1 includes a sensitivity study over five random seeds, which supports all of the original findings, and Table S2 extends our models to test on multiple animals.

Our original use of spontaneous recordings in our demonstration was not motivated by accuracy considerations. We chose those recordings in the reported experiments because (i) spontaneous recordings are performed almost ubiquitously in the field, enabling immediate practical utility, (ii) such recordings are prominently available in the Bugeon et al dataset. Indeed, it is perhaps likely that more complicated stimuli or tasks will require more complicated neuronal computations and reveal intricate, type-specific computations, thus increasing the accuracy of cell type identification. 

In order to illustrate on this point, and also demonstrate the robustness of our approach. We have performed multiple new experiments, and tested our model based on parts of the Bugeon et al dataset that we have not used in the initial submission, including testing on recordings with 3 sets of visual stimuli (drifting gratings and two different natural scenes), following the same reporting protocol. The summary table below suggests an increase in cell type identification accuracy from in-vivo activity as stimulus complexity increases, which is a novel insight:

|      NeuPRINT     | Spontaneous | Drifting Grating | Natural Scenes 1 | Natural Scenes 2 |
|:-----------------:|:-----------:|:----------------:|:----------------:|:----------------:|
| Subclass Accuracy |    0.756    |       0.780      |       0.805      |       0.782      |
|   Class Accuracy  |    0.807    |       0.809      |       0.836      |       0.836      |

We will incorporate a more detailed table including these results (see attached pdf) to the revision.

Meta-review: This paper proposed an original method for predicting neurons' transcriptomic identity from physiological recordings, which generated a high level of interest and discuss amongst the reviewers. I congratulate the authors for their detailed replies to reviewer comments, which were critical in convincing them to improve their assessment of the paper's contributions and raise their scores. Although none of the reviewers moved far enough to strongly endorse the paper for acceptance, they were unanimous in rating it as "above threshold".  In light of this agreement, I am inclined to agree that it should be accepted. Congratulations! Please attend carefully to reviewer comments and suggestions when revising the final manuscript.
Decision: Accept (poster)