PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, Ricky T. Q. Chen

arXiv:2409.08861v5cs.LGmath.OCstat.ML

TL;DR

Reward fine-tuning of dynamical generative models lacks simple, theoretically grounded methods that account for dependence between noise variables and generated samples. The paper formulates fine-tuning as stochastic optimal control with a memoryless noise schedule and introduces Adjoint Matching, reporting improvements across multiple axes, including consistency, preference-model generalization, and diversity.

  • Problem

    Reward fine-tuning of dynamical generative models lacks simple, theoretically grounded methods that account for dependence between noise variables and generated samples.

    Source: §1 · p. 2

  • Method

    The paper formulates reward fine-tuning as stochastic optimal control with a memoryless noise schedule and casts control optimization as regression in its Adjoint Matching algorithm.

    Source: §5.2 · p. 10, §8 · pp. 15–16

  • Results

    The authors report that their memoryless SOC formulation works better than multiple existing diffusion fine-tuning works, while Adjoint Matching outperforms related gradient-based methods and generalizes to unseen human preference rewards with better consistency and good diversity.

    Source: §1 · p. 2, §8 · p. 16

  • Takeaways & Limitations

    The memoryless schedule supports convergence to the same tilted distribution used in learning from human feedback, corresponding to DDPM sampling for diffusion and a new process for flow models.

    Source: §1 · p. 2, §8 · pp. 15–16

    Adjoint methods are gradient-based optimization on a highly non-convex problem, and prior work has reported empirical instability.

    Source: §5.1.1 · p. 9

Abstract

from arXiv · show

Dynamical generative models that produce samples through an iterative process, such as Flow Matching and denoising diffusion models, have seen widespread use, but there have not been many theoretically-sound methods for improving these models with reward fine-tuning. In this work, we cast reward fine-tuning as stochastic optimal control (SOC). Critically, we prove that a very specific memoryless noise schedule must be enforced during fine-tuning, in order to account for the dependency between the noise variable and the generated samples. We also propose a new algorithm named Adjoint Matching which outperforms existing SOC algorithms, by casting SOC problems as a regression problem. We find that our approach significantly improves over existing methods for reward fine-tuning, achieving better consistency, realism, and generalization to unseen human preference reward models, while retaining sample diversity.

1 Introduction

Flow Matching and diffusion models are widely used, but their base samples often need improvement. This paper frames reward fine-tuning as stochastic optimal control and introduces a memoryless noise schedule and Adjoint Matching.

1 Introduction

  • Prior diffusion fine-tuning often neglects the base distribution, while reward networks can invite adversarial artifacts; existing regularization lacks a simple provable tilted-distribution guarantee.The paper motivates fine-tuning toward the reward-tilted distribution while accounting for the base model.Source: §1 · p. 2
  • The memoryless noise schedule removes dependency between noise variables and generated samples, correcting value-function bias and enabling provable convergence to the tilted distribution, including for noiseless Flow Matching.The paper identifies this as the first fine-tuning of noiseless Flow Matching models.Source: §1 · p. 2
  • Adjoint Matching solves stochastic optimal control problems using gradient-based scalability and a least-squares regression objective, and applies beyond reward fine-tuning.The authors characterize the method as orthogonal to the reward fine-tuning application.Source: §1 · p. 2
  • Comparisons find that the proposed method generalizes to unseen human-preference reward models, improves text-to-sample consistency, and retains good diversity.The comparisons assess realism, consistency, and diversity.Source: §1 · p. 2

2 Preliminaries on dynamical generative models

These dynamical generative models turn initial noise into samples through iterative stochastic-process simulation. The section introduces Flow Matching and denoising diffusion models before unifying them in a common notation.

2 Preliminaries on dynamical generative models

  • Flow Matching and diffusion processes share the reference flow’s time marginals, so simulating from Gaussian noise yields samples distributed as the data.The reference flow interpolates between noise and data using time-dependent coefficients with α0 = β1 = 0 and α1 = β0 = 1.Source: §2 · p. 3

2.1 Flow Matching

  • Flow Matching uses an ODE whose parametric velocity matches the reference-flow derivative; an SDE variant preserves the same marginals through a compensating drift term.A commonly used schedule is αt = t and βt = 1 − t; the SDE allows an arbitrary state-independent diffusion coefficient σ(t).Source: §2.1 · p. 3

2.2 Denoising Diffusion Models

  • DDIM sampling iteratively transforms Gaussian noise x0 into a denoised sample xK, using a noise predictor trained to estimate each iterate’s noise.The sequence ᾱk increases from 0 to 1, while the update’s noise sequence σk is arbitrary; DDPM is presented as a particular case of DDIM (Song et al. (2021a); Ho et al. (2020)).Source: §2.2 · p. 3

3 Flow Matching and diffusion models from a common perspective

The section places Flow Matching and diffusion models in a shared continuous-time framework, expressing their generative processes as SDEs unified by a reference flow. This formulation is intended to support fine-tuning across both model families.

3 Flow Matching and diffusion models from a common perspective

  • The authors use continuous-time formulations to simplify notation and enable fine-tuning methods inspired by the continuous-time paradigm, which perform better than discrete-time counterparts in their empirical validations.The unified notation covers models following the reference flow’s time marginals, whether pretrained with Denoising Diffusion or Flow Matching.Source: §3 · p. 4
  • With uniform time discretization, the DDIM update is first-order equivalent to an Euler–Maruyama discretization, and matches reference-flow marginals when αt = √¯αt and βt = √1 −¯αt (Ho et al., 2020).The continuous-time identification sets t = k/K and relates the discrete noise predictor to its continuous-time counterpart.Source: §3 · p. 4
  • Up to first-order approximation, the DDPM sampling scheme is a special case of DDIM when σ(t) = ˙¯αt/¯αt (Ho et al. (2020); Algorithm 2).Source: §3 · p. 4
  • Under a Normal-distributed X0 and affine reference flow, score-function expressions for Flow Matching velocity and diffusion noise prediction unify both generative processes as an SDE parameterized by reference-flow coefficients.The shared process is dXt = b(Xt, t) dt + σ(t) dBt.Source: §3 · p. 4

4 Fine-tuning as “memoryless” stochastic optimal control

Reward fine-tuning via stochastic optimal control can recover the reward-tilted sample distribution without statistical bias, but only when fine-tuning uses a specific memoryless noise schedule. The section explains the SOC formulation, the bias from naïve KL regularization, and why this schedule resolves it.

4.1 Preliminaries on the stochastic optimal control problem formulation

  • SOC optimizes a control vector field for a stochastic process, and its affine control cost is equivalent to a KL divergence from the base process, yielding a KL-regularized MaxEnt RL interpretation (Todorov (2006); Ziebart et al. (2008)).The KL term encourages the controlled process to remain close to the base distribution, while negative state costs correspond to intermediate and terminal rewards.Source: §4.1 · p. 5

4.2 The initial value function bias problem

  • Naïve KL-regularized fine-tuning is biased because its optimal distribution includes the initial value function V(X0, 0); with noiseless sampling, it has no effect.The bias arises because the initial noise X0 and generated sample X1 are dependent, rather than yielding the desired tilted distribution.Source: §4.2 · p. 6

4.3 The memoryless noise schedule for fine-tuning dynamical generative models

  • Fine-tuning with the SOC reward objective must use the memoryless noise schedule σ(t) = √2η_t to generate samples from the tilted distribution; afterward, sampling can use any noise schedule.The memoryless condition makes the initial noise X0 independent of the final sample X1. The schedule approaches zero at t = 1 and is infinite at t = 0.Source: §4.3 · pp. 6–8

5 Adjoint Matching for control-affine stochastic optimal control

The section develops Adjoint Matching by combining continuous adjoints with least-squares regression for control-affine stochastic optimal control. Its key design removes importance weighting and unnecessary adjoint terms while retaining a theoretically grounded regression objective.

5.1.1 The adjoint method

  • Discrete adjoints store the solver graph and consume substantial memory, while discrete and continuous adjoints scale to high dimensions but can be empirically unstable.The discrete approach may require gradient checkpointing to reduce memory use.Source: §5.1.1 · p. 9

5.1.2 Importance-weighted matching objectives for regressing onto the optimal control

  • Cross-entropy and SOCM (Domingo-Enrich et al. (2023)) use path-integral importance-weighted regression, but weight variance grows exponentially with process dimension.This makes importance-weighted matching impractical for high-dimensional fine-tuning, despite its convex functional landscape and potential for stable training.Source: §5.1.2 · pp. 9–10

5.2 Adjoint Matching

  • Adjoint Matching regresses the current control toward a target vector field, and its objective has the optimal control as its unique critical point.The target is −σ(t)^T∇xJ(u; x, t); the basic objective’s gradient matches the continuous adjoint method, and the regression enforces a consistency relation.Source: §5.2 · pp. 10–11
  • The lean adjoint removes terms with zero expectation at the optimum, yielding a different expected gradient that may reduce variance and computational cost.It avoids computing the control Jacobian, while the final Adjoint Matching objective retains the optimal control as its unique critical point.Source: §5.2 · p. 11
  • Adjoint Matching is an unweighted least-squares regression that avoids high-variance importance weights while retaining a target-vector-field interpretation.Unlike importance-weighted matching, it scales like adjoint methods without requiring importance sampling.Source: §5.2 · p. 11

6 Related work

The related work distinguishes reward-based fine-tuning from direct preference optimization and situates reward fine-tuning among diffusion methods, inference-time optimization, and optimal-control approaches. It also connects conditional sampling in inverse problems to the reward fine-tuning framework.

Fine-tuning from human feedback

  • Reward-based RLHF learns a reward model before maximum-entropy reinforcement learning, whereas DPO fine-tunes directly from preferences; differentiable rewards in flow and diffusion models favor reward-based approaches.DPO is typically applied with filtered data and does not explicitly use a reward model.Source: §6 · p. 12

Fine-tuning for diffusion models

  • Diffusion fine-tuning methods use policy gradients, proximal policy optimization, reward gradients, or SOC, but differ in their use of reward differentiability and connection to the tilted distribution.Fan and Lee (2023) use policy gradients; Black et al. (2024) use proximal policy gradients without leveraging reward differentiability. Fan et al. (2023) consider KL-regularized rewards but do not connect them to the tilted distribution; Xu et al. (2023) and Clark et al. (2024) use reward gradients and heuristics. Uehara et al. (2024b) target the tilted distribution but require a value function and two SOC problems. Bruna and Han (2024) study quadratic rewards, while Zhang et al. (2024) address GFlowNets.Source: §6 · p. 12

Inference-time optimization methods

  • Inference-time methods guide velocities, optimize initial noise, or estimate controls per sample, rather than amortizing control into a fine-tuned generative model.Velocity guidance is heuristic and its generated distribution is not well understood; initial-noise optimization takes an opposite approach to the initial value bias problem.Source: §6 · p. 12

Optimal control in generative modeling

  • Optimal control has trained ODE-, SDE-, and jump-process generative models, while simulation-free matching objectives became favored for simplified constructions; these developments motivate SOC for reward fine-tuning.SOC has also supported sampling from unnormalized distributions and transport problems with important intermediate path distributions.Source: §6 · p. 12

Conditional sampling in inverse problems

  • Conditional inverse-problem sampling with likelihood p(y|x) is a special case of reward fine-tuning using r(x) = log p(y|x).Denker et al. (2024) formulate SOC and use log-variance and moment losses; Wu et al. (2023a) propose an SMC sampler with increased inference-time compute, while Du et al. (2024) approximately solve a Lagrangian formulation using Gaussian paths.Source: §6 · pp. 12–13

7 Experiments

Experiments fine-tune a Flow Matching model with reward scaling and compare memoryless Adjoint Matching with reward-fine-tuning and SOC baselines. Results show strong overall tradeoffs, while λ and classifier-free guidance expose consistency–diversity tradeoffs.

Evaluation setup

  • Evaluation measures consistency with ClipScore and PickScore, unseen-preference generalization with HPSv2, and diversity with DreamSim; baselines include DPO (Wallace et al., 2023a), ReFL (Xu et al., 2023), and DRaFT-K (Clark et al., 2024).DRaFT-1 performs best among these baselines and is used for additional ablations.Source: §7 · p. 14

Evaluation results

  • Adjoint Matching achieves the best Pareto fronts and outperforms existing baselines across reported metrics at roughly matched wall-clock time.Continuous and discrete adjoint methods yield poorer consistency and human-preference metrics.Source: Figure 5 · p. 15, §7 · p. 15

Ablation: base model vs. reward tradeoff

  • Increasing λ improves consistency and human preference while reducing diversity; Adjoint Matching remains stable across λ, unlike baselines that often rely on early stopping.λ scales the ImageReward reward against KL regularization; Figure 3 illustrates the tradeoff qualitatively.Source: §7 · pp. 13, 15, Figure 3 · p. 14

Ablation: classifier-free guidance

  • Higher classifier-free guidance weight improves text-to-image consistency and, to some extent, human preference, but reduces sample diversity.Guidance is applied after fine-tuning the conditional model, although this is not principled because the unconditional model is not fine-tuned.Source: Figure 4 · p. 14, §7 · p. 15

8 Conclusion · Appendix

The paper formulates reward fine-tuning for dynamical generative models as stochastic optimal control with a memoryless noise schedule, and introduces Adjoint Matching as a regression-based SOC algorithm. It reports improvements in consistency, generalization to unseen human preferences, and sample diversity.

8 Conclusion

  • The SOC formulation uses a memoryless noise schedule to converge to the same tilted distribution used in large language model learning from human feedback.Source: §8 · p. 15
  • The approach significantly outperforms baseline methods in text-to-image consistency, generalization to unseen human preference, and sample diversity.The conclusion reports these gains on large-scale text-to-image generation.Source: §8 · p. 16
  • Adjoint Matching outperforms related gradient-based methods for solving stochastic optimal control problems.Source: §8 · p. 16
  • The memoryless SOC formulation works better than multiple existing works on fine-tuning diffusion models.Source: §8 · p. 16
  • For diffusion models, the memoryless schedule corresponds to DDPM sampling; for flow models, it defines a new Memoryless Flow Matching generative process.Source: §8 · p. 16
  • Adjoint Matching solves SOC problems by casting them as a regression problem.Source: §8 · p. 16
  • The paper presents the first theoretically driven algorithm for fine-tuning Flow Matching models.Source: §8 · p. 16

A Additional Figures & Tables … B.2.2 Proof of Lemma 2

The appendix adds fine-tuning comparisons and sample visualizations, then develops the continuous-time and forward–backward SDE results and their proofs. It shows that arbitrary noise schedules preserve the time-reversed trajectory law, while the alternative constant schedule σ(t) = 1 has a reward and metric disadvantage except for diversity.

A Additional Figures & Tables

  • The appendix also plots ImageReward, control cost, and ClipScore against wall-clock time, averaging three fine-tuning runs and reporting standard-error confidence intervals.Evaluation uses separate test datasets of size 200.Source: A Additional Figures & Tables · p. 23, Figure 6 · p. 23
  • Supplementary tables report text-to-image metrics, additional evaluation metrics, CFG results, optimization hyperparameters, and sampling-step comparisons.The tables specify fine-tuning and inference noise schedules where applicable.Source: Table 3 · p. 24, Table 4 · p. 24, Table 5 · p. 25
  • With σ(t) = 1, initial value-function bias prevents high ImageReward, and other metrics are lower than with the memoryless schedule except diversity.ImageReward is the fine-tuning reward function.Source: Table 7 · p. 26
  • Supplementary figures show generated samples across classifier-free guidance settings and noise schedules, including prompts featuring people.They include guided samples with w = 1 and σ(t) = 0, and unguided samples with w = 0 under σ(t) = 0 or σ(t) = √2ηt.Source: Figure 7 · p. 23, Figure 8 · p. 27, Figure 9 · p. 28

B.1 The continuous-time limit of DDIM

  • DDIM’s update yields the continuous-time SDE through first-order approximations, with its discretization identified as Euler–Maruyama.The derivation introduces continuous time and a step size h = 1/K.Source: §B.1 · p. 31

B.2.1 Proof of Lemma 1

  • DDIM and Flow Matching choices of κ_t and η_t fulfill the stated conditions (50).The proof checks the conditions for each model’s schedule.Source: §B.2 · p. 31, §B.2.1 · p. 33

B.2 Forward and backward stochastic differential equations

  • For arbitrary noise schedules, forward and backward SDEs have equal marginals up to time reversal and identical trajectory distributions after flipping time.The trajectory result equates joint distributions at any sequence of times, and is stated as equivalence of the full forward and backward process measures.Source: §B.2 · pp. 32–33

B.2.2 Proof of Lemma 2

  • The proof of Lemma 2 solves the forward SDE by variation of parameters and expresses its solution using Gaussian noise.It integrates the equation and uses that the Brownian increment has the distribution of a Gaussian-noise expression.Source: §B.2.2 · pp. 33–34

B.2.3 Proof of Proposition 4 … C.1 Maximum entropy RL

These appendices derive discrete-time reverse-process and score/flow relationships, then connect continuous-time stochastic optimal control to maximum-entropy reinforcement learning, whose optimal trajectory distribution is reward-tilted.

B.2.3 Proof of Proposition 4

  • The proof discretizes the forward process on K + 1 equispaced timesteps, tracks terms through linear order in h, and neglects O(h^3/2) and higher.It corrects gaps in the earlier derivation by De Bortoli et al. (2021), and concludes that the resulting update is the Euler-Maruyama discretization of the backward process.Source: §B.2.3 · pp. 34, 36

B.3 The relationship between the noise predictor ϵ and the score function

  • The noise predictor ϵ and the score function s are related by the displayed formula, connecting DDIM noise prediction with score prediction.The derivation uses variables with the same distribution to establish the relationship.Source: §B.3 · p. 36

B.4 The relationship between the vector field v and the score function

  • For an arbitrary diffusion coefficient, the drift needs a correction term; the resulting SDE and the original ODE yield matching Fokker-Planck equations.This derivation gives the Flow Matching SDE with an arbitrary noise schedule.Source: §B.4 · p. 37

C Stochastic optimal control as maximum entropy RL in continuous space and time

  • Taking the stepsize to zero in KL-regularized RL with Euclidean actions and conditional-Gaussian transitions yields an SOC problem, and RL algorithms have diffusion fine-tuning analogs.The paper also points to a continuous-time connection in which expected control cost equals the KL divergence between trajectory distributions conditioned on starting points.Source: C Stochastic optimal control as… · p. 37
  • The continuous-time SOC–MaxEnt RL connection equates expected control cost with the KL divergence between trajectory distributions conditioned on their starting points.The paper identifies this result in Appendix C.4 as a direct, rigorous connection.Source: C Stochastic optimal control as… · p. 37

C.1 Maximum entropy RL

  • MaxEnt RL adds policy entropy to each step’s reward to encourage exploration and robustness, while KL regularization against a base policy limits deviation.The classical setup models an agent taking actions under a policy and receiving rewards along a trajectory.Source: §C.1 · pp. 37–38
  • The optimal MaxEnt RL policy induces the tilted trajectory distribution with initial-state marginal p0.This follows from recasting the objective over trajectory distributions and implies that the optimal policy’s distribution is q∗.Source: §C.1 · p. 38

C.2 From maximum entropy RL to stochastic optimal control … D Proofs of Section 4.3: memoryless noise schedule and fine-tuning recipe

The appendix establishes the continuous-time connection between maximum-entropy reinforcement learning and stochastic optimal control, proves the associated loss convergence, and applies Girsanov’s theorem to the control-cost regularizer. The supplied passages cover Sections C.2–C.4 but contain no evidence for Section D.

C.2 From maximum entropy RL to stochastic optimal control

  • Under vector-valued states and actions, vector-field policies, and conditional Gaussian transitions, MaxEnt RL converges to the SOC problem as the timestep count grows.The policies use a control vector field, and the Gaussian transition has drift b + σu and covariance hσσ⊤.Source: §C.2 · p. 38
  • The correspondence lets SOC inherit generic MaxEnt RL losses and additionally permits differentiation of state and terminal costs.The optimal- and base-policy distributions translate to distributions induced by the optimal control and null control.Source: §C.2 · pp. 38–39

C.3 Proof of Proposition 5: from MaxEnt RL to SOC

  • As the step count grows, controlled and uncontrolled trajectories become Euler–Maruyama discretizations of their respective SDE processes.The controlled process uses u; the base policy corresponds to the uncontrolled process with u = 0.Source: §C.3 · p. 39
  • With rewards scaled by the timestep, the limiting MaxEnt RL loss is the SOC loss.The running and terminal rewards correspond to the SOC running and terminal costs in the limit.Source: §C.3 · p. 39
  • The proof also shows that augmenting each policy with its fixed transition kernel preserves the per-step KL regularizer.The lemma states equality between the augmented-kernel and policy KL terms.Source: §C.3 · pp. 39–40

C.4 Proof of equation (18): the control cost is a KL regularizer

  • The proof of the control-cost KL regularizer applies Girsanov’s theorem to the probability measures of the controlled and reference SDE solutions.The passage identifies these measures as those of the solutions used to establish equation (18).Source: §C.4 · p. 41

D.1 Proof of Proposition 1: the memoryless noise schedule … E.1 Derivation of the Continuous Adjoint method

The appendices prove when the memoryless noise condition holds, derive fine-tuning recipes for general noise schedules, and establish the continuous-adjoint gradient used to optimize the control.

D.1 Proof of Proposition 1: the memoryless noise schedule

  • The proof derives the score from the flow-matching interpolation distribution, whose forward and backward trajectory laws agree after a time flip.Their marginals therefore satisfy p⃗_t = p_{1-t}.Source: §D.1 · p. 42
  • Assumption (25) is sufficient for memorylessness: it makes the forward process endpoints independent, and trajectory time-flip symmetry transfers that independence to the backward process.The proof concludes that assumption (25) is also necessary: if the stated limit is nonzero at some time, endpoint independence fails in general.Source: §D.1 · p. 43

D.2 Proof of Theorem 1: fine-tuning recipe for general noise schedules

  • Theorem 1's derivation uses the HJB equation to identify the reward-tilted process's optimal control and recover the fine-tuned inference dynamics.The proof sets p*(x) proportional to p_base(x) exp(r(x)) and relates the control to the gradient of the log-density difference.Source: §D.2 · pp. 43–45
  • For DDIM, the derived fine-tuned inference SDE matches equation (6) with ϵ = ϵ*.This identifies the fine-tuned noise prediction with ϵ*.Source: §D.2 · p. 45
  • For Flow Matching, the derived fine-tuned inference SDE matches equation (4) with v = v*.This identifies the fine-tuned velocity with v*.Source: §D.2 · p. 45

E.1 Derivation of the Continuous Adjoint method

  • Proposition 6 expresses the parameter gradient of the adjoint loss as equation (32), with its derivation using Lagrange multipliers and stochastic integration by parts.The approach mirrors and extends Domingo-Enrich et al. (2023, Lemma 8).Source: §E.1 · p. 46
  • Lemma 5 gives the unique solution of the adjoint ODE for a fixed vector field and also covers a control parameterized by θ.The derivation establishes these statements using the ODE and its parameter dependence.Source: §E.1 · pp. 46–47

E.2 Proof of Proposition 2: Theoretical guarantees of the basic Adjoint Matching loss … F.2 Adapting Diffusion-DPO (Wallace et al., 2023a) to flow matching

The appendices establish that both basic and full Adjoint Matching losses have the optimal control as their only critical point, and give a DDIM fine-tuning procedure. They also adapt reward-feedback and Diffusion-DPO approaches to flow matching, with the reported on-policy DPO results no better than the base model overall.

E.2 Proof of Proposition 2: Theoretical guarantees of the basic Adjoint Matching loss

  • The basic Adjoint Matching loss has only the optimal control as a critical point, using the HJB equation and its unique solution to establish optimality.A control satisfying u(x, t) = −σ(t)⊤∇xJ(u; x, t) makes J satisfy the HJB equation; uniqueness then identifies the optimal control.Source: §E.2 · pp. 48–49

E.3 Theoretical guarantees of the Adjoint Matching loss

  • The full Adjoint Matching loss likewise has the optimal control as its only critical point, by reducing its critical points to those of the basic loss.The proof uses uniqueness of an integral equation for the expected lean adjoint to establish the reduction.Source: §E.3 · pp. 49–50

E.4 Pseudo-code of Adjoint Matching for DDIM fine-tuning

  • The DDIM fine-tuning procedure solves the lean adjoint ODE backward along each trajectory, computes the Adjoint Matching objective, then updates model parameters by gradient descent.The algorithm initializes the fine-tuned denoiser from the pretrained denoiser and computes the trajectory and lean adjoint without gradients.Source: §E.4 · p. 51

F.1 Adapting ReFL (Xu et al., 2023) to flow matching

  • ReFL (Xu et al., 2023) increases reward on denoised samples; its flow-matching adaptation expresses the denoiser map in terms of the vector field.The described ReFL update uses detached trajectories.Source: §F.1 · p. 51

F.2 Adapting Diffusion-DPO (Wallace et al., 2023a) to flow matching

  • The adapted Diffusion-DPO loss for flow matching uses ranked sample pairs, and the authors also describe a reward-based version when a reward function is available.The adapted loss is expressed in terms of the flow-matching vector field.Source: §F.2 · p. 52
  • The on-policy DPO model performs like the base model or somewhat worse on some metrics, so the authors conclude it is not competitive when the base model is not already good.Their preference pairs are sampled from the current model and are imperfect, unlike the off-policy, curated high-quality pairs used by Wallace et al. (2023a).Source: §F.2 · pp. 52–53

G Experimental details … G.6 Remarks on number of sampling timesteps

The appendix specifies shared fine-tuning settings and implementation choices, then reports computational costs and how inference timestep counts affect performance.

G Experimental details

  • Fine-tuning used K = 40, effective batch size 40 on two A100 GPUs, and 40k licensed text prompts per epoch, with shared Adam settings except for Discrete Adjoint’s learning rate.The shared Adam settings included learning rate 2 × 10^-5; Discrete Adjoint used 1 × 10^-5 in other reported tables after instability with the shared rate.Source: G Experimental details · p. 53

G.1 Noise schedule details

  • With K = 40, the noise schedule used small offsets to avoid an infinite coefficient at t = 0 and to keep diffusion non-negligible near t = 1.For the lean adjoint terminal state, the method uses a final noiseless update to reduce gradient bias from noise in the regular final iterate; this agrees in the continuous-time limit.Source: §G.1 · p. 53

G.2 Selection of gradient evaluation timesteps

  • For computational efficiency, gradient evaluation sampled ten timesteps uniformly from 0 to 0.725 and always included the final ten timesteps, 0.75 to 0.975.The authors state that tuning the final 25% of steps is critical for empirical performance, while initial steps are less important.Source: §G.2 · p. 54

G.3 Loss function clipping: the LCT hyperparameter

  • Adjoint Matching clipped loss terms at LCT = 1.6 × λ^2, while Continuous Adjoint used LCT = 1600 × λ^2.Clipping limits high-magnitude early-step terms that can muffle the important late-step gradient signal; the threshold scales quadratically with λ.Source: §G.3 · p. 54

G.4 Computation of evaluation metrics

  • ClipScore diversity was the variance of embeddings from 40 generations per prompt averaged across 25 prompts; PickScore and Dreamsim diversity used the analogous computation.The appendix also states that PickScore and HPS were computed using their respective libraries.Source: §G.4 · p. 54

G.5 Remarks on computational costs

  • Adjoint Matching took 156 seconds per iteration versus 152 seconds for Discrete Adjoint, while Continuous Adjoint took 204 seconds.Adjoint Matching and Discrete Adjoint perform similar forward and backward passes; Continuous Adjoint adds backward passes to compute gradients of the control norm during adjoint ODE solving.Source: §G.5 · p. 55

G.6 Remarks on number of sampling timesteps

  • At 100 and 200 timesteps, Adjoint Matching metrics were statistically equal to its 40-timestep results, while 10 and 20 timesteps performed worse, especially 10.At 100 and 200 timesteps, its performance gap over DRaFT-1 increased; at 10 timesteps, that gap vanished except for diversity, where Adjoint Matching remained clearly better.Source: §G.6 · p. 55

Related papers

Loading 2409.08861v5…