Feature-guided score diffusion for sampling conditional densities (2410.11646v1)

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 guides the diffusion with a projected score. The projection pushes the image feature vector towards the feature vector centroid of the target class. The projected score and the feature vectors are learned by the same network. Specifically, the image feature vector is defined as the spatial averages of the channels activations in select layers of the network. Optimizing the projected score for denoising loss encourages image feature vectors of each class to cluster around their centroids. It also leads to the separations of the centroids. We show that these centroids provide a low-dimensional Euclidean embedding of the class conditional densities. We demonstrate that the algorithm 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.

• The paper proposes a novel diffusion method that leverages neural network embeddings to guide sampling from conditional densities.
• It modifies standard diffusion by projecting neural activations into a feature space, efficiently aligning samples with target distributions.
• Experimental results on Gaussian mixtures and image classes demonstrate high-quality, robust sampling in complex, high-dimensional settings.

Feature-Guided Score Diffusion for Sampling Conditional Densities

The paper under discussion presents a novel approach to sampling from conditional densities using score-based diffusion models. It introduces a feature-guided score diffusion technique that adapts the standard diffusion framework by incorporating information from a carefully constructed embedding space. This modification addresses key challenges in the estimation of conditional scores, particularly when applied to complex probability distributions.

Core Contributions

The authors propose an iterative diffusion sampling strategy that modifies the score at each step using features derived from neural network activations. This feature vector influences the diffusion by guiding it toward specific conditional densities through a projection in the feature space. The two central tenets of this methodology are:

1. Feature Space Projection: The feature vector is defined as the spatial averages of activations from certain layers of the network, acting as a summarizing statistical representation. This projected score directs the diffusion process to align samples with the target conditional distribution.
2. Single Network Parameterization: Both the feature vector and the score are computed by the same neural network. This implementation advantageously reduces computational demands and unifies the optimization under a single denoising loss criterion, leading to efficient learning of the conditional densities.

Methodological Insights

The proposed method excels in avoiding direct estimation of conditional scores, which is mathematically and computationally intensive. Instead, it leverages learned embeddings to differentiate between classes by ensuring the feature space achieves both concentration of within-class features and separation of class centroids. This is analytically supported by defining a Euclidean embedding where these feature vectors correspond closely to the actual statistical distances between conditional densities.

Experimental Validation

Through empirical evaluation on both Gaussian mixtures and natural image classes, the feature-guided diffusion demonstrates its proficiency in accurately sampling from target distributions. Notable results include:

• Gaussian Mixtures: It was shown that the proposed method successfully samples from Gaussian components within a mixture, circumventing the pitfalls encountered by previous likelihood-based guidance methods.
• Image Class Generation: High-quality and diverse image samples were generated, showcasing the method's effectiveness in practical visual domains. The interpolation capability within the embedding space further suggests robustness in handling out-of-distribution tasks.

Implications and Future Directions

This work presents significant implications for the design of conditional generative models. The feature-guided approach not only provides an accurate mechanism for sampling from specified conditional densities but also introduces a scalable solution likely to be applicable to other domains in AI. Future research may explore:

• Exploration of Alternative Feature Representations: Refinement of feature vector construction could further enhance the model, particularly for spatially variable datasets.
• Theoretical Investigations: Deeper understanding of the mathematical framework supporting stochastic interpolants in projected scores could illuminate further practical applications.

In conclusion, feature-guided score diffusion constitutes an innovative stride in conditional density sampling, achieving quality sampling without the inefficiencies associated with likelihood estimation. This framework's promising results advocate for its applicability in broader AI applications.