ConceptAttention: Diffusion Transformers Learn Highly Interpretable Features

TL;DR

This work tackles the interpretability of multi-modal diffusion transformers by introducing ConceptAttention, which repurposes DiT attention layers to create contextualized concept embeddings and high-quality saliency maps without retraining. By projecting image patch outputs into the attention output space of these concept embeddings, the method yields sharper, open-set saliency maps that can localize textual concepts in images and generalize to video. Across Flux DiT, Stable Diffusion 3.5 Turbo, and CogVideo X, ConceptAttention achieves state-of-the-art zero-shot image segmentation on ImageNet-Segmentation and Pascal VOC, outperforming CLIP-, DINO-, and UNet-based interpretability baselines. The results demonstrate that DiT representations are highly transferable to downstream vision tasks, offering a path toward more transparent and controllable diffusion models.

Abstract

Do the rich representations of multi-modal diffusion transformers (DiTs) exhibit unique properties that enhance their interpretability? We introduce ConceptAttention, a novel method that leverages the expressive power of DiT attention layers to generate high-quality saliency maps that precisely locate textual concepts within images. Without requiring additional training, ConceptAttention repurposes the parameters of DiT attention layers to produce highly contextualized concept embeddings, contributing the major discovery that performing linear projections in the output space of DiT attention layers yields significantly sharper saliency maps compared to commonly used cross-attention maps. ConceptAttention even achieves state-of-the-art performance on zero-shot image segmentation benchmarks, outperforming 15 other zero-shot interpretability methods on the ImageNet-Segmentation dataset. ConceptAttention works for popular image models and even seamlessly generalizes to video generation. Our work contributes the first evidence that the representations of multi-modal DiTs are highly transferable to vision tasks like segmentation.

Figures (17)

• Figure 1: ConceptAttention produces saliency maps that precisely localize the presence of textual concepts in images. We compare Flux raw cross attention, DAAM tang_what_2022 with SDXL, and TextSpan gandelsman_interpreting_2024 for CLIP.
• Figure 2: ConceptAttention augments multi-modal DiTs with a sequence of concept embeddings that can be used to produce saliency maps. (Left) An unmodified multi-modal attention (MMAttn) layer processes both prompt and image tokens. (Right) ConceptAttention augments these layers without impacting the image appearance to create a set of contextualized concept tokens.
• Figure 3: ConceptAttention can generate high-quality saliency maps for multiple concepts simultaneously. Additionally, our approach is not restricted to concepts in the prompt vocabulary.
• Figure 4: ConceptAttention produces higher fidelity raw scores and saliency maps than alternative methods, sometimes surpassing in quality even the ground truth saliency map provided by the ImageNet-Segmentation task. Top row shows the soft predictions of each method and the bottom shows the binarized predictions.
• Figure 5: (a) MMAttn combines cross and self attention operations between the prompt and image tokens. (b) Our ConceptAttention allows the concept tokens to incorporate information from other concept tokens and the image tokens, but not the other way around.