Title: Learning Time-Invariant Representations for Individual Neurons from Population Dynamics
Venue: NeurIPS 2023
Accepted: yes
Keywords: population dynamics

neuronal representation

calcium imaging

cell types
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

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.

1 Introduction

Population recordings of neuronal activity enable relating behaviorally-relevant dynamics to the summary activity of the recorded population. While this has produced numerous insights into how the brain works [1], the activity and identity of individual neurons should be analyzed to achieve a mechanistic understanding at the implementation level [2], which may hold the key to new biologically-inspired algorithms [3, 4]. Moreover, emerging experimental evidence suggests that neurons have diverse yet stable molecular identities, which can dictate their computational roles [5, 6].

Joint (i.e., multimodal) profiling and alignment of electrophysiological features and gene expression of individual neurons suggest a good correspondence between these two modalities in slice experiments [7-11]. 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 a popular method, these experiments do not allow for studying the concurrent responses of different neuron types to stimuli.

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 a 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, a 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 a 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 a well-known yet simple connectional and functional property of neuronal circuits [12-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 [1, 15-20]. When interpretability 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 a 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 [28] 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 a set of features [7, 9, 29, 10, 11] or mechanistic model parameters [30], instead of time series models, to 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.

A large-scale study that focuses on the in-vivo activity of the individual neurons only, and not on their interaction, uses a neural network with attention layers in a supervised setting to predict the neuronal subclass, based on transgenic mouse lines [31].

Other predictive models: Implementing nonlinear autoregressive models with exogenous inputs (NARX models) using recurrently connected neural networks has a long history [32]. Modern implementations typically use gated units to address the vanishing gradients problem [33]. Recently, the transformer architecture [28], has produced impressive results in sequence-to-sequence 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 as a dynamical model to integrate with both individual and population dynamics, and the representation of individual neuron is assigned with a time-invariant embedding.

Phenomenological Hodgkin-Huxley equation: For point neuron models, which ignore the spatial distribution of the different ion channels, the celebrated Hodgkin-Huxley (HH) equations provide a phenomenological description of the neuronal membrane potential via coupled differential equations [36, 37]. The membrane potential controls the generation of spiking outputs and calcium

imaging provides an indirect measurement of this voltage over time [38]. 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 as a sum over the synapses. The HH equations also include multiple neuron-specific, 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,etc. can (i) modulate the neuron-specific parameter set, (ii) add more parameters to 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 ( i ) t denote the calcium activity of neuron i at time t and consider the following equation for dynamics:

where ¯ P ( -i ) t denotes the activity of all the neurons that provide (synaptic or extra-synaptic) input to neuron i at time t . Φ ( i ) denotes a time-invariant representation for neuron i , which implicitly captures the intrinsic parameters of HH equations and other such time-invariant aspects of neuronal identity.

Permutation-invariant 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 t (Eq. 1) in a way that is invariant to permutations, number, and detailed identity of the neurons contributing to it. To achieve this, we propose to replace ¯ P t with multiple (asymptotically) invariant statistics of the activity of a neighboring population of neurons, such as average activity [40].

Behavioral observations, such as pupil diameter, can serve as indirect readouts on the activity of unobserved neurons. Concatenating these observations with the aforementioned statistics will enrich the exogenous input observed by the dynamical model. We call this new concatenated variable P t .

Center-surround partition: Synaptic connection probability between neurons depends on distance [12, 13]. Similarly, neuronal co-variability correlates with spatial distance [14]. To include this neurobiological insight while maintaining permutation-invariance of our model, we propose a simple 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].

Discrete-time dynamical model of neuronal activity: We re-write neuronal dynamics in discrete time as

where X ∈ R N × T 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.e., 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 C are calculated over the center partition (neurons whose distance to neuron i is at most ∆ ) and the statistics in S are calculated over the surround partition (neurons whose distance to neuron i is larger than ∆ ), see Appendix E.1. Here, ∆ is a hyperparameter. B denotes the contribution of time-resolved behavioral observations. Hence, the triplet ( C t , S t , B t ) corresponds to ¯ P t . Φ ( i ) ∈ R K denotes a K -dimensional

back propagation

behavior

inputs

information

gradient

prediction loss

sndinos

session1

running speedpupil size

session1

④

Vw

neuronindex

333

M

repreated

neurons

neuronindex

Transformer

X(i)

t-W+1:t

Pt-W+1:t

implicit

mouseprimary

dynamical

Φ(i)

visualcortex(V1)

model

time

N x K

time

W

masked calcium

learnedtime-invariant

predicted calcium

overviewof

fluorescence traceinputs

representationΦ

fluorescencetraceoutputs

dynamical model f

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 Φ ( i ) :

where || · || denotes a norm. It is 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 of the optimization, while summarizing neuronal variability with Φ ( i ) , the second argument.

The idea of inferring invariant representations for neurons directly from in vivo recordings by fitting a dynamical model (i.e., predicting activities in the next time step) is 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 a 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 a 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 ˜ X ( i ) t +1 , where neuronal activity at time t +1 is zero-out. This masked activity and the past activities ˜ X ( i ) t -W +1: t are concatenated with the time-invariant representations Φ ( i ) and permutation-invariant summary of population dynamics ¯ P t -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 ˆ X ( i ) t +1 and ground truth activity X ( i ) t +1 . The dynamical model f and time-invariant representation Φ ( i ) 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 1.

4 Experiments

4.1 A Multimodal Dataset

We use a recent, public multimodal dataset to train and demonstrate our model: Bugeon et al. [6] obtained population activity recordings from the mouse primary visual cortex (V1) 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, Sncg, 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.e., 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 a new two-step benchmark: (i) self-supervised learning of time- and permutation-invariant 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 permutation-invariant 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-surround activity to predict the individual neurons' activity in the next time step. After training, we fix the dynamical model f and only optimize the time-invariant representations Φ for training, validation, and test neurons under the same self-supervised 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 1). 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, gatedrecurrent 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.

Linear model: We use a linear dynamical system where the activity at the last step is predicted from a linear combination of the activities and statistics of the population activities, behavioral information from previous steps inside a temporal window and the time-invariant representation (which is repeated

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, a nonlinear activation was 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 as 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 a 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 to train the model. For KNN, we use the scikit-learn implementation [46] with the number of nearest neighbors k = 5 .

4.4 Baselines

LOLCAT and its variants: LOLCAT [31] is a supervised framework for predicting cell types from individual neuronal activities using a 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 of [28] and use a special classification token to represent the classification of the entire neuronal activity [47]. Transformer+ISI operates on the inter-spike interval (ISI) distributions input summarized from non-overlapping subwindows of continuous 2-photon calcium recordings. We use suite2p package [48] to infer spikes from raw calcium traces and compute the ISI distributions. On the other hand, Transformer+Raw operates directly on the raw calcium traces. Unlike NeuPRINT, LOLCAT and the two supervised 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 a 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.

Principal component analysis: We project the raw calcium activities to a low-dimensional representation 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 to PCA, we project the raw calcium activities to a low-dimensional representation using Uniform Manifold Approximation and Projection (UMAP) [49], and evaluate the resulted time-variant representation by downstream classifiers.

Random: We further generate random representations for individual neurons, and train a supervised classifier on the random representations to 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 1: (i) time-invariant vs. time-variant; (ii) self-supervised representation learning vs. end-to-end

1.0

1.0

NeuPRINT EI

1.0

0.2

0.3

Lamp5

0.80

NeuPRINTsubclass

0.07

0.13

0.00

0.8

Lamp5

0.85

LOLCATsubclass

0.00

0.15

0.00

0.8

labe

True

EC

0.81

0.19

0.8

0.6

0.4

IN

0.29

0.71

0.1

0.2

label

Vip

0.08

0.75

0.17

0.00

0.6

label

Vip

0.62

0.15

0.23

0.00

0.6

EC

IN

0.0

0.0

Predictedlabel

Lamp5 Vip Pvalb Sst

True

LOLCATEI

1.0

Pvalb

0.20

0.10

0.70

0.00

0.4

Pvalb

0.47

0.06

0.47

0.00

0.4

0.8

label

EC

0.90

0.10

0.4

0.6

0.2

Sst

0.67

0.00

0.00

0.33

0.2

Sst

0.50

0.17

0.33

0.00

0.2

True

IN

0.74

0.26

0.4

0.2

0.0

EC

IN

Lamp5

Predictedlabel

Vip

Pvalb

Sst

0.0

Lamp5

Predicted label

Vip

Pvalb

Sst

0.0

Predictedlabel

EC

IN

0.0

supervised learning vs. unsupervised learning; (iii) activity of the neuron of interest ('Individual') as the only input to the dynamics model f vs. permutation-invariant representation of population dynamics provided as exogenous input to f . Under the data-limited scenario (i.e., a small amount of labeled samples), which describes a vast majority of neuroscience datasets, we find that our selfsupervised representation learning model NeuPRINT with a supervised downstream MLP classifier outperforms the current state-of-the-art 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 an 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.

Ablation studies (Table 1 and Figure 3) show that using a permutation-invariant summary of population dynamics is critical to improving NeuPRINT's performance on the downstream subclass prediction.

1.0

withoutpopulation

1.0

with population

0.8

0.756

Top-1 Class Accuracy

0.8

0.7910.807

0.6640.680

0.669

0.6

0.5610.526

0.537

0.6

0.600 0.608

0.616

Subclass

0.474

0.474

0.4

0.404

0.386

0.4

0.2

0.2

Top-1

0.0

0.0

TransformerTransformer

TransformerTransformer

LOLCAT

NeuPRINT

LOLCAT

NeuPRINT

+ISI

十

Raw

+ISI

Raw

1.2

1.0

1.4

1.0

1.0 4

AblationStudies

Top-1 Subclass accuracy

Negative Log Likelihood

Top-1 Subclass accuracy

Top-1 Subclass accuracy

MeanSquaredError

1.0

0.8

1.2

0.8

0.8

0.8

0.6

1.0

0.6

0.6

0.6

0.8

0.4

0.4

0.4

0.4

MSE

0.6

NLL

subclass acc

0.2

subclass acc

F 0.2

0.2

0.2

class acc

0.4

class acc

0.0

0.0

0.0

Linear NonlinearRecurrent Transformer

w/o

I!dnd+

size

+frame

+running

state

punouns-

uoejndod

speed

uoendod+

activity

On the other hand, this effect is 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 over 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 in 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 Permutation-Invariant Summary of Population Dynamics Enhances the Time-invariant Representation of Individual Neurons

To further investigate the role of each component in the permutation-invariant summary of population dynamics, we perform a series of ablation studies as shown in Figure 4. 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 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 to external inputs (hence the summary variables proposed here) forms a part of the neuron identity. Overall, we find using all of the available permutation-invariant components of the summary of the population recording improves the accuracy by 22% in subclass prediction.

Top-1 Subclass Accuracy

1.0

Spontaneous

0.8

DriftingGratings

Natural Scenes1

0.6

0.4

0.474

Natural Scenes 2

0.421

0.491 0.491

0.561

0.509

0.579 0.561

0.474 0.509

0.439

0.544

0.2

0.0

Transformer

Transformer

LOLCAT

NeuPRINT

1.0 1

+ISI

+Raw

ClassAccuracy

0.807 0.809 0.836 0.836

0.8

0.608 0.600

0.648 0.648

0.680

0.640 0.672 0.664

0.669 0.648

0.696

0.656

0.6

0.4

Top-1(

0.2

0.0

Transformer

Transformer

LOLCAT

NeuPRINT

ISI+

+Raw

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 3 sets of visual stimuli (drifting gratings and two different natural scene image sets [6]). The results summarized in Figure 5 suggest an increase in 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 1 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 2. While the performance of LOLCAT increases on the class prediction task with this larger dataset, it 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 on 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 of a neuron has two distinct components (the neuron itself, and the synaptic and modulatory inputs it receives), we presented a self-supervised approach to infer identity 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 permutation-invariant statistics of the activity of the neighboring neurons.

In the field of natural language processing (NLP), learning the dynamics of the underlying system has been shown to enable success in 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 a self-supervised task. We trained simple classifiers with limited ground-truth labels to subsequently predict the molecular label, based on 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, 43], the transformer architecture with its self-attention mechanism provides the best results among the architectures we considered. Due to 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-invariant inputs and information 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 permutation-invariance property. In particular, motivated by insights on circuit motifs in neuronal networks [12-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 recordings by learning extra embeddings that represent different brain regions, sessions or 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 an initial attempt. Ever-growing neuroscience datasets represent a straightforward 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 as 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 a 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 ( https://github.com/lumimim/NeuPRINT/ ) for reproducibility.

References

Shreya Saxena and John P Cunningham. Towards the neural population doctrine. Current opinion in neurobiology , 55:103-111, 2019.

David Marr. Vision: A computational investigation into the human representation and processing of visual information . MIT press, 2010.

Hannah Bos, Anne-Marie Oswald, and Brent Doiron. Untangling stability and gain modulation in cortical circuits with multiple interneuron classes. bioRxiv , pages 2020-06, 2020.

Yuhan Helena Liu, Stephen Smith, Stefan Mihalas, Eric Shea-Brown, and Uygar Sümbül. Celltype-specific neuromodulation guides synaptic credit assignment in a spiking neural network. Proceedings of the National Academy of Sciences , 118(51):e2111821118, 2021.

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(3):522-539, 2017.

Stephane Bugeon, Joshua Duffield, Mario Dipoppa, Anne Ritoux, Isabelle Prankerd, Dimitris Nicoloutsopoulos, David Orme, Maxwell Shinn, Han Peng, Hamish Forrest, et al. A transcriptomic axis predicts state modulation of cortical interneurons. Nature , 607(7918):330-338, 2022.

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):935953, 2020.

Federico Scala, Dmitry Kobak, Matteo Bernabucci, Yves Bernaerts, Cathryn René Cadwell, Jesus Ramon Castro, Leonard Hartmanis, Xiaolong Jiang, Sophie Laturnus, Elanine Miranda, et al. Phenotypic variation of transcriptomic cell types in mouse motor cortex. Nature , 598(7879):144-150, 2021.

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.

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 Series C: Applied Statistics , 70(4):980-1000, 2021.

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.

Jennifer S Lund, Alessandra Angelucci, and Paul C Bressloff. Anatomical substrates for functional columns in macaque monkey primary visual cortex. Cerebral cortex , 13(1):15-24, 2003.

Robert B Levy and Alex D Reyes. Spatial profile of excitatory and inhibitory synaptic connectivity in mouse primary auditory cortex. Journal of Neuroscience , 32(16):5609-5619, 2012.

Robert Rosenbaum, Matthew A Smith, Adam Kohn, Jonathan E Rubin, and Brent Doiron. The spatial structure of correlated neuronal variability. Nature neuroscience , 20(1):107-114, 2017.

Byron M Yu, John P Cunningham, Gopal Santhanam, Stephen Ryu, Krishna V Shenoy, and Maneesh Sahani. Gaussian-process factor analysis for low-dimensional single-trial analysis of neural population activity. Advances in neural information processing systems , 21, 2008.

Yuanjun 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.

Yuan Zhao and Il Memming Park. Variational latent gaussian process for recovering single-trial dynamics from population spike trains. Neural computation , 29(5):1293-1316, 2017.

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-922. PMLR, 2017.

Daniel Durstewitz. A state space approach for piecewise-linear recurrent neural networks for identifying computational dynamics from neural measurements. PLoS computational biology , 13(6):e1005542, 2017.

Felix Pei, Joel Ye, David Zoltowski, Anqi Wu, Raeed H Chowdhury, Hansem Sohn, Joseph E O'Doherty, Krishna V Shenoy, Matthew T Kaufman, Mark Churchland, et al. Neural latents benchmark'21: evaluating latent variable models of neural population activity. arXiv preprint arXiv:2109.04463 , 2021.

Chethan Pandarinath, Daniel J O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D Stavisky, Jonathan C Kao, Eric M Trautmann, Matthew T Kaufman, Stephen I Ryu, Leigh R Hochberg, et al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods , 15(10):805-815, 2018.

Mohammad Reza Keshtkaran, Andrew R Sedler, Raeed H Chowdhury, Raghav Tandon, Diya Basrai, Sarah L Nguyen, Hansem Sohn, Mehrdad Jazayeri, Lee E Miller, and Chethan Pandarinath. A large-scale neural network training framework for generalized estimation of single-trial population dynamics. Nature Methods , pages 1-6, 2022.

Steffen Schneider, Jin Hwa Lee, and Mackenzie Weygandt Mathis. Learnable latent embeddings for joint behavioral and neural analysis. arXiv preprint arXiv:2204.00673 , 2022.

Tze Hui Koh, William E Bishop, Takashi Kawashima, Brian B Jeon, Ranjani Srinivasan, Yu Mu, Ziqiang Wei, Sandra J Kuhlman, Misha B Ahrens, Steven M Chase, et al. Dimensionality reduction of calcium-imaged neuronal population activity. Nature Computational Science , 3(1):71-85, 2023.

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 neurophysiology , 93(2):1074-1089, 2005.

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-999, 2008.

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-2391, 2022.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

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):120-127, 2021.

Yves Bernaerts, Michael Deistler, Pedro J Goncalves, Jonas Beck, Marcel Stimberg, Federico Scala, Andreas S Tolias, Jakob H Macke, Dmitry Kobak, and Philipp Berens. Combined statistical-mechanistic modeling links ion channel genes to physiology of cortical neuron types. bioRxiv , pages 2023-03, 2023.

Aidan Schneider, Mehdi Azabou, Louis McDougall-Vigier, David F Parks, Sahara Ensley, Kiran Bhaskaran-Nair, Tomasz Nowakowski, Eva L Dyer, and Keith B Hengen. Transcriptomic cell type structures in vivo neuronal activity across multiple timescales. Cell Reports , 42(4), 2023.

Tsungnan Lin, Bill G Horne, Peter Tino, and C Lee Giles. Learning long-term dependencies in narx recurrent neural networks. IEEE Transactions on Neural Networks , 7(6):1329-1338, 1996.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation , 9(8):1735-1780, 1997.

Joel Ye and Chethan Pandarinath. Representation learning for neural population activity with neural data transformers. Neurons, Behavior, Data analysis, and Theory , 5(3):1-18, 2021.

Trung Le and Eli Shlizerman. Stndt: Modeling neural population activity with spatiotemporal transformers. In Advances in Neural Information Processing Systems , 2022.

Alan L Hodgkin and Andrew F Huxley. A quantitative description of membrane current and its application to conduction and excitation in nerve. The Journal of physiology , 117(4):500, 1952.

Wulfram Gerstner, Werner M Kistler, Richard Naud, and Liam Paninski. Neuronal dynamics: From single neurons to networks and models of cognition . Cambridge University Press, 2014.

Christine Grienberger and Arthur Konnerth. Imaging calcium in neurons. Neuron , 73(5):862885, 2012.

Eric Jonas and Konrad Kording. Automatic discovery of cell types and microcircuitry from neural connectomics. Elife , 4:e04250, 2015.

Jakub M Tomczak, Maximilian Ilse, and Max Welling. Deep learning with permutation-invariant operator for multi-instance histopathology classification. arXiv preprint arXiv:1712.00310 , 2017.

Stephen W Kuffler. Discharge patterns and functional organization of mammalian retina. Journal of neurophysiology , 16(1):37-68, 1953.

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-1259, 1998.

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:1877-1901, 2020.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32 , pages 8024-8035. Curran Associates, Inc., 2019.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

F. 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 E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research , 12:2825-2830, 2011.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

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.

Leland McInnes, John Healy, Nathaniel Saul, and Lukas Großberger. Umap: Uniform manifold approximation and projection. Journal of Open Source Software , 3(29):861, 2018.

A Sensitivity Analysis

To quantify the sensitivity of our methods and baselines on 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 3.

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 of our evaluation framework and the significance of the performance gaps.

B A Spectrum of Implicit Dynamical Models

We explore a 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 (NLL) when the model predicts both mean and standard deviation of output distribution to train the models. The results are summarized in Table 4.

C Roles of Permutation-Invariant Summary of Population Dynamics

To investigate the role of each component in the permutation-invariant summary of population dynamics, we perform a series of ablation studies as shown in Table 5. 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 of the proposed components contribute to the success of the time-invariant representations as evaluated by the downstream transcriptomic classification task.

D Extensions to Other Visual Stimulus Conditions

We further test our model NeuPRINT and other baselines on recordings with 3 sets of visual stimuli (drifting gratings and two different natural scenes). The results in Table 6 suggest an increase in cell type identification accuracy from in-vivo activity as stimulus relevance and complexity increases.

E Implementation Details

E.1 NeuPRINT

Permutation-invariant summary of population dynamics: We use C , S to represent D , D dimensional permutation- and size-invariant (i.e., N ) surrogates for brain activity at each time point. C and S compute identical statistics (mean and standard deviation) of population activity except that the statistics in C are calculated over the center partition (neurons whose distance to neuron i is at most ∆ ) and the statistics in S are calculated over the surround partition (neurons whose distance to neuron i is larger than ∆ ). Here, ∆ is a hyperparameter:

Batch sampling: For each batch of data we randomly sample 512 labeled and unlabeled neurons to be included in the batch. For each neuron we further randomly sample 2 sub-windows of 512 timesteps from its continuous calcium flourescene traces. Each resulting sample for the i th neuron is denoted X ( i ) .

Multihead attention: We use the transformer encoder architecture [28] and the masking strategy as originally proposed in [47]. Random timesteps in X ( i ) are masked (zero-out) with probability 0 . 25 and concatenated with the permutation-invariant population summary ¯ P and time-invariant representation ϕ ( i ) along the feature dimension to form input ¯ X ( i ) . Note that the same learnable ϕ ( i ) is repeated at every timestep, enforcing time-invariance. Similar to [47], we embed input ¯ X ( i ) and

employ sinusoidal positional embedding to encode the temporal order in the input sequence, resulting in ˜ X ( i ) = Emb ( ¯ X ( i ) + E.

For each input ˜ X ( i ) , a set of weights W Q ∈ R T × d q , W K ∈ R T × d k , W V ∈ R T × d v are learned to transform input ˜ X ( i ) to a set of query, key, and value ( Q,K,V ) , where Q = ˜ X ( i ) W Q , K = ˜ X ( i ) W K , V = ˜ X ( i ) W V . Attention between temporal tokens for one attention head is computed as:

Each head will find a different pattern in the data and produce an output of size d v . The final attention output will be a concatenation of these single-head outputs. We use 2 heads in our model.

Feedforward layers and residual connections are subsequently applied to attention output:

where MSA represents the multihead attention operation, FF represents the feedforward layer with ReLU activation, and Z ( i ) represents the reconstructed calcium trace with masked timesteps recovered.

Computational cost: NeuPRINT 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.1 of the main paper.

E.2 LOLCAT and Its Variants

We use the publicly available implementation of LOLCAT from [31]. 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]. For LOLCAT and Transformer+ISI, continuous time series of calcium traces are divided into non-overlapping sub-windows of size 64, within which the Inter-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 to represent the classification output [47]. No positional embedding is added to the input sequence. Therefore the transformer output at the CLS token position will represent the pooling operation as in [31], and all sub-windows are treated as if they are independent trials, which enables Transformer+ISI to apply to any number of observed trials - an important design choice for LOLCAT.

We further implement Transformer+Raw, a variant of LOLCAT where the transformer operates directly on the raw calcium traces rather than the ISI distribution of calcium traces. Transformer+Raw

follows the same architecture as Transformer+ISI, except that now the positional embedding is added to 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, we train and evaluate downstream classifiers using random vectors of size 64 as representations for individual neurons, equivalent to the 64-dim ϕ ( i ) . To compare NeuPRINT with unsupervised methods PCA and UMAP, we first project the data to a lower dimensional space using 64 components, then train and evaluate downstream classifiers on the low-dimensional representation.

E.4 Downstream Classifiers

We use 5 nearest neighbors for KNN downstream classifier. For the downstream MLP classifier, we use a multi-layer perceptron network with a single hidden layer of size 2048 and ReLU activation.

E.5 Hyperparameters

We include all of the important hyperparameters (representation dim, window size, number of epochs, learning rate, batch size, etc.) for our NeuPRINT and other models (supervised-learning baselines LOLCAT, Transformer+ISI and Transformer+Raw, unsupervised representation learning baselines UMAP and PCA, chance-level prediction based on random features) in Table 7.

F Pseudo Code

Our NeuPRINT framework includes three main components: an 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; a supervised learning framework to train the downstream classifiers for subclass and class prediction, taking the learned time-invariant representations as inputs. The pseudo code for these three components is listed as follows:

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nThe 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\".\n\nQUESTIONS:\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\nLIMITATIONS:\nThe authors adequately addressed the ethical limitations of their paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nStrengths:\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\nWEAKNESSES:\nWeakness.\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\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe limitations of NeuPRINT could receive a more thorough treatment (see points above).", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nUnderstanding 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nWhile 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.\n\nQUESTIONS:\nCan the approach be directly used on other kinds of neural recordings like electrophysiological data?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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\nQUESTIONS:\nWhile 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\nLIMITATIONS:\nThe paper addresses limitations (although that could be expanded) and has a broader impact section.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}