Modelos y plataformas de IA

Entendiendo los Modelos de Difusión: Una Inmersión Profunda en la Inteligencia Artificial Generativa

mm
Publicado el
Actualizado el
Añade Unite.AI a tus fuentes preferidas en Google

Los modelos de difusión han surgido como un enfoque poderoso en la inteligencia artificial generativa, produciendo resultados de vanguardia en la generación de imágenes, audio y video. En este artículo técnico en profundidad, exploraremos cómo funcionan los modelos de difusión, sus innovaciones clave y por qué han tenido tanto éxito. Cubriremos los fundamentos matemáticos, el proceso de entrenamiento, los algoritmos de muestreo y las aplicaciones de vanguardia de esta emocionante nueva tecnología.

Introducción a los Modelos de Difusión

Los modelos de difusión son una clase de modelos generativos que aprenden a desenoizar gradualmente los datos invirtiendo un proceso de difusión. La idea central es comenzar con ruido puro y refinarlo iterativamente en una muestra de alta calidad de la distribución objetivo.

Este enfoque se inspiró en la termodinámica no equilibrada – específicamente, el proceso de revertir la difusión para recuperar la estructura. En el contexto del aprendizaje automático, podemos pensar en ello como aprender a revertir la adición gradual de ruido a los datos.

Algunas ventajas clave de los modelos de difusión incluyen:

  • Calidad de imagen de vanguardia, superando a los GAN en muchos casos
  • Entrenamiento estable sin dinámica adversaria
  • Altamente paralelizable
  • Arquitectura flexible – cualquier modelo que mapee entradas a salidas de la misma dimensionalidad puede ser utilizado
  • Base teórica sólida

Sumérjamonos más a fondo en cómo funcionan los modelos de difusión.

Fuente: Song et al.

Fuente: Song et al.

Las ecuaciones diferenciales estocásticas gobiernan los procesos directo e inverso en los modelos de difusión. La EDE directa agrega ruido a los datos, transformándolos gradualmente en una distribución de ruido. La EDE inversa, guiada por una función de puntuación aprendida, elimina progresivamente el ruido, lo que lleva a la generación de imágenes realistas a partir de ruido aleatorio. Este enfoque es clave para lograr un rendimiento generativo de alta calidad en espacios de estado continuo.

El Proceso de Difusión Directo

El proceso de difusión directo comienza con un punto de datos x₀ muestreado de la distribución de datos real, y agrega gradualmente ruido gaussiano durante T pasos de tiempo para producir versiones cada vez más ruidosas x₁, x₂, …, xT.

En cada paso de tiempo t, agregamos una pequeña cantidad de ruido según:

x_t = √(1 - β_t) * x_{t-1} + √(β_t) * ε

Donde:

  • β_t es un calendario de varianza que controla cuánto ruido se agrega en cada paso
  • ε es ruido gaussiano aleatorio

Este proceso continúa hasta que xT sea casi puro ruido gaussiano.

Matemáticamente, podemos describir esto como una cadena de Markov:

q(x_t | x_{t-1}) = N(x_t; √(1 - β_t) * x_{t-1}, β_t * I)

Donde N denota una distribución gaussiana.

El calendario β_t se elige generalmente para ser pequeño para los primeros pasos de tiempo y aumentar con el tiempo. Las opciones comunes incluyen calendarios lineales, cosenos o sigmoideos.

El Proceso de Difusión Inverso

El objetivo de un modelo de difusión es aprender el inverso de este proceso – comenzar con ruido puro xT y desenoizarlo progresivamente para recuperar una muestra limpia x₀.

Modelamos este proceso inverso como:

p_θ(x_{t-1} | x_t) = N(x_{t-1}; μ_θ(x_t, t), σ_θ^2(x_t, t))

Donde μ_θ y σ_θ^2 son funciones aprendidas (generalmente redes neuronales) parametrizadas por θ.

La innovación clave es que no necesitamos modelar explícitamente la distribución inversa completa. En su lugar, podemos parametrizarla en términos del proceso directo, que conocemos.

Específicamente, podemos demostrar que la media óptima del proceso inverso μ* es:

μ* = 1/√(1 - β_t) * (x_t - β_t/√(1 - α_t) * ε_θ(x_t, t))

Donde:

  • α_t = 1 – β_t
  • ε_θ es una red de predicción de ruido aprendida

Esto nos da un objetivo simple – entrenar una red neuronal ε_θ para predecir el ruido que se agregó en cada paso.

Objetivo de Entrenamiento

El objetivo de entrenamiento para los modelos de difusión se puede derivar de la inferencia variacional. Después de algunas simplificaciones, llegamos a una pérdida L2 simple:

L = E_t,x₀,ε [ ||ε - ε_θ(x_t, t)||² ]

Donde:

  • t se muestrea uniformemente de 1 a T
  • x₀ se muestrea de los datos de entrenamiento
  • ε se muestrea como ruido gaussiano
  • x_t se construye agregando ruido a x₀ según el proceso directo

En otras palabras, estamos entrenando el modelo para predecir el ruido que se agregó en cada paso de tiempo.

Arquitectura del Modelo

La arquitectura U-Net es central en el paso de desenoización del modelo de difusión. Presenta una estructura codificador-decodificador con conexiones de salto que ayudan a preservar detalles finos durante el proceso de reconstrucción. El codificador muestrea progresivamente la imagen de entrada mientras captura características de alto nivel, y el decodificador muestrea las características codificadas para reconstruir la imagen. Esta arquitectura es particularmente efectiva en tareas que requieren una localización precisa, como la segmentación de imágenes.

La red de predicción de ruido ε_θ puede utilizar cualquier arquitectura que mapee entradas a salidas de la misma dimensionalidad. Las arquitecturas de estilo U-Net son una opción popular, especialmente para tareas de generación de imágenes.

Una arquitectura típica podría parecerse a:


<p>class DiffusionUNet(nn.Module):
def __init__(self):
super().__init__()</p>

<p># Muestreo
self.down1 = UNetBlock(3, 64)
self.down2 = UNetBlock(64, 128)
self.down3 = UNetBlock(128, 256)</p>

<p># Cuello de botella
self.bottleneck = UNetBlock(256, 512)</p>

<p># Muestreo inverso
self.up3 = UNetBlock(512, 256)
self.up2 = UNetBlock(256, 128)
self.up1 = UNetBlock(128, 64)</p>

# Salida
self.out = nn.Conv2d(64, 3, 1)

<p>def forward(self, x, t):
# Incorporar paso de tiempo
t_emb = self.time_embedding(t)</p>

<p># Muestreo
d1 = self.down1(x, t_emb)
d2 = self.down2(d1, t_emb)
d3 = self.down3(d2, t_emb)</p>

<p># Cuello de botella
bottleneck = self.bottleneck(d3, t_emb)</p>

<p># Muestreo inverso
u3 = self.up3(torch.cat([bottleneck, d3], dim=1), t_emb)
u2 = self.up2(torch.cat([u3, d2], dim=1), t_emb)
u1 = self.up1(torch.cat([u2, d1], dim=1), t_emb)</p>

# Salida
return self.out(u1)

Los componentes clave son:

  • Arquitectura de estilo U-Net con conexiones de salto
  • Incorporación de tiempo para condicionar en el paso de tiempo
  • Profundidad y anchura flexibles

Algoritmo de Muestreo

Una vez que hemos entrenado nuestra red de predicción de ruido ε_θ, podemos utilizarla para generar nuevas muestras. El algoritmo de muestreo básico es:

  1. Comenzar con ruido gaussiano puro xT
  2. Para t = T a 1:
    • Predecir ruido: ε_θ(x_t, t)
    • Calcular media: μ = 1/√(1-β_t) * (x_t - β_t/√(1-α_t) * ε_θ(x_t, t))
    • Muestrear: x_{t-1} ~ N(μ, σ_t^2 * I)
  3. Devolver x₀

Este proceso desenoiza gradualmente la muestra, guiada por nuestra red de predicción de ruido aprendida.

En la práctica, existen varias técnicas de muestreo que pueden mejorar la calidad o la velocidad:

  • Muestreo DDIM: Una variante determinista que permite menos pasos de muestreo
  • Muestreo ancestral: Incorpora la varianza aprendida σ_θ^2
  • Muestreo truncado: Detiene el muestreo temprano para una generación más rápida

Aquí hay una implementación básica del algoritmo de muestreo:


<p>def sample(model, n_samples, device):
# Comenzar con ruido puro
x = torch.randn(n_samples, 3, 32, 32).to(device)</p>

<p>for t in reversed(range(1000)):
# Agregar ruido para crear x_t
t_batch = torch.full((n_samples,), t, device=device)
noise = torch.randn_like(x)
x_t = add_noise(x, noise, t)</p>

<p># Predecir y eliminar ruido
pred_noise = model(x_t, t_batch)
x = remove_noise(x_t, pred_noise, t)</p>

<p># Agregar ruido para el siguiente paso (excepto en t=0)
if t &gt; 0:
noise = torch.randn_like(x)
x = add_noise(x, noise, t-1)</p>

return x

Las Matemáticas detrás de los Modelos de Difusión

Para comprender verdaderamente los modelos de difusión, es crucial profundizar en las matemáticas que los sustentan. Exploraremos algunos conceptos clave con más detalle:

Cadena de Markov y Ecuaciones Diferenciales Estocásticas

El proceso de difusión directo en los modelos de difusión se puede ver como una cadena de Markov o, en el límite continuo, como una ecuación diferencial estocástica (EDE). La formulación de la EDE proporciona un marco teórico poderoso para analizar y ampliar los modelos de difusión.

La EDE directa se puede escribir como:

dx = f(x,t)dt + g(t)dw

Donde:

  • f(x,t) es el término de deriva
  • g(t) es el coeficiente de difusión
  • dw es un proceso de Wiener (movimiento browniano)

Diferentes opciones de f y g llevan a diferentes tipos de procesos de difusión. Por ejemplo:

  • Explosión de varianza (VE) EDE: dx = √(d/dt σ²(t)) dw
  • Preservación de varianza (VP) EDE: dx = -0.5 β(t)xdt + √(β(t)) dw

Entender estas EDE nos permite derivar estrategias de muestreo óptimas y ampliar los modelos de difusión a nuevos dominios.

Emparejamiento de Puntuación y Emparejamiento de Puntuación de Desenoización

La conexión entre los modelos de difusión y el emparejamiento de puntuación proporciona otra perspectiva valiosa. La función de puntuación se define como el gradiente de la densidad de probabilidad logarítmica:

s(x) = ∇x log p(x)

El emparejamiento de puntuación de desenoización apunta a estimar esta función de puntuación entrenando un modelo para desenoizar ligeramente puntos de datos perturbados. Este objetivo resulta ser equivalente al objetivo de entrenamiento del modelo de difusión en el límite continuo.

Esta conexión nos permite aprovechar técnicas del modelado generativo basado en puntuación, como la dinámica de Langevin annealada para muestreo.

Técnicas de Entrenamiento Avanzadas

Muestreo de Importancia

El muestreo de importancia es una técnica que permite enfocar el entrenamiento en los pasos de tiempo más informativos. Una aproximación es utilizar una distribución no uniforme sobre los pasos de tiempo, ponderada por la norma L2 esperada de la puntuación:

p(t) ∝ E[||s(x_t, t)||²]

Esto puede conducir a un entrenamiento más rápido y a una mejor calidad de las muestras.

Destilación Progresiva

La destilación progresiva es una técnica para crear modelos de muestreo más rápidos sin sacrificar la calidad. El proceso funciona de la siguiente manera:

  1. Entrenar un modelo de difusión base con muchos pasos de tiempo (por ejemplo, 1000)
  2. Crear un modelo estudiante con menos pasos de tiempo (por ejemplo, 100)
  3. Entrenar al modelo estudiante para que coincida con el proceso de desenoización del modelo base
  4. Repetir los pasos 2-3, reduciendo progresivamente los pasos de tiempo

Esto permite una generación de alta calidad con significativamente menos pasos de desenoización.

Innovaciones Arquitectónicas

Modelos de Difusión Basados en Transformadores

Aunque las arquitecturas U-Net han sido populares para los modelos de difusión de imágenes, trabajos recientes han explorado el uso de arquitecturas de transformadores. Los transformadores ofrecen varias ventajas potenciales:

  • Mejor manejo de dependencias de largo alcance
  • Mecanismos de condicionamiento más flexibles
  • Escalabilidad más fácil a tamaños de modelo más grandes

Modelos como DiT (Diffusion Transformers) han mostrado resultados prometedores, lo que podría ofrecer un camino hacia una generación aún más de alta calidad.

Modelos de Difusión Jerárquicos

Los modelos de difusión jerárquicos generan datos a múltiples escalas, permitiendo tanto coherencia global como detalles finos. El proceso generalmente implica:

  1. Generar una salida de baja resolución
  2. Mejorar progresivamente la resolución

Este enfoque puede ser particularmente efectivo para la generación de imágenes de alta resolución o contenido de larga duración.

Temas Avanzados

Guía Libre de Clasificador

La guía libre de clasificador es una técnica para mejorar la calidad de las muestras y el control. La idea clave es entrenar dos modelos de difusión:

  1. Un modelo incondicional p(x_t)
  2. Un modelo condicional p(x_t | y) donde y es alguna información de condicionamiento (por ejemplo, una descripción de texto)

Durante el muestreo, interpolamos entre estos modelos:

ε_θ = (1 + w) * ε_θ(x_t | y) - w * ε_θ(x_t)

Donde w > 0 es una escala de guía que controla cuánto enfatizar el modelo condicional.

Esto permite un condicionamiento más fuerte sin necesidad de volver a entrenar el modelo. Ha sido crucial para el éxito de modelos de texto a imagen como DALL-E 2 y Stable Diffusion.

Difusión Latente

El Modelo de Difusión Latente (LDM) implica codificar los datos de entrada en un espacio latente donde ocurre el proceso de difusión. El modelo agrega progresivamente ruido a la representación latente de la imagen, lo que lleva a la generación de una versión ruidosa, que luego se desenoiza utilizando una arquitectura U-Net. La U-Net, guiada por mecanismos de atención cruzada, integra información de diversas fuentes de condicionamiento como mapas semánticos, texto e imágenes, reconstruyendo finalmente la imagen en el espacio de píxeles. Este proceso es fundamental para generar imágenes de alta calidad con estructura controlada y atributos deseados.

Esto ofrece varias ventajas:

  • Entrenamiento y muestreo más rápidos
  • Mejor manejo de imágenes de alta resolución
  • Incorporación de condicionamiento más fácil

El proceso funciona de la siguiente manera:

  1. Entrenar un autoencoder para comprimir imágenes en un espacio latente
  2. Entrenar un modelo de difusión en este espacio latente
  3. Para la generación, muestrear en el espacio latente y decodificar a píxeles

Este enfoque ha sido muy exitoso, impulsando modelos como Stable Diffusion.

Modelos de Coherencia

Los modelos de coherencia son una innovación reciente que apunta a mejorar la velocidad y la calidad de los modelos de difusión. La idea clave es entrenar un solo modelo que pueda mapear desde cualquier nivel de ruido directamente a la salida final, en lugar de requerir desenoización iterativa.

Esto se logra a través de una función de pérdida cuidadosamente diseñada que impone coherencia entre predicciones a diferentes niveles de ruido. El resultado es un modelo que puede generar muestras de alta calidad en una sola pasada hacia adelante, acelerando significativamente la inferencia.

Consejos Prácticos para Entrenar Modelos de Difusión

Entrenar modelos de difusión de alta calidad puede ser desafiante. Aquí hay algunos consejos prácticos para mejorar la estabilidad del entrenamiento y los resultados:

  1. Recorte de gradientes: Utilice recorte de gradientes para prevenir gradientes explosivos, especialmente al comienzo del entrenamiento.
  2. Promedio móvil de los pesos del modelo: Mantenga un promedio móvil de los pesos del modelo para el muestreo, lo que puede conducir a una generación más estable y de mayor calidad.
  3. Aumento de datos: Para modelos de imágenes, simples aumentos como flips horizontales aleatorios pueden mejorar la generalización.
  4. Programación de ruido: Experimente con diferentes programaciones de ruido (lineal, coseno, sigmoide) para encontrar lo que funciona mejor para sus datos.
  5. Entrenamiento de precisión mixta: Utilice entrenamiento de precisión mixta para reducir el uso de memoria y acelerar el entrenamiento, especialmente para modelos grandes.
  6. Generación condicional: Incluso si su objetivo final es la generación incondicional, entrenar con condicionamiento (por ejemplo, en clases de imágenes) puede mejorar la calidad general de las muestras.

Evaluación de Modelos de Difusión

Evaluar adecuadamente los modelos generativos es crucial pero desafiante. Aquí hay algunas métricas y enfoques comunes:

Distancia de Inception de Fréchet (FID)

FID es una métrica ampliamente utilizada para evaluar la calidad y la diversidad de las imágenes generadas. Compara las estadísticas de las muestras generadas con los datos reales en el espacio de características de un clasificador preentrenado (generalmente InceptionV3).

Puntuaciones FID más bajas indican una mejor calidad y distribuciones más realistas. Sin embargo, FID tiene limitaciones y no debe ser la única métrica utilizada.

Puntuación de Inception (IS)

Inception Score mide tanto la calidad como la diversidad de las imágenes generadas. Utiliza una red Inception preentrenada para computar:

IS = exp(E[KL(p(y|x) || p(y))])

Donde p(y|x) es la distribución condicional de clase para la imagen generada x.

Puntuaciones IS más altas indican una mejor calidad y diversidad, pero tiene limitaciones conocidas, especialmente para conjuntos de datos muy diferentes de ImageNet.

Logaritmo Negativo de la Probabilidad Predictiva (NLL)

Para los modelos de difusión, podemos computar el logaritmo negativo de la probabilidad de los datos retenidos. Esto proporciona una medida directa de cuán bien el modelo se ajusta a la distribución de datos real.

Sin embargo, NLL puede ser computacionalmente costoso de estimar con precisión para datos de alta dimensionalidad.

Evaluación Humana

Para muchas aplicaciones, especialmente las creativas, la evaluación humana sigue siendo crucial. Esto puede involucrar:

  • Comparaciones lado a lado con otros modelos
  • Evaluaciones de tipo prueba de Turing
  • Evaluaciones específicas de tarea (por ejemplo, generación de subtítulos de imágenes para modelos de texto a imagen)

Aunque subjetiva, la evaluación humana puede capturar aspectos de la calidad que las métricas automatizadas pasan por alto.

Modelos de Difusión en Producción

Desplegar modelos de difusión en entornos de producción presenta desafíos únicos. Aquí hay algunas consideraciones y prácticas recomendadas:

Optimización para Inferencia

  1. Exportación ONNX: Convertir modelos a formato ONNX para una inferencia más rápida en diferentes hardware.
  2. Cuantización: Utilizar técnicas como la cuantización INT8 para reducir el tamaño del modelo y mejorar la velocidad de inferencia.
  3. Almacenamiento en caché: Para modelos condicionales, almacenar en caché resultados intermedios del modelo incondicional para acelerar la guía libre de clasificador.
  4. Procesamiento por lotes: Aprovechar el procesamiento por lotes para hacer un uso eficiente de los recursos de GPU.

Escalabilidad

  1. Inferencia distribuida: Para aplicaciones de alto rendimiento, implementar inferencia distribuida en múltiples GPUs o máquinas.
  2. Muestreo adaptativo: Ajustar dinámicamente el número de pasos de muestreo según el equilibrio deseado entre calidad y velocidad.
  3. Generación progresiva: Para salidas grandes (por ejemplo, imágenes de alta resolución), generar progresivamente desde baja a alta resolución para proporcionar resultados iniciales más rápidos.

Seguridad y Filtrado

  1. Filtrado de contenido: Implementar sistemas de filtrado de contenido robustos para prevenir la generación de contenido dañino o inapropiado.
  2. Marca de agua: Considerar la incorporación de marcas de agua invisibles en el contenido generado para rastreabilidad.

Aplicaciones

Los modelos de difusión han encontrado éxito en una amplia gama de tareas generativas:

Generación de Imágenes

La generación de imágenes es donde los modelos de difusión ganaron prominencia por primera vez. Algunos ejemplos notables incluyen:

  • DALL-E 3: El modelo de texto a imagen de OpenAI, que combina un codificador de texto CLIP con un decodificador de imagen de difusión
  • Stable Diffusion: Un modelo de difusión latente de código abierto para generación de texto a imagen
  • Imagen: El modelo de texto a imagen de Google (GOOGL )

Estos modelos pueden generar imágenes muy realistas y creativas a partir de descripciones de texto, superando enfoques basados en GAN anteriores.

Generación de Video

Los modelos de difusión también se han aplicado a la generación de video:

  • Modelos de Difusión de Video: Generar video tratando el tiempo como una dimensión adicional en el proceso de difusión
  • Make-A-Video: El modelo de texto a video de Meta
  • Imagen Video: El modelo de texto a video de Google

Estos modelos pueden generar clips de video cortos a partir de descripciones de texto, abriendo nuevas posibilidades para la creación de contenido.

Generación 3D

Trabajos recientes han extendido los modelos de difusión a la generación 3D:

  • DreamFusion: Generación de texto a 3D utilizando modelos de difusión 2D
  • Point-E: El modelo de nube de puntos de OpenAI para generación de objetos 3D

Estos enfoques permiten la creación de activos 3D a partir de descripciones de texto, con aplicaciones en juegos, realidad virtual/aumentada y diseño de productos.

Desafíos y Direcciones Futuras

Aunque los modelos de difusión han demostrado un éxito notable, aún existen varios desafíos y áreas para investigación futura:

Eficiencia Computacional

El proceso de muestreo iterativo de los modelos de difusión puede ser lento, especialmente para salidas de alta resolución. Enfoques como la difusión latente y los modelos de coherencia apuntan a abordar esto, pero aún se necesitan mejoras en la eficiencia.

Controlabilidad

Aunque técnicas como la guía libre de clasificador han mejorado el control, aún hay trabajo por hacer para permitir un control más fino sobre las salidas generadas. Esto es especialmente importante para aplicaciones creativas.

Generación Multi-Modal

Los modelos de difusión actuales excel en la generación de una sola modalidad (por ejemplo, imágenes o audio). Desarrollar modelos de difusión verdaderamente multi-modales que puedan generar sin problemas a través de modalidades es una dirección emocionante para el trabajo futuro.

Comprensión Teórica

Aunque los modelos de difusión tienen resultados empíricos sólidos, aún hay más que entender sobre por qué funcionan tan bien. Desarrollar una comprensión teórica más profunda podría conducir a mejoras y nuevas aplicaciones.

Conclusión

Los modelos de difusión representan un paso adelante en la inteligencia artificial generativa, ofreciendo resultados de alta calidad en una variedad de modalidades. Al aprender a revertir un proceso de adición de ruido, proporcionan un enfoque flexible y teóricamente fundamentado para la generación.

Desde herramientas creativas hasta simulaciones científicas, la capacidad de generar datos complejos y de alta dimensionalidad tiene el potencial de transformar muchos campos. Sin embargo, es importante abordar estas tecnologías poderosas de manera reflexiva, considerando tanto su inmenso potencial como los desafíos éticos que plantean.

mm

He dedicado los últimos cinco años sumergiéndome en el fascinante mundo de Machine Learning y Deep Learning. Mi pasión y experiencia me han llevado a contribuir a más de 50 proyectos de ingeniería de software diversos, con un enfoque particular en AI/ML. Mi curiosidad en curso también me ha llevado hacia el Procesamiento de Lenguaje Natural, un campo que estoy ansioso por explorar más a fondo.