DualHeads: Reward-tilted few-step diffusion made simple

Yuanzhi Zhu; Eleftherios Tsonis; Simo Alami Chehboune; Xi Wang; Stéphane Lathuilière; Vicky Kalogeiton

cs.CVOct 20, 2026

DualHeads: Reward-tilted few-step diffusion made simple

Authors: Yuanzhi Zhu, Eleftherios Tsonis, Simo Alami Chehboune, Xi Wang, Stéphane Lathuilière, Vicky Kalogeiton

Organizations: LIX, École Polytechnique, CNRS, IPP · Inria at Univ. Grenoble Alpes, CNRS, LJK

Abstract

Distilled diffusion models generate high-quality images in a few sampling steps, but aligning them with task requirements and human preferences is difficult. Methods developed for multi-step diffusion models rely on differentiable rewards, or on score fields and stochastic trajectories that distilled generators may lack. We ask whether a few-step generator can be adapted using only external reward evaluations and reference samples, and show that two binary classifiers suffice. Our method, DualHeads, implements them as lightweight heads on a frozen diffusion backbone. The reward head learns which samples of the same prompt score higher, the reference head separates reference images from generated ones, and the generator maximizes a weighted sum of their logits. When the heads are optimal, this objective is KL-regularized maximization of the learned reward, and the ratio of head weights sets the regularization strength. Empirically, with both rule-based and model-based rewards, a few-step DualHeads generator reaches task rewards close to those of 40-step reinforcement learning methods while preserving image quality, at no extra inference cost. We hope that this simple and elegant dual-head formulation will help advance few-step reward-tilted generation.

Figures & tables

Figure 1: Four-step samples from DualHeads. Z-Image-Turbo ( Jiang et al., 2026 ) fine-tuned with DualHeads using the HPSv3 reward ( Ma et al., 2025 ) .

Figure 2: Generator update of DualHeads. The generator $f_{\theta}$ first rolls out $N$ steps from noise without gradient; the endpoint $x_{1}$ is re-noised to $x_{t_{i}}$ , from which $f_{\theta}$ predicts $\hat{x}_{1}$ in one differentiable step. The prediction is noised to $x_{t}$ and sent to a frozen feature extractor $F$ followed by a reference head $D_{\phi}$ and a reward head $R_{\psi}$ ; the generator ascends the utility $s=\lambda_{D}D_{\phi}+\lambda_{R}R_{\psi}$ through $\hat{x}_{1}$ , while the backbone and both heads stay fixed during its update.


Method	NFE	GenEval $\uparrow$	OCR $\uparrow$	Aes. $\uparrow$	ImgRwd $\uparrow$	Pick $\uparrow$	CLIP $\uparrow$	HPS $\uparrow$
		Task Metric		Held-out Metric
SD3.5-M (base)	80	0.63	0.59	5.39	0.87	22.34	27.99	0.279
SD3.5-M TDM ( Luo et al., 2025 )	4	0.62	0.51	5.58	0.86	22.14	27.18	0.290
Flow-GRPO ( Liu et al., 2025 ) w/o KL	80	0.95	—	4.93	0.44	21.16	—	—
Flow-GRPO ( Liu et al., 2025 ) w/ KL	80	0.95	—	5.25	1.03	22.37	29.25	0.274

Table 1: Comparison on compositional image generation ( GenEval ), visual text rendering ( OCR ), and human preference metrics. Image quality and preference scores are computed on DrawBench ( Saharia et al., 2022 ) prompts. Baseline numbers are from Flow-GRPO ( Liu et al., 2025 ) where reported, and from our reimplementation otherwise. Abbreviations: Aes. (Aesthetics), ImgRwd (ImageReward), Pick (PickScore), CLIP (CLIPScore), HPS (HPSv2.1). NFE : number of function evaluations. Few-step RL baselines are quoted from their papers; RTDMD reports only GenEval for its GenEval -only run.

Figure 3: Qualitative comparison on OCR and GenEval test prompts. DualHeads + NB signifies training with the Nano Banana image pool.

		Verifiable metrics		Model-based metrics
Model	NFE	GenEval $\uparrow$	OCR $\uparrow$	HPSv3 ${}^{\ast}\uparrow$	Aes. $\uparrow$	ImgRwd $\uparrow$	Pick $\uparrow$	CLIP $\uparrow$	HPSv2.1 $\uparrow$
Z-Image (base; Z-Image Team et al., 2025 )	100	0.66 †	0.74 †	9.93	5.36	0.92	22.39	28.47	0.285
Z-Image-Turbo (DMDR; Jiang et al., 2026 )	4	0.73 †	0.78 †	11.23	5.42	0.96	22.86	27.42	0.296
TDM-R1 ( Luo et al., 2026a )	4	0.77 †	0.79 †	12.67	5.39	1.16	22.82	28.52	0.307
DualHeads	2	0.76	0.64	12.91	5.27	1.10	22.77	27.31	0.307
DualHeads	4	0.74	0.75	13.93	5.36	1.07	23.03	27.12	0.314

Table 2: Reinforcing Z-Image ( Z-Image Team et al., 2025 ) with HPSv3 ( Ma et al., 2025 ) as the training reward ( ∗ ). Z-Image-Turbo and TDM-R1 ( Luo et al., 2026a ) are evaluated at

1024^{2}

, except the entries marked † , which are reported in Table 3 of TDM-R1 (at

512^{2}

, with Z-Image-Turbo at 4 NFE); the DualHeads rows are EMA weights of one run (reward weight 25) after 120 training iterations, sampled with two and four steps. NFE : number of function evaluations; GenEval : official overall accuracy; OCR : mean OCR score; ImgRwd : ImageReward; CLIP : CLIPScore. The best entry per column is in bold ; shaded rows denote our method.

Panel of Figure 4 — Figure 4: Ablations on OCR. Evaluation OCR score during training, (a) for several values of $\beta_{\mathrm{eq}}$ with a warmed-up reward head, and (b) for a warmed-up versus a freshly initialized reward head.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

		OCR	GenEval	Multi-reward	HPSv3
Model	Base model	SD3.5-M	SD3.5-M	SD3.5-M	Z-Image
Model	LoRA rank / $\alpha$	$32$ / $64$
Sampling	Resolution	$512^{2}$
	Sampler; steps; CFG	LCM; $4$ ; $1$
	Time grid	$(0,0.1,0.25,0.5)$
Rollouts	Group size $K$	$24$

Appendix

Table 3: Training configuration of the main runs for each reward.

Method	NFE	Aes. $\uparrow$	ImgRwd $\uparrow$	Pick ${}^{\ast}\uparrow$	CLIP ${}^{\ast}\uparrow$	HPS ${}^{\ast}\uparrow$
SD3.5-M (base)	80	5.39	0.87	22.34	27.99	0.279
SD3.5-M w/o CFG †	80	5.15	$-$ 0.48	20.66	23.91	0.211
SD3.5-M w/ CFG †	80	5.38	0.93	22.46	28.80	0.282
Flow-GRPO ( Liu et al., 2025 ) w/ CFG †	80	5.28	1.02	22.48	29.12	0.266
DiffusionNFT ( Zheng et al., 2026 ) w/o CFG †	40	5.96	1.41	23.64	28.89	0.324
Choi et al. (2026b) ‡	40	6.06	1.45	23.68	29.60	0.325

Appendix

Table 4: Comparison on multi-reward training with PickScore, CLIPScore and HPSv2.1 as training rewards ( ∗ ). Image quality and preference scores are computed on DrawBench ( Saharia et al., 2022 ) prompts. Rows marked † are quoted from Table 1 of MeanFlowNFT ( Huang et al., 2026b ) , and the row marked ‡ from Choi et al. (2026b) ; the other rows are our evaluations. Abbreviations: Aes. (Aesthetics), ImgRwd (ImageReward), Pick (PickScore), CLIP (CLIPScore), HPS (HPSv2.1).

Diffusion and flow-matching models
Model	Single Obj.	Two Obj.	Counting	Colors	Position	Attr. Bind.	Overall
SD-XL ( Podell et al., 2024 )	0.98	0.74	0.39	0.85	0.15	0.23	0.55
FLUX.1 Dev ( Black Forest Labs, 2024 )	0.98	0.81	0.74	0.79	0.22	0.45	0.66
SD3.5-M ( Esser et al., 2024 )	0.98	0.78	0.50	0.81	0.24	0.52	0.63
TDM-SD3.5-M ( Luo et al., 2025 ) (4 NFE)	0.99	0.77	0.49	0.79	0.23	0.44	0.61
Multi-step RL (SD3.5-M backbone)

Appendix

Table 5: GenEval results by category. Baseline numbers are taken from Flow-GRPO ( Liu et al., 2025 ) or the original papers, and reimplemented where unavailable. Among the RL-tuned models, the best score per column is in bold and the second best is underlined ; the shaded row is our method. Obj.: object; Attr. Bind.: attribute binding.

Figure 5: Visualization of samples from the ablation in Fig. 4(a) .

Panel of Figure 6 — Figure 6: Visual comparison of four-step and two-step generation from the four-step generators trained with the OCR reward (top) and the GenEval reward (bottom).

Figure 7: Effect of the reference pool. Same prompts with base-model (top) and Nano Banana (bottom) references: the style changes, the content does not.

Figure 8: Qualitative comparison for Z-Image-based models.

Figure 9: Uncurated samples on GenEval test prompts.

Figure 10: Uncurated Z-Image samples after HPSv3 training.

Figure 11: Uncurated samples on DrawBench prompts after multi-reward training (PickScore + CLIPScore + HPSv2.1).

OCR
			All samples			References
Source	Prompts	Samples	Mean $r$	$r=1$	$r\leq 0$	Kept	Mean $r$	$r=1$
Nano Banana	19,621	19,625	0.63	51.9%	28.8%	61.6%	0.97	84.3%
Teacher	5,000	40,000	0.57	21.2%	20.5%	96.8%	0.98	81.1%
GenEval
Nano Banana	33,199	50,000	0.47	58.0%	20.2%	63.3%	0.93	81.9%

Appendix

Table 6: Reward statistics of the reference data. All samples: every generated image (for the teacher, the first eight draws per prompt). References: each prompt’s highest-reward sample with

r\geq 0.6

(OCR) or

r\geq 0.5

(GenEval); Kept is the share of prompts that have one.

Figure 12: Random Nano Banana samples. Images for the OCR (top) and GenEval (bottom) training prompts, drawn at random. Above each image is its reward $r$ , the OCR reward or the dense GenEval score; below is the target text (OCR) or the prompt (GenEval).

Figure 13: Random teacher-simulated samples. SD3.5-M samples for the OCR (top) and GenEval (bottom) training prompts, drawn at random; annotations as in Fig. 12 .

Figure 14: Lowest-reward Nano Banana samples. OCR (top): eight images drawn at random among those with $r=0$ , the minimum, which about 29% of images reach. GenEval (bottom): the eight images with the lowest dense GenEval score.

Figure 15: Lowest-reward teacher-simulated samples. OCR (top): eight samples drawn at random among those with $r=0$ , which about 20% of samples reach. GenEval (bottom): the eight lowest dense GenEval scores, at most one sample per prompt.

Figure 16: Ablations of the heads and of the generator loss type on OCR.

Diffusion-Based Image Generation

Diffusion-based image generation creates images by progressively denoising noise, with controls for content, style, and structure.

See all 328 papers

Reinforcing Few-step Generators via Reward-Tilted Distribution Matching
May 25, 2026Yushi Huang, Xiangxin Zhou, Ruoyu Wang +3Diffusion Model Alignment Diffusion Model Distillation
Tencent Hunyuan · Hong Kong University of Science and Technology · Westlake University
TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward
Mar 8, 2026Yihong Luo, Tianyang Hu, Weijian Luo +1RL for Diffusion Models Reinforcement Learning
Hong Kong University of Science and Technology · The Chinese University of Hong Kong, Shenzhen · hi-Lab, Xiaohongshu Inc +1
DRM: Diffusion-based Reward Model With Step-wise Guidance
May 25, 2026Jaxon Zhang, Binxin Yang, Hubery Yin +2Diffusion Model Alignment Diffusion Model Guidance
Peking University · Work done during an internship at WeChat Vision, Tencent Inc. · WeChat Vision, Tencent Inc.

May 25, 2026J/K move · Enter open · S save

Reinforcing Few-step Generators via Reward-Tilted Distribution Matching

Yushi Huang, Xiangxin Zhou, Ruoyu Wang, Chi Zhang+2

Tencent Hunyuan · Hong Kong University of Science and Technology · Westlake University

Diffusion Model Alignment Diffusion Model Distillation Few-Step Diffusion Sampling RL for Diffusion Models Distribution Matching Distillation T2I Diffusion Models

Open paper AskPDF arXiv

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.

Abstract

Diffusion-Based Image Generation

Explore similar work

Reinforcing Few-step Generators via Reward-Tilted Distribution Matching

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

DRM: Diffusion-based Reward Model With Step-wise Guidance