Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control

We thank the reviewer for their comments, and provide answers below.

• “it would strengthen the claim to include experiments validating the numerical effectiveness of the new SOC objective.”

We did run experiments using the toy settings and the models of Domingo-Enrich et al. [3]. Namely, these are experiments where the optimal control is either known analytical or can be computed, which allows for the evaluation of the control $L^2$ error between the learned control and the optimal one. The results we obtained showed that on the simple settings of Domingo-Enrich et al. [3], SOCM (their algorithm) outperforms all others, including Adjoint Matching, which attains a control $L^2$ error on par with Continuous Adjoint. These settings are relatively low-dimensional ($d=10,20$), and have low-magnitude state and terminal cost. The reason that Domingo-Enrich et al. [3], restrict their focus to such problems is that SOCM has an importance weight $\alpha$ that blows up when either the dimension or the magnitudes of the costs are higher; see Fig. 4 of Domingo-Enrich et al., [3] and the explanation below it, or the discussion in Appendix F.1.2 of our work. Thus, we found that such experimental settings are not reflective of the realistic SOC problems one must tackle in high-dimensional machine learning, such as reward fine-tuning generative models. Given that our work is already extensive and multifaceted, we decided not to include such experiments in it. Adding such experiments would distract the reader from the single-minded goal of reward fine-tuning. A complete theoretical analysis between different SOC algorithms is worth a more focused look in a future work.

We thank the reviewer for their comments, and provide answers below.

• “Weaknesses of the memory-less noise schedule: The statement “however, there does not yet exist a simple approach which actually provably generates from the tilted distribution” is, to the best of my knowledge not true [...]”

Our work is not the first one that tackles the problem of reward fine-tuning for diffusion processes, and it is not the first one that uses memoryless noise schedules (previous works did not use this terminology). However, the previous works that the reviewer refers to consider specific interpolation paths for the base model: [2] focuses on the denoising diffusion path ($\alpha_t = \sqrt{t}$, $\beta_t = \sqrt{1-t}$) and [3] focuses on the Föllmer drift path ($\alpha_t = t$, $\beta_t = \sqrt{t(1-t)}$). We are the first work to single out memorylessness as the unified condition to guarantee that the optimal process of the SOC problem will generate samples from the right distribution: the initial and final samples of the base process need to be independent. In our view, this is an important contribution, because it really simplifies the fine-tuning pipeline when compared to approaches that do not take it into account, such as the one of Uehara et al. [1]. While the fine-tuning approach of [1] is also generic in that it can work for multiple generative paths, it is arguably not simple, because it involves learning a value function, and then solving two SOC problems.

Moreover, the insights of our Prop. 1 and Thm. 1 are clearly novel with respect to [2,3,4,5]: we show that the base process is memoryless provided that the noise schedule is “large” enough, but if we want to be able to perform inference on the fine-tuned model using a different noise schedule than the fine-tuning one (such as the commonly used zero noise schedule), we “must” use a unique memoryless noise schedule, which is the one that corresponds to the time reversal of the explicit Ornstein Uhlenbeck process (in [2]) or Brownian bridge process (in [3]). The proof techniques for our results are very different from techniques used in the literature, which makes them interesting from a theoretical perspective as well. Finally, we show samples using the zero noise schedule and the memoryless noise schedule in Figures 10 & 11, so we have shown that the memoryless noise schedule samples are both qualitatively and quantitatively worse. It is thus important to convert back to the zero noise schedule after fine-tuning, which is actually the main reason we need to use the memoryless noise schedule. This is a novel contribution of ours. We emphasize that the contributions we make in Subsec. 3.3 are highly novel and are not mere rewriting of previous works.

For these reasons, we have reformulated this sentence to be more precise: “there does not yet exist a simple, generic approach [...]”.

• “Did the authors compare the proposed approach to the work by Uehara et. al. [1]?”

We do compare our approach to the work of Uehara et al. [1] (see App. B, lines 1318-1320) in words, but not experimentally. The work by Uehara et al. [1] tackles the same problem as ours, but instead of starting from a memoryless base process, it learns the initial value function $V(X_0,0)$, then learns a sampler for $\exp(-V(X_0,0))$, and then solves the fine-tuning SOC problem by setting $p_0$ to be the approximation of $\exp(-V(X_0,0))$ that has been learned. We did not compare our method with theirs experimentally, as their pipeline requires a lot more steps, and their memory usage is substantially higher because they need to store the value function model and the sampler for $\exp(-V(X_0,0))$ in memory as well. Hence, an apples-to-apples comparison with our method and all other baselines is very challenging.

We thank the reviewer for their comments, and provide answers below.

• “there lacks a baseline in experiments in directly optimizing objective (17) (though theorectically there is value bias as shown in the paper). [...] There is comparison for this in the synthetic examples in Figure 2, but more practical downstream tasks evaluations are needed [...].”

We tried different noise schedules to fine-tune in Figure 2 to exemplify that such noise schedules do not yield the desired target distribution. Thus, our theoretical framework provides theoretical evidence that the noise schedule that we need to use is the memoryless one.

Apart from the synthetic experiments in Figure 2, we would like to remark that the initial version of our paper already contained experiments at different noise schedules: Table 6 in App. A shows that when we fine-tune with $\sigma(t)=1$, and perform inference at either $\sigma(t)=1$ or $\sigma(t)=0$, the quality metrics are significantly worse across the board (except for diversity). Following the intuition from Figure 2, using non-memoryless schedules yields a generated distribution which is closer to the initial one than desired. In case the author referred to baselines from existing papers, we want to point out that there are no existing works that solve the SOC problem with wrong noise schedules without further corrections: existing works either use arbitrary schedules, but learn additional objects to eliminate the value function bias (Uehara et al., 2024), or just avoid KL regularization (Clark et al., 2024, Xu et al., 2023).

• "1) On line 229-230, the expression "Dividing (14) by (15)" is odd as (14) is not an equality, "Plug in normalization constant (15) to (14)" might be better."

We thank the reviewer for pointing this out. We changed this sentence to: “Dividing the RHS of (14) by (15)...”. We also added a previous sentence that mentions “the right-hand side (RHS) of (14)”. Hence, the terminology RHS should be clear to the reader.

• "2）The adjoint method in Section 4 needs more introduction or discussion, or an earlier pointer to the part in appendix, including more clarification on what is adjoint on line 350-352, what is this loss on Equation (21); The reference in Appendix should be referred earlier instead of being deferred to until line 422."

We agree with the reviewer. We added an earlier pointer to the appendix: “There are two approaches which yield the same gradient in the small step size limit (see \Cref{sec:adjoint_method} for more detail): [...]”. We added clarification on the loss on Equation (21) by changing the previous sentence to “The gradient of the continuous adjoint loss is given by [...]”.

• "3）In proposition 2, it is better to first define earlier of the adjoint matching objective instead of combining the definition with the proposition."

Observe that Proposition 2 introduces the Basic Adjoint Matching loss, a preliminary to the Adjoint Matching loss (equations 25-27), which is our main algorithmic contribution. The Basic Adjoint Matching loss shares the gradient with the Continuous Adjoint loss and with the Discrete Adjoint loss (up to numerical errors in the latter case). We introduce it because it offers a novel perspective on an existing algorithm, and it opens the door to proving the convergence guarantee in Prop. 2. Given the space constraints, we do not want to further emphasize the role of the Basic Adjoint Matching loss, and opted for a compact explanation which combines the definition with the proposition.

We thank the reviewer for their comments, and provide answers below.

• “it would be more convincing if results on denoising diffusion models are provided.”

As we mention in Section 2, both Flow Matching and continuous-time DDIM are described by Equation 5 and 6. While our particular models were trained using the particular scheduler usually associated with Flow Matching, our unified notation includes all pre-training schedulers. In our initial experiments, we saw good initial results on Stable Diffusion 2 model, and then switched to using our own pretrained model. A reason for this choice is that flow matching and its choice of noise scheduler is increasingly adopted as the standard approach, e.g., in Stable Diffusion 3 or Flux, since there is evidence that it is superior to DDIM. Another reason for this is to use only kosher data for training; our models are trained only on data with proper licensing. This is needed to protect individual privacy and to adhere to ethical standards in data usage. As the reviewer points out, we also provide an algorithm box (Alg. 2) translating our approach to denoising diffusion models, and we do not expect significant differences.

• “The experiments with classifier-free guidance do not seem comprehensive. It would be better if there are similar quantitative comparisons with other baselines than selected DRaFT-1.”

We would appreciate additional explanation from the reviewer to clarify their remarks. In our reply below, (i) we explain why we believe that our experiments on classifier-free guidance are appropriate within the scope of our work, and (ii) we explain why we only include classifier guidance comparisons against DRaFT-1. If our reply does not satisfy the reviewer, we will be happy to follow up with a response once more detail has been provided.

(i) We would appreciate some more detail on how the experiments on classifier-free guidance do not seem comprehensive. For reference, we explain our experiments relating to classifier-free guidance in the next two paragraphs, and would like to note that in our first submission we already included pointers and a brief discussion of the guidance results in Section 5 of the paper.

In Figure 4 (App. A), we evaluate three different guidance settings: $w=0$ (no guidance), $w=1$, $w=4$. The numerical values can be found in Table 4 (App. A), which also shows guidance results for the pre-trained model. Higher guidance values yield higher Clipscore (prompt alignment) and HPS v2 (human quality perception) values at the expense of sample diversity. In fact, models fine-tuned with adjoint matching perform much better than those fine-tuned with DRaFT-1 at high guidance values ($w=4$), relative to their respective performance at $w=0$ and $w=1$. This is not surprising because Adjoint Matching preserves the theoretical connection of the fine-tuned vector field to the score of the fine-tuned distribution, which means that classifier-free guidance on the fine-tuned model is as principled as on the base model. That does not hold for directly fine-tuned models.

Moreover, in Figure 5 (App. A) we also present a qualitative comparison of images generated with the same prompt and the same random seed at different guidance values. The generated images confirm that classifier-free guidance can be used successfully on models that have been fine-tuned using adjoint matching; and that high guidance values yield less diverse images that fit the prompt more narrowly. Also, in Figures 7 and 8 (App. A) we compare images generated with the pre-trained model, Adjoint Matching, and DRaFT-1.

(ii) The reviewer suggests that it would be good to have guidance comparisons to other baselines beyond DRaFT-1. We want to point out that according to our experiments, DRaFT-1 is the best performing baseline fine-tuning algorithm, and that all other baselines perform worse, both quantitatively and qualitatively. Moreover, note that the classifier-free guidance results are answering a simple “what-if”, and classifier-free guidance is completely orthogonal to the fine-tuning method; we see that any model with poorer fine-tuned results also has poorer results with classifier-free guidance. For these two reasons, we thought that it was not necessary to run guidance experiments on alternative baselines.

We thank the reviewers for their work. We made some changes and additions to our manuscript following their advice, which are summarized in the "General response" above, and discussed in detail in our individual responses to the reviewers. We would like to remind the reviewers which have not yet replied to our responses (4cB8, 7e67, pgS7) that the discussion period extension will come to an end soon. We encourage them to increase their score if they find our responses and revisions satisfactory.