Attention Heatmaps for Diffusion Models: Turning the Black‑Box into Explainability

Text‑to‑image diffusion models typically consist of three components: a CLIP text encoder that turns a prompt into token embeddings, a U‑Net that iteratively denoises latent images, and a VAE that decodes the final latent representation into pixels.

Model interpretability—explaining why a trained network produces a particular output—can be achieved with attention heatmaps because attention provides a normalized, human‑readable map of "what attends to what" without adding significant computational cost.

To obtain the attention matrices, the author recomputes the attention from the hidden Q, K, V tensors of each U‑Net layer. The logits are calculated as S = QKᵀ / √d (where d is the per‑head dimension) and then passed through softmax to get the attention matrix A. Both the pre‑softmax logits S and the post‑softmax matrix A are retained for later analysis.

Self‑attention uses queries and keys derived from image latent patches, producing a P×P matrix that maps each patch to every other patch. Cross‑attention uses image patches as queries and CLIP token embeddings as keys, yielding a P×T matrix that maps patches to prompt tokens. Because different resolutions r produce matrices of different sizes, the author aggregates results per resolution by averaging over attention heads, diffusion steps, and layers of the same resolution.

Method 1: Image‑to‑Text – Fix an image patch p₀ and read the corresponding row of the aggregated cross‑attention matrix Āᵣ. After discarding non‑word tokens, the remaining token scores are normalized and visualized by coloring each token according to its weight. For the prompt "A cat sitting on a red sofa," clicking the cat’s chest highlights the token cat , clicking the sofa’s center highlights sofa , and clicking near the cat‑sofa boundary shows both cat and sitting with high intensity.

Method 2: Text‑to‑Image – Fix a token t₀ and read the corresponding column of the pre‑softmax logits S̄ᵣ</>. The column is reshaped to an <code>r×r grid, min‑max normalized to [0,1], and bilinearly up‑sampled to the full image size. The resulting heatmap is overlaid on the image, producing the familiar “word cat lights up the cat” visualization. Examples show the token cat lighting the cat body, red lighting the sofa, and sitting lighting both the cat and the sofa surface.

Method 3: Image‑to‑Image – Use self‑attention rows for a selected patch p₀, reshape the row to an r×r grid, and render it with the same up‑sampling pipeline as Method 2. The heatmaps reveal that patches belonging to the same object or texture receive mutually high attention (e.g., both eyes light up together, or the entire wall lights up when a background point is selected).

All numeric values (scores, averaged matrices Āᵣ and S̄ᵣ, and min‑max normalized results) are computed at the native latent resolution without any scaling. To overlay a heatmap on the full‑size image, each output pixel (i,j) is mapped back to the low‑resolution grid using the scale factor s = N/r, and the four nearest grid values are combined with weights derived from the fractional offsets fᵢ = i/s - ⌊i/s⌋ and fⱼ = j/s - ⌊j/s⌋. The weights sum to one, ensuring a smooth convex combination; higher‑resolution layers produce sharper edges, while lower‑resolution layers yield softer, blocky regions.

The author also provides a lightweight Flask application that hosts three interactive modes: Image→Text (click a pixel to recolor token labels), Text→Image (click a token to display its heatmap), and Image→Image (click a pixel to highlight related regions). The web service demonstrates the practical utility of the visualizations for debugging prompt failures and building user trust before deployment.

In summary, generating these attention heatmaps requires modest effort but yields concrete diagnostic evidence: when a prompt fails, the heatmap directly points to the offending token or image region, turning vague "generation went wrong" symptoms into clear, actionable insights about model reliance and bias.