Time-series attribution maps with regularized contrastive learning

TL;DR

The paper addresses the lack of identifiability in gradient-based time-series attribution by introducing xCEBRA, a regularized contrastive learning framework that uses Inverted Neuron Gradient to recover the data-generating Jacobian $oldsymbol{J}_{oldsymbol{g}}$ up to a linear indeterminacy. It formalizes identifiability concepts for time-series attribution maps and proves two theorems: a goodness-of-fit result for the contrastive objective and an identifiability result for the Jacobian, linking the pseudo-inverse Jacobian $oldsymbol{J}_{oldsymbol{f}}^{+}(oldsymbol{x})$ to the ground-truth map $oldsymbol{A}_{oldsymbol{g}}$. Through extensive synthetic and RatInABox neural-data experiments, the method consistently outperforms gradient-based and Shapley baselines and demonstrates reliable dimensionality identification and applicability to neural dynamics. The approach offers a principled, data-generating–centered framework for time-series attribution with potential for broad use in neuroscience and related domains. Key constructs include the ground-truth attribution map $oldsymbol{A}_{oldsymbol{g}}$ tied to $oldsymbol{J}_{oldsymbol{g}}$, subspace identifiability of encoders, and the inverted neuron gradient $oldsymbol{J}_{oldsymbol{f}}^{+}(oldsymbol{x})$ as a practical attribution tool.

Abstract

Gradient-based attribution methods aim to explain decisions of deep learning models but so far lack identifiability guarantees. Here, we propose a method to generate attribution maps with identifiability guarantees by developing a regularized contrastive learning algorithm trained on time-series data plus a new attribution method called Inverted Neuron Gradient (collectively named xCEBRA). We show theoretically that xCEBRA has favorable properties for identifying the Jacobian matrix of the data generating process. Empirically, we demonstrate robust approximation of zero vs. non-zero entries in the ground-truth attribution map on synthetic datasets, and significant improvements across previous attribution methods based on feature ablation, Shapley values, and other gradient-based methods. Our work constitutes a first example of identifiable inference of time-series attribution maps and opens avenues to a better understanding of time-series data, such as for neural dynamics and decision-processes within neural networks.

Figures (10)

• Figure 1: Identifiable attribution maps for time-series data. Using time-series data (such as neural data recorded during navigation, as depicted), our inference framework estimates the ground-truth Jacobian matrix $\mathbf{J}_\mathbf{g}$ (i.e., $\mathbf{x}$ is the observed neural data linked to latents $\mathbf{z}$ and $\mathbf{c}$, where $\mathbf{c}$ is the explicit [auxiliary] behavioral variable that would be linked to grid cells) by identifying the inverse data generation process up to a linear indeterminacy $\mathbf{L}$. Then, we estimate the Jacobian $\mathbf{J}_\mathbf{f}$ of the encoder model ($\mathbf{f}$) by minimizing a generalized InfoNCE objective. Inverting this Jacobian $\mathbf{J}_\mathbf{f}^+$, which approximates $\mathbf{J}_\mathbf{g}$, allows us to construct the attributions.
• Figure 2: Left: Graphical model for the data generating process where $\mathbf{z}_2$ is observed through $\mathbf{c}_2$. The attribution map needs to be computed with respect to $\mathbf{z}_2$, which is inferred with supervised (contrastive) learning. Note, practically, this means $\mathbf{x}_2$ is behaviorally linked to $\mathbf{c}_2$ (denoted by dashed line). Related to Table \ref{['tbl:summary']}. Right: Graphical model for the data generating process where $\mathbf{z}_1$ is observed through $\mathbf{c}_1$. Since $\mathbf{z}_2$ is not observed, the attribution map can only be estimated through the time-contrastive component in xCEBRA. Related to Table \ref{['table:latent-attribution']}.
• Figure 3: Hybrid Regularized Contrastive Learning+Inverted Neuron Gradient (xCEBRA; Ours, black) and supervised baselines auROC vs. dimension of latent factors. Two latent factors are observable as auxiliary variables in all experiments.
• Figure 4: Attribution scores of synthetic cell types.a, the synthetic 4-cell type neural data, the simulated navigation and computed speed/head direction. b, embedding space is jointly trained with behavioral information about animal position (first 4 dimensions, top) and additional time-varying latent information (the remaining 10 dimensions) with our regularized hybrid contrastive learning setting. The position information was decoded as indicated by cross-validated $R^2$ score on held-out data. Training embedding is shown. c, time-series attribution map, showing high scores (lighter) for position. d, Attribution scores, zero-centered & standardized across cells. e, auROC across training.
• Figure 5: Synthetic Data Generation Process. We generate two sets of latent variables, $z_1$ and $z_2$, each consisting of 100,000 samples drawn from Brownian motion within a box $[-1,1]^d$. In this example, $z_1$ is connected to both $x_1$ and $x_2$, while $z_2$ is connected only to $x_2$. Additionally, we use an injective mixing function consisting of $g_1$ and $g_2$. Function $g_1$ takes 3 (denoted $d_1$) latent variables as input and outputs 25 neurons (denoted $n_1$), whereas $g_2$ takes 6 ($d_1+d2$) latent variables as input and outputs 25 neurons (denoted $n_2$). The final data $x$ is constructed by concatenating $x_1$ and $x_2$, resulting in a data matrix $x$ with a shape of 100,000 by 50.

Theorems & Definitions (13)

• Definition 1: Data generating process
• Definition 2: Subspace Identifiability
• Definition 3: Identifiability of connectivity in attribution maps
• Definition 4: Ground truth attribution map of the mixing function
• Theorem 1
• proof
• Theorem 2
• Proposition 1: restated from schneider2023cebra
• proof
• Proposition 2: restated from Proposition 6 in schneider2023cebra