<figure><img src="image_1.png" /></figure>

HAL
open science

# Learning To Sample From Diffusion Models Via Inverse Reinforcement Learning

Constant Bourdrez, Alexandre Verine, Olivier Cappé

► To cite this version:

Constant Bourdrez, Alexandre Verine, Olivier Cappé. Learning To Sample From Diffusion Models Via Inverse Reinforcement Learning. 2026. hal-05523140

HAL Id: hal-05523140

https://hal.science/hal-05523140v1

Preprint submitted on 23 Feb 2026

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_2.png" /></figure>

Distributed under a Creative Commons CC BY 4.0 - Attribution - International License
Learning To Sample From Diffusion Models Via Inverse Reinforcement
Learning

Constant Bourdrez

Alexandre Vérine

Olivier Cappé

DI ENS, Ecole normale supérieure, Université PSL, CNRS, 75005 Paris, France

constant.bourdrez@ens.psl.eu

alexandre.verine@ens.psl.eu

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.

# 1 Introduction

Diffusion models and related generative processes based on differential equations have become a standard tool for modeling high-dimensional data, with applications ranging from image synthesis to scientific and biological modeling. Sample generation proceeds by simulating a reverse-time stochastic or deterministic process that gradually transforms a simple Gaussian distribution into a complex data distribution. This evolution is governed by a vector field implemented by a large neural network, commonly referred to as the denoiser.

While the training of diffusion models has been extensively studied, the sampling procedure itself remains a critical bottleneck. A growing body of work proposes heuristic modifications of the sampling dynamics, including restart sampling [45], guidance mechanisms [19, 11], noise schedule tuning [27], and stochasticity injection [23]. Although often effective, these approaches are typically designed in an ad

hoc manner and rely on manually tuned hyperparameters, which makes them difficult to adapt systematically to new objectives, datasets, or computational constraints. In this work, we propose to learn such parameters automatically by framing diffusion sampling as a sequential decision-making problem.

Reinforcement learning has recently been applied to generative modeling, most prominently in the context of large language models through reinforcement learning from human feedback [44, 31, 4, 40]. Several works extend this paradigm to diffusion models by defining explicit reward functions and fine-tuning either the denoiser or the sampling process [7, 13, 16, 41, 5]. In parallel, diffusion models have also been studied through the lens of stochastic control. Recent work on adjoint matching and related control-based formulations proposes to fine-tune diffusion or flow models by matching backward dynamics or solving variational control problems in continuous time [3, 12, 34]. These approaches are naturally formulated in continuous time and rely on differentiating through the generative dynamics. A related inverse reinforcement learning approach for diffusion models is proposed by Yoon et al. [46], who learn an energy-based reward to guide sampling.

In this work, we take a different perspective and cast diffusion sampling as an imitation learning problem. Given access to an expert diffusion sampler e.g., samples from the forward diffusion process at different noise levels, we learn a policy that adapts the sampling process of a pre-trained diffusion model to better match expert behavior. While our approach builds on state-marginal inverse reinforcement learning, adapting these ideas to diffusion sampling requires addressing structural constraints that are absent in generic MDPs, including absorbing terminal dynamics, and expert supervision limited to marginal distributions at selected noise levels. In contrast to prior IRL and control-based approaches for diffusion models, our method (i) operates purely at sampling time, (ii) does not learn rewards or energies, and (iii) exploits the structure of noise levels to define expert occupancy without requiring expert trajectories.
## 2 Sample Generation with Diffusion Models

Diffusion models generate data in $\mathbb{R}^d$ by transporting samples from a Gaussian distribution to a target data distribution $P_0$, via a reverse-time stochastic or deterministic process [36, 37]. The reverse dynamics rely on a neural network $F$ (the denoiser) trained across noise levels to estimate the score of the noised data distribution. In practice, $F$ is trained either to predict the added noise [20] or the denoised sample [11], and may be conditioned on auxiliary information $c$ such as class labels or text prompts.

While diffusion models are commonly described in continuous time, practical implementations rely on a finite discretization. We can see the generation process as a sequence of states $s_t = (x_t, \sigma_t)$, where $x_t$ is the generated sample at time $t$ and $\sigma_t$ is the corresponding noise level defined by the noise schedule $\{\Sigma_N, \dots, \Sigma_0\}$. Typically, $\Sigma_0 = 0 < \Sigma_1 < \dots < \Sigma_N$ and $\Sigma_N$ is large enough for $P_N$ to approximate pure Gaussian noise. The sequence starts at $t=0$ from $s_0 = (x_0, \Sigma_N)$ with $x_0 \sim \mathcal{N}(0, \Sigma_N I_d)$ and ends at $t=T$ with $s_T = (x_T, \Sigma_0)$, where $x_T$ is expected to follow the target distribution $P_0$. Each step of the reverse process maps $(x_t, \sigma_t)$ to $(x_{t+1}, \sigma_{t+1})$ according to

$$ (x_{t+1}, \sigma_{t+1}) = H(x_t, \sigma_t, \dots), $$

where $H$ denotes a generic update operator that uses the denoiser $F$, evaluated at the current state $(x_t, \sigma_t)$ and possibly additional inputs such as conditioning information $c$. We treat $\sigma_t = \Sigma_0$ as an absorbing condition, i.e., once reached, the process remains at this noise level for all subsequent steps. The operator $H$ can represent a variety of samplers: it is stochastic for DDPM [20] and Score-SDE [38], and deterministic for probability-flow ODE (PF-ODE) based samplers, e.g., first-order Euler's method [37], second-order Heun's method [21, 23], or higher-order solvers such as DPM-Solver [28]. Notably, $H$ depends on hyperparameters that influence the sampling dynamics. In this work, we focus on three popular techniques to modify the sampling process: stochasticity injection, classifier-free guidance, and restart sampling.

**Stochasticity Injection:** The EDM approach of Karras et al. [23], which serves as a strong baseline for image generation, employs a hyperparameter $\gamma_{EDM}$ to introduce controlled stochasticity. Using $H_{Heun}$ to denote the deterministic Heun update operator, the modified update rule relies on an amplified noise level $\hat{\sigma}_t := \sigma_t (1 + \gamma_{EDM})$ and a noise perturbation $z \sim \mathcal{N}(0, I)$:

$$ H_{EDM}(x_t, \sigma_t) = H_{Heun}(x_t + \sqrt{\hat{\sigma}_t^2 - \sigma_t^2} z, \hat{\sigma}_t). $$

For some datasets such as CIFAR-10 [1] or FFHQ [22], $\gamma_{EDM} = 0$ yields optimal results, while for more complex datasets such as ImageNet [10], non-zero values (applied at intermediate noise levels) can improve sample quality.

**Classifier-Free Guidance:** Classifier-free guidance [19] is one of the most popular techniques to enhance sample quality in conditional diffusion models. For any $H$, where the model $F$ has been trained to be conditioned on $c$ or unconditioned ($c = \emptyset$), the update rule is kept identical to $H$ but the denoiser is modified as

$$ F_{CFG}(x_t, \sigma_t, c) = (1 + \omega)F(x_t, \sigma_t, c) - \omega F(x_t, \sigma_t, \emptyset). $$

Depending on the guidance scale $\omega$, this technique can significantly improve sample fidelity, at the cost of diversity [19, 33].

**Restart Sampling:** Another approach consists in restarting the sampling process to escape poor local generations [45]. Let $i < j$ be two integers such that $\Sigma_j > \Sigma_i$. When the sample reaches $\Sigma_i$, it is sent back $K$ times to $\Sigma_j$ by adding a random Gaussian perturbation $z \sim \mathcal{N}(0, I)$. Let $k_t$ denote the number of restarts performed up to time $t$; the update rule usually relies on any PF-ODE-based operator $H_{ODE}$ and becomes

$$ H(x_t, \sigma_t) = \begin{cases} H_{ODE}(x_t + \sqrt{\Sigma_j^2 - \Sigma_i^2} z, \Sigma_j), & \text{if } k_t < K \text{ and } \sigma_t = \Sigma_i, \\ H_{ODE}(x_t, \sigma_t) & \text{otherwise.} \end{cases} $$

This technique can be generalized in several ways. The variant we consider in this paper allows renoising from any $\Sigma_i$ to one of the $M$ previous higher noise levels $\Sigma_{i+1}, \dots, \Sigma_{i+M}$ upon reaching a restart level $\Sigma_i$. In the experiments, we will fix $M=4$ and refer to this variant as *Renoise*. Additionally, a fixed budget of function evaluations is used instead of a fixed number of restarts.

## 3 State Distribution Matching with IRL

In this section, we review reinforcement learning (RL) and inverse reinforcement learning (IRL) tools that motivate our formulation in Section 4. Our goal is to cast the diffusion sampling algorithm from Section 2 as a decision-making problem, in order to select sampler hyperparameters (e.g., guidance scales, injected noise, or renoise rules). Crucially, we do *not* assume access to an expert policy or expert actions. Instead, we only observe samples from the data distribution and their noisy counterparts across noise levels, which naturally aligns with IRL settings based on matching state distributions.

**MDP formulation** We formalize diffusion sampling as a finite-horizon Markov Decision Process (MDP) $(S, A, P, \rho, T)$. We consider discrete-time policies with a fixed horizon $T$, initial states sampled from distribution $\rho$, and transition dynamics defined by the kernel $P$. The state space $S$ is continuous, and the action space $A$ is discrete. Let $\pi_\theta(a_t | s_t)$
be a parametrized policy over actions conditioned on states, and let $s_{1:T} = (s_1, \dots, s_T)$ denote a trajectory generated by this policy. In Section 4, we will specify how each component of the MDP is instantiated for diffusion samplers.

**State occupancy measures** In our setting, the expert policy is unknown and we do not observe expert actions or full trajectories; the only supervision comes from expert states (data samples and their noised versions). This state-only regime motivates comparing policies through their state occupancy measures, i.e., the marginal distribution over states visited along trajectories.

**Definition 3.1 (Occupancy Measure).** The policy $\pi_\theta$ induces a state occupancy measure $\mu_\theta$ defined as

$$ \forall \Phi \in \mathcal{B}(S), \quad \int \Phi(s) \mu_\theta(s) = \mathbb{E}_{s_{1:T} \sim p_\theta} \left[ \frac{1}{T} \sum_{t=1}^{T} \Phi(s_t) \right], $$

where $\mathcal{B}(S)$ is the set of bounded measurable functions on $S$, and $p_\theta$ is the trajectory distribution induced by $\pi_\theta$.

When $S$ is discrete, $\mu_\theta$ admits the classical interpretation of an expected visitation frequency:

$$ \mu_\theta(s) = \mathbb{E}_{s_{1:T} \sim p_\theta} \left[ \frac{1}{T} \sum_{t=1}^{T} \mathbb{1}_{\{s_t=s\}} \right]. $$

We will not assume $S$ is discrete in what follows; the discrete case is only meant as intuition. We denote by $\mu_E$ the expert occupancy measure, induced by the (unknown) expert policy and the environment dynamics.

**State-marginal matching via f-divergences** State-marginal matching IRL seeks a policy whose state occupancy measure $\mu_\theta$ matches the expert occupancy measure $\mu_E$. Concretely, one minimizes a discrepancy between these probability measures. Several approaches derive objectives specialized to the Kullback-Leibler (KL) divergence, typically via maximum-entropy RL and reward learning [47, 29]. More generally, one can use the broader class of f-divergences [9], which provides a unified way to quantify distribution mismatch.

**Definition 3.2 (f-divergence).** The f-divergence between two probability measures $Q$ and $P$, induced by a convex function $f$ satisfying $f(1) = 0$, is

$$ D_f(Q \| P) := \mathbb{E}_P \left[ f \left( \frac{dQ}{dP} \right) \right]. \quad (1) $$

This family includes KL, Reverse-KL (rKL), and total variation distance. The choice of $f$ matters: different divergences induce different trade-offs (e.g., mode-covering (KL) vs. mode-seeking behavior (rKL)), which can change the qualitative properties of the learned policy [30, 42].

In this work, we consider objectives of the form

$$ \min_{\pi_\theta} D_f(\mu_E \| \mu_\theta), \quad (2) $$

subject to the MDP dynamics induced by the sampler.

**Prior work and our approach** State-marginal matching IRL has often been approached through reward learning under maximum-entropy RL [47, 29]. In these methods, one first estimates a reward (or energy) from expert states and then relies on MaxEnt RL to induce a policy whose state distribution matches the expert. This indirect route can be undesirable in our setting, as it ties performance to the quality of the learned reward and to the ability of the MaxEnt policy to explore diverse states. By contrast, Ghasemipour et al. [15] show how to optimize policies directly by minimizing f-divergences via convex duality, but their formulation targets state-action occupancy measures. To the best of our knowledge, an analogous reward-free, direct minimization of $D_f(\mu_E \| \mu_\theta)$ for state marginals has not been developed in the context we consider.

Our approach addresses two gaps: (i) we avoid explicit reward learning by directly optimizing a sampling policy to match expert state distributions, and (ii) we instantiate this objective for diffusion samplers, enabling principled optimization of sampler hyperparameters. In Section 4, we make the diffusion-to-MDP mapping explicit and derive the resulting optimization procedure.

# 4 Improving Diffusion with IRL

In this section, we present our theoretical contributions and introduce a general framework for controlling generative sequential processes. We first formalize diffusion sampling as a finite-horizon MDP in Section 4.1. We then describe how to learn sampling policies by minimizing f-divergences between state occupancy measures in Section 4.2.

## 4.1 The MDP Model for Diffusion Sampling

To apply our inverse reinforcement learning algorithm, we frame each previously defined component as an element of diffusion sampling. The state space $S$ is continuous, and each state is defined as

$$ s_t = (x_t, \sigma_t), $$

where $\sigma_t$ denotes the noise level at time step $t$ and $x_t \in \mathbb{R}^d$ is the current sample. The action space $A$ is discrete and encodes control decisions applied to the sampling dynamics. For instance, actions can correspond to selecting different discretized guidance scales, noise injection levels, or whether to renoise and how much at each timestep. The initial distribution $\rho$ corresponds to the marginal distribution at the highest noise level, the Gaussian distribution $N(0, \Sigma_N^2 I_d)$.
Given a state-action pair $(s_t, a_t)$, the transition kernel $\mathcal{P}$ is induced by the discretized diffusion dynamics together with the selected control action. The policy does not explicitly depend on time but the available actions may depend on the current timestep $t$. We will treat $\sigma_t = \Sigma_0$ as an absorbing condition, so that once the process reaches this noise level, it remains there for all subsequent timesteps. In practice, this can be implemented by defining $H$ such that $s_{t+1} = s_t$ whenever $\sigma_t = \Sigma_0$. In this setup, the average time required to reach the terminal noise level may be incorporated in the performance metric of a given policy.

As the state space is continuous, we rely on the definition 3.1. Exploiting the structure of the state space, we further decompose the occupancy measure as

$$ \mu_{\theta}(s) = w_{\theta}(\sigma) \times p_{\theta}(x|\sigma), $$

where $w_{\theta}(\sigma)$ denotes the weight assigned by the policy $\pi_{\theta}$ to noise level $\Sigma$, and $p_{\theta}(x|\sigma)$ is the marginal distribution of $x$ at noise level $\Sigma$ under the policy. For the samplers such as noise injection or guidance, $w_{\theta}$ is uniform over noise levels, since the process always progresses from $\Sigma_N$ to $\Sigma_0$ in $T$ steps. For samplers with renoising, $w_{\theta}$ reflects the distribution over noise levels induced by the policy, which may spend more time at certain noise levels depending on the restart strategy.

**Expert occupancy measure** In our framework, the expert is specified through samples from the target distribution at selected noise levels. We do not assume access to expert trajectories, actions, or transition dynamics. The expert occupancy measure $\mu_E$ is

$$ \mu_E(s) = w_E(\sigma) \times p_E(x|\sigma). $$

As the expert does not provide temporal information, the weights $w_E(\sigma)$ are a design choice that favors trajectories reaching quickly states $s_t = (x_t, \sigma_t)$ such that $\sigma_t = \Sigma_0$. There are various ways to define $w_E$ to reflect this preference. A simple choice is to set it uniformly over $\{\Sigma_N, \dots, \Sigma_1\}$ and assign all remaining weight to $\Sigma_0$. Our primary objective is to match the expert distribution at this terminal noise level. The absorbing condition $\{\sigma_t = \Sigma_0\}$ allows all states at noise level $\Sigma_0$ to contribute to the occupancy measure, yielding a denser learning signal that reflects how closely the generated samples match the expert distribution at the lowest noise level. $w_E$ can also be interpreted as a proxy for the time spent at each noise level. This motivates us to interpret $w_E(\Sigma_0)$ as a measure of the mean time spent at the terminal noise level, which we want the policy $w_{\theta}(\Sigma_0)$ to match.

**On the leverage of $w_E$** Given the definitions above, minimizing an $f$-divergence between the expert and policy-induced occupancy measures implies a corresponding lower bound on the divergence between the marginal noise-level weights:

$$ \mathcal{D}_f(\mu_E \| \mu_\theta) \geq \mathcal{D}_f(w_E \| w_\theta). \quad (3) $$

This follows from the data processing inequality for f-divergences with convex generators. Consequently, the choice of $w_E$ directly influences the optimization problem. Moreover, for the Kullback-Leibler and reverse Kullback-Leibler divergences, the objective admits an explicit decomposition into a term depending only on $w_E$ and $w_\theta$, and a term involving the conditional marginals at each noise level:

$$ D_{\text{KL}}(\mu_E \| \mu_\theta) = D_{\text{KL}}(w_E \| w_\theta) + \sum_{\sigma} w_E(\sigma) D_{\text{KL}}(p_E(\cdot | \sigma) \| p_\theta(\cdot | \sigma)), \quad (4) $$

$$ \mathcal{D}_{\text{rKL}}(\mu_E \| \mu_\theta) = \mathcal{D}_{\text{rKL}}(w_E \| w_\theta) + \sum_{\sigma} w_\theta(\sigma) \mathcal{D}_{\text{rKL}}(p_E(\cdot | \sigma) \| p_\theta(\cdot | \sigma)). \quad (5) $$

This decomposition highlights how the optimization jointly controls the distribution over noise levels and the quality of samples generated at each noise level.

## 4.2 Learning Policies by Minimizing f-Divergences

We now describe how to optimize sampling policies within the MDP framework introduced above. We directly optimize the policy to match the expert occupancy measure by optimizing the following objective:

$$ \min_{\theta} \mathcal{L}_f(\theta) = \min_{\theta} \mathcal{D}_f(\mu_E \| \mu_\theta). \quad (6) $$

**Policy gradient optimization** We show in Appendix A.1 that the gradient of the $f$-divergence objective with respect to the policy parameters, denoted $\nabla_{\theta}\mathcal{L}_f(\theta)$, can be expressed in the form of a policy gradient.

**Theorem 4.1.** Let $\pi_{\theta}$ be a policy parameterized by $\theta$ and let $\mathcal{L}_f(\theta) = \mathcal{D}_f(\mu_E \| \mu_\theta)$ be the $f$-divergence between the expert occupancy measure $\mu_E$ and the occupancy measure $\mu_\theta$ induced by the policy $\pi_{\theta}$. Then, the gradient of $\mathcal{L}_f(\theta)$ with respect to $\theta$ is given by

$$ \nabla_{\theta} \mathcal{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) A_t \right], \quad (7) $$

where $A_t = \sum_{t' \ge t} h_f(\frac{\mu_E(s_{t'})}{\mu_\theta(s_{t'})})$ is the learning signal, $h_f(x) = f(x) - x f'(x)$ for all $x$ and $f$ is the generator function of the $f$-divergence (see Appendix A for details).

This result enables the use of standard policy gradient methods to optimize the sampling policy using Monte Carlo estimates obtained from rollouts of the current policy. While this result seems simple and similar to other policy gradient theorems, the main novelty lies in the direct optimization of the policy instead of a reward function.
# 5 Practical Optimization

In this section, we describe how to estimate the density ratios $\frac{\mu_E}{\mu_\theta}$ appearing in the policy gradient expression in equation 4.1, and how to optimize the policy in practice.

**Discriminator-based density ratio estimation** Estimating density ratios in continuous spaces is a classical problem, with approaches ranging from kernel density estimation to direct ratio fitting and probabilistic modeling [39]. In this work, we adopt a classifier-based approach, which has been shown to scale effectively to high-dimensional settings [18, 14, 36, 32, 2] and has successfully been applied to diffusion models [24, 43].

We train a binary classifier $D_\varphi(x, \sigma)$ to discriminate between states sampled from the expert and from the policy, conditioning on the noise level $\sigma$. Since the expert provides only marginal distributions at each noise level, we sample $\sigma$ uniformly and draw states independently from $p_E(\cdot | \sigma)$ and $p_\theta(\cdot | \sigma)$. The classifier is trained by minimizing the binary cross-entropy loss

$$ \mathcal{L}_D(\varphi) = \mathbb{E}_{\sigma \sim \mathcal{U}[\Sigma_0, \dots, \Sigma_N]} \left[ -\mathbb{E}_{x \sim p_E(\cdot|\sigma)}[\log D_\varphi(x, \sigma)] - \mathbb{E}_{x \sim p_\theta(\cdot|\sigma)}[\log(1 - D_\varphi(x, \sigma))] \right]. $$

Assuming an optimal discriminator, the density ratio is

$$ \frac{p_E(x | \sigma)}{p_\theta(x | \sigma)} = \frac{D_\varphi(x, \sigma)}{1 - D_\varphi(x, \sigma)} = \exp(\tilde{D}_\varphi(x, \sigma)), \quad (8) $$

where $\tilde{D}_\varphi(x, \sigma)$ denotes the logit output of the classifier.

**Computing the occupancy measure ratio** Because expert trajectories are unavailable, the expert is specified only through state marginals at each noise level. As a result, the discriminator estimates a density ratio under a uniform distribution over noise levels rather than under the policy-induced state occupancies. In addition, while the policy accumulates occupancy mass at the terminal noise level, the expert does not provide temporal occupancy information.

We compute the full occupancy measure ratio by using importance sampling corrections to account for the mismatch between sampling and target proportions.

$$ \frac{\mu_E(x, \sigma)}{\mu_\theta(x, \sigma)} = \frac{w_E(\sigma)}{w_\theta(\sigma)} \times \frac{p_E(x|\sigma)}{p_\theta(x|\sigma)}. \quad (9) $$

As previously discussed, the term $w_E$ is a design choice while $w_\theta$ is induced by the policy and is estimated empirically from the simulated trajectories.

**Sample efficiency and practical considerations** Generating trajectories is computationally expensive, as each rollout requires simulating a full diffusion sampling trajectory. To improve sample efficiency, we reuse trajectories over multiple policy updates and add importance sampling corrections to the policy gradient. While this modification is theoretically sound, we found it to be unstable since it

leads to high-variance gradient estimates. We attempted to mitigate this issue by subtracting a baseline, a standard technique for reducing variance in policy gradient methods. However, in our setting, there is no clear choice of baseline. We therefore use the empirical average of the learning signals

$$ \bar{A}_t = \frac{1}{n_{\text{traj}}} \sum_{i=1}^{n_{\text{traj}}} A_t^{(i)} \text{ as a baseline.} $$

**Algorithm 1 Policy optimization algorithm**

**Require:** Update Sampler Operator $H$, pretrained denoiser $F$, expert state marginals $\{\mu_E(\cdot | \sigma)\}_\sigma$, divergence $f$

1: Initialize sampling policy $\pi_\theta$

2: Initialize state density-ratio estimator $D_\varphi$

3: **repeat**

4: Generate trajectories using $\pi_\theta$ choosing parameters of $H$ and frozen denoiser $F$

5: Sample expert and policy states uniformly across noise levels

6: Train $D_\varphi$ to discriminate expert from policy states

7: Compute policy learning signals $A_t$ using $D_\varphi$ and occupancy weights

8: Update $\pi_\theta$ using a $\nabla_\theta \mathcal{L}_f^{\text{PPO}}(\theta)$

9: **until** convergence

**Ensure:** Optimized sampling policy $\pi_\theta$

To further reduce variance, we adopt a Proximal Policy Optimization (PPO)-style clipped objective [35], which constrains the policy updates at each iteration as well as reward normalization. Under the same assumptions as Theorem 4.1, the PPO-style clipped gradient for minimizing an $f$-divergence is given by

$$ \nabla_{\theta} \mathcal{L}_{f}^{\text{PPO}}(\theta) = \mathbb{E}_{p_{\theta_{0}}} \left[ \sum_{t=1}^{T} \nabla_{\theta} \max(r_{t}(\theta) \hat{A}_{t}, \tilde{r}_{t}(\theta) \hat{A}_{t}) \right] $$

where $r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}$ is the probability ratio between the current and sampling policies, $\tilde{r}_t(\theta) = \text{clip}(r_t(\theta), 1 - \epsilon, 1 + \epsilon)$ is the clipped probability ratio, and $\hat{A}_t = \frac{1}{T}(A_t - \bar{A}_t)$ denotes the recentered learning signal at time $t$.

Finally, following common practice, we approximate $h_f(\mu_E(s_t)/\mu_\theta(s_t))$ by $h_f(\mu_E(s_t)/\mu_{\theta_0}(s_t))$ to avoid recomputing the density ratio at every policy update. Algorithm 1 summarizes the high-level structure of the optimization procedure. We provide a detailed Algorithm 2 in Appendix B.1.

# 6 Experiments

We evaluate the proposed framework across several diffusion sampling strategies to address four questions: (i) how broadly the method applies to different modifications of the sampling dynamics, (ii) whether conditioning the policy on the current sample $x_t$ is beneficial, (iii) how sensitive the approach is to initialization, and (iv) how the choice of divergence controls the precision-recall trade-off.
Experiments are conducted on CIFAR-10, FFHQ, and ImageNet using pretrained EDM models [23]. The expert conditional distribution $p_E(x|\sigma)$ is defined as the empirical dataset distribution at noise level $\Sigma_0 = 0$ with additional Gaussian noise at level $\sigma$. All policy generations are evaluated using FID [17] with 50k samples, Precision and Recall [26] using the improved topological method of Kim et al. [25], and the average Number of Function Evaluations (NFE) required to reach the terminal noise level (when relevant).

**What strategies can be applied?** A key strength of the framework is its flexibility: it can be applied to a wide range of sampling-time interventions without modifying the denoiser. We consider adaptive classifier-free guidance, Renoise strategy that revisits higher-noise states and learning the stochasticity injection profile $\gamma_{EDM}$. Despite their differences, all strategies are learned using the same state-occupancy matching objective.

**Computational overhead** Our method introduces negligible overhead at sampling time. The learned policies are lightweight neural networks with small action spaces, and their evaluation cost is insignificant compared to denoiser evaluations. The main computational cost arises during training, where multiple sampling trajectories are required to estimate gradients and train the discriminator. This cost remains far below that of retraining the denoiser itself.

## 6.1 Learning Adaptive Classifier-Free Guidance

Classifier-free guidance trades diversity for fidelity via a guidance scale $\omega$, but fixed guidance schedules have been shown to collapse to the mean of the conditional distribution [8]. In this experiment, we explore learning an adaptive, state-dependent guidance policy within our framework. Experiments are conducted on CIFAR-10 32 × 32 and FFHQ 64 × 64 using pretrained EDM models. We compare the learned adaptive guidance against standard fixed guidance scales. We also use this experiment as a way to compare policies dependent on the image state $\pi_\theta(x, \sigma)$ against simpler policies dependent only on the noise level $\pi_\theta(\sigma)$ while using the same initialization for both policies. Quantitative results are reported in Table 1. The details about x-dependent policies are provided in Appendix B.2. Although KL and rKL achieve comparable FID scores, they do so through qualitatively different guidance profiles. KL places more emphasis on covering the support of the data distribution, improving recall and encouraging sample diversity. In contrast, rKL concentrates probability mass on high-density regions, favoring precision and higher-fidelity samples. Figure 2b illustrates how these two different options impact the learned guidance strength profiles. The behavior exhibited by the learned policies consists in using high guidance for a small proportion of the images at each timestep, while leaving the

<table><caption>Table 1: Performance comparison for learning the guidance force $\omega(x, \sigma)$. Adaptive guidance improves FID as well as Precision and Recall compared to fixed guidance scales.</caption><thead><tr><th>Method</th><th>Policy</th><th>FID ↓</th><th>Precision ↑</th><th>Recall ↑</th></tr></thead><tbody><tr><td colspan="5" style="text-align: center; font-weight: bold;">CIFAR-10 (32 × 32)</td></tr><tr><td rowspan="2">KL</td><td>π(x, σ)</td><td><b>2.05</b></td><td>80.7</td><td><b>71.4</b></td></tr><tr><td>π(σ)</td><td>2.26</td><td>80.2</td><td>70.9</td></tr><tr><td rowspan="2">rKL</td><td>π(x, σ)</td><td>2.53</td><td><b>82.4</b></td><td>69.8</td></tr><tr><td>π(σ)</td><td>2.23</td><td>80.2</td><td>70.9</td></tr><tr><td>EDM + CFG<br>(ω = 0.2)</td><td><b>X</b></td><td>2.07</td><td>80.6</td><td>70.9</td></tr><tr><td colspan="5" style="text-align: center; font-weight: bold;">FFHQ (64 × 64)</td></tr><tr><td>KL</td><td>π(x, σ)</td><td>3.05</td><td>90.4</td><td><b>88.6</b></td></tr><tr><td>rKL</td><td>π(x, σ)</td><td><b>3.04</b></td><td><b>90.7</b></td><td>88.3</td></tr><tr><td>EDM + CFG<br>(ω = 0.1)</td><td><b>X</b></td><td>3.11</td><td><b>90.7</b></td><td>87.8</td></tr></tbody></table>

majority of images unguided. The policy increases the proportion of highly guided samples in the medium-to-low noise levels, corresponding to the speciation phase of sampling proposed by Biroli et al. [6]. Around medium noise levels, the guidance strength exhibits higher mean and variance, indicating a greater need for guidance during this phase, which coincides with the emergence of topological structures in the images. Moreover, the dependence in $x$ leads to policies with improved performance. This experiment highlights two important points. First, conditioning the policy on the current state $(x, \sigma)$ is crucial, as a single scalar guidance value cannot capture the varying trade-offs required at different noise levels. Second, the choice of divergence provides a direct and interpretable mechanism for controlling the precision-recall

<figure><img src="image_3.png" /><figcaption>Figure 1: FFHQ samples generated with fixed guidance $\omega = 0.1$ (left) and with adaptive guidance optimized with rKL (right).</figcaption></figure>
<figure><img src="image_4.png" /></figure>

<figure><img src="image_5.png" /><figcaption>Figure 2: Learned classifier-free guidance profiles $\omega(x, \sigma)$ for different $f$-divergences. Higher guidance improves precision but reduces diversity. We report the mean as well as the 0.25 and 0.75 quantiles across samples.</figcaption></figure>

<table><caption>Table 2: Performance comparison for learning renoise strategies. Best results are highlighted in bold.</caption><thead><tr><th>Method</th><th>Policy</th><th>FID</th><th>Precision</th><th>Recall</th><th>NFE</th></tr></thead><tbody><tr><td colspan="6"><strong>CIFAR-10 (32 x 32)</strong></td></tr><tr><td rowspan="2">KL</td><td>$\pi(x, \sigma)$</td><td><strong>3.18</strong></td><td>77.3</td><td><strong>72.5</strong></td><td>17.7</td></tr><tr><td>$\pi(\sigma)$</td><td>3.45</td><td>77.4</td><td>71.7</td><td>17.0</td></tr><tr><td rowspan="2">rKL</td><td>$\pi(x, \sigma)$</td><td>3.21</td><td><strong>78.1</strong></td><td>71.6</td><td>30.5</td></tr><tr><td>$\pi(\sigma)$</td><td>3.46</td><td>76.5</td><td>72.2</td><td>26.5</td></tr><tr><td>EDM</td><td><strong>X</strong></td><td>3.42</td><td>76.7</td><td><strong>72.5</strong></td><td>17</td></tr><tr><td colspan="6"><strong>FFHQ (64 x 64)</strong></td></tr><tr><td>KL</td><td>$\pi(x, \sigma)$</td><td>3.04</td><td><strong>90.9</strong></td><td>88.1</td><td>87.4</td></tr><tr><td>rKL</td><td>$\pi(x, \sigma)$</td><td>3.07</td><td><strong>90.9</strong></td><td>87.9</td><td>68.2</td></tr><tr><td>EDM</td><td><strong>X</strong></td><td><strong>2.60</strong></td><td>90.5</td><td><strong>89.0</strong></td><td>59</td></tr><tr><td colspan="6"><strong>ImageNet (64 x 64)</strong></td></tr><tr><td>KL</td><td>$\pi(x, \sigma)$</td><td>2.96</td><td>85.4</td><td><strong>92.3</strong></td><td>51.8</td></tr><tr><td>rKL</td><td>$\pi(x, \sigma)$</td><td><strong>2.92</strong></td><td>85.4</td><td>92.1</td><td>52.1</td></tr><tr><td>EDM</td><td><strong>X</strong></td><td>3.01</td><td><strong>85.8</strong></td><td>91.8</td><td>350</td></tr></tbody></table>

trade-off during sampling.

## 6.2 Learning Adaptive Renoising

We next use the empirical dataset distribution as the expert to improve the EDM sampler through renoising and stochasticity injection. Experiments are conducted on CIFAR-10 32 × 32, FFHQ 64 × 64, and ImageNet 64 × 64.

**Learning to renoise** Renoising consistently improves performance on CIFAR-10 and ImageNet, highlighting the fact that going beyond fixed noise schedules can be beneficial. On CIFAR-10, performance depends strongly on the choice

<figure><img src="image_6.png" /><figcaption>Figure 3: Learned stochasticity injection profiles $\gamma_{EDM}(x, \sigma)$ on CIFAR-10 32 × 32 for different $f$-divergences.</figcaption></figure>

of divergence. KL-based objectives improve sample quality while reducing NFE, whereas Reverse-KL tends to allocate additional computation to high-noise regions that have limited influence on the final samples, which results in increased NFE without gains in FID. This observation shows that the choice of divergence influences not only the objective being optimized, but also the points along the sampling trajectory where the learned policy is most active.

**Learning $\gamma_{EDM}$** We now study learning the stochasticity injection profile $\gamma_{EDM}(x, \sigma)$. Learning $\gamma_{EDM}$ consistently matches or improves upon the baseline FID across datasets. More importantly, the behavior of the divergences is consistent with previous experiments. In this experiment as well, conditioning on $x$ reduces policy stochasticity; however, the induced behavior does not exhibit a simple or easily predictable structure, highlighting the importance of $x$-dependence for this task. On CIFAR-10, the Reverse-KL
<table><caption>Table 3: Performance comparison for learning the stochasticity injection profile $\gamma_{\text{EDM}}(x, \sigma)$. Best results are highlighted in bold.</caption><thead><tr><th>Method</th><th>Policy</th><th>FID</th><th>Precision</th><th>Recall</th></tr></thead><tbody><tr><td colspan="5"><strong>CIFAR-10 (32 x 32)</strong></td></tr><tr><td rowspan="2">KL</td><td>$\pi(x, \sigma)$</td><td>1.99</td><td>78.5</td><td><strong>73.2</strong></td></tr><tr><td>$\pi(\sigma)$</td><td>2.03</td><td>78.7</td><td>72.7</td></tr><tr><td rowspan="2">rKL</td><td>$\pi(x, \sigma)$</td><td>2.03</td><td><strong>79.7</strong></td><td>71.9</td></tr><tr><td>$\pi(\sigma)$</td><td>2.03</td><td>78.6</td><td>72.4</td></tr><tr><td>EDM</td><td><strong>X</strong></td><td><strong>1.98</strong></td><td>78.7</td><td>72.9</td></tr><tr><td colspan="5"><strong>FFHQ (64 x 64)</strong></td></tr><tr><td>KL</td><td>$\pi(x, \sigma)$</td><td>2.53</td><td>90.4</td><td>88.2</td></tr><tr><td>rKL</td><td>$\pi(x, \sigma)$</td><td><strong>2.49</strong></td><td><strong>90.9</strong></td><td>87.9</td></tr><tr><td>EDM</td><td><strong>X</strong></td><td>2.56</td><td>90.1</td><td><strong>89.4</strong></td></tr><tr><td colspan="5"><strong>ImageNet (64 x 64)</strong></td></tr><tr><td>KL</td><td>$\pi(x, \sigma)$</td><td>2.14</td><td>85.7</td><td><strong>92.0</strong></td></tr><tr><td>rKL</td><td>$\pi(x, \sigma)$</td><td>2.24</td><td><strong>86.3</strong></td><td>91.1</td></tr><tr><td>EDM</td><td><strong>X</strong></td><td><strong>2.12</strong></td><td>85.9</td><td>91.7</td></tr></tbody></table>

and KL have totally different behaviors with the Reverse-KL adding less noise with low variance at low-noise regions to increase precision. The KL is the opposite adding more noise with higher variance to increase diversity. Figure 3 illustrates the learned profiles.

## 6.3 Impact of Temperature on the Policy

Temperature is a standard mechanism for controlling the trade-off between sample diversity and fidelity, notably in the context of large language models. In our framework, the sampling policy is stochastic and draws actions from a categorical distribution, which allows us to directly control the entropy of the policy by adjusting its temperature at inference time. Letting $\beta$ denote the temperature parameter, the action $a_t$ at time $t$ is sampled according to $a_t \sim \pi_{\theta}(a_t | s_t)^{1/\beta}/Z$, where $Z$ is the normalizing constant. Lower temperatures concentrate probability mass on high-likelihood actions, yielding more deterministic behavior, while higher temperatures flatten the distribution and encourage exploration. We study the impact of temperature on learned policies for noise injection on CIFAR-10 32 x 32. As shown in Figure 4, increasing the temperature leads to a higher number of function evaluations (NFE), as the policy more frequently selects actions that inject noise and revisit higher-noise states. This increased exploration results in longer trajectories, but has a limited effect on sample quality. In particular, Precision and Recall remain relatively stable across a wide range of temperatures, indicating that the learned policy is robust to moderate changes in entropy at sampling time. A similar behavior is observed for policies learned to control the stochasticity

injection parameter $\gamma_{\text{EDM}}$ (see Appendix B.3). In contrast, it is observed that trying to enforce a more deterministic behavior of the learned policy (i.e., using smaller $\beta$) is detrimental in terms of image quality (as measured by FID). This observation suggests that deterministic schedules, that are commonly used in practice, may be inherently sub-optimal.

<figure><img src="image_7.png" /><figcaption>Figure 4: Impact of temperature on Precision, Recall, and NFE for learned renoise policies on CIFAR-10 32 x 32. Higher temperature increases diversity and NFE.</figcaption></figure>

## 7 Conclusion

We introduced an inference-time framework for adapting diffusion sampling by optimizing sampling policies via inverse reinforcement learning, without retraining the denoiser or learning explicit rewards. By matching expert state occupancy measures through f-divergence minimization, the method provides an interpretable way to control sampling behavior. Experiments on CIFAR-10, FFHQ, and ImageNet indicate improved cost-quality trade-offs across multiple sampling strategies.

While our framework provides a flexible and effective approach for improving diffusion sampling, it currently has some limitations. First, our experiments show sensitivity to initialization, as poorly initialized policies can converge to suboptimal local minima, highlighting the need for better heuristics or pretraining. Second, performance depends strongly on the design of the discrete action space, with poorly chosen actions degrading sampling quality. Exploring continuous action spaces may alleviate this issue for some strategies and is left for future work.
# Acknowledgements

This work was granted access to the HPC resources of IDRIS under the allocations 2025-AD011016159R1, 2025-A0181016159 and 2025-AD011017138 made by GENCI. This project was partly funded by the French government's Agence Nationale de la Recherche under the "France 2030" program (ANR-23-IACL-0008, PRAIRIE-PSAI and ANR-22- PESN-0014).

# References

[1] Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.

[2] Samaneh Azadi, Catherine Olsson, Trevor Darrell, Ian Goodfellow, and Augustus Odena. Discriminator Rejection Sampling, February 2019. URL http://arxiv.org/abs/1810.06758. arXiv:1810.06758 [cs, stat].

[3] Iskander Azangulov, Peter Potaptchik, Qinyu Li, Eddie Aamari, George Deligiannidis, and Judith Rousseau. Adaptive diffusion guidance via stochastic optimal control, 2025. URL https://arxiv.org/abs/2505.19367.

[4] Mohammad Gheshlaghi Azar, Mark Rowland, Bilal Piot, Daniel Guo, Daniele Calandriello, Michal Valko, and Rémi Munos. A general theoretical paradigm to understand learning from human preferences, 2023. URL https://arxiv.org/abs/2310.12036.

[5] Roberto Barceló, Cristóbal Alcázar, and Felipe Tobar. Avoiding mode collapse in diffusion models fine-tuned with reinforcement learning, 2024. URL https://arxiv.org/abs/2410.08315.

[6] Giulio Biroli, Tony Bonnaire, Valentin de Bortoli, and Marc Mézard. Dynamical Regimes of Diffusion Models, February 2024. URL http://arxiv.org/abs/2402.18491. arXiv:2402.18491.

[7] Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305.13301, 2023.

[8] Muthu Chidambaram, Khashayar Gatmiry, Sitan Chen, Holden Lee, and Jianfeng Lu. What does guidance do? A fine-grained analysis in a simple setting. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, *Advances in Neural Information Processing Systems*, volume 37, pages 84968–85005. Curran Associates, Inc., 2024. doi: 10.5220/079017-2698.

[9] Imre Csiszár. Information-type measures of difference of probability distributions and indirect observations.

Studia Scientiarum Mathematicarum Hungarica, 2:299–318, 1967.

[10] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In *2009 IEEE Conference on Computer Vision and Pattern Recognition*, pages 248–255, June 2009. doi: 10.1109/CVPR.2009.5206848. URL https://ieeexplore.ieee.org/document/5206848. ISSN: 1063-6919.

[11] Prafulla Dhariwal and Alex Nichol. Diffusion Models Beat GANs on Image Synthesis, June 2021. URL http://arxiv.org/abs/2105.05233. arXiv:2105.05233 [cs, stat].

[12] Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, and Ricky T. Q. Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control, 2025. URL https://arxiv.org/abs/2409.08861.

[13] Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Reinforcement learning for fine-tuning text-to-image diffusion models. In *Thirty-seventh Conference on Neural Information Processing Systems*, 2023.

[14] Chelsea Finn, Paul Christiano, Pieter Abbeel, and Sergey Levine. A Connection between Generative Adversarial Networks, Inverse Reinforcement Learning, and Energy-Based Models, November 2016. URL http://arxiv.org/abs/1611.03852. arXiv:1611.03852 [cs].

[15] Seyed Kamyar Seyed Ghasemipour, Richard Zemel, and Shixiang Gu. A Divergence Minimization Perspective on Imitation Learning Methods, November 2019. URL http://arxiv.org/abs/1911.02256. arXiv:1911.02256 [cs].

[16] Jianshu Guo, Wenhao Chai, Jie Deng, Hsiang-Wei Huang, Tian Ye, Yichen Xu, Jiawei Zhang, Jenq-Neng Hwang, and Gaoang Wang. Versat2i: Improving text-to-image models with versatile reward. arXiv preprint arXiv:2403.18493, 2024.

[17] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium. In *Advances in Neural Information Processing Systems*, volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper/2017/hash/8a1d694707eb0fefe65871369074926d-Abstract.html.
[18] Jonathan Ho and Stefano Ermon. Generative Adversarial Imitation Learning, June 2016. URL http://arxiv.org/abs/1606.03476. arXiv:1606.03476 [cs].

[19] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance, 2022. URL https://arxiv.org/abs/2207.12598.

[20] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. arXiv, December 2020. URL http://arxiv.org/abs/2006.11239. arXiv:2006.11239, NeurIPS 2020.

[21] Alexia Jolicoeur-Martineau, Ke Li, Rémi Piché-Taillefer, Tal Kachman, and Ioannis Mitliagkas. Gotta Go Fast When Generating Data with Score-Based Models, May 2021. URL http://arxiv.org/abs/2105.14080. arXiv:2105.14080 [cs, math, stat].

[22] Tero Karras, Samuli Laine, and Timo Aila. A Style-Based Generator Architecture for Generative Adversarial Networks, March 2019. URL http://arxiv.org/abs/1812.04948. arXiv:1812.04948 [cs, stat].

[23] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the Design Space of Diffusion-Based Generative Models, October 2022. URL http://arxiv.org/abs/2206.00364. arXiv:2206.00364, NeurIPS 2022.

[24] Dongjun Kim, Yeongmin Kim, Se Jung Kwon, Wanmo Kang, and Il-Chul Moon. Refining Generative Process with Discriminator Guidance in Score-based Diffusion Models. April 2023. URL http://arxiv.org/abs/2211.17091. arXiv:2211.17091, ICML 2023.

[25] Pum Jun Kim, Yoojin Jang, Jisu Kim, and Jaejun Yoo. TopP&R: Robust Support Estimation Approach for Evaluating Fidelity and Diversity in Generative Models, June 2023. URL http://arxiv.org/abs/2306.08013. arXiv:2306.08013 [cs].

[26] Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved Precision and Recall Metric for Assessing Generative Models. In 33rd Conference on Neural Information Processing Systems (NeurIPS 2019), Vancouver, Canada., October 2019. arXiv: 1904.06991.

[27] Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence for score-based generative modeling with polynomial complexity, May 2023. URL http://arxiv.org/abs/2206.06227. arXiv:2206.06227 [cs, math, stat].

[28] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps, October 2022. URL http://arxiv.org/abs/2206.00927. arXiv:2206.00927 [cs].

[29] Tianwei Ni, Harshit Sikchi, Yufei Wang, Tejus Gupta, Lisa Lee, and Benjamin Eysenbach. f-IRL: Inverse Reinforcement Learning via State Marginal Matching, December 2020. URL http://arxiv.org/abs/2011.04709. arXiv:2011.04709, CoRL 2020.

[30] Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-GAN: Training Generative Neural Samplers using Variational Divergence Minimization, June 2016. URL http://arxiv.org/abs/1606.00709. arXiv:1606.00709 [cs, stat].

[31] Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D. Manning, and Chelsea Finn. Direct Preference Optimization: Your Language Model is Secretly a Reward Model, December 2023. URL http://arxiv.org/abs/2305.18290. arXiv:2305.18290 [cs].

[32] Kevin Roth, Aurelien Lucchi, Sebastian Nowozin, and Thomas Hofmann. Stabilizing Training of Generative Adversarial Networks through Regularization. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper/2017/hash/7bccfde7714a1ebadf06c5f4cea752c1-Abstract.html.

[33] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S. Sara Mahdavi, Rapha Gontijo Lopes, Tim Salimans, Jonathan Ho, David J. Fleet, and Mohammad Norouzi. Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding, May 2022. URL http://arxiv.org/abs/2205.11487. arXiv:2205.11487 [cs].

[34] Riccardo De Santi, Marin Vlastelica, Ya-Ping Hsieh, Zebang Shen, Niao He, and Andreas Krause. Provable maximum entropy manifold exploration via diffusion models, 2025. URL https://arxiv.org/abs/2506.15385.

[35] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal Policy Optimization Algorithms, August 2017. URL http://arxiv.org/abs/1707.06347. arXiv:1707.06347 [cs].

[36] Jiaming Song and Stefano Ermon. Bridging the Gap Between f-GANs and Wasserstein GANs. In Proceedings of the 37th International Conference on Machine Learning, pages 9078–9087. PMLR, November 2020. URL https://proceedings.mlr.press/v119/song20a.html. ISSN: 2640-3498.

[37] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising Diffusion Implicit Models. October 2022. URL http://arxiv.org/abs/2010.02502. arXiv:2010.02502, ICLR 2022.
[38] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. SCORE-BASED GENERATIVE MODELING THROUGH STOCHASTIC DIFFERENTIAL EQUATIONS. 2021.

[39] Masashi Sugiyama, Taiji Suzuki, and Takafumi Kanamori. Density Ratio Estimation: A Comprehensive Review. 2010.

[40] Hao Sun and Mihaela van der Schaar. Inverse reinforcement learning meets large language model post-training: Basics, advances, and opportunities, 2025. URL https://arxiv.org/abs/2507.13158.

[41] Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diament, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine-tuning of continuous-time diffusion models as entropy-regularized control. arXiv preprint arXiv:2402.15194, 2024.

[42] Alexandre Verine. Quality and Diversity in Generative Models through the lens of f-divergences. PhD thesis, January 2024. URL https://theses.fr/279916922.

[43] Alexandre Verine, Ahmed Mehdi Inane, Florian Le Bronnec, Benjamin Negrevergne, and Yann Chevaleyre. Improving Discriminator Guidance in Diffusion Models, June 2025. URL http://arxiv.org/abs/2503.16117. arXiv:2503.16117 [cs].

[44] Chaoqi Wang, Yibo Jiang, Chenghao Yang, Han Liu, and Yuxin Chen. Beyond Reverse KL: Generalizing Direct Preference Optimization with Diverse Divergence Constraints, September 2023. URL https://arxiv.org/pdf/2405.15624. arXiv:2309.16240.

[45] Yilun Xu, Mingyang Deng, Xiang Cheng, Yonglong Tian, Ziming Liu, and Tommi Jaakkola. Restart Sampling for Improving Generative Processes, November 2023. URL http://arxiv.org/abs/2306.14878. arXiv:2306.14878, NeurIPS 2023.

[46] Sangwoong Yoon, Himchan Hwang, Dohyun Kwon, Yung-Kyun Noh, and Frank C. Park. Maximum entropy inverse reinforcement learning of diffusion models with energy-based models, 2024. URL https://arxiv.org/abs/2407.00626.

[47] Brian D Ziebart, Andrew Maas, J Andrew Bagnell, and Anind K Dey. Maximum Entropy Inverse Reinforcement Learning.
# Appendix

## A Proofs

### A.1 Gradient formula proof

*Proof.* Let's consider our objective function:

$$ \mathcal{L}_f(\theta) = \mathcal{D}_f(\mu_E \| \mu_\theta) $$

We can compute the gradient of $\mathcal{L}_f(\theta)$ with respect to $\theta$ using the chain rule and the properties of the f-divergence:

$$ 
\begin{aligned}
\nabla_{\theta} \mathcal{L}_f(\theta) &= \nabla_{\theta} \mathcal{D}_f(\mu_E \| \mu_{\theta}) \\
&= \nabla_{\theta} \int_S f \left( \frac{\mu_E(s)}{\mu_{\theta}(s)} \right) \mu_{\theta}(s) ds \\
&= \int_S \nabla_{\theta} \mu_{\theta}(s) f \left( \frac{\mu_E(s)}{\mu_{\theta}(s)} \right) - \mu_{\theta}(s) \frac{\nabla_{\theta} \mu_{\theta}(s)}{\mu_{\theta}(s)^2} \mu_E(s) \dot{f} \left( \frac{\mu_E(s)}{\mu_{\theta}(s)} \right) ds \\
&= \int_S \nabla_{\theta} \mu_{\theta}(s) h_f \left( \frac{\mu_E(s)}{\mu_{\theta}(s)} \right) ds \\
&= \int_S \nabla_{\theta} \log \mu_{\theta}(s) h_f \left( \frac{\mu_E(s)}{\mu_{\theta}(s)} \right) \mu_{\theta}(s) ds
\end{aligned}
$$

with $h_f(x) = f(x) - x\dot{f}(x)$.

Now, using the property of the state occupancy measure, we can write:

$$ 
\begin{aligned}
\nabla_{\theta} \mathcal{L}_f(\theta) &= \mathbb{E}_{s_{1:T} \sim p_{\theta}} \left[ \frac{1}{T} \sum_{t'=1}^{T} \nabla_{\theta} \log \mu_{\theta}(s_{t'}) h_f \left( \frac{\mu_E(s_{t'})}{\mu_{\theta}(s_{t'})} \right) \right] \\
&= \mathbb{E}_{s_{1:T} \sim p_{\theta}} \left[ \frac{1}{T} \sum_{t'=1}^{T} \left( \sum_{t=1}^{T'} \nabla_{\theta} \log \pi_{\theta}(a_t|s_t) \right) h_f \left( \frac{\mu_E(s_{t'})}{\mu_{\theta}(s_{t'})} \right) \right] \\
&= \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]
\end{aligned}
$$

by reindexing the sums.

<figure><img src="image_8.png" /></figure>

### A.2 Adding importance sampling

*Proof.* Adding importance sampling to the gradient formula derived in Theorem 4.1 gives:

$$ \nabla_{\theta} \mathcal{L}_f(\theta) = \mathbb{E}_{s_{1:T} \sim p_{\theta_0}} \left[ \frac{p_{\theta}(s_{1:T})}{p_{\theta_0}(s_{1:T})} \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 $\theta_0$ is the parameter of the policy used to sample the trajectories.

<figure><img src="image_9.png" /></figure>

## B Experimental Details

We parameterize the sampling policy and the state density estimator using a shared neural architecture, implemented as a dual-head network based on an EDM-style Encoder-UNet. The resulting model enables joint representation learning while keeping the policy and density objectives decoupled.

**Backbone and Preconditioning** Given a noisy state $x$ and noise level $\sigma$, we apply the standard EDM preconditioning and map $\sigma$ to the corresponding normalized sigma variable $\tau$. The network is conditioned on $\tau$ and, when applicable, on
class labels for conditional generation. For non-stationary policies, the discrete trajectory step index is additionally provided as an explicit conditioning variable. We also use a pretrained ADM encoder as a frozen feature extractor. The input image is first processed by the ADM encoder, and the resulting latent representation is used as input to the policy and density heads. All ADM parameters are kept fixed throughout training.

**Policy Head** The policy head outputs logits over a discrete action space of size $A$, where each action corresponds to a sampling-time control operation (e.g., continuing denoising or restarting to a previous noise level). The policy head is implemented as convolutional layers with output channels equal to $A$, sharing the same architectural design and conditioning mechanisms as the density head. For stationary policies, the policy is conditioned only on the current state and noise level. For non-stationary policies, the trajectory step index is included as an additional conditioning signal. During training and inference, actions are sampled from the resulting categorical distribution.

**Density Head** The density head estimates a scalar quantity associated with the state distribution, used for density-ratio-based training objectives. It is implemented as convolutional layers with a single scalar output. The density head shares the same backbone structure and conditioning as the policy head but has independent parameters in the output layers. The output layers of the density head are initialized using Xavier initialization to stabilize early training.

## B.1 Training Details

**Initialization** The initialization of both the policy and density networks plays an important role in practice. Since the learning problem is non-convex, we do not guarantee convergence to an unique policy, and different initializations may lead to different local optima. We initialize the policy network using a task-specific heuristic corresponding to a reasonable sampling strategy (for instance favoring continuation of denoising at high noise levels and allowing restarts at later stages). This initialization provides a stable starting point and significantly improves training stability compared to random initialization.

The density network is initialized using expert trajectories and rollouts generated by the initialized policy. These expert samples are used to initialize the density-ratio estimator before joint training with the policy, which helps stabilize the early stages of optimization.

**Optimization** For both the policy and density networks, we apply an exponential moving average (EMA) to the network parameters during training. EMA is used for evaluation and inference, as it consistently yields more stable behavior and improved performance. Finally, when training the density estimator, we apply label smoothing at high noise levels. At large noise scales, the input states contain limited information, and hard labels can lead to unstable gradients. Label smoothing in this regime improves robustness and reduces variance during training.

We provide details of the complete training procedure in table 4.

<table><caption>Table 4: Parameters used for training the sampling policies across different datasets and strategies.</caption><thead><tr><th rowspan="2">Parameters</th><th colspan="3">Renoise</th><th colspan="3">Gammas</th><th colspan="2">Guidance</th></tr><tr><th>CIFAR-10</th><th>FFHQ 64×64</th><th>ImageNet 64×64</th><th>CIFAR-10</th><th>FFHQ 64×64</th><th>ImageNet 64×64</th><th>CIFAR-10</th><th>FFHQ 64×64</th></tr></thead><tbody><tr><td>Cond. model</td><td>×</td><td>×</td><td>✓</td><td>×</td><td>×</td><td>✓</td><td>✓</td><td>✓</td></tr><tr><td>LR</td><td>1e-4</td><td>1e-5</td><td>1e-5</td><td>1e-4</td><td>1e-5</td><td>1e-5</td><td>1e-4</td><td>1e-5</td></tr><tr><td># traj.</td><td>1024</td><td>1536</td><td>2048</td><td>1024</td><td>1536</td><td>2048</td><td>1024</td><td>1536</td></tr><tr><td># DRE Init.</td><td>100</td><td>200</td><td>200</td><td>100</td><td>200</td><td>200</td><td>100</td><td>200</td></tr><tr><td>Action dim.</td><td>4</td><td>4</td><td>4</td><td>10</td><td>10</td><td>20</td><td>20</td><td>20</td></tr><tr><td>Res.</td><td>32×32</td><td>64×64</td><td>64×64</td><td>32×32</td><td>64×64</td><td>64×64</td><td>32×32</td><td>64×64</td></tr><tr><td># params.</td><td>5.8M</td><td>5.8M</td><td>6.8M</td><td>5.8M</td><td>5.8M</td><td>6.8M</td><td>5.8M</td><td>5.8M</td></tr><tr><td># Noise levels</td><td>9</td><td>20</td><td>20</td><td>18</td><td>40</td><td>256</td><td>18</td><td>40</td></tr><tr><td>GPU hrs</td><td>30</td><td>48</td><td>80</td><td>60</td><td>80</td><td>160</td><td>100</td><td>150</td></tr><tr><td>GPU</td><td>V100</td><td>V100</td><td>H100</td><td>V100</td><td>V100</td><td>H100</td><td>V100</td><td>V100</td></tr></tbody></table>
**Algorithm 2** Policy Optimization via State-Marginal f-Divergence Minimization

**Require:** Pretrained denoiser F, policy $\pi_\theta$, discriminator $D_\phi$, replay buffer B, number of epochs $n_{\text{epoch}}$, batch size $N_{\text{batch}}$, expert state samples $S_E$, clipping parameter $\epsilon$, learning rate $\eta$, and policy update frequency K.

1: Initialize policy $\pi_\theta$ and discriminator $D_\phi$.

2: **for** $k = 1$ **to** $n_{\text{epoch}}$ **do**

3: **if** $k \bmod K = 0$ **then**

4: Generate trajectories

$$s_{1:T} = \{(s_t, a_t)\}_{t=1}^T$$

using policy $\pi_\theta$ and denoiser F.

5: Train discriminator $D_\phi$ to distinguish expert states $S_E$ from policy-generated states $\{s_t\}$.

6: Estimate density ratios using

$$\frac{\mu_E(s)}{\mu_\theta(s)} \approx \frac{D_\phi(s)}{1 - D_\phi(s)}.$$

7: Compute advantage estimates

$$A_t = \frac{1}{T} \sum_{t'=t}^{T} h_f \left( \frac{\mu_E(s_{t'})}{\mu_\theta(s_{t'})} \right), \quad h_f(x) = f(x) - x f'(x),$$

and store $(s_t, a_t, A_t)$ in buffer B.

8: Save current policy parameters: $\theta_0 \leftarrow \theta$.

9: **end if**

10: **for** each minibatch $\{(s_i, a_i, A_i)\}_{i=1}^{N_{\text{batch}}}$ sampled from B **do**

11: Compute PPO-style objective:

$$\hat{\mathcal{L}}_f(\theta) = \frac{1}{N_{\text{batch}}} \sum_{i=1}^{N_{\text{batch}}} \max(r_i(\theta) (A_i - \bar{A}), \text{clip}(r_i(\theta), 1-\epsilon, 1+\epsilon) (A_i - \bar{A})),$$

where

$$r_i(\theta) = \frac{\pi_\theta(a_i | s_i)}{\pi_{\theta_0}(a_i | s_i)}, \quad \bar{A} = \frac{1}{N_{\text{batch}}} \sum_{i=1}^{N_{\text{batch}}} A_i.$$

12: Update policy parameters:

$$\theta \leftarrow \theta - \eta \nabla_{\theta} \hat{\mathcal{L}}_f(\theta).$$

13: **end for**

14: **end for**

**Ensure:** Trained sampling policy $\pi_\theta$.

## B.2 Is it necessary to have $x$-dependent policies for CFG?

In this appendix, we analyze whether conditioning the sampling policy on the image state $x$ is necessary when learning adaptive classifier-free guidance (CFG). In our framework, this corresponds to comparing policies of the form $\pi_\theta(\omega | x, \sigma)$ against simpler policies that depend only on the noise level, $\pi_\theta(\omega | \sigma)$.

A noise-level-dependent policy learns a single guidance scale per noise level, which can be interpreted as a global trade-off between fidelity and diversity at each stage of the sampling process. While this is sufficient to recover standard hand-tuned guidance schedules, it cannot adapt guidance strength across samples. In contrast, an $x$-dependent policy can modulate the guidance scale based on the current state of the sample, allowing different images to receive different amounts of guidance at the same noise level.

Empirically, we observe that $x$-dependent policies consistently outperform $\sigma$-only policies in terms of sample quality. As reported in Table 1, conditioning on $(x, \sigma)$ improves FID, Precision, and Recall across datasets and divergences. Qualitatively, the learned policies exhibit sparse behavior: at a given noise level, only a subset of samples receives strong guidance, while the majority remain weakly guided. This behavior cannot be captured by a single scalar guidance value shared across all
samples.

This observation aligns with recent analyses of CFG, which suggest that excessive guidance can lead to mode collapse and loss of diversity, while insufficient guidance degrades fidelity. An *x*-dependent policy enables the sampler to selectively apply strong guidance where it is most beneficial, while preserving diversity elsewhere. In practice, this adaptivity appears crucial for achieving favorable precision–recall trade-offs without manual tuning.

<figure><img src="image_10.png" /><figcaption>Figure 5: Comparison of *x*-dependent and noise-level-dependent policies for adaptive CFG on CIFAR-10 using KL and Reverse-KL divergence minimization. The *x*-dependent policy achieves better FID, Precision, and Recall by selectively applying guidance based on the current sample state. The left profile corresponds to a policy depending only on noise level *σ*, while the right profile uses both *x* and *σ*. It is clear that the policy independent of *x* produces suboptimal results.</figcaption></figure>

We also trained an *x*-independent policy for the Gamma strategy on CIFAR-10, but it also underperformed compared to the *x*-dependent version, suggesting that adaptivity to the current sample state is generally beneficial across different sampling interventions.

<table><caption>Table 5: Comparison of *x*-dependent and noise-level-dependent policies for adaptive Gamma sampling on CIFAR-10 using KL and Reverse-KL divergence minimization. The *x*-dependent policy achieves better FID, Precision, and Recall by selectively adjusting the noise schedule based on the current sample state.</caption><thead><tr><th>Method</th><th>FID ↓</th><th>Precision ↑</th><th>Recall ↑</th><th><i>x</i>-dependent</th></tr></thead><tbody><tr><td rowspan="3">KL</td><td>1.99</td><td>78.5</td><td><b>73.2</b></td><td>×</td></tr><tr><td>2.03</td><td>78.6</td><td>72.7</td><td>✓</td></tr><tr><td>2.06</td><td>79.0</td><td>72.4</td><td>✓</td></tr><tr><td rowspan="3">rKL</td><td>2.03</td><td><b>79.7</b></td><td>71.9</td><td>×</td></tr><tr><td>2.03</td><td>78.6</td><td>72.4</td><td>✓</td></tr><tr><td>2.06</td><td>79.2</td><td>72.5</td><td>✓</td></tr><tr><td>EDM</td><td><b>1.98</b></td><td>78.7</td><td>72.9</td><td>×</td></tr></tbody></table>

## B.3 Temperature

We experimented with adjusting the temperature of the policy during evaluation to study its effect on the trade-off between sample quality and diversity. Lowering the temperature sharpens the action distribution, leading to more deterministic behavior, while increasing it encourages exploration by flattening the policy over noise-injection actions. As discussed in Section 5.2, temperature therefore provides an inference-time mechanism to control the exploration-exploitation balance of the learned policy.

Figure 6 shows that increasing temperature leads to higher Recall but also higher FID. For the Gamma strategy, this behavior can be attributed to more frequent stochasticity injection at higher noise levels, which increases sample diversity
<figure><img src="image_11.png" /><figcaption>Figure 6: Effect of temperature scaling on the learned Gamma sampling policy for CIFAR-10 using KL divergence minimization.</figcaption></figure>

but also perturbs trajectories that would otherwise converge to high-fidelity modes. As a result, diversity gains come at the cost of reduced sample sharpness.

Notably, we observe that Precision and Recall are jointly optimized near the temperature used during training ($\beta = 1$). This suggests that the learned policy is already calibrated to balance exploration and exploitation under the training objective, and that deviating significantly from this temperature primarily trades off fidelity for diversity rather than uncovering strictly better operating points.

## B.4 Controlling the Distribution of Generated Samples

We also considered settings where the output distribution of a pretrained model is biased and must be corrected at inference time. We focused on renoising strategies. On CIFAR-10, although the dataset is balanced, pretrained EDM models overproduce animal classes relative to vehicle classes. We defined a target class distribution that increases the mass of vehicle classes while preserving normalization. Figure 7 shows that the learned renoise policy shifts the generated class distribution toward the target. While the match is not exact, the bias is substantially reduced, indicating that the policy captures the intended correction direction.

A key modeling choice is whether the control policy depends explicitly on the trajectory step of the induced MDP. We distinguish between non-stationary policies $\pi_{\theta}(x, \sigma, t)$, which depend on the step index, and stationary policies $\pi_{\theta}(x, \sigma)$, which don't. Although the optimal policy for a finite sampling horizon is generally non-stationary, we find that stationary policies are significantly easier to train and often perform better in practice. Across all strategies, stationary policies consistently match or outperform non-stationary ones, yielding more stable optimization and better trade-offs between distributional control and sample quality at comparable NFE.
<figure><img src="image_12.png" /><figcaption>Figure 7: Class distribution of samples generated by a pretrained EDM model on CIFAR-10, before and after applying a learned renoise policy to correct class imbalance. The target distribution increases the mass of vehicle classes (automobile, truck) relative to animal classes. The learned policy shifts the generated distribution closer to the target, reducing bias while maintaining sample quality.</figcaption></figure>