Seg4Diff: Unveiling Open-Vocabulary Semantic Segmentation in Text-to-Image Diffusion Transformers

Generalizability to different models

Due to the character limit, we kindly refer the reviewer to the "Generalizability to different models" in the rebuttal for revewer FeAF. We thank the reviewer for the understanding.

Clarification on expressions

We greatly appreciate the reviewer’s constructive suggestion. Below, we have conducted broader benchmarking on unsupervised segmentation settings and compared it with the suggested baseline. In addition to the SA-1B–trained model presented in the manuscript, we also include COCO-trained variants for comparison. The results are summarized below:

These results confirm that the insights from our layer-selection analysis translate into tangible improvements across multiple datasets, and they also invite future work on more sophisticated supervisory designs.

In light of the reviewer’s comment, we will replace “ultimately boosts both segmentation and generation performance” with a more precise statement: “The proposed supervision improves segmentation accuracy without degrading, and in some cases modestly enhances, image generation quality.” We apologise for any earlier overstatement.

[1] Couairon, Paul, et al. "Diffcut: Catalyzing zero-shot semantic segmentation with diffusion features and recursive normalized cut." Advances in Neural Information Processing Systems 37 (2024): 13548-13578.

Regarding improvement in generation quality

We evaluate the CLIP score to assess the alignment between generated images and their corresponding text prompts. To ensure robustness against randomness, we conduct evaluations on three diverse prompt sets: (1) 500 prompts from Pick-a-Pic [2], with five images generated per prompt; (2) 5,000 images generated using prompts from the COCO caption validation set; and (3) captions generated by CogVLM for 1,000 validation images sampled from SA-1B, which do not overlap with our training data.

Based on the reviewer’s feedback, we also evaluated our model on T2I-Compbench++ [3].

Overall, our method outperforms the baseline on nearly every metric. The only exceptions are the reasoning-focused benchmarks—specifically the 2-D/3-D spatial and non-spatial tests—where our scores trail the baseline. This shortfall is unsurprising: those tasks demand explicit spatial reasoning and action understanding, whereas our supervisory signal was designed to strengthen perceptual and segmentation cues, not high-level reasoning. In light of the reviewer’s comment, we will also rephrase the claims by slightly toning down, if the reviewer agrees. We apologise for any earlier overstatement.

Also, we will conduct and include a user study to better assess the generation quality of ours in the final version.

[2] Kirstain, Yuval, et al. "Pick-a-pic: An open dataset of user preferences for text-to-image generation." Advances in neural information processing systems 36 (2023): 36652-36663.

[3] Huang, Kaiyi, et al. "T2i-compbench: A comprehensive benchmark for open-world compositional text-to-image generation." Advances in Neural Information Processing Systems 36 (2023): 78723-78747.

Comparison with ZestGuide

We thank the reviewer for suggesting discussions regarding ZestGuide [4]. While both ZestGuide and ours leverage similar mask loss design, we notice that as ZestGuide targets layout generation, its goal is to enforce the diffusion model to generate the instance in the mask area. In contrast, our motivation is to improve multi-modal interaction within certain layers, rather than enforcing certain objects to be grounded to certain locations. This also differentiates our work with ZestGuide in terms of training configuration, where we generally fine-tune the model in a set of images, whereas ZestGuide requires per-image optimization. Nonetheless, we appreciate the reviewer’s pointer and remain open to further suggestions.

[4] Couairon, Guillaume, et al. "Zero-shot spatial layout conditioning for text-to-image diffusion models." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

Regarding the choice of timestep

We apologize for the typo in L157: the text should read “T = 14 out of 28”, not “T = 14 and T = 28”. We will correct this in the revision.

Initially, we chose t=14 as our early experiments on ADE20K showed the best result for t=14 when we conducted a coarse search over the timesteps. To further explore the effects of timesteps, we evaluate over all 28 timesteps in VOC and COCO-Object, and present the results below. Note that we only show some of the best and worst performing layers for brevity.

The results reveal that t=8 in fact consistently outperforms our original choice of t=14 for both VOC and COCO-Object, with t=6 and 7 similarly showing strong results. While we default the timestep to t=14 for additional experiments presented in the rebuttals to avoid confusion, we thank the reviewers for pointing this out and promise to update the scores as well as include the ablation on timesteps in the revised manuscript.

Typos

Thank you for pointing out the typos. We apologize for any confusion caused by this and will update our manuscript accordingly.

Limitations of our method

We apologize for not including the limitations of our method in the supplementary material. To address this oversight, we provide the limitations below, which will certainly be included in our revised version.

Our study deliberately focuses on dense perception and image‐generation tasks, leaving reasoning-centric evaluations (e.g., action recognition or spatial-relations QA) to future work because our supervision signal targets segmentation semantics rather than high-level reasoning. In addition, we attach the loss to a single “sweet-spot” layer per backbone for maximal simplicity; while this already delivers strong gains, a more exhaustive layer/timestep sweep—or the addition of auxiliary heads such as optical flow or pose—could yield incremental improvements without altering our core conclusions.

Generalizability to different models

To demonstrate that our insights are not tied to SD3, we replicated the entire analysis on two additional MM‑DiT variants—SD 3.5 [1] and Flux [2]. Summarizing the results, we observed consistent trends across all models (Note that the exact layer‑wise statistics can differ between models).

Because the rebuttal cannot include external media, we present the key quantitative results here and will supply the corresponding plots and qualitative examples in the camera‑ready version. Following the methodology in the main paper, we use three complementary evaluation protocols:

• Attention‑norm analysis – We compute the value‑vector norms for image, text, and actual prompt tokens at every layer, highlighting only the most informative layers for readability.
• Attention‑perturbation analysis – For each layer, we perturb its cross‑attention and measure DINO, CLIP‑I, and CLIP‑T similarities between the original and perturbed images; larger drops indicate layers that are critical to image fidelity.
• Segmentation evaluation – Using cross‑attention maps as pseudo‑masks, we evaluate mIoU on Pascal VOC to assess mask quality.

In the table, bold indicates the highest entry for value norm and segmentation metrics, and the lowest for perturbation metrics.

Stable Diffusion 3.5 Medium

Flux.1-dev

From the table above, we can similarly observe the correlation between value norm and segmentation performance among layers for both Flux and Stable Diffusion 3.5. While SD3.5 appears to have layer 9, identical to SD3, to exhibit strong semantic grounding, Flux shows layer 12 and 17 to have a similar tendency. This hints that our observation and methodology are not proprietary for SD3, but can be applied to other DiT-based diffusion models with multi-modal attention, highlighting generalizability of our insights and findings.

Comparison with previous attention-based methods

While we acknowledge that performing segmentation based on attention maps of diffusion models is an existing idea, we highlight that our exploration mainly focuses on understanding MM-DiT-based diffusion models and explore emergent properties and reveal that certain layers are able to show strong segmentation capabilities. On the other hand, [1,2] focuses on gathering information within the attention maps to identify segments, without fully regarding the image-text interaction within the diffusion models. This requires iterative algorithms [1] or tuning highly sensitive hyperparameters [2], whereas our method can simply obtain segmentation directly from a single attention layer and show competitive results. We also highlight that despite being simple, we outperform [2] in terms of IoU, while having a deficit in accuracy.

We wish to highlight that our key contribution lies more on the in-depth analysis and investigation of learned representations of MM-DiT, exploration of critical lays for segmentation capabilities, and demonstration of strong zero-shot segmentation performance with further possible enhancement using localization capabilities of the critical layers.

[1] Tian, Junjiao, et al. "Diffuse attend and segment: Unsupervised zero-shot segmentation using stable diffusion." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

[2] Wang, Jinglong, et al. "Diffusion model is secretly a training-free open vocabulary semantic segmenter." IEEE Transactions on Image Processing (2025).

Unified expressions

Thank you for pointing out the inconsistency in expressions. We will update our manuscript to unify those accordingly.

Limitations of our method

We apologize for not including the limitations of our method in the supplementary material. To address this oversight, we provide the limitations below, which will certainly be included in our revised version.

Our study deliberately focuses on dense perception and image‐generation tasks, leaving reasoning-centric evaluations (e.g., action recognition or spatial-relations QA) to future work because our supervision signal targets segmentation semantics rather than high-level reasoning. In addition, we attach the loss to a single “sweet-spot” layer per backbone for maximal simplicity; while this already delivers strong gains, a more exhaustive layer/timestep sweep—or the addition of auxiliary heads such as optical flow or pose—could yield incremental improvements without altering our core conclusions.

Generalizability to different models

Due to the character limit, we kindly refer the reviewer to the "Generalizability to different models" in the rebuttal for reviewer FeAF. We thank the reviewer for the understanding.

Clarification on the method and the task

Our framework offers both a training-free path and a fine-tuning stage, neither of which do not require ground-truth masks—at training or inference. This lets us evaluate in zero-shot and open-vocabulary settings without sacrificing scope. When adaptation is desired, we fine-tune on pseudo-labels generated automatically by Segment Anything (SAM), so the method remains label-free and readily transferable to new domains and datasets.

Extracting high-quality segments directly from diffusion models is particularly valuable: it turns the creation of otherwise costly segmentation labels into a by-product of sampling foundation models. Our study therefore focuses on uncovering the emergent segmentation behavior of MM-DiT diffusion models and on showing how a lightweight fine-tune can efficiently harness these properties for downstream tasks.

Comparison with more baselines

We thank the reviewer for emphasizing the need to include strong baselines in our comparison. Below, we clarify our evaluation choices:

Segment Anything (SAM) We did not include SAM [1] because its prompt-driven, open-set mask-generation paradigm is fundamentally different from our class-conditional diffusion approach and thus out of scope as a peer baseline. Instead, SAM serves as an upper bound on achievable mask quality.

Diffuse, Attend, and Segment We kindly invite the reviewer to revisit Table 2, where the “Diffuse, Attend and Segment [2]” is already one of our core comparators: its performance appears as a name of “DiffSeg”.

CLIP-based non-diffusion methods We would like to highlight that CLIP-based methods are not fair to compare with ours, where CLIP-based evaluations [3,4] assume no ground-truth class labels given, whereas diffusion-based methods are evaluated with ground-truth class names provided. This difference in setting comes from the inherent characteristics of CLIP and diffusion models, where CLIP needs distractors, while diffusion models need highly descriptive information regarding the image.

Although this setting mismatch precludes a perfectly fair head-to-head comparison, we provide a comparison below:

[1] Kirillov, Alexander, et al. "Segment anything." Proceedings of the IEEE/CVF international conference on computer vision. 2023.

[2] Tian, Junjiao, et al. "Diffuse attend and segment: Unsupervised zero-shot segmentation using stable diffusion." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

[3] Lan, Mengcheng, et al. "Proxyclip: Proxy attention improves clip for open-vocabulary segmentation." European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2024.

[4] Zhang, Dengke, Fagui Liu, and Quan Tang. "Corrclip: Reconstructing correlations in clip with off-the-shelf foundation models for open-vocabulary semantic segmentation." arXiv preprint arXiv:2411.10086 (2024).

Benefit of our method for from-scratch training

We thank the reviewer for suggesting a valuable point for discussion. Although we can not afford the computational budget to train a full diffusion model from scratch, we believe that our training strategy could benefit the model even when training from scratch. Given that we are able to improve both generation and segmentation performance, we consider our method to provide training signals that align with the diffusion model’s knowledge. A recent work, REPA [5] trains diffusion networks from scratch without LoRA, while enforcing a feature‑alignment loss against DINO. This direct supervision accelerates convergence and improves generation fidelity, showing that alignment losses can be effective even when the entire model is updated. Nonetheless, REPA also shows that the layer selection for applying the alignment loss is crucial for both sample quality and training speed. Similarly, we expect layer-selection to still be of importance as it seems crucial to maintain alignment with the layer that naturally shows emergent properties, in terms of training efficiency, generation quality, and segmentation capabilities. We will add such discussions to the paper upon revision.

[5] Yu, Sihyun, et al. "Representation alignment for generation: Training diffusion transformers is easier than you think." arXiv preprint arXiv:2410.06940 (2024).

Additional results on the perturbation experiment

To probe the importance of each cross‑attention block, we repeated the perturbation study but fully masked the image‑to‑text (I2T) channels one layer at a time. Qualitatively, removing I2T attention produces artifacts that mirror those seen when we previously applied a Gaussian blur to the same pathway, indicating comparable disruption mechanisms. To assess the effects in a quantitative manner, we adopted DreamBooth’s protocol [7] for generating 50 images per perturbation and reported the average change in DINO, CLIP‑I, and CLIP‑T scores relative to the unperturbed baseline. All three metrics dip most sharply when layer 9 is masked, underscoring its pivotal role in aligning text prompts with visual content.

We thank the reviewer for the valuable suggestion, and we will include the qualitative and quantitative results as well as the discussions in the final version.

[7] Ruiz, Nataniel, et al. "Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2023.

Limitations of our method

We apologize for not including the limitations of our method in the supplementary material. To address this oversight, we provide the limitations below, which will certainly be included in our revised version.

Our study deliberately focuses on dense perception and image‐generation tasks, leaving reasoning-centric evaluations (e.g., action recognition or spatial-relations QA) to future work because our supervision signal targets segmentation semantics rather than high-level reasoning. In addition, we attach the loss to a single “sweet-spot” layer per backbone for maximal simplicity; while this already delivers strong gains, a more exhaustive layer/timestep sweep—or the addition of auxiliary heads such as optical flow or pose—could yield incremental improvements without altering our core conclusions.

Generalizability to different models

To demonstrate that our insights are not tied to SD3, we replicated the entire analysis on two additional MM‑DiT variants—SD 3.5 [1] and Flux [2]. Summarizing the results, we observed consistent trends across all models (Note that the exact layer‑wise statistics can differ between models).

Because the rebuttal cannot include external media, we present the key quantitative results here and will supply the corresponding plots and qualitative examples in the camera‑ready version. Following the methodology in the main paper, we use three complementary evaluation protocols:

• Attention‑norm analysis – We compute the value‑vector norms for image, text, and actual prompt tokens at every layer, highlighting only the most informative layers for readability.
• Attention‑perturbation analysis – For each layer, we perturb its cross‑attention and measure DINO, CLIP‑I, and CLIP‑T similarities between the original and perturbed images; larger drops indicate layers that are critical to image fidelity.
• Segmentation evaluation – Using cross‑attention maps as pseudo‑masks, we evaluate mIoU on Pascal VOC to assess mask quality.

In the table, bold indicates the highest entry for value norm and segmentation metrics, and the lowest for perturbation metrics.

Stable Diffusion 3.5 Medium

Flux.1-dev

From the table above, we can similarly observe the correlation between value norm and segmentation performance among layers for both Flux and Stable Diffusion 3.5. While SD3.5 appears to have layer 9, identical to SD3, to exhibit strong semantic grounding, Flux shows layer 12 and 17 to have a similar tendency. This hints that our observation and methodology are not proprietary for SD3, but can be applied to other DiT-based diffusion models with multi-modal attention, highlighting the generalizability of our insights and findings.

Ablation on the choice of timestep

Initially, we chose t=14 as our early experiments on ADE20K showed the best result for t=14 when we conducted a coarse search over the timesteps. To further explore the effects of timesteps, we evaluate over all 28 timesteps in VOC and COCO-Object, and present the results below. Note that we only show some of the best and worst performing layers for brevity.

The results reveal that t=8, in fact, consistently outperforms our original choice of t=14 for both VOC and COCO-Object, with t=6 and 7 similarly showing strong results. While we default the timestep to t=14 for additional experiments presented in the rebuttals to avoid confusion, we thank the reviewers for pointing this out and promise to update the scores as well as include the ablation on timesteps in the revised manuscript.

Comparison with COCO-trained model

We agree that the masks predicted by SAM can be overly fine‑grained, which can hinder training by enforcing the attention maps within the diffusion model to very small regions. In this regard, we carry out the same training with COCO images and masks, and report the results in weakly and unsupervised settings, as well as the generation performance with Clipscore below.

We can observe that not only both weakly and unsupervised performance showed improvements, but also the generation performance has improved when compared to being trained on SA-1B. This suggests the importance of mask granularity in our fine‑tuning framework. We appreciate the reviewer’s constructive suggestion and will incorporate the COCO‑trained model into the manuscript upon revision.

Comparison with ODISE and SemFlow

We appreciate the reviewer for suggesting ODISE [1] and SemFlow [2] for comparison. While both models are based on diffusion models, we notice that both ODISE and SemFlow are trained with different levels of supervision and evaluated in different settings, which prevents a fair comparison to our work. SemFlow is trained in a closed-set manner, which cannot perform open-vocabulary segmentation, and ODISE requires full supervision of having class-labeled masks. By contrast, our model operates in a weakly supervised, zero-shot setting and—without any class labels, while lightweight fine-tuning with inexpensive SA-1B or COCO masks pushes the scores. Nonetheless, we show comparisons in the table above as a reference, and we will add clarifications regarding the different settings of different methods to the manuscript.

[1] Xu, Jiarui, et al. "Open-vocabulary panoptic segmentation with text-to-image diffusion models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2023.

[2] Wang, Chaoyang, et al. "Semflow: Binding semantic segmentation and image synthesis via rectified flow." Advances in Neural Information Processing Systems 37 (2024): 138981-139001.

Limitations of our method

We apologize for not including the limitations of our method in the supplementary material. To address this oversight, we provide the limitations below, which will certainly be included in our revised version.

Our study deliberately focuses on dense perception and image generation tasks, leaving reasoning-centric evaluations (e.g., action recognition or spatial-relations QA) to future work because our supervision signal targets segmentation semantics rather than high-level reasoning. In addition, we attach the loss to a single “sweet-spot” layer per backbone for maximal simplicity; while this already delivers strong gains, a more exhaustive layer/timestep sweep—or the addition of auxiliary heads such as optical flow or pose—could yield incremental improvements without altering our core conclusions.