Classification Diffusion Models: Revitalizing Density Ratio Estimation

We thank the reviewer for the insightful and interesting comments.

Comparing the one step likelihood computation to an integral along the sampling trajectory

In theory Eq. (10) should be normalized, but this is a great idea. We’ll definitely add it to the final version. We’d just like to note that while our one-step likelihood computation depends only on the model's output, the integral approach depends also on the sampling method. Specifically, different noise schedulers correspond to different ODE/SDE discretizations, and may affect the quality of the generated samples. This is even if each denoiser by itself is optimal. For a scheduler that leads to low quality samples, the likelihood of the generated images obtained with our one-step approach would be low (high NLL). For example, as we show in Sec. 4.3, we can choose a scheduler that improves the NLL over the test set but leads to a worse FID of the generated samples. We argue that this doesn’t imply the model learned a worse approximation of $p(x)$, but only that the sampling trajectory is worse than that of DDPM.

The DDPM scheduler leads to good generated samples, though. This suggests that the discretization associated with this scheduler leads to an accurate solution of the ODE. And therefore, in this case, we agree that your suggested experiment is interesting and we will include it in our final version. Thanks!

Likelihood evaluation is performed in a single step, but sampling is not

We completely agree. However, there may be applications that only require likelihood evaluation for real images, and not sampling of synthetic images. This is the case e.g. in out-of-distribution detection. Therefore, we believe there is merit in separately reporting the NFEs required for likelihood evaluation and for sampling.

The ability of our method to evaluate the likelihood of real images in one step stands in contrast to methods that calculate the likelihood by solving an ODE through a sampling-like process (with repeated evaluations of the network), as in DDMs.

In any case, note from Table 3 that our method is competitive also with methods that evaluate the likelihood using many NFEs.

The role of t=T+1

This is a very interesting question. Please note that both in Equation (3) in [34] and in Proposition 1 in Choi’s work [7], the likelihood is extracted from the ratio $\log\frac{p(x|0)}{p(x|T+1)}$, since $p(x|T+1)$ is known (a Gaussian in this case). In order to achieve a correct likelihood estimation, all the intermediate ratios need to be accurate for the same input $x$, and specifically, the ratios between densities which represent high noise levels need to be accurate when the input is a clean image. For this reason, DRE methods failed to learn distributions of datasets which are more complicated than MNIST. When you incorporate the MSE loss using our theorem, it enforces the log ratio $\log\frac{p(t|x)}{p(T+1|x)}$ to be accurate for any $t$, which solves this problem. However, it would be very interesting in future work to try solving it in different ways.

Using only the MSE loss

When training the model using only the MSE loss, the CE is very high (please see Table 1) and Figure 3 is very noisy, which indicates that the model fails to learn the distribution $p(t|x)$ in this case. Your observation is correct and very interesting. The MSE indeed doesn’t change if we add a function of $t$ to the model’s output. This is why the CE loss is important for estimating $p(t|x)$. Without the CE loss, the model can function as a denoiser but is useless for the purpose of outputting the likelihood in a single step. This indeed seems to be the only degree of freedom that the CE loss fixes because the MMSE predictor corresponds to $\nabla_x \log p(t|x)$, which uniquely defines $\log p(t|x)$ up to an additive constant that may depend on $t$.

Note that the difference between the true and predicted $\log p(t|x)$ may depend on $x$. However, in this case it is the same function of $x$ for all $t$. This is true because this function should cancel out after the softmax and also when we take the difference between $\log p(T+1|x)$ and $\log p(t|x)$ in Equation (8).

To conclude, the MSE loss alone is invariant to adding a function of $t$ to the model’s output, and the CE loss removes this degree of freedom. In addition, both losses are invariant to adding a function of $x$ to the model’s output, but the softmax removes this degree of freedom. Importantly, when we calculate the NLL using Equation (10), this function of $x$ is canceled out even before the softmax, since it is added both to $\log p(t|x)$ and $\log p(T+1|x)$.

Why is the prediction of DDM poor for large $t?$

We don’t know why diffusion models fail here, and it would be very interesting to explore it in the future. However, we do know why CDM is particularly good at the higher noise levels. Please note that in Theorem 3.1, we show that the optimal denoiser depends on the model's prediction for $p(t|x)$ and $p(T+1|x)$. At the high noise levels, $t$ is close to $T+1$, and as we show in Fig. 3, in this case the model’s prediction for $p(t|x)$ and $p(T+1|x)$ is accurate even without incorporating the MSE loss. This accuracy is further improved when we include the MSE loss.

I am happy that you find my suggestions helpful, and that they might improve the paper. I reiterate that this paper should be accepted for the reasons outlined above.

t = T+1

In Choi et al., I was referring to what the authors of this paper call the "pathwise method" (section 5.2). The resulting integral has the property that the (infinitesimal) density ratios (as well as the scores) need only be accurate on typical samples from $p(x|t)$ for their respective timestep $t$. Although it seems that they evaluate this more sophisticated method only on Gaussian data.

Using only the MSE loss

The reason I asked this question is because this statement is deceptive: "$\nabla \log p$ uniquely defines $\log p$ up to an additive constant". This is true in a strict sense, but not in an approximate sense. In particular, one might have a good approximation of the score on the support of $p$ without having a good approximation of the energy on said support, when the support is disconnected. E.g., for a mixture of two well-separated Gaussians, the relative probability of each mode is not captured by the score (it is encoded in the value of the score in-between the two modes, where there is no training data). This is known as the lack of a Poincaré inequality in the maths literature. So this means that in practice, only using the MSE loss might lead to estimated values of $\log p(t|x)$ whose error also depends on $x$.

Likelihood estimation is the primary motivation; the paper lacks illustrations of its use

Our primary motivation was to analyze why existing DRE methods fail to capture distributions of complex high-dimensional data, and to develop theory and a practical DRE method that doesn’t suffer from this limitation. The single-step likelihood evaluation is actually a by-product of our analysis. Nevertheless, we agree with the reviewer that it would be interesting to explore the single step likelihood computation in the context of downstream tasks. We see this as an intriguing direction for future work. For example [R2] showed that a generative model which is able to calculate the likelihood, can be used for zero-shot classification. However, in their case they need to perform many NFEs in order to calculate the likelihood. These applications are out of scope for our paper, but we believe that our method has the potential to be used for solving these tasks in future work.

Additional discussion about the importance of the classification loss may seem required

That's an important point. As shown in Table 1, when training using only MSE loss, the CE is very high. In this case, we can use the model as a denoiser, but it will not function as a noise level classifier and specifically will be useless for calculating the likelihood of a given input using Theorem 3.2. The explanation is that the MSE is invariant to an addition of a function of $t$ to the classifier’s output (since when we take the derivative with respect to $x$ in Eq. 8, it zeros out). Therefore, we must combine the CE loss to remove this degree of freedom.

[R2] Zimmermann, Roland S., et al. Score-based generative classifiers. In: NeurIPS Workshops ‘21.‏

The architecture has to be changed

We agree with the reviewer that the architecture can't be exactly the same since DDM is a denoiser and CDM is a classifier, which may impact the comparisons a little bit. However, as the reviewer notes, the change in the architecture is minor – we replaced only the last layer of the DDM architecture by a convolution, followed by max-pooling and a linear layer. This has a negligible effect on the number of parameters. It is important to note that this architecture was optimized for DDPM and not for CDM, yet CDM still achieves comparable or better results with it. As we mentioned in the text, we believe that an important and interesting future direction would be to optimize the architecture for CDM.

Old diffusion models

We want to emphasize that our main goal is not improving denoising diffusion models, but rather making DRE-based methods work for the first time on datasets more challenging than MNIST. Therefore, we mostly focus on analyzing the reasons that DRE-based methods failed, and show how our theorem leads to a practical algorithm that overcomes these issues. Moreover, we used the same architecture of DDPM only for a fair comparison. In future work, it would be interesting to investigate different architectures which are more suitable for CDM (and in general may be different than the standard denoiser architectures).

Upper bound on NLL

Our method indeed doesn’t upper bound the NLL but rather estimates it. Yet, this is true for all methods in Table 3, which calculate the likelihood and not the ELBO. For example, DDPM-SDE and DDPM++ (which estimate the likelihood and not ELBO) numerically solve an ODE in which each step can accumulate errors.

As for the fact that $p(t|x)$ appears with both signs, this is a key problem with existing DRE based methods, which our method solves. Without incorporating the MSE loss, the model’s prediction is inaccurate for $t$ values that are far from the real $t$ as we show in Fig. 3. Specifically, when calculating the likelihood of a clean image $x$, the model needs to be accurate for both $p(0|x)$ and $p(T+1|x)$. This is the reason that DRE-based methods failed to date. In our case, the addition of the MSE loss solves this problem and makes the NLL calculation more accurate (see also the answer to the last question below). In App. D we show that our NLL calculation is exact on high-dimensional toy examples for which we can calculate the ground-truth NLL analytically.

Function evaluation complexity varies wildly between methods

While the evaluation complexity depends on the architecture, in all cases shown in the table, the number of NFEs is at least 100 and in some cases even thousands. In our case, the calculation requires only a single NFE which is much cheaper computationally. In addition, we use the same architecture as DDPM which appears in the table and requires approximately 200 NFEs, and as we show, we achieve a better NLL.

Validating the log likelihood evaluation on a downstream task

We agree with the reviewer that it would be interesting to validate the effectiveness of our single step likelihood computation in the context of downstream tasks. We see the use of our method in applications as an intriguing direction for future work. For example [R1] showed that a generative model which is able to calculate the likelihood, can be used for zero-shot classification. However, they need to perform many NFEs in order to calculate the likelihood, whereas our method can do this with a single NFE.

Following the reviewer’s concern, to further validate our NLL evaluation, we will add to the final version more high-dimensional distribution examples (in addition to those already appearing in App. D), for which we can compute the ground-truth NLL analytically.

The architecture has no timestep conditioning

Exactly; you understood it correctly. Our theorem suggests that, given the optimal classifier that predicts the timestep (and therefore can't be conditioned on it), we can extract the optimal MMSE denoiser. This effect should persist independently of the architecture and specifically for stronger architectures. It is a unique property of CDMs that doesn't exist in DDMs, since in CDMs the time condition implicitly appears by taking the corresponding classifier's output (i.e., if we want to be conditioned on timestep $t$, we will take the derivative with respect to the $t$-th entry of the classifier outputs, as shown in Theorem 3.1).

Your interpretation that our architecture essentially outputs the result for all time-steps at once is correct, but note that it does so in a very efficient manner. Specifically, if we wanted to do so in a diffusion model, it would have to output $T$ predicted noise maps at once. Our model, on the other hand, only outputs $T$ scalars. It is the gradient of each output that gives the predicted noise map.

Why does MSE help in Fig. 3?

That is a very important point. When training the model using only the CE loss, the model’s prediction $p(t|x)$ may be accurate for the correct noise level but not for more distant noise levels (as shown in Fig. 3). This is why DRE-based methods have failed to capture the distribution of images to date. As Theorem 3.1 shows, the optimal denoiser for timestep $t$ depends on the model prediction for both $\log p(t|x)$ and $\log p(T+1|x)$. Therefore, when we add the MSE loss on the derivative of the difference between them (using the formula from Theorem 3.1), we enforce the classifier’s output to be accurate not only in it’s $t$-th entry (as can be achieved using only CE) but also in its $(T+1)$-th entry. By enforcing this, we obtain an optimal classifier that predicts the correct probability $p(t|x)$ both locally around $t$ (thanks to the CE) and globally for distant $t$ values (thanks to the MSE).

[R1] Zimmermann, Roland S., et al. Score-based generative classifiers, NeurIPS Workshops ‘21.‏

The potential of DRE-based methods to learn data distributions of larger scale datasets

We acknowledge that the Celeb-A 64x64 and CIFAR-10 32x32 datasets we experimented with are not very large and high-dimensional by today’s standards, and we certainly agree that extending the evaluation to datasets like ImageNet would be valuable. Unfortunately, however, we lack the computational resources to do so. For instance, as reported in [9], Table 10, achieving satisfactory results (in terms of FID) on ImageNet 128x128 requires a large amount of computation, reaching roughly 521 days on a single V100 GPU. We can run on eight A6000 GPUs, with which we approximate it will take us 56 days to train a standard DDM with the same architecture and number of iterations, using a batch size which is smaller by a factor of 2.

Nevertheless, we believe that the preliminary exploration we present on CIFAR-10 and Celeb-A is still valuable for the research community, as it analyzes why DRE-based methods have failed to capture distributions of high-dimensional data to date, and suggests a practical algorithm based on our theoretical result for overcoming this problem. Furthermore, we truly believe that with sufficient computational resources, future works will be able to train our method on larger-scale datasets.

CDMs can be more computationally challenging than DDPMs in sampling

We agree with the reviewer and have mentioned it as a limitation in the main text. Indeed, while a DDM requires a single forward pass for each denoising step, a CDM requires both a forward pass and a backward pass. Nevertheless, the computational cost of performing a forward and a backward pass through a network depends on its architecture. In this work, we chose to use the same architecture as that used by DDPM [19] to isolate the impact of our algorithmic approach from the choice of the model architecture when comparing it to DDMs. However, it is an important future direction to explore architectures that are particularly optimized for CDMs and that alleviate the gap in computational complexity. Such architectures should have the property that performing a forward pass and a backward pass through them is computationally similar to performing only a forward pass in a regular DDM. This could potentially be achieved e.g., by relying only on the encoder part of the UNet. However, we leave this exploration for future work.

Minor grammar typos

Thanks. We will fix them in the final version.