Schedule On the Fly: Diffusion Time Prediction for Faster and Better Image Generation
This work addresses efficiency and quality issues in diffusion-based image generation for AI practitioners, offering a plug-and-play solution that is incremental but impactful.
The paper tackles the problem of inefficient predetermined noise schedules in diffusion models for text-to-image generation by introducing an adaptive scheduler that predicts optimal noise levels per instance, resulting in a 50% reduction in denoising steps while improving aesthetic and human preference scores to 5.44 and 29.59, respectively.
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.