Schedule On the Fly: Diffusion Time Prediction for Faster and Better Image Generation

TL;DR

This paper introduces TPDM, an adaptive diffusion scheduling method that predicts optimal denoising steps on the fly, improving image quality and efficiency in text-to-image generation.

Contribution

The paper proposes a plug-and-play Time Prediction Module trained with reinforcement learning to adaptively determine diffusion steps, enhancing image quality and reducing computation.

Findings

Achieves higher aesthetic and human preference scores.

Uses around 50% fewer denoising steps.

Improves image quality and efficiency in diffusion models.

Abstract

Diffusion and flow matching models have achieved remarkable success in text-to-image generation. However, these models typically rely on the predetermined denoising schedules for all prompts. The multi-step reverse diffusion process can be regarded as a kind of chain-of-thought for generating high-quality images step by step. Therefore, diffusion models should reason for each instance to adaptively determine the optimal noise schedule, achieving high generation quality with sampling efficiency. In this paper, we introduce the Time Prediction Diffusion Model (TPDM) for this. TPDM employs a plug-and-play Time Prediction Module (TPM) that predicts the next noise level based on current latent features at each denoising step. We train the TPM using reinforcement learning to maximize a reward that encourages high final image quality while penalizing excessive denoising steps. With such an adaptive scheduler, TPDM not only generates high-quality images that are aligned closely with human preferences but also adjusts diffusion time and the number of denoising steps on the fly, enhancing both performance and efficiency. With Stable Diffusion 3 Medium architecture, TPDM achieves an aesthetic score of 5.44 and a human preference score (HPS) of 29.59, while using around 50% fewer denoising steps to achieve better performance.