Feature-guided score diffusion for sampling conditional densities

Authors: Zahra Kadkhodaie, Stéphane Mallat, Eero P Simoncelli
Primary area: generative models
Keywords: Conditional sampling, Representation learning, Score diffusion

Abstract:
Score diffusion methods can learn probability densities from samples. The score of the noise-corrupted density is estimated using a deep neural network, which is then used to  iteratively transport a Gaussian white noise density to a target density. Variants for conditional densities have been developed, but correct estimation of the corresponding scores is difficult.
We avoid these difficulties by introducing an algorithm that operate by projecting the score onto the target class mean in a learned feature space.
The features and the projected score are computed using the same network, which is trained by optimizing a single denoising loss. 
Learned feature vectors of same-class images are tightly clustered relative to those of different classes. 
We show that feature class centroids provide a low-dimensional Euclidean embedding of the class conditional densities. 
We demonstrate that, when trained on a dataset of mixed image classes,
this projected score can generate high quality and diverse samples from the conditioning class.
Conditional generation can be performed using feature vectors interpolated between those of the training set, demonstrating out-of-distribution generalization.

Decision: reject
Reviewer scores: 6, 8, 5, 3

Reviews and discussion:
<thread>
creation date:
1729930392617

thread_type:
official_review

<thread_object>
creator: reviewer_epc9

<content>
summary:
This paper focuses on the problem of conditional generation with diffusion models, and specifically class-conditional generation. The authors claim that commonly used class-conditional sampling methods are inexact, and thus propose a new method for class-conditional sampling with diffusion models. Their method, taking inspiration from a GMM distribution, attempts to train a diffusion model that is biased to a certain class based on the features of some given conditioning image. The authors demonstrate several examples of their method.

strengths:
- Conditional sampling is a key problem, where improvements may have a high impact.
- The proposed solution is novel, to the best of my knowledge.

weaknesses:
- The paper's motivation is not sufficiently clear. While adding noise and conditioning may not commute in the general case, I believe the operation do commute in the case where each sample image has only one corresponding class, which is the common case in class-conditional sampling.
- Many conditional diffusion models in the literature are trained with the condition as an input to the network, alleviating the concern regarding the inexactness of using Eq. (4) for sampling. While such method require training a conditional model, but so does the method proposed in the paper, making the advantage here unclear.
- The experiment section is missing a meaningful comparison with alternative conditional generation methods, despite those methods being well established and mentioned in the text. Could the authors to compare their method to existing conditional diffusion models, and highlighting specific advantages or differences?
- The proposed method relies on access to an example image from a given class for sampling, making generation more cumbersome. this concern is not addressed in the paper. Perhaps this could be meaningfully discussed in a limitations section.
- Please address several typos in the discussion section.

questions:
Please see concerns under Weaknesses.
- Why have the authors chosen not to use an established class conditional dataset, such as CIFAR10 or ImageNet for the experiments in the paper?
- Can the authors elaborate why adding noise and conditioning do not commute? I believe such a claim requires proof.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730581962242

thread_type:
official_review

<thread_object>
creator: reviewer_hvba

<content>
summary:
The most common approach to guide diffusion models is by using score guidance, whether in the form of classifier guidance or in the form of classifier free guidance. The paper claims that this approach relies on an approximation and is thus inaccurate. As an alternative, it proposes a simple method for training a conditional diffusion model. The approach is motivated intuitively and tested empirically.

strengths:
- Classifier free guidance (CFG) is the most dominant approach for guiding diffusion models today, even though it is known to lead to biased densities. Several recent papers analyzed the drawbacks of the approach from a theoretical standpoint. However the topic of designing good practical alternatives to CFG is still under-explored. This paper attempts to fill this gap, which is undoubtedly an important goal.
- The paper presents clear intuition and empirically validates that the assumptions underlying the proposed construction hold.

weaknesses:
- The whole motivation of the paper is to propose an alternative to existing guidance methods. However, it does not provide theoretical guarantees that the approach samples from the conditional distribution. And it also does not provide any empirical evidence that the proposed approach outperforms the standard way of conditioning diffusion models. Specifically, it does not compare the sampling quality to that obtained with a conditional denoiser (with the common conditioning mechanism for U-Nets), neither without nor with CFG. It also does not present any quantitative measure of sample quality (e.g. FID) on natural images, and does not present results on popular datasets like CIFAR or ImageNet, which prevents from evaluating whether the proposed method leads to SOTA results. As such, it is impossible to evaluate the effectiveness of the proposed method.
- I believe that the motivation presented in the paper is inaccurate and somewhat misleading. Using score guidance (as in Eq. (4)) should theoretically lead to accurate results. What is inaccurate in CFG is that when taking the guidance parameter w to be greater than 1, it does not sample from the titled distribution p(x)p(y|x)^(1+w) from which it attempts to sample, but rather from a different distribution. See for instance the closed form expressions for the Gaussian setting in [1] (Eqs. (11),(12)), where the distributions of samples obtained with CFG-DDPM and CFG-DDIM are incorrect only when the CFG parameter (gamma) is greater than 1, but are correct when it equals 1. In particular, the statement in L147-148 that Chidambaram et al. (2024) proved the inaccuracy of guidance on Gaussian mixtures is incorrect. They proved the inaccuracy of CFG only when the guidance parameter w is sufficiently large. When w=1 the process is theoretically accurate.


[1] Bradley and Nakkiran, "Classifier-Free Guidance is a Predictor-Corrector", arxiv 2408.09000.

questions:
- Can the authors comment on how the method compares to regular guidance on common datasets, like CIFAR-10, in terms of FID?
- Can the authors comment on the motivation for suggesting an alternative to regular guidance (second weakness stated above)?
- L216: Why are the activation layer averages close to the first principal components of the network channels? Can the authors provide a proof or a reference to a paper showing this?

Typos:
- L216: principle -> principal
- L360: "some some"
- L429: missing parentheses around (m2-m1)
- L517: remove "And most"
- L518: mixeed -> mixed
</content>
</thread_object>

</thread>
<thread>
creation date:
1730705188626

thread_type:
official_review

<thread_object>
creator: reviewer_wjvb

<content>
summary:
This paper introduces a guided score diffusion method that samples from class conditional distributions by calculating a projected score based on feature vectors, avoiding the need to directly estimate the scores of those densities. The method employs a neural network trained to minimize a single denoising loss, with the feature vector defined as spatial averages from selected layers, leading to a Euclidean embedding of class conditional probabilities. It effectively clusters learned features around their centroids, allowing for accurate sampling from target conditional distributions and facilitating smooth transitions between classes through linear combinations of mean feature vectors.

strengths:
This is a well-written paper. Both score and feature vectors are represented with the same network. The learned feature vectors cluster around their centroids, which enhances the accuracy of sampling rom the conditional probability density. The method enables gradual transitions of the images between classes through linear interpolation of mean feature vectors. The experimental results show that a diffusion algoriothm based on the projected score provides an accurate sampling of conditional probabilities.

weaknesses:
The authors provided a way to build the feature vectors that share the same network weights as the score function. It is not clear how to determine the feature vector dimension.

questions:
How does the feature dimension affects the learning and generation?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730708716040

thread_type:
official_review

<thread_object>
creator: reviewer_9mda

<content>
summary:
This paper introduces a novel guided score-based diffusion model that does not require any additional structural modifications. Instead, it extracts features from key layers of its own score estimation model (i.e. Unet) and applies spatial averaging, enabling guidance control in the feature space. Unlike methods such as classifier-free guidance, this feature-guided diffusion does not rely on likelihood estimation, allowing for a more accurate estimation of conditional density.

strengths:
1. The paper is well-structured, easy to read, and highly innovative, introducing a projected score embedded in feature space. This embedding is not only straightforward to obtain (directly extracted from the score estimation model) but also adheres to Euclidean interpolation properties.

2. The model successfully achieves conditional generation in a mixture of Gaussian distributions, demonstrating that the feature-guided score diffusion model can accurately capture conditional density—an ability lacking in many other approaches.

3. Compared to mixture models, the feature-guided score diffusion model shows stronger concentration and separation across different classes within the embedding space.

weaknesses:
The dataset used in experiments is overly simple. The training dataset is derived by cropping 1700 images into 234k patches. Although the patches are non-overlapping, the data distribution for each class lacks sufficient diversity. Experiments on a more complex dataset, like ImageNet, would strengthen the paper’s validity.

questions:
1. What is the difference between feature space and latent space in stable diffusion? While the latter requires an additional encoder, how does the former feed $x$ and $x'$ into the U-Net?

2. In Fig. 3 (left), the feature-guided model shows strong performance across all noise levels. Does this suggest that fewer NFE (Number of Function Evaluations) could be used during conditional generation? Designing experiments to verify this could strengthen the paper.

3. In Fig. 4 (middle), were the two metrics normalized? Otherwise, this comparison might not be on the same scale. Additionally, the variance in image feature vectors appears large—is this due to insufficient training? Please clarify further.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732897099779

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global comment

comment:
While all reviewers appreciated the novelty of our work, the unresolved criticism by three reviewers was a lack of precise comparison with SOTA on sample quality (FID). We acknowledge that we did not present FID scores and that scaling up the model to do so is a valuable future direction. However, contributions of this work are aimed at other facets of the problem, orthogonal to improving SOTA on sample quality. More precisely, we demonstrated **two advantages** of our method over previous work, which are important in scientific applications: 
1. Previous method (both classifier-guidance and classifier-free guidance) have a trade-off between high quality samples vs sampling from the correct conditional density (reduced diversity). Specifically when $\omega = 1$ they sample from the correct density but the quality is poor. When $\omega \neq 1 $ the sample quality is good but they do not sample from the correct density. **Our method does not exhibit this trade-off**: It samples from the correct conditional in the tractable case of a Gaussian (fig 6) AND it returns high quality samples in the case of real images (figs 7, 11, 14, 16).
2. Additionally, our method is the first to **learn and use an embedding space within a denoising diffusion model**. One of the main disadvantages of diffusion models is that unlike VAEs they do not produce an explicit latent space representation. This is particularly important for scientific application where understanding the implicit learned density matters. We show the effectiveness of the learned representation by showing that it successfully guides the conditional sampling. The success is because the classes are concentrated and separated in the embedding space (fig 4 and 15) and define a Euclidean metric such that densities closer in KL divergence correspond to closer embeddings in feature space  (fig 5).

We believe these contributions - learning and using a representation in diffusion models - are significant and aligned with the topic of the International Conference on Learning Representations.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors present a method for score-based diffusion that doesn't depend on score guidance approximations. The four reviewers agreed that the problem is very relevant to current research directions and addresses a known issue in the field. The reviews were split two positive and two negative, and the paper overall was borderline, but below threshold for acceptance. In general, while the reviewers appreciate the relevance of the proposal, it wasn't clear to them that this particular way of addressing it provides a significant enough advance for ICLR. Issues were raised on the motivation and advantage over conditional generative methods. It was felt that the experiments did not provide enough compelling evidence in favor of the method as a practical alternative. The reviewers all engaged with the rebuttal, but three of four did not feel that their concerns were alleviated.