Research map / Image generation / Fast diffusion sampling
Learning to Sample From Diffusion Models via Inverse Reinforcement Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Frames diffusion sampler tuning (guidance, stochasticity injection, renoising) as a finite-horizon MDP and learns a sampling policy with state-marginal f-divergence policy gradients (PPO-style) while keeping the EDM denoiser frozen.
Problem
Sampler hyperparameters of diffusion models (guidance scales, noise schedules, stochasticity) are tuned by costly grid search; the goal is to learn them automatically without retraining the denoiser and without expert trajectories or rewards.
Approach
Cast diffusion sampling as a discrete-action MDP and minimize an f-divergence between expert and policy state occupancy measures (expert specified by data marginals at noise levels) using a policy gradient with learning signal h_f of density ratios estimated by a discriminator, avoiding explicit reward learning.
Key results
- On ImageNet-64, a single learned-sampler training run costs 160.6 EFLOPs versus 1437.4 EFLOPs for an 8^3 three-hyperparameter grid search (up to 9x lower cost), with 16% inference overhead.
- Adaptive CFG with KL on CIFAR-10 reaches FID 2.02 and Recall 71.9 versus EDM+CFG at FID 2.07 and Recall 70.9.
- For renoising on ImageNet-64, KL reduces NFE from 350 (EDM renoise baseline) to 51.8 with FID 2.96 vs 3.01.
Surprising findings
- Using a pretrained ADM encoder as a frozen feature extractor matters: removing it raises FID from 2.05 to 3.04.
- On FFHQ renoising, the deterministic EDM sampler remains stronger in FID than learned renoising policies.
Limitations
- Depends on discrete action-space design and is sensitive to policy initialization (heuristic init required).
- Renoising underperforms the deterministic EDM sampler in FID on FFHQ.
- Evaluated only on 32x32 and 64x64 EDM models; extension to stronger pretrained models left to future work.
Experimental setup
Datasets: CIFAR-10 FFHQ ImageNet 64x64
Benchmarks: CIFAR-10 FFHQ 64x64 ImageNet 64x64
Metrics: FID Precision Recall NFE EFLOPs
Baselines: EDM EDM+CFG grid search EDM stochasticity hyperparameters f-PG
Builds on: EDM PPO GAIL f-PG classifier-free guidance
Closest papers by meaning
- Hierarchical Variational Policies for Reward-Guided Diffusion (NeurIPS 2026 · similarity 0.91)
- Training Diffusion Models with Reinforcement Learning (ICLR 2024 · similarity 0.91)
- A Variational Perspective on Solving Inverse Problems with Diffusion Models (ICLR 2024 · similarity 0.90)
- Prior-Guided Diffusion Planning for Offline Reinforcement Learning (NeurIPS 2025 · similarity 0.90)
- Maximum Entropy Inverse Reinforcement Learning of Diffusion Models with Energy-Based Models (NeurIPS 2024 · similarity 0.90)
- Towards Better Alignment: Training Diffusion Models with Reinforcement Learning Against Sparse Rewards (CVPR 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.