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

# Value-IG: per-dim Integrated Gradients — Σ_j ρ_VIG(l,h,j) = ρ_IG(l,h)
# Heatmap: A[CLS→p] · v[h,j,p]   (clean forward pass)
top_c, hmaps, _ = crp.explain(img, scoring='value_integrated_gradients', heatmap='attention', top_k=TOP_K,
                               layers=LAYERS, class_idx=pred_class, steps=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-IG — class {pred_class}: {info.index_to_description(pred_class)}')
plt.show()

  L 5 H5 dim 41  +0.0604
  L 5 H1 dim 51  +0.0549
  L 5 H5 dim 25  -0.0502
  L 5 H1 dim 43  +0.0479
  L 5 H1 dim 32  -0.0477
  L 5 H1 dim  0  +0.0382

# Value-Occlusion: occlusion scores + A[CLS→p]·‖v_p‖ spatial maps
# Same ρ(l,h) as occlusion; heatmap also encodes value magnitude
top_c, hmaps, _ = crp.explain(img, scoring='value_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'Value-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

# Value-KernelSHAP: SHAP scores + A[CLS→p]·‖v_p‖ spatial maps   (~30 s)
top_c, hmaps, _ = crp.explain(img, scoring='value_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'Value-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

# Value-Causal Tracing: causal scores + A[CLS→p]·‖v_p‖ spatial maps
top_c, hmaps, _ = crp.explain(img, scoring='value_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'Value-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

Score Comparison

[L × H] score matrices for all 8 base methods side-by-side. Red = head promotes the predicted class; blue = head suppresses it. (Value-modulated variants share the same scores — only their heatmaps differ.)

base_methods = {
    'gxi':                  {},
    'occlusion':            {},
    'integrated_gradients': {'steps': 50},
    'smoothgrad':           {'n_samples': 50},
    'kernelshap':           {'n_samples': 512},
    'causal_tracing':       {},
    'value_norm':           {},
    'lrp':                  {'eps': 0.1},
}

scores_dict = {}
for m, kw in base_methods.items():
    s = crp.scores(x, method=m, class_idx=pred_class, layers=LAYERS, **kw).float().cpu().numpy()
    scores_dict[m] = s
    print(f'{m:25s}  Σ={s.sum():.3f}  |max|={abs(s).max():.3f}')

fig, axes = plt.subplots(2, 4, figsize=(16, 7))
for ax, (name, s) in zip(axes.flat, scores_dict.items()):
    vmax = max(abs(s).max(), 1e-6)
    im = ax.imshow(s, cmap='RdBu_r', aspect='auto', vmin=-vmax, vmax=vmax)
    ax.set_title(name, fontsize=10)
    ax.set_xlabel('Head')
    ax.set_ylabel('Layer')
    plt.colorbar(im, ax=ax, shrink=0.8)
plt.suptitle(
    f'[L×H] scores — class {pred_class}: {info.index_to_description(pred_class)}',
    fontsize=11)
plt.tight_layout()
plt.show()

gxi                        Σ=0.155  |max|=0.172
occlusion                  Σ=0.868  |max|=0.365
integrated_gradients       Σ=0.060  |max|=0.221
smoothgrad                 Σ=0.301  |max|=0.228
kernelshap                 Σ=1.845  |max|=0.983
causal_tracing             Σ=-0.724  |max|=0.506
value_norm                 Σ=26.346  |max|=7.054
lrp                        Σ=1.335  |max|=2.095

---

Head input-occlusion

Instead of using the CLS-attention row (ρ × attention), this method occludes each input patch in turn while restricting the forward pass through a single head (all others zeroed). The per-patch score is baseline_logit − occluded_logit: positive regions are what the head relies on for predicting class c.

from methods import head_input_occlusion as hio

# Pick the top GxI head
top_c_gxi, hmaps_gxi, _ = crp.explain(img, scoring='gxi', heatmap='attention', top_k=TOP_K,
                                        layers=LAYERS, class_idx=pred_class)
l_star, h_star, gxi_score = top_c_gxi[0]
print(f'Target head: L{l_star} H{h_star}  gxi_score={gxi_score:+.4f}')

# Side-by-side: attention heatmap vs input-occlusion heatmap for the same head
occ_out = hio.score_head(crp.model, img, l_star, h_star, class_idx=pred_class)

fig, axes = plt.subplots(1, 2, figsize=(7, 3.5))
axes[0].imshow(visualize.overlay_heatmap(img, hmaps_gxi[0]))
axes[0].set_title(f'L{l_star} H{h_star}\nGxI attention heatmap', fontsize=8)
axes[0].axis('off')
axes[1].imshow(visualize.overlay_heatmap(img, occ_out['heatmap']))
axes[1].set_title(f'L{l_star} H{h_star}\nInput-occlusion (bottleneck)', fontsize=8)
axes[1].axis('off')
fig.suptitle(f'class {pred_class}: {info.index_to_description(pred_class)}', fontsize=9)
plt.tight_layout(); plt.show()

Target head: L5 H1  gxi_score=+0.1717

# Top-k concept grid with input-occlusion heatmaps
top_c, hmaps, _ = crp.explain(
    img, scoring='gxi', heatmap='input_occlusion',
    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'Head input-occlusion — 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