Diffusion Models

Los modelos de difusión son modelos de IA generativa que crean datos aprendiendo a invertir un proceso gradual de adición de ruido. Durante el entrenamiento, se corrompen imágenes u otros datos con ruido gaussiano en muchos pasos pequeños, y una red neuronal aprende a predecir el ruido en cada paso. Para generar contenido, el modelo parte de ruido aleatorio puro y lo elimina paso a paso hasta que aparece una muestra nueva. Jascha Sohl-Dickstein y sus colaboradores propusieron este enfoque en 2015 en Deep Unsupervised Learning using Nonequilibrium Thermodynamics, y Jonathan Ho, Ajay Jain y Pieter Abbeel lo hicieron competitivo en 2020 con Denoising Diffusion Probabilistic Models (DDPM), que estableció un FID de 3,17, entonces puntero, en CIFAR-10. En 2021, investigadores de OpenAI demostraron que los modelos de difusión superaban a las GANs en la síntesis de imágenes de ImageNet, y ahora los modelos de difusión y sus variantes de emparejamiento de flujos impulsan sistemas de generación de imágenes a partir de texto como Stable Diffusion, generadores de vídeo, editores de imágenes y modelos científicos.

Cómo funcionan los modelos de difusión#

El proceso directo es fijo, no se aprende. Añade un poco de ruido gaussiano en cada uno de los T pasos; DDPM utiliza T = 1.000, con una varianza del ruido que aumenta linealmente de 0.0001 a 0.02 por paso. Como la suma de ruido gaussiano sigue siendo gaussiana, una imagen de entrenamiento puede saltar directamente a cualquier paso t mediante una fórmula cerrada, xt = sqrt(alpha_bar_t) * x0 + sqrt(1 - alpha_bar_t) * noise, donde alpha_bar_t es el producto de 1 - beta a lo largo de los primeros t pasos. El entrenamiento selecciona repetidamente una imagen, un paso aleatorio y ruido nuevo, construye xt y enseña a la red a predecir ese ruido con un error cuadrático medio, el objetivo simplificado de DDPM. Este código aplica el proceso directo a una foto con el calendario de DDPM:

import urllib.request

import numpy as np
import torch
from PIL import Image

# DDPM's linear schedule: 1,000 steps whose noise variance rises from 0.0001 to 0.02
betas = torch.linspace(1e-4, 0.02, 1000)
alpha_bar = torch.cumprod(1 - betas, dim=0)

# Load a photo and scale its pixels to [-1, 1]
image = Image.open(urllib.request.urlopen("https://ultralytics.com/images/bus.jpg")).convert("RGB").resize((256, 256))
x0 = torch.tensor(np.array(image), dtype=torch.float32).permute(2, 0, 1) / 127.5 - 1

# Jump straight to step t in closed form: x_t = sqrt(alpha_bar) * x0 + sqrt(1 - alpha_bar) * noise
frames = []
for t in (0, 99, 299, 599, 999):
    xt = alpha_bar[t].sqrt() * x0 + (1 - alpha_bar[t]).sqrt() * torch.randn_like(x0)
    frames.append(((xt.clamp(-1, 1) + 1) * 127.5).byte().permute(1, 2, 0).numpy())
    print(f"step {t + 1}: image weight {alpha_bar[t].sqrt():.3f}, noise weight {(1 - alpha_bar[t]).sqrt():.3f}")
Image.fromarray(np.hstack(frames)).save("forward_diffusion.jpg")

El peso de la imagen baja de 1.000 en el paso 1 a 0.947 en el paso 100, 0.630 en el paso 300, 0.161 en el paso 600 y 0.006 en el paso 1.000, mientras que el peso del ruido sube hasta 1.000. El archivo forward_diffusion.jpg guardado muestra el autobús nítido al principio, después granulado, luego tenue y, finalmente, desaparecido.

La generación recorre el proceso a la inversa. Parte de ruido gaussiano puro y, en cada paso, la red predice el ruido de la muestra actual, un algoritmo de resolución elimina parte de ese ruido y el resultado, algo más limpio, pasa al siguiente paso. Una instrucción de texto o una etiqueta de clase se introduce en la red como condición, a menudo mediante atención cruzada. La guía sin clasificador, presentada por Jonathan Ho y Tim Salimans, entrena una red tanto con la condición como sin ella, y en el momento del muestreo extrapola desde la predicción incondicional hacia la condicional; una escala de guía más alta, como el valor predeterminado de 7,5 en los flujos de Stable Diffusion 1.x, sigue más de cerca la instrucción, pero reduce la diversidad.

Acelerar el muestreo#

El muestreo DDPM requiere 1.000 evaluaciones de la red por imagen, así que buena parte del progreso del campo se ha dedicado a reducir esa cifra:

• Mejores algoritmos de muestreo: DDIM reutiliza un DDPM entrenado con un algoritmo de muestreo determinista y genera muestras de alta calidad entre 10 y 50 veces más rápido en tiempo real. Los algoritmos de resolución de EDO posteriores reducen la generación habitual a unos 20-50 pasos.
• Difusión latente: un modelo de difusión latente elimina el ruido de una representación comprimida por un autocodificador, en lugar de hacerlo directamente sobre los píxeles. El autocodificador de Stable Diffusion transforma una imagen de 512×512×3 en una representación latente de 64×64×4, con 48 veces menos valores.
• Destilación: los modelos de consistencia aprenden a transformar directamente cualquier muestra ruidosa en una limpia, lo que permite generar contenido en un solo paso; los modelos de consistencia latente aplican la misma idea a Stable Diffusion en 2-4 pasos.
• Trayectorias más rectas: el emparejamiento de flujos y el flujo rectificado entrenan con trayectorias rectas del ruido a los datos que los algoritmos de resolución pueden seguir en menos pasos.

La propia red de eliminación de ruido ha pasado de la U-Net convolucional de DDPM y Stable Diffusion 1.x al transformador de difusión (DiT), que utilizan la mayoría de los generadores actuales de imágenes y vídeo.

Modelos de difusión y otros modelos generativos#

• Redes generativas antagónicas (GANs): una GAN genera una muestra en una sola pasada hacia delante, pero su entrenamiento adversario entre el generador y el discriminador puede provocar inestabilidad y colapso de modos, situación en la que el generador solo cubre una parte de los datos. El entrenamiento de difusión es una regresión estable y cubre la distribución de datos de forma más completa, a costa de necesitar muchos pasos de muestreo.
• Autocodificadores variacionales (VAEs): un VAE decodifica un vector latente para generar una muestra en una sola pasada, y sus muestras tienden a ser más borrosas. Los modelos de difusión latente solo utilizan un autocodificador de estilo VAE para comprimir imágenes y dejan la generación en manos del modelo de difusión.
• Modelos autorregresivos: los modelos de lenguaje y los modelos visuales autorregresivos generan un token o una escala cada vez, condicionados por lo que se ha generado antes. Un modelo de difusión refina en paralelo todas las partes de la muestra en cada paso, y los modelos de difusión para lenguaje aplican la misma eliminación de ruido al texto.

Aplicaciones en el mundo real#

• Generación y edición de imágenes: los modelos de texto a imagen crean imágenes a partir de instrucciones; el relleno de imágenes vuelve a generar una región enmascarada para que encaje con el entorno; ControlNet restringe la composición mediante mapas de bordes, profundidad o pose; y los escaladores de difusión realizan superresolución.
• Generación de vídeo: los modelos de generación de vídeo eliminan el ruido de representaciones latentes comprimidas que abarcan muchos fotogramas a la vez, por lo que cada fotograma se genera junto con los contiguos.
• Datos sintéticos de entrenamiento: las imágenes generadas cubren lagunas en los conjuntos de datos de visión, como defectos poco frecuentes, condiciones meteorológicas inusuales o puntos de vista atípicos. En un conjunto de datos de Ultralytics Platform, la acción Generar imágenes similares crea de 1 a 16 imágenes nuevas por lote inspiradas en una imagen existente; las que conserves se incorporan al conjunto de datos sin etiquetar, listas para su anotación.
• Robótica: las políticas de difusión generan secuencias de acciones robóticas eliminando su ruido, condicionadas por imágenes de cámaras.
• Ciencia: AlphaFold 3 predice las estructuras de complejos de proteínas, ácidos nucleicos y moléculas pequeñas con un módulo de difusión que elimina el ruido de las coordenadas atómicas sin procesar y sustituye al módulo de estructura de AlphaFold 2.