Selective Aggregation of Attention Maps Improves Diffusion-Based Visual Interpretation

1. Problem Statement

Text-to-Image (T2I) generative models, such as Stable Diffusion, rely heavily on cross-attention mechanisms to align text prompts with image generation. While previous research (e.g., DAAM) has utilized cross-attention maps to interpret model behavior and perform tasks like segmentation, these methods typically aggregate attention maps across all attention heads indiscriminately.

The authors identify a critical gap: different attention heads within a T2I model possess distinct characteristics. Some heads are highly specialized for specific visual concepts, while others are less relevant. Current methods that average all heads may dilute concept-specific signals with noise from irrelevant heads, leading to suboptimal interpretability and segmentation accuracy. The paper aims to determine if selectively aggregating only the most relevant attention heads can improve visual interpretation and segmentation performance.

2. Methodology

The proposed approach consists of three main stages:

A. Preliminaries & Baselines

• Cross-Attention Mechanism: In T2I models, the cross-attention map $M^{(h)}_t$ at head $h$ and timestep $t$ reveals the probability distribution of text tokens over spatial image locations.
• DAAM (Diffusion-based Attention Map): The standard baseline method averages attention maps across all $H$ heads and all timesteps $T$ to create a single aggregated map for a specific token.
• HRV (Head Relevance Vector): The authors leverage a method from prior work (HRV) to quantify the relevance of each attention head to specific visual concepts. This involves generating concept-words, encoding them, and measuring their activation strength across heads to create a relevance score vector.

B. Selective Attention Map Aggregation

Instead of averaging all heads, the authors propose a Selective Aggregation strategy:

1. Scoring: Use the HRV method to score each attention head based on its relevance to a target concept (e.g., "Animals").
2. Selection: Select the top $k$ attention heads (e.g., top 20-25%, specifically 30 heads in their experiments) with the highest relevance scores.
3. Aggregation: Compute the final attention map ($\hat{M}^{Ours}$) by averaging the maps only from this selected subset of heads ($H_n$), rather than the entire set.
   $$\hat{M}^{Ours} = \frac{1}{|H_n|} \sum_{h \in H_n} F_h \left( \frac{1}{T} \sum_{t=1}^{T} M^{(h)}_t \right)$$
   (Where $F_h$ is bicubic interpolation to match image resolution.)

C. Evaluation Protocol

• Dataset: 100 images generated using Stable Diffusion v1.4 with prompts in the format "photo of a {animal}" (10 animal types).
• Ground Truth: Semantic segmentation maps generated using Grounded-SAM.
• Metric: Mean Intersection over Union (IoU) calculated by thresholding the aggregated attention maps at 0.3, 0.4, and 0.5 and comparing them to the ground truth.

3. Key Contributions

1. Selective Aggregation Framework: The paper introduces a method to filter attention heads based on concept relevance before aggregation, demonstrating that not all heads contribute equally to a specific concept.
2. Performance Improvement: The selective approach outperforms the standard DAAM method in diffusion-based object segmentation, achieving higher IoU scores.
3. Diagnostic Capability: The method provides a tool to diagnose prompt misinterpretations. By isolating heads relevant to specific concepts (e.g., "Animals" vs. "Electronics"), the authors can visually explain why a model generated multiple interpretations of an ambiguous word (e.g., "mouse" as an animal vs. a computer device).
4. Ablation Insights: The study confirms that using the least relevant heads results in significantly poor segmentation, proving that indiscriminate aggregation is suboptimal.

4. Experimental Results

• Quantitative Results (Segmentation):

  • The proposed method consistently achieved higher Mean IoU scores than DAAM across all tested thresholds (0.3, 0.4, 0.5).
  • Example: At threshold 0.4, DAAM scored 0.7540, while the proposed method scored 0.7765.
  • Head Selection Impact: Using the top 30 most relevant heads for "Animals" yielded an IoU of 0.7765 (at 0.4 threshold), whereas using the bottom 30 least relevant heads dropped the score to 0.6172.

• Qualitative Results:

  • Visualizations show that DAAM often highlights regions of "undesired focus" or lacks focus on the target object.
  • The selective method produces cleaner masks that align more precisely with the target objects.

• Misinterpretation Diagnosis:

  • In a case study with the prompt "A mouse and other office objects," the model generated both a rodent and a computer mouse.
  • Selective Visualization: The "Animals" subset of heads focused only on the rodent, while the "Electronics" subset focused only on the computer mouse.
  • DAAM Visualization: Showed focus on both objects simultaneously, making it harder to distinguish the source of the ambiguity.

• Ablation Study:

  • Testing different numbers of selected heads (20, 30, 40) revealed that 30 heads (approx. 23% of the 128 total heads in SD v1) provided the optimal balance of performance.

5. Significance and Future Work

• Interpretability: This work shifts the paradigm from "averaging everything" to "selecting the best," offering a more precise lens into how T2I models process semantic information.
• Controllability: By understanding which heads drive specific concepts, developers can better control generation or debug model failures.
• Generalizability: While the current study focuses on Stable Diffusion v1.4 and the "Animals" category, the authors argue that since attention mechanisms are fundamental to T2I models, these findings likely generalize to other architectures and concepts.
• Future Directions: The authors plan to expand the concept vocabulary beyond animals and apply the method to diverse T2I architectures to further validate the approach.

Conclusion

The paper demonstrates that selective aggregation of attention maps is a superior strategy for visual interpretation in diffusion models. By filtering out irrelevant attention heads and focusing only on those most responsive to a target concept, the method significantly improves segmentation accuracy (IoU) and provides a powerful diagnostic tool for understanding and correcting model misinterpretations.