Score as Action: Fine Tuning Diffusion Generative Models by Continuous-time Reinforcement Learning

We sincerely thank you for providing thoughtful suggestions to improve our paper. Please find our responses to your questions below:

Q1: intuitive explanations for the value function design in Equation (20)

A1: The design of our value function in Equation (20) is motivated to capture the structural property of score-based diffusion models (we explain in (a) below) and to meet the boundary condition of the value function (we explain in (b) below).

• (a) Recall that the value function $V(t,x)$ is defined as the expected reward of final generated image (if regularization penalty parameter $\beta$ is 0) conditional on the current time $r$ and state/latent $x$, so an natural and naive approximation of it can be the reward of the denoised image $\hat{x}_\theta(t,x)$ of the current state $x$, as the denoised image can be understood as "a prediction of the image at the end of SDE trajectory". This illustrate our first term reward mean predictor; of course such approximation to the value function is not accurate, so we use the second term with the same network architecture to further approximate the possible residual term which can improve the performance by reducing the mean square error.

• (b) The value function has a boundary condition that at the final time $T$, $V(T,x)=R(x)$, so we choose two proper coefficient functions $c_{\text{skip}}(t)$ and $c_{\text{out}}(t)$ to satisfy such needs: we let $c_{\text{skip}}(T)=1$ and $c_{\text{out}}(T)=0$, thus the boundary condition is naturally satisfied.

We hope that this clarifies our design space and we are happy to take further questions.

Q2: How does the closed-form advantage-rate function to generate images with higher quality?

A2: We illustrate this mainly through the RL prospective as images with higher quality can be reflected as higher average reward for RL. As far as we know, the most relevant work about the conventional continuous-time PPO approaches is in [1], which estimates the advantage rate function through a model-free manner, but requires sampling many more examples to make the estimation accurate and algorithms stable. In our method, thanks to the closed-form advantage-rate function and further subtraction of second order term (which is hard to compute) as the baseline function (please find more detailed discussion about this in Line 309-322 after Theorem 1), we get a much simpler way to compute the advantage function which is also much more sample efficient, since we do not need to query the generative model for the next state to go. We believe the leverage of such structural property is one of the key reasons we get quite stable training reward curves as in illustrate in orange curve of Figure 9.

[1]. Zhao et al. 2023. Policy optimization for continuous reinforcement learning. NeurIPS 2023

We sincerely thank you for providing detailed and thoughtful feedbacks to improve our paper. Please find our responses to your questions below:

Q1: Clarification on the motivation in Figure 1 and Continuous-RL framework

A1: We are sorry if our wording has caused confusion, but we didn't claim on such of "mode collapse". What we claimed or observed is: the discrete RL trained policy, as a diffusion model, "overfits" to using 50 time steps as inference steps, as illustrated in Figure 1. Diffusion models has a nice natural tradeoff of generation quality and speed, as more inference steps could typically lead to more satisfactory images while causing extra time for inference. The most advanced models typically still haves such tradeoffs. We are motivated to let the fine-tuned models not lose such property (which is a dimension of freedom for users of different purposes), as it's pitiful that the performance improvement brought by RL fine-tuning like DDPO when inferencing the fine-tuned models on other time discretization steps like 25 and 100 could be discounted compared to the improvement typically reported only on fixed time steps like 50.

Comparing our framework to [1] and [2], we agree that the continuous-time control formulation already exists and is similar to our learning objective, but 1) we adopt different parameterizations that focus on treating the score function as the first class citizen and as the primary action space, which motivates us to incorporate the structural property in value function design, 2) the methods in these two works are pure stochastic optimal control methods, which are fundamentally different from our RL-based approaches. The effectiveness and scalability of RL methods have been largely investigated in LLM fine tuning which makes them quite popular, and we believe that our work can help improve the study of RL-based fine-tuning methods for diffusion models by the continuous-time framework with better leverage of structural property and thus resulting scalable, more efficient and more stable algorithm.

Q2: The implications and interpretations of Figure 2

A2: To summarize, Figure 2 and Table 1 mainly ablate on how to smartly incorporate the diffusion model policy network in the value function design space and partial effects of the link function. As we already theorectically explained our motivations after Equation (20), we conduct experiments in Figure 2 and Table 1 to empirically showcase how things differ when putting the denoised mean $\hat{x}_{\theta}(t,x)$ in the first reward mean predictor term or the second residual corrector term. Our experiments indeed showcase the promise of incorporating such structural property in our design space, as shown by the reduced MSE of green line compared to the orange line. Incoporating the denoised mean in the second part almost does not improve performance while introducing extra computational burden. The performance of the green line and purple line also showcases that the choice of link function also matters. We agree that there could be other link functions that might outperform sin and cos, however it's impossible for us to exhaust all. We would like to pursue possible theorectical guidance of chosing this as future work.

Q3: Comparisons against common RL-finetuning baselines such as DRaFT and AlignProp.

A3: We would like to remark that DRaFT and AlignProp are algorithms for fine-tuning diffusion models, but they are not RL-finetuning methods, which might partically deviate from our primary study on RL-based methods. Nevertheless, we agree that they are baselines for fine-tuning diffusion models and we conduct additional experiments comparing our method, DDPO, DRaFT and AlignProp. Since DRaFT code is not released, we implemented it by revising the codebase of AlignProp. Please find the reward curves comparisons in the link here. Our experimental results show that DRaFT and AlignProp perform similarly (the time costs are also similar), but they both underperform our CTRL approach, or even DDPO when using ImageReward as both the reward signal and the evaluation metric. This demonstrates the advantage of our proposed method.

Q4: If continuous RL can further boost a reward direct propagation-based method

A4: Thank for raising up a direction that could be insightful. Since DRaFT and AlignProp are not RL-based algorithms, we agree and believe that such methods could have continuous time extensions. This is by high level related to our paper of utilizing continuous-time RL over discrete RL, but we believe that the overall framework will be quite different and will not be RL-related anymore. We would add this point to possible direction which can further showcase the insights brought by our framework; we would like to pursue this as future work.

We sincerely thank you for providing thoughtful suggestions to improve our paper. Please find our responses to your questions below:

Q1: add pseudocode, figure of method

A1: Please find the pseucode of our algorithm here. We will update this part to the revised version.

Q2: provide detaild description and more evidence of "Continuous-time RL outperforms Discrete-time RL baseline methods in both efficiency and stability because CTRL only require time discretization for estimating the policy gradient."

A2: As illustrated in the earlier Section 4, we are esstentially motivated to approximate the policy gradient formula in Equation (23), since PPO can be understood as more efficient implementations of such policy gradient by introducing the surrogate loss that shares the same gradient and clipping mechanism. This is shown in the definition of loss objective in the pseucode.

Also we showed in the experiments of fine-tuning stable diffusion v1.5 (the second environment) as our method yields higher rewards training on the same amount of samples compared to DDPO, and the variance of the reward curves also appear to be much smaller. That's why we claim "efficiency and stability", and we will add this heuristic explanation part in the revised version.

Q3: only compared with 2 models, and add more set of the number of discretization timesteps compared with DPPO.

A3: We conduct additional experiments comparing our method, DDPO, DRaFT [1] and AlignProp [2]. Since DRaFT code is not released, we implemented it by revising the codebase of AlignProp. Please find the reward curves comparisons in the link here. Our experimental results show that DRaFT and AlignProp perform similarly, but they both underperform our CTRL approach, or even DDPO when using ImageReward as both the reward signal and the evaluation metric. Due to time constraints, we haven't tested the more set of the number of discretization timesteps of our methods compared with DPPO since they would require more hyparameter tuning for a fair comparison, but we will also add the corresponding results and comparisons in the later version.

Q4: How long does it take of finetune SD v1.5 with CTRL?

A4: The current version of our implemented CTRL takes twice as much time as DDPO for our current implementation, since it requires the value network update and inference which is similar to PPO vs REINFORCE methods in LLMs.

References:

[1] DRaFT: Directly fine-tuning diffusion models on differentiable rewards

[2] AlignProp: Aligning text-to-image diffusion models with reward backpropagation

We sincerely thank you for providing thoughtful suggestions to improve our paper. Please find our responses to your questions below:

Q: Comparisons to other important baselines, such as RAFT and DRaFT

A: We conduct additional experiments comparing our method, DDPO, DRaFT [1] and AlignProp [2]. Since DRaFT code is not released, we implemented it by revising the codebase of AlignProp. Please find the reward curves comparisons in the link here. Our experimental results show that DRaFT and AlignProp perform similarly, but they underperform our CTRL approach, or even DDPO when using ImageReward as both the reward signal and the evaluation metric.

We would like to comment that we agree these are important diffusion models fine-tuning baseline methods, but our paper primarily focuses on RL-based methods, so we didn't pursue these baseline in our current submitted version. Nevertheless, we agree that these comparisons could strengthen our results, and we will update the experiments of these non-RL baseline methods in the revised version. Due to the time constraints, we haven't tested the RAFT [3] based on ImageReward, but we will also add the corresponding results and comparisons in the later version.

References:

[1] DRaFT: Directly fine-tuning diffusion models on differentiable rewards

[2] AlignProp: Aligning text-to-image diffusion models with reward backpropagation

[3] RAFT: Reward rAnked Fine-Tuning Algorithm