Paper: Learning Energy-Based Models by Cooperative Diffusion Recovery Likelihood (ICLR-2024)

Reviewer comment (Presentation):
The related work section is placed in the Appendix, which seems quite odd. As a result, the previous work on cooperative training is sort of "hidden" in the main paper. This should be described before Section 3.3.

Author response:
Thank you for your advice. Our arrangement was primarily influenced by the page limit constraints. We believed it would be more beneficial to allocate more space to experimental results in the main body of the paper, while relegating a detailed discussion of related work to the appendix. Following your suggestions, we have expanded our discussion in Section 3.2 to further explore the connection between our work and cooperative training. Additionally, we guide readers to the appendix for more comprehensive discussions.

Reviewer comment (Presentation):
It might be helpful to illustrate or give a schematic overview of the proposed approach in some kind of figure (e.g., a sequence of noise levels, one EBM and initializer model per level). This could help illustrate the general idea.

Author response:
Thank you for your advice. We have added an overall figure to the revised version of our paper. Please see Figure 5 (a) and (b) in the modified version of our paper.

Reviewer comment (Experiments):
For the experiment in Figure 7, it would be helpful to report FID scores for (a) and (b), i.e., for K=0 and K=15 steps of Langevin refinement. This experiment could perhaps also be repeated for K=1, 2, 4, 8 steps of Langevin refinement.

Author response:
Certainly. In accordance with your advice, we tested the FID for Figure 7 (current Figure 8) using various Langevin steps. The outcomes are presented in the table below. It's important to note that for this particular setting, we merely reduced the number of Langevin steps without adjusting their step size as detailed in Section 4.2.

Table: FID scores for ImageNet32 unconditional generation with different number of Langevin steps

|  Total MCMC steps |  FID |
| ----------------- |------|
|K=0| 38.94|
|K=1| 34.97|
|K=2| 30.89|
|K=4| 24.75|
|K=8| 16.21|
|K=15| 9.35|

Reviewer comment (Writing):
Section 3.1: "...constrains the conditional energy landscape to be localized around y_t": Should y_t be x_{t+1}?

Author response:
Thank you very much for helping us find the typos, we have modified in the current version of our paper.

Reviewer comment (Writing):
Abstract: "noisy versons": versons → versions.

Author response:
Thank you very much for helping us find the typos, we have modified in the current version of our paper.

Reviewer comment (Writing):
Section 3.2: "which may again requires MCMC sampling": requires → require?

Author response:
Thank you very much for helping us find the typos, we have modified in the current version of our paper.

Reviewer comment (Writing):
Algorithm 1: noise level → noise levels, sampling step → sampling steps.

Author response:
Thank you very much for helping us find the typos, we have modified in the current version of our paper.

Reviewer comment (Writing):
Algorithm 2: sampling steps L → sampling steps K? (also some inconsistency with this in the Appendix).

Author response:
Thank you very much for helping us find the typos, we have modified in the current version of our paper.

Reviewer comment (Writing):
Section 4.4: CDLR strong → CDLR achieves strong?

Author response:
Thank you very much for helping us find the typos, we have modified in the current version of our paper.

Reviewer comment (Writing):
Section G.5: doesn't gives → doesn't give (or "does not give").

Author response:
Thank you very much for helping us find the typos, we have modified in the current version of our paper.

Reviewer comment (Writing):
The clarity should be improved. For example, consider the relationships between the proposed method and related works.

Author response:
Thank you for your question. The goal of our work is to push the study of Energy Based Model. We would like to emphasize that, while the study of EBM draws more and more attentions recently, unlike the popular diffusion model, the scaling up of training for EBM is still a largely uncharted territory. The primary challenge lies in devising a more efficient sampling strategy. In pursuit of this goal, many individual efforts have been made. For instance, DRL[1] employs a series of noise levels, and cooperative learning[2, 3] introduces a generator or initializer. Though these methods have led to improvements in their own contexts, they are often regarded as separate techniques, and their individual performance may exhibit noticeable gaps when compared to other generative frameworks. We are the first in integrating these elements into a unified framework, and we are the first to successfully demonstrate that by combining these two techniques, the performance of EBM can be substantially enhanced.

Furthermore, we make many practical design choices related to noise scheduling, network architecture, noise variance reduction. These choices helps the training of EBM and we carefully conducted ablation studies on each component (please see our Appendix G). We believe these designs can greatly benefit future works in the study of EBM.

Besides, we show that EBM trained with CDRL can benefit from simple sampling adjustment techniques to greatly reduce its sampling steps. This further improves the sampling efficiency of EBM as a generative model and might inspire future works in similar directions.

Thus, we believe that our work has made great contributions. Our model and algorithm can pave the way for the scaling up of EBM training on complex distributions.

[1] Learning Energy-Based Models by Diffusion Recovery Likelihood

[2] Cooperative training of descriptor and generator networks.

[3] A tale of two flows: Cooperative learning of langevin flow and normalizing flow toward energy-based model.

Reviewer comment (Novelty):
The contribution is kind of incremental. Based on the Diffusion Recovery Likelihood (DRL) (Gao et al., 2021) in Sec. 3.1, the authors proposed in Sec. 3.2 a new "initializer model" to amortize the expensive MCMC sampling.

Author response:
Thank you for your question. The goal of our work is to push the study of Energy Based Model. We would like to emphasize that, while the study of EBM draws more and more attentions recently, unlike the popular diffusion model, the scaling up of training for EBM is still a largely uncharted territory. The primary challenge lies in devising a more efficient sampling strategy. In pursuit of this goal, many individual efforts have been made. For instance, DRL[1] employs a series of noise levels, and cooperative learning[2, 3] introduces a generator or initializer. Though these methods have led to improvements in their own contexts, they are often regarded as separate techniques, and their individual performance may exhibit noticeable gaps when compared to other generative frameworks. We are the first in integrating these elements into a unified framework, and we are the first to successfully demonstrate that by combining these two techniques, the performance of EBM can be substantially enhanced.

Furthermore, we make many practical design choices related to noise scheduling, network architecture, noise variance reduction. These choices helps the training of EBM and we carefully conducted ablation studies on each component (please see our Appendix G). We believe these designs can greatly benefit future works in the study of EBM.

Besides, we show that EBM trained with CDRL can benefit from simple sampling adjustment techniques to greatly reduce its sampling steps. This further improves the sampling efficiency of EBM as a generative model and might inspire future works in similar directions.

Thus, we believe that our work has made great contributions. Our model and algorithm can pave the way for the scaling up of EBM training on complex distributions.

[1] Learning Energy-Based Models by Diffusion Recovery Likelihood

[2] Cooperative training of descriptor and generator networks.

[3] A tale of two flows: Cooperative learning of langevin flow and normalizing flow toward energy-based model.

Reviewer comment (Theory):
In the paragraph following Eq. (5), why "maximizing recovery likelihood still guarantees an unbiased estimator of the true parameters of the marginal distribution of the data?"

Author response:
The recovery likelihood formulation uses EBM to model the marginal distribution of the data and uses classic MLE loss to optimize the parameter. According to the classical analysis of MLE, the point estimate given by optimizing recovery likelihood model is an unbiased estimator of the true parameter. For more details, please check the appendix A.2 in [1].

[1]Learning Energy-Based Models by Diffusion Recovery Likelihood

Reviewer comment (Theory):
Eqs. (6) and (8) are quite like the formula of the conventional diffusion model. Please elaborate on the relationships between them.

Author response:
Inspired by the diffusion model, DRL [1] introduces a method to learn a series of Energy-Based Models (EBMs) across various noise levels denoted as t. Our research expands on DRL by incorporating an initializer at each noise level, enhancing the EBM learning process. Like the diffusion model, our initializer is based on Gaussian distributions. A notable difference, however, is that unlike diffusion models that let the Unet to learn from the data, in our cooperative algorithm, the U-Net is trained using the revised sample of the EBM, i.e. $\tilde{y_t}$ not merely the original data augmented with Gaussian noise; it is the product of a Langevin sampling process, guided by the energy function. This enables the unet to armortized the Langevin sampling process. We delve deeper into this difference and present an ablation study in Appendix G.2, offering further insights into this aspect.

Reviewer comment (Theory):
It seems that Eq. (9) indicates a deterministic noising process, right? If so, why "there’s still variance for xt given x0 and xt+1?"

Author response:
With the noise variance reduction technique, at each iteration, we use the same noise to inject $x_t$ and $x_{t+1}$. However, our model contains both an EBM and its initializer. During the sampling process, at each noise level, the initializer proposes an initialization and the EBM modifies it with MCMC. In both the initial proposal stage and MCMC stage, random noise is injected. During the training process, as discussed in the previous question, our initializer learns from the revised sample of the MCMC process instead of directly learning from the data. This implies that $\tilde{y_t}$ is not completely deterministic given $x_0$ and $x_{t+1}$. On the other hand, we define the EBM on the marginal distribution of $x_t$ (more precisely $y_t$), so updating the EBM only requires drawing samples from the marginal distribution $p(x_t)$, not the joint distribution $p(x_t, x_{t+1})$. The noise variance reduction technique does not alter the marginal distributions of $x_t$ and $x_{t+1}$, thereby ensuring it does not introduce extra bias into EBM training.

Reviewer comment (Writing):
How to interpret Eq. (11), especially the $\tilde p_{\theta}$?

Author response:
Equation (11) is the equation of classifier-free guidance [2]. The classifier-free guidance technique aims to improve the quality of samples in conditional generation, building upon the concept of classifier guidance [1]. In this method, an additional classifier $p(y|x)$ is used in conjunction with eiter a conditional model $p(x|y)$ or an unconditional model $p(x)$  to enhance the estimation of conditional generation scores. For detailed insights, please refer to section 4 in [1]. Further, [2] proposes that instead of employing an additional classifier $p(y|x)$, it might be more effective to estimate the guided score by analyzing the difference between the conditional model $p(x|y)$ and the unconditional model $p(x)$. Equation (11) is derived following the methodologies in [1, 2]. Under our notation system, in Equation 11, $f_\theta (y_t; c, t) = \log p_\theta(y_t|c)$ denotes the score for conditional generation, and $f_\theta (y_t; t) = \log p_\theta(y_t)$ represents the score for the unconditional model. Instead of directly using $\log p_\theta(y_t|c)$ in the sampling process, we utilize an enhanced version $\tilde{p_\theta}(y_t|c)$, which combines aspects of both the conditional and unconditional models, in a manner similar to [2].

[1] Diffusion Models Beat GANs on Image Synthesis

[2] Classifier-Free Diffusion Guidance

Reviewer comment (Novelty):
The proposed method is a straightforward combination of two known techniques, DRL [1] and MCMC amortization [2,3]. While the simplicity of the idea is practically appealing, the idea lacks theoretical novelty. Moreover, the tasks demonstrated in the paper are already well-explored in the diffusion model literature, and the algorithms for the tasks are straightforward extensions of diffusion-based ones, as the score is just the gradient of the energy. For instance, classifier-free guidance is explored in [4], and compositional generation is explored in [5].

Author response:
Thank you for your question. The primary aim and contribution of our work is advancing the study of EBM. Although EBMs are gathering increasing attention, scaling up their training remains largely unexplored. The key challenge lies in developing a more efficient sampling strategy. While DRL [1] made strides by using multiple noise levels to ease sampling, it established a strong baseline but still falls short when compared to other generative frameworks. Narrowing this gap has proven challenging, raising questions about the potential of EBMs as generative models. While amortized inference is a known technique in prior studies [2, 3], our work is the first to show that integrating amortized sampling with the DRL framework significantly enhances EBM performance. Concurrent effort [3], explores amortizing MCMC in a EBMs fitted within the latent space of a very simple VAE model. The EBM structure as well as the VAE latent space are not very complex to model. Their focus is different and they require running a 100-step diffusion model alongside 30 or 50 steps of MCMC. In contrast, our model addresses more complex data distributions with just a single-step initialization proposal and 15-step MCMC sampling, achieving much more superior sample quality compared to [3]. This underscores that our scaling-up effort is far from trivial. Regarding tasks like classifier guidance and compositional generation, we do not claim to be the first to explore these areas. Instead, our focus is to demonstrate that our model, while delivering strong performance, is compatible with these techniques, broadening its applicability in diverse applications.

Furthermore, we make many practical design choices related to noise scheduling, network architecture, noise variance reduction. These choices helps the training of EBM and we carefully conducted ablation studies on each component (please our Appendix G). We believe these designs can greatly benefit future works in the study of EBM.

Besides, we show that EBM trained with CDRL can benefit from simple sampling adjustment techniques to greatly reduce its sampling steps. This further improves the sampling efficiency of EBM as a generative model and might inspire future works in similar directions.

Thus, we believe that our work has made great contributions. Our model and algorithm can pave the way for the scaling up of EBM training on complex distributions.

[1] Learning Energy-Based Models by Diffusion Recovery Likelihood, Gao et al., ICLR, 2021.

[2] Approximate Inference with Amortized MCMC, Li et al., 2017.

[3] Learning Energy-Based Prior Model with Diffusion-Amortized MCMC, Yu et al., NeurIPS, 2023.

[4] Classifier-Free Diffusion Guidance, Ho et al., 2022.

[5] Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC, Du et al., ICML, 2023.

Reviewer comment (Evaluation):
The paper lacks a comparison of inference time for CDRL and the baselines.

Author response:
Please refer to the below comparison of our CDRL with other strong EBM baselines. Our CDRL, employing 90 steps (6 x 15), achieves high-quality generation with a relatively short sampling time. Furthermore, as discussed in Section 4.2 of our paper, a simple adjustment in the sampling process can be applied to further reduce the sampling duration without significantly impacting the quality of the generated samples. This comparison is presented in Table 9 of our Appendix.


Table: Time for generating 100 Cifar10 samples

| Model Name |  Total MCMC steps | Time (s) | FID |
|:----------:| ----------------- |------|-----|
|CoopFlow| 30 | 2.50|15.80|
|VAEBM| 12 |21.3|12.16|
|DRL| 6 x 30 = 180 |23.9|9.58|
|CDRL| 6 x 15 = 90 |12.2|4.31|
|CDRL (8 steps)| 6 x 8 = 48 |6.48|4.58|
|CDRL (5 steps)| 6 x 5 = 30 |4.14|5.37|
|CDRL (3 steps)| 6 x 3 = 18 |2.60|9.67|

Reviewer comment (Experiments):
The paper lacks a comparison of training cost (e.g., required VRAM) for CDRL and the baselines. I expect CDRL training is more expensive than diffusion or DRL training, as the former requires two networks while the latter only one.

Author response:
As detailed in Appendix A2, we utilized 8 A100 GPUs to train our model over 400k iterations, spanning approximately six days. Compared to DRL, although CDRL includes an additional initializer model. As outlined in Appendix A.1., we only employ a small U-Net as the initializer, resulting in the MCMC sampling process predominantly dictating the sampling time. Notably, with CDRL using only half the MCMC sampling steps at each noise level, it requires fewer computational resources than DRL.

More specifically, DRL’s standard configuration, which involves a 6-noise schedule and 30 Langevin steps, demands 240,000 iterations. In contrast, CDRL, utilizing half the Langevin steps and extending over 400,000 iterations, reduces training time to just 80% of that required for DRL. It's also important to note that even with increased computational investment in DRL, such as longer training durations or more MCMC steps, significant performance enhancements remain elusive (as demonstrated in Table 2 of DRL). In comparison, CDRL achieves superior performance with less computational expenditure, demonstrating its efficiency advantage over DRL.

Reviewer comment (Experiments):
The paper lacks results on higher resolution (e.g., 256x256) images.

Author response:
In terms of handling high-resolution images, the recent trend in generative modeling involves either utilizing the latent space of a VAE, as in Latent Diffusion[5], or initially generating a low-resolution image and then scaling up, as demonstrated by techniques like Imagen. This process might reduce the modeled space to dimensions such as 32 × 32 or 64×64, aligning with the resolution on which we conducted our experiments. Following your suggestion, we have carried and extra experiment on CelebA-HQ dataset (256 x 256) in the latent space of latent diffusion.  Our preliminary results have achieved an FID score of 10.74. We have included the new results in our modifed paper (see Figure 9 and Table 7 in the Appendix). In the following Table, we compare our preliminary results with some baselines on this dataset.

Table FID score for CelebA-HQ (256 x 256) dataset
|Model | FID score|
|------| ---------|
|GLOW[3]| 68.93|
|VAEBM[1] |20.38|
|ATEBM [2]|17.31 |
|VQGAN+Transformer[4]| 10.2|
|LDM[5]|5.11|
|CDRL(ours)  |10.74|

[1] VAEBM: A Symbiosis between Variational Autoencoders and Energy-based Models

[2] Learning energy-based models with adversarial training

[3] Glow: Generative Flow with Invertible 1x1 Convolutions

[4] Taming Transformers for High-Resolution Image Synthesis

[5] High-Resolution Image Synthesis with Latent Diffusion Models

Reviewer comment (Evaluation):
The empirical performance of CDRL is good but not very strong. It is clear that CDRL is an improvement over DRL, but it still falls behind other models in multiple tasks. The initializer is very similar to DDPM, but CDRL's unconditional CIFAR-10 FID is worse than DDPM (Table 1). Also in Table 3, CDRL is outperformed by DDPM++.

Author response:
Thank you for your question. The primary objective of our research is to push forward the study of EBM. As we discuss in our paper, EBMs have gathered increasing interest in recent research. However, unlike the widely-studied diffusion models, scaling up EBMs remains an area less explored. Prior studies like DRL have adopted noise level concepts from diffusion models into EBM, setting a strong baseline. Despite this, a significant performance gap remains when compared to diffusion models. Our work builds upon DRL, aiming to narrow this gap.

Our initializer, while drawing inspiration from DDPM's design, serves a distinct purpose. We employ the initializer only for providing an initial rough estimate at each step, with the subsequent generation process guided by MCMC under the EBM framework. This approach is illustrated in (current) Figure 8 of our Appendix. As a result, we utilize a smaller initializer, which is also cooperatively learned from the revised samples of EBM. The difference between learning directly from data and this cooperative learning approach is detailed in Appendix G.2.

Furthermore, we have implemented other practical strategies for EBM training. Our findings show that the sampling efficiency of CDRL can be greatly improved using simple sampling techniques. While we acknowledge the successes of diffusion models and recognize the current performance disparity with energy-based models, we still believe in the potential of EBMs as powerful generative models and in estimating the likelihood of complex distributions. Our work significantly advances the generative capabilities of EBM.

Reviewer comment (Theory):
It would be nice if the authors could comment on the consistency of the model after adding the initializer model. Would it alter the consistency?

Author response:
In our framework, the EBM is estimated by maximimizing recovery likelihood, where the samples are generated by starting a few MCMC sampling steps from samples given by the initializer model. It remains an unbiased estimator of the true model parameters, which means given enough data, a rich enough EBM and exact synthesis, it learns $p_\theta = p_{\rm data}$. On the other hand, the initializer model is learned by maximizing log-likelihood, treating the samples from the EBM as the training data. Therefore, it is an unbiased estimator of the EBM model distribution $p_\theta$.