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

2026-05-11

AI Summary

This work proposes a reinforcement learning (RL) post-training method for diffusion and flow matching models that preserves the efficient regression structure established during pretraining. Unlike existing approaches that rely on costly SDE trajectories, reward gradients, or surrogate losses—thereby disrupting this structure—the proposed method performs KL-regularized reward maximization by adjusting only the clean endpoint distribution while leaving the noise perturbation mechanism unchanged. It introduces a consistency loss derived from adjoint matching conditions and the REINFORCE identity, requiring merely a single-step noising operation and reward evaluation. This approach naturally extends the supervised regression framework to RL alignment without necessitating SDE rollouts, adjoint backpropagation, or reward gradients. Evaluated on Stable Diffusion 3.5M, it achieves state-of-the-art rewards on compositionality, text rendering, and human preference tasks, with up to 50× fewer training steps compared to Flow-GRPO.

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.

Problem

• Reinforcement Learning
• Diffusion Models
• Flow Matching
• Post-Training
• Reward Alignment

Innovation

• Reinforce Adjoint Matching
• RL post-training
• diffusion models
• flow-matching
• consistency loss