TDM-R1: Reforzamiento de Modelos de Difusión de Pocos Pasos con Recompensa No Diferenciable

TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward

March 8, 2026

Autores: Yihong Luo, Tianyang Hu, Weijian Luo, Jing Tang

cs.AI

Resumen

Si bien los modelos generativos de pocos pasos han permitido una potente generación de imágenes y videos a un costo significativamente menor, los paradigmas genéricos de aprendizaje por refuerzo (RL) para modelos de pocos pasos siguen siendo un problema sin resolver. Los enfoques de RL existentes para modelos de difusión de pocos pasos dependen en gran medida de la retropropagación a través de modelos de recompensa diferenciables, excluyendo así la mayoría de las señales de recompensa importantes del mundo real, por ejemplo, recompensas no diferenciables como la preferencia binaria humana, el recuento de objetos, etc. Para incorporar adecuadamente recompensas no diferenciables y así mejorar los modelos generativos de pocos pasos, presentamos TDM-R1, un novedoso paradigma de aprendizaje por refuerzo basado en un modelo líder de pocos pasos, el Emparejamiento de Distribución de Trayectorias (TDM). TDM-R1 desacopla el proceso de aprendizaje en el aprendizaje de recompensas sustitutas (surrogate) y el aprendizaje del generador. Además, desarrollamos métodos prácticos para obtener señales de recompensa por paso a lo largo de la trayectoria de generación determinista del TDM, dando como resultado un método unificado de post-entrenamiento por RL que mejora significativamente la capacidad de los modelos de pocos pasos con recompensas genéricas. Realizamos extensos experimentos que abarcan desde la representación de texto, la calidad visual y la alineación de preferencias. Todos los resultados demuestran que TDM-R1 es un paradigma de aprendizaje por refuerzo potente para modelos de texto a imagen de pocos pasos, logrando rendimientos de vanguardia en aprendizaje por refuerzo tanto en métricas dentro del dominio como fuera de él. Además, TDM-R1 también se escala eficazmente al reciente y potente modelo Z-Image, superando consistentemente tanto a sus variantes de 100 evaluaciones de función directa (NFE) como a las de pocos pasos con solo 4 NFE. Página del proyecto: https://github.com/Luo-Yihong/TDM-R1

English

While few-step generative models have enabled powerful image and video generation at significantly lower cost, generic reinforcement learning (RL) paradigms for few-step models remain an unsolved problem. Existing RL approaches for few-step diffusion models strongly rely on back-propagating through differentiable reward models, thereby excluding the majority of important real-world reward signals, e.g., non-differentiable rewards such as humans' binary likeness, object counts, etc. To properly incorporate non-differentiable rewards to improve few-step generative models, we introduce TDM-R1, a novel reinforcement learning paradigm built upon a leading few-step model, Trajectory Distribution Matching (TDM). TDM-R1 decouples the learning process into surrogate reward learning and generator learning. Furthermore, we developed practical methods to obtain per-step reward signals along the deterministic generation trajectory of TDM, resulting in a unified RL post-training method that significantly improves few-step models' ability with generic rewards. We conduct extensive experiments ranging from text-rendering, visual quality, and preference alignment. All results demonstrate that TDM-R1 is a powerful reinforcement learning paradigm for few-step text-to-image models, achieving state-of-the-art reinforcement learning performances on both in-domain and out-of-domain metrics. Furthermore, TDM-R1 also scales effectively to the recent strong Z-Image model, consistently outperforming both its 100-NFE and few-step variants with only 4 NFEs. Project page: https://github.com/Luo-Yihong/TDM-R1