La Sovrapposizione dei Modelli di Diffusione Utilizzando l'Estimatore della Densità di Itô

The Superposition of Diffusion Models Using the Itô Density Estimator

December 23, 2024

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

cs.AI

Abstract

L'esplosione cambriana dei modelli di diffusione pre-addestrati facilmente accessibili suggerisce una domanda di metodi che combinino diversi modelli di diffusione pre-addestrati senza incorrere nel significativo onere computazionale del ri-addestramento di un modello combinato più grande. In questo articolo, formuliamo il problema di combinare diversi modelli di diffusione pre-addestrati nella fase di generazione sotto un nuovo framework proposto chiamato sovrapposizione. Teoricamente, deriviamo la sovrapposizione da rigorosi principi fondamentali derivanti dall'acclamata equazione di continuità e progettiamo due nuovi algoritmi su misura per combinare modelli di diffusione in SuperDiff. SuperDiff sfrutta un nuovo stimatore scalabile della densità di It\^o per la log-verosimiglianza dell'SDE di diffusione che non comporta alcun costo aggiuntivo rispetto al ben noto stimatore di Hutchinson necessario per i calcoli di divergenza. Dimostriamo che SuperDiff è scalabile per grandi modelli di diffusione pre-addestrati poiché la sovrapposizione viene eseguita esclusivamente tramite composizione durante l'inferezza, e gode anche di una facile implementazione poiché combina diversi campi vettoriali pre-addestrati attraverso un sistema di ripesatura automatizzato. In particolare, mostriamo che SuperDiff è efficiente durante il tempo di inferenza e imita operatori di composizione tradizionali come il logico OR e il logico AND. Dimostriamo empiricamente l'utilità dell'utilizzo di SuperDiff per generare immagini più diverse su CIFAR-10, per un editing di immagini condizionato a prompt più fedele utilizzando Stable Diffusion e per un miglioramento della progettazione di strutture di proteine incondizionata e completamente nuova. 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