Emparejamiento Ponderado por Ventajas: Alineando el Aprendizaje por Refuerzo con el Pretrenamiento en Modelos de Difusión

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

Resumen

El Aprendizaje por Refuerzo (RL, por sus siglas en inglés) ha surgido como un paradigma central para avanzar en los Modelos de Lenguaje a Gran Escala (LLMs), donde el preentrenamiento y el postentrenamiento con RL comparten la misma formulación de verosimilitud logarítmica. En contraste, los enfoques recientes de RL para modelos de difusión, destacándose la Optimización de Políticas de Difusión de Ruido (DDPO), optimizan un objetivo diferente a los objetivos de preentrenamiento: la pérdida de coincidencia de puntuación/flujo. En este trabajo, establecemos un nuevo análisis teórico: DDPO es una forma implícita de coincidencia de puntuación/flujo con objetivos ruidosos, lo que aumenta la varianza y ralentiza la convergencia. Basándonos en este análisis, introducimos la Coincidencia Ponderada por Ventaja (AWM), un método de gradiente de políticas para difusión. Utiliza la misma pérdida de coincidencia de puntuación/flujo que el preentrenamiento para obtener un objetivo de menor varianza y reponderar cada muestra según su ventaja. En efecto, AWM aumenta la influencia de las muestras de alta recompensa y suprime las de baja recompensa, manteniendo el objetivo de modelado idéntico al preentrenamiento. Esto unifica conceptual y prácticamente el preentrenamiento y el RL, es consistente con la teoría de gradiente de políticas, reduce la varianza y produce una convergencia más rápida. Este diseño simple pero efectivo ofrece beneficios sustanciales: en los puntos de referencia GenEval, OCR y PickScore, AWM proporciona una aceleración de hasta 24 veces sobre Flow-GRPO (que se basa en DDPO), cuando se aplica a Stable Diffusion 3.5 Medium y FLUX, sin comprometer la calidad de la generación. El código está disponible en 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.