Voordeelgewogen Matching: Afstemming van Reinforcement Learning met Vooraf Trainen in Diffusiemodellen

Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models

September 29, 2025

Auteurs: Shuchen Xue, Chongjian Ge, Shilong Zhang, Yichen Li, Zhi-Ming Ma

cs.AI

Samenvatting

Reinforcement Learning (RL) is naar voren gekomen als een centraal paradigma voor het bevorderen van Large Language Models (LLMs), waarbij pre-training en RL post-training dezelfde log-likelihood-formulering delen. In tegenstelling hiermee optimaliseren recente RL-benaderingen voor diffusiemodellen, met name Denoising Diffusion Policy Optimization (DDPO), een doelstelling die verschilt van de pre-trainingsdoelstellingen—de score/flow matching loss. In dit werk presenteren we een nieuwe theoretische analyse: DDPO is een impliciete vorm van score/flow matching met ruisdoelen, wat de variantie verhoogt en de convergentie vertraagt. Op basis van deze analyse introduceren we Advantage Weighted Matching (AWM), een policy-gradient methode voor diffusie. Deze gebruikt dezelfde score/flow-matching loss als pre-training om een doelstelling met lagere variantie te verkrijgen en herweegt elk monster op basis van zijn voordeel. In feite verhoogt AWM de invloed van monsters met hoge beloning en onderdrukt die met lage beloning, terwijl de modelleerdoelstelling identiek blijft aan pre-training. Dit verenigt pre-training en RL zowel conceptueel als praktisch, is consistent met policy-gradient theorie, vermindert variantie en resulteert in snellere convergentie. Dit eenvoudige maar effectieve ontwerp levert aanzienlijke voordelen op: op de GenEval, OCR en PickScore benchmarks biedt AWM een versnelling tot 24 keer ten opzichte van Flow-GRPO (dat voortbouwt op DDPO), wanneer toegepast op Stable Diffusion 3.5 Medium en FLUX, zonder in te leveren op generatiekwaliteit. Code is beschikbaar op https://github.com/scxue/advantage_weighted_matching.

English

Reinforcement Learning (RL) has emerged as a central paradigm for advancing Large Language Models (LLMs), where pre-training and RL post-training share 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 Advantage Weighted Matching (AWM), a policy-gradient method for diffusion. It uses the same score/flow-matching loss as pretraining to obtain a lower-variance objective 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 unifies pretraining and RL conceptually and practically, is consistent with policy-gradient theory, reduces variance, and yields faster convergence. This simple yet effective design yields substantial benefits: on GenEval, OCR, and PickScore benchmarks, AWM delivers up to a 24times 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 https://github.com/scxue/advantage_weighted_matching.