Emparelhamento Ponderado por Vantagem: Alinhando Aprendizado por Reforço com Pré-treinamento em Modelos de Difusão

Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models

September 29, 2025

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

cs.AI

Resumo

O Aprendizado por Reforço (RL) emergiu como um paradigma central para o avanço de Modelos de Linguagem de Grande Escala (LLMs), onde o pré-treinamento e o pós-treinamento com RL compartilham a mesma formulação de log-verossimilhança. Em contraste, abordagens recentes de RL para modelos de difusão, mais notavelmente a Otimização de Política de Difusão de Ruído (DDPO), otimizam um objetivo diferente dos objetivos de pré-treinamento—a perda de correspondência de pontuação/fluxo. Neste trabalho, estabelecemos uma nova análise teórica: a DDPO é uma forma implícita de correspondência de pontuação/fluxo com alvos ruidosos, o que aumenta a variância e retarda a convergência. Com base nessa análise, introduzimos a Correspondência Ponderada por Vantagem (AWM), um método de gradiente de política para difusão. Ele utiliza a mesma perda de correspondência de pontuação/fluxo do pré-treinamento para obter um objetivo de menor variância e repondera cada amostra pela sua vantagem. Na prática, a AWM aumenta a influência de amostras de alta recompensa e suprime as de baixa recompensa, mantendo o objetivo de modelagem idêntico ao pré-treinamento. Isso unifica o pré-treinamento e o RL tanto conceitual quanto praticamente, é consistente com a teoria de gradiente de política, reduz a variância e resulta em convergência mais rápida. Este design simples, porém eficaz, traz benefícios substanciais: nos benchmarks GenEval, OCR e PickScore, a AWM oferece uma aceleração de até 24 vezes em relação ao Flow-GRPO (que é baseado na DDPO), quando aplicada ao Stable Diffusion 3.5 Medium e FLUX, sem comprometer a qualidade da geração. O código está disponível em 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.