ConceptAttention: Diffusion Transformers Learn Highly Interpretable Features
This work addresses the need for better interpretability in AI vision models, particularly for researchers and practitioners using DiTs, though it is incremental as it builds on existing DiT architectures.
The paper tackled the problem of interpretability in multi-modal diffusion transformers (DiTs) by introducing ConceptAttention, a method that generates high-quality saliency maps to locate textual concepts in images without extra training, achieving state-of-the-art performance on zero-shot image segmentation benchmarks and outperforming 15 other methods on ImageNet-Segmentation.
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.