Reinforce Adjoint Matching: Scaling RL Post-Training of Diffusion and Flow-Matching Models
For practitioners of generative image models, RAM provides a simple, scalable RL fine-tuning method that preserves pretraining's regression structure and dramatically reduces training cost.
Reinforce Adjoint Matching (RAM) enables RL post-training of diffusion and flow-matching models without costly SDE rollouts or reward gradients, achieving up to 50× faster convergence to peak reward on Stable Diffusion 3.5M while improving composability, text rendering, and human preference.
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.