Classification Diffusion Models: Revitalizing Density Ratio Estimation

Summary

This paper introduces a class of generative models based on neural networks which learn the conditional probability distribution of the noise level given a noisy image. It is trained by combining a maximum-likelihood (cross-entropy) objective with a denoising score matching objective. The resulting generative model can be sampled from by iterative denoising like in diffusion models, and also allows for direct likelihood evaluation (in principle).

Strengths

- The proposed model is original and very interesting. Such an approach has the potential to significantly improve upon diffusion models. - The experimental results are very convincing (except for NLLs, see below) and show great promise. - The analysis of the shortcomings of density ratio estimation is also on point and very clear. Although these issues were already evidenced in the previous work of Choi et al., their results were limited to MNIST while CDMs achieve satisfying results on CIFAR10 and CelebA 64x64, and have the added potential benefit of direct likelihood evaluation.

Weaknesses

The biggest weakness of the paper to me is the method used for likelihood evaluation, which has several issues: - An implicit assumption in the paper is that the marginal distribution of the noise level learned by the model matches the correct one used during training, that is, $p_\theta(t) = p(t)$. Although reasonable, it is very difficult to verify, as it requires marginalizing over the high-dimensional variable $x$. - There is no guarantee that the resulting density function (eq. (10)) is normalized (i.e., integrates to 1). This can lead to arbitrarily small _reported_ NLLs (potentially even smaller than the entropy of the data). Although the soundness of the approach is verified in a synthetic but high-dimensional example, I think it would be necessary to verify at least this property on a natural image dataset, though this is also very hard to do (but see below). - There is also no guarantee that the density function in eq. (10) is the density of samples generated by the reverse diffusion (Algorithm 2). CDMs thus contain a priori two models of $p(x)$, one for sampling and one for likelihood evaluation, which could be different. I would also like to see in the final version of the paper (not for the rebuttal) a comparison between the NLLs of these two models on natural images (the latter can be evaluated with an integral along the sampling trajectory using the change of variable formula, as pointed out in the text). If they match, this also verifies that the reported densities are correctly normalized (see previous point). - Finally, restricting NLL comparisons to "models where likelihood can be evaluated in single shot" seems slightly arbitrary, especially since here sampling is not in a single shot. It would make sense to allow a similar computational complexity for sampling from the model and evaluating the log-density at a data point. I might be missing subtle details in my assessment though, and I would be happy to update my score if the authors point out misunderstandings.

Questions

- In the proposed approach, $t = T+1$ plays a special role, which seems slightly unsatisfying to me. Having correct log-density estimates using eq. (10) requires an accurate estimation of $\log p(T+1|x)$ for $x$ a clean image, which is very challenging as natural images are very far from typical realizations of Gaussian noise (unless the explicit form of the Gaussian density is enforced here?). However, the MSE loss (which indirectly involves $\nabla_x \log p_\theta(T+1|x)$) seems to fix this problem. Still, is there a way to avoid it altogether with a more symmetric formulation? I liked how the approach of Choi et al. elegantly solved it by considering a path $(x_t)$ of typical realizations of each $p(x|t)$, though I understand that in their setting log-density estimation requires computing a (one-dimensional) integral. - Out of curiosity, what does Figure 3 look like when the model is trained with MSE loss only? A first observation is that one can add an arbitrary function of $t$ only to the model outputs without changing the MSE loss, so the predicted conditional distributions cannot be right. However, I wonder if this is the only degree of freedom: does the difference between the "true" and predicted $p(t|x)$ depend on $x$? If true, it would show that the CE objective just fixes this constant and is thus not so important. - I am very puzzled by the DDM denoising result at high noise in Figure 4. Why would they be so bad? Predicting the mean of the dataset is really easy, so this is not a capacity issue but rather a training issue? And why CDM would be specifically better at large noise levels? Is this because $T+1$ plays a special role? Minor suggestion: I understand that the "classification" term comes from historical considerations in the density ratio estimation problem. However, it seems ill-adapted here: $t$ is really a continuous variable (which has been discretized, but this does not seem necessary), and therefore $p(t|x)$ is simply a probabilistic model of a conditional distribution that is trained by maximum-likelihood. The title of the paper also sets the incorrect expectation that it is about unifying diffusion models with image classifiers.

Rating

8

Confidence

5

Soundness

3

Presentation

4

Contribution

4

Limitations

As mentioned in the discussion, an important aspect that is not studied in this paper is the architecture of the network to approximate $\log p(t|x)$. The community has extensively explored architecture to model scores $\nabla_x \log p(x|t)$, but architectures for energies have been much less studied. As rightfully pointed out, here $f$ is a UNet, so that $\nabla f$ is the gradient of a UNet and thus significantly differs from a UNet (which would be the goal to emulate score-based approaches). Let me suggest a few references that consider this problem: - Tim Salimans and Jonathan Ho. Should EBMs model the energy or the score? In Energy Based Models Workshop-ICLR, 2021. - Samuel Hurault, Arthur Leclaire, and Nicolas Papadakis. Gradient Step Denoiser for convergent Plug-and-Play. In International Conference on Learning Representations, 2021. - Regev Cohen, Yochai Blau, Daniel Freedman, and Ehud Rivlin. It has potential: Gradient-driven denoisers for convergent solutions to inverse problems. In Advances in Neural Information Processing Systems, 2021. All in all, I very much like this paper, and strongly recommend acceptance. It proposes novel ideas of interest to the generative modeling community, and the strong numerical results demonstrate that this approach has promising potential.

Summary

A common approach for likelihood estimation is the density ratio estimation (DRE). DRE is a method of modeling the density of a target distribution (or data distribution) in the form of a density ratio with a known reference distribution. The standard normal distribution is often chosen as the reference. To train the density ratio models, a classification-based method called noise contrastive estimation (NCE) is generally used. However, in high dimensions, the difference between the reference and the data distribution is large, making the classification problem very easy. Consequently, the model does not learn accurate density ratios. To overcome this saturation problem, some have proposed to create intermediate distributions between the target and reference, similar to diffusion models. For example, the transporting ratio estimation (TRE) first creates a bridge between the target and the reference distribution through their linear combinations. The density ratio between two adjacent timesteps—$t$ and $t+1$-th timestep—is trained using the NCE method. However, despite the diffusion-like modifications, TRE has failed to resolve the saturation problem in high-dimensional data. Therefore, to overcome the limitation of the previous approaches and estimate the log-likelihood efficiently for the high-dimensional datasets, the paper proposes a novel DRE-based diffusion-based generative model. Most importantly, the paper proposes to predict the timestep $t$ for a given observation $x$—$p(t | x)$—when $x$ is sampled uniformly from all intermediate distributions in the bridge. It means that the model predicts how noisy the given input is, unlike the classification problems proposed in TRE. The paper highlights two important benefits of the proposed method. First, the authors draw the analytical connection between the noise-predicting classifiers and the Bayes-optimal denoisers at each time step. This connection means conventional parameterizations in diffusion models, such as epsilon/score/denoiser, can also be written by the noise-predicting classifiers. Thus, the noise-predicting classifiers can be trained via denoising score matching (DSM) as well as the cross-entropy loss. This training overcomes the saturation problems that originated from classification-based training. Second, even compared to the TRE, regardless of its scalability issue, the noise-predicting classifiers provide efficiency at the log-likelihood evaluation. The proposed method can estimate the log-likelihood of the (clean) data in one-g, while the TRE requires $T$-number of model evaluations. In addition, the paper mentions a concurrent work that proposes a similar noise-predicting classifier, but the authors emphasize that the concurrent work didn't incorporate the DSM loss, which is the key ingredient of the proposed approach's scalability. Finally, the authors demonstrate the proposed method's effectiveness via various experiments, including generative modeling of image generation benchmark datasets. ------------------------ Update the rating from 7 to 8 after the authors' rebuttal

Strengths

One of the paper's key contributions is the introduction of a novel solution for likelihood estimation. The proposed method addresses significant challenges in the density ratio estimation and offers a fresh perspective on the parameterizations in diffusion-based models. Moreover, the authors effectively motivate the use of DSM by establishing the connection between the time-prediction classifiers and the conventional parameterizations. This training overcomes the limitations of previous DRE-based approaches. Furthermore, the proposed method exhibits comparable performance and scalability compared to popular diffusion-based generative models, further underscoring the practical effectiveness of the proposed DRE-based method.

Weaknesses

While the proposed method's originality and novelty are clear, its significance is still doubtful. While achieving efficient likelihood estimation is the primary motivation, the paper doesn’t provide sufficient discussion about it—instead, it focuses on demonstrating the improved generation performance of the proposed DRE-based method. Thus, the proposed method would be much more impactful if the authors could demonstrate what becomes available when the likelihood estimation for high-dimensional data works. Additional discussion about the importance of the classification loss may seem required. It is clear that the time-prediction classifier links to the conventional denoiser parameterization, and thus, this connection equips the use of DSM to train the classifiers, which enables the proposed model's applicability to high-dimensional data. However, this also implies the possibility of training the proposed models without any classification objective.

Questions

N/A

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

N/A

Summary

This paper introduces a clever connection between Density Ratio Estimation (DRE) and Diffusion Models (DM), showing that optimal denoisers are also optimal noise classifiers. Doing so allows them to construct a new type of loss based on noise classification. This allows DRE methods to inherit the benefits of diffusion models, like strong sampling, while adding new capabilities like single step likelihood estimation.

Strengths

- The mathematical connection is clever and elegant. - I liked the framing / connections to DRE literature. - The single step likelihood calculation gives a distinguishing feature that sets the method apart from diffusion models. - Reasonable ablations were considered, and clear and readable notation was used. - The method requires an additional backward pass for each denoising/sampling step. In the spirit of the NeurIPS guidelines, I count this clearly labeled limitation as a strength. It's intellectually interesting and also clearly distinguishes that the computation done by their model is different from regular diffusion models in a non-trivial way.

Weaknesses

The proposed approach is functionally equivalent to a diffusion model (using Eq. 8 and with necessary incorporation of MSE loss that matches typical diffusion models), so you might hope to see that you can substantially improve on diffusion models by adding the CE loss as a kind of regularizer. But there are several reasons why it isn't clear that this gives a susbtantial improvement. - The architecture has to be changed. There's no obvious way to get around that, but it does marginally affect the ease of use for these ideas, and the directness of comparisons (though on this point I agree with the authors' statement that they chose a minimal intervention). - MSE improvements (Fig. 4) are only marginally different in absolute terms, which is what matters for log likelihood. See some questions about this result. - The authors only used / compared with relatively old diffusion models, so it's not clear if modest benefits to MSE or FID in some cases would still be seen in more SOTA models. I would expect much more comprehensive experiments to support the idea that using CE loss can improve sampling (Table 2). The strongest claim of improvement is the single step likelihood calculation. However, there are three weaknesses associated with this result. Table 3 is a bit of a dubious comparison for two reasons. - Unlike the ELBO, it's not clear if the proposed estimator gives an upper bound on NLL. I suspect it doesn't. If there are errors in p(t|x), then because it shows up with both signs in Eq. 10, it could contribute to the error in either direction. So we are comparing bounds on NLL to an estimate of NLL which could look lower due only to error. - The NFE comparisons are also a bit dubious. Function evaluation complexity varies wildly between methods. We should have a fair comparison to the associated diffusion model with the same architecture, but it wasn't clear which one that is in the table. Though it's a bit moot as I wouldn't trust the result anyway, as we don't have a bound here. - A compelling application of the single step log likelihood evaluation was not discussed. If the log likelihood had been used and evaluated for some downstream application, it might have assuaged worries about the accuracy of the estimate. Minor comments: - I didn't like the abstract sentence "directly output likelihood...lacking in most generative techniques", since most common generative techniques today do output likelihood. You clarify later what you mean by "directly" so it is Ok. I just wanted to say that it was off-putting on first read.

Questions

I was really surprised by the fact that your architecture has no timestep conditioning, but matches or outperforms DDMs that do have timestep conditioning in denoising MSE. Have I understood that correctly (i.e. that the DDM in Fig. 4 does have timestep conditioning?) If so, that seems like a really interesting conclusion for diffusion research, as the final appearance of the timestep only shows up in a simple way in Eq. 8, potentially leading to a simplification of other architectures. It would be interesting to see if this effect persists for stronger architectures. On the other hand, maybe the right way to interpret this result is that your architecture essentially outputs the result for all time-steps at once. Eq. 8 selects the correct one gets a score through backprop. Still, that may be an interesting possibility to try to explore for standard diffusion models. In Fig. 3, I can understand why the classifier is limited, but it wasn't totally clear to me why MSE helps. But there may not be an easy answer to that (besides the superficial one that it is a regularizer of some sort).

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations were well addressed.

Summary

This work develops a new generative framework called the classification diffusion models (CDMs) based on the density ratio estimation (DRE), by establishing an interesting connection between the DDPM's denoiser and noise-predictive classifier, which also helps the exact likelihood computation in a single pass. As is claimed by authors, the proposed method is the first DRE-based technique to successfully generate images of the CIFAR-10 dataset.

Strengths

1. This work is well-written and clearly formulated. 2. The proposed algorithm is theoretically grounded. 3. The experiments are convincing, which verify the effectiveness of proposed methods.

Weaknesses

1. The main concern is the capability or potential of DRE-based methods to successfully learn data distributions, particularly when applied to datasets with large scales in practice. As is claimed by authors, the current development of DRE-based modeling is up to the CIFAR-10 dataset (MNIST before), which is too small to indicate the validity compared to more standard methods (e.g. DDPM). 2. Another major concern is that CDMs can be more computationally challenging than DDPMs in sampling. According to Algorithm 2, CDMs require an extra backward propagation (BP) than DDPMs for each time step when denoising, which is quite expensive.

Questions

1. Please provide more details of questions raised in the weaknesses section above. 2. Check minor grammar typos, e.g. in Line 258, " trains a classifiers...".

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

As is stated by authors, it is worthy to explore BP-friendly architectures to alleviate the computation inefficiency of CDMs.

Response to rebuttal

Thank you for the responses. I will maintain my score - I think this is nice theoretical foundation work but as "the main goal is not improving [the state of the art models], but rather making DRE-based methods work for the first time on datasets more challenging than MNIST", it somewhat limits the impact. FYI, while bidding for papers I noticed a concurrent similar work that you may want to check out/cite as concurrent: "diffusion models as noise classifiers" or something similar.