Learning To Sample From Diffusion Models Via Inverse Reinforcement Learning
This work addresses a practical bottleneck in diffusion models for researchers and practitioners by enabling more flexible and efficient sampling, though it is incremental as it builds on existing diffusion and reinforcement learning methods.
The paper tackles the problem of improving sample quality and efficiency in diffusion models without retraining the denoiser by introducing an inverse reinforcement learning framework that formulates sampling as a Markov Decision Process and uses policy gradient techniques to optimize action scheduling, resulting in improved sample quality and automatic hyperparameter tuning.
Diffusion models generate samples through an iterative denoising process, guided by a neural network. While training the denoiser on real-world data is computationally demanding, the sampling procedure itself is more flexible. This adaptability serves as a key lever in practice, enabling improvements in both the quality of generated samples and the efficiency of the sampling process. In this work, we introduce an inverse reinforcement learning framework for learning sampling strategies without retraining the denoiser. We formulate the diffusion sampling procedure as a discrete-time finite-horizon Markov Decision Process, where actions correspond to optional modifications of the sampling dynamics. To optimize action scheduling, we avoid defining an explicit reward function. Instead, we directly match the target behavior expected from the sampler using policy gradient techniques. We provide experimental evidence that this approach can improve the quality of samples generated by pretrained diffusion models and automatically tune sampling hyperparameters.