Skip to content

English · Español

01 — El bloque transformer: anatomía Pre-LN

🇪🇸 El bloque es asombrosamente simple cuando lo ves bien: dos sublayers, dos LayerNorms, dos sumas residuales. La sutileza es el orden — LN antes de la sublayer, no después. Esa pequeña diferencia se ganó la guerra contra Post-LN porque entrena más estable.

El bloque, en un diagrama

        ┌─────────── residual stream (d_model) ─────────────┐
        │                                                   │
   x ──>┤                                                   │
        │                                                   │
        ├──> LayerNorm ──> MultiHeadAttention ──┐            │
        │                                       ▼            │
        ├──────────────────────────────────────(+)──> z ─────┤
        │                                                    │
        ├──> LayerNorm ──> FFN ────────────────┐             │
        │                                       ▼            │
        └──────────────────────────────────────(+)──> y ─────┘

En ecuaciones:

\[ \begin{aligned} z &= x + \text{MHA}(\text{LN}_1(x)) \\ y &= z + \text{FFN}(\text{LN}_2(z)) \end{aligned} \]

Eso es el bloque entero. Dos normalizaciones, dos sublayers, dos sumas. La salida \(y\) tiene la misma forma que \(x\): \((T, d_\text{model})\) — longitud de secuencia \(T\), ancho residual \(d_\text{model}\). Apilar bloques es directo porque las formas coinciden.

Pre-LN vs Post-LN — la única pieza de historia que importa

El transformer original (Vaswani et al. 2017) usó Post-LN:

\[ \begin{aligned} z &= \text{LN}_1(x + \text{MHA}(x)) \\ y &= \text{LN}_2(z + \text{FFN}(z)) \end{aligned} \]

LN va fuera de la suma residual. Esto causaba inestabilidad en el entrenamiento: la varianza del residual stream puede crecer sin límite con la profundidad, ya que cada bloque añade una salida de sublayer no normalizada. Las pilas profundas (12+ capas) necesitaban schedules de warmup cuidadosos para evitar explosiones de pérdida.

Pre-LN (Xiong et al. 2020 "On Layer Normalization in the Transformer Architecture") mueve LN dentro de la suma residual, normalizando la entrada a cada sublayer en lugar de la suma:

\[y = x + \text{sublayer}(\text{LN}(x))\]

Esto acota el crecimiento de la varianza y hace el warmup casi innecesario. Cada transformer moderno (de GPT-2 en adelante, LLaMA, PaLM, Gemini, Claude) usa Pre-LN. Nosotros usamos Pre-LN. Post-LN se menciona sólo por contexto.

Un detalle pequeño pero importante: con Pre-LN, la entrada a la cabeza LM no está todavía normalizada — el residual stream tras el último bloque ha acumulado salidas de sublayer sin normalizar. Por ello los transformers modernos añaden un LayerNorm final tras el último bloque y antes de la cabeza LM. No te olvides (el lab 01 lo pillará).

LayerNorm, brevemente

Implementaste la intuición de LN en la Fase 10. La fórmula:

\[\text{LN}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \varepsilon}} + \beta\]

donde \(\mu, \sigma^2\) se calculan sobre la dimensión de features (el último eje, de tamaño \(d_\text{model}\)), no sobre los ejes de batch o secuencia. \(\gamma, \beta \in \mathbb{R}^{d_\text{model}}\) se aprenden por instancia de LN. \(\varepsilon \approx 10^{-5}\) para fp32 (a veces \(10^{-6}\)).

Dos parámetros por feature: escala \(\gamma\) y desplazamiento \(\beta\). Para un Mini-GPT con 2 capas, eso son \(2 \times 2 \times d_\text{model} = 256\) parámetros para todos los LNs dentro de los bloques, más \(2 \times d_\text{model} = 128\) para el LN final. Total de parámetros de LN: 384 — un error de redondeo en el inventario completo.

A veces LayerNorm se reemplaza por RMSNorm (Zhang & Sennrich 2019) en los transformers modernos — quitar la sustracción de la media y \(\beta\). Más barato, rendimiento similar. La Fase 17 se queda con LayerNorm por fidelidad a GPT-2. RMSNorm es un cambio de una línea si hiciera falta en una fase posterior.

El forward completo, capa por capa

Sea \(T\) = longitud de secuencia, \(|V|\) = tamaño del vocabulario, \(L\) = \(n_\text{layers}\).

Input:  tokens (T,) int32
[1]     E[tokens]              shape (T, d_model)         token embedding lookup
        (no se suma PE aquí — RoPE se aplica dentro de attention)
[2]     for l in range(L):
            h = h + MHA_l(LN_a(h))    shape (T, d_model)   sublayer de attention
            h = h + FFN_l(LN_b(h))    shape (T, d_model)   sublayer de FFN
[3]     h = LN_final(h)           shape (T, d_model)       norma final
[4]     logits = h @ E.T          shape (T, |V|)           cabeza LM atada
Output: logits (T, |V|) float32

Seis tensores fluyen por este forward, todos con formas predecibles. El Lab 00 de Borja los traza a mano sobre una configuración minúscula.

Una lectura del residual stream

Tres observaciones desde la vista del residual stream:

  1. El residual stream lo lee cada sublayer y lo escribe cada sublayer. Nada más puede comunicar. No hay estado global, no hay skips entre bloques. La información fluye sólo por el stream.
  2. Las sublayers son aditivas. Un bloque dado puede o bien contribuir (la suma es no nula) o bien hacer no-op (la suma es casi cero — pasa con algunos heads en modelos entrenados). El skip está siempre disponible.
  3. Attention escribe el resultado de la mezcla entre posiciones; FFN escribe el resultado de la transformación por posición. Attention es comunicación; FFN es computación. El bloque las empareja.

Enmascaramiento causal — sigue siendo necesario, incluso con RoPE

Una confusión habitual: RoPE proporciona información posicional, pero no hace causal a attention. Sigues necesitando enmascarar la matriz de attention para que cada posición sólo pueda atender a sí misma y a las anteriores:

\[\text{mask}[i, j] = \begin{cases} 0 & \text{si } j \le i \\ -\infty & \text{en otro caso} \end{cases}\]

Se suma a los scores de attention pre-softmax. La posición 0 atiende sólo a la posición 0. La posición 7 atiende a 0–7. Sin esta máscara, el modelo ve tokens futuros — lo cual está bien en inferencia pero es ruinoso en entrenamiento (el modelo simplemente copia el siguiente token).

La máscara causal vive dentro de la implementación de MHA (Fase 15). El lab 03 verifica, por perturbación, que la máscara está correctamente cableada hasta la salida.

Cómo se ve el bloque en código (sólo esbozo — la implementación va en el lab)

class TransformerBlock:
    def __init__(self, d_model, n_heads, d_ff):
        self.ln1 = LayerNorm(d_model)
        self.attn = MultiHeadAttention(d_model, n_heads)  # de la Fase 15
        self.ln2 = LayerNorm(d_model)
        self.ffn = FFN(d_model, d_ff)

    def forward(self, x):
        # x: (T, d_model)
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

Cinco líneas no triviales. El lab 00 de la Fase 17 construye esto a mano; el lab 01 lo apila.

Qué NO cubre este archivo

  • La matemática de attention. Fase 15. Trata MHA como: entrada \((T, d_\text{model})\), salida \((T, d_\text{model})\), aplica proyecciones Q/K/V, RoPE sobre Q y K, scaled dot-product attention con máscara causal, proyección de salida.
  • El FFN. Siguiente archivo (02-ffn-and-activations.md).
  • Tied embedding + cabeza LM. Archivo 03.

Siguiente: 02-ffn-and-activations.md