Intermediate

Attention Maps & Visualization

Visualize what neural networks "see" using Grad-CAM, transformer attention weights, saliency maps, and integrated gradients.

Why Visualize Attention?

Deep neural networks process input through layers of learned representations. Attention visualization techniques let us peek inside these layers to understand which parts of the input the model focuses on when making a prediction. This is invaluable for debugging, validation, and building trust.

Grad-CAM

Gradient-weighted Class Activation Mapping (Grad-CAM) produces a coarse localization map highlighting important regions in an image for a specific prediction. It works with any CNN-based architecture.

Python - Grad-CAM with PyTorch

import torch
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from torchvision import models, transforms
from PIL import Image

# Load a pre-trained model
model = models.resnet50(pretrained=True)
model.eval()

# Hook into the last convolutional layer
activations = {}
gradients = {}

def forward_hook(module, input, output):
    activations['value'] = output.detach()

def backward_hook(module, grad_input, grad_output):
    gradients['value'] = grad_output[0].detach()

target_layer = model.layer4[-1]
target_layer.register_forward_hook(forward_hook)
target_layer.register_full_backward_hook(backward_hook)

# Forward pass
img_tensor = preprocess(image).unsqueeze(0)
output = model(img_tensor)
pred_class = output.argmax(dim=1).item()

# Backward pass for the predicted class
model.zero_grad()
output[0, pred_class].backward()

# Compute Grad-CAM
weights = gradients['value'].mean(dim=[2, 3], keepdim=True)
grad_cam = F.relu((weights * activations['value']).sum(dim=1))
grad_cam = grad_cam.squeeze().numpy()

# Normalize and resize to input image size
grad_cam = (grad_cam - grad_cam.min()) / (grad_cam.max() - grad_cam.min())
grad_cam = np.uint8(255 * grad_cam)

# Overlay on original image
plt.imshow(image)
plt.imshow(grad_cam, alpha=0.5, cmap='jet')
plt.title(f"Grad-CAM: {class_names[pred_class]}")
plt.axis('off')
plt.show()

Transformer Attention Weights

Transformer models (BERT, GPT, ViT) have built-in attention mechanisms. We can extract and visualize these attention weights to understand which tokens or patches attend to each other.

Python - Visualizing BERT attention

from transformers import BertTokenizer, BertModel
import matplotlib.pyplot as plt
import seaborn as sns

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased',
                                   output_attentions=True)

text = "The bank approved the loan application"
inputs = tokenizer(text, return_tensors='pt')
outputs = model(**inputs)

# outputs.attentions is a tuple of (num_layers,) tensors
# Each tensor has shape (batch, num_heads, seq_len, seq_len)
attention = outputs.attentions

# Visualize attention from last layer, averaged across heads
last_layer_attention = attention[-1].squeeze().mean(dim=0)
tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])

plt.figure(figsize=(10, 8))
sns.heatmap(last_layer_attention.detach().numpy(),
            xticklabels=tokens, yticklabels=tokens,
            cmap='viridis')
plt.title("BERT Attention - Last Layer (avg heads)")
plt.show()

Saliency Maps

Saliency maps compute the gradient of the output with respect to the input pixels. Pixels with large gradients are the ones that most influence the prediction.

Python - Vanilla saliency map

img_tensor = preprocess(image).unsqueeze(0)
img_tensor.requires_grad_(True)

output = model(img_tensor)
pred_class = output.argmax(dim=1).item()

# Compute gradient of predicted class w.r.t. input
output[0, pred_class].backward()
saliency = img_tensor.grad.abs().squeeze().max(dim=0)[0]

plt.imshow(saliency.numpy(), cmap='hot')
plt.title("Saliency Map")
plt.axis('off')
plt.show()

Integrated Gradients

Integrated Gradients addresses the noise problem of vanilla gradients by accumulating gradients along a path from a baseline (e.g., black image) to the actual input. This produces smoother, more reliable attributions.

Python - Integrated Gradients with Captum

from captum.attr import IntegratedGradients

ig = IntegratedGradients(model)

# Compute attributions
attributions = ig.attribute(
    img_tensor,
    target=pred_class,
    n_steps=200,
    baselines=torch.zeros_like(img_tensor)
)

# Visualize
from captum.attr import visualization as viz
viz.visualize_image_attr_multiple(
    attributions.squeeze().permute(1, 2, 0).detach().numpy(),
    original_image,
    methods=["heat_map", "blended_heat_map"],
    signs=["positive", "positive"],
    titles=["Attribution", "Overlay"]
)

Comparison of Visualization Methods

Method	Model Type	Granularity	Reliability
Grad-CAM	CNNs	Region-level	High
Saliency Maps	Any differentiable	Pixel-level	Moderate (noisy)
Integrated Gradients	Any differentiable	Pixel-level	High
Attention Weights	Transformers	Token/patch-level	Moderate

⚠

Attention is not always explanation: Research has shown that attention weights do not always correlate with feature importance. Use attention visualization as one signal among many, not as definitive proof of what the model "looks at." Combine with gradient-based methods for more reliable explanations.