fast.ai · Parte 2 · Unidad 4 de 5

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

Tres arquitecturas que, combinadas, son literalmente los tres componentes de Stable Diffusion: un autoencoder comprime imágenes a un espacio latente, un VAE hace que ese espacio se pueda samplear con sentido, y attention es lo que permite condicionar la generación con texto.

Antes de leerU3 §1, DDPM · HF cap. 1 §3, atención

🎯 Al terminar vas a poder

Explicar con tus palabras qué es un autoencoder, qué hace deconv (upsampling + conv) y por qué el latente de un autoencoder común tiene "huecos".
Escribir el forward de un VAE con reparametrization trick y su loss (reconstrucción + KL), y explicar término por término qué empuja la KL.
Predecir qué pasa si el peso de la KL es demasiado alto (posterior collapse) o demasiado bajo (vuelta al autoencoder).
Implementar self-attention sobre un mapa de activaciones [n, c, h, w] con el residual correcto, y estimar su costo en memoria en función de h·w.
Diagnosticar los errores clásicos: residual mal sumado, view sobre tensor no contiguo, sigmoid con la loss equivocada.

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:

def conv(ni, nf, ks=3, act=True):
    layers = [nn.Conv2d(ni, nf, stride=2, kernel_size=ks, padding=ks//2)]   # stride 2: divide alto y ancho por 2
    if act: layers.append(nn.ReLU())
    return nn.Sequential(*layers)

def deconv(ni, nf, ks=3, act=True):
    layers = [nn.UpsamplingNearest2d(scale_factor=2),                          # duplica alto y ancho copiando píxeles
              nn.Conv2d(ni, nf, stride=1, kernel_size=ks, padding=ks//2)]   # y una conv 3×3 aprende a suavizarlos
    if act: layers.append(nn.ReLU())
    return nn.Sequential(*layers)

ae = nn.Sequential(                        # entrada: 1×28×28
    nn.ZeroPad2d(2),                       # → 1×32×32, para que las divisiones por 2 cierren
    conv(1,2), conv(2,4),                  # encoder: 2×16×16 → 4×8×8 (el latente: 256 números)
    deconv(4,2), deconv(2,1, act=False),  # decoder: 2×16×16 → 1×32×32
    nn.ZeroPad2d(-2),                      # recorta de vuelta a 28×28
    nn.Sigmoid()                           # píxeles en [0, 1]
)
# 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 notebook lo arma con UpsamplingNearest2d seguido de una conv 3×3 normal, en vez de usar ConvTranspose2d: el upsampling "nearest" duplica cada píxel en un bloque de 2×2 y la conv aprende a convertir esos bloques en bordes suaves. Es la construcción que evita los artefactos de tablero de ajedrez típicos de las convoluciones transpuestas, y la misma que usa el decoder del U-Net de la próxima unidad. El resultado, tras entrenar, reconstruye imágenes de Fashion-MNIST razonablemente bien.

Un detalle contable: el "cuello de botella" acá es de 4 canales × 8 × 8 = 256 números para una imagen de 784 píxeles — una compresión de aproximadamente 3×. El VAE de Stable Diffusion hace lo mismo a otra escala: de 3×512×512 a 4×64×64, una reducción de 48× en cantidad de valores.

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.

Hay además un problema de escala que empeora al anterior: nada obliga a que los latentes reales vivan cerca del origen ni tengan varianza 1. Si el encoder decidió que las remeras viven alrededor de (40, −12, …) y los zapatos alrededor de (−300, 8, …), un torch.randn ni siquiera cae cerca de ninguna isla. Un VAE ataca las dos cosas a la vez: obliga a las islas a estar cerca del origen, y las ensancha hasta que se solapan.

Lo que la KL le hace al espacio latente. Un autoencoder solo garantiza los puntos exactos que vio, y nada lo obliga a ponerlos cerca del origen; la KL de un VAE acerca cada media a cero y fuerza cada desvío hacia uno, hasta que las distribuciones se solapan y cualquier randn cae en territorio que el decoder conoce.

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

El reparametrization trick es la línea de z. Samplear de N(mu, σ²) directamente sería una operación aleatoria sin derivada: el gradiente de la loss no tendría forma de llegar a mu ni a lv, y el encoder no aprendería nada. Escribirlo como mu + σ · ε con ε ~ N(0,1) mueve el azar a un input externo (el randn_like) y deja a mu y σ en operaciones determinísticas derivables. Que se use lv (log-varianza) en vez de σ es por estabilidad numérica: la red puede producir cualquier número real y (0.5·lv).exp() siempre es positivo.

📐 Derivación formal de la Evidence Lower Bound (ELBO) y divergencia KL gaussiana

En un modelo generativo latente con datos $\mathbf{x}$ y variables latentes $\mathbf{z}$ , queremos maximizar la log-verosimilitud marginal $\log p_\theta(\mathbf{x}) = \log \int p_\theta(\mathbf{x}, \mathbf{z}) \, d\mathbf{z}$ . Como la integral es intratable, introducimos una distribución variacional $q_\phi(\mathbf{z} \mid \mathbf{x})$ parametrizada por el encoder:

\log p_\theta(\mathbf{x}) = \mathbb{E}_{q_\phi(\mathbf{z} \mid \mathbf{x})}\left[ \log \frac{p_\theta(\mathbf{x}, \mathbf{z})}{q_\phi(\mathbf{z} \mid \mathbf{x})} \right] + D_{\text{KL}}(q_\phi(\mathbf{z} \mid \mathbf{x}) \parallel p_\theta(\mathbf{z} \mid \mathbf{x}))

Dado que la divergencia de Kullback-Leibler es no-negativa ( $D_{\text{KL}} \ge 0$ ), el primer término constituye la cota inferior variacional (ELBO):

\mathcal{L}_{\text{ELBO}}(\theta, \phi; \mathbf{x}) = \mathbb{E}_{q_\phi(\mathbf{z} \mid \mathbf{x})}[\log p_\theta(\mathbf{x} \mid \mathbf{z})] - D_{\text{KL}}(q_\phi(\mathbf{z} \mid \mathbf{x}) \parallel p(\mathbf{z}))

El primer término es el error de reconstrucción (MSE o Binary Cross Entropy). El segundo término es la divergencia KL contra la previa estándar $p(\mathbf{z}) = \mathcal{N}(\mathbf{0}, \mathbf{I})$ . Para un espacio latente diagonal de $d$ dimensiones con media $\mu_j$ y varianza $\sigma_j^2$ :

D_{\text{KL}}(q \parallel p) = -\frac{1}{2} \sum_{j=1}^d \left( 1 + \log(\sigma_j^2) - \mu_j^2 - \sigma_j^2 \right) = \frac{1}{2} \sum_{j=1}^d \left( \mu_j^2 + \sigma_j^2 - 1 - \log(\sigma_j^2) \right)

Las derivadas revelan el comportamiento exacto del espacio:

$\frac{\partial D_{\text{KL}}}{\partial \mu_j} = \mu_j$ : fuerza la media a cero, concentrando las representaciones alrededor del origen.
$\frac{\partial D_{\text{KL}}}{\partial \sigma_j^2} = \frac{1}{2}\left(1 - \frac{1}{\sigma_j^2}\right)$ : el gradiente se anula únicamente cuando $\sigma_j^2 = 1$ . Si $\sigma_j^2 \to 0$ (colapso a punto determinista), la penalización $-\log(\sigma_j^2) \to +\infty$ explota.

Samplear de N(μ, σ²) es una operación sin derivada: el gradiente que vuelve de la loss muere ahí y el encoder no aprende nada. Escribir el mismo sampleo como μ + σ·ε, con el ruido ε tomado de afuera, deja a μ y σ en una suma y un producto, y el gradiente pasa.

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 bce_loss(inp, x): return F.binary_cross_entropy(inp[0], x)   # reconstrucción: píxeles en [0,1]
def vae_loss(inp, x): return kld_loss(inp, x) + bce_loss(inp, x)

La fórmula es la KL entre dos gaussianas en forma cerrada, y se lee término por término. Para cada dimensión latente, KL(N(mu, σ²) ‖ N(0,1)) = ½·(mu² + σ² − 1 − log σ²). El término mu² empuja cada media hacia 0 — las islas se acercan al origen. El par σ² − log σ² tiene su mínimo exactamente en σ = 1: penaliza tanto una distribución demasiado angosta (que volvería a ser un punto, con huecos alrededor) como una demasiado ancha. El resultado es que cada imagen del dataset ya no ocupa un punto sino una nube gaussiana, y todas las nubes se solapan alrededor del origen: cualquier randn cae dentro de alguna.

El peso relativo de los dos términos es la perilla más importante del VAE, y explica los dos modos de falla clásicos. Si la KL pesa demasiado, la solución más barata para el encoder es ignorar la imagen y devolver siempre mu=0, σ=1: la KL es exactamente cero, el latente no lleva información, y el decoder aprende a producir la imagen "promedio" del dataset para cualquier z — eso es posterior collapse, y se ve como reconstrucciones borrosas e idénticas entre sí. Si la KL pesa demasiado poco, el encoder achica σ hasta cero y aleja las medias, y volvés al autoencoder de §2 con sus huecos. La variante β-VAE hace explícito ese multiplicador como hiperparámetro; en el notebook el balance sale del hecho de que ambos términos usan .mean() sobre sus propias dimensiones.

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.

🔧 El componente que ya usás sin haberlo mirado por dentro

El VAE de Stable Diffusion (y de Flux2, en tu página de NVFP4 en difusión) hace exactamente esto — comprime una imagen de 512×512 píxeles a un latente mucho más chico donde corre la difusión, y lo decodifica de vuelta al final. La razón por la que difusión latente es tan eficiente comparada con difundir directamente sobre píxeles es precisamente esta: correr mil pasos de denoising sobre un latente comprimido es muchísimo más barato que correrlos sobre la imagen a resolución completa — el VAE es lo que hace viable ese ahorro sin perder calidad perceptible. En diffusers ese componente es AutoencoderKL, y el "KL" del nombre es exactamente el término de arriba.

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):
        super().__init__()
        self.scale = math.sqrt(ni)
        self.norm = nn.GroupNorm(1, ni)
        self.q = nn.Linear(ni, ni)
        self.k = nn.Linear(ni, ni)
        self.v = nn.Linear(ni, ni)
        self.proj = nn.Linear(ni, ni)

    def forward(self, inp):
        n,c,h,w = inp.shape
        x = self.norm(inp).view(n, c, -1).transpose(1, 2)   # [batch, h*w, canales]: cada píxel es un "token"
        q,k,v = self.q(x), self.k(x), self.v(x)
        s = (q @ k.transpose(1,2)) / self.scale             # [batch, h*w, h*w]: afinidad entre cada par de píxeles
        x = s.softmax(dim=-1) @ v
        x = self.proj(x).transpose(1,2).reshape(n,c,h,w)
        return x + inp   # residual: se suma la ENTRADA original, no la salida de la atención

⚠️ El bug del residual

Una versión anterior de este snippet reasignaba x con la salida de la atención y después devolvía self.proj(x)… + x: sumaba la proyección con la propia salida de la atención, y la entrada original se perdía. El residual deja de ser residual, el bloque ya no puede aproximar la identidad, y todo lo que la Unidad 2 explicó sobre por qué los skip connections estabilizan redes profundas se evapora en silencio: no hay error de shapes, solo un modelo que entrena peor. Por eso el parámetro se llama inp y no se toca nunca.

nn.GroupNorm(1, ni) con un solo grupo normaliza cada ejemplo sobre todos sus canales y posiciones a la vez — es la LayerNorm de imágenes de la Unidad 2, con parámetros de escala y desplazamiento por canal, y sin depender del batch. La escala √ni es la misma que en texto: evita que los productos punto crezcan con la dimensión y saturen el softmax.

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. Fijate que el camino de vuelta usa reshape y no view: después del transpose el tensor ya no es contiguo en memoria, y view fallaría; reshape copia si hace falta. 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.

El costo es la parte que decide dónde se puede usar: la matriz s tiene (h·w)² entradas por ejemplo y por head. A 16×16 son 65.536 valores — nada. A 64×64 (el latente de Stable Diffusion) son 16,7 millones por head, y a 512×512 en píxeles serían 68 mil millones, imposible. Por eso la atención va en las resoluciones bajas del U-Net y por eso la difusión latente es la que hace viable condicionar con texto.

La misma atención del Capítulo 1, con posiciones espaciales en lugar de palabras. El reshape es todo el truco; el precio es que s tiene una entrada por par de píxeles, cuadrático en h·w, y eso decide dónde puede ir un bloque de atención y por qué se difunde en el latente y no en píxeles.

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:

from einops import rearrange

class SelfAttentionMultiHead(nn.Module):
    def __init__(self, ni, nheads):
        super().__init__()
        self.nheads = nheads
        self.scale = math.sqrt(ni/nheads)
        self.norm = nn.GroupNorm(1, ni)
        self.qkv = nn.Linear(ni, ni*3)       # q, k y v en una sola proyección
        self.proj = nn.Linear(ni, ni)

    def forward(self, inp):
        n,c,h,w = inp.shape
        x = self.norm(inp).view(n, c, -1).transpose(1, 2)
        x = self.qkv(x)
        x = rearrange(x, 'n s (h d) -> (n h) s d', h=self.nheads)   # cada head pasa a ser un "ejemplo" más
        q,k,v = torch.chunk(x, 3, dim=-1)
        s = (q @ k.transpose(1,2)) / self.scale
        x = s.softmax(dim=-1) @ v
        x = rearrange(x, '(n h) s d -> n s (h d)', h=self.nheads)   # y se vuelven a juntar
        x = self.proj(x).transpose(1,2).reshape(n,c,h,w)
        return x + inp

🟢 En producción, el núcleo es una sola llamada

Las cuatro líneas del medio (producto q@kᵀ, escala, softmax, producto por v) hoy se reemplazan por F.scaled_dot_product_attention(q, k, v), que en GPU despacha a un kernel fusionado (FlashAttention o equivalente) que nunca materializa la matriz (h·w)² completa en memoria. Es la misma función que usa vLLM por debajo para el prefill; los números del párrafo anterior sobre memoria cuadrática son el motivo de que esos kernels existan.

⚡ U-Net condicional: Sinusoidal Timestep Embeddings y Cross-Attention tensorial

En Stable Diffusion, el U-Net recibe dos señales de condicionamiento simultáneas: en qué momento de la difusión estamos ( $t$ ) y qué queremos generar (texto $c$ ). La mecánica tensorial de ambas piezas es rigurosa:

Sinusoidal Timestep Embeddings: Pasar el escalar $t \in [0, 1000]$ como un único float numérico es inestable: una red convolucional no tiene sesgo inductivo para interpretar magnitudes absolutas tan variables. Siguiendo a Vaswani et al. y Ho et al., el timestep se proyecta en una base ortogonal de frecuencias sinusoidales geométricas: $\text{PE}_{(t, 2i)} = \sin\left(\frac{t}{10000^{2i/d}}\right), \quad \text{PE}_{(t, 2i+1)} = \cos\left(\frac{t}{10000^{2i/d}}\right)$ Un pequeño MLP de dos capas lineales con activación SiLU proyecta este vector $[B, d]$ hacia cada bloque residual del U-Net, modulando las activaciones vía normalización adaptativa (AdaGN): $\mathbf{y} = \gamma(t) \cdot \text{GroupNorm}(\mathbf{x}) + \beta(t)$ .
Mecánica tensorial de la Cross-Attention: Para inyectar el texto en la imagen, el mapa de características del U-Net $\mathbf{X}_{\text{img}} \in \mathbb{R}^{B \times C \times H \times W}$ se aplana en una secuencia de tokens espaciales $\mathbf{X} \in \mathbb{R}^{B \times (H \cdot W) \times C}$ . El texto codificado por CLIP provee $\mathbf{C}_{\text{text}} \in \mathbb{R}^{B \times L \times d_c}$ (en SD 1.4, $L=77, d_c=768$ ): $\mathbf{Q} = \mathbf{X} \mathbf{W}_Q \in \mathbb{R}^{B \times (H \cdot W) \times d_k}, \quad \mathbf{K} = \mathbf{C}_{\text{text}} \mathbf{W}_K \in \mathbb{R}^{B \times L \times d_k}, \quad \mathbf{V} = \mathbf{C}_{\text{text}} \mathbf{W}_V \in \mathbb{R}^{B \times L \times d_v}$ La matriz de afinidad cruzada $\mathbf{A} = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_k}}\right) \in \mathbb{R}^{B \times (H \cdot W) \times L}$ asigna un peso de atención a cada token textual para cada píxel espacial de la imagen: cada parte del dibujo decide a qué palabras prestar atención durante la síntesis.

5. Errores comunes

Los tres bloques de esta unidad fallan de formas silenciosas: casi nunca hay excepción, solo un modelo que "no genera bien".

Síntoma	Causa probable	Qué hacer
El VAE entrena, la reconstrucción mejora, pero `mu` y `lv` no cambian nunca de su valor inicial.	Sampleaste `z` con `torch.normal(mu, std)` o similar sin reparametrizar: no hay gradiente a través del sampling hacia el encoder.	`z = mu + (0.5lv).exp() torch.randn_like(lv)` (§3). Verificá con `mu.grad` distinto de `None` después de un `backward`.
Reconstrucciones grises, lavadas, y la loss se estanca alto.	`Sigmoid` final combinada con MSE: la MSE penaliza poco los errores cerca de 0 y 1 justo donde la sigmoid satura, y el gradiente que llega es minúsculo.	Con salida sigmoid y targets en `[0,1]`, usar BCE (`F.binary_cross_entropy`, o mejor `binary_cross_entropy_with_logits` sin la sigmoid en el modelo). MSE va con salida lineal.
Todas las imágenes generadas por el VAE son la misma mancha borrosa.	Posterior collapse: la KL pesa demasiado y el encoder devuelve `N(0,1)` para cualquier entrada.	Bajar el peso de la KL (o subirlo gradualmente durante el entrenamiento, "KL annealing"), y mirar que `kld_loss` no esté clavada en cero.
`CUDA out of memory` al agregar atención al U-Net, aunque las convoluciones entraban bien.	Atención a resolución completa: la matriz de afinidades es `(h·w)²` por head y por ejemplo — cuadrática en píxeles.	Ponerla solo en las resoluciones bajas (16×16, 32×32 en el latente), o usar `F.scaled_dot_product_attention` que no materializa la matriz.
`RuntimeError: view size is not compatible with input tensor's size and stride`	`.view()` sobre un tensor no contiguo — típicamente justo después de un `transpose`.	Usar `.reshape()` (copia si hace falta) o `.contiguous().view()`.
El bloque de atención "no aporta": el modelo con atención entrena igual o peor que sin ella.	Residual olvidado o mal sumado (la salida consigo misma, como en el bug de §4): el bloque no puede arrancar como identidad.	Guardar la entrada en una variable que no se reasigne y devolver `salida + inp`. Chequeo rápido: con los pesos de `proj` en cero, el bloque tiene que devolver exactamente `inp`.

6. Resumen

Un autoencoder comprime y reconstruye (conv con stride para bajar, deconv = upsampling + conv para subir), pero su espacio latente tiene "huecos" que el decoder nunca aprendió a interpretar.
Un VAE fuerza ese espacio a parecerse a una distribución normal (vía divergencia KL, cuyos términos empujan mu → 0 y σ → 1) — por eso samplear ruido puro y decodificar da imágenes coherentes. El reparametrization trick es lo que permite que el gradiente atraviese el sampling.
El balance entre reconstrucción y KL es la perilla central: demasiada KL da posterior collapse, demasiado poca devuelve un autoencoder con huecos.
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.
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). Su memoria es cuadrática en h·w, y el residual tiene que sumar la entrada original.

7. Ponete a prueba

Respondé antes de abrir cada respuesta. Si fallás una, volvé a la sección indicada.

1. ¿Por qué el decoder de un autoencoder común produce basura si le das torch.randn, aunque reconstruya bien?

Porque solo fue entrenado en los puntos del latente que corresponden a imágenes reales, sin ningún incentivo para que el espacio entre ellos tenga sentido, y además nada obliga a esos puntos a vivir cerca del origen con varianza 1. Un randn cae en un hueco (o directamente lejos de todo). (§2)

2. ¿Qué rompe si reemplazás z = mu + (0.5*lv).exp() * randn_like(lv) por un sampleo directo de la normal?

El gradiente: samplear es una operación sin derivada, así que la loss no puede propagar nada hacia mu ni lv y el encoder no aprende. El reparametrization trick mueve el azar a un input externo y deja mu y σ en operaciones derivables. (§3)

3. En la KL, ¿qué término empuja las medias a cero y cuál fija el desvío en 1? ¿Qué pasa si multiplicás toda la KL por 100?

mu² empuja las medias al origen; σ² − log σ² tiene su mínimo en σ=1. Con un peso enorme, la salida más barata es mu=0, σ=1 para toda entrada: KL exactamente cero, latente sin información, decoder que produce la imagen promedio — posterior collapse. (§3)

4. Un bloque de atención sobre un mapa de 128×128 con 8 heads: ¿cuántas entradas tiene la matriz de afinidades por ejemplo, y qué hacer si no entra en memoria?

(128·128)² = 16.384² ≈ 268 millones por head, por 8 heads ≈ 2.150 millones de valores por ejemplo: en float16 son más de 4 GB por imagen del batch. Mover la atención a resoluciones más bajas del U-Net, o usar scaled_dot_product_attention con un kernel fusionado que no materializa la matriz. (§4)

5. ¿Por qué la vuelta a [n,c,h,w] usa reshape y la ida usa view?

A la ida el tensor recién salido de GroupNorm es contiguo, así que view funciona. A la vuelta se acaba de hacer un transpose, que cambia los strides sin mover datos: el tensor no es contiguo y view lanza error; reshape copia si hace falta. (§4, §5)

6. Con los pesos de self.proj puestos en cero, ¿qué tiene que devolver SelfAttention.forward(inp)? ¿Y qué devolvía la versión con el bug?

Exactamente inp: la salida de la atención proyectada es cero y el residual suma la entrada original. La versión con el bug devolvía 0 + x donde x era la salida de la atención sin proyectar — la entrada se perdía y el bloque no podía ser la identidad. (§4)

8. Para practicar

✏️ Ejercicio 1 — La KL como función de mu y σ (CPU, sin datos)

Evaluá la fórmula de kld_loss sobre una grilla de valores y confirmá que el mínimo está en mu=0, σ=1. Después mirá qué pasa con el gradiente a través del reparametrization trick.

import torch

def kl(mu, lv): return -0.5 * (1 + lv - mu**2 - lv.exp())

for mu in (0., 1., 3.):
    for sigma in (0.1, 0.5, 1., 2.):
        lv = torch.tensor(sigma**2).log()
        print(f"mu={mu:<3} sigma={sigma:<4} KL={kl(torch.tensor(mu), lv):.3f}")

# reparametrization: ¿llega gradiente a mu y lv a través del sampleo?
mu = torch.zeros(4, requires_grad=True)
lv = torch.zeros(4, requires_grad=True)
z = mu + (0.5*lv).exp() * torch.randn(4)
(z**2).sum().backward()
print("grad mu:", mu.grad, "grad lv:", lv.grad)   # ambos distintos de None y de cero

Pista / qué esperar

KL=0.000 solo en mu=0, sigma=1; crece con mu² y también al alejarse de σ=1 en cualquier dirección (mirá que σ=0.1 penaliza más que σ=2: el −log σ² explota cerca de cero). Los dos gradientes salen definidos. Si reemplazás la línea de z por torch.normal(mu, (0.5*lv).exp()) el gradiente llega igual (PyTorch reparametriza por dentro), pero con torch.normal(mu.detach(), ...) o con un sampleo con numpy, no llega — y ese es el caso que se da al escribir el sampleo "a mano" fuera del grafo.

✏️ Ejercicio 2 — Verificar la self-attention contra PyTorch y cazar el bug

Implementá SelfAttention tal como está en §4, compará el núcleo con F.scaled_dot_product_attention, y confirmá la propiedad del residual: con proj en cero, la salida es la entrada.

import torch, math, torch.nn as nn, torch.nn.functional as F
torch.manual_seed(0)

class SelfAttention(nn.Module):
    def __init__(self, ni):
        super().__init__()
        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, inp):
        n,c,h,w = inp.shape
        x = self.norm(inp).view(n, c, -1).transpose(1, 2)
        q,k,v = self.q(x), self.k(x), self.v(x)
        s = (q @ k.transpose(1,2)) / self.scale
        a = s.softmax(dim=-1) @ v
        ref = F.scaled_dot_product_attention(q, k, v)         # misma matemática, kernel fusionado
        assert torch.allclose(a, ref, atol=1e-5), "el núcleo no coincide"
        return self.proj(a).transpose(1,2).reshape(n,c,h,w) + inp

sa = SelfAttention(32)
inp = torch.randn(2, 32, 16, 16)
out = sa(inp)
print("shape:", out.shape)

with torch.no_grad():
    sa.proj.weight.zero_(); sa.proj.bias.zero_()
print("con proj=0, salida == entrada:", torch.allclose(sa(inp), inp))

# memoria de la matriz de afinidades, por ejemplo y por head, en float16:
for hw in (16, 32, 64, 128, 512):
    print(f"{hw}x{hw}: {(hw*hw)**2 * 2 / 2**20:,.0f} MiB")

Pista / qué esperar

shape: torch.Size([2, 32, 16, 16]), el assert pasa, y True para el chequeo del residual. Ahora reproducí el bug: cambiá + inp por + a.transpose(1,2).reshape(n,c,h,w) (la salida de la atención) y el chequeo da False — sin ningún error de shapes. La tabla de memoria pasa de 0.1 MiB a 16×16 a 128 GiB a 512×512: por eso la atención vive en el latente.

9. Fuentes y para seguir

fastai/course22p2 — notebooks 08_autoencoder.ipynb, 29_vae.ipynb y 27_attention.ipynb; el segundo tiene la comparación visual autoencoder vs. VAE generando desde ruido.
Practical Deep Learning for Coders — Part 2 — la Lección 24 construye la atención en vivo y la verifica contra diffusers.
Auto-Encoding Variational Bayes (Kingma & Welling, 2013) — el VAE original; el apéndice B tiene la KL en forma cerrada de §3.
Attention Is All You Need (Vaswani et al., 2017) — la sección 3.2 es toda la matemática de §4; el resto del paper es el transformer que ya conocés.
AutoencoderKL en diffusers — el VAE de Stable Diffusion tal como se usa en producción; compará su encode/decode con el de §3.
einops — la documentación de rearrange; leer la notación 'n s (h d) -> (n h) s d' con fluidez ahorra muchos bugs de reshape.