Het superponeren van diffusiemodellen met behulp van de Itô-dichtheidsschatting.

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

Samenvatting

De Cambrium-explosie van gemakkelijk toegankelijke vooraf getrainde diffusiemodellen suggereert een vraag naar methoden die meerdere verschillende vooraf getrainde diffusiemodellen combineren zonder de aanzienlijke rekenlast van het opnieuw trainen van een groter gecombineerd model. In dit artikel stellen we het probleem van het combineren van meerdere vooraf getrainde diffusiemodellen in de generatiefase voor onder een nieuw voorgesteld kader genaamd superpositie. Theoretisch leiden we superpositie af van rigoureuze eerste principes voortkomend uit de gevierde continuïteitsvergelijking en ontwerpen we twee nieuwe algoritmes op maat gemaakt voor het combineren van diffusiemodellen in SuperDiff. SuperDiff maakt gebruik van een nieuwe schaalbare It\^o-dichtheidsschatting voor de log-likelihood van de diffusie SDE die geen extra overhead met zich meebrengt in vergelijking met de bekende schatter van Hutchinson die nodig is voor divergentieberekeningen. We tonen aan dat SuperDiff schaalbaar is naar grote vooraf getrainde diffusiemodellen omdat superpositie uitsluitend wordt uitgevoerd door samenstelling tijdens inferentie, en ook gemakkelijk te implementeren is omdat het verschillende vooraf getrainde vectorvelden combineert door middel van een geautomatiseerd herwegingsmechanisme. Opmerkelijk is dat we aantonen dat SuperDiff efficiënt is tijdens inferentie en traditionele samenstellingsoperatoren nabootst zoals de logische OF en de logische EN. We tonen empirisch het nut aan van het gebruik van SuperDiff voor het genereren van meer diverse afbeeldingen op CIFAR-10, meer trouwe prompt-geconditioneerde beeldbewerking met Stable Diffusion, en verbeterd onvoorwaardelijk de novo-structuurontwerp van eiwitten. 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