Direct Consistency Optimization for Robust Customization of Text-to-Image Diffusion Models

Text-to-image (T2I) diffusion models, when fine-tuned on a few personal images, can generate visuals with a high degree of consistency. However, such fine-tuned models are not robust; they often fail to compose with concepts of pretrained model or other fine-tuned models. To address this, we propose a novel fine-tuning objective, dubbed Direct Consistency Optimization, which controls the deviation between fine-tuning and pretrained models to retain the pretrained knowledge during fine-tuning. Through extensive experiments on subject and style customization, we demonstrate that our method positions itself on a superior Pareto frontier between subject (or style) consistency and image-text alignment over all previous baselines; it not only outperforms regular fine-tuning objective in image-text alignment, but also shows higher fidelity to the reference images than the method that fine-tunes with additional prior dataset. More importantly, the models fine-tuned with our method can be merged without interference, allowing us to generate custom subjects in a custom style by composing separately customized subject and style models. Notably, we show that our approach achieves better prompt fidelity and subject fidelity than those post-optimized for merging regular fine-tuned models.

Paper

Similar papers

Peer review

Reviewer XWYq6/10 · confidence 5/52024-06-24

Summary

The paper introduces a novel training framework and novel sampling technique to enhance the state-of-the-art personalized and style-preserving adaptation of text-to-image diffusion models in the low-shot fine-tuning regime.

Strengths

* The training objective introduced in the paper is implicitly incentivized to learn a reward for the model for better consistency. * The decoupling of classifier-free guidance to include the learned consistency preserving model enhances image similarity with the reference image, thereby also improving image fidelity.

Weaknesses

* DCO is computationally expensive since the proposed sampling scheme makes use of two networks during inference. However, the authors also mention it explicitly in the limitations section. But it would be nice to be forthright about it in the paper's main text. * Multiple references to existing works are missing. Examples include SVDiff [1], Custom Diffusion [2], IP-Adapters [3], and, DiffuseKrona [4] (for controllable personalized generation), and MaPO [5] (reward fine-tuning through alignment). It also didn't include references to works that successfully applied LoRA in the context of adapting text-to-image models [6], [7], [8]. * The paper's main text reads incomplete without the minimum details on the datasets being used for training and evaluation. For example, in lines 219 - 224, the reader doesn't know what reference images were used. Similarly, in lines 251 - 254, there is no mention of Appendix D, which sheds light on more details. This is also the case for the text accompanying Table 1 (Lines 280 - 287). * While reporting the results of Figure 4, it would have been better to talk about $\omega_{text}$ as well because that substantially impacts the generation quality. From the accompanying text, the reader doesn't know how it was varied or if it was kept fixed. ## References [1] SVDiff: Compact Parameter Space for Diffusion Fine-Tuning, https://arxiv.org/abs/2303.11305. [2] Multi-Concept Customization of Text-to-Image Diffusion, https://arxiv.org/abs/2212.04488. [3] IP-Adapter: Text Compatible Image Prompt Adapter for Text-to-Image Diffusion Models, https://arxiv.org/abs/2308.06721. [4] DiffuseKronA: A Parameter Efficient Fine-tuning Method for Personalized Diffusion Models, https://arxiv.org/abs/2402.17412. [5] Margin-aware Preference Optimization for Aligning Diffusion Models without Reference, https://arxiv.org/abs/2406.06424. [6] Implicit Style-Content Separation using B-LoRA, https://arxiv.org/abs/2403.14572. [7] https://github.com/cloneofsimo/lora [8] https://huggingface.co/blog/lora

Questions

## Suggestions * Figure 1.B is not obvious when conveying the limitations. Perhaps mention that the subject is not faithfully preserved (hat is missing) and the style is not faithfully preserved (blue background is missing) in the case of DreamBooth. * Line 97 - 98: It would be better to also mention how training with classifier-free guidance is implemented in practice, i.e., we randomly make captions associated with images to null (caption dropout). * Since the authors use LoRA for fine-tuning, I think the memory overhead could still be kept to a minimum if they enable and disable the adapter layers within the denoiser network. See how it's implemented [here](https://github.com/huggingface/diffusers/blob/3fca52022fe0ea9aaf0a0ea8a0fc13308bf69a9f/examples/research_projects/diffusion_dpo/train_diffusion_dpo_sdxl.py#L1019) and [here](https://github.com/huggingface/diffusers/blob/3fca52022fe0ea9aaf0a0ea8a0fc13308bf69a9f/examples/research_projects/diffusion_dpo/train_diffusion_dpo_sdxl.py#L1035). This probably should improve the memory requirements quite a bit. Could the authors verify this approach? * It would be nice to enlist and cite the libraries that were used in the codebase. They do enable productive research and citing them goes a long way. * Provide a loss landscape comparison when varying the amount of DCO $\beta$. * Highlight the highest scores reported in Table 1. * If both GPT-4 and LlaVA were used for captioning, how were the combined captions used during training? ## Questions * If there's a connection between the guidance scheme introduced in InstructPix2Pix [1] and the one introduced in DCO, it would be nice to have that discussed. * Line 204: It's not clear what do the authors mean by "omitting. * For the DreamBooth and DreamBooth + PP experiments, do the authors use LoRA? If so, the text is not clear. * Was consistency sampling applied for the DreamBooth experiments? * Perhaps refer to the combination of LoRA DreamBooth + Textual Inversion as "pivotal tuning" as commonly referred in the community [2]? * With a higher value of $\omega_{con}$, I would assume that the image similarity would improve. But with increasing $\omega_{con}$, it seems to be dropping off as per Figure 4. Or am I looking at it wrongly? * Did the authors experiment with other merging techniques geared toward LoRAs [3]? ## Nits * "In specific" could be replaced with "specifically". * Line 46: "... policy optimization problem [19, 20], which amounts _for_ the consistency ..." * Line 75: "_To the best of our knowledge,_ this paper is the first to have studied consistency as a reward function for T2I personalization." * Line 83: Shouldn't $\lambda_{1}$ be $\lambda_{t}$? Similar comments for $\lambda_{0}$. ## References [1] InstructPix2Pix: Learning to Follow Image Editing Instructions, https://arxiv.org/abs/2211.09800 [2] https://replicate.com/cloneofsimo/lora-training/readme [3] https://huggingface.co/blog/peft_merging

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors are clear about limitations. Perhaps mentioning them a bit from the main text would be more sensible.

Reviewer ayXu7/10 · confidence 4/52024-07-03

Summary

Current personalized T2I models, like DreamBooth, are capable of generating personalized images by fine-tuning on a small set of reference images. However, these fine-tuned models often suffer from robustness issues, such as poor compositional capabilities with the pretrained model concepts and other fine-tuned models. The paper "Direct Consistency Optimization for Robust Customization of Text-to-Image Diffusion Models" introduces a novel fine-tuning method called Direct Consistency Optimization (DCO) to enhance the robustness of such personalized text-to-image (T2I) diffusion models. DCO introduces a novel training objective that minimizes the deviation between the fine-tuning and pretrained models by optimizing a consistency function. The method involves computing a loss function that measures the deviation in noise prediction errors between the fine-tuned and pretrained models, thus allowing efficient implementation using a noise-prediction loss approach.

Strengths

* Novel Objective Function: The introduction of Direct Consistency Optimization (DCO) represents a novel approach to fine-tuning text-to-image diffusion models. Unlike previous methods that primarily focus on incorporating additional datasets to retain pretrained knowledge, DCO directly controls the deviation from the pretrained model * The paper effectively combines concepts from constrained policy optimization and noise prediction to derive the DCO loss function. This is a novel contribution based on my knowledge * good amount of experiments and results visualization * this work could have big potential impact how the booming application of personalized t2i diffusion models

Weaknesses

* Incremental Improvement: While the DCO method introduces a novel fine-tuning objective, the actual improvement seem marginal (based on figure 4 and figure 7) * Limited Diversity of Experiments: The experiments primarily focus on subject and style personalization within specific datasets. A broader range of experiments, including diverse set of human images. * Reproducibility Concerns: Although the methodology is well-explained, the paper does not provide detailed information on implementation specifics such as the codebase, computing resources, or detailed training schedules. This might pose challenges for reproducibility and practical adoption by the community.

Questions

have you thought about trying a kandinsky2.2/dalle2 like model which can directly condition on the image embedding for a (or multiple) reference images?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

yes its adequately address in line 632

Reviewer aXKh5/10 · confidence 3/52024-07-12

Summary

This paper introduces a novel fine-tuning objective called Direct Consistency Optimization, which regulates the deviation between fine-tuning and pre-trained models to preserve pre-trained knowledge during the fine-tuning process. Models fine-tuned using this method can be merged seamlessly without interference.

Strengths

1. The results seem good and surpass the baseline methods. 2. The direct merging of subject and style is interesting.

Weaknesses

1. The comparison with Lora is missing. Usually the results of Lora customization are very competitive. 2. In fig1, the pose of the bear in different images are mostly the same as those in the reference image, which is overfitting. It needs to compare the diversity of generated images.

Questions

As seen in weakness.

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

The manuscript includes discussions on limitations and social impacts.

Reviewer HfJW5/10 · confidence 3/52024-07-12

Summary

This paper studies the catastrophically forgetting issue in personalizing text-to-image diffusion model. The main comparable baseline is DreamBooth, which prevents forgetting by finetuning the model on a subsample of original training data while learning new concepts. The proposed method, on the other hand, directly regularizes the deviation of reference model from original model, which resulting a new objective function that essentially penalizing overfitting to the ‘hard examples’ during personalization. Empirical evaluations are conducted on subject and style transfer datasets, where the proposed method achieves superior pareto optimality on image fidelity v.s. prompt following ability, as well as better visual results compared with DreamBooth baseline.

Strengths

1. The proposed method is conceptually simple yet yield strong empirical results (if judged solely based on the displayed results) 2. One benefit of DCO is that it relies solely on the reference data, as the user does not have to gain access to the original training data for regularization during finetuning. 3. The proposed method, in particular DCO, has its novelty and value to the community, if the strong performance can be further verified by extra experiments requested in the weakness section.

Weaknesses

The empirical evaluation can benefit from extra validations. 1. Lacks human evaluation. The paper could benefit from a proper human study from third-party judgers. Currently, the empirical comparisons seem to be confined by automatic evaluation and a subset of visualizations. 2. Compounded factors are not studied separately. The Textual Inversion and comprehensive captioning techniques are deployed on top of DB and DCO. With these enabled, it is unclear what might be the interplay between these techniques and different baselines. For instance, the compatibility of DB might be less compatible with DB than DCO, since DCO only requires the reference data during training. These ablations would not necessarily diminish the proposed method; They are just very beneficial for research purposes IMO.

Questions

1. At a high level, both DCO and DBpp are trying to learn the concept while keeping certain aspects of the new model closer to the original model. It seems (e.g. line 39 and Figure 2) that the motivation for using is based primarily on empirical performance. Could the author elaborate more on the intuition of why DBpp hurts the subject fidelity?

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

Yes

Reviewer XWYq2024-08-08

Thank you for comprehensively addressing my comments and running those experiments. Hope they were not too much trouble. I only have minor comments. Furthermore, I have raised my score to 6 from 5 after looking at your comments and rebuttal. All the best! [W2] I still think they at least deserve a mention. All of them target customization of T2I diffusion models at their very core. Maybe pick the ones that you think are most relevant and aligned. [S3] Glad it helped. [S5] Would it be included in the appendix of the main draft? I think it will be helpful for practitioners as it gives an idea about convergence. [S7] Thanks for the clarification. Perhaps this could be clarified in the main text.

Authorsrebuttal2024-08-08

Thank you for your positive feedback and for raising the score. We appreciate your thoughtful comments and are glad the additional experiments were helpful. Regarding [W2], we will ensure to mention and discuss the relevant works that align with our study, as they indeed target customization of T2I diffusion models at their core. For [S3], we are pleased that the clarification was useful. We will update the technical implementation in our code release to reflect this. In response to [S5], we will include the training loss landscape in the appendix of the main draft, as we agree it will be beneficial for practitioners to understand convergence. For [S7], we will clarify this point further in the main text to enhance clarity. Thank you once again for your valuable feedback, which has significantly improved the quality of our work. If you have any further comments or suggestions, please do not hesitate to share them.

Reviewer XWYq2024-08-09

I acknowledge reading your comment, and I wish you all the best!

Reviewer ayXu2024-08-10

I ack the rebuttal

Reviewer HfJW2024-08-14

Reply to authors

Thank you for the responses, they addressed my concerns. I've improved my ratings accordingly. I suggest the authors to include the new experiments in the revision.

Authorsrebuttal2024-08-14

Thank you for your positive feedback and for updating your ratings. We appreciate your suggestion and will incorporate the details of the additional experiments in our final revision. Thank you once again for your valuable input!

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC