TDM-R1: Versterking van Few-Step Diffusiemodellen met Niet-Differentieerbare Beloning

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

Samenvatting

Hoewel generatieve modellen met weinig stappen krachtige beeld- en videogeneratie tegen aanzienlijk lagere kosten mogelijk hebben gemaakt, blijven generieke paradigma's voor reinforcement learning (RL) voor modellen met weinig stappen een onopgelost probleem. Bestaande RL-benaderingen voor diffusiemodellen met weinig stappen zijn sterk afhankelijk van backpropagatie door differentieerbare beloningsmodellen, waardoor de meerderheid van belangrijke beloningssignalen uit de echte wereld wordt uitgesloten, zoals niet-differentieerbare beloningen zoals binaire voorkeuren van mensen, objectaantallen, enz. Om niet-differentieerbare beloningen correct te integreren om generatieve modellen met weinig stappen te verbeteren, introduceren we TDM-R1, een nieuw reinforcement learning-paradigma gebaseerd op een toonaangevend model met weinig stappen, Trajectory Distribution Matching (TDM). TDM-R1 ontkoppelt het leerproces in surrogaatbeloningsleren en generatorleren. Verder ontwikkelden we praktische methoden om per-stap beloningssignalen langs het deterministische generatietraject van TDM te verkrijgen, wat resulteert in een uniforme RL-natrainingsmethode die het vermogen van modellen met weinig stappen aanzienlijk verbetert met generieke beloningen. We voeren uitgebreide experimenten uit variërend van tekstweergave, visuele kwaliteit en voorkeursafstemming. Alle resultaten tonen aan dat TDM-R1 een krachtig reinforcement learning-paradigma is voor tekst-naar-beeldmodellen met weinig stappen, dat state-of-the-art reinforcement learning-prestaties behaalt op zowel in-domein als out-of-domein metrieken. Bovendien schaalt TDM-R1 effectief naar het recente sterke Z-Image-model en presteert het consistent beter dan zowel de 100-NFE- als de weinig-stappen varianten met slechts 4 NFE's. Projectpagina: 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