Head-CRP: Attribution Method Demonstration

All 14 scoring methods demonstrated on the same image. Each cell runs one method and shows the top-6 concept heatmaps via crp.explain().

Section	Methods
Base	gxi, occlusion, integrated_gradients, smoothgrad, kernelshap, causal_tracing, value_norm, lrp
Value-modulated	value_gxi, value_smoothgrad, value_integrated_gradients, value_occlusion, value_kernelshap, value_causal_tracing

import sys
sys.path.insert(0, '..')  # repo root

import torch
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from PIL import Image
from torchvision import transforms
from timm.data import ImageNetInfo

from crp.crp import CRP, IMAGENET_MEAN, IMAGENET_STD
from crp import visualize

crp = CRP(model_name='vit_small_patch16_224')
info = ImageNetInfo()
L = len(crp.model.blocks)
H = crp.model.blocks[0].attn.num_heads

# ── Image ───────────────────────────────────────────────────────────────────
val_dir = Path('../data/imagenet_val')
class_dirs = [d for d in sorted(val_dir.iterdir()) if d.is_dir()] if val_dir.exists() else []
if class_dirs:
    import random; random.seed(42)
    img_path = random.choice(list(class_dirs[0].glob('*.JPEG')))
    img = Image.open(img_path).convert('RGB')
    print(f'Image: {img_path}')
else:
    rng = np.random.RandomState(42)
    img = Image.fromarray(rng.randint(0, 255, (224, 224, 3), dtype=np.uint8))
    print('Image: synthetic (data/imagenet_val not found)')

# Pre-processed tensor for crp.scores()
t = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
])
x = t(img).unsqueeze(0).to(crp.device)

with torch.no_grad():
    logits = crp.model(x)
pred_class = int(logits.argmax().item())
conf = logits.softmax(-1)[0, pred_class].item()
print(f'Model: {L} layers × {H} heads = {L * H} head-concepts')
print(f'Class {pred_class}: {info.index_to_description(pred_class)} (conf={conf:.3f})')

# ── Global demo parameters ────────────────────────────────────────────────
LAYERS = [5]   # None = all layers; e.g. [9, 10, 11] for late layers only
TOP_K  = 6      # number of top concepts shown per method

plt.figure(figsize=(3, 3))
plt.imshow(img)
plt.title(f'class {pred_class}: {info.index_to_description(pred_class)[:40]}', fontsize=9)
plt.axis('off'); plt.tight_layout(); plt.show()

Image: ../data/imagenet_val/n01440764/n01440764_7591.JPEG
Model: 12 layers × 6 heads = 72 head-concepts
Class 0: tench, Tinca tinca (conf=0.943)

Base Methods

Scores [L, H] — CLS-attention spatial maps.

# GxI: gradient × activation — 1 forward + 1 backward pass
# ρ(l,h) = Σ_{n,j} z^(l,h)_{n,j} · ∂f_c/∂z^(l,h)_{n,j}
top_c, hmaps, _ = crp.explain(img, scoring='gxi', heatmap='attention', top_k=TOP_K, layers=LAYERS, class_idx=pred_class)
for l, h, s in top_c:
    print(f'  L{l:2d} H{h}  {s:+.4f}')
fig = visualize.concept_grid(img, top_c, hmaps,
    title=f'GxI — class {pred_class}: {info.index_to_description(pred_class)}')
plt.show()

  L 5 H1  +0.1717
  L 5 H0  +0.1481
  L 5 H5  -0.1281
  L 5 H2  -0.0923
  L 5 H3  +0.0844
  L 5 H4  -0.0291

# Occlusion: exact logit-drop when head is zeroed — L×H+1 forward passes
# ρ(l,h) = f_c(x) − f_c(x | z^(l,h) = 0)
top_c, hmaps, _ = crp.explain(img, scoring='occlusion', heatmap='attention', top_k=TOP_K, layers=LAYERS, class_idx=pred_class)
for l, h, s in top_c:
    print(f'  L{l:2d} H{h}  {s:+.4f}')
fig = visualize.concept_grid(img, top_c, hmaps,
    title=f'Occlusion — class {pred_class}: {info.index_to_description(pred_class)}')
plt.show()

  L 5 H5  +0.3650
  L 5 H1  +0.1920
  L 5 H0  +0.1914
  L 5 H3  +0.1102
  L 5 H4  +0.0384
  L 5 H2  -0.0285

# Integrated Gradients: path integral — completeness: Σρ = f_full − f_empty
# ρ_IG(l,h) = Σ_{n,j} z·∫₀¹ ∂f/∂z|_{α·z_actual} dα   (steps=50 Riemann sum)
top_c, hmaps, _ = crp.explain(img, scoring='integrated_gradients', heatmap='attention', top_k=TOP_K,
                               layers=LAYERS, class_idx=pred_class, steps=50)
for l, h, s in top_c:
    print(f'  L{l:2d} H{h}  {s:+.4f}')
fig = visualize.concept_grid(img, top_c, hmaps,
    title=f'Integrated Gradients — class {pred_class}: {info.index_to_description(pred_class)}')
plt.show()

  L 5 H1  +0.2210
  L 5 H3  -0.1057
  L 5 H0  -0.0431
  L 5 H5  -0.0142
  L 5 H2  +0.0095
  L 5 H4  -0.0076

# SmoothGrad: GxI averaged over Gaussian-noise inputs — variance-reduced attribution
# ρ_SG(l,h) = E_ε[ Σ_{n,j} z(x+ε)·∂f/∂z(x+ε) ]   (n_samples=50)
top_c, hmaps, _ = crp.explain(img, scoring='smoothgrad', heatmap='attention', top_k=TOP_K,
                               layers=LAYERS, class_idx=pred_class, n_samples=50)
for l, h, s in top_c:
    print(f'  L{l:2d} H{h}  {s:+.4f}')
fig = visualize.concept_grid(img, top_c, hmaps,
    title=f'SmoothGrad — class {pred_class}: {info.index_to_description(pred_class)}')
plt.show()

  L 5 H2  -0.2462
  L 5 H3  +0.2371
  L 5 H0  +0.1498
  L 5 H1  +0.1004
  L 5 H4  +0.0774
  L 5 H5  -0.0225

# KernelSHAP: Shapley values via importance-sampled weighted least squares
# Efficiency axiom: Σ φ(l,h) = f_full − f_empty   (enforced by construction)
# Runtime: ~30 s on GPU (n_samples=512 coalitions, batch_size=64)
top_c, hmaps, _ = crp.explain(img, scoring='kernelshap', heatmap='attention', top_k=TOP_K,
                               layers=LAYERS, class_idx=pred_class, n_samples=512)
for l, h, s in top_c:
    print(f'  L{l:2d} H{h}  {s:+.4f}')
fig = visualize.concept_grid(img, top_c, hmaps,
    title=f'KernelSHAP — class {pred_class}: {info.index_to_description(pred_class)}')
plt.show()

  L 5 H5  +0.9830
  L 5 H4  +0.8949
  L 5 H2  +0.8892
  L 5 H1  -0.4181
  L 5 H0  -0.2716
  L 5 H3  -0.2326

# Causal Tracing: logit-recovery fraction after patching clean head into corrupted input
# ρ_CT(l,h) = (f_c(x_corr, patch(l,h)=clean) − f_c(x_corr)) / (f_c(x) − f_c(x_corr))
# Score ≈ 1: head is causally sufficient for class prediction; ≈ 0: no effect
top_c, hmaps, _ = crp.explain(img, scoring='causal_tracing', heatmap='attention', top_k=TOP_K, layers=LAYERS, class_idx=pred_class)
for l, h, s in top_c:
    print(f'  L{l:2d} H{h}  {s:+.4f}')
fig = visualize.concept_grid(img, top_c, hmaps,
    title=f'Causal Tracing — class {pred_class}: {info.index_to_description(pred_class)}')
plt.show()

  L 5 H3  -0.5056
  L 5 H5  +0.4563
  L 5 H1  -0.2821
  L 5 H0  -0.2550
  L 5 H4  -0.0978
  L 5 H2  -0.0393

# Value Norm: gradient-free information-content proxy — single forward pass
# ρ(l,h) = Σ_p A^(l,h)[CLS→p] · ‖v^(l,h)_p‖₂   (always ≥ 0)
top_c, hmaps, _ = crp.explain(img, scoring='value_norm', heatmap='attention', top_k=TOP_K, layers=LAYERS, class_idx=pred_class)
for l, h, s in top_c:
    print(f'  L{l:2d} H{h}  {s:+.4f}')
fig = visualize.concept_grid(img, top_c, hmaps,
    title=f'Value Norm — class {pred_class}: {info.index_to_description(pred_class)}')
plt.show()

  L 5 H5  +7.0543
  L 5 H1  +6.2967
  L 5 H4  +4.5019
  L 5 H0  +4.2754
  L 5 H3  +2.9252
  L 5 H2  +1.2920

# LRP-ε: layer-wise relevance propagation — 1 forward + 1 backward (no autograd)
# ρ(l,h) = total relevance at head h of block l; initialised at logit_c, propagated via ε-rule
# eps=0.1: stable (max 1.4× per-block growth); eps=1e-6 explodes ×32k at block 9 (GELU sparsity)
top_c, hmaps, _ = crp.explain(img, scoring='lrp', heatmap='attention', top_k=TOP_K, layers=LAYERS,
                               class_idx=pred_class, eps=0.1)
for l, h, s in top_c:
    print(f'  L{l:2d} H{h}  {s:+.4f}')
fig = visualize.concept_grid(img, top_c, hmaps,
    title=f'LRP-ε (eps=0.1) — class {pred_class}: {info.index_to_description(pred_class)}')
plt.show()

  L 5 H5  +2.0946
  L 5 H1  +1.4773
  L 5 H2  -1.0972
  L 5 H3  -0.7480
  L 5 H0  -0.4876
  L 5 H4  +0.0954

Value-Modulated Variants

The spatial heatmap is replaced by A[CLS→p] · v[h,j,p] (per value-dimension) or A[CLS→p] · ‖v_p‖ (L2-norm) instead of plain CLS attention.

• Gradient methods (value_gxi, value_smoothgrad, value_integrated_gradients): scores decomposed per value dimension → [L, H, d_h]; top concepts reported as (layer, head, dim, score).
• Non-gradient methods (value_occlusion, value_kernelshap, value_causal_tracing): same [L, H] scores as the base method; spatial map uses A·‖v‖.

# Value-GxI: per-value-dim GxI — strict decomposition of ρ_GxI(l,h)
# ρ(l,h,j) = Σ_n z_{n,j}·∂f/∂z_{n,j}   →   Σ_j ρ(l,h,j) = ρ_GxI(l,h)
# Heatmap: A[CLS→p] · v[h,j,p]
top_c, hmaps, _ = crp.explain(img, scoring='value_gxi', heatmap='attention', top_k=TOP_K, layers=LAYERS, class_idx=pred_class)
for l, h, j, s in top_c:
    print(f'  L{l:2d} H{h} dim{j:3d}  {s:+.4f}')
fig = visualize.concept_grid(img, top_c, hmaps,
    title=f'Value-GxI — class {pred_class}: {info.index_to_description(pred_class)}')
plt.show()

  L 5 H0 dim 13  +0.2041
  L 5 H5 dim 55  -0.1334
  L 5 H5 dim  6  -0.1160
  L 5 H3 dim 50  -0.0997
  L 5 H5 dim 53  -0.0906
  L 5 H3 dim 30  -0.0879

# Value-SmoothGrad: per-dim SmoothGrad — Σ_j ρ_VSG(l,h,j) = ρ_SG(l,h)
# Heatmap: A[CLS→p] · v[h,j,p]   (from clean input)
top_c, hmaps, _ = crp.explain(img, scoring='value_smoothgrad', heatmap='attention', top_k=TOP_K,
                               layers=LAYERS, class_idx=pred_class, n_samples=50)
for l, h, j, s in top_c:
    print(f'  L{l:2d} H{h} dim{j:3d}  {s:+.4f}')
fig = visualize.concept_grid(img, top_c, hmaps,
    title=f'Value-SmoothGrad — class {pred_class}: {info.index_to_description(pred_class)}')
plt.show()

  L 5 H0 dim 13  +0.1982
  L 5 H5 dim 55  -0.1591
  L 5 H5 dim 53  -0.1222
  L 5 H4 dim 21  +0.1054
  L 5 H3 dim 50  -0.1054
  L 5 H3 dim 41  +0.1045