Source-linked AI summary
Learnable latent embeddings for joint behavioral and neural analysis
Steffen Schneider, Jin Hwa Lee, Mackenzie Weygandt Mathis
TL;DR
Neuroscience lacks flexible nonlinear methods that jointly leverage behavioral and neural data for latent embeddings. CEBRA addresses this gap with hypothesis- or discovery-driven contrastive learning, producing consistent embeddings with high decoding performance across behaviors and datasets.
Problem
Flexible nonlinear techniques for explicitly leveraging joint behavioral and neural data to reveal neural correlates of behavior remain limited.
Method
CEBRA uses contrastive learning with user-defined behavioral labels or time-only labels to learn consistent neural latent embeddings for visualization and decoding.
Results
CEBRA achieved high decoding performance across behaviors and decoded natural-movie frames with greater than 95% accuracy using a cross-animal model trained on Neuropixels and 2P data.
Takeaways & Limitations
CEBRA supports hypothesis- and discovery-driven analysis across sessions, animals, recording modalities, and behavioral tasks while producing consistent embeddings and mitigating nuisance variance.
Takeaways & Limitations
Relating CEBRA embeddings to ground-truth latent variables requires assumptions about the data-generating process, including injectivity and the positive distribution p.
Abstract
from arXiv · showhide
Mapping behavioral actions to neural activity is a fundamental goal of neuroscience. As our ability to record large neural and behavioral data increases, there is growing interest in modeling neural dynamics during adaptive behaviors to probe neural representations. In particular, neural latent embeddings can reveal underlying correlates of behavior, yet, we lack non-linear techniques that can explicitly and flexibly leverage joint behavior and neural data. Here, we fill this gap with a novel method, CEBRA, that jointly uses behavioral and neural data in a hypothesis- or discovery-driven manner to produce consistent, high-performance latent spaces. We validate its accuracy and demonstrate our tool's utility for both calcium and electrophysiology datasets, across sensory and motor tasks, and in simple or complex behaviors across species. It allows for single and multi-session datasets to be leveraged for hypothesis testing or can be used label-free. Lastly, we show that CEBRA can be used for the mapping of space, uncovering complex kinematic features, and rapid, high-accuracy decoding of natural movies from visual cortex.
Results
CEBRA uses behavior, time, or both to learn consistent neural embeddings for hypothesis testing, discovery, visualization, and decoding. Across hippocampal, primate motor, and mouse visual datasets, it produced structured embeddings, cross-subject consistency, and high decoding performance.
- Method: CEBRA combines user-defined behavioral labels and time-only sampling with a nonlinear contrastive encoder to produce low-dimensional neural embeddings.The framework supports hypothesis-driven, discovery-driven, and hybrid analyses.
- Benchmarking: CEBRA significantly outperformed tSNE, UMAP, and pi-VAE at reconstructing ground-truth synthetic data (one-way ANOVA, F(3, 396)=278.31, p=3.95e-97).
- Consistency: CEBRA produced more consistent embeddings across subjects and repeated runs than competing algorithms, while CEBRA-Behavior exceeded conv-pi-VAE in cross-subject correlation without requiring test-time labels.CEBRA-Time and CEBRA-Behavior also produced qualitatively similar embeddings, unlike the inconsistent conv-pi-VAE outputs.
- Hypothesis and discovery: Position-guided and time-only hippocampal embeddings revealed smooth track-related structure, while hybrid embeddings separated behavior-relevant and temporal variance.Position plus direction was the most informative label, and the converged loss provided an additional label-selection metric.
- Decoding: CEBRA-Behavior achieved approximately 5 cm median absolute positional decoding error, outperforming reported 12 cm error and competing methods.For hippocampal decoding, CEBRA-Behavior significantly outperformed pi-VAE and conv-pi-VAE, while CEBRA-Time outperformed tSNE, UMAP, and PCA.
- Generalization: Joint training across animals increased embedding consistency without reducing decoding performance, and pretrained models adapted to unseen sessions at 50–100 steps/second.One adaptation step reduced positional decoding error by 10 cm.
Discussion
CEBRA is presented as a flexible framework for producing consistent neural embeddings across animals, sessions, modalities, and analysis settings. The discussion emphasizes robust decoding, rapid adaptation, and applications ranging from behavioral mapping to natural-movie reconstruction.
- Consistency: CEBRA produces consistent embeddings across subjects and can remove variance associated with animal identity, sessions, or user-defined nuisance variables.These properties are linked to revealing common structure and topologically robust neural spaces.
- Limitations: The authors identify a remaining gap in understanding how CEBRA latent spaces map onto neural-level computations.This limits the scope of the paper’s interpretation despite the method’s utility for mapping behavior to neural populations.
- Method: CEBRA combines non-linear independent component analysis with contrastive learning and specialized sampling strategies for time-series data.The method shapes embeddings conditioned on behavior or time while extending the standard InfoNCE objective.
- Versatility: CEBRA supports supervised and self-supervised analyses, enabling both hypothesis-driven and discovery-driven use with behavioral, contextual, or time-only labels.The framework is demonstrated with continuous and discrete labels and with natural-movie features used for visual-cortex decoding.
- Results: More than 95% frame-decoding accuracy was achieved from visual-cortex recordings, including a pseudo-mouse model jointly trained across Neuropixels and two-photon data.The discussion reports that CEBRA outperformed classical methods and modified conv-pi-VAE in consistency and decodability.
- Generalization: Pretrained CEBRA models adapt to new animals within tens of steps, reaching equal or better performance than training only on the unseen animal.The discussion highlights faster training than pi-VAE and relevance to time-sensitive brain-machine-interface applications.
Methods
CEBRA is evaluated across synthetic and neural datasets using paired, synchronized signals and flexible context-conditioned sampling. Its theoretical and empirical framework targets consistent, interpretable embeddings for discovery, hypothesis testing, and decoding.
- Datasets: CEBRA was benchmarked on synthetic data and rat, mouse, and macaque neural recordings spanning electrophysiology, calcium imaging, navigation, visual stimulation, and reaching.The datasets included a 1.6 m rat linear track, mouse visual cortex recordings, and macaque center-out reaching.
- Embedding model: CEBRA encodes reference and paired samples into a shared low-dimensional space and compares them with similarity functions including normalized dot product or negative mean squared error.The encoders can operate on signals from the same or different modalities, sessions, or animals.
- Data and sampling: The method requires paired recordings with synchronized neural signals and continuous or discrete context variables, while allowing user-defined positive and negative sampling schemes.Context can include behavior, time, trial identity, rewards, or brain-area identity.
- Theoretical basis: Under sufficiently diverse data, independently trained CEBRA models are consistent up to linear transformations, and matching conditional distributions identify the context variable up to the same transformation.If no mapping exists between context and signal, CEBRA instead yields a uniform distribution of points on the hypersphere.
Supplementary Note 1
Supplementary Note 1 develops the consistency and identifiability properties of CEBRA embeddings, including how sampling and auxiliary variables shape recoverable representations. It also describes discovery- and hypothesis-driven sampling choices and applications to invariant, interpretable neural embeddings.
- Consistency and identifiability: Sufficient variability in the negative sampling distribution is required for consistency of a d-dimensional embedding across runs.For each reference sample, negative samples must vary in at least d directions.
- Consistency and identifiability: CEBRA embeddings trained on comparable data distributions agree across runs up to a linear transformation at full convergence.This provides linear consistency even when individual embedding coordinates differ.
- Consistency and identifiability: Such linear consistency preserves the performance of linear downstream analyses across different CEBRA models.Examples include linear regression and general linear models.
- Consistency and identifiability: Recovering true underlying latents requires additional assumptions about the data-generating process beyond consistency across algorithm runs.The assumptions concern injectivity and the positive distribution, with hypothesis-driven training allowing the latter to be defined from observed auxiliary variables.
- Sampling and applications: The framework permits separate encoders for neural data and behavior, rather than requiring the same encoder for both data types.This flexibility allows different aspects of a dataset to be encoded separately.
- Sampling and applications: Time-contrastive learning samples reference, positive, and negative time points to target a chosen temporal scale, while behavior-guided learning uses position, direction, and related auxiliary variables.For reaching data, sampling can be conditioned on direction when position differences depend on movement direction.
- Sampling and applications: CEBRA supports discovery-driven time-only analysis and hypothesis-driven analyses that preserve selected variables while reducing the influence of others.The resulting embeddings can support visualization and downstream decoding, including cross-animal settings where animal identity is treated as irrelevant.
Supplementary Note 2
The supplementary theory establishes when CEBRA embeddings are bijective, consistent across independently trained models, and related to ground-truth latents. These guarantees depend on the similarity function, sampling setup, and assumptions about the data-generating process.
- Flexible training modes: CEBRA controls embedding geometry through positive and negative sampling distributions and the similarity function, supporting discovery-driven, hypothesis-driven, and hybrid modes.Sampling can also enforce invariance to user-defined nuisance variables such as subject or modality identity.
- Bijectivity: At optimal generalized InfoNCE solutions, compositions of the data generator and feature encoder become bijective on the data support.The bijectivity result is used to weaken some conditions required by prior contrastive-learning identifiability theory.
- Consistency: CEBRA theory derives consistency guarantees showing independently trained models can produce embeddings related by linear transformations.The result applies under dot-product similarity and generalized InfoNCE minimization, with diversity conditions supporting the guarantees.
- Ground-truth recovery: Relating CEBRA embeddings to ground-truth latents requires assumptions about the ground-truth model, unlike consistency between model runs.The paper distinguishes guarantees relating two trained models from guarantees relating a trained model to the underlying data-generating process.
- Ground-truth recovery: Under discovery-driven conditions, the learned representation is affine to the ground-truth latent space when the generator is injective and the setup assumptions hold.The proposition covers dot-product or semi-metric similarities with specified latent-domain conditions.
- Ground-truth recovery: Hypothesis-driven CEBRA recovers observable context variables up to an affine transformation while additional unobserved latent sources are not represented.The result separates variables supplied as auxiliary information from latent sources not observed during training.
Supplementary Tables
The supplementary tables report statistical tests associated with decoding and frame-error analyses in the Allen Neuropixels dataset.
- Supplementary Table 1 reports consistency statistics for Figure 1 using data from four rats.
- Supplementary Table 2 reports decoding statistics for Figure 2 across four rats, including supervised, self-supervised, and unsupervised analyses.
- Supplementary Tables 3 and 4 report posthoc Tukey HSD tests for Allen Neuropixels decoding with 10-frame and 1-frame windows.
- Supplementary Table 5 reports posthoc Tukey HSD tests for scene classification with a 1-frame window in the Allen Neuropixels dataset.
- Supplementary Table 6 reports posthoc Tukey HSD tests for mean frame error using 10 frames in the Allen Neuropixels dataset.