The Superposition of Diffusion Models Using the Itô Density Estimator

Abstract: The Cambrian explosion of easily accessible pre-trained diffusion models suggests a demand for methods that combine multiple different pre-trained diffusion models without incurring the significant computational burden of re-training a larger combined model. In this paper, we cast the problem of combining multiple pre-trained diffusion models at the generation stage under a novel proposed framework termed superposition. Theoretically, we derive superposition from rigorous first principles stemming from the celebrated continuity equation and design two novel algorithms tailor-made for combining diffusion models in SuperDiff. SuperDiff leverages a new scalable It^o density estimator for the log likelihood of the diffusion SDE which incurs no additional overhead compared to the well-known Hutchinson's estimator needed for divergence calculations. We demonstrate that SuperDiff is scalable to large pre-trained diffusion models as superposition is performed solely through composition during inference, and also enjoys painless implementation as it combines different pre-trained vector fields through an automated re-weighting scheme. Notably, we show that SuperDiff is efficient during inference time, and mimics traditional composition operators such as the logical OR and the logical AND. We empirically demonstrate the utility of using SuperDiff for generating more diverse images on CIFAR-10, more faithful prompt conditioned image editing using Stable Diffusion, as well as improved conditional molecule generation and unconditional de novo structure design of proteins. https://github.com/necludov/super-diffusion

• The paper introduces SuperDiff, which composes multiple pre-trained diffusion models without re-training, enabling innovative joint inference.
• It leverages an Itô density estimator to bypass expensive divergence computations, thereby reducing computational load and variance.
• Empirical results in image and protein generation demonstrate enhanced diversity and fidelity, highlighting SuperDiff’s practical benefits across domains.

Superposition of Diffusion Models with the Itô Density Estimator

The paper "The Superposition of Diffusion Models Using the Itô Density Estimator" presents a novel framework for combining pre-trained diffusion models during the generation stage without re-training, termed SuperDiff. The approach leverages the concept of superposition from physical systems to enable the composition of multiple generative models through joint inference processes. This is achieved while maintaining efficiency and scalability, crucial for handling large pre-trained diffusion models prevalent in contemporary image and protein generation tasks.

Theoretical Framework

The superposition problem is addressed through a rigorous approach rooted in the continuity equation that governs diffusion processes. The authors introduce a novel density estimator, grounded in Itô's lemma, which provides an efficient method to estimate densities during the reverse-time simulation of diffusion models. Unlike traditional methods requiring expensive divergence computations of drift vector fields, the proposed estimator avoids these computational burdens. This not only reduces computational load but mitigates variance, making it particularly viable for large models.

SuperDiff: Two Modes of Diffusion Combination

The paper introduces SuperDiff, a method for the compositional inference of generative models, which can operate in two distinct modes likened to logical operators:

1. Logical OR (Mixture of Densities): This mode generates samples from a mixture of densities attributed to the component models, effectively superimposing their individual outputs. It offers flexibility in sample generation, allowing the leveraging of multiple models for enhanced performance in terms of image diversity and fidelities in datasets like CIFAR-10.
2. Logical AND (Equal Density Sampling): Contrarily, this mode aims to produce samples lying on an equal density locus for all models involved. It achieves this by solving a system of linear equations that ensure equal change in density contributions from each model, a unique approach compared to traditional ensemble methods.

Empirical Validation and Implications

Empirical results demonstrate the efficacy of SuperDiff through image generation experiments using CIFAR-10 and Stable Diffusion, as well as protein backbone generation tasks. SuperDiff outperforms individual diffusion models trained on disjoint datasets, often even surpassing models trained on the union of datasets, which underscores its potential in maximizing utility from pre-trained models.

For protein generation, SuperDiff's capability to exploit compositional benefits furthers its practical applications, enhancing designability and producing diverse and novel protein structures. This is particularly vital in contexts such as drug discovery, where the designability and novelty of protein structures are critical.

Implications and Future Work

The introduction of SuperDiff marks significant advancements in the field of generative modeling, especially for applications where utilizing existing pre-trained models optimally is preferred over starting new training procedures. The practical implications extend across various domains, including bioinformatics and computational creativity.

However, the application of SuperDiff is computationally bound by model evaluations rather than the cheap manipulation of internal weights, suggesting future work could involve developing lightweight extensions or hybrid methods that synergize with model weight efficiencies. Additionally, further exploration of the Itô density estimator's applicability across various diffusion settings could yield broader applications in generative AI.

In conclusion, the SuperDiff framework not only introduces an innovative perspective on model superposition but also paves the way for scalable and efficient generative processes, promising substantial impacts on how generative models are utilized in various scientific and industrial applications.