Add xCEBRA implementation (AISTATS 2025) #225

gonlairo · 2025-02-17T20:20:56Z

xCEBRA

eXplainable CEBRA 🔎🦓

This PR adds the following features:

multiobjective solver -> fit multiple subspaces with a new API
attribution methods (via captum), including our new method, inverted neuron gradient
regularized contrastive learning using jacobian regularization (required for identifiable attribution maps, but also useful for regularizing training more generally); aka xCEBRA

This code supports the following paper:

https://openreview.net/forum?id=aGrCXoTB4P

@inproceedings{
schneider2025timeseries,
title={Time-series attribution maps with regularized contrastive learning},
author={Steffen Schneider and Rodrigo Gonz{\'a}lez Laiz and Anastasiia Filippova and Markus Frey and Mackenzie W Mathis},
booktitle={The 28th International Conference on Artificial Intelligence and Statistics},
year={2025},
url={https://openreview.net/forum?id=aGrCXoTB4P}
}

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 (RegCL) trained on time-series data. We show theoretically that RegCL 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 better understanding of time-series data, such as for neural dynamics and decision-processes within neural networks.

Outline of the Method:

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
Jg (i.e., x is the observed neural data linked to latents z and c, where 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 L. Then, we estimate the Jacobian Jf of the encoder model (f) by minimizing a generalized
InfoNCE objective. Inverting this Jacobian J+f , which approximates Jg, allows us to construct the attributions.

* Add multiobjective solver and regularized training * Add example for multiobjective training * Add jacobian regularizer and SAM * update license headers * add api draft for multiobjective training * add all necessary modules to run the complete xcebra pipeline * add notebooks to reproduce xcebra pipeline * add first working notebook * add notebook with hybrid learning * add notebook with creation of synthetic data * add notebook with hybrid training * add plot with R2 for different parts of the embedding * add new API * update api wrapper with more checks and messages * add tests and notebook with new api * merge xcebra into attribution * separate xcebra dataset from cebra * some minor refactoring of cebra dataset * separate xcebra loader from cebra * remove xcebra distributions from cebra * minor refactoring with distributions * separate xcebra criterions from cebra * minor refactoring on criterion * separate xcebra models/criterions/layers from cebra * refactoring multiobjective * more refactoring... * separate xcebra solvers from cebra * more refactoring * move xcebra to its own package * move more files into xcebra package * more files and remove changes with the registry * remove unncessary import * add folder structure * move back distributions * add missing init * remove wrong init * make loader and dataset run with new imports * making it run! * make attribution run * Run pre-commit * move xcebra repo one level up * update gitignore and add __init__ from data * add init to distributions * add correct init for attribution pacakge * add correct init for model package * fix remaining imports * fix tests * add examples back to xcebra repo * update imports from graphs_xcebra * add setup.py to create a package * update imports of graph_xcebra * update notebooks * Formatting code for submission Co-authored-by: Rodrigo Gonzalez <[email protected]> * move test into xcebra * Add README * move distributions back to main package * clean up examples * adapt tests * Add LICENSE * add train/eval notebook again * add notebook with clean results * rm synthetic data * change name from xcebra to regcl * change names of modules and adapt imports * change name from graphs_xcebra to synthetic_data * Integrate into CEBRA * Fix remaining imports and make notebook runnable * Add dependencies, add version flag * Remove synthetic data files * reset dockerfile, move vmf * apply pre-commit * Update notice * add some docstrings * Apply license headers * add new scd notebook * add notebook with scd --------- Co-authored-by: Steffen Schneider <[email protected]>

* bump version * update dockerfile * fix progress bar * remove outdated test * rename models

MMathisLab · 2025-02-17T21:02:28Z

Let's move demo /examples to https://github.com/AdaptiveMotorControlLab/CEBRA-demos :D

MMathisLab

I left comments throughout, thank you!!

setup.cfg

tests/test_models.py

Dockerfile

MMathisLab · 2025-02-17T21:03:49Z

NOTICE.yml

@@ -35,3 +35,83 @@
    - 'tests/**/*.py'
    - 'docs/**/*.py'
    - 'conda/**/*.yml'
+


I think best just in the related files, not here?

MMathisLab · 2025-02-17T21:05:13Z

cebra/attribution/attribution_models.py

+import torch
+import torch.nn as nn
+import tqdm
+from captum.attr import NeuronFeatureAblation


Do we really need all the baselines in this package? We should write up docs for this also ...

MMathisLab · 2025-02-17T21:16:22Z

cebra/solver/util.py

@@ -29,7 +30,7 @@


 def _description(stats: Dict[str, float]):
-    stats_str = [f"{key}: {value: .4f}" for key, value in stats.items()]
+    stats_str = [f"{key}: {value:.3f}" for key, value in stats.items()]


why this change?

see above, will revert back

MMathisLab · 2025-02-17T21:16:34Z

cebra/solver/util.py

@@ -73,7 +74,9 @@ class ProgressBar:
    "Log and display values during training."

    loader: Iterable
-    log_format: str
+    logger: logging.Logger = None


more details needed

MMathisLab · 2025-02-17T21:16:55Z

examples/synthetic_data.pkl

remove and move to demos

MMathisLab · 2025-02-17T21:17:05Z

examples/train_and_evaluate.ipynb

remove and move to demos

MMathisLab · 2025-02-17T21:17:14Z

examples/train_and_evaluate_scd.ipynb

remove and move to demos

stes

I made a few edits:

The old MultiobjectiveSolver is now again accessible (important for the hybrid model in Fig 2 in CEBRA paper), it is now called LegacyMultiobjectiveSolver. For an end user using the sklearn API with hybrid=True, this can now still be used.
Reverted some changes from the research code base not important for release
added additional tests, incl. integration tests from the notebooks
(resolved some additional review comments)

tests/test_models.py

setup.cfg

stes · 2025-02-19T00:05:08Z

cebra/solver/single_session.py

@@ -130,6 +131,16 @@ def _inference(self, batch):
 class SingleSessionHybridSolver(abc_.MultiobjectiveSolver):
    """Single session training, contrasting neural data against behavior."""

+    log: Dict = dataclasses.field(default_factory=lambda: ({


I will revert some of these changes back, were part of the research code base and should not go to the package...

stes · 2025-02-19T00:05:17Z

cebra/solver/util.py

@@ -29,7 +30,7 @@


 def _description(stats: Dict[str, float]):
-    stats_str = [f"{key}: {value: .4f}" for key, value in stats.items()]
+    stats_str = [f"{key}: {value:.3f}" for key, value in stats.items()]


see above, will revert back

stes · 2025-02-19T00:31:59Z

.github/workflows/build.yml

+      # NOTE(stes): Temporarily disable, INCLUDE BEFORE MERGE!
+      #- name: Check that no binary files have been added to repo
+      #  if: matrix.os == 'ubuntu-latest'
+      #  run: |
+      #    make check_for_binary


pragmatic workaround until i moved the demo files. will remove once fixed

stes · 2025-02-19T00:34:48Z

cebra/data/datasets.py

+            negative=None,
+        )
+
+    def load_batch_contrastive(self, index: BatchIndex) -> Batch:


That issue only appears once #168 is used, so I think not of concern here. This function simply constructs the batch from indices, this should be unrelated to the batched inference issue

stes · 2025-02-19T00:35:55Z

cebra/datasets/__init__.py

@@ -96,7 +96,6 @@ def get_datapath(path: str = None) -> str:
    from cebra.datasets.gaussian_mixture import *
    from cebra.datasets.hippocampus import *
    from cebra.datasets.monkey_reaching import *
-    from cebra.datasets.synthetic_data import *


fix 5e30829

stes · 2025-02-19T01:17:39Z

next:

check docstring coverage and write docstrings
remove ratinabox==1.8 and ephysiopy==1.9.62 deps for xcebra? I think not needed for core functionality. when dropping, could be added to the docs somewhere
fix and consolidate naming of some of the newly added classes
cebra.distributions.DeltaVMFDistribution seems missing, check!

stes · 2025-02-19T01:19:28Z

tests build!

gonlairo and others added 6 commits February 17, 2025 21:16

Fix tests

cbcb229

* bump version * update dockerfile * fix progress bar * remove outdated test * rename models

Apply fixes to pass ruff tests

97ad03f

Fix typos

36282d0

Update license headers, fix additional ruff errors

5fa5b42

remove unused comment

9ef548d

stes changed the title ~~Aistats2025~~ Add xCEBRA implementation (AISTATS 2025) Feb 17, 2025

MMathisLab requested changes Feb 17, 2025

View reviewed changes

MMathisLab added the xCEBRA🦓 label Feb 17, 2025

stes assigned stes and gonlairo Feb 17, 2025

Merge branch 'main' into aistats2025

a9e7027

cla-bot bot added the CLA signed label Feb 18, 2025

stes added 9 commits February 19, 2025 00:27

rename regcl in codebase

a2f7117

change regcl name in dockerfile

761f373

Improve attribution module

2052ec9

Fix imports name naming

5e483d0

add basic integration test

4254173

temp disable of binary check

1d69957

Add legacy multiobjective model for backward compat

1d10668

add synth import back in

5e30829

Fix docstrings and type annot in cebra/models/jacobian_regularizer.py

458958f

stes reviewed Feb 19, 2025

View reviewed changes

stes added 3 commits February 19, 2025 01:52

add xcebra to tests

d906ad5

add missing cvxpy dep

6f91018

fix docstrings

df4f661

more docstrings to fix attr error

d81b93d

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add xCEBRA implementation (AISTATS 2025) #225

Add xCEBRA implementation (AISTATS 2025) #225

gonlairo commented Feb 17, 2025 •

edited by stes

Loading

MMathisLab commented Feb 17, 2025

MMathisLab left a comment

MMathisLab Feb 17, 2025

MMathisLab Feb 17, 2025

MMathisLab Feb 17, 2025

stes Feb 19, 2025

MMathisLab Feb 17, 2025

MMathisLab Feb 17, 2025

MMathisLab Feb 17, 2025

MMathisLab Feb 17, 2025

stes left a comment

stes Feb 19, 2025

stes Feb 19, 2025

stes Feb 19, 2025

stes Feb 19, 2025

stes Feb 19, 2025

stes commented Feb 19, 2025 •

edited

Loading

stes commented Feb 19, 2025

Add xCEBRA implementation (AISTATS 2025) #225

Are you sure you want to change the base?

Add xCEBRA implementation (AISTATS 2025) #225

Conversation

gonlairo commented Feb 17, 2025 • edited by stes Loading

xCEBRA

eXplainable CEBRA 🔎🦓

This code supports the following paper:

Abstract:

Outline of the Method:

MMathisLab commented Feb 17, 2025

MMathisLab left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

stes left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

stes commented Feb 19, 2025 • edited Loading

stes commented Feb 19, 2025

gonlairo commented Feb 17, 2025 •

edited by stes

Loading

stes commented Feb 19, 2025 •

edited

Loading