Resumen Técnico: Expansión de Dominio de Difusión (DDE)

Enunciado del Problema

Los modelos de difusión preentrenados han logrado resultados de vanguardia en la generación de imágenes, audio y otras modalidades de datos. Sin embargo, volver a entrenar estos modelos desde cero para manejar dominios de salida más grandes (por ejemplo, imágenes panorámicas, pistas de audio largas) o entradas de condicionamiento más complejas (por ejemplo, múltiples ubicaciones de objetos) es computacionalmente costoso. Los enfoques existentes para expandir los dominios generativos a menudo dependen de mecanismos de coordinación algorítmicos fijos (por ejemplo, MultiDiffusion) que no requieren entrenamiento adicional pero están limitados a tareas específicas y carecen de flexibilidad. Otros métodos, como el entrenamiento de módulos difusores dinámicos, a menudo requieren modalidades fijas durante el tiempo de entrenamiento o no logran generalizarse a dominios más grandes que aquellos observados durante el entrenamiento.

El desafío central abordado por este artículo es cómo extender eficientemente el dominio generativo de los modelos de difusión preentrenados para generar objetos más grandes y manejar condicionamientos más complejos sin reentrenar los modelos base, asegurando al mismo tiempo que la solución se generalice a instancias más grandes que aquellas vistas durante el entrenamiento del coordinador.

Metodología: Expansión de Dominio de Difusión (DDE)

Los autores proponen la Expansión de Dominio de Difusión (DDE), un marco que emplea una red "coordinadora" compacta y entrenable para orquestar las salidas de modelos de difusión preentrenados.

Formulación General

Dado un modelo de difusión condicional preentrenado $p(x|y)$ con una red de eliminación de ruido $D(x, y, t)$, el objetivo es generar un objeto expandido $X^{[L]}$ condicionado por una entrada expandida $Y^{[L]}$, donde $L$ representa la escala (por ejemplo, número de parches o etiquetas de condicionamiento).

1. Descomposición: El par expandido $(X^{[L]}, Y^{[L]})$ se descompone en $L$ partes más pequeñas $([x_1, \dots, x_L], [y_1, \dots, y_L])$ utilizando una función de mapeo $F$. Estas partes caen dentro del dominio original del modelo base.
2. Generación Base: Cada parte $x_i$ es procesada por el eliminador de ruido preentrenado $D(x_i, y_i, t)$.
3. Coordinación: Una red coordinadora entrenable $C[\cdot]$ toma la secuencia de salidas desruidadas $[D(x_i, y_i, t)]_{i=1}^L$ y la información de condicionamiento $[y_i]_{i=1}^L$ como entrada. Aprende a coordinar estas salidas para producir una estimación coherente del objeto completo expandido $X^{[L]}(0)$.
4. Objetivo de Entrenamiento: El coordinador se entrena minimizando el error de eliminación de ruido del eliminador de ruido compuesto $D^{[L]}$ en un conjunto de datos de tamaño $L_{train}$. Crucialmente, el método está diseñado para generalizarse a tamaños de prueba $L_{test} \ge L_{train}$.

Arquitectura del Coordinador

Los autores utilizan una arquitectura de Transformador de Visión (ViT) para el coordinador, adaptada de Peebles & Xie (2023).

• Procesamiento de Entrada: Los parches desruidados de los modelos preentrenados se "fragmentan" y se aumentan con codificaciones posicionales globales relativas al objeto más grande. Los autores emplean Incrustaciones Posicionales Rotativas (RoPE) para manejar dimensiones espaciales más grandes de manera efectiva.
• Reconciliación: Durante el proceso hacia adelante, el transformador procesa todos los parches. Para manejar parches superpuestos (utilizados para garantizar la continuidad), la salida final se construye promediando los valores en las regiones superpuestas, inspirado en el método de muestreo MultiDiffusion.
• Eficiencia: El coordinador es significativamente más pequeño que el modelo base (por ejemplo, 66M vs. 405M parámetros en la generación de música) y requiere menos épocas de entrenamiento.

Contribuciones Clave

1. Marco Novel: Introducción de DDE, un método que utiliza un coordinador pequeño y entrenable para extender las capacidades generativas de los modelos de difusión preentrenados a objetos más grandes y condicionamientos más complejos.
2. Capacidad de Generalización: Demostración de que el coordinador puede generalizarse para generar objetos y manejar entradas de condicionamiento de mayor escala que aquellas presentes en sus datos de entrenamiento ($L_{test} \ge L_{train}$).
3. Diseño de Arquitectura: Validación de una arquitectura basada en ViT como coordinador adecuado, que aprende dependencias de largo alcance mientras impone consistencia local a través del promediado de superposiciones.
4. Eficiencia: El enfoque logra una generación de alta calidad con significativamente menos parámetros y tiempo de entrenamiento en comparación con reentrenar modelos base o utilizar arquitecturas de coordinador más grandes.

Resultados Experimentales

Los autores evaluaron DDE en tres dominios:

1. Generación de Pistas de Audio Largas (Slakh2100)

• Tarea: Generar pistas musicales de longitud $10l$ (120 segundos) utilizando un modelo base entrenado en longitud $l$ (12 segundos).
• Líneas Base: Concatenación y MultiDiffusion.
• Resultados: DDE logró una Distancia de Audio Fréchet (FAD) de 2.112 para pistas de 4l y 2.142 para pistas de 10l. Esto superó significativamente a MultiDiffusion (4.796) y a la Concatenación (4.596) para la tarea de 10l. El coordinador (66M parámetros) fue mucho más pequeño que el modelo base (405M).

2. Generación de Imágenes Condicionales (CLEVR)

• Tarea: Generar imágenes condicionadas por múltiples ubicaciones de objetos ($L$ etiquetas), donde el modelo base solo maneja una ubicación.
• Líneas Base: RRR (Reducir, Reutilizar, Reciclar) y MultiDiffusion.
• Resultados: DDE demostró una precisión superior en el cumplimiento de múltiples restricciones de condicionamiento. Para 5 condicionamientos, DDE logró una precisión del 44.5% (muestreador Heun), en comparación con el 33.6% de RRR y el 36.3% de MultiDiffusion. El coordinador se generalizó desde el entrenamiento en $L_{train}=2$ hasta la prueba en $L_{test}=5$.

3. Generación de Imágenes Satelitales

• Tarea: Generar grandes imágenes satelitales condicionadas por mapas esquemáticos.
• Líneas Base: MultiDiffusion.
• Resultados: DDE logró una puntuación FID más baja que MultiDiffusion tanto para tamaños dentro de la distribución ($N=96$, FID 31.75 vs 37.81) como fuera de la distribución ($N=128$, FID 27.37 vs 35.01). La inspección visual mostró que DDE producía imágenes más coherentes sin las costuras visibles observadas en las muestras de MultiDiffusion.

Significado y Afirmaciones

El artículo afirma que DDE ofrece una solución práctica y eficiente para escalar modelos de difusión sin el costo prohibitivo de reentrenar. Al aprender a coordinar procesos preentrenados, el método permite:

• Expansión de Dominio: Generación de objetos más grandes (audio más largo, imágenes más grandes) y manejo de condiciones más complejas (múltiples etiquetas) más allá de las capacidades originales del modelo.
• Generalización: La capacidad de extrapolar a escalas más grandes que aquellas vistas durante el entrenamiento del coordinador.
• Calidad y Consistencia: El coordinador basado en ViT aprende efectivamente dependencias de largo alcance, resolviendo conflictos entre pasos de eliminación de ruido locales para producir salidas globalmente consistentes.

Los autores señalan que la limitación principal es el requisito de datos adicionales y tiempo para entrenar el modelo coordinador, aunque este costo es sustancialmente menor que entrenar un modelo de difusión completo desde cero. Se propone trabajo futuro para aplicar DDE a la generación de secuencias de proteínas y la creación de videos a partir de clips cortos.