Research map / Image generation / Fast diffusion sampling
Schedule On the Fly: Diffusion Time Prediction for Faster and Better Image Generation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes TPDM, a plug-and-play Time Prediction Module that adaptively predicts the next diffusion time per step for frozen flow-matching models, trained with PPO and an ImageReward-based step-discounted reward to cut denoising steps by about 50%.
Problem
Diffusion and flow-matching text-to-image models use a predetermined denoising schedule for all prompts, which is suboptimal for varying prompt complexity and wastes steps on simple prompts.
Approach
Treat the diffusion time schedule as an adaptive chain-of-thought: a lightweight Time Prediction Module predicts a Beta-distributed decay rate per step from latent features, and is trained with RL (PPO/RLOO-style) on the final image reward discounted by the number of steps, while the base diffusion model stays frozen.
Key results
- TPDM-SD3-Medium reaches HPSv2.1 29.59 and Aesthetic v2 5.445 with 15.28 average steps, versus 29.12 HPS for SD3-Medium at 28 steps.
- TPDM-FLUX.1-dev uses 13.57 average steps with Aesthetic v2 5.685 versus 5.622 for FLUX.1-dev at 28 steps.
- In a user study with 15 volunteers, TPDM-SD3-Medium won 47.25% of comparisons versus 26.58% for SD3-Medium at 28 steps.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use adaptive time prediction | Stable Diffusion 3 Medium · COCO 2017 validation | HPSv2.1: 29.12 -> 29.59 | Helped |
| Use adaptive time prediction | Stable Diffusion 3 Medium · COCO 2017 validation | Aesthetic v2: 5.433 -> 5.445 | Helped |
| Use adaptive time prediction | Stable Diffusion 3 Medium · COCO 2017 validation | Inference steps: 28 -> 15.28 | Helped |
| Use adaptive time prediction | Stable Diffusion 3.5 Large · COCO 2017 validation | Aesthetic v2: 5.487 -> 5.525 | Helped |
| Use adaptive time prediction | FLUX.1-dev · COCO 2017 validation | CLIP-T: 0.308 -> 0.314 | Helped |
| Use adaptive time prediction | FLUX.1-dev · COCO 2017 validation | HPSv2.1: 31.94 -> 30.77 | Hurt |
| Use first and last layer features instead of only first 2 layers as TPM inputs (ablation) | TPDM-SD3-Medium · COCO 2017 validation | Aesthetic v2: 5.400 -> 5.445 | Helped |
| Use first and last layer features instead of only last 2 layers as TPM inputs (ablation) | TPDM-SD3-Medium · COCO 2017 validation | Aesthetic v2: 5.356 -> 5.445 | Helped |
| Use first and last layer features instead of only last 2 layers as TPM inputs (ablation) | TPDM-SD3-Medium · COCO 2017 validation | Inference steps: 19.30 -> 15.28 | Helped |
Surprising findings
- Naive training that supervises loss after two steps causes the model to finish denoising in very few steps and degrades image quality, motivating the whole-trajectory RL objective.
Limitations
- The Time Prediction Module architecture is simple and may be improved with more delicate designs.
- The diffusion backbone is frozen during training; joint updating is left for future work.
- Evaluation relies largely on reward-model-based and human-preference metrics; FID does not improve.
Experimental setup
Datasets: COCO 2017 Laion-Art COYO-700M HPSv2 prompts
Benchmarks: COCO 2017 validation (5,000 prompts) HPSv2.1 benchmark prompts (3,200)
Metrics: FID CLIP-T Aesthetic v2 Pick Score HPSv2.1 user win rate
Baselines: SD3-Medium SD3.5-Large FLUX.1-dev DDPO DPOK Diffusion-DPO Xia et al. scheduler optimization Sabour et al. optimal sampling plan Wang et al. RL scheduler search
Builds on: flow matching PPO RLOO ImageReward DDPO
Closest papers by meaning
- AdaDiff: Adaptive Step Selection for Fast Diffusion Models (AAAI 2025 · similarity 0.89)
- Towards Understanding the Working Mechanism of Text-to-Image Diffusion Model (NeurIPS 2024 · similarity 0.89)
- Denoising Task Difficulty-based Curriculum for Training Diffusion Models (ICLR 2025 · similarity 0.89)
- Training-Free Adaptive Diffusion with Bounded Difference Approximation Strategy (NeurIPS 2024 · similarity 0.88)
- Adaptive Non-Uniform Timestep Sampling for Accelerating Diffusion Model Training (CVPR 2025 · similarity 0.88)
- Faster Diffusion: Rethinking the Role of the Encoder for Diffusion Model Inference (NeurIPS 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.