Source-linked AI summary
Captum: A unified and generic model interpretability library for PyTorch
Narine Kokhlikyan, Vivek Miglani, Miguel Martin, Edward Wang, Bilal Alsallakh, Jonathan Reynolds, Alexander Melnikov, Natalia Kliushkina, Carlos Araya, Siqi Yan, Orion Reblitz-Richardson
TL;DR
Captum addresses limited PyTorch support for interpreting complex black-box neural networks. It unifies generic attribution algorithms and evaluation metrics across models and modalities, while providing scalable computation. The paper reports substantial runtime reductions with multi-GPU execution and presents applications including text classification and regression.
Problem
Existing interpretability frameworks had insufficient PyTorch support, while attribution evaluation remained difficult and model-understanding tools were needed beyond computer vision.
Method
Captum provides generic gradient- and perturbation-based feature, neuron, and layer attribution algorithms, evaluation metrics, scalable computation, and an interactive Captum Insights tool.
Results
Captum supports multiple neural-network applications and modalities, while increasing GPUs from 1 to 8 reduced Feature Ablation execution time by approximately 85% in a VGG19 experiment.
Takeaways & Limitations
The library is intended for model interpretability in research and production across classification and non-classification applications.
Takeaways & Limitations
Some algorithms have model-specific requirements; for example, GradCam and GuidedGradCam only make sense for convolutional models.
Abstract
from arXiv · showhide
In this paper we introduce a novel, unified, open-source model interpretability library for PyTorch [12]. The library contains generic implementations of a number of gradient and perturbation-based attribution algorithms, also known as feature, neuron and layer importance algorithms, as well as a set of evaluation metrics for these algorithms. It can be used for both classification and non-classification models including graph-structured models built on Neural Networks (NN). In this paper we give a high-level overview of supported attribution algorithms and show how to perform memory-efficient and scalable computations. We emphasize that the three main characteristics of the library are multimodality, extensibility and ease of use. Multimodality supports different modality of inputs such as image, text, audio or video. Extensibility allows adding new algorithms and features. The library is also designed for easy understanding and use. Besides, we also introduce an interactive visualization tool called Captum Insights that is built on top of Captum library and allows sample-based model debugging and visualization using feature importance metrics.
1 Introduction
Captum addresses the need to understand complex, black-box neural networks by providing a generic PyTorch interpretability library. It combines attribution algorithms with evaluation metrics and supports applications beyond computer vision.
- Understanding black-box neural networks is especially important in healthcare, finance, and self-driving vehicles, where trust in AI systems is critical.
- Existing interpretability frameworks had insufficient support for PyTorch models, motivating Captum’s generic implementations of attribution algorithms.Captum supports gradient- and perturbation-based methods across model modalities, with extensible and scalable computation.
- Captum includes feature, neuron, and layer importance algorithms for interpreting black-box neural networks.These attribution approaches are intended to make model behavior understandable in human terms.
- Captum provides generic implementations of infidelity and maximum-sensitivity metrics because attribution evaluation is difficult and often subjective or domain-specific.The metrics can be used with any PyTorch model and most attribution algorithms.
- Captum targets multiple domains because model-understanding research has focused largely on computer vision and CV-specific implementations are not always straightforward to adapt.The library is designed for use across research and production settings.
2 An Overview of the Algorithms
Captum organizes attribution methods into primary, neuron, and layer variants, with broad gradient and perturbation-based coverage. It also provides unified usage, baseline-aware methods, and memory-efficient execution for large inputs and multi-GPU systems.
- Attribution variants: Primary attributions map model outputs to inputs, neuron attributions map hidden neurons to inputs, and layer attributions map outputs to hidden-layer neurons.
- Algorithm families: Most Captum algorithms are gradient- or perturbation-based, while some methods belong to neither category.Figure 2 distinguishes primary and neuron methods from layer variants and uses color coding for algorithm families.
- Algorithm considerations: Baseline choice affects attribution quality for Integrated Gradients, DeepLift, SHAP variants, Feature Ablation, and Occlusion.Baselines represent the absence of an input feature and must be chosen carefully by the user.
- Usability: All algorithms share a unified API and signature, making it easier to compare methods and switch between attribution approaches.
- Scalability: Captum slices large inputs into chunks, aggregates sequential attribution results, and supports multi-GPU execution to avoid memory problems and improve runtime.Feature perturbation can also process multiple features together in one batch when sufficient memory is available.
- Scalability: Approximately 85%: Feature Ablation execution time drops when increasing GPUs from 1 to 8 for VGG19 with one 3 x 224 x 224 image.The Integrated Gradients experiment also showed substantial execution-time decreases as GPUs increased, with 2990 integral approximation steps held constant.
3 Evaluation
Captum provides generic, PyTorch-compatible metrics for evaluating attribution quality, addressing the difficulty of subjective visual explanations and domain-specific quantitative measures. Infidelity evaluates attribution agreement with output changes under perturbations, while maximum sensitivity measures attribution change under small input perturbations.
- Evaluation: Infidelity and maximum sensitivity provide generic quantitative evaluation metrics for feature, neuron and layer importances across PyTorch models and most Captum algorithms.These metrics address challenges in qualitative and quantitative attribution evaluation.
- Infidelity: Infidelity measures the discrepancy between attribution-weighted perturbations and the model output change between an input and its perturbed version.For global attributions, the baseline x0 represents the perturbed input, and the indicator function identifies features included in the perturbation.
- Infidelity: The infidelity formulation covers real-valued perturbations and local attribution functions, including Saliency, GuidedBackprop, Deconvolution, GradCam, GuidedGradCam, DeepLift and unscaled Integrated Gradients.Global attribution methods instead multiply saliency maps by x − x0 before infidelity is computed.
- Maximum sensitivity: Maximum sensitivity measures how much an attribution changes when the input is slightly perturbed within an Lp ball of radius r.It uses Monte-Carlo sampling, is a more robust variant of Lipschitz Continuity, and does not require attribution continuity everywhere.
- Evaluation: The two metrics are demonstrated in different applications.
4 Applications
Captum is applied across text classification, regression, and multimodal visual question answering, with attribution visualizations supporting analysis of inputs, neurons, and modalities.
- Text classification: Integrated Gradients identifies salient IMDB tokens and visualizes whether each token contributes toward or away from the predicted class.Color intensity represents the magnitude of each token's signal.
- Regression: Layer conductance attributes Boston house-price predictions to the ten neurons in the last linear layer, alongside their learned weights.Both normalized attribution scores and weights are plotted for comparison.
- Interactive debugging: Captum Insights interactively subsamples examples and attributes output classes to model inputs using multiple attribution algorithms.The tool is designed to improve model debugging experience.
- Multimodality: For multimodal visual question answering, aggregated attribution magnitudes indicate whether text or image provides the stronger predictive signal.The visualization supports comparing contributions across input modalities.
5 Conclusion
The paper presents Captum as a unified PyTorch interpretability library supporting generic attribution algorithms, scalable computation, evaluation metrics, multimodal applications, and model debugging.
- 5 Conclusion: Captum provides generic gradient- and perturbation-based attribution algorithms for PyTorch neural networks.The library covers feature, neuron, and layer attribution and is intended for research and production use.
- 5 Conclusion: The library scales attribution computations for large inputs and includes generic infidelity and maximum-sensitivity evaluation metrics.The paper also introduces Captum Insights for model debugging and demonstrates multimodal applications.
6 Future Work
Future work extends Captum beyond its current attribution methods toward broader model-understanding capabilities, including adversarial robustness, concepts, embeddings, and neuron debugging.
- 6 Future Work: Future work includes expanding the attribution-algorithm list and investigating model understanding beyond attribution methods.The authors specifically mention adversarial robustness and its intersection with attribution research.
- 6 Future Work: The authors identify concept-based interpretability as a direction for explaining models globally with human-understandable concepts.This is presented as an avenue for future exploration.
- 6 Future Work: Future directions also include visualizing high-dimensional latent embeddings and debugging what information individual or groups of neurons encode.These goals target deeper understanding of latent representations and neuron functions.