Learning Energy-Based Prior Model with Diffusion-Amortized MCMC

Anonymous Code Submission

We thank the reviewer for your prompt reply! We are pleased to hear that our new results are very encouraging. We have asked our AC for details about anonymous code submission, since we are only allowed to submit the code link to AC according to the instructions this year. We have prepared the code and will submit the link immediately once permitted. Here we provide more details about our code submission to help you reproduce and validate our results. - Environment Specification - We use Pytorch to train our models. We did not use other third-party python packages for implementation. Please feel free to let us know if there are any dependency issues. We will be more than happy to help. - Version of packages used in the toy example experiments: - Python == 3.9.2 - Pytorch == 1.10.0 - numpy == 1.21.2 - matplotlib == 3.3.4 - Code Structure We have added comments to most functions and code blocks in our code. The code files should have the following structure: ``` toy_code/ src/ diffusion_helper_func.py diffusion_net.py toy_example.py ``` - `toy_example.py`is the main file. It includes the training algorithm, data generation pipeline and visualization functions. - `src/diffusion_net.py` contains the detailed network structure of the diffusion network. - `src/diffusion_helper_func.py` contains helper function for implementing the denoising diffusion process. - How to Run To train the model on the toy example, you can simply run the following command in the `toy_example` folder. ``` CUDA_VISIBLE_DEVICES=<DEVICE_ID> python toy_example.py --seed <RANDOM_SEED_TO_SPECIFY> ``` For example, ``` CUDA_VISIBLE_DEVICES=0 python toy_example.py --seed 0 ``` - Here `--seed` argument specifies the random seed, which basically decides the ground-truth posterior distribution. The script will automatically generate a `logs/toy/<TIMESTAMP>` folder in the `toy_code` folder, where `<TIMESTAMP>` indicates the time you started this training process. - There will be two automatically created additional folders in `logs/toy/<TIMESTAMP>` once running the script: i) `ckpt` which saves the trained weights and ii) `viz` folder that saves the visualization of ground-truth and learned posterior distributions. The image file names are `<ITERATION>_lang_post_Q` and `<ITERATION>_lang_post_gt`, which indicates visualization of the learned distribution and the ground-truth distribution respectively. These visualization results are saved every 100 iterations. - Important Tips about Training - For most random seeds, we observed that our learned sampler could achieve decent approximation of the ground-truth posterior distributions obtained by long-run langevin dynamics within 300-3000 training iterations. This would take from several minutes to an hour or so on a NVIDIA RTX A6000 GPU. The training process takes ~2GB GPU memory. It is possible that there are some extreme cases where longer training iterations are needed to produce decent results. - For some random seeds, the default 1000-step langevin dynamics for sampling ground-truth posterior distribution might not converge. You may consider using 2000 or more steps by modifying the `g_l_steps` argument in the `sample_langevin_post_z` function at L277 in the `toy_example.py`. One possible sign is that the `g_loss (avg) Q` (reconstruction error obtained by learned posterior samples) is significantly lower than `g_loss (avg) L` (reconstruction error obtained by langevin dynamics samples). Finally, we kindly ask the reviewer to not distribute our code, since we have not officially published our work. Thank you again for your time for evaluating our work, please feel free to let us know if you have any additional questions/comments/concerns.

Anonymous Code Submitted

Thank you for your patience. We have submitted the anonymous code for reproducing the results for the toy example to our AC. Please feel free to let us know if you have any problems with the code.

Thanks for your reply

Thank you for your prompt response. We would like to point out that although DDPM estimates the gradient of the log density of *noise-perturbed versions of the target density*, this gradient estimate **does not** translate to reliable estimate of the *target density (without noise)* via re-parametrization. For a given noise level $\sigma > 0$, the denoising score matching or DDPM objective is not a consistent objective for learning the underlying target density without noise. Annealing the noise level $\sigma \to 0$ can mitigate this problem, but the gradient estimation is only reliable in the immediate vicinity of the modes of the target distribution, where the density is high. For low-density regions of the target distribution, the objective may not have enough evidence to estimate score functions accurately, due to the lack of samples [1]. In contrast, EBM is typically learned through Maximum Likelihood Estimation, which theoretically is the most accurate estimate in terms of asymptotic variance [2]. This formulation together with its learning algorithm produces much more reliable explicit (unnormalized) density estimation. From this point of view, these two models are different. In our previous response, we discussed the importance of energy-based modeling, i.e., learning the density of target distribution explicitly. We hope the key differences between EBM and DDPM mentioned above in this response further clarifies the necessity of learning the explicit density of the target distribution. Specifically in the set-up of our submission, we care about modeling the latent space. The latent space EBM provides explicit prior probability density $p_\alpha(z)$ for modeling latent variables $z$, and consequently defines explicit posterior density $p_\theta(z|x) \propto p_\alpha(z)p_\beta(x|z)$ for posterior inference upon which we build our approximate MLE learning algorithm. Simply plugging a DDPM into the latent space does not provide us with this well-defined MLE learning framework and could lead to problematic learning algorithms. Finally, we would like to kindly refer to [3] for a comprehensive discussion about EBM and its connection with DDPM. We hope our response helps to address your concerns and explains why latent space EBM is considered in this work. Please feel free to let us know if you have any additional questions/comments/concerns. [1] Yang Song, and Stefano Ermon. Generative Modeling by Estimating Gradients of the Data Distribution. NeurIPS 2019. [2] Peter J. Bickel and Kjell A. Doksum, Mathematical Statistics: Basic Ideas and Selected Topics, Volume I. [3] Yang Song and Diederik P. Kingma. How to Train Your Energy-Based Models, arXiv:2101.03288

Thank you for your reply

Thank you for your further response. We provide point-to-point replies below. - > If the results in 2023 cannot show improvements over these classic ones, the proposed method might not make a real progress in learning EBMs. We respectfully disagree with this statement for the following reasons. - **The proposed method and the mentioned previous methods work for different models in completely different set-ups.** We would like to continue our discussion in our previous response to further clarify that these methods are essentially different. As mentioned in sec. 1 L16-28 in the main text, we specifically consider *learning an EBM in the latent space for latent variables $z$ (LEBM) as an informative prior* for the generator network in this work. Although our focus is also on learning an EBM, the LEBM stands on the generator to capture the regularity of its latent variables $z$, while the main generation process is done by the generator network $g$. The learning process is built upon the posterior inference of $z$ to formulate an EM-like approximate MLE learning algorithm. In this formulation, the generator is directly supervised by the reconstruction error of the observed data $x$. In contrast, in the set-up of [a,b], the EBM resides explicitly in the data space for $x$ and serves as a refiner of the initial output of the generator network. The learning process does not involve inference of $z$ or modeling of the latent space. In [a], the refined output is further used as the training sample of the generator, while in [b] the generator is trained by optimizing the learned energy score function. In sum, the key differences between our set-up and [a,b] include, at least, the following ones: - Although we all consider learning EBMs, the roles of EBMs are completely different. In our work, the light-weight LEBM is used as a flexible yet powerful prior for the generator, while in [a,b] the EBMs are used as a refiner or the main generation model; the learning processes in [a,b] do not involve inference of $z$ or modeling of the latent space. - Although both our work and [a,b] involve learning the generator, we can see in [a,b] the generator is supervised by a different objective that incorporates the training signal from the data-space EBM, rather than directly optimizing the reconstruction error of the original observed data $x$, i.e., $\log p(x | z)$ in our set-up. Therefore, we think the differences of FID scores of these methods **may not** serve as valid indicators, since the methods work for different models in completely different set-ups and are not directly comparable. - **The proposed method and the mentioned previous methods are not direct competitors, but are actually complementary to each other.** Following our discussion above, we would like to add that a more reasonable comparison would be incorporating the LEBM in the latent space of the generator used in [a,b], and jointly train the LEBM, generator and the data-space EBM to see whether the composed methods achieve higher performance. These methods are not direct competitors, but can be complementary to each other to further improve previous methods. In our initial submission here, we prefer to keep our model and learning method clean and simple, without involving extra networks and learned computations. We are happy to explore these directions in future works. - > One may doubt whether the proposed method can work or not when using a deeper architecture. We would like to point out that in sec. 4.1 L215-228 and tab. 1 and fig. 2 CelebA-HQ columns we have explicitly provided the results on higher dimension data (256x256) to show that our method can scale up with deeper architecture (supp. C tab. 2). To further verify the effectiveness of this method on more advanced architecture, in sec. 4.1 *GAN inversion* paragraph L229-254 we have provided further results using the StyleGAN network as the generator. Therefore, we believe we have provided positive evidence that our method can effectively work with deeper architectures to scale up. - > Suggest that the authors show results with a bit deeper advanced network, which is not difficult for experiments. Thank you for the constructive suggestion. Due to the time limit, we were only able to experiment on slightly deeper generators and report very preliminary results on CIFAR-10. Here we only add `conv3x3` layers to the original one; `+N` shows the number of additional layers and `w/ res` means adding residual connection to the conv layers. Please note that these models have not reached their best performances, but have already shown some initial improvements over the original model. ||+0|+1|+2|+4 w/ res| |:--:|:--:|:--:|:--:|:--:| |FID|57.72|57.63|55.50|54.59| We hope our response helps to address your concerns. Please feel free to let us know if you have any additional questions.