TDM-R1: RL for Few-Step Diffusion Models

• TDM-R1 is a reinforcement learning post-training paradigm that integrates surrogate reward modeling with deterministic trajectory matching to optimize few-step diffusion models.
• It decouples reward learning from generator optimization using a reverse-KL based regularization, ensuring sharp, realistic images even with non-differentiable reward signals.
• Empirical evaluations show that TDM-R1 significantly boosts OCR accuracy, compositional generation, and preference alignment, achieving competitive results in a 4-step denoising regime.

Searching arXiv for the TDM-R1 paper and closely related few-step diffusion RL work. TDM-R1 is a reinforcement-learning post-training paradigm for few-step text-to-image diffusion models that is designed to use generic, non-differentiable reward signals, including human binary likeness, object counts, and OCR-based text-rendering rewards. Built on Trajectory Distribution Matching (TDM), it decouples optimization into surrogate reward learning and generator learning, and exploits the deterministic denoising trajectory of TDM to obtain per-step reward signals without back-propagating through the reward function itself. In the formulation reported in "TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward" (Luo et al., 8 Mar 2026), this yields a unified RL post-training method for few-step models that improves compositionality, text rendering, and preference alignment while preserving the sharpness and realism associated with distribution-level training objectives.

1. Conceptual scope and problem setting

TDM-R1 addresses a specific gap in few-step generative modeling: few-step diffusion models are computationally efficient, but generic RL paradigms for them remained unsolved because existing approaches strongly rely on differentiable reward models. The method is motivated by the observation that many practically important rewards are non-differentiable, including humans' binary likeness, object counts, and OCR accuracy, while few-step models are especially sensitive to loss design (Luo et al., 8 Mar 2026).

The framework is situated in the lineage of distilled diffusion systems in which a student generator produces an image in a small number of denoising steps. The underlying TDM formulation uses a deterministic student trajectory that approximates the teacher diffusion’s stochastic trajectory distribution. TDM-R1 preserves that deterministic trajectory and uses it as the structural basis for RL. This is crucial because denoising-style objectives that are acceptable in many-step settings tend, in the reported experiments, to produce blurry images when naively imposed on few-step generators.

The central design choice is a decoupling. First, a surrogate reward model is learned from black-box terminal rewards. Second, the few-step generator is optimized against that surrogate while remaining close to the teacher through a marginal reverse KL term. This is analogous, in the paper’s own framing, to reward-model learning followed by policy optimization in RLHF-style pipelines, but translated into the trajectory space of few-step diffusion models (Luo et al., 8 Mar 2026).

2. Mathematical structure and relation to TDM

The starting point is the standard forward diffusion parameterization

$$q(x_t \mid x_0) = \mathcal{N}(x_t; \alpha_t x_0, \sigma_t^2 I),$$

with the usual denoising training objective

$$\mathbb{E}_{x_0,\epsilon,t}\,\lambda_t \| f_\psi(x_t,t) - \epsilon \|_2^2.$$

TDM introduces a $K$-step student with deterministic latent states $\{x_{t_1},\dots,x_{t_K}\}$. For step $k$, the student marginal at time $t$ is defined as

$$p_{\theta,k}(x_t) \triangleq \int q(x_t \mid x_{t_k})\, p_\theta(x_{t_k})\, dx_{t_k}.$$

Training aligns these student marginals with teacher marginals through an integral reverse KL. In the score-difference approximation used by TDM, the relevant gradient has the form

$$\nabla_\theta L(\theta) \approx \mathbb{E}_{k,t\ge t_k,q(x_t\mid x_{t_k})} \lambda_t\,[s_{\mathrm{fake}}(x_t)-s_\psi(x_t)]\,\frac{\partial x_{t_k}}{\partial\theta}.$$

TDM-R1 inherits this distribution-matching backbone rather than replacing it with a denoising-based RL loss. This is one of its defining features. The paper argues that few-step generators require distribution-level regularization, especially reverse KL to a strong teacher, because direct weighted denoising losses degrade visual fidelity in low-NFE regimes (Luo et al., 8 Mar 2026).

3. Surrogate reward learning on noisy trajectory states

A central difficulty in diffusion RL is that rewards are naturally defined on the clean image $x_0$, whereas optimization occurs along noisy intermediate states. TDM-R1 formalizes a noisy-state reward

$r(x_t,c) = \mathbb{E}_{p(x_0 \mid x_t)}[r(x_0,c)].$

In generic stochastic settings this is a high-variance object. In TDM-R1, however, the generation trajectory is deterministic, so the mapping from $$\mathbb{E}_{x_0,\epsilon,t}\,\lambda_t \| f_\psi(x_t,t) - \epsilon \|_2^2.$$0 to $$\mathbb{E}_{x_0,\epsilon,t}\,\lambda_t \| f_\psi(x_t,t) - \epsilon \|_2^2.$$1 becomes effectively Dirac-like along the learned path. The paper uses this to justify low-variance per-step reward assignment from terminal non-differentiable rewards (Luo et al., 8 Mar 2026).

The surrogate reward is parameterized as $$\mathbb{E}_{x_0,\epsilon,t}\,\lambda_t \| f_\psi(x_t,t) - \epsilon \|_2^2.$$2 and is implemented through a diffusion-like likelihood-ratio model relative to a reference process. In the noisy-state formulation reported in the paper,

$$\mathbb{E}_{x_0,\epsilon,t}\,\lambda_t \| f_\psi(x_t,t) - \epsilon \|_2^2.$$3

After Markov factorization, this becomes a sum of transition-level log-ratios.

Training of $$\mathbb{E}_{x_0,\epsilon,t}\,\lambda_t \| f_\psi(x_t,t) - \epsilon \|_2^2.$$4 uses group-based preference optimization. For each prompt, the generator samples a group of trajectories, terminal rewards are computed, and standardized group advantages are formed: $$\mathbb{E}_{x_0,\epsilon,t}\,\lambda_t \| f_\psi(x_t,t) - \epsilon \|_2^2.$$5 Positive and negative subsets are then compared through a Bradley–Terry probability over group rewards. The resulting objective encourages the surrogate model to assign higher relative likelihood to noisy ancestors of high-reward samples and lower relative likelihood to those of low-reward samples. The reference model is not static: TDM-R1 uses an EMA-based dynamic reference, which the ablations report as more stable than a fixed reference (Luo et al., 8 Mar 2026).

4. Generator optimization and training loop

Once the surrogate reward exists, generator learning proceeds by combining reward maximization with TDM-style reverse-KL regularization. The paper writes the objective as

$$\mathbb{E}_{x_0,\epsilon,t}\,\lambda_t \| f_\psi(x_t,t) - \epsilon \|_2^2.$$6

where $$\mathbb{E}_{x_0,\epsilon,t}\,\lambda_t \| f_\psi(x_t,t) - \epsilon \|_2^2.$$7 denotes stop-gradient through the surrogate during generator updates.

The corresponding gradient has two terms: a reward-gradient term derived from the surrogate log-likelihood ratio and the original TDM reverse-KL term involving $$\mathbb{E}_{x_0,\epsilon,t}\,\lambda_t \| f_\psi(x_t,t) - \epsilon \|_2^2.$$8. The $$\mathbb{E}_{x_0,\epsilon,t}\,\lambda_t \| f_\psi(x_t,t) - \epsilon \|_2^2.$$9 network is an online score model trained on synthetic student samples, exactly to approximate the student score needed in the TDM gradient. This preserves the original few-step distribution-matching mechanism instead of replacing it with a reward-weighted denoising proxy (Luo et al., 8 Mar 2026).

Algorithmically, the system alternates among four operations. A prompt is sampled. The current few-step generator produces a group of deterministic trajectories and terminal images. A black-box reward function scores the terminal images. Those grouped trajectories are then used to update the surrogate reward model, update the fake-score model, and finally update the few-step generator. This loop makes TDM-R1 an online RL scheme in which the reward model and generator co-evolve on the generator’s own data distribution.

5. Empirical performance

The paper reports evaluations on compositional generation, visual text rendering, and preference alignment, with both in-domain and out-of-domain metrics (Luo et al., 8 Mar 2026). The most direct comparisons are summarized below.

Setting	Baseline	TDM-R1
TDM-SD3.5-M, GenEval overall, 4 NFE	0.61	0.92
TDM-SD3.5-M, OCR accuracy, 4 NFE	0.55	0.95
SD3.5-M, GenEval overall, 80 NFE	0.63	0.92 on TDM-R1 4 NFE
GPT-4o, GenEval overall	0.84	0.92 on TDM-R1 4 NFE
Z-Image, GenEval	0.66	0.77
Z-Image, HPSv3	7.32	9.90

On GenEval, TDM-R1 raises a 4-NFE TDM-SD3.5-M model from 0.61 to 0.92 overall, with the reported category breakdown 1.00 for single object, 0.96 for two object, 0.88 for counting, 0.85 for colors, 0.93 for position, and 0.91 for attribute binding. In the same table, Flow-GRPO on SD3.5-M at 80 NFE is listed at 0.95 and DGPO at 0.97, so TDM-R1 approaches the strongest many-step RL baselines while remaining in a 4-NFE regime (Luo et al., 8 Mar 2026).

On visual text rendering, the paper reports OCR accuracy improving from 0.55 to 0.95 for the 4-NFE TDM-SD3.5-M setting. The OCR-related reward is based on

$K$0

where $K$1 is edit distance and $K$2 is the reference-string length. The paper also notes cross-task generalization: GenEval-trained TDM-R1 improves OCR, and OCR-trained TDM-R1 improves GenEval (Luo et al., 8 Mar 2026).

For Z-Image, TDM-R1 is reported to scale effectively to the recent strong model, consistently outperforming both its 100-NFE and few-step variants with only 4 NFEs. The reported HPSv3-aligned TDM-R1-Z-Image obtains GenEval 0.77, OCR 0.79, HPSv3 9.90, Aesthetic 5.49, ImageReward 0.94, PickScore 20.45, and UnifiedReward 3.75. This suggests that the method is not confined to one distilled SD3.5-based backbone but transfers to a larger, more recent diffusion-transformer setting (Luo et al., 8 Mar 2026).

6. Ablations, limitations, and place in the literature

The ablation studies define the paper’s central negative result: directly combining TDM with a diffusion RL loss such as DGPO leads to initial reward gains followed by blurred outputs and degraded performance. The paper attributes this to a mismatch between denoising-style RL objectives and the reverse-KL-based distribution matching required for few-step sharpness. TDM-R1’s decoupled surrogate-reward design is presented as the remedy (Luo et al., 8 Mar 2026).

Additional ablations report that deterministic trajectories outperform stochastic ones, and that a dynamic EMA reference outperforms a static reference. A separate comparison shows that distilling from an already RL-trained many-step teacher improves early learning but saturates earlier than TDM-R1, because the jointly updated surrogate can continue adapting to the few-step student’s evolving distribution.

The method’s limitations are also structurally clear. Training requires the simultaneous maintenance of a few-step generator, a surrogate diffusion reward model, and an online fake-score model. Reward misspecification remains possible, even though the paper monitors out-of-domain metrics such as Aesthetic Score, DeQA, ImageReward, PickScore, and UnifiedReward. The article also suggests several future directions, including extension to video and other modalities, richer reward compositions, and more refined trajectory-level credit assignment (Luo et al., 8 Mar 2026).

Within the broader few-step diffusion literature, TDM-R1 is best understood as a method that translates RLHF-style decoupling into the deterministic trajectory space of TDM. Its defining claim is not merely that non-differentiable rewards can be used, but that they can be used without sacrificing the distribution-level objectives that make few-step image synthesis viable. This suggests that TDM-R1 occupies a distinct position among diffusion RL methods: it is an RL framework specialized to the pathological sensitivity of low-NFE generators rather than a straightforward adaptation of many-step diffusion RL.