Paper: DiffEnc: Variational Diffusion with a Learned Encoder (ICLR-2024)

Reviewer comment (Experiments):
The benefit of the encoder is unclear. There is limited exploration of the learned encoded space outside of the visualization provided in Figure 1. What is gained by this formulation?

Author response:
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.

Reviewer comment (Evaluation):
In spite of the added theoretical overhead incurred by the learned encoder, there is minimal empirical improvement over existing baselines (BPD 2.65 -> 2.63 on the CIFAR10 dataset).

Author response:
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.)

Reviewer comment (Theory):
Parameterization of the learned encoder is very simplistic. Why is the encoder conditioned only on the time $\lambda_t$ and not the data $\mathbf{x}$ itself? While I do appreciate elegant derivations, the justification for this parameterization (Appendix I) -- that this makes the math nice -- does not feel entirely sufficient.

Author response:
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 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.

Reviewer comment (Evaluation):
It is still not entirely clear to me what is learned by the encoder model. Can the authors directly visualize the output of the learned encoder model at different times $t$?

Author response:
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).

Reviewer comment (Experiments):
How does Section 7 (v-prediction parameterization) relate to the rest of the paper? Is it used for experiments in Section 8?

Author response:
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).

Reviewer comment (Reproducibility):
In which experiments in Section 8 were a fixed noise schedule used, and in which were a learned noise schedule used?

Author response:
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.

Reviewer comment (Presentation):
The organization of sections is perplexing. It's challenging for me to discern whether Section 2 serves as an introductory section or is meant to highlight one of your contributions.

Author response:
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”.

Reviewer comment (Presentation):
There is a need for improvements in the organization of the content.

Author response:
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.

Reviewer comment (Experiments):
The paper has very limited evaluation and doesn't compare to some of relevant baselines that are even mentioned in the paper, like latent diffusion, and only compares on Cifar-10 and MNIST.

Author response:
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.

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.

Reviewer comment (Experiments):
The paper mentions that some methods only show improvement after longer training, hinting at potential inconsistencies in the results in case of slightly different training setups due to not training till convergence, making it hard to judge whether the proposed changes are a significant improvement.

Author response:
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.

Reviewer comment (Evaluation):
The ablation studies of the different methods are interesting but it's unclear what insight is gained from this, e.g. what does it mean that the loss is dominated by the diffusion loss?

Author response:
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.

Reviewer comment (Writing):
In section 3, the reviewer first missed the added subscript in $x_t$ in eq 13, suggesting it might make it easier to read to change the notation from $x_t$ to e.g. $e_t$ or something similar to be clear there.

Author response:
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.

Reviewer comment (Theory):
In section 6, the reviewer was not fully following the parametrisation of $x(\lambda_t)$ and asks how this relates to eq (13), e.g. does eq. (29) indicate that $q(z_t|x) = N(a_t^3 x; ...)$?

Author response:
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.