Explainable and consistent embeddings of high-dimensional recordings using auxiliary variables

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 (Urai et al., 2022; Krakauer et al., 2017). 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 contrastive learning method that jointly uses time-series data and auxiliary labels in either a (supervised) hypothesis- or (self-supervised) discovery-driven manner to produce consistent, high-performance latent spaces. This method is particularly useful for joint modeling of neural and behavioral data.

To benchmark our method, we first consider variants of a synthetic neural dataset that have different noise and firing characteristics. We find that our contrastive learning method outperforms generative models (pi-VAE; Zhou & Wei, 2020) and common tools for dimensionality reduction (tSNE and UMAP), in both the recovery of the ground truth latents and in the variance across runs (R² linear regression between true & recovered embeddings, mean/SEM: 92.0/0.02: ours, 84.1/0.62: pi-VAE, 80.9/0.04: UMAP, 78.8/0.30: tSNE).

Critically, such non-linear methods should be identifiable, and explainable. Therefore, we sought to explore the explainablity of our method. Since most ``explainable AI'' methods lack theoretical guarantees, it is difficult to judge whether they have produced the correct result. Therefore, we design another synthetic dataset with a tractable ground-truth attribution map between neurons and auxiliary variables. Then, we leverage theoretical guarantees of our learning objective to estimate this attribution map from the Jacobian matrix of the learned feature encoder. Empirically, our model estimates the connections between neurons and auxiliary variables with high accuracy and outperforms a supervised learning baseline.