SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lu  $\mathbf{M}\mathbf{i}^{1,2,*}$ , Trung  $\mathbf{L}\mathbf{e}^{2,*}$ , Tianxing  $\mathbf{H}\mathbf{e}^{2}$ , Eli Shlizerman $^{2}$ , Uygar Sumbul $^{1}$
$^{1}$ Allen Institute for Brain Science
$^{2}$  University of Washington
{lu.mi,uygars}  $@$  alleninstitute.org
{tle45, shlizee}@uw.edu
goosehe@cs.washington.edu
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.
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].
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.
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_{t}^{(i)}$  denote the calcium activity of neuron  $i$  at time  $t$  and consider the following equation for dynamics:
$$
\frac {d X _ {t} ^ {(i)}}{d t} = f \left(X _ {t} ^ {(i)}, \bar {P} _ {t} ^ {(- i)}, \Phi^ {(i)}\right), \tag {1}
$$
where  $\bar{P}_t^{(-i)}$  denotes the activity of all the neurons that provide (synaptic or extra-synaptic) input to neuron  $i$  at time  $t$ .  $\Phi^{(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  $\bar{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  $\bar{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
$$
X _ {t + 1} ^ {(i)} = f \left(X _ {t - W + 1: t} ^ {(i)}, \left(C _ {t - W + 1: t} ^ {(- i)}, S _ {t - W + 1: t} ^ {(- i)}, B _ {t - W + 1: t}\right), \Phi^ {(i)}\right), \tag {2}
$$
where  $X \in \mathbb{R}^{N \times 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^{\prime}$ -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  $\Delta$ ) and the statistics in  $S$  are calculated over the surround partition (neurons whose distance to neuron  $i$  is larger than  $\Delta$ ), see Appendix E.1. Here,  $\Delta$  is a hyperparameter.  $B$  denotes the contribution of time-resolved behavioral observations. Hence, the triplet  $(C_t, S_t, B_t)$  corresponds to  $\bar{P}_t$ .  $\Phi^{(i)} \in \mathbb{R}^K$  denotes a  $K$ -dimensional
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  $N \times K$  representation  $\Phi$ . The learned representations are later evaluated on supervised downstream tasks to predict transcriptomic class and subclass identities. In the optimization framework, neuron-specific representation  $\Phi_i$  is repeated at every time step, then concatenated with masked past neuronal activity  $\hat{X}_{t-W+1:t}^{(i)}$  and permutation-invariant population inputs  $\bar{P}_{t-W+1:t}$  to form the input. The transformer model is trained to predict neural activity  $\hat{X}_{t+1}^{(i)}$  at the masked step with a causal attention mask over the  $W$ -step 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  $\Phi^{(i)}$ :
$$
\arg \min  _ {f, \left\{\Phi^ {(i)} \right\} _ {i}} \sum_ {i, t} \mathbb {E} _ {X _ {t + 1} ^ {(i)}} \left| \left| X _ {t + 1} ^ {(i)} - f \left(X _ {t - W + 1: t} ^ {(i)}, \left(C _ {t - W + 1: t} ^ {(- i)}, S _ {t - W + 1: t} ^ {(- i)}, B _ {t - W + 1: t}\right), \Phi^ {(i)}\right) \right| \right|, \tag {3}
$$
where  $||\cdot||$  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  $\Phi^{(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  $\tilde{X}_{t+1}^{(i)}$ , where neuronal activity at time  $t+1$  is zero-out. This masked activity and the past activities  $\tilde{X}_{t-W+1:t}^{(i)}$  are concatenated with the time-invariant representations  $\Phi^{(i)}$  and permutation-invariant summary of population dynamics  $\bar{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  $\hat{X}_{t+1}^{(i)}$  and ground truth activity  $X_{t+1}^{(i)}$ . The dynamical model  $f$  and time-invariant representation  $\Phi^{(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.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.3\mathrm{Hz}$ . 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  $\Phi$  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  $\Phi$  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, gated-recurrent network (GRU), and transformer with self-attention. We optimize the parameters of the dynamics  $f$  and the neuronal representation  $\Phi$  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$ .