Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control

We thank the reviewer for their feedback. We have addressed your concerns by providing explanations on (1) the implementation with Stable Diffusion, (2) additional discussion of related work, and (3) the choice of $\alpha$. We are happy to answer more if further clarification is needed.

Q. Algorithm 1 presents a method with three stages: value function learning and solving two neural SDEs. This seems computational heavy, especially the Neural SDE solving should be hard due to the storation of computation graph when the SDE is parameterized by huge neural networks. It seems impossible for stable diffusion as usually it is only possible to forward pass with batch size = 4 on 80G GPU, let alone the computation graph along a trajectory. I wonder if more clarification can be made about the computation implementation and speed.

A. 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 have provided an extensive discussion of the implementation details and computational costs in Appendix F.2.1 and Appendix G. Below is a summary:

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

• Fine-tuning only the LoRA modules rather than the original diffusion model weights.
• 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) within hardware constraints. Specifically, leveraging these methods allows fine-tuning Stable Diffusion with a batch size of 4 to consume approximately 30 GB of GPU memory.
• Additionally, we use gradient accumulation and multiple GPUs to scale the effective batch size to 128, facilitating smoother optimization.

2. The two neural SDEs are optimized separately and sequentially. Once the initial distribution is obtained, the optimization of the second neural SDE does not require storing computational graphs for parameters from the first SDE.

[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. The pretrained diffusion model is not continuous time. How to adapt it to the framework proposed in this paper?

A. Regarding the distinction between continuous-time and discrete-time training in diffusion models, we interpret the reviewer’s comment as referring to the former as training with randomly sampled time steps and the latter as training with fixed time steps. When pre-trained models are trained in a discrete-time training way, in solving the Neural SDE within our approach, the most natural strategy is to adopt the same discretization scheme as used in the pre-trained model, which is precisely the approach taken in our algorithm.

Q. There are several other works that should be compared or mentioned

A. Thank you for your suggestion. We now acknowledge that all of the referenced works are related to ours. We have incorporated it into a new version of our draft.

[1] and [2] have addressed problems similar to ours, but their approaches still differ significantly. [1] addressed the initial bias issue discussed in Section 6 by modifying the noise schedule, whereas we tackled it by introducing an additional optimization problem. [2] approached a related problem by designing an objective function inspired by a detailed balance loss in Gflownets, while our algorithm directly solves the control problem using neural SDE. While their algorithm has certain benefits when rewards are differentiable like PPO, our primary focus is more on how to mitigate overoptimization from both theoretical and empirical perspectives. [3] has explored fine-tuning in diffusion models, framing it as a bi-level optimization problem. However, they do not appear to discuss strategies for constructing objective functions, such as incorporating KL regularization to prevent over-optimization.

We sincerely appreciate your detailed and thoughtful feedback. In this response, we have addressed your primary concerns regarding our algorithms. We are happy to answer more if further clarification is required. Regarding your suggestions for the presentation, we have incorporated many of them, and the change is in red. (For some of the suggestions, like replacing sections in the main text and appendix, we plan to do; but at this moment, to avoid confusion for other reviewers, we didn't implement it)

Weakness: My main concern with the paper is the following: The terminal reward $v_0$ for learning the drift $q$ is approximated with a neural network using a regression objective. Hence, the learned model could be arbitrarily bad in regions with no data which is also the cause for the overoptimization problem. It thus seems like the problem has just shifted to the initial value function estimation and thus not properly addressed.

While we acknowledge your observation, we think it won't be a primary concern for the following reasons. In summary, we do not encounter the distribution shift issue in this context because error guarantees are required solely within the original data distribution. Specifically, the reasoning is as follows:

First, recall from Corollary 2, the optimal initial distribution is

$$\exp(v^{ \star}\_0( \cdot)/ \alpha) \nu\_{ini}(\cdot)/C.$$

Then, leveraging the relation from Lemma 1:

$

 \exp\left(\frac{v^{\star}\_{0}(x)}{\alpha} \right)= \mathrm{E}\_{\mathbb{P}^{data}}\left [\exp \left (\frac{r(x_T)}{\alpha } \right)|x_0=x \right], 

$

our algorithm approximates the value function $v^{\star}$ in the optimal initial distribution with supervised learning:

$

 \hat g = \mathrm{argmin}\_{g \in \mathcal{G}} \sum\_{i=1}^n  \left ( \exp   \left (\frac{r(x^i_T)}{\alpha } \right )- g(x^i_0) \right )^2,  x^i_0 \sim \nu\_{ini}, x^i_T\sim \mathbb{P}^{data}(x_0),  

$

where $\mathcal{G}$ denotes a function class, such as neural networks. Under the well-specification assumption (a standard assumption as used in Wainright 2019), we achieve the following mean squared error (MSE) guarantee:

$

    \mathrm{E}\_{x_0\sim \nu\_{ini}}\left [ \left (\hat g(x_0)- g^{\star}(x_0) \right)^2 \right ] = O\left (\frac{Cap(\mathcal{G})}{n} \right)

$

where $Cap(\mathcal{G})$ represents the capacity of the function class $\mathcal{G}$ and $g^{\star}$ is $\exp(v_0(\cdot)/\alpha)$.

This result implies that we obtain the MSE guarantee within the original initial distribution $\nu\_{ini}$. Although this guarantee does not extend beyond the support of $\nu\_{ini}$, it is not an issue because we know that the support of the optimal initial distribution is restricted to $\nu\_{ini}$ (recalling Corollary 2), and in the second step of approximating the optimal initial distribution, in our algorithm, we consistently do sampling on the support of the original initial distribution $\nu\_{ini}$.

Q. While inference time is already a bottleneck for diffusion models, the proposed approach requires simulating two SDEs.The approach requires solving two stochastic optimal control problems and an additional regression problem for obtaining the initial value function.

While we ackowldge your opinion, we have an explanation that this won't be a concern in Appendix G.1. In particualr, the following paragraph demonstrates that learning the optimal initial distribution is not a practical bottleneck for our algorithm.

Computational Cost: Cost of Learning Initial Distributions: Our overhead in learning the initial distribution is low. When learning a second diffusion chain, our goal is to learn $\exp(v^{\star}_0(x))\nu\_{\mathrm{ini}}(x)$. This distribution is much simpler and smoother compared to the target distribution $\exp(r(x)/\alpha)p\_{\mathrm{data}}(x)$ in the first diffusion chain. Therefore, we require fewer epochs to learn this distribution. For instance, in image generation, we use only $200$ reward queries to learn initial distributions while the main diffusion chain takes $12000$ queries for finetuning. The wall time of learning the initial distribution (i.e., the second diffusion chain) is $30-40$ minutes in image experiments, while the wall time of learning the second chain is roughly $1800$ minutes.

Weakness. The method requires a) gradient checkpointing and b) low-rank modules to train

This technique is commonly employed when fine-tuning large models. Representative works for fine-tuning diffusion models (Black et al., 2023; Prabhudesai et al., 2023; Clark et al., 2023) also utilize these techniques, indicating that this is not a unique aspect of our work. Hence, we believe this is not a particular weakness of our methods. Additionally, we clarify that for smaller models, such as those used in experiments with biological diffusion models, employing this technique is not necessary.

Thank you for your response and engagement. Below, we provide additional clarification regarding your concerns. We are happy to clarify more if specific points remain unclear.

1. First, we would like to emphasize the assumptions outlined in Section 6:

We start with a reference SDE over the interval $t \in [-T,0]; t \in [-T,0]; d x_t = \tilde \sigma(t)dw_t, x_{-T}=x_{fix},$ such that the distribution at time $0$ follows $\nu_{ini}$.

2. Under the above assumption, we have concluded in Theorem 2 that the generated distribution by

$

t \in[-T,0]; d x_t = q^{\star}(t, x_t)dt + \tilde \sigma(t) dw_t

$

is $\hat a(x_0)\nu_{ini}(x_0)/C,$ where $q^{\star}$ is defined by the following control problem

$

q^{\star}=\arg \max_{q}\mathrm{E}_{P^{q}} \left[\hat a(x_0)-0.5 \alpha \int^0_{-T} \frac{|q(t,x_t)|^2}{\tilde \sigma^2(t) }dt \right], \quad (a)

$

Our control problem (a) incorporates the KL divergence as the second term (the integral term), which ensures that the optimization remains within the support of $\nu_{ini}$. To provide a high-level explanation of why the support is preserved, , consider a simplified problem below.

• For an optimization problem where $q$ is any distribution over $X$, $\nu_{ini}$ is a fixed distribution, and $\hat a(\cdot)$ represents a value function:

$

q^{\star} = \arg\max_{q \in \Delta(X)} E_{x \sim q}[\hat a(x)] - \mathrm{KL}(q|| \nu_{ini})

$

In this problem, we can easily show the optimal distribution $\forall x \in X, q^{\star}(x) = \hat a(x)\nu_{ini}(x)/C$ [1]. Thus, $q^*$ will retain the support of $\nu_{init}$ over $X$.

Returning to our control problem (a), the second term (i.e., the integral term) corresponds to the KL divergence between the path measures on $x_{0:T}$ induced by two SDEs—one that needs to be learned and another corresponding to $\nu_{ini}$ (formally shown using Girsanov theorem [2] in our proof). Consequently, this term similarly ensures that the distribution induced by the fine-tuned SDE remains within the distribution of $\nu_{ini}$.

• At this stage, note this theorem is unrelated to the function approximation of value functions ($\hat a(x_0)$)

3. Now, let me connect the above discussion with the function approximation error for $\hat a$. As discussed in our first rebuttal, by invoking the standard guarantees of regression, we obtain the mean squared error (MSE) guarantee for $\hat a(x_0)$ under the initial distribution:

$

\mathrm{E}\_{x_0 \sim \nu_{ini}}[ ( \hat a(x_0)- a(x_0) )^2]. 

$

Thus, there is no distributional shift, as the MSE guarantee for $\hat a$ is explicitly on the support of $\nu_{ini}$ and we are sampling from $\hat a(x_0)\nu_{ini}(x_0)$, as previously stated in the second step.

[1] Ziegler, Daniel M., et al. "Fine-tuning language models from human preferences." arXiv preprint arXiv:1909.08593 (2019).

[2] Shreve, Steven E. Stochastic calculus for finance II: Continuous-time models. Vol. 11. New York: springer, 2004.

We appreciate your positive feedback. We have addressed your concerns by providing (1) an additional explanation regarding the significance of learning initial distributions, (2) a clarification of the choice of $\alpha$, and (3) an explanation of how the generated images mitigate overoptimization.

Q. The motivation for learning a different initial distribution is unclear. Why does the initial distribution affect the finetuning of diffusion models, and why should we finetune the model on a different initial distribution? On the other hand, does it affect the model performance on general tasks? In other words, given the model finetuned with the learned initial distribution, if users sample the generation from the original Gaussian distribution in the sampling phase, how about the quality of generations?

Thank you for raising this point. Intuitively, in many diffusion models, the initial point could retain significant information about the entire trajectory, making its optimization beneficial during fine-tuning. As an extreme case, in specific diffusion models like rectified flow [1], the initial points contain all trajectory information, as trajectories do not overlap. Here is a more concrete discussion. We have added some of them in the main text.

• Empirical viewpoint: In Section 8.4, we compare the performance of our method with that of original initial distributions (Gaussian distributions). Our results demonstrate that incorporating initial distribution learning enhances performance. Notably, two independent studies [2,3] also validate the effectiveness of learning initial distributions (but in different ways from ours).

• Theoretical viewpoint: Theorem 1 demonstrates that learning an initial distribution enables sampling from the target distribution $p(x)\exp(r(x))$. However, we do not intend to cliam that good practical performance is unattainable without learning an initial distribution. Empirically, we observe that learning an initial distribution enhances performance, but certain good performance can still be achieved without it, as evidenced in Section 8.4.

[1] Liu, Xingchao, Chengyue Gong, and Qiang Liu. "Flow straight and fast: Learning to generate and transfer data with rectified flow." arXiv preprint arXiv:2209.03003 (2022).

[2] Domingo-Enrich, Carles, et al. "Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control." arXiv preprint arXiv:2409.08861 (2024).

[3] Ben-Hamu, Heli, et al. "D-Flow: Differentiating through Flows for Controlled Generation." arXiv preprint arXiv:2402.14017 (2024). Equations in Lines 372-377 are slightly confusing.

Q. Equations in Lines 372-377 are slightly confusing.

A. Thank you for catching it. We have provided additional explanations, marked in red.

Q. The connection between the added drift term and the classifier guidance is interesting. Given that they have a similar mathematical formulation, I expect a theoretical analysis of their difference in performance.

A. Thank you for your encouraging feedback. In general, classifier guidance requires differentiable value function models, which necessitates realizability within the value function class. Conversely, our fine-tuning algorithm relies on realizability within the policy class. We are leaving a more detailed analysis as future work.

Q. How was α in PPO+KL set in experiments of Figure 3?

A. Thank you for pointing this out. We experimented with several values of $\alpha$ (like $\alpha=1,0.1, 0.01$). For higher value of $\alpha$, we observed that the PPO learning curve slows significantly, making it difficult to optimize rewards. Consequently, we have included only the best result. We will provide a more detailed clarification in the next version.

We appreciate the reviewer’s positive feedback.

Q. In the write-up, the authors state that the methods is an entropy-regularized control against the pretrained diffusion model but in the equations, the data distribution is used instead of the pre-trained distribution. Could you clarify?

Your understanding is correct. We have considered a scenario where the pre-trained distribution closely approximates the data distribution, assuming the availability of a foundational pre-trained generative model trained on a large dataset of natural samples. Meanwhile, the data with feedback remains significantly more limited. This setting is practical as we mention in the introduction. We have clarified this assumption more explicitly in the updated version of the manuscript.

Approximation errors when solving of E9 and E10. Can [1] be used instead?

Thanks for the suggestion. This approach could potentially be applied. However, in practice, we observe that explicitly approximating Q-functions (e.g., critic models) is challenging in the context of fine-tuning diffusion models due to the high dimensionality of the state space. Therefore, we employ purely policy-based optimization methods to solve our optimization problem. While incorporating approaches such as SAC, SQL, or S^2AC in [1] is indeed an intriguing direction, we leave this exploration for future work.