Learning To Sample From Diffusion Models Via Inverse Reinforcement Learning

2026-02-09

AI Summary

This work investigates how to enhance the quality and efficiency of image generation from pre-trained diffusion models without retraining their denoising networks. To this end, the diffusion sampling process is formulated as a finite-horizon Markov decision process, and inverse reinforcement learning is introduced for the first time to directly infer an optimal action scheduling policy from desired sampling behaviors—bypassing the need for explicit reward function design. By leveraging this approach, the method significantly improves sample fidelity while enabling automatic tuning of sampling hyperparameters, all while preserving the original denoising network unchanged.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

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.

Problem

• diffusion models
• sampling strategies
• inverse reinforcement learning
• sample quality
• hyperparameter tuning

Innovation

• inverse reinforcement learning
• diffusion sampling
• Markov Decision Process
• policy gradient
• denoising diffusion models