Open-Vocabulary Attention Maps with Token Optimization for Semantic Segmentation in Diffusion Models

TL;DR

This paper introduces Open-Vocabulary Attention Maps (OVAM), a training-free extension of diffusion models that enables attention maps for arbitrary words, thereby producing semantic segmentation pseudo-masks beyond the prompt vocabulary. It adds a token-optimization pipeline that learns class-descriptor tokens from a single annotation per class, significantly improving mask accuracy without architectural changes or retraining. OVAM combines an attribution prompt with cross-attention and applies self-attention refinements and CRF post-processing to yield high-quality binaries suitable for training segmentation models. Experimental results show large gains in pseudo-mask quality and downstream segmentation performance, including strong improvements when using OVAM-generated data to train models under data-scarce conditions and notable gains when combined with real data. This approach provides a practical route to open-vocabulary segmentation and efficient synthetic data generation for diffusion-based methods.

Abstract

Diffusion models represent a new paradigm in text-to-image generation. Beyond generating high-quality images from text prompts, models such as Stable Diffusion have been successfully extended to the joint generation of semantic segmentation pseudo-masks. However, current extensions primarily rely on extracting attentions linked to prompt words used for image synthesis. This approach limits the generation of segmentation masks derived from word tokens not contained in the text prompt. In this work, we introduce Open-Vocabulary Attention Maps (OVAM)-a training-free method for text-to-image diffusion models that enables the generation of attention maps for any word. In addition, we propose a lightweight optimization process based on OVAM for finding tokens that generate accurate attention maps for an object class with a single annotation. We evaluate these tokens within existing state-of-the-art Stable Diffusion extensions. The best-performing model improves its mIoU from 52.1 to 86.6 for the synthetic images' pseudo-masks, demonstrating that our optimized tokens are an efficient way to improve the performance of existing methods without architectural changes or retraining.

Figures (13)

• Figure 1: (a) We introduce Open-Vocabulary Attention Maps (OVAM), a training-free extension for text-to-image diffusion models to generate text-attribution maps based on open-vocabulary descriptions. Our approach overcomes the limitations of existing methods constrained by words contained within the prompt Attn2maskDAAMSecretSegmenterdiffumask. (b) Our token optimization process enhances the creation of accurate attention maps, thereby improving the performance of existing semantic segmentation methods based on diffusion attentions Attn2maskDAAMDatasetDMli2023grounded. (c) Finally, we validate the utility of OVAM in producing synthetic images with precise pixel-level annotations.
• Figure 2: A schematic representation of the OVAM generation process (red module) utilizing the Stable Diffusion architecture sd (rest of modules). The example synthesizes an image using the generator prompt monkey with hat walking. During the OVAM generation, pixel queries $Q$ are extracted from the denoising network. These pixel queries are combined with the text embedding $K'$ corresponding to the attribution prompt mouth, constructing the OVAM heatmap $D_{X,k}(X')$, which highlights the monkey's mouth in the synthesized image.
• Figure 3: Diagram illustrating the optimization process for an $X'$ composed of two tokens: a background and a car token. OVAM generates one heatmap for each token, and the optimization updates $X'$ to align the attentions generated with the target mask.
• Figure 4: Comparative visualization of attention maps. Left images show attention using class name tokens (bird, bicycle, sofa and person) while the images on the right use optimized tokens with a training set that does not contain these images.
• Figure 5: Class-performance comparison (IoU) of pseudo-masks generated by existing methods with/without the OVAM's token optimization, for the classes of the synthetic dataset COCO-cap.