ConceptAttention: Diffusion Transformers Learn Highly Interpretable Features

TL;DR

ConceptAttention demonstrates that diffusion transformer attention layers can be repurposed to generate highly interpretable, sharp saliency maps for visual concepts, achieving state-of-the-art zero-shot segmentation performance.

Contribution

This work introduces ConceptAttention, a novel method that leverages diffusion transformer attention layers for interpretable saliency maps without additional training.

Findings

Produces sharper saliency maps than cross-attention methods

Achieves state-of-the-art zero-shot segmentation on ImageNet-Segmentation

Generalizes to video generation tasks

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.