Reforzando Generadores de Pocos Pasos mediante Coincidencia de Distribuciones Sesgada por Recompensa

Resumen

Avances recientes en destilación de difusión en pocos pasos han permitido una generación eficiente de imágenes, aunque alinear estos modelos con las preferencias humanas sigue siendo un desafío. Proponemos Destilación por Emparejamiento de Distribuciones Sesgada por Recompensa (RTDMD, por sus siglas en inglés), un marco de dos etapas que unifica la destilación por emparejamiento de distribuciones con el aprendizaje por refuerzo guiado por recompensa para generadores de flujo de pocos pasos. Demostramos que minimizar la divergencia KL con respecto a una distribución de profesor sesgada por recompensa se descompone naturalmente en un término de emparejamiento de distribuciones y un término de maximización de recompensa. En la primera etapa, introducimos Destilación por Emparejamiento de Distribuciones Consistente con el Entorno (AC-DMD), que realiza un emparejamiento de distribuciones por subintervalos y aumenta el objetivo de puntuación falsa con un regularizador de consistencia para ayudar al modelo de puntuación falsa a rastrear la distribución cambiante del generador bajo actualizaciones limitadas. En la segunda etapa, optimizamos conjuntamente ambos términos: para el término de maximización de recompensa, derivamos un gradiente de política híbrido que combina un estimador de estilo GRPO para las transiciones intermedias estocásticas con la retropropagación directa de la recompensa a través del paso final determinista, e introducimos además GRPO por subconjunto de pasos (SubGRPO) para reducir la varianza. Los experimentos en SD3, SD3.5 y FLUX.2 demuestran que RTDMD establece nuevos resultados de última generación en métricas de preferencia, estética y composición con solo 4 pasos de inferencia, superando a métodos anteriores de generación de texto a imagen en pocos pasos. El código y los modelos están disponibles en https://github.com/Harahan/RTDMD.

English

Recent advances in few-step diffusion distillation have enabled efficient image generation, yet aligning these models with human preferences remains challenging. We propose Reward-Tilted Distribution Matching Distillation (RTDMD), a two-stage framework that unifies distribution matching distillation with reward-guided reinforcement learning for few-step flow generators. We show that minimizing the KL divergence to a reward-tilted teacher distribution naturally decomposes into a distribution matching term and a reward maximization term. In the first stage, we introduce Ambient-Consistent Distribution Matching Distillation (AC-DMD), which performs subinterval-wise distribution matching and augments the fake score objective with a consistency regularizer to help the fake score model track the shifting generator distribution under limited updates. In the second stage, we jointly optimize both terms: for the reward maximization term, we derive a hybrid policy gradient that combines a GRPO-style estimator for the stochastic intermediate transitions with direct reward backpropagation through the deterministic final step, and further introduce step-subset GRPO (SubGRPO) to reduce variance. Experiments on SD3, SD3.5, and FLUX.2 demonstrate that RTDMD establishes new state-of-the-art results across preference, aesthetic, and compositional metrics with only 4 inference steps, outperforming previous few-step text-to-image generation methods. Code and models are available at https://github.com/Harahan/RTDMD.