Die Überlagerung von Diffusionsmodellen unter Verwendung des Itô-Dichteschätzers.

The Superposition of Diffusion Models Using the Itô Density Estimator

December 23, 2024

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

cs.AI

Zusammenfassung

Die kambrische Explosion leicht zugänglicher vortrainierter Diffusionsmodelle deutet auf einen Bedarf an Methoden hin, die mehrere verschiedene vortrainierte Diffusionsmodelle kombinieren, ohne die signifikante Rechenlast eines erneuten Trainings eines größeren kombinierten Modells zu verursachen. In diesem Artikel formulieren wir das Problem der Kombination mehrerer vortrainierter Diffusionsmodelle in der Generierungsphase unter einem neu vorgeschlagenen Rahmen namens Superposition. Theoretisch leiten wir Superposition aus strengen Grundprinzipien ab, die aus der bekannten Kontinuitätsgleichung stammen, und entwerfen zwei neuartige Algorithmen, die speziell für die Kombination von Diffusionsmodellen in SuperDiff entwickelt wurden. SuperDiff nutzt einen neuen skalierbaren It\^o-Dichteschätzer für die Log-Likelihood der Diffusions-SDE, der im Vergleich zum bekannten Hutchinson-Schätzer für Divergenzberechnungen keine zusätzlichen Kosten verursacht. Wir zeigen, dass SuperDiff für große vortrainierte Diffusionsmodelle skalierbar ist, da die Superposition ausschließlich durch Komposition während der Inferenz durchgeführt wird. Zudem bietet es eine problemlose Implementierung, da es verschiedene vortrainierte Vektorfelder durch ein automatisches Neugewichtungsschema kombiniert. Beachtenswert ist, dass SuperDiff effizient während der Inferenzzeit ist und traditionelle Kompositionsoperatoren wie das logische ODER und das logische UND nachahmt. Wir zeigen empirisch den Nutzen von SuperDiff bei der Erzeugung vielfältigerer Bilder auf CIFAR-10, bei der treuen prompt-konditionierten Bildbearbeitung mit Stable Diffusion und bei der verbesserten bedingungslosen de-novo-Strukturdesign von Proteinen. 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