Maximum Entropy Inverse Reinforcement Learning of Diffusion Models with Energy-Based Models
Authors: Sangwoong Yoon
Himchan Hwang
Dohyun Kwon
Yung-Kyun Noh
Frank C. Park
Venue: NeurIPS 2024 oral (2024)
Keywords: diffusion models
inverse reinforcement learning
dynamic programming
reinforcement learning
generative modeling
Primary area: generative_models
TL;DR: We present an inverse reinforcement learning framework for training diffusion models and provide a novel RL algorithm for diffusion models which leverages value functions.

Abstract: We present a maximum entropy inverse reinforcement learning (IRL) approach for improving the sample quality of diffusion generative models, especially when the number of generation time steps is small. Similar to how IRL trains a policy based on the reward function learned from expert demonstrations, we train (or fine-tune) a diffusion model using the log probability density estimated from training data. 
Since we employ an energy-based model (EBM) to represent the log density, our approach boils down to the joint training of a diffusion model and an EBM. Our IRL formulation, named Diffusion by Maximum Entropy IRL (DxMI), is a minimax problem that reaches equilibrium when both models converge to the data distribution. The entropy maximization plays a key role in DxMI, facilitating the exploration of the diffusion model and ensuring the convergence of the EBM. We also propose Diffusion by Dynamic Programming (DxDP), a novel reinforcement learning algorithm for diffusion models, as a subroutine in DxMI. DxDP makes the diffusion model update in DxMI efficient by transforming the original problem into an optimal control formulation where value functions replace back-propagation in time. Our empirical studies show that diffusion models fine-tuned using DxMI can generate high-quality samples in as few as 4 and 10 steps.  Additionally, DxMI enables the training of an EBM without MCMC, stabilizing EBM training dynamics and enhancing anomaly detection performance.

Maximum Entropy Inverse Reinforcement Learning of Diffusion Models with Energy-Based Models
Sangwoong Yoon', Himchan Hwang?, Dohyun K.on Yung-Kyun Noh fohl,4f, Frank C. Park2.5t Korea Institute for Advanced Study, "Seoul National University, *University of Seoul, "Hanyang University, 'Saige Research, 'Co-corresponding Authors swyoon@kias, re. kr. himchan@robotics. snu. ac ac.kr, dhkwon@uos.c ac. kr, nohyungdhanyang. ac kr, fcpêsnu.ar aac. .c.
Abstract
We present maximum entropy inverse reinforcement learning (IRL) approach for improving the sample quality of diffusion generative models, especially when the number of generation time steps is small. Similar to how IRL trains a policy based the reward function learned from expert demonstrations, a diffusion model using the log probability density estimated from training data. Since we employ energy-based model (EBM) to represent the log density, approach boils down to the joint training diffusion model and an EBM. Our IRL formulation, named Diffusion by Maximum Entropy IRL (DxMI), is a minimax problem that reaches equilibrium when both models converge to the data distribution. The entropy maximization plays key role in DxMI, facilitating the exploration of the diffusion model and ensuring the convergence of the EBM. We also propose Diffusion by Dynamic Programming (DxDP), a novel reinforcement learning algorithm for diffusion models, as subroutine in DxMI. DxDP makes the diffusion model update in DxMI efficient by transforming the original problem into an optimal control formulation where value functions replace back propppprppagtiion in time. Our empirical studies show that diffusion models fine-tuned using DxMI can generate high- high-quality samples in as few as 4 and 10 steps. Additionally, DxMI enables the training of an EBM without MCMC, stabilizing EBM training dynamics and enhancing anomaly detection performance.
train (or fine-tune)
Introduction
Generative modeling is a form of imitation learning. Just as an imitation learner produces an action that mimics a demonstration from an expert, generative model synthesizes sample resembling the training data. In generative modeling, the expert to be imitated corresponds to the underlying data generation process. The intimate connection between generative modeling and imitation learning is already well appreciated in the literature 2
The connection imitation learning plays a central role in diffusion models 3 which generate samples by transforming a Gaussian noise through iterative additive refinements. The training of diffusion model essentially an instance behavioral cloning (5), a widely adopted imitation learning algorithm that mimics an expert's 's action at each state. During training, diffusion model is optimized follow predefined diffusion trajectory that interpolates between noise and data. The trajectory provides a step-by-step demonstration of how to transform Gaussian noise into a sample, allowing diffusion models to achieve new state-of-the of-tth-art in many generation tasks.
Behavioral cloning is also responsible for the diffusion model's 's key limitation, the slow generation speed. A behavior-cloned policy reliable when the state distribution deviates from the expert demonstration 6 7] Likewise, the sample quality from a diffusion model degrades as the gap
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Diffusion Model
Xp
Xy-1
X7 x)
(xx)
q
(x)
Positive Sample
Negative Sample
Reward
Training Data
Energy- Energy-Based Model
Figure 1: (Left) Overview of DxMI. The diffusion model I (x) is trained using the energy of q(x) as a reward. The EBM ((x) trained using samples from (xx as negative samples. (Right) ImageNet 64 generation examples from 10-step diffusion model before DxMI fine-tuning (up) and after fine-tuning (down). Only the last six steps out of ten are shown.
between training and generation grows. A diffusion model is typically trained to follow a fine-grained diffusion trajectory of 1,000 or more steps. Since 1,000 neural network evaluations are prohibitively expensive, fewer steps are often used during generation, incurring the distribution shift from the training phase and thus degraded sample quality. Speeding up a diffusion model while maintaining its high sample quality is a problem of great practical value, becoming an active field of research [
The slow generation in diffusion models can be addressed by employing inverse reinforcement learning (IRL; [14)), another imitation learning approach. Unlike behavioral cloning, which blindly mimics the expert's action each state, the IRL approach first infers the reward function that explains the trajectory. When applied to diffusion models, IRL allows a faster generation trajectory to found by guiding a sampler using the learned reward [[] 12 13 This approach more frequently referred to as adversarial training because common choice for the reward function is a discriminator classifying the training data and diffusion model's samples. However, resembling GAN, this adversarial approach may have similar drawbacks, such as limited exploration.
The first contribution of this paper is formulating maximum entropy IRL [15 16 model. Our formulation, named Diffusion by Maximum Entropy IRL (DxMI, pronounced "di- by-me"), minimax problem that jointly optimizes diffusion model and an energy- model (EBM). In DxMI, the EBM provides the estimated log density the reward signal for the diffusion model. Then, the diffusion model is trained maximize the reward from EBM while simultaneously maximizing the entropy of generated samples. The maximization of entropy in DxMI facilitates exploration and stabilizes the training dynamics, as shown in reinforcement learning (RL) [17. IRL. [15. and EBM (18) literature. Furthermore, the entropy maximization lets the minimax problem have an equilibrium when both the diffusion model and the EBM become the data distribution.
for a diffusion
The diffusion model update step in DxMI is equivalent to maximum entropy RL. However, this step challenging to perform for two reasons. First, the diffusion model update requires the gradient of the marginal entropy of a diffusion model's 's samples, which is difficult to estimate for discrete-time diffusion models, such as DDPM [3]. Second, propggatngg the gradient through the whole diffusion model is often infeasible due to memory constraints. Even with sufficient memory, the gradient may explode or vanish during propagation through time, causing training instability [19
Our second contribution is Diffusion by Dynamic Programming (DxDP), a novel maximum entropy RL algorithm for updating a diffusion model without the above-mentioned difficulties. First, DxDP mitigates the marginal entropy estimation issue by optimizing the upper bound of the original objective. In the upper bound, the marginal entropy is replaced by conditional entropies, which are easier to compute for diffusion model. Then, DxDP removes the need for back- propagation in time by interpreting the objective as an optimal control problem and applying dynamic programming using value functions. The connection between optimal control and diffusion models is increasingly gaining attention (00 21 22), but DxDP the first instance of applying discrete-time dynamic programming directly to train a diffusion model. Compared to policy gradient methods, we empirically find that DxDP converges faster and provides stronger diffusion models. As an RI. algorithm, DxDP may have broader utility, such as fine-tuning a diffusion model from human AI feedback.

---

We provide experimental results demonstrating the effectiveness of DxMI in training diffusion models and EBMs. On image generation tasks, DxMI can train strong short-run diffusion models that generate samples in 4 or 10 neural network evaluations. Also, DxMI can be used to train strong energy-based anomaly detectors. Notably, DxMI provides a novel way to train EBM without MCMC, which is computationally demanding and sensitive to the choice of hyperparameters.
The paper is structured as follows. In Section 2, we introduce the necessary preliminaries. Section presents the DxMI framework, and Section 4 proposes the DxDP algorithm. Experimental results and related work are provided in Sections 5 and 6, respectively. Section 7 concludes the paper. The code for DxMI can be found in https: /tttps:/ggthhn github.c.... com/ htt/://o///ffoiiiio/////// Ctrprfuusiioo-bb---attttttiiin MaxE-tXEEttt.RRg..
2 Preliminaries
Diffusion Models. The diffusion model refers range of generative models trained by reversing the trajectory from the data distribution to the noise distribution. Among diffusion models, we focus on discrete-time stochastic samplers, such DDPM (3) which synthesize sample through the following iteration producing Xo, x XT E RD::
Xo 1), and X++1 0,X, + () O for t= 0, 1.....,,T-11
where (), and u(x,() )) is the output of a neural network. The coefficients Oe E R and Tr E R250 are constants. Note that reverse the time direction in a diffusion model from the convention to be consistent with RL. The final state XT is the sample generated by the diffusion model, and its marginal distribution is (xT). We will often drop the subscript T and write the distribution as π (x). The conditional distribution of transition in Eq. (I) is denoted as (Xx++ (x.). For continuous-time diffusion models [3], Eq. ① corresponds using the Euler- Maller-aruyymaa solver.
The generation process in Eq. ① defines a T-horizon Markov Decision Process (MDP) except for the reward. State s; and action a, defined s, (x,, t) and a, X(+1- The transition dynamics is defined as p( SS,, a.) = ((a,,++1). where (tx,,)) is a Dirac delta function at (x,, t). With a reward function defined, diffusion model can be trained with RL [11 [2 19 25 In this paper, we consider a case where the reward is the log data density log p(x), which is unknown.
Energy-Basee -Based Models. An energy-based model (EBM) g(x) uses scalar function called energy E(x) to represent a probability distribution:
q(x)
E: X - IR,
(2)
Z
where >> 0 is temperature, X' is the compact domain of data, and Z Jxx exp(-E(xx//)ddx < 0o the normalization constant.
The standard method for training an EBM is by minimizing KL divergence between data p(x) and EBM q(x), i.e min, K L(pllg), where K L(pllg) := fx p(x) log(p(x))q(())dd Computing the gradient of KT(pllg) with respect to EBM parameters requires MCMC sampling, which is computationally demanding and sensitive to hyperparameters. The algorithm presented in this paper serves an alternative method for training an EBM without MCMC.
EBMs have a profound connection to maximum entropy IRL, where Eq. (2 serves as a model for an expert's policy [1] 16 2 In maximum entropy IRL, x corresponds to an action (or sequence of actions), and (x) represents the expert's 's cost of the action. The expert is then assumed to generate actions following (x)). This assumption embodies the maximum entropy principle because Eq. (2)
a distribution that minimizes the cost while maximizing the entropy of the action. Here, T balances cost minimization and entropy maximization.
3 Diffusion by Maximum Entropy Inverse Reinforcement Learning
3.1 Objective: Generalized Contrastive Divergence
We aim to minimize the (reverse) KL divergence between a diffusion model (x) and the data density p(x). This minimization can improve the sample quality of =xx). particularly when T is small.
min K (L( (x)xx( max F, p(x) + H(T(x)), #elt
3

---

where II is the set of feasible (x)'s, ss and H(()) = J [x) log π (x) is the differential entropy. This minimization is a maximum entropy RL problem: The log data density log)(x) is the reward, and mxx) is the stochastic policy. However, we cannot solve Eq. (3) directly since log p(x) is
unknown in a typical generative modeling setting. Instead, training data from p(x) are available, allowing us to employ an Inverse RL. approach.
In this paper, we present Diffusson by Maximum Entropy IRL (DxMI) as an IRL approach for solving Eq. (3). We employ an EBM ((x) (Eqq ② as a surrogate for p(x) and use log g(x) as the reward for training the diffusion model instead of logp(x)) At the same time, we train q(x) to be close to p(x) by minimizing the divergence between p(x) and g(x):
min K (TT())))) fet
and
min (L(p(())l((xx)) rE
(4)
where 2 is the feasible set of EBMs. When the two KL divergences become zero, p(x) = q(x) π (x) is achieved. However, minge o K (p(x))lq(()) is difficult due to the normalization constant of (xx). Instead, we consider an alternative minimax formulation inspired by Contrastive Divergence (CD; 26)). celebrated algorithm for training an EBM.
Objective. Let p(x) be the data distribution. Suppose that Q and II are the feasible sets of the EBM (xx) and the diffusion model (xx), respectively. The learning problem of DxMI is formulated as follows:
min max {T(p(x))lg(x))) - -KL(m())ll((x)).
(5)
QE@ elll
We shall call Eq. ( Generulized CD (GCD), because Eq. generalizes CD by incorporating a gen-
eral class of samplers. CD (26) is originally given as mingeg K [[(())) L
pig (x)))((x).
Here, v/e (x) is a k-step MCMC sample distribution where Markov chains having g(x) as a stationary distribution are initialized from p(x). GCD replaces 1/p (x) with a general sampler xxx) at the expense of introducing max operator.
When the models are well-specified, i.c., p(x) E Q = 11, Nash equilibrium of GCD is p(x) = g(x) * (x), which is identical to the solution of Eq. 44p Meanwhile, there is need to compute the normalization constant, as the two KL divergences cancel the normalization constant out. Note that the objective function (Eq. (5) can be negative, allowing g(x) be closer to p(x) than π rxx).
Our main contribution is exploring the application of discrete-time diffusion models (Eq. [ as rxx) in GCD and not discovering GCD for the first time. GCD is mathematically equivalent to formulation called variational maximum likelihood or adversarial EBM, which has appeared several
times in EBM literature [27 28, 29 30 B1 32 33). However, none of them have investigated the use of a diffusion model as ग (x), where optimization and entropy estimation are challenging. We discuss
the challenges in Section 3.2 and provide a novel algorithm to address them in Section
3.2 Alternating Update of EBM and Diffuusion Model
In DxMI, we update a diffusion model and an EBM in an alternative manner to find the Nash equilibrium. We write 0 and ó as the parameters of the energy (x) and a diffusion model π (x).
respectively. While EBM update is straightforward, we require subroutine described in Section
for updating the diffusion model. The entire procedure of DxMI is summarized in Algorithm I
EBM Update. The optimization with respect EBM is written as ming Eptx)) (x))
(x)(() [x)]. During the update, also regularize the energy by the square of energies ) [xx + (x)) for y >> ensure the energy is bounded. We set y 1 unless stated otherwise. This regularizer is widely adopted in EBM [34.
Difficulty of Diffusion Model Update. Ideally, diffusion model parameter á should be updated by minimizing K KL(TTg (TTT)ll90) = ( (x)) However, this update is difficult in practice for two reasons.
First, marginal entropy (x)= is difficult to estimate. Discrete-time diffusion models (Eq. 1
do not provide an efficient way to evaluate log Tx(()) required in the computation of H( (x)).
unlike some continuous-time models, e.g., continuous normalizing flows [36] 23 Other entropy
estimators based on k-nearest neighbors [37] or variational methods [38 39 do not scale well to

---

Algorithm Diffusion by Maximum Entropy IRL
1: Input: Dataset D, Energy Ex(x), Value V (x,), and Sampler (x(:))
2: 8t <- 0;
I/ Initialize Adaptive Velocity Regularization (AVR) Minibatch dimension omitted for brevity.
3: for x D do
4: Sample Xo:T To (x
5: ming (x) Ee (xx) (E + {x(r))
I/ Energy update I/ Value update
6: fort = I,..., do
2
7: ming V (x++1)) + log T (x(+11xx) + 2=|1xx x++1l1=- (x(),)
8:
end for
9: for randomly chosen among Xo.T do
I/ Sampler update Reparametrization trick
Sample one-step: X++1 ~ TH (x++1x1+
11: ming vt+1
4 (x++1(0)) + 1 (x)++ 25|1xx x ( X++1 a function of ø.
12: end for
of f- ast + (1 a)] X++1113/D
|/ AVR update
14: end for
high-diment high-dimenniioal spaces. Second, propagating the gradient through time in a diffusion model may require significant memory. Also, the gradient may explode vanish during propagation, making the training unstable [19].
Diffusion by Dynamic Programming
In this section, we present novel RL algorithm for diffusion model, Diffusion by Dynamic Programming (DxDP), which addresses the difficulties updating a diffusion model for the reward. DxDP leverages optimal control formulation and value functions to perform the diffusion model update in DxMI efficiently. Note that DxDP can be used separately from DxMI train diffusion model for arbitrary reward.
4.1 Optimal Control Formulation
Instead of solving ming K (xx(((x)))l directly, minimize its upper bound obtained from the data processing incquality:
K (xT((x))l)xxx))) I (L(TT(Xxo))] [
(6)
Here, we introduce auxiliary distribution ((xx.T-1|xx (xxT), and the inequality holds for an arbitrary choice of ÿ ((xor-11 |xr). In this paper, we consider particularly simple case where x(xorT1| (x)) factorized into conditional Gaussians as follows:
T-1
q(xoTT-
II (x+1),, where ((x+|X++1) = (x(+,,,,, t=o
> 0.
Now we minimize the right-hand side of Eq. (6): min& KL( 1l we plug in the definitions of each distribution, multiply by T, and discard all constants, we obtain the following problem:
(ooloo(())) (x1 (x | When
T-1
min E (x(x))
Ex(xT) ) +
log n (x x.) +
Σ
|
x;|12
(8)
t= 2s
t=0
which is an optimal control problem. The controller T-(-) is optimized minimize the terminal cost Ex(xr) plus the running costs for each transition Xx+1). The first running cost log To (x(xX+11Xx (x,) responsible for conditional entropy maximization, because The second running cost regularizes the "velocity" -X,41-x11'" The temperature balances between the terminal and running costs.
E, B=[loTTT (x(+g))( = -H(Ts(x441))
We have circumvented the marginal entropy computation in GCD, all terms in Eq. 8 are easily computable. For the diffusion model considered this paper (Eq. (I)) the conditional entropy has a particularly simple expression H[TT((x+) xx)) D log a, + 0.5D log 2T.. Therefore, optimizing the entropy running cost amounts to learning o, 's in diffusion, and we treat or as a part of the diffusion model parameter & in DxMI.
5

---

True E(x)
Ex(x) T= 0.1
Table 1: Quantitative results for Gaussians experi- ment. SW denotes the sliced Wasserstein distance between samples and data. AUC computed for classification between data and uniform noise using the energy. The standard deviation is computed from 5 independent samplings. The ideal maximum value of AUC is about 0.906.
T=0
T=1
Method
T
Pretrain
SW (L)
AUC (1)
DDPM
5
0.967=00.05
DDPM
10
0.824=0.002
DDPM
100
0.241
DDPM
1000
0.123
DxMI
* 1000
O
0
0.074=
0.707
DxMI
0.01
0.074=2.017
0.751
Figure 2: 2D density estimation on 8 Gaus- sians. Red shades indicate the energy (white low), and the dots are generated samples.
DXMI
5
0.1 0.068=0.00
0.898
DxMI
1
1.030=0.004
0.842
DxMI
X
0.1 0.076=0.011
0.883
4.2 Dynamic Programming
We propose dynamic programming approach for solving Eq. 8
Dynamic programming introduces
value functions break down the problem into smaller problems each timestep, removing the need for back propagation in time. Then, policy, diffusion model in our case, is optimized through iterative alternating applications of policy evaluation and policy improvement steps.
Value Function. value function, or cost-to-go function (x+), ) is defined as the expected sum of the future costs starting from xs, following स. We write the parameters of a value function as th.
T-1
T-1
(x:) = En (x( +
Σ log T (x Xx) + 8==1
Σ p=t
T
X
X xx
(9)
25
Note that (T E A value function can be implemented with neural network, but there are multiple design choices, such as whether to share the parameters with T (x) or Exx), and also whether the parameters should be shared across time. We explore the options in our experiments.
for O....
Policy Evaluation. During policy evaluation, estimate the value function for the current diffusion model by minimizing the Bellman residual, resulting in the temporal difference update.
min Ex, Xx, (sg[V + (xx-1)] + log To (x(+1 +
T ||x:
X
(x,)),, (10)
2s}
where sg[:] denotes a stop-gradient operator indicating that gradient is not computed for the term.
Policy Improvement. The estimated value is used to improve the diffusion model. For each in a trajectory Xo.r sampled from (Xo:)). the diffusion model optimized to minimize the next-state value and the running costs.
min Ex (xc vet (x4+1) vx+1(x(+) + 7 log To (x(4 2s
||x
xx+1|l
x1
(11)
In practice, each iteration of policy evaluation and improvement involves single gradient step.
Adaptive Velocity Regularization (AVR). We additionally propose method for sys- tematically determining the hyperparameter st's of the auxiliary distribution x(x)) can optimize 8, such that the inequality Eq. (6) is as tight as possible by solving mitlsg......r_ KL KL(==(Xo-T) (1 (x)).. After calculation (details in Appendix the optimal Sf can be obtained analytically: (st)² Ex, xxxxx [||x, x In practice, we can use exponential moving average compute the expectation Ex,x, i1 - [||x, | during training: of f- as? + (1 a)][x, / where set 0.99 for all experiment.

---

4.3 Techniques for Image Generation Experiments
When using DxDP for image generation, one of the most common applications of diffusion models, we introduce several design choices to DxDP to enhance performance and training stability. The resulting algorithm is summarized Algorithm 2
Time-Independent Value Function. In image generation experiments (Section 5.2 we let the value function be independent of time, i.e., (x,) Ve (x²). Removing the time dependence reduces the number of parameters to be trained. More importantly, a time-indepennent value function can learn better representation because the value function is exposed to diverse inputs, including both noisy and clean images. On the contrary, a time- dieendepeddent value function V (xx) never observes samples having different noise levels than the noise level of Xt-
Time Cost. Also, in the value update (Eq. (T0) step of image generation experiments, we introduce time cost function R(t) > 0, which replaces the running cost terms T log To (X (xx) - r||x, x || (2s?). The time cost R(t) only depends on time t. The modified value update equation given follows:
min Ex X (xx (x + R(t) V (x.)))
(12)
Meanwhile, we retain the running cost terms in the diffusion model (policy) update step (Eq. (II) The time cost R(t) is predetermined and fixed throughout training. The introduction of time cost is motivated by the observation that the running costs can fluctuate during the initial stage of training, posing difficulty in value function learning. The time cost stabilizes the training by reducing this variability. Moreover, the time cost ensures that the value function decreases monotonically over time. Such monotonicity is known to be beneficial in IRL for episodic tasks [16].
We employ two types of R(t): "linear" and "sigmoid". A linear time cost is given as R(t) c where we use c 0.05. The linear time cost encourages the value decrease linearly time progresses. The sigmoid time cost is R((1) = a(-(+7/2)) (apl))-((--) (ot--1111++- (T/22)), where axx) = (1 exp((xx)) With the sigmoid time cost, the value function trained to follow a sigmoid function centered at T T/2 when plotted against the time. Other forms of R(t) are also possible.
Separate tuning of entropy regularization log T (x(xx+1|xx and velocity regularization | X 1//1111/(20) (2s{)), such that the resulting running cost becomes Ti log To (Xx+ (x:) + r2l|x: Xx+1 12// Typically, we found Ti > 72 beneficial, indicating the benefit of exploration. Setting TI A does not violate our maximum entropy formulation, as scaling T2 is equivalent to scaling sf ss which can be set arbitrarily.
In image generation, we assign different values of temperature for
5 Experiments
In this section, we provide empirical studies that demonstrate the effectiveness of DxMI in training diffusion model and an EBM. We first present 2D example, followed by image generation and anomaly detection experiments. More details on experiments can be found in Appendix C
5.1 2D Synthetic Data
We illustrate how DxMI works on 2D 8 Gaussians data. DxMI is applied to train a five-step diffusion model (T = 5) with a corresponding time-dependent value network, both parametrized by time- conditioned multi-layer perceptron (MLP). The last time step (? 5) of the value network is treated as the energy. The sample quality is measured with sliced Wasserstein distance (SW) to test data. Also, quantify the quality of an energy function through the classification performance to uniform noise samples (Table I..
First, we investigate the effect of maximum entropy regularization т. Setting an appropriate value for 7 greatly benefits the quality of both the energy and the samples. When 0.1, the samples from DxMI have smaller SW than the samples from a full-length DDPM do. The energy also accurately captures the data distribution, scoring high AUC against the uniform noise. Without entropy regularization (r = 0). DxMI becomes similar to GAN 40) The generated samples align moderately well with the training data, but the energy does not reflect the data distribution. When is too large (T 1), the generated samples are close to noise. In this regime, DxMI behaves similarly
7

---

Table 2: CIFAR-10 unconditional image genera-
Table 3: ImageNet x4 conditional image gener- ation. the starting point of DxMI fine-tuning.
tion. t: the starting point of DxMI fine-tuning.
NFE FID (4) Rec. (f)
NFE FID (4) Prec. (1) Rec. (t)
Score SDE (VE) 23
2000
2.20
0.59
PD 10
2.57
ADM 47
250
2.07
0.74
0.63
Consistency Model
42
2
2.93
DFNO 48
1
8.35
PD [10
8.34
PD [T0
1
15.39
0.59
0.62
2-Rectified Flow
43
4.85
0.50
BigGAN-deep
4.06
0.79
0.48
Consistency Model
3.55
Backbone: EDM
StyleGAN-XL
1.85
0.47
EDM (Heun)
50
2.44
0.71
0.67
Backbone: DDPM
EDM (Ancestral)
10
50.27
0.37
0.35
DDPM 3
1000
3.21
0.57
EDM (Ancestral)
82,95
0.26
0.25
Pa
10
35.85
0.29
Consistency Model
42
2
4.70
0.69
0.64
DDIM 45
10
13.36
Consistency Model
42
6.20
0.68
0.63
SFT-PG 11
10
4.82
0.606
DxMI
10
2.68
0.777
0.574
DxMI
10
3.19
0.625
T= 0
10
2.72
0.782
0.564
T=
10
3.77
0.613
Linear time cost
10
2.81
0.742
0.594
Linear time
10
3.39
0.595
DxMI+Value Guidance
10
2.67
0,780
0.574
No time cost
10
5.18
0.595
T=0
10
2.76
0.786
0.560
DxMI + Value Guidance
10
3.17
0.623
DxMI
3.21
0.758
0.568
0=0
10
3.72
0.613
==0
3.65
0.767
0.552
Linear time cost
4
3.40
0.762
0.554
Backhone: DDGAN
DxMI+Value Guidance
4
3.18
0.763
0.566
DDGAN' [46
4.15
0.523
4 3.67
0.770
0.541
DxMI
4
3.65
0.532
to Noise Contrastive Estimation
41
enabling energy function learning to certain extent. These
effects visualized in Fig.
Next, we experiment on whether pre-training a sampler as DDPM helps DxMI. Table suggests that the pre-training is beneficial but not necessary to make DxMI work. We also visualize the value functions Fig. and find that the time evolution value interpolates the data distribution and Gaussian distribution.
5.2
Image Generation: Training Diffusion Models with Small T
On image generation tasks, we show that DxMI can be used to fine-tune a diffusion model with reduced generation steps, such as T
Table 4: LSUN Bedroom 256 x 256 unconditional
image generation.
10.
We test DxMI on unconditional CIFAR-10 [52) (32 x 32), conditional ImageNet (53) down- sampled to 64 x 64, and LSUN Bedroom [54)
NFE FID ()) Prec. (f) Rec. (f)
StyleGAN2 [5I
2.35
0.59
0.48
(256x25) 256). using three diffusion model back- bones, DDPM [3. DDGAN
Backbone: EDM
EDM (00
79
2.44
0.71
0.67
46 and variance The results
Consistency Model DxMI
422
2
5.22
0.68
0.39
exploding version of EDM
50
5.93
0.563
0.477
can be found in Table 2 3
and
Starting from
a publicly available checkpoint of each pretrained backbone, we first adjust the noise schedule for the target sampling steps 7. When adjusting the noise, for DDPM, we follow the schedule of FastDPM [8) and for EDM, we use Eq. (5) of [50]. No adjustment is made for DDGAN, which was originally built for T 4. The adjusted models
used the starting point of DxMI training. A single CIFAR-10 10 run reaches the best FID in less than 4 hours on four A GPUs. We set ti = 0.1 and T2 = 0.01. The sigmoid time cost is used for all image generation experiments. The sample quality measured by FID (55) Precision (Prec., (56)), and Recall (Rec., (56)). ResNet is used as our value function and is trained from scratch. More experimental details in Appendix C.2
Short-run diffusion models fine-tuned by DxMI display competitive sample quality. Unlike distillation methods, which are often limited by their teacher model's performance, DxMI can surpass the pre- trained starting point. Although DxMI does not support single-step generation, DxMI offers principled approach to training high-quality generative model with moderate computation burden

---

(Appendix D Note that DDGAN does not fit the formulation of DxMI, as x(x++1|x/ in DDGAN is not Gaussian. Nevertheless. DxMI can still enhance sample quality, showing its robustness.
Furthermore, DxMI outperforms SFT-PG [1], another IRL. approach implemented with a policy gradient. For fair comparison, we have ensured that the backbone and the initial checkpoint of SFT-PG and DxMI are identical. Thus, the performance gap can be attributed to the two differences between SFT-PG and DxMI. First, DxMI uses dynamic programming instead of policy gradient. In DxDP, the value function is more directly utilized to guide the learning of the diffusion model. Meanwhile, in policy gradient, the role of the value function is variance reduction. As SFT-PG also requires value function during training, the computational overhead of DxMI is nearly identical to SFT-PG. Second. DxMI incorporates the maximum entropy principle, which facilitates exploration.
We also conduct ablation studies for the components of DxMI and append the results in Table 2 and 3 First, the temperature parameters set 7ero π = 72 = 0 in the sampler update (II). Then, compare the linear time cost to the sigmoid time cost. In both cases, We observe the increase in FII) and the decrease in Recall.
To investigate whether the trained value function captures useful information, we implement value guidance, where we shift the trajectory of generation slightly along the value function gradient, similarly to classifier guidance [47] and discriminator guidance [57]. When sampling the next step Xt+1, we add a small drift with coefficient X, i... X f- X++1 Ao Vx.., (x)-))). We observe sample quality metric improvement until À is 0.5. This observation suggests that the value function gradient is aligned well with the data density gradient.
5.3 Energy-Based Anomaly Detection and Localization
We demonstrate the ability of DxMI to train an accurate energy function on an anomaly detection task using the MVTec-AD dataset [61, which contains 224x224 RGB images of 15 object categories. We follow the multi-class problem setup proposed by (60) The training dataset con- tains normal object images from 15 categories without any labels. The test set consists of both normal and defective object images, each provided with an anomaly label and a mask indicating the defect location. The goal is to de- tect and localize anomalies, with performance measured by AUC computed per object category. This setting is challenging because the energy function should reflect the multi-modal data distribution. Following the preprocess- ing protocol in 60 59]. each image is transformed into a 272x 14x 14 vector using a pre-trained EfficientNet-bal [62]. DxMI is conducted in a 272-dimensional space, treat- ing each spatial coordinate independently. With the trained energy function, we can evaluate the energy value of 14x14 spatial features and use max pooling and bilinear interpolation for anomaly detection and localization, respectively.
Table 5: MVTec-AD multi-class anomaly detection and localization experiment. Anomaly detection (DET) and localization (LOC) performance are measured in AUC. Due to the space constraint, only the average AUC over 15 classes is presented. The full results are provided in Table 6
Model
DET
LOC
DRAEM 58
88.1
87.2
MPDR 59
96.0
96.7
UniAD 60
96.5
96.8
DxMI
97.0 10.11
90.1
67.9+5.90
84.6:14.22
We use separate networks for the energy function and the value function in this experiment, as the primary goal is to obtain an accurate energy function. We employ an autoencoder architecture for the energy function, treating the reconstruction error of sample as its energy (63) The diffusion model and the value function are five-step tine-conditioned MLPs. Unlike conventional diffusion models, DxMI allows for a flexible choice of x We set the initial distribution for the sampler to the data distribution applied with noise, aiming to identify the energy value more precisely near the data distribution. More experimennnt experimental details can be found in Appendix C,3
DxMI demonstrates strong anomaly classification and localization performance, as shown in Table 5] This result indicates that the trained energy function effectively captures the boundary of normal data. When entropy maximization is disabled by r = O, the diffusion model fails to explore and only exploits regions of minimum energy, resulting in poor performance. We observe that moderate level of 7 = 0.1 benefits both the sampler and the energy function, as encourages exploration and provides a suitable level of diversity in negative samples.

---

Related Work
Faster Difflusion Models. Significant effort has been dedicated reducing the number of genera- tion steps in diffusion models during sampling while preserving sample quality. One popular approach to keep the trained diffusion model unchanged and improve the sampling phase independently
by tuning the noise schedule 8 64 65 [9. improving differential equation solvers [50] [6] 67 [8],
and utilizing non-Markovian formulations [45]
69) 70]. While these methods are training-free free, the sample quality can be further improved when the neural network is directly tuned for short-run sampling. Distillation methods train faster diffusion sampler using training signal from longer-run
74). A distilled model usually cannot outperform the teacher model, but adversarial or IRL methods may exceed full-length diffusion models. Hybrid methods [13] [2) combine distillation with adversarial loss, while other methods [46] 75 apply adversarial training to each denoising step. DxMI and SFT-PG [T]] rely fully on adversarial training for final samples, allowing beneficial deviations from the diffusion path and reducing statistical distance from the data.
diffusion model, showing strong performance [T0] 71 72 48
73 43
RL for Diffusion Model. RL is often employed to fine- tume diffusion models for a reward function. The source of the reward signal can be a computer program [24] 76 20 77 78), or a human evaluator [1 (79 25.. DxMI focuses on setting where the estimated log data density is the reward. When RL applied to diffusion models, the policy gradient [80] is the dominant choice 24 76 (7). DxMI offers value function-baset approach an alternative to the policy gradient. Maximum entropy RL for diffusion models is investigated in [20] 81 [2) but only in the continuous- time setting. DxDP investigates the discrete-time setting, which is more suitable for accelerating generation speed.
Energy-Basee Models. DxMI provides method of utilizing a diffusion model eliminate the need for MCMC. Many existing EBM training algorithms rely on MCMC, which is computationally
expensive and difficult to optimize for hyperparameters [34 83 [8
84
85
35
63 59. Joint training
of an EBM with a separate generative model is a widely employed strategy to avoid MCMC. EBMs can be trained jointly with a normalizing flow [86] 87], a generator (30 88 89 or a diffusion model 190 33 However, none of the works use a diffusion model directly as a sampler.
DxMI shares the objective function with several prior works in EBM 27
29
30 31 33
Related Theoretical Analyses. The convergence guarantees of entropy-regularized IRL are provided in [T 992 under the assumption of linear reward and the infinite time horizon. Their guarantees are not directly applicable to practical instance of DxMI, mainly due to the nonlinearity of the reward function, the continuous state and action spaces, and the finite-horizon setting. Establishing the convergence guarantee for DxMI could be an important future research direction. On the other hand, theoretical analyses have been conducted on MaxEnt RL under finite state and action spaces 193, which is relevant for the discrete version of DxDP. More focused analysis on entropy regularized RL for diffusion models provided in [94]
Conclusion
In this paper, we leverage techniques from sequential decision making to tackle challenges in generative modeling, revealing a significant connection between these two fields. We anticipate that this connection will spur variety of algorithmic innovations and find numerous practical applications.
Broader Impacts. DxMI may facilitate deep fakes or fake news. However, trained on relatively low-resolution academic datasets, the models created during our experiments are not capable enough to cause realistic harm. Generative models trained solely using DxMI may possess fairness issues.
Training multiple components simultaneously, DxMI introduces several hyperparam- eters. To reduce the overhead of practitioners, we provide hyperparameter exploration guideline in Appendix B DxMI is not directly applicable to training single-step generator. However, diffusion model fine-tuned by DxMI can be distilled to a single-step generator. DxMI does not offer the flexibility of using different value of generation steps 7 during the test time. Direct theoretical analysis of DxMI is challenging since the models are built on deep neural networks. Theoretical analysis that rationalizes the empirical results will be an important direction for future work.
Limitations.
10

---

Acknowledgments and Disclosure of Funding
S. Yoon is supported by a KIAS Individual Grant (AP095701) via the Center for AI and Natural Sciences at Korea Institute for Advanced Study and ITP//MSIT RRS-2023-00220628). H. Hwang and Park are supported in part by IITP-MSIT grant RS-2021-I122 -12222688 (SNU AI Innovation (Training and Inference Methods for Goal Oriented AI Agents), MiitrMinistry of Science ICT). Korea, under the Global Research Support Program in the Digital Field program (ataaammmSS202--0-00 (04-60043))8) supervised by the ILTP(Institute for Information & Communications Technology Planning Evaluation), KIAT grant P0020536 (HRD Program for Industrial Innovation), SRRC NRF grant RS-2023-0020800 SNU-AIIS, SNU-IPAI, SNU-IAMD, SNU BK21+ Program in Mechanical Engineering, SNU Institute for Engineering Research, Microsoft Research Asia, and SNU Interdisciplinary Program in Artificial Intelligence. D. Kwon is partially supported by the National Research Foundation of Korea (NRF) grant funded and No. RS-2024-00408003) and the POSCO Science Fellowship of POSCO TJ Park Foundation. Y.-K. Noh was partly supported by NRE/MSITT 20..R-2024-0042111 (203)) and IITP/MSIT (2020-0-01373).
Hub), HUTP-MSIT grant 2022-220480, RS-
by the Korea government (MSIT) (No.
References
[1] Chelsea Finn, Paul Christiano. Pieter Abbeel, and Sergey Levine. connection between generative adversar- ial networks, inverse reinforcement learning. and energy- -based models. arXiv preprint arXiv:1611.0385. 2016.
[2] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
[3] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, Ranzato, Hadsell, M.F. Balcan, and H. Lin, editors, Advances Neural Information Processing Systems, volume 33, pages 6840 6851. Curran Associates, Inc., 2020.
[4] Jascha Sohl- boch--Dcckttei, Eric Weiss, Niru Maheswaranathan and Surya Ganguli. Deep unsupervised learning using uonequilibrium thermodynamics. Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Leurning Research, pages 2256-2265, Lille, France, 07-09 Jul 2015. PMLR.
[5] Dean A. Pomerleau. Alvinn: An autonomous land vehicle in a neural network. In D. Touretzky, editor, Advances in Neural Information Processing Systems, volume 1. Morgan- Kaufmaan,, 1988.
[6] Stephane Ross, Geoffrey Gordon, and Drew Bagnell. reduction imitation learning and structured prediction to no-regret online learning. In Geoffrey Gordon, David Dunson, and Miroslav Dudík, editors, Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, volume 15 of Proceedings of Machine Learning Research, pages 627-635, Fort Lauderdale. FL, USA, 11-13 Apr 2011. PMLR.
[7] Siddharth Reddy, Anca D. Dragan, and Sergey Levine. (SQIL): Imitation learning via reinforcement learning with sparse rewards. In International Conference on Learning Representations, 2020.
[8] Zhifeng Kong and Wei Ping. On fast sampling of diffusion probabilistic models. urXiv:2106.00133 00132, 2021.
arXiv preprint
[9] Robin San-Roman, Eliya Nachmani, and Lior Wolf. Noise estimation for generative diffusion models. arXiv preprint Xrxiv (04.02600, 2021.
[10] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations, 2022.
[11] Ying Fan and Kangwook Lee. Optimizing DDPM sampling with shortcut fine-tuning, In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardr, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 9623 9639, PMLR, 23 29 Jul 2023.
[12] Yanwu Xu, Yang Zhao, Zhisheng Xiao. and Tingbo Hou. Ufogen: You forward once large scale text-to- image generation via diffuusion gans. arXiv preprint arXiv:2311.09233 2023.
[13] Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion distillation. arXiv preprint 21X1422311.1704,, 2023.
11

---

[14] Andrew Y Ng and Stuart J Russell. Algorithms for inverse reinforcement learning. In Proceedings of the Seventeemth International Conference on Machine Learning, pages 663-670, 2000.
[15] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, al. Maximum entropy inverse reinforcement learning. In AAAI, volume 8, pages 1433-1438. Chicago, II., USA, 2008.
[16] Chelsea Finn, Sergey Levine, and Picter Abbeel. Guided cost learning: Deep inverse optimal control via policy optimization. In Maria Florina Balcan and Kilian Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Leurning Research, pages 49-58, New York, New York, USA, 20-22 Jun 2016. PMLR
[17] Tuomas Haamojj, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1861-1870. PMM... PMLR, 10-15 Jul 2018.
[18] Yilun Du, Shuang Li, B. Joshua Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of energy based models. In Proceedings the 38th International Conference on Machine Learning (ICML-21), 2021.
[19] Kevin Clark, Paul Vicol, Kevin Swersky, and David Fleet. Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309. 2023.
[20] Masatoshi Uchara, Yulai Zhao, Kevin Black, Ehsan Hajiramczanali, Gabriele Scalia, Nathanicl Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine-tuning of continuoustimme diffusion models entropy-regularien control. arXiv preprint 24X1v:2402.15194, 2024.
[21] Julius Bemnr,, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion-based generative modeling. arXiv preprint arXiv:2211.01. 1.01364,, 2022.
[22] Qinsheng Zhang and Yongxin Chen. Path integral sampler: A stochastic control approach for sampling. In International Conference on Learning Representations, 2022.
[23] Yang Song, Jascha Sohl- Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[24] Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Picter Abbeel, Moham- mad Ghavamzadeh, Kangwook Lee, and Kimin Lec. Reinforcement learning for fine-tuning text-to-image diffusion models. In seirth- Conference on Neural Information Processing Systems, 2023.
[25] Bram Wallace, Meihua Dang, Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages $228-8238, June 2024.
[26] Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 1478) 2002.
[27] M. Ehsan Abbasnejad, Qinfeng Shi, Anton van den Hengel, and Lingqiao Liu. A generative adversarial density estimator. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), June 2019.
[28] Bo Dai, Zhen Liu, Hanjun Dai, Niao He, Arthur Gretton, Le Song, and Dalc Schuurmans. Exponential family estimation via adversarial dynamics embedding. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buse Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[29] Zihang Dai, Amjad Almahairi, Philip Bachman, Eduard Hovy, and Aaron Courville. Calibrating energy- based generative adversarial networks. In International Conference on Learning Representations, 2017.
[30] Will Sussman Grathwohl, Jacob Jin Kelly, Milad Hashemi, Mohammad Norouzi, Kevin Swersky, and David Duvenaud. No (meme) for me: Amortized sampling for fast and stable training of energy-based models. In International Conference on Learning Representations, 2021.
[31] Rithesh Kumar, Sherjil Ozair, Anirudh Goyal, Aaron Courville, and Yoshua Bengio. Maximum entropy generators for energy- based models. arXiv preprint 191.00190..08508 2019.
12

---

[32] Cong Geng, Jia Wang, Zhiyong Gao. Jes Frellsen, and So ren Hauberg. Bounds all around: training energy-based models with bidirectional bounds. In M. Ranzato, A. Beygelzimer. Y. Dauphin, P.S. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 19808-1999 19821. Curran Associates, Inc., 2021.
[33] Cong Geng, Tian Han, Peng-Tas Tao Jiang, Hao Zhang, Jinwei Chen. Søren Hauberg, and Bo Improving adversarial energy-based model via diffusion process. arXiv preprint arXiv:2403. 2024.
[34] Yilun Du and Igor Mordatch, Implicit generation and modeling with energy based models. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d Alche-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 3608-3618. Curran Associates, Inc., 2019.
[35] Hankook Lee, Jongheon Jeong, Sejun Park, and Jinwoo Shin. Guiding energy-basedd models via contrastive latent variables. In International Conferencee on Learning Representations, 2023.
[36] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations, In S. Bengio. H. Wallach, H. Larochelle, K. Grauman, N. Cesa Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[37] Lyudmyla F Kozachenko and Nikolai N Leonenko. Sample estimate of the entropy of random vector. Problemy Peredacki Informatsii, 23(2)99-1 16, 1987.
[38] Mohamed Ishmael Belghazi. Aristide Baratin, Sai Rajeshwar, Sherjil Ozair. Yoshua Bengio, Aaron Courville, and Devon Hjelm. Mutual information neural estimation. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 531-540 15400 PMLR, 10-15 Jul 2018.
[39] XuanLong Nguyen, Martin J Wainwright, and Michael Jordan. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information Theory, 56(11):5847-586 2010.
[40] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q. Weinberger, editors, Advances Neural Information Processing Systems, vol- ume 27. Curran Associates, Inc., 2014.
[41] Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Yee Whye Teh and Mike Titterington, editors, Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume 9 of Proceedings of Machine Learning Research, pages 297-304, Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010. PMLR.
[42] Yang Song. Prafulla Dhariwal. Mark Chen, and Ilya Sutskever. Consistency models. In Andreas Krause, Emma Brunskill. Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors. Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 3221 1-32252. PMLR, 23-29 Jul 2023.
[43] Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023.
[44] Axel Sauer, Katja Schwarz, and Andreas Geiger. Stylegan-xl: Scaling stylegan to large diverse datasets. In ACM SIGGRAPH 2022 conference proceedings. pages 1-10, 2022.
[45] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
[46] Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion GANs. In International Conference Learning Representations, 2022.
[47] Prafulla Dhariwal and Alexander Nichol. Diffusson models beat gans on image synthesis. In M. Ranzato,
A. Beygelzimer, Y. Dauphin. P.S. Liang, and J. Wortman Vaughan, editors, Advances Neural Information Processing Systems, volume 34, pages 8780-8794. Curran Associates. Inc., 2021.
[48] Hongkai Zheng. Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. Andreas Krause, Emma Brunskill, Kyunghyun Cho. Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 42390 42402. PMLR, 23-29 Jul 2023.
13

---

[49] Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale GAN training for high fidelity natural image synthesis. In International Conference Learning Representations, 2019.
[50] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 26565 26577. Curran Associates, Inc., 2022.
[51] Tero Karras, Samuli Laine, Miik Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila, Analyzing and improving the image quality of stylegan. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
[52] A Krizhevsky. Learning multiple layers of features from tiny images. Master's thesis, University of Tronto, 2009.
[53] Deng. Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei. Imagenet: large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-28188 255. leee, 2009.
[54] Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song. Thomas Funkhouser, and Jianxiong Xiao. Lsun: Con- struction of large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv: 150.60366606 2015.
[55] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to local nash equilibrium. In I. Guyon, U. Von Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates. Inc., 2017.
editors,
[56] Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[57] Dongjun Kim, Yeongmin Kim, Se Jung Kwon, Wanmo Kang, and Il-Chul Moon. Refining generative process with discriminator guidance in score-based diffusion models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 16567-16595 PMLR, 23-29 Jul 2023.
[58] Vitjan Zavrtanik, Matej Kristan, and Danijel Skočaj. Draem a discriminatively trained reconstruction embedding for surface anomaly detection. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages
October 2021.
[59] Sangwoong Yoon, Young Uk Jin, Yung- Kyun Noh, Frank C. Park. Energy-based based models for anomaly detection: A manifold diffusion recovery approach. In Teirty- Conference on Neural Information Processing Systems, 2023.
[60] Zhiyuan You, Lei Cui, Yujun Shen, Kai Yang, Xin Lu, Yu Zheng, and Xinyi Le. A unified model for multi- class anomaly detection. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Syxtems, volume 35, pages 4571-4584. Curran Associates. Inc., 2022.
[61] Paul Bergmann, Kilian Batzner, Michael Fauser, David Sattlegger, and Carsten Steger. The mytec anomaly detection dataset: comprehensive real-word dataset for unsupervised anomaly detection. International Journal of Computer Vision, 129(4): 1038-1099 2021.
[62] Mingxing Tan and Quoc L.e. Efficientnet: Rethinking model scaling for convolutional neural networks. In International conference on machine learning, pages 6105-61 14. PMIRR 2019.
[63] Sangwoong Yoon, Yung- Kyun Noh, and Frank Park. Autoencoding under normalization constraints. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 12087-12097. PMLR, 18-24 Jul 2021.
[64] Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang. Maalytic-DPMa an analytic estimate of the optimal reverse variance in diffusion probabilistic models. In International Conference Learning Representations, 2022.
14

---

[65] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusson probabilistic models. In Marina Meila and Tong Zhang, editors. Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8162-8171. PMLR, 18-24 Jul 2021.
[66]
Alexia Marttnne,, Ke Li, Rémi Fiché-Taillefie Tal Kachman, and loannis Midiagkas. Gotta go fast when generating data with bose- models. arXiv preprint arXiv:2105.14080 2021.
[67] 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. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 5775-5787, Curran Associates, Inc., 2022.
[68] Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. In The Eleventh International Conference on Learning Representations, 2023.
[69] Qinsheng Zhang, Molei Tao, and Yongxin Chen. gDDIM: Generalized denoising diffusion implicit models. In The Eleventh International Conference on Learning Representations, 2023.
[70] Daniel Watson, William Chan, Jonathan Ho, and Mohammad Norouzi. Learning fast samplers for diffusion models by differentiaiing through sample quality. In International Conference on Learning Representations, 2022.
[71] David Berthelot, Arnaud Autef, Jierui Lin, Dian Ang Yap. Shuangfei Zhai, Siyuan Hu, Daniel Zheng, Walter Talbott, and Eric Gu. Tract: Denoising diffusion models with transitive closure time-distillation. arXiv preprint arXiv:2303.044 042448 2023.
[72] Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed, 2021.
[73] Wujie Sun, Defang Chen, Can Wang, Deshi Ye, Yan Feng, and Chun Chen. Accelcrating diffusion sampling with classifier- feature distillation. In 2023 IEEE International Conference on Multimedia and Expo (ICME), pages 810-815, 2023.
[74] Michael Samuel Albergo and Eric Vanden-Fijndes Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023.
[75] yanwu xu, Mingming Gong, Shaoan Xic, Wei Wei, Matthias Grundmann, Kayhan Batmanghelich, and Tingho Hou, Semi-implicit denoising diffusion models (siddms). In A. Oh, T. Naumann, A. Globerson, K. Sacnko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 17394, Curran Associates, Inc., 2023.
[76] Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikow, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305. 13301, 2023.
[77] Scung Hyun Lee, Yinxiao Li, Junjie Ke, Innfarn Yoo, Han Zhang, Jiahui Yu, Qifei Wang, Fei Deng, Glenn Entis, Junfeng He, et al. Parrot: Pareto-optimal multi-reward reinforcement lcarning framework for text-to-image generation. arXiv preprint arXiv:2401. 005675550055....11 2024.
[78] Jianshu Guo, Wenhao Chai, Jie Deng, Hsiang-Wae 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.
[79] Shu Zhang, Xinyi Yang, Yihao Feng, Can Qin, Chia-Chih Chen, Ning Yu, Zeyuan Chen, Huan Wang, Silvio Savarese, Stefano Ermon, Caiming Xiong, and Ran Xu, Hive: Hareessing human feodback for instructional visual editing, arXiv preprint arXiv:2303.099: 096118 2023.
[80] Ronald J Williams. Simple statistical sradient-following algorithms for connectionist reinforcement learning. Reinforcement learning, pages 5-32, 1992.
[81] Masatoshi Uchara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Sergey Levine, and Tommaso Biancalani. Feedback efficient online fine-tuning of diffusion models. arXiv preprint arXiv:2402. 16359, 2024.
[82] Masatoshi Uehara, Yulai Zhao, Ehsan Hajiramezanali, Gabriele Scalia, Gökeen Eraslan, Avantika Lal, Sergey Levine, and Tommaso Biancalani. Bridging model- based optimization and generative modeling via conservative tine-tuning of diffusion models. arXiv preprint 2rXiv:2405.19673 2024.
15

---

[83] Erik Nijkamp, Mitch Hill. Song-Chur Cong-Chhn Zhu, and Ying Nian Wu. Learning non-convergent non- -on-pessittent short-run meme toward energy-based model. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Bor Buc, E. Fox, and Garnell, editors, Advances in Neural Information Processing Systems, volume 32, pages Curran Associates, Inc., 2019.
[84] Zhisheng Xiao, Karsten Kreis, Jan Kautz, and Arash Vahdat. [VAEBM): A symbiosis between variational autoencoders and energy-based based models. In International Conference on Learning Representations, 2021.
[85] Ruiqi Gao, Yang Song. Ben Poole, Ying Nian Wu, and Diederik P Kingma. Learning energy- based models by diffusion recovery likelihood. In International Conference on Learning Representations, 2021.
[86] Jianwen Xie, Yaxuan Zhu, Jun Li, Ping Li. A tale two flows: Cooperative learning of langevin flow and normalizing flow toward energy-based model. In International Conference on Learning Representa- tions, 2022.
[87] Ruiqi
Gao, Erik Nijkamp. Diederik Kingma, Zhen Xu, Andrew M Dai, and Ying Nian Wu. Flow contrastive estimation of -based models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7518-7528 2020.
[88] Will Grathwohl, Kuan- Chich Wang, Joern- Henrik Jacobsen, David Duvenaud, and Richard Zemel. Learn- ing the stein discrepancy for training and evaluating energy- based models without sampling. In Hal Daumé III and Aarti Singh. editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 Proceedings of Machine Learning Research, pages 3732-3747. PMLR, 13-18 Jul 2020.
[89] Tian Han, Brik Nijkamp, Xiaolin Fang, Mitch Hill, Song- Chun Zhu, and Ying Nian Wu. Divergence triangle for joint training of generator model, energy- based model, and inferential model. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2019.
[90] Peiyu Yu, Yaxuan Zhu, Sirui Xie, Xiaojian (Shawn) Ma, Ruiqi Gao, Song-Chun Zhu, and Ying Nian Wu. Learning energy-hased prior model with diffusion-amortizzd meme. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 42717-42747. Curran Associates, Inc., 2023.
[91] Siliang Zeng, Chenliang Li, Alfredo Garcia, and Mingyi Hong. Maximum-likelihonee inverse reinforcement learning with finite-time guarantees. Advances in Neural Information Processing Systems, 15120122- 2022.
[92] Titouan Renard, Andreas Schlaginhaufen, Tingting Ni, and Maryam Kamgarpour. Convergence of model- free entropy-regularized inverse reinforcement learning algorithm. arXiv preprint arXiv:2403.1633 2024.
[93] Matthieu Geist, Bruno Scherrer, and Olivier Pietquin. theory of regularized Markov decision processes. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Muchine Learning Research, pages 2160-2169, PMLR, 09-15 Jun 2019.
[94] Wenpin Tang. Fine-tuning of diffission models via stochastic control: entropy regularization and beyond. arXiv preprint arXiv:2403.0622 062799 2024.
16

---

A Adaptive Velocity Regularization
We are interested in the following optimization problem:
min
K (L(T go (xT (x0
(13)
Plugging in our choice of log (x, (x+1) = 1xx, Ixx-Xxulll" D log Se, we can rewrite the optimization
207
problem as follows.
1|x,
Xe+1 1|
D log st
min
Ex X
(14)
2s
t=
where constant term with respect to s is omitted. Since the object function is separable, we can solve the optimization for each 8t independently.
min Ex.,x_ ~u
[|xx
x
Dog St,
== t., T
(15)
2sF
This optimization has an analytic solution: (s²)² E Ex,xx+ [[[x; Xt+1|13]/D.
B Guideline for Hyperparameters
The most crucial design decision in DxMI is how to construct the value function. This design revolves around two primary axes. The first axis is whether the value function should be time-dependent or iime-independent. The second axis is whether the value function should share model parameters with other networks, such the sampler the energy network.
Value function.
Our experiments demonstrate various combinations of these design choices. 2D experiment, the time-dependent value function shares parameters with the EBM, which is a recommended approach for smaller problems. In an image experiment, we employ time-independent value function that shares parameters with the energy network, effectively making the value function and the energy function identical. This design choice promotes monotonicity and efficient sample usage, as a single value function learns from all intermediate samples.
In an anomaly detection experiment, we use time-dependent value function that does not share parameters with any other network. This design is suitable when specific structure needs to be enforced on the energy function, such as with an autoencoder, and when making the structure time-dependent is not straightforward.
While these are the options we have explored, there are likely other viable possibilities for designing the value function.
Coefficient т. Although the coefficient 7 plays important role in DxMI, we recommend rumning DxMI with 5 = 0 when implementing the algorithm for the first time. everything in order, DxMI should function to
extent. During normal training progression, the energy values of positive and negative samples should converge as iterations proceed.
After confirming that the training is progressing smoothly, you can start experimenting with increasing the value of 7. Since y determines the absolute scale of the energy, the magnitude of r should be adjusted accordingly. We set y = 1 in all our experiments and recommend this value. In such a case, the optimal 7 likely to be less than 0.5.
Learning rate. As done in two time-scale optimization 155 we use a larger learning rate for the value function than for the sampler. When training the noise parameters de in the diffusion model, we also assign learning rate 100 times larger than that of the sampler.
C Details on Implementations and Additional Results
C.1 2D Experiment
Sample quality quantified using sliced Wasserstein-2 distance (SW) with 1,000 projections of 10k samples. The standard deviation is computed from 5 independent samplings. Density estimation
17

---

Algorithm 2 Diffusion by Maximum Entropy IRL for Image Generation
1: Input: Dataset D, Energy Ex(x), Value Vx(xx) and Sampler (x(x)))
2: 8t <-0;
AVR initialization
3: for in D do
// Minibatch dimension is omitted for brevity.
4: Sample X0:T ~ (x).
5: ming Ex(x) - E (xT) +y Ex(x))?
I/ Energy update // Value update
6: fort = T- ,,...,,000 do
7:
ming [sgnn end for
+ (x([](())(( (xx+1)] R(1) - (x:))
8:
9:
for x, randomly chosen among X
Sampler update I/ Reparametrization trick
10:
Sample one-step: X(+1 T (x |x,) ming V V (x++1(0)) T2 Tr
11:
+ 2 X+ xx11
|| X++1 is function of ø.
end for
13: of f- ass + (1 a)l|x; xx++111/DD
I/ AVR update
14:
end for
performance is measured by AUC on discriminating test data and uniform samples over the domain. AUC is also computed with 10k samples.
v(x,==0) 0)
V(xt=1 1)
vixt= 2)
V(xt= 3)
v(xt=4)
E(x) (= Vxx,=
Figure 3: Value functions at each time step (r = 0.1 case). Blue indicates a low value.
C.2 Image Generation
Datasets. CIFAR-I0 retrieved through torchvision API. ImageNet downloaded from Kaggle and downsampled to 64 64 following https: //gi thub. com/ openal/ http///gnaaaaddd.....pp diffusion We apply random horizontal flips on all images. When computing FID, the whole 50,000 training images of CIFAR-10 10 are used. We make sure that no JPEG compression is used during processing. For Ima- geNet, we use the batch stat file provided by https::/p//. //gi cww htmm/opeaai///e httpg/uudddddfffur http:/abw..ccopppppaaaauaddddddd
Models. For DDPM on CIFAR-IO, we use the checkpoint provided by the SFT-PG repository [1. For DDGAN, we utilize the official checkpoint. For EDM ImageNet 64, we use the checkpoint from the consistency model repository [42]. In all experiments, we employ the same ResNet architecture, which does not include any normalization layers, such as batch normalization.
Training. For all runs, we use batch size of 128. In the CIFAR-10 experiments, we use the Adam optimizer with learning rate of 10-7 for the sampler weights, 10-5 for the value weights, and 10-5 for the s. In the ImageNet 64 experiments, we use RAdam with learning rate of 10 for the sampler. Additionally, we utilize a mixed precision trainer handle FP16 weights, The value weights are updated with a learning rate of 10 using Adam. The a, 's are updated with a learning rate of 10 To select the best model, we periodically generate 10,000 images for CIFAR- 10 and 5,000 images for ImageNet. The checkpoint with the best FII) score is selected the final model.
Evaluation For computing FID, Precision, and Recall scores, we used the TensorFlow-bbsed evaluation script provided by Consistency Models 42 repository, which is based on the codebase of 47]. All the images are saved in a PNG format when fed to the evaluation script.
Additional Details. For optimal performance in image generation, we often tune the coefficients of
two running costs separately. Let denote the coefficient of log (x++1 xx) 71 and the coefficient
of 25 |x
x as 72- In the CIFAR-10 experiments with = 10 and = 4, we set 71 0.1
18

---

and = 0.01. In the ImageNet experiments with T = 10, we set 72 = 0.01, and for T = 4. we set 71 0.1 and T2 0.01.
We believe the optimal т values vary for each setting due to differences in noise schedules and magnitudes. For example, the DDPM backbone is a variance-preseevvngg formulation, while EDM is variance-cxploding formulation. Exploring a unified method for selecting the entropy regularization parameter is an interesting research direction.
C.3 Anomaly Detection
Dataset and Feature Extraction. MVTec AD is a dataset designed evaluate anomaly detection techniques, particularly for industrial inspection. contains over 5000 high-resolution images divided into 15 different categories of objects and textures. Each category includes of defect-free training images and test set with both defective and non-defective images.
The original dataset consists of high-resolution RGB images sized x24x224. Following the methods used similar studies (59] 66), we extract a 272 14x full feature representation each image using& pre-trained EfficientNet-b4 [62). We then train the energy function using DxMI in 272- dimensional space, treating every spatial feature from the training images as normal training data. Each data point x in R_722 projected to 5272 through normalization. This projection is effective because the direction of the feature often represents the original image better than its magnitude.
Anomaly Detection and Localization. For anomaly detection, the energy of single image is calculated by applying max pooling to the energy value of each spatial feature. For anomaly localization, the energy values of the 14x 14 spatial features are upsampled to match the original 224 224 image resolution using bilinear interpolation.
Model Design. We utilize an autoencoder architecture for the energy function, as described in 63). using the reconstruction error of sample as its energy value. Considering that the data distribution lies on g272, we appropriately narrow the function classes for the energy and sampler. Specifically, we constrain the decoder manifold of the energy function and the mean prediction of the sampler to remain on 5 We pretrain the energy function (i.e., the autoencoder) to minimize the reconstruction error of the training data. Both the sampler and the value function are trained from scratch.
Choice of (xx).. Unlike traditional diffusion models, DxMI permits a flexible choice of (xx)). To train the energy function effectively near the data distribution, we set the initial distribution for the sampler as the data distribution corrupted with noise. To apply meaningful noise the data in 8272, we use the pretrained autoencoder that is also used for the initial energy function to project the samples from the data distribution the latent space, apply perturbations, and then restore the data to produce initial samples, as suggested in [59]. To maintain a consistent initial distribution, we fix the autoencoder used for generating the initial samples.
Additional Details. We use autoencoder with latent dimension 128 the energy function. Encoder and decoder each consist of an MLP with 3 hidden layers and 1024 hidden dimensions. We use conditiniiio MLP for the sampler and value function, encoding the time information into a 1 228-dimensional vector using sinusoidal positional encoding. The input x, is concatenated with the time embedding vector. We use MLPs with 3 hidden layers of 2048 and 1024 hidden dimensions for the sampler and value function, respectively.
The model is trained for 100 epochs with a batch size of 784 (=4x 14x 14) using the Adam optimizer. We use a learning rate of 10-6 for the sampler and value function and 10-4 for the energy function.
D Implementation and Computational Complexity of DxMI
DxMI (Algorithm[1 I may seem complicated, but it largely mirrors the procedure of Max Ent IRL, with the exception of the AVR update. Notably, DxDP, the Max Ent RL subroutine within DxMI, is significantly simpler than standard actor-critic RL algorithms. Unlike these algorithms such as Soft Actor Critic (SAC) 17]. which requires training both value function and two Q-functions-Dxil only trains single value function and no Q-functions.
19

---

Table MVTec-AD detection and localization task in the unified setting. AUROC scores (percent) are computed for each class. UniAD and DRAEM results are adopted from 60 The largest value in a task marked as boldface.
Detection UniAD
Localization UniAD
DxMI
MPDR
DRAEM
DxMI
MPDR DRAEM
Bottle
100.0 +0.00
99.7+0.04
100.0
97.5
98.5
18.140004
98.5
87.6
Cable
97.1 0.37
95.2+0.4
95.5
57.8
96.6+..10
97.3++0.10
95.6
71.3
Capsule
89.8= 0.61
86,920.73
86.4
65.3
98.5+0.03
98.5=t0001
98.2
50.5
Hazelnut
100.0 0.04
99.81...1
99.9
93.7
98.41.0044
98.1100.10
98.4
96.9
Metal Nut
99.9 0.11
99.2:10.99
99.9
72.8
95.5
94.8:10.09
94.5
62.2
Pill
95.4
93.7.1.65
94.0
82.2
95.6110.7
95.0.10.16
94.9
94.4
Screw
88.9 0.51
87.5+
98.6+0.08
98.3-+0.08
98.1
95.5
Toothbrush
92.2:+1.66
94.2
89.6
90.6
98.8
98.4
98.7
97.7
Transistor
99.2:10.28
99,8+u009
98.3
74.8
96.0#0.13
97.9
95.4
65.5
Zipper
96.3:10.0
95.8+
95.3
98.8
96.7
96.8+10.4
96.2
98.3
Carpet
99.9=0.04
99.8+0.02
99.9
98.0
98.8+0.02
98.5++001
98.8
98.6
Grid
98. 98.611028
98 0..2
97.9
99.3
97.0+1.07
96.5:10.44
96.9
98.7
Leather
100.0.10.00
100.0.10.00
100.0
98.7
98.5100.03
98.810..0
98.5
97.3
Tile
100.010..0
99.31..14
100.0
95.2
95.2 10.14 91.8100.10
94.6
98.0
Wood
98.3
98.6110.88
97.9
99.8
93.810.07
93.2.10.08
93.8
96.0
Mean
97.0=t.11
96.5+0008
96.0
88.1
97.1
96.8+0.02
96.7
87.2
DxMI does not demand excessive computation, a major concern in image generation experiments, In these experiments, the only additional component beyond the diffusion model is the EBM, which shares the same network as value function. Additionally, the EBM used in DxMI is typically much smaller than the diffusion model, imposing minimal computational overhead. For instance, in our CIFAR- 10 experiment (T=10), the EBM consists of 5M parameters, compared to 36M the diffusion model. This is also significantly smaller than the critic networks used in GANs, such the 158.3M parameters BigGAN [49] In practice, our CIFAR-10 experiment completes in under 24 hours two A100 GPUs, while the ImageNet 64 experiment takes approximately 48 hours four A100 GPUs.
E
Interpretation of the policy improvement method
In this section, we show that our policy improvement method introduced in Eq. II can effectively minimize the KI. divergence between the joint distributions, K ( (Xo:T) (lgo (xx:r-1/xx)
The policy improvement at each timestep f can be expressed as
min Exx,x -(xc+1(x.)
[ [V++((x++) + T log x (x++1 (x:) - log q( [x++
const.
(16)
Here omitted the parameters & and v to avoid confusion. Using the definition of value function. above minimization can be expressed as
E log x (o0tt))
min Ex(xer)) x(X++1|80 (xc)
E(xT) +
log n (x/41 (xx)
(17)
where
T-1
Ex(Xer)
E(xT) +
Σ
log x. (xe+1 (xx+1|xx) - Σ
log q(xe (x0441)
(18)
d'=l
=TBa(xer)[ (- log q(x)) log ((xxx-((xxx)))) + log т (xxT)) log Z log T (x:)]
(19)
log +x(,)
(20)
Note that x, is fixed in this optimization problem, and the optimization variable is xx+1. Therefore the policy improvement step at time t equivalent
min
(21)
(x++1
Therefore, the policy improvement step at each time step results in the minimization of the KI. divergence between the joint distribution, x (Xx.T) and q(Xo.T).
20

---

Additional Image Generation Results
Additional samples from the image generation models are presented in Fig. and Fig.
CIFARIO Real
SFT-PG (T-10)
(T-10)
Figure 4: Randomly selected samples from CIFAR-10 training data, SFT-PG (T 10, FID: 4.32),
and DxMI (T' = 10, FID: 3.19).
ImageNet
Consistency Model (T=1)
DxMI (T=4)
DxMI (T=10)
Figure 5: Randomly selected samples from ImageNet 64x64 training data, Consistency Model
(I' = 1, FID: 6.20), DxMI (T = 4. FID: 3.21), and DxMI (T = 10, FID; 2.68). Note that the Consistency Model samples distort human faces, while the DxMI samples depict them in correct proportions.
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our contributions are the maximum entropy IRL formulation for diffusion models and learning algorithm inspired by dynamic programming. These contributions are explicitly stated in the abstract and the introduction.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We have dedicated a paragraph the discussion of limitations in Section 7
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: No formal theorem proposition is provided in the main manuscript.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We have put best effort into making the proposed method reproducible. The concrete description of the proposed algorithm is provided Algorithm] [ The detailed information on the model and implementation are described in Section 5 and Appendix
C
We release our source code and model checkpoints at whttps:/ //gi:://tttttt cwm/ccmm////// swyoon/ Our public codebase will include instructions for preparing data and evaluating performance metrics.
Di by
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide
code
httpp://github..com/s/// com/swooon///
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
C C describe experimental details including hyperparam- eters, model architecture, data preparation, and the choice of optimizers. We also provide guideline for tuning hyperparameters in Appendix B
Justification: Section and Appendix
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report error bars for 2D experiments where repetitive experiments can be readily conducted. We did not do statistical significance testing.
22

---

8. Experiments Compute Resourees
Question: For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: We provide information compute resource in Section and Appendix
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http:: http://wwrr..ccoooooooo//ccccaaa
Answer: [Yes]
Justification: We have reviewed and followed the NeurlPS Code of Ethics.
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We have included paragraph discussing the potential negative impact of our work in Section 7
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Although we have image generation models, they are not equipped with safeguards. Our models trained relatively low-resolution images, such CIFAR-IO and ImageNet 64 x64, which reduces the likelihood of causing harm in the real world.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We have cited all datasets and bascline models used in our experiment. We also provide the license information in Appendix C
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our work does provide new assets.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our work does not involve any crowd sourcing experiments with human subjects.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
23

---

Answer: [NA]
Justification: Our work does not involve experiments with human subjects.
24


Reviews:
[{"review_id": "zkKoOPC3uN", "reviewer": "Reviewer_EhWr", "summary": "The authors introduce a maximum entropy inverse reinforcement learning (IRL) approach for\nenhancing the sample quality of diffusion generative models, especially with limited generation\ntime steps. Named Diffusion by Maximum Entropy IRL (DxMI), the approach involves joint\ntraining of a diffusion model and an energy-based model (EBM). The EBM provides the\nestimated log density as a reward signal for the diffusion model, which is trained to maximize\nboth the reward from EBM and the entropy of generated samples. Additionally, the authors\npropose Diffusion by Dynamic Programming (DxDP), a novel reinforcement learning algorithm\nthat optimizes diffusion model updates efficiently by transforming the problem into an optimal\ncontrol formulation. Empirical studies demonstrate that diffusion models fine-tuned with DxMI\ncan generate high-quality samples in as few as 4 to 10 steps and improve the stability of EBM\ntraining dynamics, enhancing anomaly detection performance.", "questions": "1. Generalization to Complex Tasks: How do the authors envision the performance of DxMI in\nmore complex generative tasks, such as high-resolution image generation or text-to-image\nsynthesis?\n2. Training Time Comparison: What is the average training time for models using DxMI\ncompared to other generative model training techniques such as GANs or VAE-based\napproaches?\n3. Scalability: How scalable is the proposed DxMI approach when applied to larger datasets or\nmodels with significantly more parameters? Are there any anticipated bottlenecks?", "limitations": "The authors have clearly presented the limitations in the paper.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Innovation: The proposed DxMI methodology introduces the concept of maximum entropy IRL to the training of diffusion models, which is novel and potentially impactful in improving sample quality and inference time of diffusion model.\n2. Clarity: The manuscript is well-written and logically structured, with clear explanations of the theoretical foundations and algorithms.", "weaknesses": "1. Complexity of Implementation: The implementation of DxMI, particularly the joint training\nof diffusion models and EBMs, might be complex and require significant computational\nresources. The practical feasibility in various settings could be further elaborated.\n2. Limited Scope of Experiments: The experiments, while promising, are somewhat limited in\nscope. More diverse and complex tasks could further validate the robustness and versatility\nof the proposed approach.\n3. Comparative Analysis: While the proposed methods show improvements, a more detailed\ncomparative analysis with state-of-the-art techniques, including training time and\ncomputational cost comparisons, would strengthen the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer EhWr,\n\nWe appreciate the comprehensive feedback on our manuscript. All the comments and questions raised have been considered below. \n\n\n**Q. Complexity of Implementation**\n> The implementation of DxMI, particularly the joint training of diffusion models and EBMs, might be complex and require significant computational resources. \n\n* DxMI may seem complicated, but in fact its complexity is no larger than GANs and actor-critic RL methods.\n    * Particularly in image experiments, EBM is the only additional component over a diffusion model. EBM functions similarly to the discriminator of GAN. Furthermore, EBM used in DxMI is typically much smaller than the diffusion model, introducing minimal burden. For example, in our CIFAR-10 experiment (T=10), the EBM has 5M parameters while the diffusion model has 36M parameters. \n    * DxMI is significantly simpler than standard actor-critic RL, such as Soft Actor-Critic (SAC) [1], which trains a value function and two Q functions simultaneously. On the contrary, DxMI does not require a Q function and trains only a single value function.\n\n[1] Haarnoja, Tuomas, et al. \"Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.\" ICML 2018. https://proceedings.mlr.press/v80/haarnoja18b\n\n\n**Q. More Experiments**\n\n>The practical feasibility in various settings could be further elaborated.\n\n\n> The experiments, while promising, are somewhat limited in scope. More diverse and complex tasks could further validate the robustness and versatility of the proposed approach.\n\n> How do the authors envision the performance of DxMI in more complex generative tasks, such as high-resolution image generation or text-to-image synthesis?\n\n* The manuscript already includes experiments on a variety of tasks and data types, such as 2D density estimation, unconditional and conditional image generation, and anomaly detection using latent vectors. If you find any particular aspect of the experimental portfolio to be limited, please let us know.\n* To further demonstrate the scalability of DxMI, we will provide additional experimental results on LSUN Bedroom 256x256. DxMI (T=4) achieves competitive results (FID: 5.93, Recall: 0.477), whereas DDPM (T=1000) achieves FID: 4.89, Recall 0.45. We will augment our experiment section with the LSUN Bedroom experiment. This observation shows that DxMI is indeed scalable to high-dimensional data.\n* We believe that DxMI can be effectively applied to text-to-image synthesis. However, text-to-image synthesis typically requires a significant amount of computing resources (at least 25 A100 days) and involves numerous experimental conditions to explore (e.g., selecting text prompts). Therefore, we suggest this as an intriguing direction for future research.\n \n \n**Q. Comparative Analysis**\n> While the proposed methods show improvements, a more detailed comparative analysis with state-of-the-art techniques.\n\n* Please find that the current manuscript provides comparative analysis on Section 6 Related Work.\n* We will augment our comparative analysis by adding additional paragraphs in the introduction and sections that describe the proposed method.\n\n**Q. Average Training Time**\n\n> What is the average training time for models using DxMI compared to other generative model training techniques such as GANs or VAE-based approaches?\n\n* Our CIFAR-10 experiment takes less than 24 hours on two A100 GPUs, while our ImageNet 64 experiment takes approximately 48 hours on four A100 GPUs. The computational resources required for DxMI training are significantly lower compared to state-of-the-art GANs, which can take up to 48 hours on a Google TPU V3 Pod [2]. A TPU Pod consists of 1024 TPU chips, which are considerably more powerful than a few A100 GPUs. This difference is partly because DxMI can leverage a pre-trained diffusion model.\n\n[2] Brock et al. Large scale GAN training for high fidelity natural image synthesis. ICLR 2019. https://arxiv.org/abs/1809.11096\n\n**Q. Scalability**\n> Scalability: How scalable is the proposed DxMI approach when applied to larger datasets or models with significantly more parameters? Are there any anticipated bottlenecks?\n\n* For now, we do not observe any sign of bottleneck or inscalability. We believe DxMI is as scalable as GANs, since DxMI also leverages an EBM as a discriminator. GANs have already been shown to be scalable to very high-dimensional data [2].\n\n\nAgain, we thank you for providing valuable comments. Do not hesitate to let us know if you have further questions.\n\nBest wishes,   \nAuthors.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "URzc54JKDe", "reviewer": "Reviewer_Nqik", "summary": "This paper seeks to improve diffusion models by employing inverse reinforcement learning methods of imitation rather than (more myopic) behavioral cloning methods, which prevalent existing diffusion models can be viewed as using. It trains an energy-based model using maximum entropy inverse reinforcement learning and proposes an optimal control problem for diffusion based on minimizing an upper bound of the contrastive KL divergence. The benefits of the approach are demonstrated with a focus on generating outputs with few diffusion iterations.", "questions": "Is there potential for theoretical analysis or guarantees using this approach?\n\nAre there visual differences in the generated outputs of different models that could be highlighted in 5.2?", "limitations": "Potential abuses using deep fakes are described, along with limitations.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "The paper is motivated by the key connection between imitation and diffusion models, the characterization of existing methods corresponding to behavioral cloning approaches for imitations, and the potential for improved diffusion using more sophisticated inverse reinforcement learning imitation methods. This motivation is nicely described.\n\nThe technical content of the paper is quite dense, but the authors present it clearly. \n\nOther work exists exploring this perspective of diffusion as imitation/optimal control, but the paper’s approach is nicely constructed to avoid MCMC and policy gradient optimization, which are often bottlenecks in existing methods. \n\nExperimental results show the benefits of this approach, including comparisons with a similar approach that is reliant on policy gradient (SFT-PG) and other recent diffusion model learning methods for generation and anomaly detection.", "weaknesses": "Closely related work isn’t described in the introduction to better frame the contributions of this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer Nqik,\n\nThank you for taking the time to review our work. We deeply value your feedback and are happy to address your questions.\n\n**Q. Is theoretical analysis possible?**\n> Is there potential for theoretical analysis or guarantees using this approach?\n\nThanks for bringing up an important point. \n\n* Yes, there is significant potential for theoretical analysis in our MaxEnt IRL problem. However, directly analyzing the DxMI implementation may be highly challenging because our EBM and diffusion model comprise deep neural networks.\n* Conducting theoretical analysis in a more restricted setting might be more feasible while still providing valuable insights. Previous works [1,2] offered convergence guarantees for MaxEnt IRL in a discrete state-action space. While their results do not directly apply to DxMI due to algorithmic and other differences, their results suggest that a similar analysis could be conducted on DxMI under suitable assumptions.\n* Please consider that the main focus of this paper is to present a practical algorithm that is empirically scalable and effective. We believe providing a theoretical analysis that rationalizes the empirical results in the paper is an important future work. \n* We will add a paragraph describing the theoretical results from [1,2] in our Related Work section. We will also mention the difficulty of theoretical analysis in our limitations section.\n\n\n[1] Renard et al., Convergence of a model-free entropy-regularized inverse reinforcement learning algorithm, arxiv, 2024.     \n\n[2] Zeng et al., Maximum-Likelihood Inverse Reinforcement Learning with Finite-Time Guarantees, NeurIPS 2022. \n\n**Q. Related works are not described in the introduction.**\n\n> Closely related work isn’t described in the introduction to better frame the contributions of this paper.\n\n\n* We had to defer our discussion on prior works to the Related Work section (Section 6) due to the page limitation. \n* In the revised manuscript, we will include a paragraph in the introduction which describes related work. If there is any specific work that you want us to additionally cite, we are happy to incorporate them into the updated manuscript.\n\n**Q. Are there visual differences in the generated outputs of different models?**\n\n* Please find examples of randomly generated images from DxMI and other models in the attached PDF, highlighting the visual differences in their outputs. For example, the Consistency Model samples distort human faces, while the DxMI samples depict them in correct proportions.\n\nAgain, we thank the reviewer for acknowledging the value of our work. If you have any further questions, we are happy to address them.\n\nBest regards,  \nAuthors.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PKUoFy0mCu", "reviewer": "Reviewer_yjrf", "summary": "The authors propose learning a denoising diffusion model without using the denoising loss. Instead, they propose first training an energy based model and then treating the diffusion denoising sampler as an RL trajectory with the energy based model as the reward. To learn the energy based model, however they propose a generalized version of contrastive divergence which uses the current diffusion model as part of its objective function. Experiments are designed to validate this idea, starting from pretrained diffusion checkpoints.", "questions": "In Algorithm 1, Line 5, are the $x$ and $x_T$ completely independent samples with no relation to each other? From the way its written, $x_T$ is sampled starting from independent noise $x_0$, and $x$ is a particular data sample for that iteration.  Moreover, since for the pre-trained checkpoints, $\\pi$ is likely close to $p$, the energy objective looks like it has minimal training signal.", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper explores a creative combination of a variety of ideas all seeking to replace the simple denoising diffusion training objective. \n- The exposition is well written\n- The proposed approach involves learning both an energy based model and a diffusion model using DP where the number of steps can be small. The latter has value for fine-tuning diffusion models against other types of rewards.", "weaknesses": "- The authors introduce an elaborate scheme to do away with the score function denoising objective, but their experiments rely on pre-trained checkpoints that use the score function training to get these results so the final results are based on stacking the new methods on top of the original diffusion training.\n \n- The objective in Eq (5) seems somewhat hard to be confident of, in the sense that given any $\\pi$ and $p$, the optimal $q$ could diverge away from $p$ such that $KL(p||q) < KL(\\pi|q) >> 0$ (while still being closer than q since the optimum value of the objective has to be negative). This could theoretically make it unstable as an iteration progresses, $\\pi$ will track $q$ and then $q$ can take another step further away from both $p$ and $\\pi$ while being optimal for Eq (5), which only requires that the relative distance to $p$ be less than that of $q$.  As another point, when initializing with a $\\pi$ that is close to $p$ (e.g. in the pretrained model case), the optimal $q$ does not look like it has to be close to either given the cancellation involved.  \n\n-  The key proposal in Algorithm 1 has a multiple highly unstable procedures mixed together within each iteration -- e.g. the energy update in Line 5 and the TD bootstrap objective in line 7.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer yjrf,\n\nWe appreicate your time and effort in reviewing our work. Here, we are happy to address your concerns and questions.\n\n**Q. Is the objective function stable?** \n\n> The objective in Eq (5) seems somewhat hard to be confident of, in the sense that given any $\\pi$ and $p$, the optimal $q$ could diverge away from p such that $KL(p||q) = KL(\\pi||q) \\gg 0$.  This could make it unstable as an iteration progresses, $\\pi$ will track q and then q can take another step further away from both p and $\\pi$ while being optimal for Eq (5). \n\nWe would like to clarify the critical misunderstandings regarding the objective function. \n\nLet us write our objective function as $L(q,\\pi)=KL(p \\parallel q) - KL(\\pi \\parallel q)$, where we aim to solve $\\min_q \\max_\\pi L(q,\\pi)$.\n\n* When $p$ and $\\pi$ are fixed and $p\\neq \\pi$, the optimal $q$ does not satify $KL(p \\parallel q) = KL(\\pi \\parallel q) \\gg 0$.\n    * It is a misunderstanding that our objective function $L(q,\\pi)$ has a minimum at 0, where $KL(p||q) = KL(\\pi||q)$.\n    * When $p\\neq \\pi$, the objective function can have a negative value. For example, setting $q=p$ makes the objective negative $L(p,\\pi)=-KL(p \\parallel \\pi)<0$. Therefore, $KL(p||q) = KL(\\pi||q)$ is not optimal for $q$, as there are other values of $q$ that achieve a lower objective function value.\n    * Thus, our objective function always keep $q$ closer to $p$ than to $\\pi$, not invoking the instability questioned in the review comment.\n\n\n**Q. Is learning possible when $\\pi$ is close to $p$?**\n> As another point, when initializing with a $\\pi$ that is close to p (e.g. in the pretrained model case), the optimal q does not look like it has to be close to either given the cancellation involved.\n\n> Moreover, since for the pre-trained checkpoints, π is likely close to p, the energy objective looks like it has minimal training signal.\n\n* When $p=\\pi\\neq q$, it is true that there is no learning signal for $q$. However, $p=\\pi\\neq q$ is not a Nash equilibrium, and learning is not terminated at this point.\n* When $p=\\pi\\neq q$, our objective function drives $\\pi$ away from $p$ to be close to $q$. After $\\pi$ becomes different from $p$, the learning signal for $q$ is generated. \n    * This behavior may be undesirable in practice. However, we are most interested in the case where the number of function evaluations is small ($T=4$ or $10$). With small $T$, the initial sample quality from $\\pi$ is very bad (Figure 1 (Right) of the manuscript), indicating that $p \\neq \\pi$.\n* As $p=\\pi\\neq q$ is not a Nash equilibrium of our objective function, an optimization algorithm, if done correctly, will eventually lead to the Nash equilibirum $p=q=\\pi$.\n\n\n**Q. DxMI still uses a diffusion model as a starting point.**\n> The authors introduce an elaborate scheme to do away with the score function denoising objective, but their experiments rely on pre-trained checkpoints that use the score function training to get these results so the final results are based on stacking the new methods on top of the original diffusion training.\n\n* DxMI is not meant to be a complete replacement for denoising objective. We will update our introduction to clarify that DxMI is a complementary training algorithm for diffusion models.\n* Please note that diffusion model pre-training is not always necessary for DxMI. In our 2D experiment and anomaly detection experiment, DxMI demonstrated its ability to train a sampler without pre-training.\n\n\n**Q. The proposed algorithm has multiple unstable procedures.**\n\n> The key proposal in Algorithm 1 has a multiple highly unstable procedures mixed together within each iteration -- e.g. the energy update in Line 5 and the TD bootstrap objective in line 7.\n\n* We understand the proposed algorithm can be seemingly complicated. \n    * To deal with the complexity, we will make our code public and provide model checkpoints. Also, we disclose our hyperparameters and suggest a hyperparameter selection strategy in Appendix B.\n* However, to the best of our knowledge, there is no empirical or theoretical evidence of instability.\n    * The TD update and the energy update does not interfere with each other, as they operate on different inputs. Both updates are stable, as the energy update equation is regularized (coefficient $\\gamma$), and TD update is simply mean squared error minimization.\n* Empirically, we found the algorithm much more stable than MCMC-based EBM training algorithms, which occasionally diverge for no reason.\n* If you have a particular concern regarding the algorithm's stability, we are happy to discuss it.\n\n\n**Q. In Algorithm 1, Line 5, are the x and x_T completely independent samples with no relation to each other?**\n\n* You are correct that, in Algorithm 1, $\\mathbf{x}$ denotes a real data and $\\mathbf{x}_T$ indicates a sample from the diffusion model. We will make this point clear in the updated manuscript by adding a comment in Algorithm 1.\n\n\n\n\nThank you again for your constructive review. We believe there was a misunderstanding regarding our objective function, and we hope our response clarifies this issue. We kindly request that you reconsider the decision in light of our responses and update the score accordingly. We are also eager to address any additional concerns you may have.\n\nBest regards,  \nAuthors.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "S5VjlN4nUD", "reviewer": "Reviewer_YBX1", "summary": "This work proposes a maximum entropy inverse reinforcement learning (IRL) approach for improving the sample quality of diffusion generative models, especially when using a small number of generation steps.", "questions": "1. Upon reviewing this work, I found that the energy-based objective (Eq. (2)) and training objectives (Eqs. (3), (4), (5)) share very strong similarities/motivations to the concepts employed in existing KL-regularized RL-based fine-tuning papers (e.g., [1-4]). In particular, Eq. (6) and Eq. (7) in [2] serve precisely as the KL-based training objectives and the energy-based model, which involves sampling from unnormalized distributions. While I think the similarities might be just in terms of high-level methodologies and principles, it still becomes crucial for the authors to provide a thorough discussion highlighting the technical and methodological distinctions between this approach and the prior works in order to well-position the proposed method. For discussing the methodological relations/differences, I also recommend the authors carefully check [1] as this paper is written in a principled way, such that it becomes a good reference for understanding the methodologies.\n\n2. As mentioned, do you have any theoretical insight into the convergence properties of DxDP or approximation guarantees relating to the IRL and original objectives? Many IRL theories might be worth checking (references [1-4] also might be worth checking for this goal, probably). Even without providing rates/bounds, an intuitive discussion/remark on these points could further validate the approach.\n\n3. Can the DxMI approach be extended to other families of generative models, e.g., bridge-matching diffusion models, which are very similar to standard denoising diffusion models? What are the key challenges or requirements for the generative model?\n\n4. How do you expect this approach to scale to more complex datasets or higher-dimensional spaces? Will the sample efficiency gains be more or less pronounced?\n\n5. In experiments, how sensitive are the results to hyperparameters of DxMI/DxDP, such as the coefficient on the entropy term? What strategies did you use to tune these?\n\n[1] https://arxiv.org/abs/2403.06279\n\n[2] https://arxiv.org/abs/2402.16359\n\n[3] https://arxiv.org/abs/2305.16381\n\n[4] https://arxiv.org/abs/2405.19673", "limitations": "The authors adequately addressed the limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents a novel and principled approach for improving diffusion models by formulating the problem as maximum entropy IRL. This elegant formulation allows jointly optimizing a diffusion model and EBM to enhance sample quality, especially with fewer generation steps.\n\n2. The proposed DxDP algorithm is a key technical innovation that makes optimizing the IRL objective tractable. By leveraging ideas from optimal control and dynamic programming, DxDP enables efficient diffusion model updates without back-propagation through time, which is a significant practical benefit.\n\n3. Strong empirical results on both image generation and anomaly detection tasks demonstrate the approach's effectiveness. DxMI can generate high-quality samples with very few (4-10) steps and enables training EBMs without MCMC, such a diffusion step reduction is impressive to me.", "weaknesses": "1. While the acclaimed acceleration of the diffusion model looks effective (and is also stated as a major advantage of this algorithm), the comparison to prior diffusion model acceleration methods is somewhat limited. A more comprehensive evaluation across different speed-quality tradeoffs and more discussion of DxMI's relative strengths and weaknesses compared to other approaches would be essential.\n\n2. While the empirical results look good, the paper lacks theoretical analysis of the proposed methods, such as convergence rates for DxDP or any approximation guarantees relating the IRL objective to the original objective. Adding such analysis would help characterize this method's ups and downs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer YBX1,\n\nWe deeply appreciate your thorough and constructive comments. We will do our best to answer your questions.\n\n**Q. What is the connection to KL-regularized RL fine-tuning, particularly in papers such as [1-4]?**\n\nThanks for pointing out an interesting connection. The sampler update step of DxMI is indeed closely related to KL-regularized RL.\n\n* We will augment our manuscript to discuss this connection in detail.\n    * We will add paragraphs dedicated to KL-regularized fine-tuning. In the paragraphs, we will cite the papers [1-4] and discuss them in detail.\n    * Please note that the current manuscript already mentions some works using KL-regularized RL, including [3], in the second paragraph of the Related Work section.\n* The sampler update objective of DxMI is a special case of KL-regularized RL. However, there are three key differences that makes DxMI distinct from existing KL-regularized RL fine-tuning methods.\n    * First, DxMI employs an uninformative reference policy, such as a Gaussian distribution (Eq. (7)). Due to this choice, DxMI can deviate from the pretraining model to find a better generation trajectory more suitable for small $T$.\n    * Second, DxMI employs a novel value-based algorithm for updating a diffusion model.\n    * Third, while most KL-regularized RL works assume that the reward function is known, DxMI simultaneously learns the reward from data.\n* BRAID [4] is particularly related to DxMI, as it also considers the problem of learning a reward from offline data. However, the reward is a separate random variable in [4], while in DxMI the reward is the log data density.\n\n[1] https://arxiv.org/abs/2403.06279 Tang. Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond. 2024.\n\n[2] https://arxiv.org/abs/2402.16359 Uehara et al. Feedback Efficient Online Fine-Tuning of Diffusion Models. 2024.\n\n[3] https://arxiv.org/abs/2305.16381 Fan et al. DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models. 2023. \n\n\n[4] https://arxiv.org/abs/2405.19673 Uehara et al. Bridging Model-Based Optimization and Generative Modeling via Conservative Fine-Tuning of Diffusion Models. 2024.\n\n**Q. The comparison to prior diffusion model acceleration methods is somewhat limited.**\n\n* Our comparative discussion had to be limited due to page constraints. Currently, the first paragraph of Section 6 (\"Faster Diffusion Models\") is dedicated to reviewing and comparing prior diffusion acceleration methods.\n* In the updated manuscript, we will enhance the \"Faster Diffusion Models\" section to highlight the key differences between DxMI and previous methods. Additionally, we will incorporate this comparison into the introduction and the experiments sections. The enhancements will include the following points:\n    * The key distinction between DxMI and existing diffusion acceleration methods is that DxMI does not use the intermediate states in the trajectory of a diffusion model. Most diffusion distillation methods focus on preserving or learning from these intermediate states. In contrast, DxMI directly aims to match the final state of a sampler to the data distribution. The promising performance of DxMI indicates that deviating from the original diffusion trajectory may be beneficial for sample quality when the generation has to be performed in a very few steps. Among existing methods, only SFT-PG employs a similar approach; however, DxMI outperforms SFT-PG by using dynamic programming instead of a policy gradient.\n\n\n**Q. A more comprehensive evaluation across different speed-quality tradeoffs would be essential.** \n\n* In the updated manuscript, we will include a figure demonstrating the speed-quality trade-off with more data points, such as $T$=2, 20, and 40.\n* Qualitatively, the best trade-off for DxMI is achieved in the mid-range of $T$, from 4 to 10.\n    * If $T$ is too small, the sampler is less capable, and the data processing inequality becomes less tight, making our MaxEnt regularization less effective.\n    * If $T$ is too large, the sampler's capability increases, but the value function learning becomes more challenging.\n\n**Q. More discussion of DxMI's relative strengths and weaknesses compared to other approaches would be essential.**\n\nWe will augment our discussion on strengths and weaknesses of DxMI in the update manuscript. Currently, some of weaknesses are mentioned in our Limitation section. Focusing on diffusion acceleration application, our strengths and weaknesses can be summarized as follows.\n* Strengths\n    * Unlike other diffusion distillation methods where the performance is bounded by the teacher model, in principle DxMI may achieve better sample quality than the pretrained model (see our CIFAR-10 case).\n    * The dynamic programming-based in DxMI is more effective than the policy gradient-based algorithm (e.g., SFT-PG).\n    * DxMI produces an EBM as a byproduct, which can be used in other applications such as anomaly detection or transfer learning.\n* Weaknesses \n    * When $T=1$, DxMI reduces to GAN, offering no additional advantage.\n    * DxMI does not offer the flexibility of using a different value of $T$ during the test time.\n\n----\nDue to the character limit, we will continue answering your question in the comment.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## General Comments to AC and All Reviewers

We appreciate all reviewers for their thoughtful comments and remarks. We thank the reviewers for their insightful feedback and constructive comments and for providing suggestions that would improve our paper.

First of all, we are encouraged that the reviewers found the following points:

(i) Our approach and formulation are novel, principled, and elegant in improving sample quality and inference time of a diffusion model. (EhWr, YBX1)   

(ii) The motivation is nicely described, the algorithm is well constructed to avoid bottlenecks in existing methods (Nqik), and it uses a creative combination of ideas (yjrf).  

(iv) The experimental results show the benefits of the proposed approach (Nqik), which are strong and impressive (YBX1).  

(iii) The manuscript is well-structured, and the presentation is clear. (EhWr, Nqik, yjrf)   


Based on the feedback from all reviewers, the most significant shared concerns or major points to address are as follows. 

**On theoretical analysis (Nqik, YBX1):** We agree with the reviewers that a theoretical analysis of our MaxEnt IRL problem would be valuable. To the best of our knowledge, there are few theoretical results available for MaxEnt IRL. Previous works [1,2] have provided convergence guarantees for MaxEnt IRL in a discrete state-action space with an infinite horizon. In contrast, our approach considers a continuous state-action space with a finite horizon. Additionally, there are algorithmic differences regarding the reward functions. Consequently, their results do not directly apply to DxMI. However, we believe a similar analysis could be performed on DxMI under appropriate assumptions and a simplified setting, particularly with linear reward functions and a tabular policy.

[1] Renard et al., Convergence of a model-free entropy-regularized inverse reinforcement learning algorithm, arxiv, 2024.   

[2] Zeng et al., Maximum-Likelihood Inverse Reinforcement Learning with Finite-Time Guarantees, NeurIPS 2022. 


**On scalability and complexity of the algorithm (EhWr, YBX1):** Because we formulate the problem using MaxEnt IRL, DxMI might initially seem complex. However, its complexity is comparable to Generative Adversarial Networks (GANs) and actor-critic reinforcement learning methods. We believe that DxMI is as scalable as conventional diffusion models and GANs, which have already demonstrated scalability to very high-dimensional data. To demonstrate the scalability of DxMI, we run DxMI (T=4) on LSUN Bedroom 256x256 and achieve competitive results (FID: 5.93, Recall: 0.477), where DDPM (T=1000) achieves FID: 4.89, Recall 0.45. Additionally, we have not observed any signs of bottlenecks or scalability issues. For instance, DxMI is not particularly sensitive to the entropy regularization coefficient. We provide guidance for hyperparameter tuning in Appendix B. More specific answers can be found in individual comments.


The attached PDF provides examples of generation from DxMI on CIFAR-10 and ImageNet 64x64, along with samples from competitive baselines. 

We hope the above answers some common questions from the reviewers. We also respond to individual comments from each reviewer below.

Meta-review: The paper presents a novel and efficient methodology to improve diffusion generative models based on the formulation of an maximum entropy inverse RL problem.

All reviewers agree this paper presents a novel and significant contribution which is methodologically strong.

During the author's rebuttal several points were discussed in detail.
In order of importance, the points revolved around:

(1) The complexity and scalability of the algorithm.

(2) The existence of theoretical guarantees / formal statements to ensure the stability of the objective function.

(3) Additional comparative analysis to prior diffusion model acceleration methods with relative strengths and weaknesses.

(4) The need for additional experiments.

(5) A more comprehensive evaluation across different speed-quality trade-offs.

(6) Misunderstandings related to the objective function, the need of a pre-trained diffusion model.

The rebuttal was convincing in clarifying all these points and no further discussion was required.

I think this is a very strong paper which will have significant impact and thus recommend acceptance.

While the primary focus of this work is to introduce a practical algorithm, I suggest incorporating the discussion regarding the above points as much as possible.
Decision: Accept (oral)