La superposition des modèles de diffusion en utilisant l'estimateur de densité d'Itô

The Superposition of Diffusion Models Using the Itô Density Estimator

December 23, 2024

Auteurs: Marta Skreta, Lazar Atanackovic, Avishek Joey Bose, Alexander Tong, Kirill Neklyudov

cs.AI

Résumé

L'explosion cambrienne des modèles de diffusion pré-entraînés facilement accessibles suggère une demande de méthodes combinant plusieurs modèles de diffusion pré-entraînés différents sans supporter la charge computationnelle significative de ré-entraîner un modèle combiné plus grand. Dans cet article, nous formulons le problème de combiner plusieurs modèles de diffusion pré-entraînés à l'étape de génération sous un nouveau cadre proposé appelé superposition. Théoriquement, nous dérivons la superposition de premiers principes rigoureux découlant de l'équation de continuité célèbre et concevons deux nouveaux algorithmes sur mesure pour combiner des modèles de diffusion dans SuperDiff. SuperDiff tire parti d'un nouvel estimateur d'intensité Itô scalable pour la log-vraisemblance de l'EDS de diffusion qui n'entraîne aucun surcoût supplémentaire par rapport à l'estimateur bien connu de Hutchinson nécessaire pour les calculs de divergence. Nous démontrons que SuperDiff est scalable pour de grands modèles de diffusion pré-entraînés car la superposition est effectuée uniquement par composition lors de l'inférence, et bénéficie également d'une implémentation sans douleur car elle combine différents champs vectoriels pré-entraînés grâce à un schéma de réaffectation automatisé. Notamment, nous montrons que SuperDiff est efficace pendant le temps d'inférence, et imite des opérateurs de composition traditionnels tels que le OU logique et le ET logique. Nous démontrons empiriquement l'utilité de l'utilisation de SuperDiff pour générer des images plus diverses sur CIFAR-10, une édition d'image conditionnée par prompt plus fidèle en utilisant Stable Diffusion, et une conception améliorée de structures de protéines de novo inconditionnelles. https://github.com/necludov/super-diffusion

English

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, and improved unconditional de novo structure design of proteins. https://github.com/necludov/super-diffusion