TDM-R1: Verstärkung von Few-Step-Diffusionsmodellen mit nicht-differenzierbaren Belohnungen

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

March 8, 2026

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

cs.AI

Zusammenfassung

Während generative Modelle mit wenigen Schritten eine leistungsstarke Bild- und Videogenerierung zu deutlich geringeren Kosten ermöglicht haben, bleiben generische Reinforcement-Learning (RL)-Paradigmen für solche Modelle mit wenigen Schritten ein ungelöstes Problem. Bestehende RL-Ansätze für Diffusionsmodelle mit wenigen Schritten stützen sich stark auf das Backpropagieren durch differenzierbare Belohnungsmodelle, wodurch die Mehrheit der wichtigen Belohnungssignale aus der realen Welt ausgeschlossen wird, z. B. nicht-differenzierbare Belohnungen wie binäre Präferenzen von Menschen, Objektanzahlen usw. Um nicht-differenzierbare Belohnungen zur Verbesserung generativer Modelle mit wenigen Schritten sinnvoll einzubinden, führen wir TDM-R1 ein, ein neuartiges Reinforcement-Learning-Paradigma, das auf einem führenden Modell mit wenigen Schritten basiert, dem Trajectory Distribution Matching (TDM). TDM-R1 entkoppelt den Lernprozess in Surrogat-Belohnungslernen und Generatorlernen. Darüber hinaus haben wir praktische Methoden entwickelt, um belohnungsrelevante Signale pro Schritt entlang der deterministischen Generierungstrajektorie von TDM zu erhalten, was zu einer einheitlichen RL-Nachschulungsmethode führt, die die Fähigkeiten von Modellen mit wenigen Schritten im Umgang mit generischen Belohnungen signifikant verbessert. Wir führen umfangreiche Experimente durch, die Textdarstellung, visuelle Qualität und Präferenzausrichtung abdecken. Alle Ergebnisse demonstrieren, dass TDM-R1 ein leistungsstarkes Reinforcement-Learning-Paradigma für Text-zu-Bild-Modelle mit wenigen Schritten ist, das state-of-the-art RL-Leistungen sowohl in domäneninternen als auch domänenexternen Metriken erreicht. Darüber hinaus skaliert TDM-R1 auch effektiv mit dem recent starken Z-Image-Modell und übertrifft konsistent sowohl seine 100-NFE- als auch seine Varianten mit wenigen Schritten bei nur 4 NFEs. Projektseite: 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