Source-linked AI summary
Reinforce Adjoint Matching: Scaling RL Post-Training of Diffusion and Flow-Matching Models
Andreas Bergmeister, Stefanie Jegelka, Nikolas Nüsken, Carles Domingo-Enrich, Jakiw Pidstrigach
TL;DR
Diffusion models’ supervised-regression structure does not by itself address reward optimization in post-training. RAM derives a reward-corrected regression target and achieves the highest reward on three image-generation tasks, matching Flow-GRPO’s peak reward in up to 50× fewer training steps.
Problem
Diffusion and flow-matching pretraining uses supervised regression, while post-training seeks to optimize generation toward a reward rather than merely match data.
Source: §1 · p. 2Method
Results
RAM achieves the highest training reward on all three tasks and matches Flow-GRPO’s peak reward in 50×, 48×, and 34× fewer training steps.
Source: §1 · p. 2Takeaways & Limitations
The simple regression objective that scales pretraining also scales post-training, without reward gradients or SDE rollouts.
Source: §1 · p. 2RAM drops the path-cost correction; retaining it through random-time estimation has high variance, while trajectory-based differentiation trades variance against compute.
Source: §5 · p. 7
Abstract
from arXiv · showhide
Diffusion and flow-matching models scale because pretraining is supervised regression: a clean sample is noised analytically, and a model regresses against a closed-form target. RL post-training aligns the model with a reward. In image generation, this makes samples compose objects correctly, render text legibly, and match human preferences. Existing methods rely on costly SDE rollouts, reward gradients, or surrogate losses, sacrificing pretraining's regression structure. We show that the structure extends to RL post-training. Under KL-regularized reward maximization, the optimal generative process tilts the clean-endpoint distribution towards samples with higher reward and leaves the noising law unchanged. Combining this with the adjoint-matching optimality condition and a REINFORCE identity, we derive Reinforce Adjoint Matching (RAM): a consistency loss that corrects the pretraining target with the reward. At each step, we draw a clean endpoint from the current model, evaluate its reward, noise it as in pretraining, and regress. No SDE rollouts, backward adjoint sweeps, or reward gradients are required. Like the pretraining objective, RAM is simple and scales. On Stable Diffusion 3.5M, RAM achieves the highest reward on composability, text rendering, and human preference, reaching Flow-GRPO's peak reward in up to $50\times$ fewer training steps.
1 Introduction
The paper extends diffusion and flow-matching models’ supervised-regression structure to reward-based post-training. RAM uses a reward-corrected regression target and achieves the highest reward on three image-generation tasks, matching Flow-GRPO’s peak in fewer training steps.
1 Introduction
- Pretraining scales through supervised regression: analytically noise a clean data sample, then regress against a closed-form target.Reward-based post-training can align image generation with complex prompts, legible text, and human preferences.Source: §1 · p. 2
- Existing diffusion and flow-matching post-training methods use costly SDE rollouts, reward gradients, or surrogate objectives because sample likelihoods are intractable.Policy-gradient methods use Gaussian transition densities to optimize a variational lower bound, while stochastic optimal-control estimators backpropagate reward gradients through an adjoint ODE.Source: §1 · p. 2
- RAM derives a reward-corrected regression target from an optimal-control fixed point and a REINFORCE identity, then samples endpoints, evaluates rewards, and analytically noises them.The optimal process changes which clean samples are likely but preserves the noising rule; RAM needs no SDE rollouts, adjoint sweeps, or reward gradients.Source: §1 · p. 2
- RAM achieves the highest reward on GenEval, OCR, and PickScore, matching Flow-GRPO’s peak in 50×, 48×, and 34× fewer training steps, respectively.The comparisons are reported for Stable Diffusion 3.5M; per-step compute cost is slightly lower for RAM.Source: §1 · p. 2
2 Diffusion and Flow-Matching Models
Diffusion and flow-matching models use a forward process that corrupts data into noise and a backward process that generates samples. Their shared training structure constructs noisy states analytically and learns a velocity field by regression.
2 Diffusion and Flow-Matching Models
- Diffusion models specify forward and backward processes explicitly, whereas flow-matching models specify the noising kernel and generate by integrating a velocity field.The associated forward and backward SDEs are well-defined; the probability-flow ODE is the default flow-matching sampler described here.Source: §2 · pp. 2–3
- The forward process interpolates from data at t=0 to a Gaussian prior at t=1, yielding a noising kernel for a clean sample.The paper uses linear interpolation, while noting that nonlinear schedules recover alternatives such as DDPM and that its results generalize.Source: §2 · p. 2
- Velocity matching learns the velocity field as a conditional expectation and regresses against a random target constructed from clean data and independent Gaussian noise.Because the noisy state is an affine function of the clean sample and noise, each training pair is cheap to construct.Source: §2 · p. 3
- Predicting clean data or noise instead of velocity is equivalent because any two of the velocity, clean data, and noise determine the third.Source: §2 · p. 3
3 RL Post-Training as Optimal Control
The paper formulates KL-regularized reward post-training as stochastic optimal control of a pretrained generative process. Its optimal process tilts the clean-endpoint distribution while preserving the conditional noising law, and adjoint matching expresses optimality as a fixed-point condition.
3 RL Post-Training as Optimal Control
- The KL term keeps the post-trained model close to the pretrained reference and within its support, while reward shifts mass toward desirable samples.The reward is scaled to control regularization strength.Source: §3 · p. 3
- The paper casts KL-regularized reward maximization as stochastic optimal control, steering the pretrained backward SDE with a drift correction while penalizing control cost.The path-space control problem is equivalent to KL-regularized reward maximization by Girsanov’s theorem.Source: §3 · p. 3
- The optimal controlled process tilts the clean-endpoint distribution toward higher-reward samples while leaving the conditional law of noisy states given the endpoint unchanged.The characterization is stated under standard regularity assumptions.Source: §3 · p. 4
- The unchanged conditional noising law means the optimal process uses the same analytic Gaussian kernel as pretraining; only the endpoint distribution changes.The conditional trajectory law given a clean endpoint depends on the interpolation schedule, not on the data distribution.Source: §3 · p. 4
- Adjoint matching makes optimality a self-consistency condition on the control, with the adjoint defined as the spatial gradient of a Bellman value function.The condition holds under standard regularity assumptions.Source: §3 · p. 4
4 Reinforce Adjoint Matching
Reinforce Adjoint Matching (RAM) turns the adjoint optimality condition into a reward-corrected regression objective, using endpoint sampling and analytic noising rather than SDE rollouts. Its fixed point agrees with the KL-regularized optimum to first order in reward, and is exact under specified linearity conditions.
4 Reinforce Adjoint Matching
- RAM uses a REINFORCE identity to replace the reward gradient with a reward-weighted backward-bridge score in a regression target for the velocity field.This avoids requiring a differentiable reward function.Source: §4 · p. 4
- RAM drops the path-cost correction because retaining it entails either a high-variance estimate or costly differentiation along a controlled SDE rollout.The approximation coincides with the full adjoint at initialization and is exact for a Gaussian reference under a linear reward.Source: §4 · pp. 4–5
- At the optimum, the bridge from a clean endpoint to a noisy state follows the pretraining noising kernel, enabling ODE-sampled endpoints with analytic noising.Each endpoint can yield K conditionally independent training states from one model sample and reward query.Source: §4 · p. 5
4.1 Relating the RAM fixed point to the KL optimum
- The KL-optimal velocity can be characterized along a path of exponentially tilted endpoint distributions, with a conditional reward-velocity covariance as the path integrand.RAM uses a single on-policy evaluation at the path endpoint in place of averaging that covariance over the path.Source: §4.1 · pp. 6–7
- RAM and the KL optimum agree to first order in reward; the rule is exact when velocity varies linearly along the tilt path, including Gaussian references with linear rewards.The first-order result follows by scaling reward by η and expanding around η = 0.Source: §4.1 · p. 7
5 Retaining the Path-Cost Correction
The paper examines estimators that retain the path-cost correction, trading variance against computation, and compares them with RAM on a 2D example. All estimators recover the tilted density there, while RAM has the lowest variance overall and is the only estimator reported to scale to text-to-image post-training.
5 Retaining the Path-Cost Correction
- Retaining the path-cost correction trades a high-variance random-time estimate against controlled-SDE simulation with pathwise differentiation.The length of pathwise differentiation trades variance against compute.Source: §5 · p. 7
- A generalized adjoint splits into a REINFORCE term over the prefix and a pathwise gradient over the suffix; practical choices are full-horizon or local, while intermediate choices require separate sweeps per target time.Full-horizon computes targets for every time in one backward sweep; local avoids the sweep by collapsing the suffix to one discretization step.Source: §5 · pp. 7–8
5.1 Estimator variance on a 2D example
- All estimators recover the tilted density, but RAM has the lowest variance overall and is the only estimator that scales to text-to-image post-training.At convergence, regression loss equals the variance of the control-space target.Source: Figure 3 · p. 8, §5.1 · p. 8
- Among path-cost-corrected estimators, full-horizon has substantially lower variance than local or random-jump, and the Bayes bridge score beats Malliavin within full-horizon.The comparison uses the regression loss at convergence as a readout of estimator variance.Source: §5.1 · p. 8
6 Text-to-Image Experiments
On Stable Diffusion 3.5M, RAM achieves the highest training reward on all three tasks, with image-quality trade-offs varying by task. It also reaches Flow-GRPO’s peak rewards in substantially fewer steps.
Experimental setup
- The evaluation post-trains Stable Diffusion 3.5 Medium separately for GenEval compositional correctness, OCR text rendering, and PickScore human preference.Models are evaluated on held-out benchmark prompts, with an independent prompt set used to check for reward hacking; DrawBench prompts assess image quality.Source: §6 · pp. 8–9
Training efficiency
- With classifier-free guidance, RAM remains compatible after post-training, while CFG lowers the task reward of DiffusionNFT and AWM.Figure 4 compares training-reward curves for RAM and the two efficiency-oriented baselines, including lighter CFG re-evaluations of the baselines.Source: Figure 4 · p. 9, §6 · p. 9
Results
- RAM achieves the highest training reward on all three tasks, with comparable or better image quality for compositional generation and preference alignment.For visual text rendering, RAM avoids the severe quality collapse of DiffusionNFT and AWM while matching or exceeding their OCR reward; Flow-GRPO has higher image-quality metrics but lower OCR reward.Source: §6 · p. 9, Table 2 · p. 10
Training efficiency
- Roughly 50× fewer steps on GenEval (0.95), 48× fewer on OCR (0.92), and 34× fewer on PickScore (23.31) let RAM match Flow-GRPO’s peak rewards.Per-step compute is comparable, so the step ratios translate directly to wall-clock training time; RAM noises each endpoint K = 8 times with independent noise vectors.Source: §6 · p. 9
7 Related Work
Related work places RAM at the intersection of pretraining-like regression and adjoint methods, while the appendix characterizes the control objective through path-space KL regularization and a data-independent conditional path law.
7 Related Work
- Prior diffusion and flow-matching work establishes regression on analytically noised data, while RAM reuses that structure with a target built from the current model and scalar reward.Earlier policy-gradient approaches learn from stochastic rollouts; RAM replaces the rollout with regression against a closed-form target.Source: §7 · p. 10
- Adjoint Matching and related control methods target KL-regularized objectives but require reward gradients and a backward adjoint ODE; RAM instead uses a closed-form Bayes bridge target.Tilt Matching reaches the same velocity-change equation by annealing through intermediate models, whereas RAM targets the fixed point directly.Source: §7 · p. 11
- DiffusionNFT uses positive/negative guidance and AWM uses a flow-matching ELBO, whereas RAM derives its update from KL-regularized control.Differentiable-reward methods backpropagate through the sampling chain; bridge estimators connect to RAM’s path-cost-corrected estimators, while its main objective uses a cheaper Bayes bridge approximation.Source: §7 · p. 11
- Extending RAM’s gradient-free regression structure to flow maps is identified as future work.Existing flow-map reward-alignment methods cited here require reward gradients.Source: §7 · p. 11
A.1 Optimal control as path-space KL regularization
- Under Novikov’s condition and a finite exponential reward moment, the stochastic control problem is equivalent to a Gibbs variational problem over absolutely continuous path laws.The stated path-law tilt has density proportional to the exponential of the endpoint reward.Source: §A.1 · p. 15
A.2 Path-law factorization
- The conditional reference trajectory law given a clean endpoint is independent of the data distribution, yielding the path-space factorization behind the memoryless property.At t = 1, the noising kernel is standard Gaussian regardless of the clean endpoint, so the source marginal is also data-independent.Source: §A.2 · p. 15
A.3 Optimal controlled process
The optimal controlled process changes the clean-endpoint distribution according to the reward while preserving the conditional path law. Its time marginals follow by applying the forward noising process to the tilted endpoint distribution.
A.3 Optimal controlled process
- The optimal path law tilts the reference clean-endpoint distribution by the exponential reward while leaving the conditional path law given that endpoint unchanged.The resulting process is equivalently the time reversal of forward noising started from the tilted endpoint distribution, which determines its time-t marginals.Source: §A.3 · pp. 15–16
A.4 Optimal velocity and score identities
- The appendix derives optimal velocity and score identities by applying the standard velocity–score relation to the optimal time marginals, and relates the control to the score difference.The derivation compares the reference and optimal drift decompositions.Source: §A.4 · p. 16
A.5 Verification and self-consistency
The verification result characterizes optimal controls through a self-consistency condition, which is both necessary and sufficient under the stated regularity and Markov assumptions. The proof connects the condition to the HJB equation.
A.5 Verification and self-consistency
- The self-consistency condition (12) is necessary and sufficient for optimality, allowing a candidate control to be certified by checking one PDE.The verification theorem assumes a unique C1,2 HJB solution and an optimal Markov feedback control.Source: §A.5 · p. 16
- The paper includes the proof as standard HJB and adjoint-matching theory to keep the argument self-contained.Source: §A.5 · p. 16
- Self-consistency turns the linear value-function PDE into the HJB equation that singles out the optimum.The proof uses a linear PDE valid for any admissible control and the nonlinear HJB equation for the optimum.Source: §A.5 · p. 17
- Because the HJB equation has a unique solution, a control satisfying the condition is optimal.At the fixed point, substitution yields the HJB equation and identifies the value function with the optimal value function.Source: §A.5 · p. 17
B Derivations for RAM
The derivations establish the Bayes bridge score and its role in RAM’s control-space loss, then characterize the optimum along a path of tilted endpoint distributions. In a Gaussian-reference, linear-reward case, the RAM target exactly matches the optimal adjoint.
B.1 Bayes bridge score derivation
- The Bayes bridge score is used by both the jump estimator and the pathwise family, with a simplified form when s=0.For s=0, the derivation sets a_t,0 = 1−t and β_t,0 = t.Source: §B.1 · p. 17
B.2 Control-space form of RAM and time-dependent weighting
- The bridge-score prefactor combines with the control relation to give a timestep weighting that downweights times near the source.Dropping that factor yields the uniformly weighted velocity loss used in the main text.Source: §B.2 · p. 18
B.3 Path-integral characterization of the optimum
- Differentiating the conditional endpoint density with respect to the tilt parameter gives a reward-centered conditional term, which is applied to f in the derivation.The derivation applies this identity with f = ϵ−X0.Source: §B.3 · p. 18
B.4 Exact case: Gaussian reference under a linear reward
- With a Gaussian reference and linear reward, the tilted endpoint remains Gaussian, with its mean shifting linearly in the tilt and covariance unchanged.In this case the integrand is constant in the tilt parameter, so the right-endpoint rule is exact and the RAM target matches the optimal adjoint at every x.Source: §B.4 · p. 18
- In the same exact case, the optimal control is independent of position, recovering the path-cost-vanishing instance.Source: §B.4 · p. 19
C Derivations for Retaining the Path-Cost Correction
These derivations develop estimators that retain the path-cost correction, from random-jump and bridge-score identities to local and full-horizon procedures. They also specify the approximations and computation used by the variants.
C.1 Random-jump identity and the jump estimator
- The random-jump estimator combines a single-point control-cost evaluation with a Bayes bridge score, using an unbiased single-point cost estimate.The derivation replaces a controlled-process sample with a two-step construction and substitutes the Bayes bridge score.Source: §C.1 · p. 19
C.2 Generalized adjoint identity
- The generalized adjoint identity expresses the value gradient through a bridge-score term and a suffix path-cost gradient.The estimator uses one SDE rollout to estimate the identity’s right-hand side.Source: §C.2 · p. 20
C.3 Malliavin bridge-score estimator
- The Malliavin bridge-score estimator can be accumulated using an equivalent reverse-mode adjoint SDE, requiring vector–Jacobian products rather than explicit Jacobians.The full-horizon specialization produces endpoint-indexed accumulators in a single backward sweep.Source: §C.3 · pp. 20–21
C.4 Local one-step Gaussian score
- The local one-step Gaussian bridge score differs from the exact continuous-time score by the O(δ) error already incurred by Euler–Maruyama.It is derived from a single Euler–Maruyama step from t to t−δ.Source: §C.4 · p. 21
C.5 Cancellation of ∇xut terms in the local estimator
- At the fixed point, score and path-cost contributions involving ∇xut cancel to first order in δ, yielding the local target.Source: §C.5 · p. 21
C.6 Discrete-time full-horizon training procedure
- The full-horizon procedure rolls out the controlled SDE once, then computes targets for all interior grid points in a discrete reverse-mode sweep.The grid endpoint tK remains below 1 because the integrated variance diverges at the noise endpoint.Source: §C.6 · p. 22
D Experimental Details
The experiments post-train SD3.5M with LoRA and compare RAM against baselines using a shared prompt count per step. The section reports model, hardware, sampling, optimizer, averaging, and reward-scaling settings.
D Experimental Details
- SD3.5M is post-trained with rank-32 LoRA on four NVIDIA H100 GPUs with 96 GB memory each.Source: D Experimental Details · p. 23
- Each update uses 48 prompts, 24 samples per prompt, and 8 noisy training samples per image, for an effective batch of 9,216.Baselines use 48 prompts per step for fair comparison; other baseline hyperparameters and implementations are retained.Source: D Experimental Details · p. 23
- Training uses 20-step Euler sampling and classifier-free guidance scale 2.0, while evaluation uses 40 steps and task-dependent guidance scales.Evaluation guidance is 4.5 for GenEval and OCR and 2.0 for PickScore; training samples timesteps from p(t)=2t, which the passage says enhances stability over uniform sampling.Source: D Experimental Details · pp. 23–24
- The setup uses AdamW with learning rate 3e−4 and β2=0.95, EMA decay 0.9, and task-dependent reward coefficients of 100 or 1000.The reward coefficients are 100 for GenEval and OCR, and 1000 for PickScore.Source: D Experimental Details · p. 24