DiffEnc: Variational Diffusion with a Learned Encoder

Thank you very much for your review which helped up improve the article.

The benefit of the encoder is unclear.

In a sense we are asking the question: “Can a diffusion model benefit from a non-trivial time-dependent encoder”. Since we achieve state of the art results on CIFAR-10, we say that the answer is “yes”. Since the encoded image is still in the full image space, the encoding is not directly applicable for representation learning. However, as we mention in the new section 7 (old section 9), it would be an interesting direction for future research to explore how distillations of the time-dependent encodings could be used for representation learning.

In spite of the added theoretical overhead incurred by the learned encoder, there is minimal empirical improvement over existing baselines

Our large models achieve state of the art on CIFAR-10 showing a significant improvement over a model that uses the exact same architecture for the generative model. This is an extremely well-studied benchmark so it can be expected that it is difficult to improve SOTA with a lot. Furthermore, we continued training our models until 8 million steps and saw a further performance increase to 2.62 as seen in the updated Table 1. (The VDM paper uses 10 million steps.)

Empirical validation is slim

We are currently training models on downscaled ImageNet 32x32. Due to the size of the models we do not have results ready for this response, however, we expect to have them on Monday, where we will share them in a new response.

Is there improvement in terms of any other metrics? There is a lot of research on how likelihood can be a poor proxy for generative model performance.

As mentioned in your reference [1] likelihood and visual quality of samples are not directly linked and therefore the metric used to train a model is important when choosing the application. Since we show that our model can achieve good results on likelihood, and likelihood is important in the context of semi-supervised learning, it would be an interesting direction for future research to use this kind of model for classification in a semi-supervised setting. We have added this possible application to section 7. FID scores are a quite noisy metric, but we have added them to appendix R, and they show that FID scores for the VDM model with v parametrization and DiffEnc are not statistically different.

Parameterization of the learned encoder is very simplistic. Why is the encoder conditioned only on the time?

The encoder is dependent on both the time and the original image. As we wrote in the previous version of the article in section 3: $x_t = f_{\phi}(x, t)$. We have made this more clear in the beginning of the section. In our new figure 1, we also visualize the setup and show how the encoder takes both an image and time as inputs.

...the justification for this parameterization (Appendix I) -- that this makes the math nice -- does not feel entirely sufficient.

The reason for using $\alpha_t$ multiplied with the data in the encoder is that we want the encoder to be the identity for t = 0 and to send the mean to 0 for t = 1. This is important for minimizing the reconstruction and latent losses as we mention in the beginning of the new section 4 (old section 6). The argument about making the math nicer is only to justify that instead of using just $\alpha_t$, we use $\alpha_t^2$. We have rewritten appendix I to make this clearer.

Can the authors directly visualize the output of the learned encoder model at different times

Good idea. In our new figure 1, we have some examples of the encoded images. In general, the encoder seems to enhance edges for lower t and bring the mean of the image closer to 0 for higher t, however while keeping more signal for higher t than the non-trainable encoder would. We have clarified this in results in the new section 5 (old section 8.2).

How does Section 7 (v-prediction parameterization) relate to the rest of the paper?

As we write in the new section 4 (old section 7), we use the v parametrization in our experiments. The exact equations we use for our loss are using the new numbers (19) and (20), the old numbers were (38) and (39).

In which experiments in Section 8 were a fixed noise schedule used, and in which were a learned noise schedule used?

Thank you for bringing this lack of clarity to our attention. As we mention in the experimental setup of the new section 5 (old section 8.1), we used fixed endpoints for the noise schedule for the large models and experimented with both fixed and learned endpoints for the small models. Table 2 has results of the large models, so these are with a fixed noise schedule, and we have added this to the caption. In table 3 and 4, we write in the column “noise” which results are for a fixed noise schedule and which are for a trainable schedule. In the caption of table 5, we write that these results are for a fixed noise schedule.

Thank you very much for your review. It has been very helpful for making revisions.

The presentation of the method, namely section 3 and 6 could be improved significantly.

Thank you for this valuable feedback. We have made a thorough rewrite of the article and feel that this makes the presentation of the method easier to follow.

Moving the figure provided in Appendix A to the main text might be helpful...

Thank you very much for this suggestion. We have made a visualization of how DiffEnc works and have added it as our figure 1. We feel this clarifies the overall idea.

...the performance improvement is not overwhelming...

The result of the large models is new state of the art on CIFAR-10 showing a significant improvement over a model that uses the exact same architecture for the generative model. This is an extremely well-studied benchmark so it can be expected that it is difficult to improve SOTA by a substantial amount. Furthermore, we continued training our models until 8 million steps and saw a further performance increase to 2.62 as seen in the updated Table 1. (The VDM paper uses 10 million steps.)

... an improvement on par with what might be possible through better hyperparameter tuning…

Since we train a baseline VDM using V parametrization, which uses the exact same hyperparameters as the DiffEnc model, if better results could be gained for the VDM with better hyperparameter tuning, then it would likely be possible to get the same improvements for DiffEnc.

... or by increasing the model capacity (which is implicitly part of what this method is doing, at least in the case where the encoder has learnable parameters).

The encoder is not used during generation. Thus, our work shows that using an encoder enables us to find a better generative model without increasing the number of parameters of the generative model.

The results being limited to CIFAR-10...

We are currently training models on downscaled ImageNet 32x32. Due to the size of the models we do not have results ready for this response, however, we expect to have them on Monday, where we will share them in a new response.

Do you have a sense for what the weaknesses of your gradient approximation (Section 6) are? It seems that the approximation will be very bad for t near 1.

The approximation of the gradient will be bad for the trainable encoder if $y_{\phi}(\lambda_t) - dy_{\phi}(\lambda_t)/dt$ is large. One might think that this would force $y_{\phi}(\lambda_t)$ to be close to $dy_{\phi}(\lambda_t)/dt$, however inspection of some examples show that this is not the case. The reason might be that the normed part of the loss is scaled with $\alpha_t^2$, and therefore the difference $(y_{\phi}(\lambda_t) - dy_{\phi}(\lambda_t)/dt)$ can be rather large without affecting the total loss notably. Of course, one might get even better results with a better approximation, and this would be an interesting direction for future research.

What is the intended take-away from Figure 1?

(In the new version of the paper this is figure 2) In short, the take-away is that the encoder learns to do something different from how it was initialised, and what it learns is different for different timesteps. We have reworded the figure caption and added a sentence in the results of the new section 5 (old section 8.2) to clarify. We did not have room for the details in the paper, but we have added them to appendix S.

Thank you for your response and your changes to the submission. I think these changes have substantially improved the clarity and credibility of the work, and I will be raising my score to a 6.

Thank you very much for your review, especially your helpful suggestion of clarifying the main equation.

It's challenging for me to discern whether Section 2 serves as an introductory section or is meant to highlight one of your contributions.

Good point. Section 2 is meant as an introductory section, to introduce the notation used for variational diffusion models. We have added a sentence at the beginning of the section to clarify this and changed the title to “Preliminaries on Variational Diffusion Models”.

The absence of a central theorem throughout the paper poses a difficulty for readers in anticipating the direction of the derivations

Thank you for this suggestion. We have added in the introduction which equation is the main expression of the DiffEnc loss - Equation (18) - and sketched how to arrive at it.

it's clear that there's a need for improvements in the organization of the content.

We have taken your comment to heart and made a thorough restructuring of the sections. We think that this made the organization more intuitive by making the separation between DiffEnc and the standard diffusion approach clearer.

Thank you very much for your review. In the new version of the article, we have done our best to address your concerns.

The paper has very limited evaluation...

We are currently training models on ImageNet 32x32. However, SOTA models are large and require access to extensive compute. We will share results later in the rebuttal period.

...doesn't compare to some of relevant baselines that are even mentioned in the paper, like latent diffusion...

In Table 1, we compare to all state of the art likelihood-based models for CIFAR-10 without data augmentation. Since we did not use data augmentation on CIFAR-10, we have only compared with other results on non-augmented CIFAR-10. In the latent diffusion paper by Rombach et al 2022 “High-Resolution Image Synthesis with Latent Diffusion Models”, they do not report any likelihoods on CIFAR-10, therefore they have not been included in the table.

…it mentions that some methods only show improvement after longer training

By design, the DiffEnc and VDMv models have the same initial loss. However, the trends differ, especially from around 2 million training steps. Which means that it is only after more than 2 million steps the mean losses differ significantly. We would not say that this shows any kind of inconsistency in the results. We have made this clearer in the description of the results in table 5 which was in the old section 8.2. However, we have had to move table 5 to appendix Q due to the page restriction.

...what does it mean that the loss is dominated by the diffusion loss?

By construction, the VDM model has very small latent and reconstruction loss when using the fixed signal to noise ratio endpoints from the original article. Thus, the diffusion loss is where there is the biggest potential for improvement. We have added a mention of this in results in the new section 5 (old section 8.2) when discussing tables 3-5.

In section 3, I first missed the added subscript in x_t in eq 13

Thank you for pointing us towards this point of potential confusion. We use $x_t =x_{\phi}(\lambda_t)$, to remind the reader that the encoded image is dependent on both x and t through $\lambda_t$. We think that changing the notation to not use x would obfuscate this dependency. We have added “dependent on x and t” when we introduce the encoder in section 3 to clarify.

In section 6, I was not fully following the parametrisation of $x(\lambda_t)$ how does this related to eq (13)?

From the old section 5 and onwards we write $x_t = x_{\phi}(\lambda_t)$ to include the dependence of the encoder on the learned parameters, $\phi$, and the log of the signal to noise ratio, $\lambda_t$. In the new version we introduce $x_{\phi}(\lambda_t)$ from the beginning of section 3, hopefully, this makes it clearer. When given x, we have for the non-trainable encoder $z_t = \alpha_t^3x$ and for the trainable encoder $z_t = \alpha_t(\alpha_t^2x + \sigma_t^2 y_{\phi}(\lambda_t))$ where $y_{\phi}$ is the trainable part of the encoder.