Maximum Entropy Inverse Reinforcement Learning of Diffusion Models with Energy-Based Models

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.

Dear Reviewer YBX1,

We deeply appreciate your thorough and constructive comments. We will do our best to answer your questions.

Q. What is the connection to KL-regularized RL fine-tuning, particularly in papers such as [1-4]?

Thanks for pointing out an interesting connection. The sampler update step of DxMI is indeed closely related to KL-regularized RL.

• We will augment our manuscript to discuss this connection in detail.
  • 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.
  • 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.
• 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.
  • 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$.
  • Second, DxMI employs a novel value-based algorithm for updating a diffusion model.
  • Third, while most KL-regularized RL works assume that the reward function is known, DxMI simultaneously learns the reward from data.
• 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.

[1] https://arxiv.org/abs/2403.06279 Tang. Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond. 2024.

[2] https://arxiv.org/abs/2402.16359 Uehara et al. Feedback Efficient Online Fine-Tuning of Diffusion Models. 2024.

[3] https://arxiv.org/abs/2305.16381 Fan et al. DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models. 2023.

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

Q. The comparison to prior diffusion model acceleration methods is somewhat limited.

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

Q. A more comprehensive evaluation across different speed-quality tradeoffs would be essential.

• 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.
• Qualitatively, the best trade-off for DxMI is achieved in the mid-range of $T$, from 4 to 10.
  • If $T$ is too small, the sampler is less capable, and the data processing inequality becomes less tight, making our MaxEnt regularization less effective.
  • If $T$ is too large, the sampler's capability increases, but the value function learning becomes more challenging.

Q. More discussion of DxMI's relative strengths and weaknesses compared to other approaches would be essential.

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

• Strengths
  • 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).
  • The dynamic programming-based in DxMI is more effective than the policy gradient-based algorithm (e.g., SFT-PG).
  • DxMI produces an EBM as a byproduct, which can be used in other applications such as anomaly detection or transfer learning.
• Weaknesses
  • When $T=1$, DxMI reduces to GAN, offering no additional advantage.
  • DxMI does not offer the flexibility of using a different value of $T$ during the test time.

---

Due to the character limit, we will continue answering your question in the comment.

Dear Reviewer yjrf,

We appreicate your time and effort in reviewing our work. Here, we are happy to address your concerns and questions.

Q. Is the objective function stable?

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

We would like to clarify the critical misunderstandings regarding the objective function.

Let 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)$.

• 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$.
  • It is a misunderstanding that our objective function $L(q,\pi)$ has a minimum at 0, where $KL(p||q) = KL(\pi||q)$.
  • 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.
  • Thus, our objective function always keep $q$ closer to $p$ than to $\pi$, not invoking the instability questioned in the review comment.

Q. Is learning possible when $\pi$ is close to $p$?

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.

Moreover, since for the pre-trained checkpoints, π is likely close to p, the energy objective looks like it has minimal training signal.

• 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.
• 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.
  • 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$.
• 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$.

Q. DxMI still uses a diffusion model as a starting point.

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.

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

Q. The proposed algorithm has multiple unstable procedures.

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.

• We understand the proposed algorithm can be seemingly complicated.
  • 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.
• However, to the best of our knowledge, there is no empirical or theoretical evidence of instability.
  • 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.
• Empirically, we found the algorithm much more stable than MCMC-based EBM training algorithms, which occasionally diverge for no reason.
• If you have a particular concern regarding the algorithm's stability, we are happy to discuss it.

Q. In Algorithm 1, Line 5, are the x and x_T completely independent samples with no relation to each other?

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

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

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.

Dear Reviewer Nqik,

Thank you for taking the time to review our work. We deeply value your feedback and are happy to address your questions.

Q. Is theoretical analysis possible?

Is there potential for theoretical analysis or guarantees using this approach?

Thanks for bringing up an important point.

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

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

Q. Related works are not described in the introduction.

Closely related work isn’t described in the introduction to better frame the contributions of this paper.

• We had to defer our discussion on prior works to the Related Work section (Section 6) due to the page limitation.
• 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.

Q. Are there visual differences in the generated outputs of different models?

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

Again, we thank the reviewer for acknowledging the value of our work. If you have any further questions, we are happy to address them.

Best regards,
Authors.

Dear Reviewer EhWr,

We appreciate the comprehensive feedback on our manuscript. All the comments and questions raised have been considered below.

Q. Complexity of Implementation

The implementation of DxMI, particularly the joint training of diffusion models and EBMs, might be complex and require significant computational resources.

• DxMI may seem complicated, but in fact its complexity is no larger than GANs and actor-critic RL methods.
  • 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.
  • 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.

[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

Q. More Experiments

The practical feasibility in various settings could be further elaborated.

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.

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?

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

Q. Comparative Analysis

While the proposed methods show improvements, a more detailed comparative analysis with state-of-the-art techniques.

• Please find that the current manuscript provides comparative analysis on Section 6 Related Work.
• We will augment our comparative analysis by adding additional paragraphs in the introduction and sections that describe the proposed method.

Q. Average Training Time

What is the average training time for models using DxMI compared to other generative model training techniques such as GANs or VAE-based approaches?

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

[2] Brock et al. Large scale GAN training for high fidelity natural image synthesis. ICLR 2019. https://arxiv.org/abs/1809.11096

Q. Scalability

Scalability: How scalable is the proposed DxMI approach when applied to larger datasets or models with significantly more parameters? Are there any anticipated bottlenecks?

• 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].

Again, we thank you for providing valuable comments. Do not hesitate to let us know if you have further questions.

Best wishes,
Authors.