Advantage Weighted Matching: Allineamento dell'Apprendimento per Rinforzo con il Pretraining nei Modelli di Diffusione

Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models

September 29, 2025

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

cs.AI

Abstract

Il Reinforcement Learning (RL) è emerso come un paradigma centrale per il progresso dei Large Language Models (LLMs), dove il pre-training e il post-training con RL condividono la stessa formulazione di log-verosimiglianza. Al contrario, i recenti approcci RL per i modelli di diffusione, in particolare il Denoising Diffusion Policy Optimization (DDPO), ottimizzano un obiettivo diverso da quelli del pre-training—la perdita di score/flow matching. In questo lavoro, stabiliamo una nuova analisi teorica: il DDPO è una forma implicita di score/flow matching con target rumorosi, il che aumenta la varianza e rallenta la convergenza. Basandoci su questa analisi, introduciamo l'Advantage Weighted Matching (AWM), un metodo di policy-gradient per la diffusione. Esso utilizza la stessa perdita di score/flow matching del pre-training per ottenere un obiettivo a varianza ridotta e ripesca ogni campione in base al suo vantaggio. In pratica, l'AWM aumenta l'influenza dei campioni ad alta ricompensa e sopprime quelli a bassa ricompensa, mantenendo l'obiettivo di modellazione identico al pre-training. Ciò unifica concettualmente e praticamente il pre-training e il RL, è coerente con la teoria del policy-gradient, riduce la varianza e produce una convergenza più rapida. Questo design semplice ma efficace offre notevoli vantaggi: sui benchmark GenEval, OCR e PickScore, l'AWM fornisce un'accelerazione fino a 24 volte rispetto al Flow-GRPO (che si basa su DDPO), quando applicato a Stable Diffusion 3.5 Medium e FLUX, senza compromettere la qualità della generazione. Il codice è disponibile all'indirizzo 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.