Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence

Thank you for your feedback. We are glad that you found our theoretical results strong and rigorous. Below please find our response to your questions.

I have to say that the theory in this paper is very solid. However, the lack of empirical results is also the biggest drawback of this paper.

Thank you for your suggestion. To demonstrate the practical efficiency and effectiveness of the proposed PI-FT method, we conduct thorough numerical experiments, which are included in https://anonymous.4open.science/r/ICML-rebuttal-741C/figures.pdf.

Experiment Set-up. We fine-tune Stable Diffusion v1.5 for text-to-image generation using LoRA and ImageReward. Our prompts—“A green-colored rabbit,” “A cat and a dog,” “Four wolves in the park,” and “A dog on the moon”—assess color, composition, counting, and location. During fine-tuning, we generate 10 trajectories (50 transitions each) to calculate the gradient with 1K gradient steps (AdamW, LR = 3e-4), and KL regularization ($\beta= 0.01$).

Evaluation. We first compare ImageReward scores of images generated by the pre-trained model, DPOK and PI-FT. To ensure a fair evaluation, we perform 10 gradient steps with LR = 1e-5 per sampling step in DPOK. Each gradient step is computed using 50 randomly sampled transitions from the replay buffer. Consequently, 1K sampling steps and 10K gradient steps in DPOK result in a computational cost similar to that of PI-FT. As shown in Figure 1 in the link, PI-FT consistently outperforms both the pre-trained model and DPOK across all four prompts. Figure 2 further shows that PI-FT more accurately generates the correct number of wolves, and correctly places the dog on the moon. Additionally, PI-FT avoids mistakes such as depicting a rabbit with a green background as appeared in the pre-trained model and DPOK. The texture of the generated animals is also more natural-looking under PI-FT compared to other baselines.

Effect of KL regularization. KL regularization is known to improve fine-tuning. We analyze its effect in the PI-FT method using the prompt “Four wolves in the park,” varying $\beta \in$ {0.01, 0.1, 1.0}. Figure 3 shows that the gradient norm decreases to 0 in all three settings, illustrating the convergence of our algorithm. In Figure 4, we see that the ImageReward score increases and eventually stabilizes when $\beta$ is small. In contrast, the score exhibits limited improvement with large $\beta$. This observation is consistent with Figure 5, where the KL divergence remains large throughout training for $\beta = 0.01$ and stays significantly at a smaller level for $\beta \in$ {0.1, 1.0}. Figure 6 illustrates that smaller $\beta$ produces images with nearly four wolves, while larger $\beta$ yields fewer. These results underscore the importance of the KL coefficient in our framework.

… Could authors show more details about the comparison and similarity of their methods and your method.

In the literature, DPOK and DDPO use generic RL algorithms such as PPO or REINFORCE to optimize the policy/score. Despite showing empirical promise on certain tasks, existing methods overlook the underlying structure of diffusion models, leaving significant room for improvement in efficiency and lacking any theoretical justification. In contrast, by fully utilizing the problem structure, we formulate fine-tuning as a KL-regularized stochastic control problem (with linear dynamics), and propose the PI-FT algorithm. This structure allows us to derive regularity and convergence guarantees, closing the gap left by prior works. We will revise lines 90-96 to clarify these points.

... show the convergence and the regularity of the optimal value function & the optimal control policy of your approach based on above mentioned empirical settings.

In the newly added experiments, we do observe convergence of both the control policy and the value function during training, supporting our theoretical claims.

… I am curious could the conditions in above lemmas and theorems be realized in empirical settings…

Empirically, we observe numerical results consistent with our theoretical claims. In Figure 3 (in the link), we observe that the algorithm converges linearly and the gradient norm of the U-Net stabilizes after approximately 200 sampling steps. Additionally, we compare PI-FT and DPOK under the same computational budgets; as shown in Figure 1, PI-FT achieves consistently higher ImageReward scores across all four prompts. This performance suggests that our approach converges more efficiently in practice. Although our theory assumes Lipschitz conditions and a specific range of the regularization coefficient, our experiments demonstrate that convergence remains robust even for small $\beta = 0.01$ and relatively large learning rates.

Thank you for your detailed feedback. We hope our response addresses your concerns. If so, we would appreciate it if you could reflect it in your evaluation of our paper.

Thank you for your constructive feedback. We are delighted that you found our theoretical results rigorous and our presentation clear. Below is our point-to-point response to your comments:

Proper evaluation is missing: Algorithm 1 is defined and justified clearly but is purely theoretical; however, the practicality of the method is questionable…

A significant weakness is the absence of any empirical validation…

As mentioned, Algorithm 1 appears overly idealized, and its practical utility in realistic scenarios remains highly questionable. In practical high-dimensional environment, taking expectations to obtain accurate value function can have difficulty…

We included a new set of thorough experiments in this rebuttal. Please see https://anonymous.4open.science/r/ICML-rebuttal-741C/figures.pdf for the results and see response to Reviewer 1H7D for detailed explanations on the experiments.

I find following the flow of this work pretty easy and smooth. One suggestion would be: all theoretical results before Thm 2.8 are standard in the literature, that is to say, all contents up until Page 5. Results such as Lemma 2.5 and Eq. 8 (many remarks are also well-known) have long been discovered, see [1], [2]. These standard results take up too much space in the main text. I would suggest improving the presentation, for example, simply referring to these works instead of repeating all the details in the main text. Otherwise, it is also possible to only present the most important parts in the main text but defer burdensome calculation details and remarks to the appendix in the next version.

Thank you for your constructive feedback. To the best of our knowledge, we are the first to use KL divergence over transition dynamics (on path space) to control the deviation of the fine-tuned model from the pre-trained model, whereas formulations in the literature consider KL between the terminal state distributions. For this reason, the results before Theorem 2.8 are novel and valuable to readers although not technically significant. Therefore, we decide to follow your suggestion and defer calculation details and remarks to the appendix in the revised manuscript.

…since this work is purely theoretical and does not have any kind of simulation studies. It is advised to spend more efforts in demonstrating the technical hardness of Thm 2.8 and Thm 3.1. It is confusing to me what the major difficulty is in presenting these results. It feels like Thm 2.8 is a direct calculation based on the Lipschitzness assumptions, which might not present enough technical novelty. I might be wrong, but can the authors comment on this?

We appreciate the opportunity to clarify the technical novelty of Theorem 2.8 and Theorem 3.1.

Regarding Theorem 2.8, the core challenge stems from the fact that Eq. (17) is an implicit equation, in which both sides involve the optimal control $u_t^*$. This prevents a direct application of Lipschitz assumptions on model parameters. In contrast, we prove the Lipschitz property of $u_t^*$ through the detailed and nontrivial derivations in Line 726 - 736, where the choice of $\beta_t$ ensures the invertibility of the coefficient. Additionally, the differentiability of $u_t^*$ relies on the Gaussian smoothing effect; see Eq. (30) – (32), which is novel analysis to the RL literature. The Lipschitz condition of $\nabla u_t^*$ is especially non-trivial as it involves $\nabla^2 V_{t+1}^*$, which may not be Lipschitz. We address this using the integration by parts formula in Eq. (34), which is a key technical step. We will discuss these in the manuscript.

In terms of Theorem 3.1, the convergence analysis of the PI-FT algorithm relies on maintaining the regularity of the value function throughout training, which is highly nontrivial. For example, prior work such as [3] proves convergence under the assumption that the value function remains regular during training (which is non-tractable); see Assumption 2 in [3]. In contrast, our result establishes the desirable regularity through a connection between $V_t^{(m)}$ and the optimal value function $V_t^*$.

Thanks to the reviewer’s question, we will revise the manuscript to better emphasize the technical challenges and contributions underlying these results.

Thank you for your time and for your valuable feedback. Your suggestions definitely helped us improve the quality of this manuscript and highlight our contributions. We hope our response addresses your questions and clarifies the challenges in our work. We also hope that the effort we invested in conducting thorough numerical experiments demonstrates the promising practical performance of our proposed algorithm. If so, we would be grateful if you could reflect it in your evaluation of our paper.

[3] Zhou, Mo, and Jianfeng Lu. "A policy gradient framework for stochastic optimal control problems with global convergence guarantee." arXiv preprint arXiv:2302.05816 (2023).

Thank you for your detailed feedback. We are glad to receive your positive feedback on our theoretical contributions and that you found our placement of the work clear and honest. Below please find our response to your questions.

It is not clear to me if the policy iteration algorithm presented is supposed to work only on discrete-spaces, or, more precisely, how the control $u$ is represented within that algorithm. What settings does this algorithm actually cover? I am asking this since Sec. 4 seems to extend the control representation to be parametrized, but in doing so it seems to propose another algorithm. Did I misinterpret something?

While we consider a discrete-time model, the PI-FT algorithm is formulated on continuous state and action spaces, with both state and control taking values in $\mathbb{R}^d$. This setting captures the essence of fine-tuning and indeed disparts us from the usual RL literature. Specifically, we focus on Markovian control policy $u_t: \mathbb{R}^d \to \mathbb{R}^d$, and our algorithm iteratively updates $u_t(y)$ for each $y \in \mathbb{R}^d$. Section 4 introduced a linear parametrization of the control policy to enable an efficient implementation in practice. This parameterized setting should be viewed as a special version of the PI-FT algorithm, rather than a separate algorithm. Moreover, we believe the theoretical results regarding regularity and convergence established in Section 3 apply to this parameterized setting and a roadmap has been provided in Section 4. Please let us know if this explanation clarifies your concerns.

Sec. 4 seems to propose an algorithm for continuous spaces where the control is represented via a linear parametrization. While a linear approximation would make sense to represent the forward process drift, this seems an excessive approximation for the reverse process control, which should have an order of complexity comparable to a score network. It seems to me that this object would be non-linear even in the easiest examples possible. Hence the question: are there some relevant examples (even simple) where this assumption holds?

We refer to the parameterization in Section 4 as linear in the sense that the control policy is a linear combination, through parameter $K_t$, of (nonlinear) basis functions. Hence the resulting control can be possibly nonlinear in the state variable $y$. This is a flexible, powerful, and tractable parameterization often used in control and RL, where depending on the basis, it can have high expressivity. In particular, our parameterization includes a wide class of score approximators such as random features, kernel methods and even overparameterized neural networks under the NTK regime. Hence, despite being linear in parameter $K_t$, the policy class can be sufficiently rich to approximate highly nonlinear score functions. Moreover, our experimental results (which we discuss below) suggest our principled framework remains valid for generic neural networks. We will clarify these points in our revision.

At times it is very unclear whether the algorithms are introduced within this work (e.g., the policy gradient one), or if this work aims to only perform theoretical analysis on them. Please clarify.

Apologies for the confusion and thank you for raising this point. To clarify, both the policy iteration method (PI-FT) and its policy gradient variant are new and proposed in this work. Existing methods such as DPOK and DDPO rely on generic RL developments such as PPO or REINFORCE and fail to fully leverage the fine-tuning structure. On the contrary, the specific setting considered, with linear dynamics and entropy-regularization, is tailored towards the development of efficient fine-tuning diffusion models. For this reason, the PI-FT algorithm and its parametric extension directly computes the policy gradient of a KL-regularized control objective. This principled design leads to a more efficient implementation in practice compared to prior works; see our experiments below. We will clarify all these points in our final revision. In particular, we will specifically highlight that both algorithms are proposed in this work and are our contribution.

Not presenting experimental evaluations renders impossible to evaluate the practical relevance of proposed algorithmic ideas.

We included a new set of thorough experiments in this rebuttal. Please see link https://anonymous.4open.science/r/ICML-rebuttal-741C/figures.pdf for the results and see response to Reviewer 1H7D for detailed explanations on the experiments.

Line 434, 'replies' should be 'relies', I guess.
Line 397, 'the' is repeated twice.

Thank you for pointing out these typos. We will fix them in the revised manuscript.

Finally, we would like to thank you for your constructive feedback. We hope our response answers your questions. If so, we would greatly appreciate it if you could reflect it in your evaluation of our paper.

 "This parameterized setting should be viewed as a special version of the PI-FT, rather than a separate algorithm"? … this setting is formally tackled with another algorithm in Sec. 4 (namely the PG scheme). 

…a linear dynamics for the sake of analysis,...diffusion process acts on a learned latent space…a linear or kernelized approximation…a fair abstraction…

… experimental comparison…e.g., [1,2] among others. Hence the question: why DPOK, which does not seem to rely on the classic duality for KL-regularized control/RL…