TDM-R1: 비미분 가능 보상을 활용한 Few-Step 확산 모델 강화
TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward
March 8, 2026
저자: Yihong Luo, Tianyang Hu, Weijian Luo, Jing Tang
cs.AI
초록
소수-스텝 생성 모델이 훨씬 낮은 비용으로 강력한 이미지 및 동영상 생성을 가능하게 했지만, 소수-스텝 모델을 위한 일반적인 강화 학습 패러다임은 여전히 해결되지 않은 문제로 남아 있습니다. 기존의 소수-스텝 확산 모델에 대한 강화 학습 접근법은 미분 가능한 보상 모델을 통한 역전파에 크게 의존하여, 인간의 이진적 선호도나 객체 개수 등과 같은 비미분 가능한 보상을 포함한 대부분의 중요한 실제 세계 보상 신호를 배제해 왔습니다. 비미분 가능한 보상을 적절히 통합하여 소수-스텝 생성 모델을 개선하기 위해, 우리는 주요 소수-스텝 모델인 궤적 분포 매칭(TDM)을 기반으로 한 새로운 강화 학습 패러다임인 TDM-R1을 소개합니다. TDM-R1은 학습 과정을 대리 보상 학습과 생성기 학습으로 분리합니다. 더 나아가, 우리는 TDM의 결정론적 생성 궤적을 따라 스텝별 보상 신호를 얻는 실용적인 방법을 개발하여, 소수-스텝 모델의 일반적인 보상 처리 능력을 크게 향상시키는 통합된 강화 학습 사후 학습 방법을 도출했습니다. 텍스트 렌더링, 시각적 품질, 선호도 정렬에 이르는 광범위한 실험을 수행했습니다. 모든 결과는 TDM-R1이 소수-스텝 텍스트-이미지 모델을 위한 강력한 강화 학습 패러다임이며, 도메인 내 및 도메인 외 메트릭 모두에서 최첨단 강화 학습 성능을 달성함을 보여줍니다. 또한 TDM-R1은 최근 강력한 Z-Image 모델에도 효과적으로 확장되어, 단 4 NFE로 100-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