CoMat: Aligning Text-to-Image Diffusion Model with Image-to-Text Concept Matching

We sincerely appreciate your valuable comments. We found them extremely helpful in improving our draft. We address each comment in detail, one by one below.

Comment 1. Complexity of Implementation

Thank for your feedback. Indeed, we acknowledge the intricate nature of our approach. However, we argue that the complex design is necessary. Our method's multiple components work as a whole to address the misalignment problem while maintaining high-quality image generation:

1. Misalignment Solution: a) Concept Activation Module: Utilizes an image-to-text model to activate individual concepts within the prompt. b) Attribute Concentration: Employs a segmentation model to guide the activated concepts to their appropriate spatial locations in the generated image.
2. Quality Preservation: a) Fidelity Preservation Module: Incorporates a pre-trained diffusion model as a discriminator to mitigate potential quality degradation resulting from the image-to-text model's guidance. b) Mixed Latent Strategy: Leverages ground truth images to provide additional guidance during the diffusion process.

As for the resource overhead, instead of full finetuning, our method uses LoRA to reduce the computational cost. Please refer to the response to Comment 3 for further discussion.

Comment 2. Insufficient Comparative Analysis

Thanks for your advice. We provide more quantitative comparisons with state-of-art methods like TokenCompose [1], Structure Diffusion [2], and Attend-and-Excite [3] in the T2I-CompBench benchmark, as shown in Table 1 in the author rebuttal PDF. We will include these results in our final draft.

Comment 3. Resource Intensiveness of Loss Computation

Sorry for the confusion caused. We respectfully disagree for the following two reasons:

1. The denoising process is a must for all the training losses, not only $\mathcal{L}_{i2t}$.

   Our training process indeed needs multiple denoising steps to generate the image. However, we add multiple supervision during this process to foster the model to align the generated image with the prompt.

   The multiple denoising process starts from a pure noise $x_T$, and then we iteratively denoise the noise to produce an image $\mathcal{I}$. Of all the denoising steps, we uniformly sample $K$ steps to enable the gradient. We supervise the attention maps during denoising in these steps, which corresponds to $\mathcal{L}\_\text{pos}$ and $\mathcal{L}\_\text{neg}$. The gradient of $\mathcal{L}\_{i2t}$ and $\mathcal{L}\_{adv}$ is back-propagated through the image to the LoRA for the sampled $K$ steps.

2. More efficient and beneficial proven in previous works

   Compared with directly supervised fine-tuning the diffusion model, our training process is identical to the inference process of the diffusion model. This training and test alignment contributes to a more efficient learning process. In fact, our method only needs 2000 iters to achieve a good performance, while directly finetuning typically requires much more iteration (Please refer to the table in Limitation 1). This paradigm of training has been proven more efficient and beneficial in previous works like DPOK [5].

Question 1. Can image-to-text model offer effective pixel-level guidance?

We provide the visualization of the gradient of $\mathcal{L}\_{i2t}$ on the generated image in the Fig. 1 in the author rebuttal PDF.

Question 2. Ablation studies that demonstrate the effects of using $\mathcal{L}_{i2t}$ for training

We conduct ablation studies on the T2I-CompBench to verify the effectiveness of the $\mathcal{L}\_{i2t}$ losses in Table 4 in the paper. The $\mathcal{L}\_{i2t}$ greatly enhances the text-image alignment of the diffusion model. Please refer to Section 5.4 for more details.

Limitation 1. Increase model's complexity and computational demands

Thanks for your feedback. Although our method does introduce multiple components, these components bring about abundant supervision to foster the fast training speed and also excellent performance. Besides, we list the performance on T2I-CompBench and training cost compared with GORS [4]:

Our method well balances the training iterations and the performance. Comat-SD1.5 generally achieves better performance with roughly 2% of iterations.

Besides, all these modules are removed during inference. Therefore, no inference cost is introduced.

[1] Wang, Zirui, et al. "TokenCompose: Text-to-Image Diffusion with Token-level Supervision." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

[2] Feng, Weixi, et al. "Training-free structured diffusion guidance for compositional text-to-image synthesis." arXiv preprint arXiv:2212.05032 (2022).

[3] Chefer, Hila, et al. "Attend-and-excite: Attention-based semantic guidance for text-to-image diffusion models." ACM Transactions on Graphics (TOG) 42.4 (2023): 1-10.

$4$ 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.

$5$ Fan, Ying, et al. "Reinforcement learning for fine-tuning text-to-image diffusion models." Advances in Neural Information Processing Systems 36 (2024).

We sincerely appreciate your valuable comments. We found them extremely helpful in improving our draft. We address each comment in detail, one by one below.

Comment 1. Differences between Class-specific Prior Preservation Loss with Fidelity Preservation

Thank you for your advice. Indeed, our Fidelity Preservation (FP) module shares a similar high-level idea with the Class-specific Prior Preservation Loss (CPP Loss) introduced in [1], i.e., preserving the generation quality while finetuning. However, our method is very different in the following two aspects:

1. Target Task and Preserve Domain

   DreamBooth seeks to personalize image generation for specific objects. While the introduced CPP Loss primarily maintains generative capabilities within a narrow domain—specifically, the object class present in the training data—our proposed FP module operates within the context of text-image alignment. FP aims to preserve general generative capabilities by computing adversarial loss across the entire training dataset, encompassing a diverse range of text prompts.

2. Methodology

   DreamBooth finetunes the diffusion model with the pretraining loss, i.e., the squared error denoising loss on a certain timestamp. CPP Loss follows its form.

   In contrast, our fine-tuning procedure simulates the inference process of the diffusion model to conduct a full-step inference. We aim to directly supervise the generated image to achieve the training-test alignment. Therefore, we propose the novel FP module to leverage a discriminator to adversarially preserve its quality. The applied discriminator is also updated along with the fine-tuning process, enabling finer control of the image quality.

We will add this discussion to the final draft for clarification.

Comment 2. Potential risk of overfitting the training data

Thank you for your feedback. We conduct zero-shot evaluation on long and complex prompt in DPG-Bench in the Table 2 in our main paper.

All of our training data contains only short text prompts similar to COCO captions. The length of sentences in the training set is 12.13 words on average. The average prompt length in DPG-Bench is 78.23 words.

As shown in the Table 2 in the main paper, our method significantly enhances the alignment by over 10 points on SD1.5 and over 2 points in SDXL. The result proves that our proposed method improves the general prompt-following capability instead of overfitting to the training data. The qualitative example on DPG-Bench is shown in Fig. 11 in the Appendix.

Besides, TIFA in Table 2 in the main paper is another alignment benchmark, which our training data also does not overlap. We witnessed a 7.4 points improvement in Comat-SD1.5 and a 1.6 points improvement in Comat-SDXL. The result further reveals the generalization.

Comment 3. Identify remaining misalignment issues in the fine-tuned model

Thank you for your advice. This is a very interesting suggestion. We investigate our method proposed in [2]. We will add an extra section in the appendix to discuss our findings.

[1] 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.

[2] Du, Chengbin, et al. "Stable diffusion is unstable." Advances in Neural Information Processing Systems 36 (2024).

There is some controversy regarding the issue of generalization, but this is not a critical flaw. The experimental results with short prompts do show that the model fine-tuning has improved text-image alignment to some extent (although it is difficult to determine the extent of this improvement given the limited amount of data). Therefore, I have decided to raise my score to a Borderline accept.

Thanks for your response.

We also conduct zero-shot evaluations on the DPG-Bench, which consists of long and complex prompts. Our method exhibits excellent zero-shot performance both qualitatively and quantitatively. Please refer to our rebuttal for Comment 2 for details. We believe this could address the concern of the generalization ability to the long prompts. Besides, we visualize some successful attacks, and we find that the generated prompt contains many meaningless letters, which could largely affect the generation process. An example is: "these small, colorful fiddler crab are known for their distinctively asymmetric claws, with one being much larger than the other. males use their enlarged claw to attract mates and defend their territory., hair, joheat col< troll < shepherds a < ; children "" < lake rest <". Therefore, we argue that the result in Experiment 3 may not fully represent the result of our method's generalization ability.

We sincerely appreciate your valuable comments. We found them extremely helpful in improving our draft. We address each comment in detail, one by one below.

Comment 1. Training cost is disproportionate

Thank you for your feedback. We will open-source our training code for reproducibility.

We test on more training iterations in our pilot study but no apparent gain is witnessed. So we keep 2000 iterations as default. Also, we test the different settings for training parameters: rank 128, 256 for LORA, and full parameter finetuning. We find no apparent gain for larger rank or full parameter fintuning. We finally choose LoRA with rank 128 for efficient training. LoRA of relatively small rank could already convert the model with different abilities. Various previous works adopt this setting and achieve good performance [1, 2].

Comment 2. Real training cost

Thank you for your feedback.

1. Pixel space supervision

   During the full steps inference, we uniformly sample $K$ steps to perform the pixel space supervision, namely calculating the $\mathcal{L}\_\text{pos}$ and $\mathcal{L}\_\text{neg}$ loss. The $K$ is 5 in our experiments. Therefore, only 10% of steps in the inference require gradient and pixel space supervision. This may partially account for the fast training speed.

2. MLLM inference

   In practice, we leverage BLIP trained on captioning tasks instead of MLLM (Please refer to Appendix A.3 in the main paper for detailed information). This choice accelerates the computation of the $\mathcal{L}\_{i2t}$ loss.

Comment 3. Assign attributes to multiple same-name objects

This is a really interesting question. Currently, it is difficult to assign the attribute to multiple same-name objects by only using the segmentation model. However, we argue that the image-to-text model with advanced image-text understanding ability may distinguish the attribute for each same-name object and offer valid guidance. As shown in the Fig. 1 in the author rebuttal PDF, the image-to-text model could also provide pixel-level supervision to the image.

We will work on to solve this question in our future work.

Question 1. How to use image-to-text LLVM to provide a matching score?

Sorry for the confusion caused. The scoring is conducted by calculating the loglikelihood for the LLVM to output the prompt as the caption, given the generated image.

In practice, we leverage the LLVM specifically trained for the captioning tasks (e.g., BLIP) to do the scoring. Their scoring capability comes with their training nature. These captioning models are trained with the negative loglikelihood loss, namely, the model needs to maximize the probability for generating the caption given the corresponding image. Therefore, whenever the generated image does not align with the text prompt, the LLVM will output low loglikelihood, namely, the $\mathcal{L}\_{i2t}$ is big. We treat the $-\mathcal{L}\_{i2t}$ as the alignment score of LLVM and try to maximize it, i.e., minimize $\mathcal{L}\_{i2t}$.

Question 2. More details about Mixed Latent Strategy

Sorry for the confusion caused. We detail our mixed latent strategy below:

The mixed latent strategy contains two types of latents in the fine-tuning procedure, i.e., the latent starting from the pure noise and the noisy latent from the GT Images.

1. Latent starting from the pure noise

   This serves as the main branch in our pipeline.

   Our fine-tuning process shares the same procedure to generate an image as the diffusion model does in the inference time. We uniformly sample $K$ steps from all the inference steps to enable the gradient. Therefore, the latent is sampled from the pure noise $\mathcal{N}(0,I)$. We iteratively denoise it to obtain the generated image. The image is then used to calculate the $\mathcal{L}\_{i2t}$ and $\mathcal{L}\_{adv}$ loss. It is also sent to the segmentation model to provide the object mask for computing the $\mathcal{L}\_\text{pos}$ and $\mathcal{L}\_\text{neg}$.

   The latent starting from the pure noise corresponds to the upper left part in Fig. 4. in the main paper.

   Please refer to [2] for how to receive the gradient from the loss.

2. Noisy latent from the GT Images

   We also aim to inject information from the GT images to stabilize the fine-tuning process.

   We randomly sample a timestamp $\tau$ from a pre-defined range $[T\_1, T\_2]$. Then we obtain $x\_{\tau}$ by adding the timestamped noise $\epsilon\_{\tau}$ on the latent of the GT Image $x_0$. We also iterativaly denoise this noisy GT latent to get $\hat{x}\_0$ as we do for the latents starting from the pure noise. This $\hat{x}\_0$ is only used to calculate the $\mathcal{L}\_{i2t}$ loss.

   The latent starting from the noisy GT corresponds to the upper left part in Fig. 4 in the main paper.

Question 3. Will the code be open-sourced?

Thanks for your feedback. We will open-source the training code.

[1] Sun, Jiao, et al. "Dreamsync: Aligning text-to-image generation with image understanding feedback." Synthetic Data for Computer Vision Workshop@ CVPR 2024. 2023.

[2] Wu, Xiaoshi, et al. "Deep Reward Supervisions for Tuning Text-to-Image Diffusion Models." arXiv preprint arXiv:2405.00760 (2024).

We sincerely appreciate your valuable comments. We found them extremely helpful in improving our draft. We address each comment in detail, one by one below.

Comment 1. Limited technical novelty

Thank you for your feedback. We respectfully disagree with the reviewer's assessment. Indeed, our method follows the basic reward fine-tuning pipeline explored in [1, 2]. We argue that we only adopt this basic paradigm for fine-tuning the diffusion model. We introduce a variety of novel designs in our procedure to specifically solve the underlying problems in the image-text alignment issue. It is inappropriate to simply classify our method as the same approach of TODO(AlignProp, Draft). As stated in the strength section of reviewer 9f7Y, it is innovative to introduce the image-to-text model to solve the image-text alignment task. And the other modules are well-designed for tackling specific challenges in image-text alignment task.

We summarize our technical novelty as follows:

1. Image-to-text model as reward model

   We investigate how to use the hidden knowledge of the pre-trained image-to-text model to supervise the diffusion model for alignment. We extend the basic reward fine-tuning pipeline to the image-text alignment task and prove that the pre-trained image-to-text model is a suitable reward model for a large setting of training.

2. Attribute Concentration module

   Our visualisation results reveal the activated concepts often fails to mapping to the correct area in the image, which causes the misalignment. We introduce the novel attribute concentration module to supervise on the pixel level to address this issue.

   Previous works like [1, 2] only supervises the generated image, omitting guiding the middle steps when generating. How to add supervision for finer-grained control in these steps has not been explored before our method.

3. Fidelity Preservation module

   [1, 2] either focus on the aesthetic task or is limited to a very small setting (tens of prompts for the training data). How to effectively preserve the generation capability in a larger setting remains unexplored. We introduce the novel Fidelity Preservation module to leverage the underlying knowledge of the pre-trained diffusion model to be the discriminator. This adversarial design effectively solves the image corruption.

4. Mixed Latent Strategy

   We also novelly propose to mix the noisy GT latent with the latents starting from pure noise. This design facilitates the information from the GT image and stabilizes the fine-tuning procedure. While [1, 2] only focus on the noise from pure noise.

Comment 2. Thorough user study

Thanks for your advice. We extend our original user preference study introduced in Appendix A.1.

We categorize the alignment metric in the following aspects as in DSG1K: entities, attributes, relations and global. Entities contain the whole object and the part of the object. Attributes contain color, type, etc. Relations contain spatial and actions. Global contains other aspects like illumination, etc.

5 participants are asked to evaluate the 100 image pairs generated by SDXL and CoMat-SDXL. For each alignment metric, we set 1 as aligned and 0 as not aligned. We show the result in Table 2 in the author rebuttal PDF. Our method enhance the baseline model in all the metrics with the most significant improvement on entities.

Comment 3. Weak Contribution

Sorry for the confusion caused. We respectfully disagree the reviewer's assessment for the following reasons:

1. No fundamental challenges solved

   First of all, we argue that the fundamental challenges of the text following ability span various domains. Both the reviewer's mentioned tasks (i.e., the counting-aware prompts and spatial-aware prompts) and other tasks like object existence and attribute binding are of the same importance. If the object concept in the prompt is not activated, the object will not even appear on the image, let alone the spatial, counting problem. Therefore, it is inappropriate to prior the counting-aware or spatial-aware tasks to the other tasks, and hence didiminish the importance of our work.

   Besides, the results in Table 1 in the main paper reveals that our method successfully addresses the spatial-aware prompts and bring significant improvement to both SD1.5 and SDXL, where nearly 80% increase is witnessed for SD1.5. We also add the evaluation result for the counting-aware prompt (numeracy) introduced in [3]. The result is shown in Table 3 in author rebuttal PDF. Our training data does not contain prompts specifically designed to solve the counting-aware task, the improvement further testify the effectiveness of our method.

   We will investigate how to specifically address the spatial-aware and counting-aware tasks in the future work.

2. Focus on improving benchmark performance

   We provide both quantitative and qualitative evaluation for the zero-shot performance on long and complex text prompts in Table 2 and Fig. 11 in the main paper. The huge improvement proves our method's generalizability to various scenarios in the alignment tasks.

   Besides, we provide visualization results in Fig. 12 to 14 in the main paper to prove our method improve the prompt following ability across different tasks. These results includes various challenges like spatial-aware prompts (the bottom of Fig. 13), attribute-aware prompts (the spider and the cabin in Fig. 12), complex prompts (the lighthouse and the man in Fig.12), etc.

[1] Aligning Text-to-Image Diffusion Models with Reward Backpropagation, arXiv 2023

[2] Directly Fine-Tuning Diffusion Models on Differentiable Rewards, ICLR 2024

[3] T2I-CompBench++: An Enhanced and Comprehensive Benchmark for Compositional Text-to-image Generation