Attribution Methods#

Attribution methods for CEBRA.

This module was added in v0.6.0 and contains attribution methods described and benchmarked in [Schneider2025].

[Schneider2025]

Schneider, S., González Laiz, R., Filippova, A., Frey, M., & Mathis, M. W. (2025). Time-series attribution maps with regularized contrastive learning. The 28th International Conference on Artificial Intelligence and Statistics. https://openreview.net/forum?id=aGrCXoTB4P

Different attribution methods#

class cebra.attribution.attribution_models.AttributionMap(model, input_data, output_dimension=None, num_samples=None, seed=9712341)#

Bases: object

Base class for computing attribution maps for CEBRA models.

Parameters:

  • model (Module) – The trained CEBRA model to analyze

  • input_data (Tensor) – Input data tensor to compute attributions for

  • output_dimension (Optional[int]) – Output dimension to analyze. If None, uses model’s output dimension

  • num_samples (Optional[int]) – Number of samples to use for attribution. If None, uses full dataset

  • seed (int) – Random seed which is used to subsample the data. Only relevant if num_samples is not None.

compute_attribution_map()#

Compute the attribution map for the model.

Returns:

Attribution maps and their variants

Return type:

dict

Raises:

NotImplementedError – Must be implemented by subclasses

compute_metrics(attribution_map, ground_truth_map)#

Compute metrics comparing attribution map to ground truth.

This function computes various statistical metrics to compare the attribution values between connected and non-connected neurons based on a ground truth connectivity map. It separates the attribution values into two groups based on the binary ground truth, and calculates summary statistics and differences between these groups.

Parameters:

  • attribution_map – Computed attribution values representing the strength of connections between neurons

  • ground_truth_map – Binary ground truth connectivity map where True indicates a connected neuron and False indicates a non-connected neuron

Returns:

Dictionary containing the following metrics:

  • max/mean/min_nonconnected: Statistics for non-connected neurons

  • max/mean/min_connected: Statistics for connected neurons

  • gap_max: Difference between max connected and max non-connected values

  • gap_mean: Difference between mean connected and mean non-connected values

  • gap_min: Difference between min connected and min non-connected values

  • gap_minmax: Difference between min connected and max non-connected values

  • max/min_jacobian: Global max/min values across all neurons

Return type:

dict

compute_attribution_score(attribution_map, ground_truth_map)#

Compute ROC AUC score between attribution map and ground truth.

Parameters:

  • attribution_map – Computed attribution values

  • ground_truth_map – Binary ground truth connectivity map

Returns:

ROC AUC score

Return type:

float

check_moores_penrose_conditions(jacobian, jacobian_pseudoinverse)#

Check Moore-Penrose conditions for Jacobian matrices.

Parameters:

  • jacobian (ndarray) – Jacobian matrices of shape (num samples, output_dim, num_neurons)

  • jacobian_pseudoinverse (ndarray) – Pseudoinverse matrices of shape (num samples, num_neurons, output_dim)

Return type:

ndarray

Returns:

Boolean array of shape (num samples, 4) indicating satisfied conditions

class cebra.attribution.attribution_models.JFMethodBased(model, input_data, output_dimension=None, num_samples=None, seed=9712341)#

Bases: AttributionMap

Compute the attribution map using the Jacobian of the model encoder.

compute_attribution_map()#

Compute the attribution map for the model.

Returns:

Attribution maps and their variants

Return type:

dict

Raises:

NotImplementedError – Must be implemented by subclasses

class cebra.attribution.attribution_models.JFMethodBasedBatched(model, input_data, output_dimension=None, num_samples=None, seed=9712341)#

Bases: JFMethodBased

Compute an attribution map based on the Jacobian using mini-batches.

See also

JFMethodBased

compute_attribution_map(batch_size=1024)#

Compute the attribution map for the model.

Returns:

Attribution maps and their variants

Return type:

dict

Raises:

NotImplementedError – Must be implemented by subclasses

class cebra.attribution.attribution_models.NeuronGradientMethod(model, input_data, output_dimension=None, num_samples=None, seed=9712341)#

Bases: AttributionMap

Compute the attribution map using the neuron gradient from Captum.

Note

This method is equivalent to Jacobian-based attributions, but uses a different backend implementation.

compute_attribution_map(attribute_to_neuron_input=False)#

Compute the attribution map for the model.

Returns:

Attribution maps and their variants

Return type:

dict

Raises:

NotImplementedError – Must be implemented by subclasses

class cebra.attribution.attribution_models.NeuronGradientMethodBatched(model, input_data, output_dimension=None, num_samples=None, seed=9712341)#

Bases: NeuronGradientMethod

As NeuronGradientMethod, but using mini-batches.

See also

NeuronGradientMethod

compute_attribution_map(attribute_to_neuron_input=False, batch_size=1024)#

Compute the attribution map for the model.

Returns:

Attribution maps and their variants

Return type:

dict

Raises:

NotImplementedError – Must be implemented by subclasses

class cebra.attribution.attribution_models.FeatureAblationMethod(model, input_data, output_dimension=None, num_samples=None, seed=9712341)#

Bases: AttributionMap

Compute the attribution map using the feature ablation method from Captum.

compute_attribution_map(baselines=None, feature_mask=None, perturbations_per_eval=1, attribute_to_neuron_input=False)#

Compute the attribution map for the model.

Returns:

Attribution maps and their variants

Return type:

dict

Raises:

NotImplementedError – Must be implemented by subclasses

class cebra.attribution.attribution_models.FeatureAblationMethodBAtched(model, input_data, output_dimension=None, num_samples=None, seed=9712341)#

Bases: FeatureAblationMethod

As FeatureAblationMethod, but using mini-batches.

See also

FeatureAblationMethod

compute_attribution_map(baselines=None, feature_mask=None, perturbations_per_eval=1, attribute_to_neuron_input=False, batch_size=1024)#

Compute the attribution map for the model.

Returns:

Attribution maps and their variants

Return type:

dict

Raises:

NotImplementedError – Must be implemented by subclasses

class cebra.attribution.attribution_models.IntegratedGradientsMethod(model, input_data, output_dimension=None, num_samples=None, seed=9712341)#

Bases: AttributionMap

Compute the attribution map using the integrated gradients method from Captum.

compute_attribution_map(n_steps=50, method='gausslegendre', internal_batch_size=None, attribute_to_neuron_input=False, baselines=None)#

Compute the attribution map for the model.

Returns:

Attribution maps and their variants

Return type:

dict

Raises:

NotImplementedError – Must be implemented by subclasses

class cebra.attribution.attribution_models.IntegratedGradientsMethodBatched(model, input_data, output_dimension=None, num_samples=None, seed=9712341)#

Bases: IntegratedGradientsMethod

As IntegratedGradientsMethod, but using mini-batches.

See also

IntegratedGradientsMethod

compute_attribution_map(n_steps=50, method='gausslegendre', internal_batch_size=None, attribute_to_neuron_input=False, baselines=None, batch_size=1024)#

Compute the attribution map for the model.

Returns:

Attribution maps and their variants

Return type:

dict

Raises:

NotImplementedError – Must be implemented by subclasses

class cebra.attribution.attribution_models.NeuronGradientShapMethod(model, input_data, output_dimension=None, num_samples=None, seed=9712341)#

Bases: AttributionMap

Compute the attribution map using the neuron gradient SHAP method from Captum.

compute_attribution_map(baselines, n_samples=5, stdevs=0.0, attribute_to_neuron_input=False)#

Compute the attribution map for the model.

Returns:

Attribution maps and their variants

Return type:

dict

Raises:

NotImplementedError – Must be implemented by subclasses

class cebra.attribution.attribution_models.NeuronGradientShapMethodBatched(model, input_data, output_dimension=None, num_samples=None, seed=9712341)#

Bases: NeuronGradientShapMethod

As NeuronGradientShapMethod, but using mini-batches.

See also

NeuronGradientShapMethod

compute_attribution_map(baselines, n_samples=5, stdevs=0.0, attribute_to_neuron_input=False, batch_size=1024)#

Compute the attribution map for the model.

Returns:

Attribution maps and their variants

Return type:

dict

Raises:

NotImplementedError – Must be implemented by subclasses

Jacobian-based attribution#

Tools for computing attribution maps.

cebra.attribution.jacobian_attribution.get_attribution_map(model, input_data, double_precision=True, convert_to_numpy=True, aggregate='mean', transform='none', hybrid_solver=False)#

Estimate attribution maps using the Jacobian pseudo-inverse.

The function estimates Jacobian matrices for each point in the model, computes the pseudo-inverse (for every sample) and then aggregates the resulting matrices to compute an attribution map.

Parameters:

  • model (Module) – The neural network model for which to compute attributions.

  • input_data (Tensor) – Input tensor or numpy array to compute attributions for.

  • double_precision (bool) – If True, use double precision for computation.

  • convert_to_numpy (bool) – If True, convert the output to numpy arrays.

  • aggregate (Literal[‘mean’, ‘sum’, ‘max’]) – Method to aggregate attribution values across samples. Options are "mean", "sum", or "max".

  • transform (Literal[‘none’, ‘abs’]) – Transformation to apply to attribution values. Options are "none" or "abs".

  • hybrid_solver (bool) – If True, handle multi-objective models differently.

Returns:

A tuple containing the Jacobian matrix of shape (num_samples, output_dim, input_dim) and the pseudo-inverse of the Jacobian matrix.