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

TL;DR

TDM-R1 is an RL post-training paradigm for few-step diffusion students that can incorporate non-differentiable reward signals (binary human preference, OCR accuracy, object counts) — closing the “reward-must-be-differentiable” hole left open by earlier few-step RL methods (Diff-Instruct++, Hyper-SD, RewardFlow). It builds on Trajectory Distribution Matching (TDM): because TDM’s sampling trajectory is deterministic (ODE), a single reward at the clean endpoint provides an unbiased estimate of the reward at every intermediate noisy step, removing the variance penalty that has forced prior diffusion RL to back-propagate through differentiable reward heads. A diffusion-parameterized Surrogate Reward, trained group-wise (Bradley–Terry on positive/negative noisy groups with an EMA reference), supplies the per-step differentiable signal the generator actually optimizes against. On SD3.5-M the recipe pushes GenEval from 0.61 (4-NFE TDM base) to 0.92 — surpassing the 80-NFE SD3.5-M base (0.63), GPT-4o (0.84), and matching the previous SOTA RL-on-many-step diffusion (Flow-GRPO 0.95, DGPO 0.97) while using 20× fewer NFEs. Scaled to the recent 6B Z-Image, TDM-R1 also beats both its 100-NFE and Turbo (4-NFE) variants on HPSv3, GenEval, and five out-of-domain quality metrics. ICML 2026.

Key claims

Existing RL methods for few-step diffusion (Diff-Instruct++, Hyper-SD, RewardFlow) require the reward to be differentiable so that the student’s output can backprop through it; this excludes the majority of real-world rewards (binary human preference, OCR correctness, object count) [§1].
For a deterministic (Dirac) student trajectory, a single reward at the clean endpoint x_0 is an unbiased single-sample estimate of the marginal reward $r(x_t \mid c) = \mathbb{E}_{x_0 \sim p(x_0 \mid x_t, c)}[r(x_0, c)]$ for every intermediate $x_t$ on the trajectory — variance vanishes because the posterior is a delta [§3.1, Eq. 6]. Stochastic (SDE) trajectories introduce posterior variance that biases the per-step estimate; this formally explains Mix-GRPO’s observation that converting SDE steps into ODE steps in RL rollout improves performance.
Directly applying a standard diffusion-RL loss (DGPO-style) on top of a TDM distillation loss produces blurry outputs and degrades in late training — the denoising-style RL objective is structurally incompatible with reverse-KL minimization at few steps [§3.2, Fig. 4, Fig. 5].
The fix is a Surrogate Reward parameterized by a diffusion model and trained with group-preference optimization on Bradley–Terry pairs of (positive group, negative group) noisy samples; positivity/negativity is decided per-trajectory by the advantage w.r.t. the group mean reward [§3.2, Eq. 8, Eq. 9].
An EMA reference model in the BT objective relaxes the KL regularizer more gracefully than either a frozen reference or a hard-reset reference at fixed intervals; the latter risks freezing in a model that has overfit to noisy intermediate signals [§3.2, “Dynamic Reference Model”].
The generator loss is reward-maximization on the Surrogate Reward plus a marginal-level reverse-KL to the pretrained TDM teacher; this is weaker than the trajectory-level KL used by standard diffusion RL (which constrains every point along the path) and is therefore easier to satisfy at few steps [§3.3, Eq. 10–11].
On SD3.5-M, 4-NFE TDM-R1 reaches GenEval 0.92 (vs 0.61 for 4-NFE TDM base, 0.63 for 80-NFE SD3.5-M base, 0.84 for GPT-4o); Flow-GRPO @ 80-NFE is 0.95 and DGPO @ 80-NFE is 0.97 — TDM-R1 is within 3–5 points at 20× lower NFE [Table 1].
TDM-R1 also improves image-quality metrics (Aesthetic, DeQA, ImageReward, PickScore, UnifiedReward) above both the few-step base and the 80-NFE base, while Flow-GRPO and DGPO at 80-NFE show out-of-domain quality regressions vs the base — the reward hacking that has plagued differentiable-reward RL on diffusion does not manifest here [Table 2].
Training on one verifiable metric (GenEval or OCR) transfers to the other and to out-of-domain quality metrics — suggesting that a well-chosen verifiable proxy can broadly improve instruction-following [§4.2].
Scaling to the 6B Z-Image with HPSv3 as reward, TDM-R1-Zimage at 4 NFE beats both Z-Image @ 100 NFE and Z-Image-Turbo (DMD-R) @ 4 NFE on HPSv3, GenEval, OCR Acc., Aesthetic, ImageReward, PickScore, and UnifiedReward [Table 3].
The Surrogate Reward parameterizing diffusion model that trained TDM-R1 is itself a strong 80-NFE generator — but TDM-R1 at 4 NFE outperforms it on both in-domain and out-of-domain metrics. Improving a model as a reward signal does not require improving it as a generator [§4.2, “Comparison with the Surrogate Reward Model”].
Ablation: deterministic ODE sampling beats stochastic SDE sampling during rollout (faster convergence, higher final reward), exactly as predicted by §3.1 [Fig. 4].
Ablation: dynamic (jointly trained) Surrogate Reward beats using a frozen DGPO-trained reward — the dynamic reward both identifies per-step regions where the student is weak/strong and stays in the student’s evolving output distribution, which a frozen reward cannot do [Fig. 8].

Method

The student is a TDM few-step generator $G_\theta$ (4 NFEs on SD3.5-M; same on Z-Image-Turbo). Three modules are trained jointly: $G_\theta$ , an online “fake-score” network $\phi$ (denoising-score- matching on student samples; standard TDM/DMD machinery), and a Surrogate Reward $\psi$ (also a diffusion network).

Per-step reward estimation. At each training iteration, the current $G_\theta$ produces a group of $G$ samples conditioned on prompt $c$ via the deterministic 4-step trajectory; the entire trajectory (every intermediate $x_t$ ) is retained. The clean reward $r(x_0, c)$ — non-differentiable, e.g. GenEval/OCR accuracy or HPSv3 score — is computed once per trajectory and assigned as the unbiased per-step reward to every $x_t$ along that trajectory (justified by Eq. 6 because the trajectory is deterministic).

Surrogate Reward learning (§3.2). For each prompt’s group of $G$ noisy samples at a given step, advantages $A_g = r_g - \bar{r}_{\text{group}}$ split the group into positive and negative subgroups. The Surrogate Reward $\psi$ is parameterized as in DGPO (diffusion-model-as-reward) but with a group-level Bradley–Terry loss: the positive group should be preferred to the negative group under $\psi$ , where each sample’s contribution is weighted by $|A_g|/\sum |A_{g'}|$ . A tractable upper bound (Eq. 9 of the paper) reduces this to weighted MSE-style terms between $\psi$ ‘s noise prediction and an EMA-reference noise prediction — a diffusion-DPO-style training objective. The EMA reference replaces either a frozen reference (too restrictive, slows reward learning) or a fixed-interval reset reference (risk of locking in an over-fit intermediate model).

Generator learning (§3.3). $G_\theta$ is trained against Surrogate-Reward maximization with a marginal-level reverse-KL to the pretrained TDM teacher: $\mathcal{L}_G = -\mathbb{E}[\psi(x_t, c)] + \beta \cdot \mathrm{KL}\bigl(p_\theta(x_t \mid c) \,\|\, p_{\text{teacher}}(x_t \mid c)\bigr)$ . The KL is implemented via the standard “online fake score” trick: $\nabla_\theta \mathrm{KL} \approx \mathbb{E}[(s_\phi(x_t) - s_{\text{teacher}}(x_t)) \nabla_\theta x_t]$ , where $s_\phi$ is the fake-score network (trained by denoising score matching on student samples) and $s_{\text{teacher}}$ is the frozen TDM teacher. Gradients of $\psi$ are stopped when flowing into the generator (memory and stability).

Three nested losses, alternating updates. The fake-score $\phi$ is trained by DSM on current student samples; the Surrogate Reward $\psi$ is trained by the group-BT objective; the generator $G_\theta$ is trained by the reward + KL objective. This is analogous to a GAN-style adversarial framework on self-generated samples: $\psi$ identifies per-step weak and strong regions of $G_\theta$ ‘s output; $G_\theta$ moves to improve those regions; $\psi$ adapts to keep providing meaningful gradients. The recipe is summarized in Algorithm 1.

Recipes (from the GitHub README). Default config.train.beta_dpo = 1.0, config.train.beta = 0.001, clip_range = 1e-3, tdm_weight = 0.3; a “stronger regularization” preset (beta_dpo=100, beta=0.01) is offered for more stable training and better out-of-domain metrics. Reward backbones for the three released tasks: ImageReward, HPSv2, PaddleOCR, Mask2Former + CLIP for GenEval. Frozen reference replaces the EMA reference in the public release for implementation simplicity (no slow-EMA copy of Surrogate Reward).

Results

SD3.5-M @ 4 NFE, GenEval (Compositional Image Generation, Table 1).

Model	NFE	GenEval	Single	Two	Counting	Colors	Position	Attr. Binding
SD3.5-M (base)	80	0.63	0.98	0.78	0.50	0.81	0.24	0.52
SD3.5-M + Flow-GRPO	80	0.95	1.00	0.99	0.95	0.92	0.99	0.86
SD3.5-M + DGPO	80	0.97	1.00	0.99	0.97	0.95	0.99	0.91
TDM-SD3.5-M (base)	4	0.61	0.99	0.77	0.49	0.79	0.23	0.44
SD3.5-M + TDM-R1 (ours)	4	0.92	1.00	0.96	0.88	0.85	0.93	0.91
GPT-4o (closed)	—	0.84	0.99	0.92	0.85	0.92	0.75	0.61

The 4-NFE TDM-R1 student gains +0.31 GenEval over its own 4-NFE base and +0.29 over the 80-NFE SD3.5-M base — closing 19/20 of the gap to DGPO-@-80-NFE.

SD3.5-M, out-of-domain quality (Table 2).

For both Compositional Image Generation (GenEval-trained) and Visual Text Rendering (OCR-trained):

TDM-R1 (4 NFE) improves Aesthetic, DeQA, ImageReward, PickScore, and UnifiedReward above the 80-NFE base.
Flow-GRPO and DGPO at 80-NFE regress on Aesthetic and DeQA relative to the base (5.39 → 5.25 / 5.31; 4.07 → 4.01 / 4.03) while gaining on ImageReward and PickScore — the classic reward- hacking pattern.
Transfer: a GenEval-trained TDM-R1 lifts OCR Accuracy from 0.55 to 0.59 without OCR being in the reward loop; an OCR-trained TDM-R1 lifts GenEval from 0.61 to 0.67 [§4.2].

Z-Image (6B) with HPSv3 reward, 4 NFE (Table 3).

Model	NFE	GenEval	OCR	HPSv3	Aesthetic	ImageReward	PickScore	UnifiedReward
Z-Image	100	0.66	0.74	7.32	5.35	0.62	19.98	3.64
Z-Image-Turbo (DMD-R)	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

Two notable patterns: (a) the 4-NFE TDM-R1 strictly dominates both the 100-NFE base and the strongest existing 4-NFE student on every column, including the held-out quality metrics; (b) the gain over the already-Turbo’d few-step student is roughly the same magnitude as the Turbo step itself (4 NFE → 4 NFE w/ R1 is the same delta on HPSv3 as 100 NFE → 4 NFE Turbo), suggesting non-differentiable-reward RL is the “second free lunch” after distillation.

Ablations (§4.3).

Direct combination of TDM loss + DGPO RL loss (the obvious baseline) improves early but degrades late and produces blurry images [Fig. 4, Fig. 5].
Dynamic Surrogate Reward beats a frozen DGPO-trained reward in both convergence speed and final reward [Fig. 8].
Distilling a DGPO-RL’d many-step diffusion teacher into a few-step TDM student converges fast early then plateaus, capped by the teacher — TDM-R1 surpasses this baseline by training the reward jointly with the few-step student [Fig. 9].
Deterministic sampling beats stochastic sampling during rollout, confirming the §3.1 unbiased-estimate argument empirically [Fig. 4].

Why it’s interesting

This is the first filed evidence that non-differentiable rewards work directly on few-step diffusion students — collapsing two separate research threads (“RL on many-step diffusion with non-differentiable rewards”, Flow-GRPO / DGPO; and “few-step distillation”, TDM / DMD / ArcFlow: Unleashing 2-Step Text-to-Image Generation via High-Precision Non-Linear Flow Distillation / Transition Matching Distillation for Fast Video Generation / Diversity-Preserved Distribution Matching Distillation for Fast Visual Synthesis) into a single pipeline that pays the 4-NFE inference cost and still gets GenEval ≈ 0.92, GPT-4o-class. For Luma’s few-step image and video serving stack this is directly actionable: the same recipe used to distill a 50-NFE teacher into a 4-NFE student can now be the substrate for verifiable-metric RL post-training, without needing a differentiable reward head or 80-NFE inference during rollout. The deterministic-trajectory unbiased-estimate argument (§3.1) is the theoretically clean part — it explains why TDM’s structure is what lets the RL work, and predicts that any few-step generator with a deterministic sampler (consistency models, MeanFlow, SoFlow, ArcFlow, TMD) is a candidate substrate for the same recipe. Contrasts cleanly with Diversity-Preserved Distribution Matching Distillation for Fast Visual Synthesis (DP-DMD), which also separates per-step objectives in a few-step distillation pipeline but for a different purpose (preserving diversity via a v-prediction first step), and with Causal Forcing: Autoregressive Diffusion Distillation Done Right for High-Quality Real-Time Interactive Video Generation (Causal Forcing), which fixes the initialization of a few-step AR distillation but does not introduce reward signals at all.