Learnable latent embeddings for joint behavioral and neural analysis

TL;DR

CEBRA can be used for the mapping of space, uncovering complex kinematic features, for the production of consistent latent spaces across two-photon and Neuropixels data, and can provide rapid, high-accuracy decoding of natural videos from visual cortex.

Abstract

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.

Figures (15)

• Figure 1: CEBRA for consistent and interpretable embeddings(a): CEBRA allows for self-supervised, supervised, and hybrid approaches for hypothesis-driven and discovery-driven analysis. Overview of pipeline: collect data (e.g., pairs of behavior (or time) and neural data (x,y)), determine positive and negative pairs, train CEBRA, and produce embeddings. (b): Left: True 2D latent, where each point is mapped to spiking rate of 100 neurons. (Middle): CEBRA embedding after linear regression to the true latent. Right: Reconstruction score is $R^2$ of linear regression between the true latent and resulting embedding from each method. The "behavior label" is a 1D random variable sampled from uniform distribution of [0, 2$\pi$] that is assigned to each time bin of synthetic neural data, visualized by the color map. The orange line is the median,and each black dot is an individual run (n=100). CEBRA-Behavior shows significantly higher reconstruction score compare to pi-VAE, tSNE and UMAP (one-way ANOVA, F(3, 396)=278.31, p=3.95e-97 with post hoc Tukey HSD p<0.001). (c): Rat hippocampus data from grosmark2016diversity. Electrophysiology data collected during a task where the animal transverse a 1.6m linear track "leftwards" or "rightwards". (d): We benchmarked CEBRA against conv-pi-VAE (both with labels and without (self-supervised mode)), tSNE, and unsupervised UMAP. Note, for performance against the original pi-VAE see Extended Data Fig. \ref{['fig:CEBRAintroData']}. We plot the 3 latents (note, all CEBRA embedding figures show the first 3 latents). The dimensionality (D) of the latent space is set to the minimum and equivalent dimension per method (3D for CEBRA and 2D for others) to fairly compare. Note, higher dimensions for CEBRA can give higher consistency values (see Extended Data Fig. \ref{['fig:multiSession']}). (e): Correlation matrices depict the $R^2$ after fitting a linear model between behavior-aligned embeddings of two animals, one as the target one as the source (mean, n=10 runs). Parameters were picked by optimizing average run consistency across rats.
• Figure 2: Hypothesis-driven and discovery-driven analysis with CEBRA(a): CEBRA can be used in three modes: hypothesis-driven, discovery-driven, or in a hybrid mode, which allows for weaker priors on the latent embedding. (b): CEBRA with position-hypothesis derived embedding, shuffled (erroneous), time-only, and Time+Behavior (hybrid; here, a 5D space was used, where first 3D is guided by both behavior+time, and last 2D is guided only by time, and the first 3 latents are plotted). (c): Embeddings with position-only, direction-only, and shuffled position-only, direction-only for hypothesis testing. The loss function can be used as a metric for embedding quality. (d): We utilized the hypothesis-driven (position+direction) or the shuffle (erroneous) to decode the position of the rat, which produces a large difference in decoding performance: position+direction $R^2$ is 73.35% vs. -49.90% shuffled and median absolute error 5.8 cm vs 44.7 cm. Purple line is decoding from the 32-dimensional hypothesis-based latent space, dashed line is shuffled. Right is the performance across additional methods (The orange line indicates the median of the individual runs (n=10) that are indicated by black circles. Each run is averaged over 3 splits of the dataset). (e): Schematic of how persistent co-homology is computed. Each data point is thickened to a ball of gradually expanding radius $r$, while tracking birth and death of "cycles" in each dimension ($H^0$ counts number of connected components or 0-dim cycles, $H^1$ counts the number of loops (1-dim cycles), $H^2$ counts the number of voids (2-dim cycles)). The prominent lifespans, indicated as pink and purple arrows, are considered to determine Betti numbers. (f): Visualization of the neural embeddings computed with different input dimensions, and the related persistent co-homology lifespan diagrams below. (g): Betti numbers from shuffled embeddings (Sh.) and across increasing dimensions (d) of CEBRA, and the topology preserving circular coordinates using the first co-cycle from persistent co-homology analysis (see Methods).
• Figure 3: Forelimb movement behavior in a primate(a): Behavioral setup: monkey makes either active movements in 8 directions with the manipulandum, or the arm is passively moved via the manipulandum (real behavioral trajectories shown, with cartoon depicting the task setup). Behavior and neural recordings are from area 2 of the primary somatosensory cortex from Chowdhury et al. chowdhury2020area. (b): Comparison of embeddings of active trials generated with CEBRA-Behavior, CEBRA-Time, conv-pi-VAE variants, tSNE, and UMAP. The embeddings of trials (n=364) of each direction are post-hoc averaged. (c): CEBRA-Behavior trained with x,y position of the hand. Left panel is color-coded to x position and right panel is color-coded to y position, as in d. (d): CEBRA-Time without any external behavior variables. As in c, left and right are color-coded to x and y position, respectively. (e): Left, CEBRA-Behavior embedding trained with a 4D latent space, with discrete target direction as behavior labels, trained and plotted separately for active and passive trials. (f): Left, CEBRA-Behavior embedding trained with a 4D latent space, with discrete target direction and active and passive trials as behavior labels, plotted separately, active vs. passive trials. (g): CEBRA-Behavior embedding trained with a 4D latent space using active and passive trials with continuous (x,y) position as behavior labels, plotted separately, active vs. passive trials. The trajectory of each direction is averaged across trials (n=18--30 each, per directions) over time. Each trajectory represents 600ms from -100ms before the start of the movement. (h): Left to right: Decoding performance of: position using CEBRA-Behavior trained with x,y position (active trials); target direction using CEBRA-Behavior trained with target direction (active trials); or active vs. passive accuracy using CEBRA-Behavior trained with both active and passive movements. For each case, we trained and evaluated 5 seeds represented by black dot and the orange line represents median. (i): Decoded trajectory of hand position using CEBRA-Behavior trained on active trial with x,y position of hand. Grey line is true trajectory and red line is decoded trajectory.
• Figure 4: Spikes and calcium signaling reveal similar CEBRA embeddings(a): CEBRA-Behavior can use frame-by-frame video feature as a label of sensory input to extract neural latent space of visual cortex of mice watching a movie. (b): tSNE visualization of the DINO features of the movie frames from four different DINO configurations (latent size, model size) commonly show continuous evolution of the movie frames over time. (c, d): Visualization of trained 8D latent CEBRA-Behavior embeddings with Neuropixels data or calcium imaging, respectively. The numbers on top of each embedding is the number of neurons subsampled from the multi-session concatenated dataset. Color map is the same as in b. (e): Linear consistency between embeddings trained with either calcium imaging data or Neuropixels data. (f, g): Visualization of CEBRA-Behavior embedding (8D) trained with Neuropixels and calcium imaging, jointly. Color map is the same as in b. (h): Linear consistency between embeddings of calcium imaging and Neuropixels which were trained jointly using a multi-session CEBRA model. (i): Diagram of mouse primary visual cortex (V1, VIsp), PPC (VIsrl) and higher visual areas. (j): CEBRA-Behavior 32D model jointly trained with 2P+NP with 400 neurons then consistency measured within or across areas (2P vs. NP) across 2 unique sets of disjoint neurons for 3 seeds and averaged. (k): Models trained as in h, with intra-V1 consistency measurement vs. all inter-area vs. V1 comparison. Purple dots indicate mean of V1 intra-V1 consistency (across n=12 runs) and inter-V1 consistency (n=60). Intra-V1 consistency is significantly higher than inter-area consistency (Welch's t-test, T(19,53)=4.55, p=0.00019).
• Figure 5: Decoding of natural movie features from mouse visual cortical areas.(a): Schematic of the CEBRA encoder and kNN (or naive Bayes) decoder. (b): Examples of original frames (top row) and frames decoded from CEBRA embedding of V1 calcium recording using kNN decoding (bottom row). The last repeat among 10 repeats was used as the held-out test. (c): Decoding accuracy measured by considering a predicted frame being within 1 sec to the true frame as a correct prediction using CEBRA (NP only), jointly trained (2P+NP), or a baseline population-vector plus kNN or naive Bayes decoder using either a 1 frame (33 ms) receptive field or 10 frames (330 ms); results shown for Neuropixels dataset (V1 data). (d): Decoding accuracy measured by the correct scene prediction using either CEBRA (NP only), jointly trained (2P+NP), or baseline population-vector plus kNN or Bayes decoder using a 1 frame (33 ms) receptive field (V1 data). (e): Single frame ground truth frame ID vs predicted frame ID for Neuropixels using a CEBRA-Behavior model trained with a 330 ms receptive field (1,000 V1 neurons across mice used). (f): The mean absolute error of the correct frame index. Shown for baseline and CEBRA models as computed in c, d, e. (g): Diagram of the cortical areas considered, and decoding performance from CEBRA (NP only), 10 frame receptive field. (h): V1 decoding performance vs. layer category using 900 neurons with a 330 ms receptive field CEBRA-Behavior model.

Theorems & Definitions (26)

• Definition 1: Data generating process and encoder
• Definition 2: Generalized InfoNCE objective
• Definition 3: Generalized InfoNCE objective with limited batch size
• Proposition 1
• proof
• Definition 4: Diversity condition for bijectivity
• Proposition 2
• proof
• proof
• Proposition 3: CEBRA models are consistent.