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_cc_2024_Conference
Published: 2024-05-15 04:02:05

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.

Main paper: 
 introduction: 
 Generative modeling is a form of imitation learning. Just as an imitation learner produces an action that mimics a demonstration from an expert, a generative model synthesizes a 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 [1, 2] .The connection to imitation learning plays a central role in diffusion models [3, 4] , which generate samples by transforming a Gaussian noise through iterative additive refinements. The training of a diffusion model is essentially an instance of behavioral cloning [5] , a widely adopted imitation learning algorithm that mimics an expert's action at each state. During training, a diffusion model is optimized to follow a 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 a new state-of-the-art in many generation tasks.Behavioral cloning is also responsible for the diffusion model's key limitation, the slow generation speed. A behavior-cloned policy is not reliable when the state distribution deviates from the expert demonstration [6, 7] . Likewise, the sample quality from a diffusion model degrades as the gap 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 [8, 9, 10, 11, 12, 13] .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 at 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 be found by guiding a sampler using the learned reward [11, 12, 13] . This approach is more frequently referred to as adversarial training because a 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, 2] for a diffusion model. Our formulation, named Diffusion by Maximum Entropy IRL (DxMI, pronounced "diby-me"), is a minimax problem that jointly optimizes a diffusion model and an energy-based model (EBM). In DxMI, the EBM provides the estimated log density as the reward signal for the diffusion model. Then, the diffusion model is trained to 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.The diffusion model update step in DxMI is equivalent to maximum entropy RL. However, this step is challenging to perform for two reasons. First, the diffusion model update requires the gradient of the marginal entropy of a diffusion model's samples, which is difficult to estimate for discrete-time diffusion models, such as DDPM [3] . Second, back-propagating 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 a 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 [20, 21, 22] , but DxDP is 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 RL algorithm, DxDP may have broader utility, such as fine-tuning a diffusion model from human or 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 3 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://github.com/swyoon/Diffusion-by-MaxEntIRL.git. 
preliminaries: 
 Diffusion Models. The diffusion model refers to a 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 as DDPM [3] , which synthesize a sample through the following iteration producing x 0 , x 1 , . . . , x T ∈ R D :EQUATIONwhere ϵ t ∼ N (0, I) and µ(x , t) is the output of a neural network. The coefficients a t ∈ R and σ t ∈ R >0 are constants. Note that we reverse the time direction in a diffusion model from the convention to be consistent with RL. The final state x T is the sample generated by the diffusion model, and its marginal distribution is π(x T ). We will often drop the subscript T and write the distribution as π(x). The conditional distribution of a transition in Eq. ( 1) is denoted as π(x t+1 |x t ).For continuous-time diffusion models [23] , Eq. ( 1) corresponds to using the Euler-Maruyama solver.The generation process in Eq. ( 1) defines a T -horizon Markov Decision Process (MDP) except for the reward. State s t and action a t are defined as s t = (x t , t) and a t = x t+1 . The transition dynamics is defined as p(s t+1 |s t , a t ) = δ (at,t+1) , where δ (xt,t) is a Dirac delta function at (x t , t). With a reward function defined, a diffusion model can be trained with RL [11, 24, 19, 25] . In this paper, we consider a case where the reward is the log data density log p(x), which is unknown.Energy-Based Models. An energy-based model (EBM) q(x) uses a scalar function called an energy E(x) to represent a probability distribution:EQUATIONwhere τ > 0 is temperature, X is the compact domain of data, and Z = X exp(-E(x)/τ )dx < ∞ is 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 q KL(p||q), where KL(p||q) := X p(x) log(p(x)/q(x))dx. Computing the gradient of KL(p||q) with respect to EBM parameters requires MCMC sampling, which is computationally demanding and sensitive to hyperparameters. The algorithm presented in this paper serves as 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 [15, 16, 2] . In maximum entropy IRL, x corresponds to an action (or a sequence of actions), and E(x) represents the expert's cost of the action. The expert is then assumed to generate actions following q(x). This assumption embodies the maximum entropy principle because Eq. ( 2) is a distribution that minimizes the cost while maximizing the entropy of the action. Here, τ balances cost minimization and entropy maximization.3 Diffusion by Maximum Entropy Inverse Reinforcement Learning 
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 π(x), particularly when T is small.EQUATIONwhere Π is the set of feasible π(x)'s, and H(π) = -π(x) log π(x)dx is the differential entropy. This minimization is a maximum entropy RL problem: The log data density log p(x) is the reward, and π(x) 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 Diffusion by Maximum Entropy IRL (DxMI) as an IRL approach for solving Eq. ( 3). We employ an EBM q(x) (Eq. ( 2)) as a surrogate for p(x) and use log q(x) as the reward for training the diffusion model instead of log p(x). At the same time, we train q(x) to be close to p(x) by minimizing the divergence between p(x) and q(x):EQUATIONwhere Q is the feasible set of EBMs. When the two KL divergences become zero, p(x) = q(x) = π(x) is achieved. However, min q∈Q KL(p(x)||q(x)) is difficult due to the normalization constant of q(x). Instead, we consider an alternative minimax formulation inspired by Contrastive Divergence (CD; [26] ), a celebrated algorithm for training an EBM.Objective. Let p(x) be the data distribution. Suppose that Q and Π are the feasible sets of the EBM q(x) and the diffusion model π(x), respectively. The learning problem of DxMI is formulated as follows:EQUATIONWe shall call Eq. ( 5) Generalized CD (GCD), because Eq. ( 5) generalizes CD by incorporating a general class of samplers. CD [26] is originally given as min q∈Q KL(p(x)||q(x)) -KL(ν k p,q (x)||q(x)). Here, ν k p,q (x) is a k-step MCMC sample distribution where Markov chains having q(x) as a stationary distribution are initialized from p(x). GCD replaces ν k p,q (x) with a general sampler π(x) at the expense of introducing a max operator.When the models are well-specified, i.e., p(x) ∈ Q = Π, Nash equilibrium of GCD is p(x) = q(x) = π(x), which is identical to the solution of Eq. (4). Meanwhile, there is no 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 q(x) be closer to p(x) than π(x).Our main contribution is exploring the application of discrete-time diffusion models (Eq. ( 1)) as π(x) in GCD and not discovering GCD for the first time. GCD is mathematically equivalent to a formulation called variational maximum likelihood or adversarial EBM, which has appeared several times in EBM literature [27, 28, 29, 30, 31, 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 4.In DxMI, we update a diffusion model and an EBM in an alternative manner to find the Nash equilibrium. We write θ and ϕ as the parameters of the energy E θ (x) and a diffusion model π ϕ (x), respectively. While EBM update is straightforward, we require a subroutine described in Section 4 for updating the diffusion model. The entire procedure of DxMI is summarized in Algorithm 1.The optimization with respect to EBM is written asmin θ E p(x) [E θ (x)] - E π ϕ (x) [E θ (x)].During the update, we also regularize the energy by the square of energiesγ(E p(x) [E θ (x) 2 ] + E π ϕ (x) [E θ (x) 2 ]) for γ > 0 to ensure the energy is bounded. We set γ = 1 unless stated otherwise. This regularizer is widely adopted in EBM [34, 35] .ϕ ||q θ ) = E π ϕ (x) [E θ (x)/τ ] -H(π ϕ (x)).However, this update is difficult in practice for two reasons.First, marginal entropy H(π ϕ (x)) is difficult to estimate. Discrete-time diffusion models (Eq. ( 1)) do not provide an efficient way to evaluate log π ϕ (x) 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 1 Diffusion by Maximum Entropy IRL 1: Input: Dataset D, Energy E θ (x), Value V t ψ (x t ), and Sampler π ϕ (x 0:T ) 2: s t ← σ t // Initialize Adaptive Velocity Regularization (AVR) 3: for x in D do // Minibatch dimension is omitted for brevity.Sample x 0:T ∼ π ϕ (x 0:T ).5: min θ E θ (x) -E θ (x T ) +γ(E θ (x) 2 + E θ (x T ) 2 ) //EQUATIONend for 9:for x t randomly chosen among x 0:T do // Sampler update 10:Sample one-step: x t+1 ∼ π ϕ (x t+1 |x t ) //min ϕ V t+1 ψ (x t+1 (ϕ)) + τ log π(x t+1 |x t ) + τ 2s 2 t ||x t -x t+1 (ϕ)|| 2 // x t+1 is a function of ϕ.12:end for 13:s 2 t ← αs 2 t + (1 -α)||x t -x t+1 || 2 /D// AVR update 14: end for high-dimensional spaces. Second, propagating the gradient through time in a diffusion model may require significant memory. Also, the gradient may explode or vanish during propagation, making the training unstable [19] . 
diffusion by dynamic programming: 
 In this section, we present a novel RL algorithm for a diffusion model, Diffusion by Dynamic Programming (DxDP), which addresses the difficulties in 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 to train a diffusion model for an arbitrary reward.Instead of solving min ϕ KL(π ϕ (x T )||q θ (x T )) directly, we minimize its upper bound obtained from the data processing inequality:EQUATION) Here, we introduce an auxiliary distribution q(x 0:T -1 |x T ), and the inequality holds for an arbitrary choice of q(x 0:T -1 |x T ). In this paper, we consider a particularly simple case where q(x 0:T -1 |x T ) is factorized into conditional Gaussians as follows:EQUATIONNow we minimize the right-hand side of Eq. ( 6):min ϕ KL(π ϕ (x 0:T )||q θ (x T )q(x 0:T -1 |x T )).When we plug in the definitions of each distribution, multiply by τ , and discard all constants, we obtain the following problem:EQUATIONwhich is an optimal control problem. The controller π ϕ (•) is optimized to minimize the terminal cost E θ (x T ) plus the running costs for each transition (x t , x t+1 ).The first running cost log π ϕ (x t+1 |x t ) is responsible for conditional entropy maximization, becauseE π [log π ϕ (x t+1 |x t )] = -H(π ϕ (x t+1 |x t )).The second running cost regularizes the "velocity" ||x t+1 -x t || 2 . The temperature τ balances between the terminal and running costs.We have circumvented the marginal entropy computation in GCD, as all terms in Eq. ( 8) are easily computable. For the diffusion model considered in this paper (Eq. ( 1)), the conditional entropy has a particularly simple expression H(π ϕ (x t+1 |x t )) = D log σ t + 0.5D log 2π. Therefore, optimizing the entropy running cost amounts to learning σ t 's in diffusion, and we treat σ t 's as a part of the diffusion model parameter ϕ in DxMI. Table 1 : Quantitative results for 8 Gaussians experiment. SW denotes the sliced Wasserstein distance between samples and data. AUC is 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. We propose a dynamic programming approach for solving Eq. ( 8). Dynamic programming introduces value functions to break down the problem into smaller problems at each timestep, removing the need for back-propagation in time. Then, a policy, a diffusion model in our case, is optimized through iterative alternating applications of policy evaluation and policy improvement steps.Value Function. A value function, or cost-to-go function V t ψ (x t ), is defined as the expected sum of the future costs starting from x t , following π. We write the parameters of a value function as ψ.EQUATIONfor t = 0, . . . , T -1. Note that V T (x T ) = E(x T ). A value function can be implemented with a neural network, but there are multiple design choices, such as whether to share the parameters with π(x) or E(x), and also whether the parameters should be shared across time. We explore the options in our experiments.During policy evaluation, we estimate the value function for the current diffusion model by minimizing the Bellman residual, resulting in the temporal difference update.EQUATIONwhere 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 x t in a trajectory x 0:T sampled from π ϕ (x 0:T ), the diffusion model is optimized to minimize the next-state value and the running costs.EQUATIONIn practice, each iteration of policy evaluation and improvement involves a single gradient step.We additionally propose a method for systematically determining the hyperparameter s t 's of the auxiliary distribution q(x 0:T -1 |x T ). We can optimize s t such that the inequality Eq. ( 6) is as tight as possible by solving min s0,...,s T -1 KL(π ϕ (x 0:T )||q θ (x T )q(x 0:T -1 |x T )). After calculation (details in Appendix A), the optimal s * t can be obtained analytically:(s * t ) 2 = E xt,xt+1∼π [||x t -x t+1 || 2 ]/D.In practice, we can use exponential moving average to compute the expectationE xt,xt+1∼π [||x t -x t+1 || 2 ] during training: s 2 t ← αs 2 t + (1 -α)||x t -x t+1 || 2 /Dwhere we set α = 0.99 for all experiment.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 in Algorithm 2.Time-Independent Value Function.In image generation experiments (Section 5.2), we let the value function be independent of time, i.e., V t ψ (x t ) = V ψ (x t ). Removing the time dependence reduces the number of parameters to be trained. More importantly, a time-independent 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-dependent value function V t ψ (x t ) never observes samples having different noise levels than the noise level of x t .Also, in the value update (Eq. ( 10)) step of image generation experiments, we introduce time cost function R(t) > 0, which replaces the running cost terms τ log π ϕ (xEQUATIONt ). The time cost R(t) only depends on time t. The modified value update equation is given as follows:EQUATIONMeanwhile, we retain the running cost terms in the diffusion model (policy) update step (Eq. ( 11)).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 to decrease linearly as time progresses.Thesigmoid time cost is R(t) = σ(-t+T /2))-σ(-t-1+T /2)), where σ(x) = (1+exp(-x)) -1 .With the sigmoid time cost, the value function is trained to follow a sigmoid function centered at T /2 when plotted against the time. Other forms of R(t) are also possible.Separate tuning of τ .In image generation, we assign different values of temperature τ for entropy regularization log π ϕ (x t+1 |x t ) and velocity regularization ||x t -x t+1 || 2 /(2s 2 t ), such that the resulting running cost becomes τ 1 log π ϕ (x t+1 |x t ) + τ 2 ||x t -x t+1 || 2 /(2s 2 t ). Typically, we found τ 1 > τ 2 beneficial, indicating the benefit of exploration. Setting τ 1 ̸ = τ 2 does not violate our maximum entropy formulation, as scaling τ 2 is equivalent to scaling s 2 t 's, which can be set arbitrarily. 
experiments: 
 In this section, we provide empirical studies that demonstrate the effectiveness of DxMI in training a diffusion model and an EBM. We first present a 2D example, followed by image generation and anomaly detection experiments. More details on experiments can be found in Appendix C.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 timeconditioned multi-layer perceptron (MLP). The last time step (T = 5) of the value network is treated as the energy. The sample quality is measured with sliced Wasserstein distance (SW) to test data. Also, we quantify the quality of an energy function through the classification performance to uniform noise samples (Table 1 ).First, we investigate the effect of maximum entropy regularization τ . Setting an appropriate value for τ 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 (τ = 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 (τ = 1), the generated samples are close to noise. In this regime, DxMI behaves similarly to Noise Contrastive Estimation [41] , enabling energy function learning to a certain extent. These effects are visualized in Fig. 2 .Next, we experiment on whether pre-training a sampler as DDPM helps DxMI. Table 1 suggests that the pre-training is beneficial but not necessary to make DxMI work. We also visualize the value functions in Fig. 3 and find that the time evolution of value interpolates the data distribution and a Gaussian distribution. (256×256), using three diffusion model backbones, DDPM [3] , DDGAN [46] , and variance exploding version of EDM [50] . The results can be found in Table 2 , 3, and 4. Starting from a publicly available checkpoint of each pretrained backbone, we first adjust the noise schedule for the target sampling steps T . 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 are used as the starting point of DxMI training. A single CIFAR-10 run reaches the best FID in less than 4 hours on four A100 GPUs. We set τ 1 = 0.1 and τ 2 = 0.01. The sigmoid time cost is used for all image generation experiments. The sample quality is 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 are in Appendix C. Furthermore, DxMI outperforms SFT-PG [11] , another IRL approach implemented with a policy gradient. For a 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 a 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 are set to zero τ 1 = τ 2 = 0 in the sampler update (11) . Then, we compare the linear time cost to the sigmoid time cost. In both cases, We observe the increase in FID 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 x t+1 , we add a small drift with coefficient λ, i.e., x t+1 ← x t+1 -λσ t ∇ xt+1 V ψ (x t+1 ). 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. 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 224×224 RGB images of 15 object categories. We follow the multi-class problem setup proposed by [60] . The training dataset contains 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 detect 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 preprocessing protocol in [60, 59] , each image is transformed into a 272×14×14 vector using a pre-trained EfficientNet-b4 [62] . DxMI is conducted in a 272-dimensional space, treating 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.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 a sample as its energy [63] . The diffusion model and the value function are five-step time-conditioned MLPs. Unlike conventional diffusion models, DxMI allows for a flexible choice of π(x 0 ). 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 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 τ = 0, the diffusion model fails to explore and only exploits regions of minimum energy, resulting in poor performance. We observe that a moderate level of τ = 0.1 benefits both the sampler and the energy function, as it encourages exploration and provides a suitable level of diversity in negative samples. 
related work: 
 Faster Diffusion Models. Significant effort has been dedicated to reducing the number of generation steps in diffusion models during sampling while preserving sample quality. One popular approach is 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, 66, 67, 68] , and utilizing non-Markovian formulations [45, 69, 70] . While these methods are training-free, the sample quality can be further improved when the neural network is directly tuned for short-run sampling. Distillation methods train a faster diffusion sampler using training signal from a longer-run diffusion model, showing strong performance [10, 71, 72, 48, 73, 43, 42, 74] . A distilled model usually cannot outperform the teacher model, but adversarial or IRL methods may exceed full-length diffusion models. Hybrid methods [13, 12] combine distillation with adversarial loss, while other methods [46, 75] apply adversarial training to each denoising step. DxMI and SFT-PG [11] rely fully on adversarial training for final samples, allowing beneficial deviations from the diffusion path and reducing statistical distance from the data.RL for Diffusion Model. RL is often employed to fine-tune 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 [19, 79, 25] . DxMI focuses on a setting where the estimated log data density is the reward. When RL is applied to diffusion models, the policy gradient [80] is the dominant choice [24, 76, 11, 77] . DxMI offers a value function-based approach as an alternative to the policy gradient. Maximum entropy RL for diffusion models is investigated in [20, 81, 82] but only in the continuous-time setting. DxDP investigates the discrete-time setting, which is more suitable for accelerating generation speed.Energy-Based Models. DxMI provides a method of utilizing a diffusion model to 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, 18, 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 [90, 33] . DxMI shares the objective function with several prior works in EBM [27, 28, 29, 30, 31, 33] . However, none of the works use a diffusion model directly as a sampler.The convergence guarantees of entropy-regularized IRL are provided in [91, 92] under the assumption of a linear reward and the infinite time horizon. Their guarantees are not directly applicable to a 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 [93] , which is relevant for the discrete version of DxDP. More focused analysis on entropy regularized RL for diffusion models is 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 a 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. 
 
 Appendix: Training multiple components simultaneously, DxMI introduces several hyperparameters. To reduce the overhead of practitioners, we provide a hyperparameter exploration guideline in Appendix B. DxMI is not directly applicable to training a single-step generator. However, a diffusion model fine-tuned by DxMI can be distilled to a single-step generator. DxMI does not offer the flexibility of using a different value of generation steps T 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.We are interested in the following optimization problem: min s0,...,s T -1 KL(π ϕ (x 0:T )||q θ (x T )q(x 0:T -1 |x T )).Plugging in our choice of log q(x t |x t+1 ) = - EQUATIONwhere constant term with respect to s t is omitted. Since the object function is separable, we can solve the optimization for each s t independently.EQUATIONThis optimization has an analytic solution:(s * t ) 2 = E xt,xt+1∼π ||x t -x t+1 || 2 /D.Value function. 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 time-independent. The second axis is whether the value function should share model parameters with other networks, such as the sampler or the energy network.Our experiments demonstrate various combinations of these design choices. In a 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 a 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 a time-dependent value function that does not share parameters with any other network. This design is suitable when a 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 τ plays an important role in DxMI, we recommend running DxMI with τ = 0 when implementing the algorithm for the first time. If everything is in order, DxMI should function to some 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 τ . Since γ determines the absolute scale of the energy, the magnitude of τ should be adjusted accordingly. We set γ = 1 in all our experiments and recommend this value. In such a case, the optimal τ is likely to be less than 0.5.Learning rate. As done in two time-scale optimization [55] , we use a larger learning rate for the value function than for the sampler. When training the noise parameters σ t in the diffusion model, we also assign a learning rate 100 times larger than that of the sampler.Sample quality is 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 Sample x 0:T ∼ π ϕ (x).5:min θ E θ (x) -E θ (x T ) +γ(E θ (x) 2 + E θ (x T ) 2 )// Energy update 6:for t = T -1, . . . , 0 do // Value update and τ 2 = 0.01. In the ImageNet experiments with T = 10, we set τ 1 = τ 2 = 0.01, and for T = 4, we set τ 1 = 0.1 and τ 2 = 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-preserving formulation, while EDM is a variance-exploding formulation. Exploring a unified method for selecting the entropy regularization parameter is an interesting research direction.Dataset and Feature Extraction. MVTec AD is a dataset designed to evaluate anomaly detection techniques, particularly for industrial inspection. It contains over 5000 high-resolution images divided into 15 different categories of objects and textures. Each category includes a set of defect-free training images and a test set with both defective and non-defective images.The original dataset consists of high-resolution RGB images sized 224×224. Following the methods used in similar studies [59, 60] , we extract a 272×14×14 full feature representation of each image using a pre-trained EfficientNet-b4 [62] . We then train the energy function using DxMI in 272dimensional space, treating every spatial feature from the training images as normal training data. Each data point x in R 272 is projected to S 272 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 a single image is calculated by applying max pooling to the energy value of each spatial feature. For anomaly localization, the energy values of the 14×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 a sample as its energy value. Considering that the data distribution lies on S 272 , 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 S 272 . 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 π(x 0 ). Unlike traditional diffusion models, DxMI permits a flexible choice of π(x 0 ).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 to the data in S 272 , we use the pretrained autoencoder that is also used for the initial energy function to project the samples from the data distribution to 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 as the energy function. Encoder and decoder each consist of an MLP with 3 hidden layers and 1024 hidden dimensions. We use a time-conditional MLP for the sampler and value function, encoding the time information into a 128-dimensional vector using sinusoidal positional encoding. The input x t 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 (=4×14×14) using the Adam optimizer. We use a learning rate of 10 -5 for the sampler and value function and 10 -4 for the energy function.DxMI (Algorithm 1) 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 a value function and two Q-functions-DxMI only trains a single value function and no Q-functions. 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 the 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 in the diffusion model. This is also significantly smaller than the critic networks used in GANs, such as the 158.3M parameters in BigGAN [49] . In practice, our CIFAR-10 experiment completes in under 24 hours on two A100 GPUs, while the ImageNet 64 experiment takes approximately 48 hours on four A100 GPUs.In this section, we show that our policy improvement method introduced in Eq. ( 11) can effectively minimize the KL divergence between the joint distributions, KL(π ϕ (x 0:T )||q θ (x T )q(x 0:T -1 |x T )).The policy improvement at each timestep t can be expressed as minπ(xt+1|xt)E xt,xt+1∼π V t+1 (x t+1 ) + τ log π(x t+1 |x t ) -τ log q(x t |x t+1 ) + const.Here omitted the parameters ϕ and ψ to avoid confusion. Using the definition of value function, above minimization can be expressed asEQUATIONwhereEQUATION= τ E π(x t:T ) [-log q(x T ) -log q(x t:T -1 |x T ) + log π(x t:T ) -log Z -log π(x t )] (19) = τ KL(π(x t:T )||q(x T )q(x t:T -1 |x T )) -τ log Z + τ H(x t ) (20) Note that x t is fixed in this optimization problem, and the optimization variable is x t+1 . Therefore the policy improvement step at time t is equivalent to min π(xt+1|xt) KL(π(x t:T )||q(x T )q(x t:T -1 |x T ))Therefore, the policy improvement step at each time step t results in the minimization of the KL divergence between the joint distribution, π(x 0:T ) and q(x 0:T ). 


Decision: 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.

Reviews: review: paper_summary: paper_summary: The authors introduce a maximum entropy inverse reinforcement learning (IRL) approach for
enhancing the sample quality of diffusion generative models, especially with limited generation
time steps. Named Diffusion by Maximum Entropy IRL (DxMI), the approach involves joint
training of a diffusion model and an energy-based model (EBM). The EBM provides the
estimated log density as a reward signal for the diffusion model, which is trained to maximize
both the reward from EBM and the entropy of generated samples. Additionally, the authors
propose Diffusion by Dynamic Programming (DxDP), a novel reinforcement learning algorithm
that optimizes diffusion model updates efficiently by transforming the problem into an optimal
control formulation. Empirical studies demonstrate that diffusion models fine-tuned with DxMI
can generate high-quality samples in as few as 4 to 10 steps and improve the stability of EBM
training dynamics, enhancing anomaly detection performance.
strength_weakness: strength_weakness: 1. Complexity of Implementation: The implementation of DxMI, particularly the joint training
of diffusion models and EBMs, might be complex and require significant computational
resources. The practical feasibility in various settings could be further elaborated.
2. Limited Scope of Experiments: 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.
3. Comparative Analysis: While the proposed methods show improvements, a more detailed
comparative analysis with state-of-the-art techniques, including training time and
computational cost comparisons, would strengthen the paper.
questions: questions: 1. Generalization to Complex Tasks: 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?
2. Training Time Comparison: What is the average training time for models using DxMI
compared to other generative model training techniques such as GANs or VAE-based
approaches?
3. Scalability: How scalable is the proposed DxMI approach when applied to larger datasets or
models with significantly more parameters? Are there any anticipated bottlenecks?
limitations: limitations: The authors have clearly presented the limitations in the paper.
paper_summary: paper_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.
strength_weakness: strength_weakness: Closely related work isn’t described in the introduction to better frame the contributions of this paper.
questions: questions: Is there potential for theoretical analysis or guarantees using this approach?

Are there visual differences in the generated outputs of different models that could be highlighted in 5.2?
limitations: limitations: Potential abuses using deep fakes are described, along with limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
 
- 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.  

-  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.
questions: 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: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

2. 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.
questions: 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.

2. 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.

3. 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?

4. 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?

5. 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?

[1] https://arxiv.org/abs/2403.06279

[2] https://arxiv.org/abs/2402.16359

[3] https://arxiv.org/abs/2305.16381

[4] https://arxiv.org/abs/2405.19673
limitations: limitations: The authors adequately addressed the limitations.

Comments: title: Thanks for your reply
Thank you for your reply.
Reply
Continued Response for Reviewer YBX1
comment: Thank you for addressing my concerns.
Thanks for the detailed responses. All of my concerns have been addressed properly.
Thanks for acknowledging the strength of this work. Also, thanks again for bringing up the interesting connection to existing work. We will make sure the connection is described well in the updated version.

Best regards,

Authors.
Thank you very much for taking the time to revisit the manuscript. We truly appreciate your reconsideration of the score. We will make sure your comments are reflected in the manuscript well and try to come up with some formal statements that we can make to ensure the stability of the objective function.

Best regards,  
Authors.
Thanks for the detailed and informative response. The strengths of this work and connections with relevant topics are made much clearer.
> Thus, our objective function always keep closer to than to, not invoking the instability questioned in the review comment.

Thanks for the clarification, I understand the proposal better after re-reading the paper and have updated the initial review accordingly.  

While my comment related to the objective was not really trying to claim a formal counter example for demonstrating pathologies of the proposed objective, it would be quite impressive if you could convert your intuition in the above response into a formal argument backing the proposal.
**Q. Theoretical analysis of DxMI**

* We agree that theoretical analysis of our MaxEnt IRL problem would be very interesting. However, direct analysis of DxMI may be highly challenging because our EBM and diffusion model comprise deep neural networks.
* Conducting theoretical analysis in a more restricted setting might be more feasible while still providing valuable insights. 
* To the best of our knowledge, there is not many theoretical results for MaxEnt IRL. Previous works [5,6] 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 similar analysis could be conducted on DxMI under suitable assumptions.
* Please consider that the main focus of this paper is to present a practical algorithm that is empirically scalable and effective. We believe providing theoretical analysis that rationalizes the empirical results in the paper is an important future work. 
* We will add a paragraph describing the theoretical results from [5,6] in our Related Work section. We will also mention the difficulty of theoretical analysis in our limitations section.


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

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

**Q. Can the DxMI approach be extended to other families of generative models, e.g., bridge-matching diffusion models?**

* Thanks for the interesting suggestion. We do believe the value-based learning presented in DxMI can be extended to other types of generative modeling problems, such as finding a bridge between distributions. 

**Q. 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?**

* Currently, we do not see a particular obstacle that prevents DxMI from scaling to larger datasets. We believe DxMI is at least as scalable as GANs, which have already demonstrated its feasiblity in a very high-dimensional datasets. As an empirical evidence, we will add LSUN Bedroom 256x256 experiments in the updated manuscript.

**Q. 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?**

* DxMI is not sensitive to the entropy regularization coefficient $\tau$, which can be safely set to 0.01 or 0.1. This insensitivity arises because the energy, which competes with the entropy, is regularized by the coefficient $\gamma$ to maintain a narrow range of values close to zero. 
* We provide the guide for hyperparameter tuning in Appendix B. 
* Probably the most critical hyperparameter is learning rates, for which we assign a larger value for the energy and a smaller value for the diffusion model. This is included in Appendix B.



Thank you once more for your insightful feedback. We hope our responses have addressed your concerns. Please feel free to reach out if you have any further questions or need additional information.

Best regards,  
Authors.