Know Thyself by Knowing Others: Learning Neuron Identity from Population Context

TL;DR

NuCLR introduces a self-supervised, population-contextual framework to learn neuron-level representations from large-scale neural activity. By using a permutation-equivariant spatiotemporal transformer and a two-view contrastive objective, it yields stable, discriminative neuron embeddings that transfer zero-shot across sessions and animals. The approach achieves state-of-the-art zero-shot decoding of cell type and brain region on multiple datasets and demonstrates data-efficient labeling and scalable gains with more animals. Ablations confirm the critical role of population context and dropout regularization. These results suggest that large, diverse, unlabeled neural datasets can support general-purpose neuron identity estimation across modalities and subjects.

Abstract

Neurons process information in ways that depend on their cell type, connectivity, and the brain region in which they are embedded. However, inferring these factors from neural activity remains a significant challenge. To build general-purpose representations that allow for resolving information about a neuron's identity, we introduce NuCLR, a self-supervised framework that aims to learn representations of neural activity that allow for differentiating one neuron from the rest. NuCLR brings together views of the same neuron observed at different times and across different stimuli and uses a contrastive objective to pull these representations together. To capture population context without assuming any fixed neuron ordering, we build a spatiotemporal transformer that integrates activity in a permutation-equivariant manner. Across multiple electrophysiology and calcium imaging datasets, a linear decoding evaluation on top of NuCLR representations achieves a new state-of-the-art for both cell type and brain region decoding tasks, and demonstrates strong zero-shot generalization to unseen animals. We present the first systematic scaling analysis for neuron-level representation learning, showing that increasing the number of animals used during pretraining consistently improves downstream performance. The learned representations are also label-efficient, requiring only a small fraction of labeled samples to achieve competitive performance. These results highlight how large, diverse neural datasets enable models to recover information about neuron identity that generalize across animals. Code is available at https://github.com/nerdslab/nuclr.

Figures (7)

• Figure 1: Overview of the NuCLR framework. (A) The model takes as input the activity of a neural population over a fixed context window. Each neuron's activity is treated as a patched-token sequence and encoded across time using temporal transformer blocks. These temporally encoded tokens are then passed through spatio-temporal transformer layers that attends across neurons to incorporate population context. Finally, the model outputs one vector for each neuron in the population. (B) The resulting representations are trained with a sample-wise contrastive objective. We sample two temporally-spaced views of the same population and encode them into neuron-level representations. Neurons are also randomly dropped before encoding to build robustness to partial observations. Representations of the same neuron in both views are pulled closer, and all other neurons in the population are considered negatives and pushed apart. As a result, each neuron's representation is encouraged to be stable across time and distinguishable from other neurons in the population.
• Figure 2: Representation visualization and scaling trends. (A) A 2-D UMAP visualization of NuCLR's neuron representations for the IBL Brainwide Map dataset colored by brain region. (B) Data scaling trends of inductive zero-shot decoding performance for cell type (Allen VC) and brain region (IBL) tasks. The number of animals used in pretraining are varied along with the amount of labels used to train the linear classifier heads (label ratio). Performance improves with larger pretraining data and amount of supervision. In many cases, increasing the amount of unlabeled pretraining data is more effective than increasing the amount of labeled data.
• Figure 3: Results of bin-size sweep. Classification performance variation for different classes upon sweeping the bin-size hyperparameter of the encoder. "Aggregate" refers to the overall macro-F1 score. Error polygons represents standard error of mean (SEM) measured across 3 pretraining seeds.
• Figure 4: Spike-tokenization based encoder. (A) A spike-tokenization layer for the NuCLR encoder. (B) Pretraining-time metrics comparison for spike-tokenization version and binning version of NuCLR. These plots are for pretraining on the IBL dataset.
• Figure 5: Classification performance of intermediate layer representations. Plots present the classification performance at the output of all layers in our 6 layer model for a single seed.