TDM-R1: Усиление диффузионных моделей с малым числом шагов с помощью недифференцируемой функции вознаграждения
TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward
March 8, 2026
Авторы: Yihong Luo, Tianyang Hu, Weijian Luo, Jing Tang
cs.AI
Аннотация
Хотя малошаговые генеративные модели позволили достичь мощной генерации изображений и видео при значительно меньших затратах, общие парадигмы обучения с подкреплением (RL) для малошаговых моделей остаются нерешенной проблемой. Существующие RL-подходы для малошаговых диффузионных моделей сильно зависят от обратного распространения через дифференцируемые модели вознаграждения, что исключает использование большинства важных сигналов вознаграждения из реального мира, например, недифференцируемых наград, таких как бинарная оценка предпочтения человеком, количество объектов и т.д. Для корректного включения недифференцируемых наград с целью улучшения малошаговых генеративных моделей мы представляем TDM-R1 — новую парадигму обучения с подкреплением, построенную на основе передовой малошаговой модели, сопоставления распределений траекторий (Trajectory Distribution Matching, TDM). TDM-R1 разделяет процесс обучения на этапы обучения суррогатной функции вознаграждения и обучения генератора. Кроме того, мы разработали практические методы получения пошаговых сигналов вознаграждения вдоль детерминированной траектории генерации TDM, что привело к созданию унифицированного метода пост-обучения с подкреплением, который значительно улучшает способности малошаговых моделей работать с общими наградами. Мы проводим обширные эксперименты, охватывающие рендеринг текста, визуальное качество и согласование с предпочтениями. Все результаты демонстрируют, что TDM-R1 является мощной парадигмой обучения с подкреплением для малошаговых тексто-изобразительных моделей, достигая передовых показателей RL как по внутридоменным, так и по внедоменным метрикам. Более того, TDM-R1 также эффективно масштабируется на недавно появившуюся мощную Z-Image модель, последовательно превосходя как её 100-NFE, так и малошаговые варианты, используя всего 4 NFE. Страница проекта: 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