Head-CRP Concept Viewer¶
This notebook answers the question: what visual concepts have individual attention heads learned?
The idea¶
Head-CRP assigns a score ρ(l,h) to every attention head (layer l, head h) for every image.
A positive score means the head supported the predicted class; negative means it suppressed it.
By running this over a reference corpus we can ask: which images most strongly activate head (l,h)? Those are its concept exemplars — the visual evidence for what the head detects.
For dim-level methods (value_gxi, value_smoothgrad, value_integrated_gradients) scores
decompose further to ρ(l,h,j) per value-channel j, giving sub-head concepts.
Notebook structure¶
| Section | Question answered |
|---|---|
| 1. Head × Class Affinity | Which heads matter most for which classes overall? |
| 2. Per-class top heads | Which specific heads are most class-relevant? |
| 3. Browse head concepts | What visual pattern does a given head detect? |
| 4. Browse value-dim concepts | Finer per-channel view (dim methods only) |
| 5. Concepts for a query image | Which concepts explain one image's prediction? |
Prerequisites¶
# Build the score cache first (run from repo root, ≈7 min on GPU)
python scripts/build_exemplars.py # value_gxi, dim-level
python scripts/build_exemplars.py --method gxi # head-level, faster
python scripts/build_exemplars.py --method gxi --save-heatmaps # also save CLS attention
# maps for fast heatmap
# overlay on exemplars
# Then run this notebook top-to-bottom.
Heatmap overlays on exemplars require either
--save-heatmaps(fast, uses cached CLS attention maps) or will fall back to re-running the model per exemplar image (slow).
import sys
sys.path.insert(0, '..')
from crp.concepts import ConceptStore
store = ConceptStore('../out/concepts')
Cache : 3925 images, 10 classes Method : value_smoothgrad (dim level) Scores : scores_head [3925, 12, 6], scores_dim [3925, 12, 6, 64] Decomp : max|Σ_j ρ_dim − ρ_head| = 0.00e+00 (should be 0) Heatmaps: not cached (pass --save-heatmaps to build_exemplars.py)
# ── Global settings — edit these before running the notebook ──────────────
LAYERS = None # None = all layers; e.g. [10, 11] for the last two layers only
TOP_K = 5 # top concepts shown per query image
K_EXEMPLARS = 6 # corpus exemplars shown per concept
1. Head × Class Affinity¶
A[l·H + h, c] = mean |ρ(l,h)| across all corpus images of GT class c.
A bright cell means that head is reliably important for images from that class.
What to look for:
- Late-layer concentration — strongly attributed heads tend to cluster in the last 1–2 layers, where the model assembles its final class representation.
- Class-selective heads — a head bright for one class and near-zero for others is a class-specific feature detector.
- General heads — heads bright across many classes respond to shared low-level structure (edges, textures) rather than class identity.
store.plot_affinity()
2. Per-class top heads¶
Text ranking of the affinity matrix: for each GT class, the TOP heads with the highest
mean |ρ|. Use this to identify which specific (l, h) pairs to inspect in Section 3.
store.print_top_heads(top=5)
n01440764 (tench, Tinca tinca) #1 L11 H0 mean|ρ|=1.2018 #2 L11 H5 mean|ρ|=0.8779 #3 L 0 H1 mean|ρ|=0.8035 #4 L 2 H0 mean|ρ|=0.7403 #5 L 1 H5 mean|ρ|=0.7250 n02102040 (Japanese spaniel) #1 L11 H0 mean|ρ|=1.6600 #2 L 0 H1 mean|ρ|=0.9350 #3 L 2 H0 mean|ρ|=0.8959 #4 L11 H5 mean|ρ|=0.7800 #5 L10 H2 mean|ρ|=0.7206 n02979186 (tape player) #1 L11 H5 mean|ρ|=1.9398 #2 L 0 H1 mean|ρ|=1.3093 #3 L 2 H0 mean|ρ|=0.9157 #4 L 1 H5 mean|ρ|=0.8509 #5 L11 H0 mean|ρ|=0.8232 n03000684 (chain saw, chainsaw) #1 L 1 H5 mean|ρ|=1.2310 #2 L 0 H1 mean|ρ|=1.1154 #3 L 2 H0 mean|ρ|=1.0019 #4 L 5 H5 mean|ρ|=1.0004 #5 L 1 H1 mean|ρ|=0.9854 n03028079 (church, church building) #1 L11 H5 mean|ρ|=1.4664 #2 L 8 H3 mean|ρ|=0.8912 #3 L 0 H1 mean|ρ|=0.8662 #4 L 1 H5 mean|ρ|=0.8608 #5 L11 H3 mean|ρ|=0.7131 n03394916 (trombone) #1 L11 H5 mean|ρ|=1.6443 #2 L11 H0 mean|ρ|=1.0445 #3 L 0 H1 mean|ρ|=0.8874 #4 L 1 H1 mean|ρ|=0.7968 #5 L11 H1 mean|ρ|=0.7833 n03417042 (trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi) #1 L 1 H5 mean|ρ|=1.1027 #2 L 0 H1 mean|ρ|=0.9771 #3 L 2 H0 mean|ρ|=0.9196 #4 L11 H5 mean|ρ|=0.8874 #5 L 3 H4 mean|ρ|=0.8607 n03425413 (gas pump, gasoline pump, petrol pump, island dispenser) #1 L 2 H0 mean|ρ|=1.0702 #2 L 1 H5 mean|ρ|=1.0023 #3 L 0 H1 mean|ρ|=0.9922 #4 L 5 H5 mean|ρ|=0.9566 #5 L 1 H1 mean|ρ|=0.9426 n03445777 (golf ball) #1 L11 H5 mean|ρ|=2.0759 #2 L11 H0 mean|ρ|=1.0493 #3 L 0 H1 mean|ρ|=0.9213 #4 L 1 H5 mean|ρ|=0.7933 #5 L 5 H5 mean|ρ|=0.7544 n03888257 (parachute, chute) #1 L11 H5 mean|ρ|=1.2197 #2 L 0 H1 mean|ρ|=1.1429 #3 L 1 H5 mean|ρ|=1.0914 #4 L11 H0 mean|ρ|=0.9972 #5 L 1 H1 mean|ρ|=0.9047
3. Browse head concepts¶
show_head_concept(l, h) shows the k corpus images with the largest |ρ(l,h)| —
the concept exemplars for that head.
signed=False(default) — rank by magnitude; finds the most strongly activated images regardless of direction.signed=True— rank by positive ρ; shows images where the head actively promotes the predicted class.layers=[...]— instead of a single(l, h), iterate over all heads in the given layers. Combine withh=to pin one head across several layers.
The first example below finds the most class-selective head for the first corpus class.
The second uses layers= to browse all heads in the last two layers without a manual loop.
# Most class-selective head for class 0, restricted to LAYERS (or globally if LAYERS is None)
l, h = store.top_head_for_class(class_idx=0, layers=LAYERS)
print(f'Most selective head for class 0: L{l} H{h}')
store.show_head_concept(l, h, k=8, show_heatmaps=True)
Most selective head for class 0: L11 H0
# All heads in LAYERS — set LAYERS above to focus on specific layers
store.show_head_concept(layers=LAYERS, k=8, show_heatmaps=True)
