TDM-R1: Rinforzo di Modelli di Diffusione a Pochi Passi con Ricompense Non Differenziabili

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

March 8, 2026

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

cs.AI

Abstract

Sebbene i modelli generativi a pochi passi abbiano reso possibile una potente generazione di immagini e video a costi significativamente ridotti, i paradigmi generici di apprendimento per rinforzo (RL) per i modelli a pochi passi rimangono un problema irrisolto. Gli approcci RL esistenti per i modelli di diffusione a pochi passi si basano fortemente sulla retropropagazione attraverso modelli di ricompensa differenziabili, escludendo così la maggior parte degli importanti segnali di ricompensa del mondo reale, ad esempio ricompense non differenziabili come il gradimento binario umano, il conteggio di oggetti, ecc. Per incorporare correttamente le ricompense non differenziabili al fine di migliorare i modelli generativi a pochi passi, introduciamo TDM-R1, un nuovo paradigma di apprendimento per rinforzo basato su un modello leader a pochi passi, il Trajectory Distribution Matching (TDM). TDM-R1 disaccoppia il processo di apprendimento in apprendimento di una ricompensa surrogata e apprendimento del generatore. Inoltre, abbiamo sviluppato metodi pratici per ottenere segnali di ricompensa per ogni passo lungo la traiettoria di generazione deterministica del TDM, dando vita a un metodo unificato di post-addestramento RL che migliora significativamente la capacità dei modelli a pochi passi di gestire ricompense generiche. Conduciamo esperimenti approfonditi che spaziano dal rendering del testo, alla qualità visiva e all'allineamento alle preferenze. Tutti i risultati dimostrano che TDM-R1 è un potente paradigma di apprendimento per rinforzo per i modelli text-to-image a pochi passi, raggiungendo prestazioni all'avanguardia nell'apprendimento per rinforzo sia su metriche in-dominio che out-of-domain. Inoltre, TDM-R1 scala efficacemente anche al recente e potente modello Z-Image, superando costantemente sia le sue varianti a 100 NFE che quelle a pochi passi con soli 4 NFE. Pagina del progetto: 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