Organizations: Hong Kong University of Science and Technology · The Chinese University of Hong Kong, Shenzhen · hi-Lab, Xiaohongshu Inc · Hong Kong University of Science and Technology (Guangzhou)
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
Figures & tables
Figure 1: Samples generated by TDM-R1 using only 4 NFEs, obtained by reinforcing the recent powerful Z-Image model ( Team et al., 2025 ) .
Figure 2: TDM-R1 rapidly boosts GenEval score of few-step TDM, notably outperforming its many-step base model and GPT-4o. This is achieved without sacrificing out-of-domain metrics.
Figure 3: Qualitative comparisons of TDM-R1 against competing methods.
Figure 4: Compare the training performance and speed of TDM-R1 and potential baselines.
Figure 5: Compare TDM-R1 with the direct combination of TDM and RL loss.
Model
Overall
Single Obj.
Two Obj.
Counting
Colors
Position
Attr. Binding
Autoregressive Models:
Show-o ( Xie et al., 2024a )
0.53
0.95
0.52
0.49
0.82
0.11
0.28
Emu3-Gen ( Wang et al., 2024 )
0.54
0.98
0.71
0.34
0.81
0.17
0.21
JanusFlow ( Ma et al., 2025 )
0.63
0.97
0.59
0.45
0.83
0.53
0.42
Janus-Pro-7B ( Chen et al., 2025 )
0.80
0.99
0.89
0.59
0.90
0.79
0.66
GPT-4o ( Hurst et al., 2024 )
0.84
0.99
0.92
0.85
0.92
0.75
0.61
Table 1: GenEval Result. We highlight the best scores among specific category. Obj.: Object; Attr.: Attribution.
Model
NFE
Verifiable Metric
Image Quality
Preference Score
GenEval
OCR Acc.
Aesthetic
DeQA
ImgRwd
PickScore
UniRwd
SD3.5-M ( Esser et al., 2024 )
80
0.63
0.59
5.39
4.07
0.87
22.34
3.33
Flow-GRPO † ( Liu et al., 2025 )
80
0.95
—
5.25
4.01
1.03
22.37
3.51
DGPO † ( Luo et al., 2025d )
80
0.97
—
5.31
4.03
1.08
22.41
3.60
Flow-GRPO † ( Liu et al., 2025 )
80
—
0.92
5.32
4.06
0.95
22.44
3.42
DGPO † ( Luo et al., 2025d )
80
—
0.96
5.37
4.09
1.02
22.52
3.48
Table 2: Performance on Compositional Image Generation and Visual Text Rendering Benchmarks . ImgRwd: ImageReward; UniRwd: UnifiedReward. We highlight the metric adopted for the training signal. † Results are taken from the original paper.
Figure 6: Qualitative comparisons of TDM-R1 with GenEval signals against competing methods. It can be seen that our proposed TDM-R1 can accurately follow the instructions, while keeping a strong visual quality. The same initial noise is used to generate all images.
Model
NFE
Verifiable Metric
Model-based Metric
GenEval
OCR Acc.
HPSv3
Aesthetic
ImgRwd
PickScore
UniRwd
Z-Image
100
0.66
0.74
7.32
5.35
0.62
19.98
3.64
Z-Image-Turbo (DMDR)
4
0.73
0.78
9.13
5.40
0.78
20.28
3.70
TDM-R1-Zimage (Ours)
4
0.77
0.79
9.90
5.49
0.94
20.45
3.75
Table 3: Performance on reinforcing Z-Image . ImgRwd: ImageReward; UniRwd: UnifiedReward. We highlight the metric adopted for the training signal.
Figure 7: TDM-R1 Performance in human preference alignment.
Figure 8: Comparison of visual text alignment across variants.
Figure 9: Comparison between TDM-R1 with direct distillation.
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Figure 10: Qualitative comparisons of TDM-R1 against competing baselines. The same initial noise is used to generate all images.
Recent advances in one-step text-to-image generation have enabled real-time synthesis with remarkable efficiency and quality. Previous reinforcement learning methods for one-step generators combine image-space reward optimization with diffusion noisy-space distribution matching. This paradigm brings challenges due to a mismatch between terminal reward optimization and the underlying generative dynamics. As a result, optimization tends to exploit stochastic degrees of freedom, often improving reward at the expense of image fidelity. To address this issue, we propose Diff-Instruct with Diffused Reward (DIDR), a data-free trajectory-level alignment framework derived from Integral KL minimization. DIDR propagates the RLHF-optimal reward-tilted clean-image distribution across all noise levels along the diffusion trajectory. We show that this objective admits the same minimizer as clean-image RLHF, while naturally inducing the Diffused Reward Score (DRS), which acts as a reward-driven correction to the reference score function. To make this practical, we further introduce the Diffused Reward Proxy (DRP), an efficient estimator of DRS based on differentiable short-step denoising. Extensive experiments demonstrate that DIDR consistently Pareto-dominates existing one-step SDXL baselines. Moreover, when transferred to a 6B DiT backbone (Z-Image), DIDR surpasses its 50-step teacher in preference alignment while requiring only a single generation step.
Diffusion models offer superior generation quality at the expense of extensive sampling steps. Distillation methods, with Distribution Matching Distillation (DMD) as a popular example, can mitigate this issue, but performance degradation remains pronounced when sampling steps are limited. Reinforcement learning (RL) has been leveraged to improve the few-step generation quality during distillation, with the potential to even surpass the performance of the teacher model. However, existing approaches are combinatorial in nature, merely integrating an RL process with the distillation process, which introduces unnecessary complexities. To address this gap, we propose AdvDMD, a method that seamlessly unifies DMD distillation and RL. Specifically, AdvDMD employs the adversarially trained discriminator from DMD2 as the reward model, which assigns low scores to generated images and high scores to real ones. It is trained on both intermediate and final states of the denoising process and updated online with the distilled model, enabling a holistic supervision of the sampling trajectories and mitigating reward hacking. We adopt a unified SDE backward simulation and a different training schedule for DMD and RL to enable a more stable and efficient training. Experimental results demonstrate that the 4-step AdvDMD outperforms the original 40-step model for SD3.5 on DPG-Bench, while achieving significant performance gains for SD3 on the GenEval. On Qwen-Image, our 2-step AdvDMD achieves superior performance over TwinFlow.
Recent advances in few-step diffusion distillation have enabled efficient image generation, yet aligning these models with human preferences remains challenging. We propose Reward-Tilted Distribution Matching Distillation (RTDMD), a two-stage framework that unifies distribution matching distillation with reward-guided reinforcement learning for few-step flow generators. We show that minimizing the KL divergence to a reward-tilted teacher distribution naturally decomposes into a distribution matching term and a reward maximization term. In the first stage, we introduce Ambient-Consistent Distribution Matching Distillation (AC-DMD), which performs subinterval-wise distribution matching and augments the fake score objective with a consistency regularizer to help the fake score model track the shifting generator distribution under limited updates. In the second stage, we jointly optimize both terms: for the reward maximization term, we derive a hybrid policy gradient that combines a GRPO-style estimator for the stochastic intermediate transitions with direct reward backpropagation through the deterministic final step, and further introduce step-subset GRPO (SubGRPO) to reduce variance. Experiments on SD3, SD3.5, and FLUX.2 demonstrate that RTDMD establishes new state-of-the-art results across preference, aesthetic, and compositional metrics with only 4 inference steps, outperforming previous few-step text-to-image generation methods. Code and models are available at https://github.com/Harahan/RTDMD.
Yushi Huang, Xiangxin Zhou, Ruoyu Wang +3
Tencent Hunyuan · Hong Kong University of Science and Technology · Westlake University