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

• URL: http://arxiv.org/abs/2605.10759v2
• Date: Sun, 17 May 2026 07:26:07 GMT
• Title: Reinforce Adjoint Matching: Scaling RL Post-Training of Diffusion and Flow-Matching Models
• Abstract summary: Diffusion and flow-matching models scale because pretraining is supervised regression.<n>We show that the structure extends to RL post-training.<n>Under KL-regularized reward, the optimal generative process tilts the clean-endpoint distribution towards samples with higher reward.<n>We derive Reinforce Adjoint Matching (RAM): a consistency loss that corrects the pretraining target with the reward.
• Score: 44.082049677105964
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• MARBLE: Multi-Aspect Reward Balance for Diffusion RL [71.6241143519038]
  Reinforcement learning has become the dominant approach for aligning diffusion models with human preferences.<n>Existing practice deal with multiple rewards by training one specialist model per reward.<n>We propose MARBLE, a gradient-space optimization framework that maintains independent advantage estimators for each reward.
  arXiv  Detail & Related papers  (2026-05-07T16:20:42Z)
• Diffusion Reinforcement Learning via Centered Reward Distillation [35.979608265594685]
  We present textbf Reward Distillation (CRD), a diffusion RL framework derived from KL-regularized reward models built on forward-process fine-tuning.<n>To enable reliable text-to-image fine-tuning, we introduce techniques that explicitly control distribution drift.<n>Experiments on text-to-image post-training with textttGenEval and textttOCR rewards show that competitive SOTA reward optimization results with fast convergence and reduced reward hacking preference.
  arXiv  Detail & Related papers  (2026-03-14T21:29:33Z)
• GDRO: Group-level Reward Post-training Suitable for Diffusion Models [55.948229011478304]
  Group-level rewards successfully align the model with the targeted reward.<n>Group-level Direct Reward Optimization (GDRO) is a new post-training paradigm for group-level reward alignment.<n>GDRO supports full offline training that saves the large time cost for image rollout sampling.<n>It is diffusion-sampler-independent, which eliminates the need for the ODE-to-SDE approximation to obtainity.
  arXiv  Detail & Related papers  (2026-01-05T11:47:18Z)
• Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models [35.36024202299119]
  We introduce textbfAdvantage Weighted Matching (AWM), a policy-gradient method for diffusion.<n>AWM raises the influence of high-reward samples and suppresses low-reward ones while keeping the modeling objective identical to pretraining.<n>This unifies pretraining and RL conceptually and practically, is consistent with policy-gradient theory, reduces variance, and yields faster convergence.
  arXiv  Detail & Related papers  (2025-09-29T17:02:20Z)
• Inference-Time Scaling of Diffusion Language Models with Particle Gibbs Sampling [70.8832906871441]
  We study how to steer generation toward desired rewards without retraining the models.<n>Prior methods typically resample or filter within a single denoising trajectory, optimizing rewards step-by-step without trajectory-level refinement.<n>We introduce particle Gibbs sampling for diffusion language models (PG-DLM), a novel inference-time algorithm enabling trajectory-level refinement while preserving generation perplexity.
  arXiv  Detail & Related papers  (2025-07-11T08:00:47Z)
• Flow-GRPO: Training Flow Matching Models via Online RL [80.62659379624867]
  We propose Flow-GRPO, the first method to integrate online policy reinforcement learning into flow matching models.<n>Our approach uses two key strategies: (1) an ODE-to-SDE conversion that transforms a deterministic Ordinary Differential Equation into an equivalent Differential Equation (SDE) that matches the original model's marginal distribution at all timesteps; and (2) a Denoising Reduction strategy that reduces training denoising steps while retaining the original number of inference steps.
  arXiv  Detail & Related papers  (2025-05-08T17:58:45Z)
• RED: Unleashing Token-Level Rewards from Holistic Feedback via Reward Redistribution [50.171320156632866]
  Reinforcement learning from human feedback offers a promising approach to aligning large language models with human preferences.<n>Current reward models operate as sequence-to-one models, allocating a single, sparse, and delayed reward to an entire output sequence.<n>We propose a more fine-grained, token-level guidance approach for RL training.
  arXiv  Detail & Related papers  (2024-11-13T02:45:21Z)
• PRDP: Proximal Reward Difference Prediction for Large-Scale Reward Finetuning of Diffusion Models [13.313186665410486]
  Reward finetuning has emerged as a promising approach to aligning foundation models with downstream objectives. Existing reward finetuning methods are limited by their instability in large-scale prompt datasets. We propose Proximal Reward Difference Prediction (PRDP) to enable stable black-box reward finetuning for diffusion models.
  arXiv  Detail & Related papers  (2024-02-13T18:58:16Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.