🤖 machine learning

Learning To Sample From Diffusion Models Via Inverse Reinforcement Learning

This paper introduces an inverse reinforcement learning framework that formulates diffusion model sampling as a Markov Decision Process to automatically learn optimal sampling strategies without retraining the denoiser, achieving performance comparable to fine-tuned samplers at a significantly lower cost than traditional grid search.

Original authors: Constant Bourdrez, Alexandre Vérine, Olivier Cappé

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Constant Bourdrez, Alexandre Vérine, Olivier Cappé

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a master chef (the Denoiser) who is incredibly talented at turning a bowl of plain, scrambled eggs (random noise) into a perfect, gourmet steak (a high-quality image). This chef has already been trained for years and knows exactly how to cook.

However, there's a catch: the chef needs a specific set of instructions on how to cook. Should you stir the pan gently? Should you add a pinch of salt at the very beginning or the very end? Should you turn the heat up or down at specific moments?

In the world of AI image generation, these instructions are called sampling strategies. Traditionally, finding the perfect set of instructions is like trying to find a needle in a haystack. Researchers have to manually test thousands of different combinations of heat, timing, and seasoning (a process called "grid search"). This is incredibly expensive and time-consuming, often costing as much energy as training the chef in the first place.

The Paper's Big Idea
The authors of this paper propose a new way to teach the chef how to cook without retraining the chef or hiring a new one. Instead of manually guessing the instructions, they use a method called Inverse Reinforcement Learning (IRL).

Think of it this way:

The Old Way: You try to guess the perfect recipe by tasting 1,000 different versions of the dish and hoping one is good.
The New Way: You watch the chef cook a few times. You don't tell the chef what to do; you just show the chef the final, perfect steak. The chef then learns, "Ah, to get a result that looks like this, I need to stir here and add salt there."

How It Works (The Metaphor)
The paper treats the image generation process as a long journey or a video game level.

The Journey: The AI starts with pure static (noise) and slowly transforms it into an image. This happens step-by-step, like moving through levels in a game.
The Player: The "policy" is the player controlling the game. At every step, the player can choose to do things like:
- Add a little chaos (Stochasticity): Shake the pan a bit to see if it helps.
- Give a nudge (Guidance): Tell the chef, "Make it look more like a dog," at specific moments.
- Hit the rewind button (Renoising): If the image looks weird, go back a few steps and try again.
The Goal: The player doesn't get a score for every move. Instead, the player's only goal is to make sure the final destination (the finished image) looks exactly like the target data (the real photos).

The AI learns to make these decisions by comparing its own path to the "expert" path (the real data). It uses a mathematical tool (called an f-divergence) to measure how different its path is from the expert's path and adjusts its strategy to minimize that difference.

What They Found
The researchers tested this on famous image datasets (like CIFAR-10, FFHQ, and ImageNet). Here are the key takeaways:

It's Smarter than Manual Tuning: The AI learned to change its instructions based on the specific moment in the process. For example, it learned to add "chaos" only when the image was blurry and to be very precise when the image was almost done. This is much better than using a fixed rule for the whole process.
It Saves Massive Energy: On the ImageNet dataset, finding the best manual settings required testing thousands of combinations, costing a huge amount of computer power. Their method found a better solution with a single training run, saving up to 9 times the computational cost.
It's a Small Price to Pay: Once the AI learns these strategies, using them to generate images only adds about 16% extra work to the computer. This is a tiny price to pay to avoid the massive cost of manual testing.
Control over Quality vs. Variety: The method allows users to choose a "flavor" of learning. They can tune it to be very precise (making images look exactly like the training data) or very diverse (making images that are more varied but maybe slightly less perfect).

In Summary
This paper introduces a "smart coach" for AI image generators. Instead of manually tweaking knobs and dials to find the perfect settings, the AI learns the perfect settings by watching the target data and adjusting its own behavior in real-time. It's faster, cheaper, and produces better results than the old way of guessing and checking.

Technical Summary: Learning To Sample From Diffusion Models Via Inverse Reinforcement Learning

Problem Statement

Diffusion models have become a standard for high-dimensional data generation, yet the sampling procedure remains a critical bottleneck. While the denoising neural network is typically fixed after training, the sampling algorithm itself requires careful tuning of hyperparameters (e.g., noise schedules, guidance scales, stochasticity profiles). Currently, this tuning is performed via costly empirical grid searches. For complex models like those trained on ImageNet-64, the computational cost of exhaustive hyperparameter search can approach or exceed the cost of training the denoiser itself. Furthermore, existing reinforcement learning (RL) or control approaches for diffusion models often require explicit reward functions or fine-tuning of the denoiser, which contradicts the goal of optimizing the sampler without modifying the pretrained generator.

Methodology

The authors propose framing diffusion sampling as a sequential decision-making problem solvable via Inverse Reinforcement Learning (IRL). The core innovation is learning a sampling policy that matches target state distributions without access to expert trajectories, explicit rewards, or denoiser retraining.

1. MDP Formulation

The diffusion sampling process is formalized as a finite-horizon Markov Decision Process (MDP):

State ( $s_t$ ): Defined as $(x_t, \sigma_t)$ , where $x_t$ is the current sample and $\sigma_t$ is the noise level.
Action ( $a_t$ ): Discrete control decisions applied to the sampling dynamics, such as selecting guidance scales, noise injection levels, or deciding whether to renoise (restart) to a previous noise level.
Dynamics: The transition kernel is induced by the discretized diffusion dynamics and the selected control action. The process terminates when the noise level reaches $\Sigma_0$ (an absorbing state).

2. State-Only Inverse Reinforcement Learning

Unlike traditional IRL which often assumes expert actions or trajectories, this method operates under a state-only regime. The "expert" is defined solely by the marginal distributions of data samples at selected noise levels.

Objective: The goal is to find a policy $\pi_\theta$ whose state occupancy measure $\mu_\theta$ (the marginal distribution of states visited) matches the expert occupancy measure $\mu_E$ .
Optimization: Instead of learning a reward function, the authors directly minimize the $f$ -divergence between the expert and policy occupancy measures:
$\min_\theta D_f(\mu_E \parallel \mu_\theta)$
The paper specifically explores Kullback-Leibler (KL) and Reverse-KL (rKL) divergences, noting that the choice of $f$ -divergence dictates the trade-off between mode-covering (diversity) and mode-seeking (fidelity).

3. Policy Gradient and Optimization

The authors derive a policy gradient theorem for minimizing $f$ -divergences directly. The gradient is expressed as:
$\nabla_\theta L_f(\theta) = \mathbb{E}_{s_{1:T} \sim p_\theta} \left[ \frac{1}{T} \sum_{t=1}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \sum_{t' \ge t} h_f\left(\frac{\mu_E(s_{t'})}{\mu_\theta(s_{t'})}\right) \right]$
where $h_f$ is a function derived from the generator of the $f$ -divergence.

Density Ratio Estimation: Since $\mu_E$ and $\mu_\theta$ are unknown, the density ratio is estimated using a discriminator trained to distinguish between expert states (noised data) and policy-generated states.
Stabilization: To handle high variance in gradient estimates, the method employs a PPO-style clipped objective, importance sampling corrections, and a global batch mean baseline.

4. Sampling Strategies Evaluated

The framework is applied to three specific inference-time interventions:

Adaptive Classifier-Free Guidance: Learning a state-dependent guidance scale $\omega(x, \sigma)$ .
Stochasticity Injection: Learning a state-dependent noise injection level $\gamma_{EDM}(x, \sigma)$ .
Renoising (Restart Sampling): Learning when to revert the process to a previous noise level to escape poor local generations.

Key Contributions

MDP Formulation for Sampling: The paper casts diffusion sampling as an MDP to adapt samplers at inference time without retraining the denoiser.
Direct State-Marginal Matching: It derives a policy gradient objective that minimizes $f$ -divergences between state occupancy measures, bypassing the need for explicit reward modeling or expert action data.
Interpretable Trade-offs: The framework exposes a controllable trade-off between sample diversity (Recall) and fidelity (Precision) through the choice of $f$ -divergence (KL vs. rKL).
State-Dependent Control: The learned policies act as inference-time controls that selectively intervene (e.g., applying high guidance only at specific noise levels or for specific sample types) rather than applying fixed global schedules.

Experimental Results

The method was evaluated on CIFAR-10, FFHQ, and ImageNet-64 using pretrained EDM models.

Performance: The learned policies match or exceed the performance of the best hand-tuned baselines (found via grid search) across FID, Precision, and Recall metrics.
- Adaptive Guidance: Improved the quality-diversity trade-off compared to fixed guidance. KL divergence improved Recall (diversity), while rKL improved Precision (fidelity).
- Stochasticity Injection: The learned state-dependent profiles recovered the qualitative patterns of optimal EDM schedules (low stochasticity at the end of the trajectory) and outperformed fixed baselines on complex datasets like ImageNet.
- Renoising: Reduced the Number of Function Evaluations (NFE) significantly on ImageNet while maintaining or improving FID compared to standard restart strategies.
Computational Efficiency:
- Training: A single training run replaces exhaustive grid search. On ImageNet-64, this reduces the optimization budget by up to 9 $\times$ (from ~1437 EFLOPs for a 3-hyperparameter grid search to ~161 EFLOPs for the IRL training).
- Inference: The policy adds a modest overhead of 16% to 40% in FLOPs compared to the denoiser alone, which is considered negligible given the elimination of manual tuning.
Ablation Findings:
- Conditioning: Policies conditioned on the full state $(x, \sigma)$ significantly outperform those conditioned only on noise level $\sigma$ , and perform nearly as well as those with explicit class labels $(x, \sigma, c)$ , suggesting the policy implicitly learns class-discriminative features from the image content.
- Divergence Choice: KL and rKL are shown to be the most effective for controlling the precision-recall trade-off, while symmetric divergences (TV, Hellinger) produced less interpretable results.

Significance and Claims

The paper claims that this framework provides a systematic, automated alternative to manual hyperparameter tuning for diffusion sampling. By treating sampling as an IRL problem, it allows for the optimization of complex, state-dependent sampling strategies that are difficult to design manually.

The authors emphasize that the method is particularly valuable for large-scale models where the cost of per-trial evaluation in a grid search is prohibitive. The approach does not modify the generative model itself but rather learns an optimal "control policy" for the inference trajectory, offering a practical path to improving sample quality and diversity without the computational burden of retraining or exhaustive search. The paper remains modest, noting that the benefits are dataset-dependent and that the current framework relies on discrete action spaces, suggesting future work in continuous action spaces and more robust initialization.