Learning to Schedule in Diffusion Probabilistic Models

Recently, the field of generative models has seen a significant advancement with the introduction of Diffusion Probabilistic Models (DPMs). The Denoising Diffusion Implicit Model (DDIM) was designed to reduce computational time by skipping a number of steps in the inference process of DPMs. However, the hand-crafted sampling schedule in DDIM, which relies on human expertise, has its limitations in considering all relevant factors in the sampling process. Additionally, the assumption that all instances should have the same schedule is not always valid. To address these problems, this paper proposes a method that leverages reinforcement learning to automatically search for an optimal sampling schedule for DPMs. This is achieved by a policy network that predicts the next step to visit based on the current state of the noisy image. The optimization of the policy network is accomplished using an episodic actor-critic framework, which incorporates reinforcement learning. Empirical results demonstrate the superiority of our approach over various datasets with different timesteps. We also observe that the trained sampling schedule has a strong generalization ability across different DPM baselines.