Diff4Seg: Unveiling Open-Vocabulary Semantic Segmentation in Text-to-Image Diffusion Transformers
Chaehyun Kim, Heeseong Shin, Heeji Yoon, Eunbeen Hong, Anurag Arnab, Paul Hongsuck Seo, Sunghwan Hong, Seungryong Kim
Text-to-image diffusion models excel at translating language prompts into photorealistic images by implicitly grounding textual concepts in their cross-modal attention mechanisms. Prior work has exploited these attention maps for downstream tasks such as editing, inpainting, and zero-shot open-vocabulary semantic segmentation (OVSS), but a detailed understanding of how these maps contribute to image generation remains limited. Recent architectural advances like Multi-Modal Diffusion Transformers (MM-DiTs) introduce joint self-attention over concatenated image and text tokens, enabling richer and more scalable cross-modal alignment. In this work, we systematically analyze the attention structures of MM-DiT, focusing on how specific heads and layers propagate semantic information and influence generation quality. By decomposing attention score distributions and attention norms across layers, we identify a subset of heads that consistently align text tokens with spatially coherent image regions, naturally yielding high-quality zero-shot segmentation masks. We then introduce a lightweight LoRA-based fine-tuning method to enhance the semantic grouping capabilities of these heads without degrading— and often improving—image fidelity. Our findings demonstrate that semantic alignment is an emergent property of diffusion transformers and can be selectively amplified to improve both dense recognition and generative performance, paving the way toward unified models that bridge generation and perception.