Difusión: de ruido a imagen

El proceso hacia delante

Con un calendario de ruido $\beta_1,\dots,\beta_T$, sea $\alpha_t = 1-\beta_t$ y $\bar\alpha_t = \prod_{s\le t}\alpha_s$. Cada paso es $q(x_t\mid x_{t-1}) = \mathcal N\big(\sqrt{\alpha_t}\,x_{t-1},\ \beta_t I\big)$, y al encadenarlos sale una fórmula cerrada para cualquier paso:

La barra de la lección muestra $\bar\alpha_t$ (la parte de la «energía» de la mezcla que todavía es imagen) y $1-\bar\alpha_t$ (la que es ruido). La lección usa $T = 50$ y el calendario coseno de Nichol y Dhariwal (2021), $\bar\alpha(t) \propto \cos^2\!\big(\tfrac{t/T + s}{1+s}\cdot\tfrac{\pi}{2}\big)$ con $s = 0.008$, recortando cada $\beta_t$ a 0.35 como máximo (con tan pocos pasos, los últimos $\beta$ del coseno se acercan a 1). Así $\bar\alpha$ vale 0.899 en el paso 10, 0.494 en el 25, 0.094 en el 40 y 0.003 en el 50. Los píxeles se pasan antes a la escala $-1\ldots 1$; en pantalla se recorta a $0\ldots 1$.

Por qué no de un salto

La red se entrena con error cuadrático, y lo que minimiza el error cuadrático es el promedio condicionado $\mathbb E[x_0\mid x_t]$. Cuando $t = T$ la estática casi no conserva información, así que $\mathbb E[x_0\mid x_T]\approx \mathbb E[x_0]$: el promedio de todos los dígitos, la mancha de la lección. En cambio, si cada $\beta_t$ es pequeño, el paso inverso $q(x_{t-1}\mid x_t)$ es casi gaussiano, con la misma forma que el paso hacia delante (Sohl-Dickstein et al., 2015, que citan a Feller, 1949): basta estimar su media.

Qué aprende la red y cómo se entrena

La red $\varepsilon_\theta(x_t, t, c)$ estima el ruido. La pérdida simplificada de Ho et al. (2020) es

Gracias a la fórmula cerrada no hace falta recorrer los pasos al entrenar: se salta directo a un $t$ al azar. La red de píxeles de la lección tiene 53 056 parámetros (tres capas de 128 neuronas con conexiones residuales y SiLU; entradas: los 64 píxeles, 16 números que codifican el paso y la etiqueta en 11 casillas, la última para «sin pista») y se entrenó 30 000 rondas con lotes de 256 imágenes: 7 680 000 versiones con ruido de los 1500 dígitos de entrenamiento de la lección 15. Con la imagen del siete en el paso 20, su estimación del ruido tiene un error cuadrático medio de 0.135 por píxel; si dijera «no hay ruido», el error sería 0.864.

El muestreo

El muestreo original de DDPM parte de $x_T\sim\mathcal N(0,I)$ y, para $t=T,\dots,1$, calcula

con un poco de ruido nuevo $z$ en cada paso (salvo el último). La escena usa la variante determinista de Song et al. (2021), DDIM con $\eta = 0$, con la misma red entrenada como en DDPM:

Es decir: con la estimación del ruido se calcula la imagen limpia que implica y se vuelve a mezclar con la parte de ruido que le toca al paso anterior. Sin ruido nuevo, cada punto del plano sigue exactamente su flecha y la misma estática da siempre el mismo resultado; con otra estática sale otro.

Guiado sin clasificador

Se entrena una sola red con la etiqueta $c$ y, con probabilidad 0.15, sin ella ($\varnothing$). Al muestrear se combina

Con $s=0$ es la estimación sin pista (la pista se ignora); con $s=1$, la condicional, que sigue la pista tal cual; con $s>1$ se exagera la dirección de la pista. Es la misma convención de la lección 18. Ho y Salimans (2022) escriben lo mismo con $w = s-1$: $\tilde\varepsilon = (1+w)\,\varepsilon_\theta(x_t,c) - w\,\varepsilon_\theta(x_t)$. Cada paso guiado cuesta dos evaluaciones de la red. En el plano de la lección, con la pista «siete» y doscientos puntos (lector: la regresión logística de la lección 15), se leen como siete 26 con $s = 0$, 199 con $s = 1$ y 200 con $s = 5$. A la vez, los sietes reales que se quedan sin ningún punto nuevo a menos de 0.3 son 4 de 154 con $s = 1$, 48 con $s = 2$, 95 con $s = 3$ y 144 con $s = 5$: los puntos se amontonan en la orilla del grupo más alejada de los demás dígitos. La red de píxeles genera sus sietes con $s = 3$.

Difusión latente

Un autocodificador comprime la imagen a un código $z=E(x)$; la difusión ocurre en $z$ y al final se decodifica $x=D(z)$ (Rombach et al., 2022). En la lección, $z$ son los dos números del VAE de la lección 15, multiplicados por 1.6 antes de difundir (Stable Diffusion también reescala sus códigos con un factor fijo), y la red del plano tiene 10 370 parámetros. En Stable Diffusion v1, el autocodificador reduce cada lado 8 veces y entrega 4 canales: una imagen de $512\times512\times3 = 786\,432$ números se vuelve un código de $64\times64\times4 = 16\,384$, 48 veces menos.

En código

import numpy as np

T, s = 50, 0.008
f = lambda t: np.cos((t / T + s) / (1 + s) * np.pi / 2) ** 2
a_full = f(np.arange(T + 1)) / f(0)
betas = np.clip(1 - a_full[1:] / a_full[:-1], 1e-4, 0.35)   # calendario de ruido
alphas = 1 - betas
abar = np.cumprod(alphas)
print('imagen en la mezcla (%):', [round(float(100 * abar[t - 1]), 1) for t in (10, 25, 40, 50)])

# Datos de juguete: un solo número que vale -1 o +1. Para ellos, el mejor
# estimador del ruido se puede escribir a mano (en la lección lo aprende una red).
def ruido_estimado(x, i):
    a, b = np.sqrt(abar[i]), np.sqrt(1 - abar[i])
    x0 = np.tanh(a * x / b**2)          # promedio de x0 dado x_t
    return (x - a * x0) / b

rng = np.random.default_rng(0)
x = rng.standard_normal(10_000)        # estática pura
for i in reversed(range(T)):           # cincuenta pasos de regreso (DDPM)
    e = ruido_estimado(x, i)
    x = (x - betas[i] / np.sqrt(1 - abar[i]) * e) / np.sqrt(alphas[i])
    if i > 0:
        x += np.sqrt(betas[i]) * rng.standard_normal(x.size)

print('a menos de 0.1 de -1 o de +1:', round(float(np.mean(np.abs(np.abs(x) - 1) < 0.1)), 3))
print('fracción en +1:', round(float(np.mean(x > 0)), 3))
# imagen en la mezcla (%): [89.9, 49.4, 9.4, 0.3]
# a menos de 0.1 de -1 o de +1: 1.0
# fracción en +1: 0.499

Lo que viene

La lección 18 muestra cómo una frase se convierte en la pista que guía la difusión (CLIP y el condicionamiento por texto). Más adelante, el transformer (lecciones 21 y 22) reemplaza a la red convolucional en muchos modelos de difusión modernos.