Appariemment Pondéré par les Avantages : Alignement de l'Apprentissage par Renforcement avec le Prétraitement dans les Modèles de Diffusion

Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models

September 29, 2025

papers.authors: Shuchen Xue, Chongjian Ge, Shilong Zhang, Yichen Li, Zhi-Ming Ma

cs.AI

papers.abstract

L'apprentissage par renforcement (Reinforcement Learning, RL) est devenu un paradigme central pour faire progresser les modèles de langage de grande taille (Large Language Models, LLMs), où le pré-entraînement et le post-entraînement par RL partagent la même formulation de log-vraisemblance. En revanche, les approches récentes de RL pour les modèles de diffusion, notamment l'optimisation de politique par débruitage (Denoising Diffusion Policy Optimization, DDPO), optimisent un objectif différent de ceux du pré-entraînement—la perte d'appariement de score/flux. Dans ce travail, nous établissons une nouvelle analyse théorique : DDPO est une forme implicite d'appariement de score/flux avec des cibles bruitées, ce qui augmente la variance et ralentit la convergence. Sur la base de cette analyse, nous introduisons l'appariement pondéré par avantage (Advantage Weighted Matching, AWM), une méthode de gradient de politique pour la diffusion. Elle utilise la même perte d'appariement de score/flux que le pré-entraînement pour obtenir un objectif à variance réduite et repondère chaque échantillon par son avantage. En pratique, AWM augmente l'influence des échantillons à haute récompense et supprime ceux à faible récompense tout en maintenant l'objectif de modélisation identique au pré-entraînement. Cela unifie conceptuellement et pratiquement le pré-entraînement et le RL, est cohérent avec la théorie du gradient de politique, réduit la variance et permet une convergence plus rapide. Cette conception simple mais efficace offre des avantages substantiels : sur les benchmarks GenEval, OCR et PickScore, AWM offre jusqu'à une accélération de 24 fois par rapport à Flow-GRPO (qui s'appuie sur DDPO), lorsqu'il est appliqué à Stable Diffusion 3.5 Medium et FLUX, sans compromettre la qualité de génération. Le code est disponible à l'adresse 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.