PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yihong Luo, Tianyang Hu, Weijian Luo, Jing Tang

arXiv:2603.07700v1cs.CVcs.AI

TL;DR

Few-step diffusion RL is limited by reliance on differentiable rewards, leaving important non-differentiable feedback underused. TDM-R1 combines deterministic-trajectory reward estimation with surrogate reward learning and generator optimization, and its 4-step models reach 92% GenEval versus 61% before reinforcement while surpassing an 80-NFE base at 63%.

  • Problem

    Existing few-step diffusion RL methods rely on differentiable rewards, excluding important non-differentiable signals such as human preferences and object counts.

  • Method

    TDM-R1 uses deterministic TDM trajectories for intermediate reward estimation and decouples surrogate reward learning from few-step generator optimization.

  • Results

    92% GenEval improves from 61%, exceeding the 80-NFE base model’s 63%; TDM-R1 also outperforms the 80-NFE surrogate reward model using 4 NFEs.

  • Takeaways & Limitations

    TDM-R1 provides a unified post-training approach for non-differentiable rewards and achieves strong in-domain and out-of-domain few-step text-to-image performance.

  • Takeaways & Limitations

    Directly applying standard diffusion RL methods to few-step models can produce blurry outputs because denoising-loss updates are incompatible with few-step generation.

Abstract

from arXiv · show

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

1. Introduction

Few-step generative models offer rapid, high-fidelity generation, but existing RL methods generally require differentiable rewards, excluding important real-world feedback. TDM-R1 addresses this gap with deterministic trajectory reward estimation and decoupled surrogate-reward and generator learning, achieving strong few-step results.

  • Motivation: Few-step generative models accelerate image and video generation substantially while maintaining high fidelity.The introduction describes accelerations of up to 50 times compared with diffusion models.
  • Problem: Existing RL methods for few-step generative models require differentiable rewards so gradients can backpropagate through generated outputs.
  • Problem: Non-differentiable rewards such as human binary preferences, object counts, and OCR-based text-rendering correctness are therefore important but excluded by this requirement.
  • Approach: TDM-R1 uses deterministic TDM trajectories to estimate rewards for intermediate denoising steps and trains a diffusion-parameterized surrogate reward for stable supervision.The method does not require additional ground-truth image data.
  • Results: 92% GenEval improves from 61%, surpassing the 80-NFE base model at 63% and GPT-4o at 84%.The method uses only 4 NFEs and also scales to Z-Image, outperforming its 100-NFE and few-step variants across in-domain and out-of-domain metrics.

2. Preliminaries

The preliminaries introduce diffusion models, TDM, and reward modeling foundations used by TDM-R1. They cover forward corruption and reverse denoising, trajectory distribution matching, pairwise preference rewards, and RLHF’s reward-regularized generator objective.

  • Diffusion Models: Diffusion models corrupt clean samples with Gaussian noise across T discrete timesteps and learn a reverse denoising process.The supplied preliminaries describe noisy samples and a parameterized reverse transition, though some equations are truncated.
  • Diffusion Models: A trained diffusion model estimates the score of the noisy data distribution using its denoising network.
  • Trajectory Distribution Matching: TDM aligns student and teacher trajectories by minimizing an integral reverse-KL objective at each step of the K-step student trajectory.The student trajectory distribution and pretrained teacher diffusion model are used together, with an online fake score approximating the student score.
  • Reward Modeling: Pairwise reward modeling represents a preferred sample and a less-preferred sample under condition c, with Bradley-Terry preference probability determined by their reward difference.
  • RLHF: RLHF optimizes a conditional generator to maximize a reward while remaining close to a reference distribution through a regularization term.The parameter β controls the regularization strength.

3. Method

TDM-R1 separates surrogate reward learning from few-step generator learning and exploits deterministic TDM trajectories to estimate intermediate rewards accurately. It then combines group-based surrogate optimization, reward maximization, and distribution-level regularization for post-training.

  • Problem setup: TDM-R1 uses an online RL setup in which the current few-step model generates grouped trajectory samples that are scored by the reward function.Samples from each denoising step are collected into training datasets for subsequent reward-based optimization.
  • Accurate intermediate reward estimation: Deterministic TDM trajectories provide unbiased, lower-variance reward estimates for intermediate denoising samples.When the path from a noisy state to the clean image is deterministic, the endpoint reward can estimate rewards along the trajectory without the variance introduced by stochastic transitions.
  • Surrogate reward learning: The method learns a Surrogate Reward because standard diffusion RL methods are incompatible with few-step generation and can produce blurry outputs.The incompatibility arises because such methods often reduce to weighted denoising losses, which tend to blur images with few sampling steps.
  • Surrogate reward learning: Group-based Bradley-Terry learning captures preference relations between positive and negative groups of noisy samples using fine-grained within-group signals.Samples receive weights based on the absolute normalized advantage, emphasizing relatively better or worse members of each group.
  • Training procedure: The generator and surrogate reward are jointly trained in an adaptive loop, with the generator improving samples and the reward model refining its guidance over iterations.The procedure alternates optimization of the generator, surrogate reward, and fake score components.

4. Experiments

TDM-R1 is evaluated on compositional generation, visual text rendering, human preference alignment, and larger-model reinforcement, alongside ablations of its core components. Across these settings, it improves instruction following and quality, including 4-NFE performance that surpasses many-step baselines.

  • Experimental Setup: The evaluation covers compositional image generation, visual text rendering, human preference alignment, and out-of-domain image-quality metrics.GenEval assesses object counting, spatial relationships, and attribute binding; five independent metrics help evaluate quality beyond the training rewards.
  • Main Results: TDM-R1 achieves competitive GenEval performance while improving several out-of-domain quality metrics, unlike standard diffusion RL methods that trade quality for higher GenEval scores.The reported comparisons include Aesthetic Score, DeQA, PickScore, and ImageReward, with TDM-R1 exceeding both many-step and few-step base-model quality metrics.
  • Main Results: Training on GenEval or OCR signals can improve a different verifiable metric, suggesting that proxy-task rewards may enhance broader instruction-following ability.The paper reports this cross-task improvement in both target and out-of-domain metrics for visual text rendering.
  • Main Results: Qualitative comparisons show more accurate instruction following than the few-step baseline and the 80-NFE base model while preserving high generation quality.The comparisons use GenEval-trained models and present matched visual samples.
  • Main Results: TDM-R1 improves human-preference alignment and scales to the 6B-parameter Z-Image model, outperforming both many-step and few-step variants across reported metrics.For Z-Image, the method uses HPSv3 and shows consistent improvement over competing baselines.
  • Main Results: The surrogate reward model itself requires 80 NFE, whereas TDM-R1 samples with 4 NFE and consistently outperforms it on in-domain and out-of-domain metrics.The paper attributes this difference to the quality of per-step reward signals rather than the surrogate model’s standalone generation capability.
  • Ablation Study: Ablations show that dynamic surrogate rewards, deterministic trajectories, and a dynamic reference model improve optimization, while direct RL-loss combination causes blur and later degradation.The direct-loss baseline conflicts with the reverse-KL distillation objective; stochastic sampling converges more slowly, and static references reduce stability.
  • Ablation Study: Compared with RL-distilled diffusion initialization, TDM-R1 avoids teacher-limited plateaus and directly reinforces the trained few-step student for stronger final performance.The alternative also introduces an additional teacher-training and distillation pipeline.

5. Related Works

Prior work on few-step text-to-image diffusion spans acceleration, alignment, and reinforcement learning, but existing few-step RL methods remain restricted mainly to differentiable rewards.

  • Few-Step Text-to-Image Diffusion Sampling: Diffusion distillation remains essential for high-quality few-step generation, with trajectory matching and distribution matching as major strategy categories.
  • Text-Image Alignment: Reinforcement learning has been explored to improve text-image alignment and specialized capabilities in diffusion and few-step generative models.
  • Few-Step Text-to-Image Diffusion RL: Existing few-step diffusion RL methods are restricted to differentiable reward functions, limiting their application to non-differentiable rewards.
  • Few-Step Text-to-Image Diffusion RL: Some methods theoretically support non-differentiable rewards, but their effectiveness at scale has not been validated, while standard diffusion RL can produce blurry, suboptimal few-step generations.

6. Conclusion

TDM-R1 enables few-step diffusion models to use non-differentiable rewards through surrogate reward learning and generator optimization. It achieves state-of-the-art few-step RL performance, including 4-step models surpassing 80-NFE baselines on GenEval.

  • TDM-R1 enables few-step diffusion models to leverage non-differentiable reward feedback by decoupling surrogate reward learning from generator optimization.
  • 4-step TDM-R1 models improve GenEval from 61% to 92%, exceeding the 80-NFE base model’s 63%.
  • TDM-R1 incorporates large-scale online non-differentiable reward signals without requiring additional ground-truth image data.

A.1. Derivation of Eq. (9)

The derivation simplifies the surrogate-reward objective by canceling partition-function terms, factorizing trajectory probabilities, applying Jensen’s inequality, and rewriting log-ratios with KL divergences.

  • Partition-function terms cancel because the weights sum to zero over the full group.
  • The trajectory-level log-ratio is decomposed into transition probabilities using the diffusion process’s Markov property.
  • Jensen’s inequality produces a tractable upper bound by moving timestep and conditional-state expectations inside the sigmoid.
  • The resulting log-ratios are rewritten using KL-divergence identities to obtain the final upper bound.
  • These transformations complete the derivation of Eq. (9).

A.2. Derivation of Eq. (11)

The derivation of Eq. (11) separately computes gradients for surrogate-reward maximization and KL regularization, then combines them into the final objective gradient.

  • The surrogate reward is reformulated by converting its timestep summation into an expectation over a uniform distribution.
  • Marginalizing the conditional trajectory distribution yields the relevant variables for the reformulated reward expression.
  • The reward-term gradient is derived through the reparameterized sampling process and the chain rule.
  • For the Gaussian forward process, x_t is parameterized as α_t|t_k x_tk + σ_t|t_k ε, with ε sampled from N(0, I).
  • The KL-regularization gradient is derived separately, and combining both gradients yields the final expression in Eq. (11).

B. Experiment details

The evaluation measures text rendering fidelity using a standardized sign-prompt protocol and a minimum-edit-distance-based metric. Training and inference use four-step Flow-DPM-Solver rollouts with fixed LoRA and regularization settings.

  • Text-rendering evaluation: Text-rendering evaluation uses the TextDiffuser protocol and Flow-GRPO framework with prompts targeting text inside generated images.Prompts follow the template “A sign that says ‘text’”, where the quoted string is the target content.
  • Text-rendering evaluation: Text fidelity is based on the minimum edit distance between the rendered output and the target string.N_e denotes minimum edit distance, while N_ref denotes the character length of the quoted prompt string.
  • Training configuration: 24 samples per group and four sampling steps are used for training rollouts and inference, with rank-32 LoRA fine-tuning at 512 resolution.The default β is 100, and β_g is chosen so the reward-gradient to KL-gradient ratio is 2:1.

Details of the out-of-domain evaluation metrics

Out-of-domain evaluation covers aesthetic appeal, image-quality degradation, visual quality, text-image coherence, and multimodal preference assessment. The paper also provides qualitative comparisons using identical initial noise across generated images.

  • Quality metrics: The aesthetic score uses a CLIP-based linear regression model to quantify the visual appeal of generated images.
  • Quality metrics: DeQA measures perceived image-quality degradation from distortions, textural degradation, and low-level visual artifacts using a multimodal large language model.
  • Preference and multimodal metrics: ImageReward evaluates visual quality and text-image generation fidelity, while UnifiedReward evaluates multimodal understanding and generation tasks.UnifiedReward is described as a state-of-the-art unified reward framework for these evaluations.
  • Qualitative comparison: Figure 10 compares TDM-R1 with competing baselines under the same initial noise for all generated images.Additional visual comparisons are provided in Fig. 10.
Loading 2603.07700v1…