SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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.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
Input	—	Individual without Population	Individual with Population
Supervision	Lower Bound	Data-Limited Supervised	Unsupervised	Self -supervised	Data-Limited Supervised	Self -supervised
Task	Model	Random	LOLCAT	Trans Trans +ISI +Raw	PCA	UMAP	NeuPRINT	LOLCAT	Trans Trans +ISI +Raw	NeuPRINT
Subclass	KNN	0.260	—	—	0.263	0.281	0.415	—	—	0.610
Linear	0.256	0.404	0.474	0.474	0.316	0.404	0.537	0.474	0.683
MLP	0.302	—	0.561	0.439	0.330	0.340	0.512	—	0.756
Class	KNN	0.488	—	—	0.536	0.584	0.652	—	—	0.711
Linear	0.526	0.600	0.664	0.669	0.544	0.576	0.697	0.608	0.793
MLP	0.523	—	0.640	0.664	0.565	0.520	0.752	—	0.807
Table 1: Top-1 accuracy of transcriptomic label prediction based on representations learned by (i) our proposed self-supervised representation learning from neural dynamics framework NeuPRINT, (ii) the supervised learning method LOLCAT [31] and its variants Transformer+ISI and Transformer+Raw, (iii) unsupervised baselines PCA and UMAP, (iv) random representations (to determine the chance-level). Note that this experiment corresponds to a data-limited regime due to limited labeled data. We performed classification using three classifiers (KNN, Linear, MLP) and two tasks: predicting the subclass from the set {Lamp5, Pvalb, Vip, Sncg, Sst}, and predicting the cell class from the set {excitatory, inhibitory}. We study the performance of different models using only individual neuronal activity vs adding population statistics as input.
Figure 2: Left: Relative abundances of subclass and class labels. Right: Confusion matrices of our self-supervised representation learning framework NeuPRINT and supervised learning method LOLCAT [31] based on predicting the cell class and subclass labels. While both of the self-supervised and supervised steps are learned with all available subclasses, we excluded the Sncg population from the confusion matrices because it represents a negligible fraction of the test set with the  $80\% : 10\% : 10\%$  split, so that quantification for this population would not be reliable.
supervised learning vs. unsupervised learning; (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 self-supervised 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.
Figure 3: Accuracy of transcriptomic subclass and class prediction of NeuPRINT and baselines on single-mouse spontaneous activity recordings, with and without inputs from population statistics.
Figure 4: Left: Top-1 accuracy (or loss) of learned representations with different dynamical models (linear, nonlinear, recurrent, transformer) in the subclass prediction task. Right: Ablation studies to dissect the impact of the different components of the permutation-invariant summary of population dynamics including running speed, pupil size, frame state, population activity, center-surround activity in improving the accuracy, as in Table 5. One component is added at a time from left to right.
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.
Figure 5: Accuracy of transcriptomic subclass and class prediction of NeuPRINT and baselines on single-mouse recordings during spontaneous activity and visual stimuli-driven (drifting gratings, natural scenes) activity.
Learned Representation	Time-Variant	Time-Invariant
Lower Bound	Data-Limited Supervised	Unsupervised	Self-supervised
Inputs	—	Individual	Individual + Population
Task	Classifier	Random	LOLCAT	Trans +ISI	Trans +Raw	PCA	UMAP	NeuPRINT without population
Subclass	KNN	0.174	—	—	—	0.333	0.348	0.419
Linear	0.340	0.457	0.449	0.493	0.304	0.384	0.552
MLP	0.362	—	0.442	0.423	0.384	0.406	0.552
Class	KNN	0.581	—	—	—	0.670	0.613	0.659
Linear	0.667	0.700	0.710	0.675	0.645	0.660	0.746
MLP	0.660	—	0.702	0.707	0.682	0.667	0.770
Table 2: Extensions to multiple animals: Top-1 accuracy of transcriptomic label prediction based on (i) the representations learned by our proposed self-supervised representation learning from neural dynamics framework NeuPRINT, (ii) the supervised learning method LOLCAT and its variants Transformer+ISI and Transformer+Raw, (iii) unsupervised baselines PCA and UMAP, (iv) random representations (to determine the chance-level). Note that this experiment corresponds to a data-limited regime due to limited labeled data. We performed classification using three classifiers (KNN, Linear, MLP) and two tasks: predicting the cell class from the set {excitatory, inhibitory} and predicting the subclass from the set {Lamp5, Pvalb, Vip, Sncg, Sst}.
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 ( $\sim 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  $\sim 23\%$ , and also by  $\sim 10\%$  in class prediction (See discussion on the preliminary nature of this study below).
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/luminim/NeuPRINT/) for reproducibility.
LM is supported by the Shanahan Foundation Fellowship. LM and US thank the Allen Institute founder, Paul G Allen, for his vision, encouragement, and support. ES and TL acknowledge the support in part by A3D3 National Science Foundation grant OAC-2117997 and the Department of Electrical Computer Engineering at the University of Washington. ES also acknowledges the support in part by the Washington Research Foundation Fund and the Department of Applied Mathematics at the University of Washington.
Broader Impact Since we studied a basic neuroscience problem using recordings from mice, we do not expect this paper to have an immediate societal impact. We hope that, in the long term, variants of our method will be helpful to better understand how the brain works and the physiological consequences of the various diseases of the brain. It is nevertheless important to keep in mind that technical progress, such as reported in our study, could potentially be adapted without much difficulty to malicious use.
[1] Shreya Saxena and John P Cunningham. Towards the neural population doctrine. Current opinion in neurobiology, 55:103-111, 2019.
[2] David Marr. Vision: A computational investigation into the human representation and processing of visual information. MIT press, 2010.
[3] Hannah Bos, Anne-Marie Oswald, and Brent Doiron. Untangling stability and gain modulation in cortical circuits with multiple interneuron classes. bioRxiv, pages 2020-06, 2020.
[4] Yuhan Helena Liu, Stephen Smith, Stefan Mihalas, Eric Shea-Brown, and Uygar Sumbul. Cell-type-specific neuromodulation guides synaptic credit assignment in a spiking neural network. Proceedings of the National Academy of Sciences, 118(51):e2111821118, 2021.
[5] Anirban Paul, Megan Crow, Ricardo Raudales, Miao He, Jesse Gillis, and Z Josh Huang. Transcriptional architecture of synaptic communication delineates gabaergic neuron identity. Cell, 171(3):522-539, 2017.
[6] Stephane Bugeon, Joshua Duffield, Mario Dipoppa, Anne Ritoux, Isabelle Prankerd, Dimitris Nicoloutsopoulos, David Orme, Maxwell Shinn, Han Peng, Hamish Forrest, et al. A transcriptomic axis predicts state modulation of cortical interneurons. Nature, 607(7918):330-338, 2022.
[7] Nathan W Gouwens, Staci A Sorensen, Fahimeh Baftizadeh, Agata Budzillo, Brian R Lee, Tim Jarsky, Lauren Alfiler, Katherine Baker, Eliza Barkan, Kyla Berry, et al. Integrated morphoelectric and transcriptomic classification of cortical gabaergic cells. Cell, 183(4):935-953, 2020.
[8] Federico Scala, Dmitry Kobak, Matteo Bernabucci, Yves Bernaerts, Cathryn René Cadwell, Jesus Ramon Castro, Leonard Hartmanis, Xiaolong Jiang, Sophie Laturnus, Elanine Miranda, et al. Phenotypic variation of transcriptomic cell types in mouse motor cortex. Nature, 598(7879):144-150, 2021.
[9] Rohan Gala, Nathan Gouwens, Zizhen Yao, Agata Budzillo, Osnat Penn, Bosiljka Tasic, Gabe Murphy, Hongkui Zeng, and Uygar Sumbul. A coupled autoencoder approach for multi-modal analysis of cell types. Advances in Neural Information Processing Systems, 32, 2019.
[10] Dmitry Kobak, Yves Bernaerts, Marissa A Weis, Federico Scala, Andreas S Tolias, and Philipp Berens. Sparse reduced-rank regression for exploratory visualisation of paired multivariate data. Journal of the Royal Statistical Society Series C: Applied Statistics, 70(4):980-1000, 2021.
[11] Yeganeh Marghi, Rohan Gala, Fahimeh Baftizadeh, and Uygar Sumbul. Joint inference of discrete cell types and continuous type-specific variability in single-cell datasets with mmidas. bioRxiv, pages 2023–10, 2023.
[12] Jennifer S Lund, Alessandra Angelucci, and Paul C Bressloff. Anatomical substrates for functional columns in macaque monkey primary visual cortex. *Cerebral cortex*, 13(1):15-24, 2003.
[13] 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.
[14] 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.
[47] 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.
[48] Marius Pachitariu, Carsen Stringer, Sylvia Schröder, Mario Dipoppa, L Federico Rossi, Matteo Carandini, and Kenneth D Harris. Suite2p: beyond 10,000 neurons with standard two-photon microscopy. *BioRxiv*, page 061507, 2016.
[49] Leland McInnes, John Healy, Nathaniel Saul, and Lukas Großberger. Umap: Uniform manifold approximation and projection. Journal of Open Source Software, 3(29):861, 2018.
Learned Representation	Time-Variant	Time-Invariant
Lower Bound	Data-Limited Supervised	Unsupervised	Self-supervised
Inputs	—	Individual	Individual + Population
Task	Classifier	Random	LOLCAT	Trans +ISI	Trans +Raw	PCA UMAP	NeuPRINT without population
Subclass	KNN	0.196 ± 0.041	—	—	—	0.345 ± 0.070	0.424 ± 0.051
Linear	0.326 ± 0.055	0.449 ± 0.027	0.414 ± 0.036	0.436 ± 0.021	0.316 ± 0.046	0.595 ± 0.041
MLP	0.382 ± 0.062	—	0.456 ± 0.068	0.422 ± 0.017	0.392 ± 0.044	0.606 ± 0.054
Class	KNN	0.509 ± 0.061	—	—	—	0.569 ± 0.015	0.568 ± 0.057
Linear	0.536 ± 0.035	0.616 ± 0.017	0.632 ± 0.026	0.652 ± 0.009	0.584 ± 0.053	0.741 ± 0.035
MLP	0.565 ± 0.030	—	0.624 ± 0.010	0.651 ± 0.007	0.596 ± 0.015	0.782 ± 0.018
Table 3: Sensitivity analysis across 5 runs with different random seeds: Top-1 accuracy (mean±standard deviation) of transcriptomic label prediction based on (i) the representations learned by our proposed self-supervised representation learning from neural dynamics framework NeuPRINT, (ii) the supervised learning method LOLCAT and its variants Transformer+ISI and Transformer+Raw (abbreviated Trans+ISI and Trans+Raw, respectively) (Note that this experiment corresponds to a data-limited regime due to the size of the dataset), (iii) unsupervised baselines PCA and UMAP, (iv) random representations (to determine the chance-level). We performed classification using three different simple classifiers (KNN, Linear, MLP) and two tasks: predicting the cell class from the set {excitatory, inhibitory} and predicting the subclass from the set {Lamp5, Pvalb, Vip, Sncg, Sst}.
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.
Task	Linear	Nonlinear	Recurrent	Transformer
Objective	NLL	MSE	NLL	MSE	NLL	MSE	NLL	MSE
Loss	1.422	0.984	1.412	1.022	1.231	0.887	0.308	0.221
Subclass	0.488	0.512	0.537	0.512	0.610	0.610	0.707	0.756
Class	0.706	0.716	0.724	0.679	0.743	0.743	0.752	0.807