Adding Conditional Control to Diffusion Models with Reinforcement Learning

Q. This paper does not perform an ablation analysis of the proposed method.

We have indeed conducted an ablation analysis on the guidance strength parameter $\gamma$, as detailed in Appendix H.3. Selecting an optimal KL weight is critical. In practice, we found that $\gamma$ values between 5 and 20 are generally effective, enabling accurate conditional generations for the compressibility task. Additionally, guidance strength $\gamma$ can be freely adjusted during inference (see Appendix D). Due to page limitations, we have included this ablation study in the appendix.

Weakness: How to formalize the problem into MDP, the state, the action, and the transfer function?

Thanks for reminding this! In the updated manuscript, we have included a detailed introduction to ​​the formulation of fine-tuning diffusion models via RL in Appendix B. If you find any part unclear or need any further clarification, please let us know.

Q: Why does the experimental verification compare mainly with methods without classifiers, despite CTRL being a splitter-based method?

Thank you for your question. While we are unsure what the reviewer specifically means by splitter-based method is (and would appreciate further clarification), we understand the concern about comparisons with a classifier-based baseline. We would be happy to clarify more if we have a misunderstanding.

We would like to emphasize that the traditional classifier-based diffusion models rely on mapping $x_t \to y$ (as noted in Table 1). A key drawback of this approach is the classifier’s accuracy in predicting $y$ from intermediate $x_t$​, which can lead to cumulative errors during the diffusion process. Reconstruction guidance methods have been proposed to address this. These methods approximate intermediate states $x_t$​ back to the original input space $x_0$​, allowing classifiers to be trained on $x_0 \to y$ directly.

In the experimental part, we have included a classifier-based diffusion baseline as reconstruction guidance, providing a fair and comprehensive comparison.

Q. “The proposed method can be adapted to any off-the-shelf RL algorithms.”

Thank you for raising this concern. We agree with you. Our intention was simply to point out that many types of algorithms can be applied (though not all RL algorithms). We have revised our wording to “many types of off-the-shelf RL algorithms” accordingly.

Q. Citing the original paper proposing the PPO method for completeness.

Thank you for pointing this out. We have added a citation of the original PPO paper in the revised version.

We sincerely appreciate the reviewer’s acknowledgment of the novelty, theoretical contributions, and reproducibility of our work. Your detailed and thoughtful feedback is greatly valued. In this response, we have addressed your questions and concerns by:

1. providing additional comparisons with more advanced conditional guidance methods and clarifying the ablation analysis we have conducted
2. adding Appendix B, which provides a detailed overview of the formulation for fine-tuning diffusion models using reinforcement learning
3. providing clarifications on the comparison between our algorithm and methods without a classifier
4. updating the wording of “any off-the-shelf RL algorithms”
5. adding a citation to the original PPO paper

Please note that we have updated a new version of the manuscript with the changes highlighted in red. We are happy to answer more if further clarification is needed.

Q. This paper lacks a comparison with the latest conditional methods.

Thank you for pointing this out! We have included additional baselines for the 4-class compressibility task (Section 6.1) and the multi-task conditional generation (Section 6.2). Below, we briefly introduce the added baseline methods, all of which have used the same trained classifier adopted in CTRL for guidance:

1. Sequential Monte Carlo (SMC)-based method: Recent works [1,2] leverage resampling techniques in SMC to approximate distributions in diffusion models across a batch of samples (i.e., particle filtering). This method is training-free.

2. SVDD [3]: A decoding-based method that iteratively selects preferable samples during each diffusion step based on reward signals. For our conditioning tasks, the compressibility classifier and aesthetic score classifier provide such rewards. This method is also training-free and operates at inference time.

3. MPGD [4]: This method refines the predicted clean data $x_0$​ during inference based on a manifold hypothesis. While such refinement incurs additional computational costs during inference, it does not fine-tune diffusion model weights, remaining a training-free approach.

We report the results below, showing that CTRL achieves good performance, outperforming all baselines in both tasks. While SVDD and MPGD surpass SMC, they still fall short of CTRL. These results highlight CTRL’s superiority in achieving accurate and controllable generation, even compared to the most advanced conditioning methods.

• 4-class compressibility task (Section 6.1)

• multi-task conditional generation (Section 6.2)

[1] Wu, Luhuan, et al. "Practical and asymptotically exact conditional sampling in diffusion models." Advances in Neural Information Processing Systems 36 (2024).

[2] Phillips, Angus, et al. "Particle Denoising Diffusion Sampler." arXiv preprint arXiv:2402.06320 (2024).

[3] Li, Xiner, et al. "Derivative-free guidance in continuous and discrete diffusion models with soft value-based decoding." arXiv preprint arXiv:2408.08252 (2024).

[4] He, Yutong, et al. "Manifold preserving guided diffusion." arXiv preprint arXiv:2311.16424 (2023).

Thank you for the thoughtful and positive review! We’re glad you found our approach novel and appreciated our work’s contributions.

In this response, we have: (a) agreed that further exploration of convergence properties and theoretical guarantees could strengthen the framework, which serves as a valuable direction for future work; (b) added more recent conditional methods as baselines according to your suggestion. We have updated a new version of the manuscript with the changes highlighted in red.

Q. Could these three types of errors be more precisely quantified?

In Section 4.3, we outline three potential sources of error that CTRL may encounter, primarily to set the stage for comparisons with existing methods in Section 5. Specifically, our goal is to demonstrate that CTRL achieves greater sample efficiency than classifier-free guidance in certain scenarios. We acknowledge the importance of providing more precise quantification of these errors and will pursue this as a focus in future research. Thank you again for your constructive insights!

Q. Including more recent conditional methods as baselines.

Thank you for raising this! We have included three additional baselines for both (1) the 4-class compressibility task (Section 6.1) and (2) the multi-task conditional generation (Section 6.2). Below, we briefly introduce the added baseline methods, all of which have used the same trained classifier adopted in CTRL for guidance:

1. Sequential Monte Carlo (SMC)-based method: Recent works [1,2] leverage resampling techniques in SMC to approximate distributions in diffusion models across a batch of samples (i.e., particle filtering). This method is training-free.

2. SVDD [3]: A decoding-based method that iteratively selects preferable samples during each diffusion step based on reward signals. For our conditioning tasks, the compressibility classifier and aesthetic score classifier serve as the reward functions. This method is also training-free and operates at inference time.

3. MPGD [4]: This method refines the predicted clean data $x_0$​ during inference based on a manifold hypothesis. While such refinement incurs additional computational costs during inference, it does not fine-tune diffusion model weights, remaining a training-free approach.

We report the results below, showing that CTRL achieves good performance, outperforming all baselines in both tasks. While SVDD and MPGD surpass SMC, they still fall short of CTRL. These results highlight CTRL’s superiority in achieving accurate and controllable generation, even compared to the most advanced conditioning methods.

• 4-class compressibility task (Section 6.1)

• multi-task conditional generation (Section 6.2)

[1] Wu, Luhuan, et al. "Practical and asymptotically exact conditional sampling in diffusion models." Advances in Neural Information Processing Systems 36 (2024).

[2] Phillips, Angus, et al. "Particle Denoising Diffusion Sampler." arXiv preprint arXiv:2402.06320 (2024).

[3] Li, Xiner, et al. "Derivative-free guidance in continuous and discrete diffusion models with soft value-based decoding." arXiv preprint arXiv:2408.08252 (2024).

[4] He, Yutong, et al. "Manifold preserving guided diffusion." arXiv preprint arXiv:2311.16424 (2023).

Reviewer acmr: if possible, can you reply to the response?

We are thankful for your appreciation of our work’s presentation and theoretical contribution. In this response, we have addressed your concerns by providing explanations on: (a) the implementation with Stable Diffusion, (b) the complexity of the training process, and (c) potential deviations from the pre-trained models.

We are happy to answer more if further clarification is required.

Q. Backpropagating the gradients demands high GPU memory.

Thank you for pointing this out. We utilized standard memory-efficient fine-tuning techniques for diffusion models, as described in [1] and [2]. Hence, memory efficiency is not a substantial practical bottleneck in our experiments. While the main paper does not go into these details due to space limitations, we extensively discuss the implementation details and computational costs in Appendix C and Appendix H.

As outlined in Appendix C and following [1, 2], we employed two strategies to reduce the memory footprint of direct reward backpropagation:

1. Fine-tuning only the LoRA modules rather than the original diffusion model weights.
2. Utilizing gradient checkpointing to compute derivatives on the fly.

These techniques enable backpropagation through all 50 diffusion steps of Stable Diffusion (using a DDIM scheduler) while adhering to hardware constraints. Specifically, these strategies make it possible to fine-tune Stable Diffusion with a batch size of 4, requiring approximately 30 GB of GPU memory. Additionally, we use gradient accumulation and multiple GPUs to scale the effective batch size to 128, enabling smoother optimization.

In Appendix H.1 and H.2, we discussed further strategies to reduce computational cost and memory requirements.

[1] Clark, Kevin, et al. "Directly fine-tuning diffusion models on differentiable rewards." ICLR 2024 [2] Prabhudesai, Mihir, et al. "Aligning text-to-image diffusion models with reward backpropagation." ICLR 2024

Q: CTRL needs to train both classifier and diffusion models. This complicates the overall procedure.

Thank you for noticing this point. We acknowledge that our method requires training both a classifier and fine-tuning the diffusion model. However, we would like to highlight the unique advantages of CTRL:

1. Inference Efficiency: While classifier-based guidance avoids diffusion model fine-tuning, incorporating classifier gradients into the diffusion process significantly slows down inference. In contrast, CTRL retains the inference speed of pre-trained diffusion models, making the upfront fine-tuning costs worthwhile for achieving guidance-directed sampling with faster inference compared to classifier-based methods.
2. Sample Efficiency: As outlined in Table 1, classifier-free methods generally rely on large volumes of offline data for effective training. Although CTRL requires training a separate classifier, which classifier-free methods do not, this training process is more sample-efficient than diffusion model fine-tuning. Consequently, CTRL’s two-step approach leverages limited offline data more effectively than classifier-free methods.

Finally, from a performance perspective, as validated in Sections 6.1 and 6.2, CTRL achieves more accurate samples than both baselines. In light of these considerations, we believe the two-step nature of CTRL does not complicate training but instead balances efficiency and performance.

Q: Following reinforcement learning, the forward process might no longer correspond to the original stochastic differential equation (SDE).

If we understand correctly, the reviewer is concerned about deviation from the pre-trained models induced by fine-tuning. We are happy to clarify more if the concern is different.

Please note that fine-tuning, as applied in CTRL, naturally introduces deviations from the pre-trained model and its associated SDE, which is common in diffusion model alignment. Such deviation is inevitable as there is a mismatch between the pre-training data distribution and the target controllable distribution. However, if the deviation is significant, it may result in reward over-optimization.

To address this, we follow [3] to enforce an explicit KL constraint to control such deviations. This approach ensures a balance between aligning with the pre-trained model and optimizing task-specific performance, making any deviations purposeful and beneficial rather than detrimental.

As a general point, we note that in the context of diffusion model alignment and fine-tuning, many works have explored mitigation strategies for reward over-optimization [3,4,5].

[3] https://arxiv.org/abs/2402.15194

[4] https://arxiv.org/abs/2402.08552

[5] https://arxiv.org/abs/2405.18881

Reviewer FPqT: if possible, can you reply to the authors' response?

We are deeply thankful for your recognition of our work’s significance and performance. In this response, we have provided additional evaluation metrics to illustrate the image quality generated by CTRL. We also provide a rationale for our experimental setup. Please note that we have updated a new version of the manuscript to incorporate changes, which are highlighted in red.

Q. Evaluation metrics are insufficient; for instance, the author only verifies the controllability of CTRL without providing results on the image quality generated by CTRL.

Thank you for pointing this out! To address this, we provide additional evaluation metrics for CTRL and baseline methods, specifically BRISQUE [1] (lower values indicate better image quality) and CLIP score [2] (higher values reflect better text-image alignment).

As shown, CTRL achieves better image quality than Classifier-Free Guidance and Reconstruction Guidance while achieving the highest alignment score.

[1] Mittal, A., Moorthy, A. K., & Bovik, A. C. (2012). No-reference image quality assessment in the spatial domain. IEEE Transactions on image processing, 21(12), 4695-4708.

[2] https://lightning.ai/docs/torchmetrics/stable/multimodal/clip_score.html

Q. The author only uses compressibility and aesthetic pleasingness as conditional controls.

We appreciate your suggestion to explore more complex conditions, such as those used in ControlNet. While these are indeed compelling directions, these works focus on introducing new neural network architectures with an existing guidance method.

However, the focus of this work is developing a new principled approach to add conditional guidance with strong theoretical guarantees. Therefore, our current experimental tasks—compressibility and aesthetic pleasingness—were carefully chosen to validate the fundamental aspects of CTRL. Specifically, our contributions are as follows:

1. Framing Conditional Generation as an RL Problem: We first frame conditional generation as an RL problem, providing a novel perspective. As noted by Reviewers ahRC and acmr, this framing is innovative and represents a significant advancement in guided/conditional generation. We also formally prove that executing the soft-optimal policy in this RL would enable sampling from the target conditional distribution during inference. This finding provides the solid grounding of CTRL, as pointed out by the Reviewer FPqT.

2. Introducing the CTRL Algorithm: We propose the CTRL algorithm, which enjoys certain benefits over existing guidance methods. Specifically, CTRL eliminates the need to learn classifiers at multiple noise scales, as required in standard classifier guidance, and avoids fundamental approximations made by certain variants to bypass these challenges. In Section 5.2, we also highlighted CTRL’s distinct advantage in leveraging conditional independence, which leads to two benefits: (a) more sample-efficient training compared to classifier-free guidance and (b) could be applied for multi-task conditional generation task, which we argue is new in the literature as far as we are concerned.

3. Connecting CTRL to existing guidance methods theoretically: We establish a close theoretical and practical connection between CTRL, classifier guidance, and classifier-free guidance, as pointed out by the Reviewer FPqT and acmr.

Finally, we would like to note that our tasks with compressibility and aesthetic scores are inherently challenging. For instance, generating aesthetically pleasing images often conflicts with achieving high compressibility (Fig. 2(a)), making our second experiment—synthesizing images conditioned on these two scores—particularly demanding.

We observe that CTRL effectively handles all 2 $\times$ 2 conditions (Fig. 2(c)) and outperforms baselines. This directly validates CTRL’s advantage of leveraging conditional independence, as outlined earlier.

In summary, we think our experiments are sufficient to showcase the principled contributions of CTRL and its benefits over existing conditional methods.