TDM-R1 : Renforcement des modèles de diffusion à faible nombre d'étapes avec une récompense non différentiable

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

March 8, 2026

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

cs.AI

Résumé

Bien que les modèles génératifs à faible nombre d'étapes aient permis une génération d'images et de vidéos puissante à un coût significativement réduit, les paradigmes génériques d'apprentissage par renforcement (RL) pour les modèles à faible nombre d'étapes restent un problème non résolu. Les approches de RL existantes pour les modèles de diffusion à faible nombre d'étapes reposent fortement sur la rétropropagation à travers des modèles de récompense différentiables, excluant ainsi la majorité des signaux de récompense importants du monde réel, par exemple, les récompenses non différentiables telles que l'appréciation binaire humaine, le décompte d'objets, etc. Pour intégrer correctement les récompenses non différentiables afin d'améliorer les modèles génératifs à faible nombre d'étapes, nous présentons TDM-R1, un nouveau paradigme d'apprentissage par renforcement construit sur un modèle à faible nombre d'étapes de premier plan, l'Appariement de la Distribution de Trajectoire (TDM). TDM-R1 découple le processus d'apprentissage en apprentissage de récompense surrogate et apprentissage du générateur. De plus, nous avons développé des méthodes pratiques pour obtenir des signaux de récompense par étape le long de la trajectoire de génération déterministe du TDM, aboutissant à une méthode unifiée de post-entraînement par RL qui améliore significativement la capacité des modèles à faible nombre d'étapes avec des récompenses génériques. Nous menons des expériences approfondies couvrant le rendu de texte, la qualité visuelle et l'alignement des préférences. Tous les résultats démontrent que TDM-R1 est un paradigme d'apprentissage par renforcement puissant pour les modèles texte-image à faible nombre d'étapes, atteignant des performances de pointe en apprentissage par renforcement sur des métriques internes et externes au domaine. De plus, TDM-R1 s'adapte également efficacement au récent modèle puissant Z-Image, surpassant constamment à la fois ses variantes à 100 NFE et à faible nombre d'étapes avec seulement 4 NFE. Page du projet : 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