Upgrading VAE Training With Unlimited Data Plans Provided by Diffusion Models

Thank you for your detailed feedback! We are glad that you think the “paper is easy to follow” and the three gaps/metrics are “well defined and nicely motivated”. We address your questions below, and we updated the PDF accordingly.

Q1 - [...] lack of any theoretical investigation [...]

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.

Q2 - [diffusion models can be interpreted as data augmentation, e.g., in image-to-image translation]

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.

Q3 - [...] diffusion models are quite data hungry

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.

Q4 - [computational time]

Note that our method only increases training time but not test time, which is more important in many applications. For further comments, please refer to item (2) in our “Overall response” above.

Q5 - [...] [results are] highly dependent on the accuracy of the diffusion model [...]

Indeed, as we discuss in the last paragraph of Section 4.1, diffusion models on other data types than images are less explored, and we might not have accurate models. But 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.

Q6 - [explain amortized inference and adversarial robustness in the context of VAEs]

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.

Q7 - [$\mathcal G_\text{g} \geq 0$ is not guaranteed]

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

Q8 - [...] other types of generative models to bootstrap a VAE?

A major reason why diffusion models have become so popular (and why we use them here) is because they avoid difficulties in the training objectives of GANs which makes them easier to train.

Thank you again for your review! We would appreciate it if you could let us know whether our above responses resolve your concerns, and change your assessment of our work accordingly.

Thank you very much for your thorough review. We are pleased that you think our proposed method is “novel”, “the concept presented in the paper is intriguing”, and the results are “interesting”. We address your questions below, and we updated the PDF accordingly.

Q1 - [...] paper primarily focuses on image generation [...] for which the diffusion model already excels [...]

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.

Q2 - [...] further evidence demonstrating [...] VAE's performance for specific, necessary tasks.

We agree that Sections 5.2-5.4 evaluate the performance of the VAE somewhat indirectly. Therefore, we added more direct evaluations of practical downstream applications 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 enhances the VAE’s 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.

Q3 - The computational overhead [...] seems substantial.

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

Q4 - Figure 2: [...] it’s hard to see which method is doing the best [...] Figure 3: [...] it seems like the augmented approach does just as well but hard to see

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.

Q5 - [Figure 2:] Why is there a zoom in on the CIFAR-10 plot what is that supposed to show?

We only show this inset because the gap here is so small that it might otherwise be difficult to see that the plot does indeed show two separate green lines.

Q6 - [...] how did you pick hyperparameters for each method?

Thank you for raising this important point. Hyperparameters are listed in Appendix D. For BinaryMNIST and FashionMNIST, we chose the hyperparameters of the VAE models by consulting the literature. For CIFAR-10, we manually tried out a few hyperparameters, and chose an architecture where overfitting occurs, as we are investigating how to alleviate overfitting in VAEs only from the training data. We have updated the Appendix D with this information.

Q7 - [...] running for different seeds [...]

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.

Thank you again for your review! We would appreciate it if you could let us know whether our above responses resolve your concerns, and change your assessment of our work accordingly.

Thank you very much for your review! We address your questions below, and we updated the PDF accordingly.

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

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.

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

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.

Q3 - 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?

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.

Thank you again for your review! We hope that our response and the fixed Table 3 resolve any remaining doubts about our paper, and would appreciate it if you update the assessment accordingly.

Thank you for your insightful feedback! We are glad that you think the proposed idea is “novel”, and the paper is “well written and easy to follow”! We address your questions below, and we updated the PDF accordingly.

Q1 - [The diffusion model might focus on a specific subset of the training data.] Is the distribution of the ELBO on the test-set similar for both approaches?

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.

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

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.

Q3 - [compare sample quality between diffusion model and VAE]

We kindly refer the reviewer to item (1) in our “Overall response regarding sample quality and computational cost” above, which addresses the question of sample quality. In brief, we see little value in using VEAs for purely data generative downstream tasks as diffusion models are known to produce much better samples. However, VAEs can be used for many other downstream tasks that require semantically meaningful representations, which diffusion models don’t provide. Indeed, our paper focuses on the encoder of the VAE since this is the one that tends to overfit.

Q4 - [computational cost]

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.

Q5 - [performance gains only on MNIST and FMNIST, but hardly for CIFAR10]

Please note that, as we stress in Section 5.2, we are interested not only in the test ELBOS but also in the generalization gap (Eq. 8): small generalization gaps are very useful as they imply that ELBOs evaluated on the training set can be used for predicting the final performance of the VAE. The proposed DMaaPx has a significantly smaller generalization gap than other methods in all settings.

Q6 - diffusion models are usually also trained with the same data augmentation techniques. Does it mean that in this case it was omitted?

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.

Q7 - Are the results presented in Figures 2 and 6 statistically significant? [...] it seems like results from a single run

Thank you for this valuable point! In the updated version of the paper, we now show means and standard deviations in Figures 2, 3, and 4 on CIFAR-10. We find that our results are consistent across different random seeds. Experiments for the other data sets and for the amortization gap are still running, and we will update the paper as soon as they finish.

Thank you again for your review. We would appreciate it if you could let us know whether our above responses resolve your concerns, and change your assessment of our work accordingly.