Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models

Reinforcement Learning (RL) has emerged as a central paradigm for advancing Large Language Models (LLMs), where both pre-training and RL post-training stages are grounded in the same log-likelihood formulation. In contrast, recent RL approaches for diffusion models, most notably Denoising Diffusion Policy Optimization (DDPO), optimize an objective different from the pretraining objectives--score/flow matching loss. In this work, we establish a novel theoretical analysis: DDPO is an implicit form of score/flow matching with noisy targets, which increases variance and slows convergence. Building on this analysis, we introduce \textbf{Advantage Weighted Matching (AWM)}, a policy-gradient method for diffusion. It uses the score/flow-matching loss and reweights each sample by its advantage. In effect, AWM raises the influence of high-reward samples and suppresses low-reward ones while keeping the modeling objective identical to pretraining. This simple yet effective design yields substantial benefits: on the GenEval, OCR, and PickScore benchmarks, AWM delivers up to a $\mathbf{34}\times$ speedup over Flow-GRPO (which builds on DDPO), when applied to Stable Diffusion 3.5 Medium and FLUX, without compromising generation quality. Code is available at \url{https://github.com/scxue/advantage_weighted_matching}.

Many image-generating AI systems work by gradually turning noise into a picture. Researchers often improve these systems by giving feedback on the final image, such as whether it follows the prompt, draws readable text, or looks appealing. However, current training methods use this feedback through many small intermediate steps, which can make learning noisy, slow, and expensive. We study why this happens and propose Advantage Weighted Matching (AWM), a simpler way to train these models from feedback. Instead of focusing on every intermediate step, AWM gives more weight to final images that receive high scores and less weight to those that receive low scores. This keeps the additional training closer to the way these image generators were originally trained. In experiments with popular open-source image models, AWM improved instruction following, text rendering, and preference scores while using much less training time than a strong existing method. This can make it cheaper and more practical to build image-generation systems that better match what people ask for.