Суперпозиция моделей диффузии с использованием оценщика плотности Ито

The Superposition of Diffusion Models Using the Itô Density Estimator

December 23, 2024

Авторы: Marta Skreta, Lazar Atanackovic, Avishek Joey Bose, Alexander Tong, Kirill Neklyudov

cs.AI

Аннотация

Взрыв Кембрийской эры легко доступных предварительно обученных моделей диффузии подразумевает потребность в методах, которые объединяют несколько различных предварительно обученных моделей диффузии без значительной вычислительной нагрузки на повторное обучение более крупной объединенной модели. В данной статье мы формулируем проблему объединения нескольких предварительно обученных моделей диффузии на этапе генерации в рамках новой предложенной концепции, названной суперпозиция. Теоретически мы выводим суперпозицию из строгих первопринципных положений, вытекающих из знаменитого уравнения непрерывности, и разрабатываем два новых алгоритма, специально созданных для объединения моделей диффузии в SuperDiff. SuperDiff использует новый масштабируемый оценщик плотности Ито для логарифма правдоподобия уравнения диффузии SDE, что не влечет за собой дополнительных накладных расходов по сравнению с хорошо известным оценщиком Хатчинсона, необходимым для вычислений расхождений. Мы демонстрируем, что SuperDiff масштабируется на большие предварительно обученные модели диффузии, поскольку суперпозиция выполняется исключительно через композицию во время вывода, и также обладает простой реализацией, поскольку объединяет различные предварительно обученные векторные поля через автоматизированную схему повторного взвешивания. Особенно отметим, что SuperDiff эффективен во время вывода и имитирует традиционные операторы композиции, такие как логическое ИЛИ и логическое И. Мы эмпирически демонстрируем полезность использования SuperDiff для генерации более разнообразных изображений на CIFAR-10, более точного редактирования изображений с использованием Stable Diffusion и улучшенного безусловного проектирования структуры белков. 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