TDM-R1: Reforçando Modelos de Difusão de Poucos Passos com Recompensa Não Diferenciável

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

Resumo

Embora os modelos generativos de poucos passos tenham possibilitado uma poderosa geração de imagens e vídeos a um custo significativamente menor, os paradigmas genéricos de aprendizagem por reforço (RL) para modelos de poucos passos permanecem um problema não resolvido. As abordagens de RL existentes para modelos de difusão de poucos passos dependem fortemente da retropropagação através de modelos de recompensa diferenciáveis, excluindo assim a maioria dos sinais de recompensa importantes do mundo real, por exemplo, recompensas não diferenciáveis, como a semelhança binária humana, contagens de objetos, etc. Para incorporar adequadamente recompensas não diferenciáveis para melhorar os modelos generativos de poucos passos, introduzimos o TDM-R1, um novo paradigma de aprendizagem por reforço construído sobre um modelo líder de poucos passos, o Trajectory Distribution Matching (TDM). O TDM-R1 desacopla o processo de aprendizagem em aprendizagem de recompensa substituta e aprendizagem do gerador. Além disso, desenvolvemos métodos práticos para obter sinais de recompensa por passo ao longo da trajetória determinística de geração do TDM, resultando em um método unificado de pós-treinamento por RL que melhora significativamente a capacidade dos modelos de poucos passos com recompensas genéricas. Realizamos extensos experimentos que variam desde renderização de texto, qualidade visual e alinhamento de preferências. Todos os resultados demonstram que o TDM-R1 é um paradigma de aprendizagem por reforço poderoso para modelos de texto para imagem de poucos passos, alcançando desempenhos de RL de última geração em métricas tanto dentro quanto fora do domínio. Além disso, o TDM-R1 também escala efetivamente para o recente e forte modelo Z-Image, superando consistentemente tanto as suas variantes de 100-NFE quanto as de poucos passos com apenas 4 NFEs. Página do projeto: 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