Reinforce Adjoint Matching: Scaling RL Post-Training of Diffusion and Flow-Matching Models

Abstract: 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.

• The paper proposes a reinforcement learning post-training method that leverages analytic noise and closed-form REINFORCE-adjoint estimation for diffusion models.
• It achieves superior reward alignment with up to 50× fewer training steps while maintaining compatibility with classifier-free guidance.
• RAM simplifies optimal control by transforming costly rollouts into scalable regression tasks, enhancing both theoretical insights and practical model performance.

Reinforce Adjoint Matching: Scaling RL Post-Training for Diffusion and Flow-Matching Models

Introduction

Diffusion and flow-matching models have emerged as foundational tools for high-fidelity generative modeling in continuous domains. Their scalability is attributable to the simple supervised regression paradigm used during pretraining: analytic corruption of clean samples with Gaussian noise, followed by minimization of a closed-form regression target. Yet, to align these powerful generators with nuanced downstream desiderata—such as compositional semantics, text fidelity, or human preference rewards—reinforcement learning (RL) post-training is essential. Traditional RL post-training for these models faces challenges due to intractable sample likelihoods and high computational cost for reward alignment, especially compared to policy-gradient approaches for models with tractable likelihoods.

The "Reinforce Adjoint Matching" (RAM) method introduces a principled, scalable approach for RL post-training of diffusion and flow-matching models. RAM provides a consistency loss that extends the analytic and supervised structure of pretraining to the RL setting by leveraging a fixed-point characterization of the KL-regularized reward maximization optimum, sidestepping computational bottlenecks of SDE rollouts, reward gradient backpropagation, or surrogate losses.

Methodology

Problem Formulation and Optimality Characterization

Given a pretrained model with endpoint distribution and reward function $r(x)$, the canonical RL post-training objective is to maximize expected reward while constraining posterior drift by the KL-divergence from the reference distribution. This induces the solution: $p_0^* \propto p_0^{\mathrm{ref}}(x)\exp(r(x))$ where the optimal process tilts the endpoint distribution towards higher-reward samples but retains the conditional noising law unchanged.

This leads to a stochastic optimal control problem, steering the backward SDE with a drift correction $u_t$, and introduces a value (Bellman) function $V_t^u(x)$ whose spatial gradient defines the adjoint $A_t^u(x)$. The optimality condition can be recast as a "self-consistency" equation on the control, whereby the control must match the negative scaled adjoint: $u_t(x) = -\sigma_t A_t^u(x)$ Practically, velocity regression is preferred over direct control parametrization—a distinction addressed in the RAM formulation.

Reinforce Adjoint Matching (RAM) Algorithm

RAM exploits the structure whereby post-training modifies the endpoint distribution but not the analytic noising law. Training proceeds as follows:

• Draw an endpoint sample using an ODE sampler from the current model.
• Evaluate its reward under the task reward function.
• Analytically noise the endpoint to produce intermediate training states, following the same distribution used during pretraining.
• Form batches by re-noising each endpoint multiple times, thus amortizing reward evaluation.

The value function's adjoint is estimated using a REINFORCE (score function) identity, yielding an estimator that avoids the need for reward gradients (even with non-differentiable $r$): $$A_t^u(x) \approx \mathbb{E}[r(X_0^u) \nabla_x\log p_{0|t}^u(X_0^u\mid x)\mid X_t^u=x]$$ Using Bayes' rule and model velocity, the bridge score admits a closed-form: $$\nabla_{x_t} \log p_{0|t}(x_0|x_t) = \frac{1-t}{t}(v_t(x_t) - (\epsilon - x_0))$$ RAM trains the velocity field via a consistency loss: $$\mathcal{L}_\mathrm{RAM}(\theta) = \mathbb{E}\|\; v_t^\theta(X_t) - \mathrm{sg}\Big(v_t^{\mathrm{ref}}(X_t) + r(X_0)((\epsilon - X_0) - v_t^\theta(X_t))\Big)\; \|^2$$ where $p_0^* \propto p_0^{\mathrm{ref}}(x)\exp(r(x))$0 denotes stop-gradient.

Figure 1: RAM training flow: ODE-sampled endpoint is rewarded, noised analytically to produce batched training pairs, bypassing SDE rollouts.

Correction Terms and Scalability

RAM omits the path-cost correction in the adjoint (justified by its variance-cost tradeoff and analytic tractability), but the paper discusses generalized estimators that retain the correction at the expense of computation and stability. Empirical analysis on toy examples demonstrates that RAM's estimator achieves the lowest variance and hence best scalability for high-dimensional tasks.

Figure 2: Path-cost-corrected estimators on a 2D toy: variance at convergence is lowest for RAM, confirming its suitability for image-scale tasks.

Experimental Results

Benchmarks and Metrics

RAM was evaluated against Flow-GRPO, DiffusionNFT, and AWM on Stable Diffusion 3.5M (SD3.5M) across three prompt-aligned image-generation tasks:

• GenEval: compositional correctness (object relationships)
• OCR: text rendering legibility
• PickScore: human preference alignment

Training reward and off-the-shelf image quality metrics (Aesthetic, DeQA, ImageReward, HPSv2, PickScore) were reported both to measure task performance and to detect reward hacking.

Performance and Efficiency

RAM consistently achieves the highest training reward on all three tasks and matches or improves image quality metrics compared to strong baselines. Notably, RAM matches Flow-GRPO's peak reward in up to 50× fewer training steps and with similar per-step compute. Qualitative samples confirm improved compositionality, OCR fidelity, and preference alignment.

Importantly, post-trained RAM remains compatible with classifier-free guidance, unlike baselines where guidance can lower reward—a significant practical consideration.

Figure 3: Training-reward curves: RAM achieves faster and higher task reward compared to DiffusionNFT and AWM, and remains compatible with classifier-free guidance post-training.

Theoretical and Practical Implications

RAM's derivation connects RL post-training for diffusion models to classical entropy-regularized stochastic optimal control, with a path-space factorization that preserves the analytic noising law. This both generalizes and simplifies prior work, recasting the expensive or unstable optimization pipelines of previous approaches as scalable regression tasks using closed-form, unbiased targets.

Practically, RAM's amortized endpoint batching and pretraining-like consistency loss allow for orders-of-magnitude greater learning signal per step, making high-fidelity reward alignment with limited compute feasible for image-scale flows.

Theoretically, RAM establishes that structure-preserving analytic noising suffices for RL alignment under a broad class of objectives—ensuring identical fixed points to more complex rollout-based estimators when regularization and reward structure align (e.g., Gaussian priors and linear rewards).

Future Directions

• Non-Gaussian and Implicit Priors: Extending RAM to settings with complex priors or data-dependent noising kernels.
• Preference Learning beyond Scalar Rewards: Incorporating pairwise and more complex human preference structures in the RL objective.
• Efficient Path-Cost Correction: Developing low-variance scalable estimators for the path-cost correction may further refine post-training alignment.
• Applications to Video and Multimodal Generation: Deploying RAM for temporally extended or multimodal generative models.

Conclusion

Reinforce Adjoint Matching provides a theoretically grounded, scalable, and computationally efficient approach to RL post-training of diffusion and flow-matching models. By leveraging the analytic structure of pretraining and a closed-form REINFORCE-adjoint estimator, RAM enables high-quality alignment to downstream rewards with minimal compute overhead and without reward model hacks or loss of image quality. This framework stands to streamline further alignment advances for deep generative models in vision and beyond.