Captum Tutorial: Integrated Gradients for PyTorch Models

Contents

Figure 1: Captum Tutorial: Integrated Gradients for PyTorch Models — a heat map that looks convincing, and the convergence delta that decides whether it deserves belief

A trained network tells you what it predicts, not why. Captum is PyTorch's interpretability library, and its best-known method, Integrated Gradients, answers one specific question: how much did each input feature contribute to this prediction, relative to a reference input?

This tutorial covers how IG works, runs it on tabular, image and text models, checks that the attributions are numerically sound, and finishes with the pitfalls that make attribution maps misleading.

One update first: Captum 0.9.0 shipped on April 17, 2026. It adds NumPy 2.x support, retires Captum Insights along with the captum[insights] install extra, raises the minimum versions to Python 3.10 and PyTorch 1.13, and makes cross-tensor attribution default-on for perturbation methods such as FeatureAblation and FeaturePermutation. In my explainability tools guide I listed Captum Insights as an interactive option, and that is no longer current: you would have to pin the older 0.8.0 release to use it. The repository now lives under meta-pytorch/captum.

What Integrated Gradients Does

Plain gradients, often called saliency, have a flaw: a ReLU network can saturate, so the gradient at the actual input is near zero even for a feature that matters enormously. A saturated unit reports no sensitivity exactly where the model is most committed.

IG avoids that by accumulating gradients along a path. Pick a baseline x′, meaning a no-information reference, draw a straight line from the baseline to your input x, and integrate the gradients of the output along that line:

IGᵢ(x) = (xᵢ − x′ᵢ) × ∫₀¹ ∂F(x′ + α(x − x′)) / ∂xᵢ dα

Captum's documentation describes it as an axiomatic method that assigns each input feature an importance score by approximating this integral, with baselines supplied as arguments. In practice the integral is approximated by a sum over n_steps points, and Captum's default is 50.

Two properties make IG attractive:

• Completeness: the attributions sum to the difference between the model's output at the input and at the baseline, F(x) − F(x′). Nothing is lost or invented, and you can check it numerically.
• Sensitivity and implementation invariance: if changing a feature changes the output, that feature gets nonzero attribution, and two functionally identical networks get identical attributions.

The convergence delta Captum returns measures how far the approximation is from completeness. Per its documentation, the lower the absolute value, the better the approximation, which makes it a built-in health check rather than an optional extra.

Install

pip install captum torch torchvision transformers pillow matplotlib

You need Python 3.10 or newer with Captum 0.9.0. Pin your Captum, PyTorch and Python versions together, because the supported range has changed across releases.

Tabular Example: IG From Scratch to Verification

I train a small network where the ground truth is known: features 0 and 1 drive the label, features 2 and 3 matter only through their product. That lets me check whether IG recovers a structure I planted.

import torch
import torch.nn as nn
from captum.attr import IntegratedGradients

torch.manual_seed(0)
X = torch.randn(4000, 4)
y = ((2 * X[:, 0] - X[:, 1] + 0.5 * X[:, 2] * X[:, 3]) > 0).long()

model = nn.Sequential(nn.Linear(4, 32), nn.ReLU(), nn.Linear(32, 2))
opt = torch.optim.Adam(model.parameters(), lr=0.01)
for _ in range(300):
    opt.zero_grad()
    loss = nn.functional.cross_entropy(model(X), y)
    loss.backward()
    opt.step()
model.eval()
print("train accuracy:", (model(X).argmax(1) == y).float().mean().item())

ig = IntegratedGradients(model)
x = X[:5]
baseline = torch.zeros_like(x)          # the average row, for standardized data

attr, delta = ig.attribute(
    x, baselines=baseline, target=1, n_steps=100, return_convergence_delta=True
)
print("attributions:\n", attr)
print("convergence delta:", delta)

with torch.no_grad():
    output_gap = model(x)[:, 1] - model(baseline)[:, 1]
print("max completeness error:", (attr.sum(dim=1) - output_gap).abs().max().item())

Several details in that snippet matter more than the code:

1. target=1 selects which output you explain. For classifiers you attribute a specific class, usually the predicted one. Attributions for different targets are different explanations.
2. The example attributes logits, because the model outputs raw scores. If you want attributions in probability space, wrap the model in a function that applies softmax. Softmax saturation shrinks attributions when the model is already confident, so attributing logits is often the more informative choice.
3. The completeness error should be near zero. If it is not, raise n_steps until it is.
4. Average the absolute attributions over many rows to get a global view. Here features 0 and 1 should dominate, with 2 and 3 smaller, which tells you whether the method matches what you know about the data.

global_importance = ig.attribute(
    X[:500], baselines=torch.zeros_like(X[:500]), target=1, n_steps=50
).abs().mean(dim=0)
print(global_importance)

The completeness check is the reason I start every attribution project with data where I know the answer. It is the only moment when you can verify the method against ground truth instead of against your expectations.

Choosing a Baseline

The baseline is the most consequential choice in IG, and it is the easiest one to underestimate. Attributions answer "compared to what?", and a different baseline gives a different answer.

Modality	Common baseline	The trap
Tabular	Zeros after standardization, means, medians, a typical row	For raw features such as age or income, zero is absurd
Images	A zeros tensor, a blurred copy of the input, uniform noise	After mean and standard deviation normalization, zeros is the mean colour, not black
Text	Padding tokens, keeping special tokens in place	Removing the start and end markers changes what the model sees

• Tabular: zeros only make sense if zero is a meaningful reference. After standardization it represents the average row, but on raw data it usually does not.
• Images: a black image is the common default, but it is rarely what you think it is once normalization is applied. A blurred version of the input and uniform noise are the other two worth trying.
• Text: typically padding tokens, with the special tokens preserved.
• Multiple baselines: averaging over a distribution of baselines, as Captum's GradientShap does, reduces dependence on any single choice.

A practical habit is to run IG with two or three baselines and check whether the conclusions agree. A feature that looks important under one baseline and irrelevant under another is telling you the explanation depends on the reference, not only on the model.

Image Example: A Pretrained ResNet

This example explains an ImageNet classifier's prediction on one image. Put any JPEG at dog.jpg:

import numpy as np
import torch
from PIL import Image
from torchvision import models, transforms
from captum.attr import IntegratedGradients, NoiseTunnel
from captum.attr import visualization as viz

weights = models.ResNet18_Weights.DEFAULT
model = models.resnet18(weights=weights).eval()
preprocess = weights.transforms()

img = Image.open("dog.jpg").convert("RGB")
x = preprocess(img).unsqueeze(0)

with torch.no_grad():
    pred = model(x).argmax(1).item()
print("prediction:", weights.meta["categories"][pred])

ig = IntegratedGradients(model)
nt = NoiseTunnel(ig)                        # SmoothGrad-style noise averaging
attr = nt.attribute(
    x, nt_type="smoothgrad_sq", nt_samples=8, stdevs=0.1,
    target=pred, baselines=torch.zeros_like(x), n_steps=32, internal_batch_size=8,
)

attr_np = np.transpose(attr.squeeze(0).detach().numpy(), (1, 2, 0))
display = transforms.Compose([transforms.Resize(256), transforms.CenterCrop(224)])(img)
orig_np = np.asarray(display) / 255.0       # matches ResNet18's resize and crop

viz.visualize_image_attr_multiple(
    attr_np, orig_np,
    methods=["original_image", "heat_map"],
    signs=["all", "positive"],
    show_colorbar=True,
)

Notes on this example:

• NoiseTunnel wraps an attribution method and averages it over noisy copies of the input, which cleans up the speckle typical of raw gradient-based maps. Its cost multiplies the forward and backward passes, so keep nt_samples and n_steps modest at first and use internal_batch_size to control memory.
• The zero baseline here is the mean colour in normalized space. Try a blurred-input baseline too and compare.
• Make sure the displayed image matches the preprocessing crop, or the heat map will not line up with the picture.
• signs="positive" shows only evidence for the class. Look at negative attributions as well, because regions that argued against the class matter when you are debugging.
• model.eval() is essential. Training-mode batch norm and dropout change both the predictions and the gradients.

Text: Layer Integrated Gradients on Token Embeddings

Token IDs are discrete, so you cannot interpolate between them. Instead, LayerIntegratedGradients runs IG on the output of a layer, typically the embedding layer, and then you sum across the embedding dimension to get one score per token:

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from captum.attr import LayerIntegratedGradients

name = "distilbert-base-uncased-finetuned-sst-2-english"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForSequenceClassification.from_pretrained(name).eval()

def forward(input_ids, attention_mask):
    return model(input_ids=input_ids, attention_mask=attention_mask).logits

text = "The plot was predictable, but the acting was wonderful."
enc = tok(text, return_tensors="pt")

baseline_ids = torch.full_like(enc.input_ids, tok.pad_token_id)
baseline_ids[0, 0] = tok.cls_token_id          # keep the special tokens
baseline_ids[0, -1] = tok.sep_token_id

lig = LayerIntegratedGradients(forward, model.distilbert.embeddings)
attr = lig.attribute(
    inputs=enc.input_ids, baselines=baseline_ids,
    additional_forward_args=(enc.attention_mask,),
    target=1, n_steps=50,                      # class 1 = POSITIVE for this model
)

scores = attr.sum(dim=-1).squeeze(0)
scores = scores / scores.norm()
for token, s in zip(tok.convert_ids_to_tokens(enc.input_ids[0]), scores.tolist()):
    print(f"{token:>14} {s:+.3f}")

You will typically see words like "wonderful" carrying large positive scores while "predictable" pulls the other way, which is a quick sanity check on whether the model reads the sentence the way you do. Remember that attribution on embeddings describes how sensitive the output is to that token's embedding along the path from the baseline, not comprehension.

For generative language models, Captum added dedicated wrappers from version 0.7: LLMAttribution works with most attribution methods and lets you define interpretable text features with baselines and masks, so you can see how parts of a prompt affect a predicted response. Version 0.9.0 expands what LLMAttribution can attribute over, including remote attribution through vLLM. The API is evolving, so follow the current documentation rather than copying an older tutorial.

Other Captum Methods Worth Knowing

IG is one tool among many, and comparing methods is good practice:

Method	Needs gradients	Cost	Use it when
Saliency, InputXGradient	Yes	Single pass	You want a fast first look and accept noise
DeepLift, GradientShap	Yes	Baseline-based, cheap	You want similar quantities with other baselines
Occlusion, FeatureAblation	No	Many forward passes	The pipeline is not differentiable
LayerGradCam	Yes	Cheap	You want class-activation maps for CNNs
LayerConductance, NeuronConductance	Yes	Moderate	You need to know which hidden units matter

• Saliency and InputXGradient are cheap, single-pass gradient methods, and they are prone to noise and saturation, which is the problem IG was designed to reduce.
• DeepLift and GradientShap are baseline-based methods that approximate similar quantities. DeepLift needs a single pass and returns its own convergence delta.
• Occlusion and FeatureAblation remove or mask features and watch the output change. They need no gradients, so they work with non-differentiable pipelines, at the price of more forward passes. In 0.9.0, cross-tensor attribution is on by default for these methods, so multi-input models are no longer attributed tensor by tensor unless you opt out.
• LayerGradCam produces class-activation-style maps for convolutional networks.
• LayerConductance and NeuronConductance attribute to internal layers and neurons, which helps answer which hidden units matter.

If several methods agree, you can trust the picture more. If they disagree wildly, investigate before presenting any of them.

How to Check That Attributions Are Trustworthy

An attribution map is a hypothesis about the model, and you should test it.

1. Completeness and convergence. Check the convergence delta, and increase n_steps until it is small.
2. Baseline sensitivity. Rerun with different baselines and compare the conclusions.
3. Parameter randomization sanity check. Randomize the model's weights layer by layer and recompute the attributions. If the maps barely change, the method is reflecting the structure of the input, such as edges, rather than what the model learned. Research on saliency sanity checks showed that some methods fail this test outright.
4. Quantitative metrics. Captum includes infidelity, which measures how well attributions predict the effect of perturbing the input, and sensitivity_max, which measures how much attributions change under small input noise. Lower is better for both:

from captum.metrics import infidelity, sensitivity_max

def perturb_fn(inputs):
    noise = torch.randn_like(inputs) * 0.05
    return noise, inputs - noise                   # (perturbation, perturbed inputs)

x = X[:20]
b = torch.zeros_like(x)
attr = ig.attribute(x, baselines=b, target=1, n_steps=50)

infid = infidelity(model, perturb_fn, x, attr, target=1)
sens = sensitivity_max(ig.attribute, x, baselines=b, target=1, n_steps=50)
print("infidelity:", infid.mean().item(), "| sensitivity:", sens.mean().item())

5. Deletion and insertion tests. Remove the top-attributed features and check that the prediction drops faster than when you remove random ones. This is the most intuitive faithfulness check, and it needs no extra library.
6. Ground-truth tests. Plant a known signal, as in the tabular example above, or a spurious feature such as a watermark, and verify that the attribution finds it. This is also a good way to catch shortcut learning in real datasets.

Use the metrics to compare methods and baselines on the same inputs, not as absolute scores. An infidelity value has no units you can defend; a ranking between two candidate explanations does.

What Attributions Don't Tell You

• They explain the model, not the world. High attribution means the model used the feature, not that the feature causes the outcome.
• Correlated features share credit unpredictably, so inspect correlations before interpreting individual features, because the same dilution problem affects permutation importance and SHAP values.
• They are local. Attributions for one input do not describe global behaviour. Aggregate across a representative sample for a global view.
• They depend on the baseline and the target, so report both alongside the map.
• Pretty maps can mislead. A smooth heat map can look convincing while being unfaithful. Do not trust appearance; verify.
• Attributing probabilities rather than logits gives a different picture when the model is saturated.

If you want a second opinion on the same model, SHAP and LIME answer different questions, and the disagreement between methods is often the most useful signal you will get.

Practical Tips

1. Always call model.eval() and make sure the model is on the same device as the inputs and baselines.
2. Match the baseline's shape and preprocessing to the input.
3. Use internal_batch_size to avoid out-of-memory errors with large n_steps.
4. Start with small n_steps for exploration, then raise it until the convergence delta is small.
5. For multi-input models, pass inputs and baselines as tuples, and check the 0.9.0 cross-tensor behaviour for perturbation methods.
6. Avoid in-place activations such as ReLU(inplace=True) with backward-hook-based methods like DeepLift, which can throw errors or return wrong results.
7. Pin your Captum, PyTorch and Python versions, since both the API and the supported range have changed.
8. Save the raw attribution tensors alongside the visualizations, so you can recompute summaries later without re-running the model.

Common Pitfalls

1. Treating a single attribution map as proof of what the model looks at. It is one explanation for one input against one baseline.
2. Using a zero baseline in normalized image space and believing it is a black image.
3. Ignoring the convergence delta, which is the cheapest faithfulness signal Captum gives you.
4. Comparing attributions across inputs without a consistent baseline and target.
5. Skipping the randomization sanity check, the one test that separates model behaviour from input structure.
6. Following older tutorials that rely on captum.insights, which has been removed in 0.9.0.
7. Interpreting correlated-feature attributions as independent effects.
8. Attributing a saturated softmax output and concluding that nothing matters.

Check the edition date before you buy. Methodology ages slowly; library APIs do not age at all.

Unlock AI That Actually Works

Get lifetime access to GPT-6 Astra, Claude Fable 5.1, Gemini 3.5, Grok 4.5, and more — all in one platform. Build websites, apps, videos, content, and digital products from a single command. No monthly fees. No tool-hopping.

Click here to get GPTAstra Max now — one-time payment, lifetime access.

Frequently Asked Questions

What does Integrated Gradients actually compute?

Integrated Gradients attributes a prediction to its inputs by integrating the model's gradients along a straight path from a baseline input to the actual input. Each feature's attribution is the input difference from the baseline multiplied by the average gradient along that path. Because of the completeness property, the attributions sum to the difference between the model output at the input and at the baseline, which means you can check the explanation numerically instead of trusting the picture. Captum approximates the integral by sampling points along the path, and 50 is its default.

What baseline should I use for Integrated Gradients?

A baseline that represents no information for your data and your preprocessing. For standardized tabular data a row of zeros is the average row, but for raw features such as age or income a zero baseline is nonsense, so use means, medians or a typical row instead. For images, remember that a zeros tensor after normalization is the mean colour, not a black image, so a blurred copy of the input is often more informative. For text, use padding tokens while keeping the special tokens in place. Whatever you choose, run the method with two or three baselines.

How many n_steps do I need for Integrated Gradients?

Enough that the convergence delta is small relative to the attributions you are reading, which in Captum means the absolute value of the returned delta stays near zero. Captum defaults to 50 steps, which is a reasonable exploration setting. Raise n_steps for a final report and check that the numbers settle rather than drifting with each increase. Higher step counts cost more forward and backward passes, so pair them with internal_batch_size to keep memory under control.

How do I check that attributions are trustworthy?

Test them rather than admiring them. Check the convergence delta and raise n_steps until it is small, rerun with different baselines, and compare. Then randomize the model's weights layer by layer and recompute: if the maps barely change, you are seeing structure in the input such as edges, not what the model learned. Captum ships infidelity and sensitivity_max for quantitative comparison, and deletion or insertion tests are the most intuitive faithfulness check. If several methods agree, the picture is more believable.

Do Integrated Gradients work on text and images, not just tabular data?

Yes, with different plumbing for each. Tabular models take IntegratedGradients directly. Images work with IntegratedGradients wrapped in NoiseTunnel for smoother maps, but you must make sure the displayed crop matches the preprocessing or the heat map will not line up. Text needs LayerIntegratedGradients, because token IDs are discrete and cannot be interpolated, so you attribute at the embedding layer and sum across the embedding dimension to get one score per token. For generative models, Captum added LLMAttribution wrappers from version 0.7 and expanded them in 0.9.0.

Is Captum Insights still available?

No. Captum 0.9.0 retires Captum Insights, the interactive web interface, along with the captum[insights] installation extra, so to use it you would have to pin the older 0.8.0 release. If you read an explainability guide that lists Insights as an interactive option, including my own earlier tools guide, treat that entry as out of date. The library itself is unaffected, and the repository now lives under meta-pytorch/captum. Current attribution APIs such as IntegratedGradients, NoiseTunnel and LayerIntegratedGradients are all still there.

Wrapping Up

Integrated Gradients attributes a prediction to its inputs by integrating gradients along a path from a baseline to the input, and its completeness property gives you a built-in check in the convergence delta. Captum makes that a few lines of code for tabular, image and text models, with NoiseTunnel for cleaner image maps, LayerIntegratedGradients for token embeddings, and LLMAttribution for generative models.

The craft is in the choices around it: baselines, targets, step counts, and above all validating the explanation with sanity checks and metrics rather than with how convincing the heat map looks.

Where should you start? Run the tabular example, where you know the ground truth, and confirm that IG recovers it. Then try your own model with two different baselines and check whether the story survives. If it does, you have an explanation worth showing. If it does not, you have learned something important about both the model and the method.

Which baseline are you using right now, and would your conclusion survive a second one?

Related Articles

• SHAP vs LIME: Which Explainability Method Should You Use?
• Permutation Feature Importance Explained with Scikit-learn
• Best Explainable AI and Model Interpretability Tools
• Partial Dependence and ICE Plots: Feature Effects Explained
• SHAP Values in Python: Explain Your Machine Learning Model Predictions