A Superposição de Modelos de Difusão Usando o Estimador de Densidade de Itô

The Superposition of Diffusion Models Using the Itô Density Estimator

December 23, 2024

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

cs.AI

Resumo

A explosão cambriana de modelos de difusão pré-treinados facilmente acessíveis sugere uma demanda por métodos que combinem vários modelos de difusão pré-treinados diferentes sem incorrer no significativo ônus computacional de re-treinar um modelo combinado maior. Neste artigo, formulamos o problema de combinar múltiplos modelos de difusão pré-treinados na etapa de geração sob um novo framework proposto denominado superposição. Teoricamente, derivamos a superposição a partir de princípios rigorosos derivados da célebre equação de continuidade e projetamos dois novos algoritmos feitos sob medida para combinar modelos de difusão no SuperDiff. O SuperDiff aproveita um novo estimador de densidade de Itô escalável para a log-verossimilhança da EDS de difusão, o que não gera nenhum custo adicional em comparação com o estimador bem conhecido de Hutchinson necessário para cálculos de divergência. Demonstramos que o SuperDiff é escalável para grandes modelos de difusão pré-treinados, pois a superposição é realizada exclusivamente por meio de composição durante a inferência, e também desfruta de uma implementação sem complicações, pois combina diferentes campos vetoriais pré-treinados por meio de um esquema automatizado de reponderação. Notavelmente, mostramos que o SuperDiff é eficiente durante o tempo de inferência e imita operadores de composição tradicionais, como o OR lógico e o AND lógico. Demonstramos empiricamente a utilidade do uso do SuperDiff para gerar imagens mais diversas no CIFAR-10, edição de imagem condicionada por prompt mais fiel usando Diffusion Estável e melhoria no design de estruturas de proteínas incondicionalmente de novo. 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