Paper: Upgrading VAE Training With Unlimited Data Plans Provided by Diffusion Models (ICLR-2024)

Reviewer comment (Theory):
The concern in the paper is the lack of any theoretical investigation into why these methods work. While this is difficult to do in general, it would be better if there was any mention of it at least with simplified models or any assumptions.

Author response:
Indeed a general theory would be difficult. In Section 4.2, we mention that data augmentation generates new data points by sampling $\mathbf x'\sim p_\text{aug}(\mathbf x'|\mathbf x)$ conditioned on a _single_ data point $\mathbf x$, whereas our method generates new data points conditioned on the _entire_ training set.

For a simplified model, consider a data set where the underlying data distribution $p_\text{data}(\mathbf x)$ is a mixture of Guassians. Traditional augmentation would be limited to adding small noise to each data point (ensuring that it stays within its mode) when $p_\text{data}(\mathbf x)$ is unknown. But if we have knowledge about $p_\text{data}(\mathbf x)$, e.g. if the mixture components lie on a regular grid, then we can use discrete translations as augmentations. By contrast, a diffusion model trained on the entire data set could detect such properties automatically.

Reviewer comment (Experiments):
Diffusion methods can also be interpreted as a form of augmentation. For example, you can use them to generate more examples in different resolutions or even "image-to-image" translation methods. Thus, the argument of treating augmentations visually different in Figure 1 is not very clear. While it is mentioned that augmentations are not "accurate", this will largely depend on the domain/data. Another aspect to consider is the computational time.

Author response:
Good point! _Conditional_ diffusion models are indeed similar to traditional data augmentation. However, we propose to use an _unconditional_ diffusion model, i.e., a model that samples random data without taking an input. We clarified this as a footnote in Section 4.1 in the paper.

Reviewer comment (Theory):
This paper is highly dependent on the accuracy of the diffusion model and while a theoretical investigation would further clarify this, it is worth mentioning. For example, in tabular data settings, diffusion models may not work so effectively or perhaps when data to train the diffusion model itself is scarce. Noting that diffusion models are quite data hungry, it is uncertain whether it is a good idea to use them to bootstrap in general.

Author response:
Thank you for bringing this up as it highlights why our results are nontrivial! We show that, whether or not a given model is data hungry is a relative statement. While the training sets were _too small_ for fitting good VAEs, as all baseline VAEs in our experiments strongly overfit, the same training sets are _large enough_ to fit useful diffusion models. One might obtain better diffusion models with larger training sets, but our diffusion models were evidently already at least good enough to be useful within our proposed method.

The fact that this observation is counterintuitive highlights why it should be communicated to a broader audience: one might think diffusion models require lots of data, and this is true if the goal is to generate high-quality samples. But our paper demonstrates a surprising new application for diffusion models in the regime where one does have too little data to train a good VAE.

Reviewer comment (Writing):
“Amortized inference, and adversarial robustness” - can you explain what these are and how VAEs are used? It may be obvious for amortized inference, however it is unclear for adversarial robustness. It would be a good idea to introduce those problems beforehand and explain how VAEs are used for them.

Author response:
Amortized inference is a fundamental principle of VAEs. We review this term briefly at the beginning of the paragraph “Amortization Gap” (Section 2). Expanding on this explanation, amortized inference refers to the fact that VAEs use a neural network (the “inference network”) to map data points to latent representations which speeds up training and deployment. One can remove the inference network and instead obtain latent representations by iteratively optimizing the ELBO for each data point (as we do for $q^*$ in Section 5.3), but this is very expensive.
The end of Section 2 discusses in detail the kind of adversarial robustness that we analyze. We are happy to provide further explanation if something is not clear.

Reviewer comment (Theory):
The statement “Therefore, Gg ≥ 0” is quite strong, however this is not guaranteed. While this may be the expected outcome, it could be the case that the test data is very simple to learn and the training data is complicated, resulting in a smaller ELBO over train. It would be better to amend this statement to say “we anticipate after training G_g \geq 0” as opposed to keeping it as a statement by itself.

Author response:
The statement $\mathcal G_\text{g} \geq 0$ relies on the assumption that training and test data come from the same distribution (see below Eq. 8). The remark below Eq. 8 discusses the case where this assumption is violated, and where the test data is “simpler” than the training data (i.e., it comes from a distribution with lower entropy). As discussed below Eq. 9, this can indeed lead to an inversion of the ELBOs (i.e., to $\mathcal G_\text{g} < 0$).

Reviewer comment (Novelty):
The concept presented in the paper is intriguing, but it is difficult to discern scenarios where the proposed method would be essential in the context of a VAE.

Author response:
We agree that diffusion models already excel at image generation and that there is little benefit to using VAEs for this task. However, while our method uses a generative model (diffusion model) to generate images, our evaluations do not focus on image generation but instead on density estimation and representation learning, which are domains where VAEs provide an advantage over diffusion models. Also, as discussed in the introduction, the goal of our method is mainly to improve the encoder of the VAE, which does not play a role in sample generation. For more clarifications how our evaluations are not about generative performance, please refer to item (1) in our “Overall response regarding sample quality and computational cost” posted above.

We agree, however, that our paper focuses on the image domain since this is the domain where both diffusion models and VAEs are best understood. While diffusion models on other data types are currently less explored, research in this area currently receives a lot of attention, and we expect that our method is compatible with any models proposed in the future.

Reviewer comment (Novelty):
While it is conceivable that there are applications in biological data or reinforcement learning where one might need to map an image to a latent representation and subsequently utilize this representation, the paper primarily focuses on image generation — a task for which the diffusion model already excels.

Author response:
We agree that diffusion models already excel at image generation and that there is little benefit to using VAEs for this task. However, while our method uses a generative model (diffusion model) to generate images, our evaluations do not focus on image generation but instead on density estimation and representation learning, which are domains where VAEs provide an advantage over diffusion models. Also, as discussed in the introduction, the goal of our method is mainly to improve the encoder of the VAE, which does not play a role in sample generation. For more clarifications how our evaluations are not about generative performance, please refer to item (1) in our “Overall response regarding sample quality and computational cost” posted above.

We agree, however, that our paper focuses on the image domain since this is the domain where both diffusion models and VAEs are best understood. While diffusion models on other data types are currently less explored, research in this area currently receives a lot of attention, and we expect that our method is compatible with any models proposed in the future.

Reviewer comment (Experiments):
The computational overhead introduced by the proposed method seems substantial, and it is unclear what advantages it offers to justify this overhead. Further evidence is needed demonstrating situations where a diffusion model sampler enhances the VAE's performance for specific, necessary tasks.

Author response:
We provide numbers for the computational cost in Appendix C.
We stress that our method only increases training time but not test time, which is more important in many applications. 
Also, the generated samples can be reused when training VAEs with different configurations for cross-validation. 
For further comments, please refer to item (2) in our “Overall response regarding sample quality and computational cost”.

Reviewer comment (Presentation):
Figure 2 is confusing: although the aim is to have the difference be the smallest, the way both are plotted makes it hard to see which method is doing the best given the scale. It is suggested to replot showing only the difference. Additionally, the reason for the zoom-in on the CIFAR-10 plot is unclear.

Author response:
Thank you for the suggestion. We added annotations to Figures 2, 3, 4, and 6 that highlight the gaps and show their numerical values. We hope that these additional annotations make it easy to see that our proposed method achieves the smallest gaps everywhere except for the robustness gap in BinaryMNIST and FashionMNIST (see Table 3 in Appendix B). The reason why we plot ELBO values rather than gaps in Figures 2, 3, 4, and 6 is because both the ELBO values and the gaps are important in practice, and plotting only the gaps would conceal the absolute values.

Reviewer comment (Presentation):
Figure 3 has a similar issue — plotting the gap directly could help, as it seems the augmented approach perhaps performs just as well, but this is hard to see from the current plots.

Author response:
Thank you for the suggestion. We added annotations to Figures 2, 3, 4, and 6 that highlight the gaps and show their numerical values. We hope that these additional annotations make it easy to see that our proposed method achieves the smallest gaps everywhere except for the robustness gap in BinaryMNIST and FashionMNIST (see Table 3 in Appendix B). The reason why we plot ELBO values rather than gaps in Figures 2, 3, 4, and 6 is because both the ELBO values and the gaps are important in practice, and plotting only the gaps would conceal the absolute values.

Reviewer comment (Evaluation):
The reviewer asks whether there are samples from the VAE plotted, or any sort of test metric comparing the proposed method with the baselines.

Author response:
Thank you very much for this important point! You are right that the plots were covering one seed. We re-ran CIFAR-10 experiments with three different random seeds and updated the Figures 2, 3, and 4 with means and standard deviation, which show consistent results across different random seeds. Experiments for the other data set are still running, and we will update the paper as soon as they finish.

Reviewer comment (Experiments):
According to figure5, it needs at least 2 times training data to be better than the normal training and 10 times data to reach the best performance. How are the computation costs comparing DMaaPx with aug tuned and aug naive, to let the model achieve the same/similar performance? (in terms of test ELBO).

Author response:
We provide computational costs in Appendix C. In brief, for any reasonable amount of samples from the diffusion models, the computational overhead of DMaaPx mainly comes from *training* the diffusion model and not from sampling from it. Therefore, we see little benefit in sampling less than about $10 \times |\mathcal D_\text{train}|$ data points, at which point sampling more has diminishing returns.

On the wider point of computational cost, note that our method only increases training time but not test time, which is more important in many applications. Also, samples can be reused when training VAEs with different configurations for cross-validation. For further comments, please refer to item (2) in our overall response above.

Reviewer comment (Experiments):
In your table 3, the method effectively reduced all three gaps the most for Binary MNIST and Fashion MNIST but not CIFAR10. Can you explain why it fails on CIFAR10? It looks to me that when the dataset becomes complex, your method loses its advantage. To convince the audience, you should include more datasets.

Author response:
Thank you very much for this comment! Unfortunately Table 3 was showing wrong values (the table was transposed). We apologize for this inconvenience and we uploaded a new version of the paper that fixes this issue. The correct version of the table shows that our method keeps its advantage and does not fail on CIFAR-10. To make it easier to see the gaps, we also added their numerical values to the right margins of Figures 2, 3, 4, and 6 in the main paper.

Reviewer comment (Experiments):
In figure 2, CIFAR10, the test data’s curves seem to keep increasing. For both augtuned and DMaaPx. If trained longer, what would be the trend?

Author response:
Thank you for this question. 
We agree that the ELBOs for data augmentation and our method are not fully converged for CIFAR-10.
The main message of this plot is that, for normal training, one should have stopped after about 200 epochs since the model already starts to overfit at this point. By contrast, we do not see any overfitting with both DMaaPx and traditional data augmentation for 5x more training epochs.

We are training the models for 2000 epochs on CIFAR-10 at the moment, and we will include the results and a plot as soon as they finish. We already observe in the existing plot that the generalization gap for data augmentation starts to widen when our proposed method still has a negligible gap. We therefore expect that the gap in our method either stays constant (as it does for BinaryMNIST and FashionMNIST) or is the last one to widen.

Reviewer comment (Evaluation):
It is known that diffusion models do not necessarily represent the whole training distribution, also in the image domain (as mentioned in the submission). Could it be that VAE better aligned to the subset of training data that was generated by diffusion? Is the distribution of ELBO on the test-set similar for both approaches?

Author response:
Thank you for this interesting question. We updated the paper and now provide an analysis of the distribution of the ELBO in Appendix G (_new_). We find that our method improves the ELBO on almost all (99.9 %) of the test points (Figure 9 (b)), and we couldn’t identify a subset of test data points where the improvement was particularly small or large.

Reviewer comment (Evaluation):
In the submission it is claimed that the main goal is to “improve the desirable functionalities of VAEs such as representation learning”, while the evaluation only tackles the quality of data modeling task. It’d be interesting to compare the samples quality between diffusion model from which the knowledge was distilled to the VAE.

Author response:
We believe that there is a misunderstanding, as our experiments in Sections 5.3 and 5.4 do evaluate representation learning. For VAEs, “representation learning” refers to the (output of) the encoder. Section 5.3 evaluates the amortization gap, where a small amortization gap corresponds to a good encoder. Section 5.4 evaluates adversarial robustness. While there are various forms of robustness in VAEs, our experiments specifically evaluate robustness of the encoder, i.e., we test whether an unnoticeable change of the input image can trick the encoder into outputting a semantically different representation. We recognize that this can be easy to miss since we evaluate the success of an attack back in image space.

However, based on reviewer feedback, we added more direct evaluations of practical downstream applications of VAEs in Appendix F (_new_), specifically representation learning tasks (classifying latent representations), reconstruction tasks, and (for completeness) also data generation tasks. While performance differences are indeed small, our proposed method achieves best performance for all tasks and all metrics except when measuring sample quality (SQ) by inception score (IS). This is consistent with our claim that the proposed method mainly fixes the encoder, which affects representation learning and reconstruction but not sample quality.

Reviewer comment (Evaluation):
The results suggest that distilling knowledge from diffusion models to VAEs introduce some performance gain for simple datasets like MNIST and FMnist, while for slightly more complex dataset - CIFAR10 the performance is comparable to the standard training with augmentations.

Author response:
We provide numbers for the computational cost in Appendix C. Note that our method only increases training time but not test time, which is more important in many applications. Also, samples can be reused when training VAEs with different configurations for cross-validation. For further comments, please refer to item (2) in our overall response above.

Reviewer comment (Evaluation):
Are the results presented in Figures 2 and 6 statistically significant? The differences in some cases are extremely small, while it seems like results from a single run (judging from the background plot).

Author response:
Thank you for bringing this up. You are right! We indeed omitted data augmentation when training the diffusion model (see details in Appendix C) to not let any augmentation leak to our generated samples.