Interpretability and Feature Attribution

Feature attribution answers a narrow question: for this trained model, this input, this target output, and this reference point, which input components most changed the score? It does not prove causality, fairness, or correctness. In a PyTorch workflow it is useful because tensors, autograd, hooks, and perturbation loops give you direct access to the computations that created a prediction.

The outcome for this lesson is practical. You should be able to compute simple gradient, integrated-gradient, and occlusion attributions; explain what each number means; choose a baseline deliberately; and diagnose common attribution failures before using a heatmap or ranked feature list in model review.

How Attribution Works

A neural network maps an input tensor to one or more output scores through a sequence of differentiable operations. Attribution methods attach scores to parts of the input. For tabular data, a part may be a column. For text, it may be a token embedding or token position. For images, it may be a pixel, channel, patch, or superpixel. The target is the scalar output being explained, such as the logit for class 3 rather than the whole probability vector.

Gradient saliency is the simplest internal mechanism. Mark the input with requires_grad=True, run a forward pass, select one scalar target, call backward, and read x.grad. The gradient says how the target would change for an infinitesimal change in each input element around the current point. Large absolute gradients indicate local sensitivity. The method is cheap because it uses one backward pass, but it can be noisy, can saturate when activations are flat, and can miss features whose importance came from moving far from a baseline.

Integrated gradients repair part of that problem by averaging gradients along a path from a baseline input to the actual input. The usual path is a straight line. The attribution for each feature is the input difference multiplied by the average gradient along the path. When computed accurately, the attributions approximately satisfy completeness: their sum equals the target score difference between the input and the baseline. The baseline is therefore part of the explanation, not an implementation detail. A black image, a zero vector, a pad token, or a population mean can each answer a different question.

Occlusion and perturbation methods do not require gradients. They repeatedly replace one feature, patch, or token with a reference value and measure the score change. If masking a feature sharply lowers the target score, that feature supported the prediction under the chosen mask. Perturbation is easy to reason about and works with nondifferentiable preprocessing, but it is slower and can create unrealistic inputs. A masked image patch, for example, may be outside the training distribution.

API Anatomy in PyTorch

The core ingredients are a model in evaluation mode, an input tensor with a stable shape, a scalar target, a baseline or mask value, and a rule for aggregating attribution dimensions. Use model.eval() so dropout and batch normalization do not inject training-time behavior. Clear old gradients before a saliency pass. Prefer logits over softmax probabilities for attribution because softmax couples classes and can hide evidence when all logits shift together.

For image explanations, raw attributions often have shape [channels, height, width]. You might aggregate channels with an absolute sum for a heatmap, or keep signed values to distinguish evidence for and against the target. For text, attribution on token ids is not differentiable; explain the embedding vectors or use perturbation over tokens. For batches, compute per-example targets rather than summing unrelated examples unless the batch itself is the intended object of explanation.

Example 1: Gradient Saliency

This example uses a linear model so the attribution is deterministic and easy to inspect. The model score is 2*x0 - x1 + 0.5*x2. Gradients are the learned weights because the function is linear.

import torch

torch.manual_seed(7)
model = torch.nn.Linear(3, 1, bias=False)
with torch.no_grad():
    model.weight.copy_(torch.tensor([[2.0, -1.0, 0.5]]))

x = torch.tensor([[1.5, 2.0, -4.0]], requires_grad=True)
y = model(x).squeeze()
y.backward()

saliency = x.grad.abs()
print("prediction", round(float(y), 3))
print("gradients", x.grad.tolist())
print("saliency", saliency.tolist())

The prediction is -1.0. The gradient vector is [2.0, -1.0, 0.5], and the absolute saliency is [2.0, 1.0, 0.5]. This says feature 0 has the highest local sensitivity. It does not say feature 0 contributed the most to this particular score; feature 2 has value -4.0, so its actual signed contribution is -2.0.

Example 2: Integrated Gradients

The next example uses a sigmoid model where gradients vary along the path. Integrated gradients average those gradients between a zero baseline and the input.

import torch

def integrated_gradients(model, x, baseline, steps=64):
    scaled = torch.stack([baseline + (i / steps) * (x - baseline) for i in range(1, steps + 1)])
    scaled.requires_grad_(True)
    outputs = model(scaled).sum()
    grads = torch.autograd.grad(outputs, scaled)[0]
    avg_grads = grads.mean(dim=0)
    return (x - baseline) * avg_grads

model = torch.nn.Sequential(
    torch.nn.Linear(2, 1, bias=False),
    torch.nn.Sigmoid(),
)
with torch.no_grad():
    model[0].weight.copy_(torch.tensor([[3.0, -2.0]]))

x = torch.tensor([1.0, 2.0])
baseline = torch.zeros_like(x)
attr = integrated_gradients(model, x, baseline)
print([round(v, 4) for v in attr.tolist()])
print("delta", round(float(attr.sum() - (model(x) - model(baseline))), 4))

The printed attribution values are approximately [0.6919, -0.9226], and the completeness delta is about 0.0004. The first feature increases the target probability relative to the zero baseline. The second feature decreases it. More steps usually improves the approximation but costs more backward work.

Example 3: Occlusion Attribution

Occlusion answers a different question: what happens if one feature is replaced by a baseline value while all other features stay fixed?

import torch

def occlusion_attribution(model, x, baseline_value=0.0):
    base_score = model(x.unsqueeze(0))[0, 0]
    scores = []
    for j in range(x.numel()):
        masked = x.clone()
        masked[j] = baseline_value
        masked_score = model(masked.unsqueeze(0))[0, 0]
        scores.append(float(base_score - masked_score))
    return scores

model = torch.nn.Linear(4, 1, bias=False)
with torch.no_grad():
    model.weight.copy_(torch.tensor([[1.0, 0.0, -2.0, 0.5]]))

x = torch.tensor([3.0, 10.0, -1.0, 2.0])
print([round(v, 3) for v in occlusion_attribution(model, x)])

The output is [3.0, 0.0, 2.0, 1.0]. Feature 1 has value 10.0, but its learned weight is zero, so occluding it changes nothing. Feature 2 has a negative value and a negative weight, so it supports the positive score by 2.0. This example also shows why raw magnitude is not enough: attribution depends on the model, the input value, and the replacement value.

Design Choices and Trade-offs

Choose the attribution method to match the review question. Gradient saliency is good for quick debugging of differentiable models. Integrated gradients is better when you need a baseline-relative account and a completeness check. Occlusion is useful when gradients are unavailable or when stakeholders understand counterfactual masking more readily. For high-dimensional inputs, perturbing every element can be too slow, so use patches, grouped features, or sampled perturbations.

Choose baselines with domain knowledge. A zero vector may mean absence for normalized tabular features, but it may be a real patient state, a meaningful sensor reading, or an invalid token. A mean baseline can make explanations look stable while hiding behavior for minority subgroups. Document the baseline in every chart title, report, and saved artifact because two baselines can reverse the apparent ranking of features.

Decide whether signed or absolute attribution is appropriate. Signed values preserve whether a feature pushes toward or away from a target. Absolute values are useful for ranking sensitivity but can mix opposing evidence. For convolutional models, smoothing and channel aggregation make visualizations easier to read, but they also discard information. Keep the raw tensor when the explanation is part of an audit trail.

Failure Modes and Troubleshooting

Symptom: every saliency value is zero or x.grad is None. Cause: the input was not marked with requires_grad=True, the selected target was detached, or attribution was run under torch.no_grad(). Diagnostic steps: print x.requires_grad, check that the target has a grad_fn, and remove inference-only context managers. Correction: clone the input, call requires_grad_(True), run a fresh forward pass, select one scalar target, and call backward.

Symptom: heatmaps change between repeated runs for the same input. Cause: the model remained in training mode, dropout is active, batch normalization is using batch statistics, or preprocessing includes randomness. Diagnostic steps: compare model.training, seed preprocessing, and run the same tensor twice. Correction: use model.eval(), freeze stochastic transforms, and attribute the exact tensor passed to the model.

Symptom: integrated gradients do not add up to the model output difference. Cause: too few integration steps, explaining probabilities instead of logits, using the wrong target, or aggregating dimensions incorrectly. Diagnostic steps: compute the completeness gap explicitly, increase steps, and verify the target index. Correction: attribute a scalar logit, raise the number of path samples, and keep feature dimensions aligned until after the check.

Symptom: occlusion says a feature is harmful even though domain experts expect it to help. Cause: the replacement value creates an unrealistic input or breaks correlated features. Diagnostic steps: inspect masked samples, compare several baselines, and evaluate whether masked examples are plausible. Correction: mask groups of correlated features, use a domain-specific reference value, or switch to a conditional perturbation method.

Reliability and Performance Implications

Attribution code can become expensive. Integrated gradients with 128 steps costs roughly 128 gradient evaluations per explained input. Occlusion over a 224 by 224 image one pixel at a time is usually impractical. Run attribution offline for model analysis when possible, batch path samples, and store model version, preprocessing version, target, baseline, and method parameters with every result.

Security and privacy matter when explanations are shown outside the training team. A feature attribution can reveal sensitive input fields, token content, or medical-image regions. Do not log raw explanations for private records unless the retention policy allows it. If explanations drive user-facing decisions, label them as model-specific evidence, not as a causal proof.

Hands-on Lab

Prerequisites: Python with PyTorch installed, a shell, and permission to run local scripts. No dataset download is required.

1. Create a temporary working directory and a file named attrib_lab.py.
2. Paste the completeness check below into the file.
3. Run python attrib_lab.py. Verification succeeds when the output is 0.0 or another value below the tolerance.
4. Change attrs to attrs * 0.5 and rerun. The script should raise an AssertionError, proving the check can catch a broken attribution.
5. Rollback by restoring the original attrs line. Cleanup by deleting the temporary directory.

import torch

def validate_completeness(model, x, baseline, attributions, tolerance=1e-2):
    gap = float((model(x) - model(baseline) - attributions.sum()).abs())
    if gap > tolerance:
        raise AssertionError(f"attribution completeness gap {gap:.4f} exceeds {tolerance}")
    return gap

model = torch.nn.Linear(2, 1, bias=False)
with torch.no_grad():
    model.weight.copy_(torch.tensor([[2.0, -1.0]]))

x = torch.tensor([3.0, 4.0])
baseline = torch.zeros_like(x)
attrs = (x - baseline) * model.weight.squeeze(0)
print(round(validate_completeness(model, x, baseline, attrs), 6))

This lab is intentionally small. The habit matters: every attribution method should have at least one invariant or sanity check. Completeness is available for integrated gradients and exact for the linear calculation above. For occlusion, a comparable check is to verify that masking an unused feature produces no score change in a controlled model.

Assessment Exercises

1. A classifier uses softmax probabilities for saliency and highlights different pixels when another unrelated class is added. Explain why logits may be a better attribution target.
2. For a normalized tabular model, compare a zero baseline with a training-set mean baseline. What question does each baseline answer, and when could each mislead?
3. An integrated-gradients report has a large completeness gap. List three diagnostics you would run before trusting the chart.
4. Design an occlusion strategy for text sentiment classification. What exactly would you mask, and how would you avoid creating inputs the model never saw during training?
5. Given signed attributions [-4, 3, 1] and absolute attributions [4, 3, 1], explain what information is lost when only absolute values are reported.

Summary

Interpretability and feature attribution in PyTorch is built from concrete mechanisms: gradients through the input, path-averaged gradients from a baseline, and score differences under controlled perturbations. The useful explanation is always tied to a model version, input tensor, target, baseline, and aggregation rule. Treat attribution as evidence for debugging and review, verify it with sanity checks, and document its limits before using it to justify model behavior.