5. Gradient-Based Saliency Map for Input Attribution

Problem

A model debugging session needs to know which input features actually drove a given prediction. The simplest attribution technique — vanilla gradient saliency — answers this directly: take the gradient of the predicted class's logit with respect to the input, and its magnitude at each input position tells you how sensitive the prediction is to perturbing that position.

It requires no architecture changes, only a single backward pass with the input itself treated as a leaf tensor that requires grad.

Source: src/5_gradient_saliency_map.py

def saliency_map(model: torch.nn.Module, x: torch.Tensor, target_indices: torch.Tensor) -> torch.Tensor: ...

>>> model = nn.Linear(3, 2, bias=False)
>>> x = torch.randn(2, 3)
>>> targets = torch.tensor([0, 1])
>>> saliency_map(model, x, targets).shape
torch.Size([2, 3])

Step-by-Step Approach

1. Make a fresh leaf tensor from the input: x.clone().detach().requires_grad_(True) — never mutate or reuse the caller's original tensor, which may already have stale gradients or a different graph attached.
2. Run the model forward to get logits of shape (batch, num_classes).
3. Select exactly one logit per row — the target class for that example — with logits.gather(1, target_indices.unsqueeze(1)).squeeze(1).
4. Call .backward() on the sum of the selected logits. Summing is safe here because each selected logit only depends on its own row of x, so no cross-example gradient leakage occurs.
5. Read the gradient off the leaf tensor with x.grad, and take its absolute value — sign doesn't matter for "how sensitive," only magnitude does.
6. Sanity check: for a purely linear model with no bias, the saliency for target class c is exactly |weight[c, :]|, independent of the actual input x — a strong, input-independent check for whether the backward pass targeted the right logit.

The key insight is the linear sanity check: since a linear layer's gradient with respect to its input is just its weight matrix, verifying saliency against the known weight row catches an incorrectly targeted gather or a wrong reduction (e.g. summing all classes instead of selecting one) immediately, before testing against a real nonlinear model where the correct answer isn't obvious by inspection.

Reference solution

import torch
import torch.nn as nn


def saliency_map(model: nn.Module, x: torch.Tensor, target_indices: torch.Tensor) -> torch.Tensor:
    # clone + detach + requires_grad_ makes a fresh leaf tensor -- never
    # mutate the caller's original x in place, and never reuse an x that
    # might already carry stale .grad from a previous call
    x = x.clone().detach().requires_grad_(True)
    logits = model(x)  # (batch, num_classes)
    # pick out exactly one logit per row (the target class) without a python loop
    selected = logits.gather(1, target_indices.unsqueeze(1)).squeeze(1)  # (batch,)
    selected.sum().backward()  # sum() is safe here: each term only depends on its own row
    return x.grad.abs()

Key Functions & Tricks

• x.clone().detach().requires_grad_(True) — the standard idiom for turning an arbitrary tensor into a fresh, gradient-tracked leaf without touching the caller's original tensor.
• logits.gather(1, target_indices.unsqueeze(1)) — vectorized per-row indexing, selecting one column index per row without a Python loop.
• selected.sum().backward() — summing independent per-example scalars before calling .backward() is the standard way to backprop a batch of unrelated scalars in one call.
• x.grad.abs() — vanilla saliency uses gradient magnitude; signed gradients are used by some attribution variants but not this one.

How to Recognize This Pattern