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

   

This is the Official Repository of "TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward", by Yihong Luo, Tianyang Hu, Weijian Luo, Jing Tang.

Samples generated by TDM-R1 using only 4 NFEs, obtained by reinforcing the recent powerful Z-Image model.

Table of Contents

• Pre-trained Model
• Usage
• Start Training
• Install
• Download models
• Launch training
• Hyper-parameter Recipes
  • Core Loss Coefficients
• Acknowledgement
• Contact
• Bibtex

🔥News

• (2026/05) The training code is released.
• (2026/04) TDM-R1 is accepted to ICML 2026 🎉!

Pre-trained Model

• TDM-R1-ZImage

Usage

import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
import torch
from diffusers import ZImagePipeline
from peft import LoraConfig, get_peft_model
def load_ema(pipeline, lora_path, adapter_name='default'):
    """Load EMA weights into the pipeline's transformer adapter"""
    pipeline.transformer.set_adapter(adapter_name)
    trainable_params = [
        p for n, p in pipeline.transformer.named_parameters()
        if adapter_name in n and p.requires_grad
    ]
    state_dict = torch.load(lora_path, map_location=pipeline.transformer.device)
    ema_params = state_dict["ema_parameters"]
    assert len(trainable_params) == len(ema_params), \
        f"Parameter count mismatch: {len(trainable_params)} vs {len(ema_params)}"
    for param, ema_param in zip(trainable_params, ema_params):
        param.data.copy_(ema_param.to(param.device))
    print(f"Loaded EMA weights for adapter '{adapter_name}' from {lora_path}")
pipeline = ZImagePipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
)
transformer_lora_config = LoraConfig(
    r=32,
    lora_alpha=64,
    init_lora_weights="gaussian",
    target_modules=["to_q", "to_k", "to_v", "to_out.0", "add_k_proj", "add_v_proj"],
)
pipeline.transformer = get_peft_model(
    pipeline.transformer,
    transformer_lora_config,
    adapter_name="tdmr1",
)
load_ema(
    pipeline,
    lora_path="./tdmr1_zimage_ema.ckpt",
    adapter_name="tdmr1",
)
pipeline = pipeline.to("cuda")
image = pipeline(
      prompt=prompt,
      height=1024,
      width=1024,
      num_inference_steps=5,  # This actually results in 4 DiT forwards
      guidance_scale=0.0, 
      generator=torch.Generator("cuda").manual_seed(xxx),
  ).images[0]
image

Start Training

Reference implementation that reinforces SD3.5-Medium with the TDM-R1 objective. The released TDM-R1-ZImage checkpoint above uses the same recipe on top of Z-Image-Turbo.

1. Install

git clone https://github.com/Luo-Yihong/TDM-R1.git
cd TDM-R1
pip install -e .

2. Download models

bash download_models.sh

This pulls the SD3.5-Medium base model plus the reward backbones used by the launch scripts under scripts/single_node/ (ImageReward, HPSv2, PaddleOCR, Mask2Former + CLIP for GenEval) into /root/models/ and the standard cache dirs. Set MODELS_DIR=... to redirect.

3. Launch training

Pick the per-task script that matches the reward you want to optimize:

bash scripts/single_node/tdmr1_ocr.sh           # OCR reward
bash scripts/single_node/tdmr1_imagereward.sh   # ImageReward
bash scripts/single_node/tdmr1_geneval.sh       # GenEval

4. Hyper-parameter Recipes

Core Loss Coefficients

TDM-R1 is controlled by a small set of keys spread across config.train.*, config.sample.*, and the top-level config.* namespace (see config/base.py and per-task presets in config/tdmr1_clean.py):

Key	Meaning
config.train.beta_dpo	DPO beta on the group-preference loss;
config.train.beta	Penalty weight on the surrogate reward against the frozen model. 0 disables the KL term.
config.clip_range	PPO-style clip range.
config.rl_cfg	CFG scale used in the RL loss.
config.train.tdm_weight	Mix between the TDM loss and the RL loss; Control the regularization strength of generator.

We provide a default configuration for the key hyper-parameters, which we found to perform reasonably well in most scenarios.

# excerpt from config/tdmr1_clean.py
config.train.beta_dpo        = 1.0
config.train.beta            = 0.001
config.clip_range            = 1e-3 # or 2e-3, 5e-4
config.train.tdm_weight      = 0.3 # or 0.2 0.4

Alternatively, you can apply stronger regularization for more stable training and better unseen metrics:

# excerpt from config/tdmr1_clean.py
config.train.beta_dpo        = 100.0
config.train.beta            = 0.01
config.clip_range            = 2e-3
config.train.tdm_weight      = 0.3 # or 0.4 0.5

We note that in this public release, we use a frozen reference model for a cleaner implementation, avoiding the need to maintain an additional slow-EMA copy of surrogate reward.

Acknowledgement

Our codebase is largely built upon TDM, DGPO and Flow-GRPO. We thank the authors for their efforts to the open-source codebase.

Contact

Please contact Yihong Luo (yluocg@connect.ust.hk) if you have any questions about this work.

Bibtex

@misc{luo2026tdmr1,
      title={TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward}, 
      author={Yihong Luo and Tianyang Hu and Weijian Luo and Jing Tang},
      year={2026},
      eprint={2603.07700},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2603.07700}, 
}