Autoencoders, VAE y atención: las piezas de Stable Diffusion

1. Autoencoder: comprimir y reconstruir

Un autoencoder es una red que aprende a comprimir su propio input a una representación chica (el latente) y reconstruirlo de vuelta, entrenada simplemente para minimizar el error entre entrada y salida:

ae = nn.Sequential(   #28x28
    conv(1,2), conv(2,4),           # encoder: comprime a 8x8
    deconv(4,2), deconv(2,1, act=False), # decoder: reconstruye a 28x28
    nn.Sigmoid()
)
# entrenado con MSE entre la reconstrucción y la imagen original

deconv (upsampling + convolución) es literalmente el proceso inverso de un conv con stride — en vez de reducir la resolución espacial, la aumenta. El resultado, tras entrenar, reconstruye imágenes de Fashion-MNIST razonablemente bien.

2. Por qué un autoencoder simple no genera nada bueno

Acá está el punto pedagógico central de la unidad. Si tomás el decoder ya entrenado y le das ruido aleatorio en vez de un latente real:

noise = torch.randn(16, nl)
generated_images = model.dec(noise).sigmoid()  # resultado: basura irreconocible

El resultado es ruido visual sin forma. La razón es estructural: un autoencoder normal solo se entrena para reconstruir los puntos exactos del espacio latente que corresponden a imágenes reales del dataset — no tiene ningún incentivo para que el espacio entre esos puntos tenga algún significado. El latente aprendido es un conjunto de islas dispersas en un espacio de alta dimensión; samplear un punto al azar casi seguro cae en un "hueco" que el decoder nunca vio y no sabe interpretar.

3. VAE: forzar un espacio latente que se pueda samplear

Un Variational Autoencoder resuelve esto con un cambio de diseño simple pero decisivo: en vez de que el encoder produzca un punto exacto en el espacio latente, produce los parámetros de una distribución (media mu y log-varianza lv), y el latente real se samplea de esa distribución antes de decodificar:

class VAE(nn.Module):
    def forward(self, x):
        x = self.enc(x)
        mu, lv = self.mu(x), self.lv(x)
        z = mu + (0.5*lv).exp() * torch.randn_like(lv)   # reparametrization trick
        return self.dec(z), mu, lv

La pieza que fuerza a que el espacio latente sea "denso" y navegable es un segundo término en la loss, la divergencia KL, que penaliza que la distribución aprendida se aleje de una normal estándar N(0,1):

def kld_loss(inp, x):
    x_hat, mu, lv = inp
    return -0.5 * (1 + lv - mu.pow(2) - lv.exp()).mean()

def vae_loss(inp, x): return kld_loss(inp, x) + bce_loss(inp, x)

El resultado práctico, verificable generando desde ruido puro con el mismo código que falló en el autoencoder simple: las imágenes generadas por un VAE son reconocibles, no ruido — porque ahora todo el espacio latente (no solo los puntos vistos en entrenamiento) fue empujado activamente a tener sentido para el decoder.

4. Self-attention, la misma que ya conocés, en imágenes

La mecánica es idéntica a la que ya viste en el Capítulo 1 de HF — Query, Key, Value, producto punto escalado, softmax — aplicada acá sobre los píxeles de un mapa de activaciones en vez de sobre tokens de texto:

class SelfAttention(nn.Module):
    def __init__(self, ni):
        self.scale = math.sqrt(ni)
        self.norm = nn.GroupNorm(1, ni)
        self.q,self.k,self.v,self.proj = [nn.Linear(ni, ni) for _ in range(4)]

    def forward(self, x):
        n,c,h,w = x.shape
        x = self.norm(x).view(n, c, -1).transpose(1, 2)   # [batch, h*w, canales]
        q,k,v = self.q(x), self.k(x), self.v(x)
        s = (q @ k.transpose(1,2)) / self.scale
        x = s.softmax(dim=-1) @ v
        return self.proj(x).transpose(1,2).reshape(n,c,h,w) + x   # residual

El truco de reshape (view(n, c, -1).transpose(1, 2)) convierte un mapa de activaciones [batch, canales, alto, ancho] en una secuencia [batch, alto×ancho, canales] — cada posición espacial de la imagen se trata como un "token", y attention deja que píxeles distantes entre sí se influyan mutuamente sin importar la distancia, algo que las convoluciones (limitadas a su receptive field local, ver Lección 8) no pueden hacer directamente. El notebook verifica esta implementación contra AttentionBlock de la librería diffusers de Hugging Face copiando los pesos manualmente entre ambas — y da resultados idénticos, confirmando que no hay ningún truco extra escondido en la versión de producción.

La versión multi-head (usada en Stable Diffusion real) divide los canales en h grupos independientes con rearrange de la librería einops, corre attention en paralelo sobre cada uno, y los vuelve a unir — el mismo patrón de multi-head attention que ya viste en el contexto de texto, sin ningún concepto nuevo.

5. Resumen

1. Un autoencoder comprime y reconstruye, pero su espacio latente tiene "huecos" que el decoder nunca aprendió a interpretar.
2. Un VAE fuerza ese espacio a parecerse a una distribución normal (vía divergencia KL) — por eso samplear ruido puro y decodificar da imágenes coherentes.
3. Este mismo VAE es el componente que hace viable la difusión latente: correr el proceso costoso de denoising sobre un espacio comprimido, no sobre píxeles.
4. Self-attention en imágenes es matemáticamente idéntica a la de texto — solo cambia qué se trata como "token" (posiciones espaciales en vez de palabras).