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

The paper addresses the lack of theoretically sound reward fine-tuning methods for Flow Matching and denoising diffusion models. It formulates fine-tuning as stochastic optimal control, enforces a memoryless noise schedule, and introduces Adjoint Matching as a regression-based solver. The approach outperforms related methods across multiple evaluation axes while retaining diversity.

  • Problem

    Existing reward fine-tuning methods do not yet provide a simple approach that provably generates the desired reward-tilted distribution while accounting for the base model.

  • Method

    The paper formulates reward fine-tuning as stochastic optimal control with a memoryless noise schedule and solves SOC through the regression-based Adjoint Matching objective.

  • Results

    The approach significantly outperforms baseline methods across multiple evaluation axes, improving consistency, realism, and generalization to unseen human preference reward models while retaining sample diversity.

  • Takeaways & Limitations

    The paper provides a theoretically driven algorithm for fine-tuning Flow Matching models and applies its formulation to diffusion-model reward fine-tuning.

  • Takeaways & Limitations

    Importance-weighted matching objectives are impractical for high-dimensional applications because importance-weight variance grows exponentially with process dimension.

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

The paper frames reward fine-tuning for dynamical generative models as stochastic optimal control and identifies a memoryless noise schedule as necessary for targeting the desired tilted distribution. It also introduces Adjoint Matching and reports improvements in consistency, generalization, and diversity.

  • Flow Matching and denoising diffusion models are widely used, but base models often do not achieve the desired sample quality.
  • Existing reward fine-tuning methods can ignore the base distribution, while neural reward models may encourage undesirable artifacts.
  • The paper formulates reward fine-tuning of dynamical generative models as a stochastic optimal control problem.
  • The authors prove that naïve fine-tuning creates value function bias and propose a memoryless noise schedule that removes noise-sample dependency.
  • Adjoint Matching casts SOC optimization as least-squares regression, combining gradient-based scalability with a simple regression objective.
  • Comparisons across realism, consistency, and diversity find better unseen-reward generalization and text-to-sample consistency while retaining good diversity.

2 Preliminaries on dynamical generative models

This section introduces dynamical generative models as stochastic processes that iteratively transform Gaussian noise into data-like samples. It reviews Flow Matching and diffusion sampling processes, including their deterministic and stochastic variants.

  • Dynamical generative models transform initial noise variables into samples through iterative stochastic-process simulation.
  • Flow Matching: A Flow Matching reference flow interpolates between Gaussian noise and data using time-dependent coefficients α_t and β_t.
  • Flow Matching and diffusion construct Markov processes whose time marginals match the reference flow, yielding the data distribution at the final time.
  • Flow Matching: Flow Matching can use an ODE with a learned velocity or an SDE with state-independent diffusion, while preserving the same time marginals.
  • Denoising diffusion: DDIM sampling repeatedly denoises an initial Gaussian variable through K updates, with σ_k controlling the injected noise.
  • The base-model pretraining analysis assumes perfect training, and the notation distinguishes discrete-time diffusion quantities from their continuous-time counterparts.

3 Flow Matching and diffusion models from a common perspective

The paper unifies Flow Matching and diffusion generative processes through a common score-based SDE framework. This shared formulation supports a common analysis of their sampling dynamics and fine-tuning.

  • The paper consolidates Flow Matching and diffusion notation using the score function of the intermediate state distribution.
  • Continuous-time formulations are used to unify models with the same reference-flow time marginals and to support continuous-time fine-tuning methods.
  • Continuous-time diffusion: DDIM’s discrete update is approximated by an Euler–Maruyama discretization under uniform time discretization.
  • Continuous-time diffusion: The diffusion reference-flow coefficients satisfy α_t = √¯α_t and β_t = √1 − ¯α_t, while DDPM is a special DDIM instance for a particular diffusion schedule.
  • Substituting score-based expressions yields a unified SDE dX_t = b(X_t,t)dt + σ(t)dB_t for both Flow Matching and continuous-time DDIM.

4 Fine-tuning as “memoryless” stochastic optimal control

The paper formulates reward fine-tuning as stochastic optimal control and identifies initial-value bias in naïve KL-regularized objectives. Enforcing the memoryless noise schedule removes this bias and yields the desired tilted distribution while permitting arbitrary sampling schedules afterward.

  • SOC formulation: Reward fine-tuning is posed as a stochastic optimal control problem related to KL-regularized maximum-entropy reinforcement learning.The control-affine formulation optimizes a controlled stochastic process with drift, diffusion, control, running cost, and terminal cost.
  • Memoryless noise schedule: Fine-tuning with the memoryless schedule provably converges without statistical bias to the tilted distribution, while sampling may later use arbitrary noise schedules.This includes conversion to noiseless sampling such as σ(t) = 0 for Flow Matching and DDIM models.
  • SOC formulation: The affine control cost corresponds to a KL divergence from the base process, encouraging the controlled process to remain close to the base distribution.This interpretation connects the SOC objective to reward fine-tuning while retaining the base process as a prior.
  • Initial-value bias: Naïvely adding KL regularization produces a biased optimal distribution because the initial value function depends on the initial noise variable.The resulting distribution is therefore not the desired reward-tilted distribution, especially when the sampling process is noiseless.
  • Memoryless noise schedule: The memoryless noise schedule removes the initial value function bias by making generated samples independent of the initial noise variable.It is defined by σ(t) = √2η_t and is both sufficient and necessary within the stated family of generative processes.

5 Adjoint Matching for control-affine stochastic optimal control

The section develops Adjoint Matching as a regression-based method for control-affine stochastic optimal control, combining adjoint calculations with least-squares objectives. It removes importance weighting and unnecessary adjoint terms to improve scalability and potentially reduce variance.

  • Adjoint Matching: Adjoint Matching combines the continuous adjoint method with least-squares objectives to optimize control-affine stochastic optimal control problems.The method is presented as a regression formulation for stochastic optimal control.
  • Comparison with prior methods: Direct differentiation through SDE simulations requires storing the numerical solver's full computational graph, creating substantial memory demands.Implementations often use gradient checkpointing to reduce memory usage.
  • Adjoint Matching: The method regresses onto the current-control target vector field, whose expected gradient matches the continuous adjoint gradient.This avoids directly constructing an importance-weighted estimator of the optimal control.
  • Adjoint Matching: The basic Adjoint Matching objective has the optimal control as its only critical point, providing theoretical grounding for gradient-based optimization.The objective can be interpreted as enforcing the optimal control's fixed-point relation.
  • Comparison with prior methods: Unlike importance-weighted matching, Adjoint Matching uses a simple least-squares objective without importance weights, avoiding high-variance weights while retaining adjoint scalability.Importance-weight variance grows exponentially with the dimension of the stochastic process, making those objectives impractical for high-dimensional applications.
  • Adjoint Matching: The lean adjoint removes zero-expectation terms from the adjoint differential equation, potentially improving convergence and lowering variance at the optimum.Removing extra terms also reduces the computational cost of evaluating the lean adjoint state.

6 Related work

Related work spans reward-based and preference-based fine-tuning, inference-time guidance, and optimal-control approaches for generative models. The paper distinguishes its reward-model setting from methods using curated preference data or specialized problem assumptions.

  • Fine-tuning from human feedback: RLHF uses reward-based fine-tuning or direct preference optimization, while DPO typically operates without an explicit reward model.Reward-based methods learn a reward model and solve a maximum-entropy reinforcement-learning problem; DPO combines preference-data use into one step.
  • Fine-tuning for diffusion models: Existing diffusion fine-tuning methods include policy-gradient, proximal-policy, KL-regularized, and stochastic-optimal-control approaches with differing theoretical and implementation requirements.Some methods require learning a value function and solving two stochastic optimal control problems.
  • Inference-time optimization methods: Inference-time methods modify generation through velocity guidance or initial-noise optimization instead of updating the base model.The paper notes that the distribution produced by heuristic velocity guidance is not well understood.
  • Optimal control in generative modeling: Optimal-control methods have been applied to ODEs, SDEs, and jump processes, while simplified generative constructions often use simulation-free matching objectives.The formalism is also used for sampling from un-normalized distributions and solving transport-related problems.
  • Conditional sampling in inverse problems: Conditional-sampling work uses pre-trained diffusion models and analytic likelihoods to target posterior distributions for inverse-problem applications.Reported applications include inpainting, super-resolution, phase retrieval, and nonlinear deblurring.
  • Evaluation setup: Table 2 compares fine-tuning methods using separate noise schedules for fine-tuning and sampling, with results averaged across three runs.The table evaluates random sets of 40,000 training prompts and 1,000 test prompts.

7 Experiments

The experiments evaluate Adjoint Matching for reward fine-tuning of a Flow Matching model, comparing it with reward-based and preference-optimization baselines across consistency, human preference, diversity, and tradeoffs. Adjoint Matching generally provides stronger and more stable tradeoffs, while classifier-free guidance increases consistency at the expense of diversity.

  • Experimental setup: The experiments fine-tune a Flow Matching base model using ImageReward and compare Adjoint Matching with DPO, ReFL, and DRaFT-K baselines.DRaFT-1 is selected for additional ablations because it performs best among the baselines.
  • Evaluation: Evaluation separates text-to-image consistency, human preference generalization, and sample diversity using ClipScore, HPS v2, and DreamSim Diversity.The metrics are designed to capture tradeoffs among these aspects of generative models.
  • Tradeoff fronts: Adjoint Matching achieves the best Pareto fronts across consistency, diversity, and generalization to unseen human preferences.The comparison varies λ for Adjoint Matching and fine-tuning iterations for DRaFT-1.
  • Main results: Across all reported metrics, the memoryless SOC formulation outperforms existing baselines, while Adjoint Matching beats continuous and discrete adjoint methods on consistency and human-preference metrics.Methods are compared at roughly the same wall-clock time, with averages over three fine-tuning runs and separate test datasets.
  • Reward tradeoff: Higher λ improves consistency and human preference but reduces diversity, while Adjoint Matching remains stable across λ values.The memoryless SOC regularization performs better than baseline procedures that often require early stopping.
  • Classifier-free guidance: Higher classifier-free guidance weight improves text-to-image consistency and sometimes human preference, but worsens diversity.The guidance analysis applies CFG after fine-tuning, although only the conditional model is fine-tuned.

8 Conclusion

The paper formulates reward fine-tuning of dynamical generative models as SOC with a memoryless noise schedule and introduces Adjoint Matching as a regression-based SOC solver. It reports improved consistency, unseen-preference generalization, and diversity relative to existing methods, including for Flow Matching fine-tuning.

  • Problem: The paper investigates reward fine-tuning for dynamical generative models such as Flow Matching.It frames the problem using stochastic optimal control.
  • Memoryless SOC: The memoryless noise schedule ensures convergence to the same tilted distribution used in human-feedback learning.For diffusion models it corresponds to DDPM sampling, while Flow Matching uses a new Memoryless Flow Matching process.
  • Adjoint Matching: Adjoint Matching casts SOC optimization as a regression problem and outperforms related gradient-based SOC methods.The method is presented as a theoretically driven approach to fine-tuning Flow Matching models.
  • Conclusion: The approach significantly outperforms baseline methods across text-to-image consistency, unseen human-preference generalization, and sample diversity.The conclusion reports these gains on large-scale text-to-image generation.

A Additional Figures & Tables

The supplementary material provides additional quantitative, qualitative, and theoretical analyses of sampling schedules, guidance, trajectories, and forward-backward SDE relationships. It documents how arbitrary noise schedules preserve marginals and trajectories while fine-tuning relies specifically on the memoryless schedule.

  • Additional evaluations: Additional experiments report reward, control cost, and consistency metrics with averages over three runs and standard-error confidence intervals.The plots evaluate ImageReward, control cost, and ClipScore over wall-clock time.
  • Qualitative samples: Figures compare classifier-free guidance samples from the pre-trained and fine-tuned Flow Matching models across varying guidance weights and selected prompts.Additional figures use fixed guidance, noise schedules, prompts, rows, and random seeds to organize qualitative comparisons.
  • Schedules and ablations: Tables compare fine-tuning and inference noise schedules, including Memoryless Flow Matching with σ(t) = √2ηt and noiseless inference with σ(t) = 0.Other supplementary tables vary guidance, optimization hyperparameters, sampling steps, and alternative schedules.
  • Schedule ablation: An alternative schedule σ(t) = 1 yields lower ImageReward and lower performance on most other metrics than memoryless fine-tuning, except for diversity.The reported explanation is initial value function bias.
  • Forward-backward SDEs: Proposition 3 establishes equality of forward and backward marginal distributions up to a time flip.The result applies at every time t.
  • Trajectory distributions: For arbitrary noise schedules, forward and backward processes also have equal trajectory distributions after reversing time.This is stronger than equality of individual-time marginals.

B.2.3 Proof of Proposition 4

The proof develops a discrete-time derivation of the forward-backward trajectory equivalence and connects SOC with KL-regularized reinforcement learning. It shows that the construction preserves the relevant process distributions and motivates the control formulation used for fine-tuning.

  • Proof status: The derivation corrects unexplained approximations in an earlier proof while remaining discrete-time rather than fully formal continuous-time analysis.The cited prior result concerns the same forward-backward process relationship.
  • Discrete-time derivation: The proof discretizes the forward process on K + 1 equally spaced timesteps and tracks terms through linear order in the stepsize.Terms of order O(h^3/2) and higher are neglected while increments scale as O(h^1/2).
  • Backward construction: Using telescoping products, discrete Ito expansions, Taylor approximations, and Bayes’ rule, the proof derives the backward transition kernels.The resulting update is identified with an Euler-Maruyama discretization.
  • Trajectory equivalence: Reversing the discretized sequence and substituting the drift relation yields the backward process with the same trajectory law up to time reversal.The proof concludes by matching the initial distribution and transition kernels.
  • Score and noise relations: The derivation relates score functions to noise predictors for DDIM and Flow Matching, supporting the conversion between diffusion coefficients and sampling formulations.The appendix explicitly derives the DDIM score-noise relation and the arbitrary-noise Flow Matching SDE.
  • SOC interpretation: The SOC formulation arises as the continuous-time limit of KL-regularized reinforcement learning with conditional Gaussian transitions.The connection also identifies control cost with KL divergence between trajectory distributions conditioned on starting points.

C.2 From maximum entropy RL to stochastic optimal control

This section establishes that continuous-time maximum-entropy reinforcement learning becomes stochastic optimal control under vector-field policies and conditional-Gaussian transitions. The correspondence transfers rewards, KL regularization, and value functions into the SOC formulation.

  • Continuous-time MaxEnt RL is equivalent to SOC as the number of discrete timesteps tends to infinity.The equivalence assumes vector-valued states and actions, vector-field policies, and conditional Gaussian transitions.
  • The discrete state sequence converges to the controlled SDE trajectory, while the base policy corresponds to the uncontrolled process.The policy uses a vector field u, whereas the base policy sets u = 0.
  • The MaxEnt RL objective maps to SOC running and terminal rewards, control cost, KL regularization, and value function.The limiting MaxEnt RL loss is identified with the SOC loss.
  • The SOC geometry permits differentiating state and terminal costs, providing losses without direct analogs in classical MaxEnt RL.This follows from the correspondence between the two optimization frameworks.
  • The trajectory distributions of controlled and base processes are related through the Girsanov theorem.The theorem supplies the change-of-measure relation underlying the SOC formulation.

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

The proof shows that a specific noise schedule is sufficient and necessary for memorylessness, removing dependence between the noise variable and generated samples. It then connects the resulting value function to fine-tuned DDIM and Flow Matching dynamics.

  • Flow Matching marginals are represented through an interpolation of independent Gaussian noise and data variables.This representation is used to derive the score function and analyze the schedule.
  • The schedule assumption is sufficient for the terminal variable to be independent of the initial noise variable.The proof uses time-reversed trajectory distributions and identifies this independence with memorylessness.
  • The same schedule condition is necessary because otherwise intermediate and terminal variables generally remain dependent.Nonzero terms in the derived expression prevent the required independence.
  • The SOC value function satisfies a Hamilton-Jacobi-Bellman equation, and its gradient determines the optimal control.The proof applies a Hopf-Cole transformation and reverses time to construct the SOC problem.
  • The resulting fine-tuned inference dynamics match the DDIM and Flow Matching forms when their predictors or vector fields equal the optimal ones.The proof explicitly establishes matches for DDIM and Flow Matching.

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

This section proves the theoretical guarantee of Adjoint Matching and situates the method among alternative fine-tuning objectives. The expected Adjoint Matching loss has the optimal control as its only critical point.

  • The basic Adjoint Matching loss has a unique critical point because its stationarity condition implies the Hamilton-Jacobi-Bellman equation.Lemma 6 identifies any control satisfying the relevant relation with the optimal control.
  • The only critical point of the expected Adjoint Matching loss is the optimal control.This is the central theoretical guarantee stated in Proposition 7.
  • The full Adjoint Matching loss inherits this guarantee because every critical point of the full loss is also a critical point of the basic loss.The uniqueness of the associated integral equation completes the argument.
  • The method is implemented by computing trajectories, solving the lean adjoint ODE backward, evaluating the objective, and updating model parameters.Algorithm 2 describes this procedure for fine-tuning DDIM.
  • The paper also adapts reward-based and preference-based diffusion objectives, including ReFL and Diffusion-DPO, to Flow Matching.The adaptations rewrite denoising or preference objectives using the Flow Matching vector field.
  • On-policy preference pairs differ from curated off-policy pairs, and the resulting DPO model performs like or slightly worse than the base model on some metrics.The authors conclude that DPO is not competitive for on-policy fine-tuning when the base model is not already good.

G Experimental details

The experiments use common optimization and data settings while adding practical adjustments for numerical stability, gradient quality, and efficient timestep sampling. Adjoint Matching has similar iteration cost to Discrete Adjoint and lower cost than Continuous Adjoint.

  • The experiments use K = 40 timesteps, learning rate 2 × 10^-5, effective batch size 40, and 40k prompts per run.Discrete Adjoint uses learning rate 1 × 10^-5 because the common setting caused instability.
  • The implementation offsets the diffusion coefficient near both endpoints to avoid infinite or near-zero values.The denominator receives a small offset at t = 0, while the numerator is changed from 1 − t to 1 − t + h.
  • The final adjoint state is initialized using a final noiseless update to remove noise-induced gradient bias.The authors note that this modification is consistent with the continuous-time limit.
  • Only a subset of early timesteps is sampled, while the final ten timesteps are always included because their fine-tuning is critical for empirical performance.The subset contains ten uniformly sampled early timesteps and the last ten timesteps.
  • Adjoint Matching uses LCT = 1.6 × λ^2 for loss clipping, with the threshold tuned to preserve more late-step terms than early-step terms.The constant is selected experimentally, and the threshold must scale quadratically with λ.
  • Adjoint Matching takes 156 seconds per iteration, close to Discrete Adjoint at 152 seconds and below Continuous Adjoint at 204 seconds.The cost difference reflects additional backward passes required by Continuous Adjoint.
  • Using 100 or 200 fine-tuning timesteps yields statistically equal metrics to 40 timesteps, whereas 10 or 20 timesteps perform worse.At 10 timesteps, Adjoint Matching remains clearly better than DRaFT-1 only on diversity.
Loading 2409.08861v5…