La superposición de modelos de difusión utilizando el estimador de densidad 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

Resumen

La explosión cámbrica de modelos de difusión pre-entrenados fácilmente accesibles sugiere una demanda de métodos que combinen múltiples modelos de difusión pre-entrenados diferentes sin incurrir en la significativa carga computacional de volver a entrenar un modelo combinado más grande. En este documento, planteamos el problema de combinar múltiples modelos de difusión pre-entrenados en la etapa de generación bajo un nuevo marco propuesto denominado superposición. Teóricamente, derivamos la superposición a partir de rigurosos principios fundamentales derivados de la célebre ecuación de continuidad y diseñamos dos nuevos algoritmos hechos a medida para combinar modelos de difusión en SuperDiff. SuperDiff aprovecha un nuevo estimador de densidad de Itô escalable para la verosimilitud del logaritmo de la EDS de difusión que no conlleva una sobrecarga adicional en comparación con el conocido estimador de Hutchinson necesario para cálculos de divergencia. Demostramos que SuperDiff es escalable para grandes modelos de difusión pre-entrenados ya que la superposición se realiza únicamente a través de composición durante la inferencia, y también disfruta de una implementación sencilla al combinar diferentes campos vectoriales pre-entrenados a través de un esquema de reponderación automatizado. Es destacable que SuperDiff es eficiente durante el tiempo de inferencia, y emula operadores de composición tradicionales como el OR lógico y el AND lógico. Demostramos empíricamente la utilidad de utilizar SuperDiff para generar imágenes más diversas en CIFAR-10, una edición de imágenes condicionada a la consigna más fiel utilizando Stable Diffusion, y una mejora en el diseño de estructuras de proteínas de novo incondicional. 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