Feature-guided score diffusion for sampling conditional densities

Thank you for your review which will help us improve the paper. We provide answers to each of your concerns below, and have updated the paper accordingly.

Weaknesses:

1. We agree that the paper motivation was not sufficiently clear. We have now re-written the conditional sampling subsection to clarify our motivation (see lines 125 to 145 in the updated paper). We removed the phrase about “conditioning and adding noise not commuting”. Although the statement is correct, we realized the presentation was not immediately clear. The argument in the paper is now described as : “In practice, one might employ neural network classifiers trained on clean data to estimate the likelihood. This however introduces an error because $p_\sigma (y|x_\sigma ) \neq p(y|x_\sigma )$. In other words, the correct likelihood of y for noisy data cannot be computed by evaluating the likelihood function $p(y|.)$ on noisy data. The likelihood function changes at different noise levels.” This is a more elaborated version of “adding noise and conditioning do not commute”. This is simply to say the Bayes rule (eq4) holds if one first adds noise and then conditions, but it does not hold if one first conditions and then adds noise. Specifically:

• Adding noise and conditioning: convolving $p(x)$ with Gaussian kernel $g$ and then conditioning results in $p_\sigma(x_\sigma |c)$
• Conditioning and adding noise: convolving $p(x|c)$ with a Gaussian kernel results in $g(z) * p(x|c) = g(z) * (p(c|x)p(x)/p(c))$
• The two are not equal. Please also see the first page of Chidambaram et. al, 2024 (cited in our paper) for more details.

2. It is true that many conditional diffusion models in the literature are trained with the condition as an input, but existing algorithms do not provide samples from the appropriate conditional distributions, as demonstrated in the literature that we cite in the Background section. This includes classifier-free guidance methods where the network is trained conditionally. Several publications (cited in our submission) have shown that the choice of weighting that leads to higher quality samples achieves this at the expense of lower diversity, for example for the mixture of two Gaussians. We show that our method does not fail for mixture of two Gaussians. At the same time, it succeeds in returning visually high quality samples of complex images (see next point).

3. To address comparisons over datasets used by other algorithms, we generated new results. We applied our method to the CIFAR dataset of small images and also to larger size images in a dataset constructed as a union of 6 commonly used datasets. Results are presented in Figures 14 to 16 (AppendixD). We used the same architecture for both of these experiments: a UNet with 3 encoder and decoder bocks. Despite the increased complexity of these datasets relative to the texture example in the original submission, we find that the model succeeds in drawing visually high quality samples from the correct class. For quantitative comparison we need to train on Imagenet dataset which is a substantially larger and more complex dataset, with a relatively small number of samples in each class. We expect we would need to increase the size of the network and modify the architecture by adding commonly used elements such as attention blocks, BigGANs down-sampling blocks, etc. This is a worthy endeavor, but not one we could attempt during the rebuttal period.

4. Our method does not “rely on access to an example image for a given class for sampling”. During synthesis, it only needs access to $\phi_y$ which is a pre-computed feature vector of the class. So each iteration involves only one forward pass and it is not cumbersome. During training, we rely on access to pairs of images from the same class. So each training iteration involves 2 forward passes. We will clarify this important distinction in the revised paper

5. The typos in the discussion section have been fixed.

Questions:

• Our motivation for choosing a texture dataset was to examine a continuous conditioning space (as opposed to discrete object classes), and to sidestep any requirement of labeled data. In this dataset, with 1700 classes and only 140 images per class, we find that a relatively small and simple fully convolutional UNet network succeeds in learning the conditional densities. Achieving the same success for complex and discrete object classes (such as those in ImageNet), will require much larger networks and datasets with many images per class. However, as discussed above, we now included two new experiments on more common datasets.

• See first reply under “Weaknesses”

Regarding the summary: in case there is any misunderstanding: the method we propose does not rely on a GMM assumption. The GMM example is provided as a numerical demonstration that we can sample from the correct conditional distribution, which has been shown to fail for previous guidance algorithms.

Thank you for your important remarks, which will help us to improve the clarity of the presentation.

Weaknesses:

Regarding the second point: We have updated the conditional sampling subsection to clarify our motivation. Below, we summarize the main points (see also lines 125 to 145 in the updated paper):

• Classifier-guided diffusion models were motivated by Eq. 4. In the context of diffusion models, it is necessary for eq 4 to hold at all noise levels. As a result, the correct likelihood at all noise levels is needed to obtain the conditional score.
• In practice, learning the likelihood correctly from data for all noise levels has proven to be challenging, which is reflected in poor quality of samples. This motivated the idea of weighting the terms, which also appears in the classifier-free guidance methods.
• Classifier-free diffusion is equivalent to Eq. 4 with no weighting. Hence in theory, the score is equivalent to the conditional score. However, in practice that leads to low quality samples. As a result, \omega is treated as a hyper-parameter and is often set to a value greater than one to improve quality at the expense of diversity.
• Our goal is to introduce an alternative strategy that does not require a trading off quality with diversity. This is achieved by relying on a feature space with properties shown in the paper: separation between classes, concentration within class, and Euclidean distance metric between classes. We show that for theoretically tractable datasets (like a mixture of Gaussians), this procedure can generate samples from the correct conditional. For real data, we show that samples are visually high quality and diverse. Note that this is different from previous work, where the algorithms which sample from correct Gaussian do not lead to high quality results, while the algorithms which return high quality samples do not sample from the correct conditional Gaussian.
• We acknowledge that for some applications, such as image generation, this tradeoff may not be so detrimental. However, for some applications (e.g. applied sciences) it can be critical not to reduce the entropy of the learned distribution. We believe our approach can be beneficial for such applications.

Regarding the first point: You are right that we do not provide theoretical guarantees, and this is an important goal for future work. However, in Sections 3 and 4 we lay out a theoretical framework to analyze the dynamics of the algorithms. We also demonstrate empirically that our algorithm samples correctly from conditionals in a Gaussian mixture, and generates high quality samples from real data. We argued that this success is due to the richness of the feature space which is learned simultaneously within the same score network.

Questions:

• We have now generated results from two more complex image datasets (Appendix D, figures 14 to 16). We used the same architecture as before (a UNet with 3 encoder and decoder bocks). We show that the model succeeds in drawing visually high quality samples from the correct class. Achieving SOTA on FID is a worthy goal, but the purpose of the current work is to demonstrate that we can learn a feature space, computed within the score network and using only the denoising objective, that supports effective conditional sampling. Reporting FID scores would require training our model on the Imagenet dataset. For that, we would need to increase the size of the architecture and include elements like attention blocks, BigGAN downsampling and other elements. This is not feasible given the time constraints of the rebuttal period, but we consider it a natural future direction for the work.

• As elaborated above, the motivation for introducing this alternative method is to avoid the trade-off between diversity and quality that is found in previous methods. Moreover, our results indicate that the learned feature embedding captures interesting properties of the representation within the score network. Given the complexity of diffusion models, we think it is valuable to have a simple working model which consists of just one network which lends itself to more analysis. This can be an avenue for better understanding the inner workings of diffusion models.

• For uncentered data (which is the case here) the correlation matrix depends upon the mean. For positive images, especially when they are stationary, the largest variance is often in the direction of the mean, which is thus the dominant principal component. Although this is not mathematically guaranteed (one can define or construct a positive random process whose correlation matrix has a principal component not aligned with the mean), it can be verified for our neural network channels.

Thanks for catching these typos, which we’ve fixed. Also, thanks for mentioning the Bradley & Nakkiran paper. We were not aware of this work, which is relevant to our motivation. We have included it in the background section.

Thank you for your constructive comments and questions. We have updated the paper, and refer below to the specific changes.

Weaknesses:

1. We agree that testing our method on a more complex dataset would strengthen the paper. Following the suggestion of reviewer hvBa, we have first applied our method CIFAR dataset of small images - results are shown in figure 16 (Appendix D) of the updated paper. We have also applied it to larger images in a dataset constructed from the union of 6 commonly used datasets: CelebaHQ (30k images), a subset of LSUN bedroom class (30k), AFHQ (16k), Flowers102(8k), Stanford cars(16k), and a subset of NA_birds (30k). These results are now shown in figure 14 and 15 (Appendix D). Note that the architecture in these experiments is unchanged from the original submission: a UNet with 3 encoder and decoder blocks. In both of these experiments, the model succeeds in drawing samples from the correct class.
2. Regarding the ImageNet dataset: this is a much larger dataset, with more complex images, but with a relatively small number of images (~1k) in each class. As such, we would need to increase the size of our network substantially, and consider addition of additional architectural features such as attention blocks, BigGANs down-sampling blocks, etc. We’ve not had sufficient time for this during the rebuttal period.

Questions:

1. We have added a new diagram of our model architecture, illustrating how $x$ and $x’$ are fed into the network (see appendix A). Note a primary feature of our method is that there is no additional encoder - the model consists of a single UNet. Comparing our feature space to the latent space of stable diffusion: 1) Our features are extracted from the score network itself (means of channels), which are then imposed within the same score network operating on a noisy input, whereas stable diffusion uses a separate encoder to map the image to a latent space in which the diffusion/denoising takes place, and 2) Our model is trained on a single denoising objective, whereas in stable diffusion the features are extracted according to a reconstruction loss function (or some other variation). Our construction gives rise to an important property of the feature space: it defines a Euclidean metric for the conditional distributions. For example, two conditional densities which are far in terms of symmetrized KL divergence have feature vectors that are far in Euclidean distance. Other semantic differences: our feature space is a multi-scale representation so it can capture differences between class features at different scales. Overall, this feature space is quite rich, and we believe offers interesting opportunities for further study.
2. The suggestion based on Fig. 3 is an interesting one. We do have initial results suggesting that conditional sampling can be done in fewer steps, but these need to be verified/elaborated with additional experiments.
3. We’ve re-built figure 4. In particular, the middle panel now shows the ratio of the between-class distances and the within-class standard deviations, which is the relevant statistical decision theory quantity for separation of the classes. Both plots are shown with the same units, and (as before) we show plots for pairs of classes (thin lines), and the average over all pairs. The variability across pairs is not due to insufficient training, but arises naturally from the variable similarity between classes.

Thank you for your encouraging comments and constructive questions.

Weaknesses/Questions: In our implementation, the feature vector is constructed from spatial averages of channels in the last layers of each of the U-Net blocks. Our UNet has the following architecture:

• 3 encoder blocks with {64,128,256} channels, respectively.
• middle block with 512 channels.
• 3 decoder blocks with {256,128,64} channels, respectively.

This gives a feature vector with a total of 1344 dimensions (please see new architecture diagram in updated appendix A).

The middle column of Figure 4 (now updated) suggests that we might be able to reduce the dimensionality of the feature space by eliminating some of the early and late channels, which contribute less to the separation of classes. More generally, we are interested in reducing to a minimal dimensionality, and in understanding the relationship between this minimal dimensionality and the complexity of the data distribution. The feature vector must be sufficiently rich in order to separate elements of different classes. With our current experiments, overparameterization of the feature vector does not seem to affect the learning or generation negatively.