La arquitectura completa: encoder, decoder y máscaras

Vista de pájaro

El Transformer original es un modelo encoder-decoder para traducción. El encoder lee la oración fuente completa y produce una representación contextual de cada token. El decoder genera la traducción token por token, consultando esa representación.

Ambos son pilas de $N = 6$ bloques idénticos (en el modelo base). Cada bloque es una composición de piezas que ya conoces más dos ingredientes nuevos: la red feed-forward por posición y la combinación residual + LayerNorm.

Bloque encoder:

1. Multi-head self-attention.
2. Feed-forward por posición.

Bloque decoder:

1. Multi-head self-attention enmascarada (causal).
2. Multi-head cross-attention sobre la salida del encoder.
3. Feed-forward por posición.

Cada sub-capa está envuelta en una conexión residual seguida de normalización.

Residual + LayerNorm

Cada sub-capa del Transformer no reemplaza su entrada: le suma su contribución.

Las conexiones residuales (residual connections), heredadas de ResNet, crean un camino directo para el gradiente a través de toda la profundidad. LayerNorm normaliza cada vector de token a media 0 y varianza 1 sobre sus $d_{\text{model}}$ componentes, y luego aplica una escala y un sesgo aprendidos:

a cada vector le restamos su promedio ($\mu$) y lo dividimos por cuánto se dispersa ($\sigma$); después $\gamma$ y $\beta$, que se aprenden, lo reescalan y lo corren. $\epsilon$ es un número chiquito para no dividir por cero

A diferencia de BatchNorm, no depende de estadísticas del batch, lo que la hace estable con secuencias de longitud variable y en inferencia token a token.

Para que la suma residual tenga sentido, todas las sub-capas producen salidas de dimensión $d_{\text{model}} = 512$. Esa dimensión constante es lo que da lugar a la idea del residual stream: un "bus" de 512 dimensiones que atraviesa todo el modelo y en el que cada sub-capa lee y escribe.

La red feed-forward por posición

Después de la atención, cada token pasa independientemente por la misma red de dos capas:

con $W_1 \in \mathbb{R}^{d_{\text{model}} \times d_{ff}}$, $W_2 \in \mathbb{R}^{d_{ff} \times d_{\text{model}}}$ y $d_{ff} = 2048$: la capa interna expande la dimensión cuatro veces.

La división de trabajo es clara: la atención mezcla información entre posiciones; la FFN procesa cada posición por separado. La FFN concentra aproximadamente dos tercios de los parámetros de cada bloque, y la investigación posterior en interpretabilidad sugiere que funciona como una especie de memoria clave-valor donde se almacena buena parte del conocimiento factual del modelo.

La máscara causal

El decoder genera la traducción de izquierda a derecha: para predecir el token $t$ solo puede usar los tokens $1, \dots, t-1$. En inferencia esto es natural, porque los tokens futuros todavía no existen. Pero en entrenamiento queremos procesar la oración objetivo completa en paralelo (esa es toda la gracia del Transformer). Si no hacemos nada, la posición $t$ podría simplemente "mirar" la respuesta en la posición $t$.

La solución es una máscara: antes del softmax, se pone $-\infty$ en todos los puntajes $S_{ij}$ con $j > i$. Como $e^{-\infty} = 0$, esas posiciones reciben peso exactamente cero:

a los puntajes les sumamos $M$: cero si la posición $j$ es la actual o una anterior, menos infinito si viene después. Así ninguna posición puede mirar el futuro

El resultado es una matriz de atención triangular inferior. Además, la salida del decoder se desplaza una posición a la derecha (se antepone un token de inicio), de modo que la posición $t$ recibe como entrada el token $t-1$ y debe predecir el token $t$. A esto se le llama teacher forcing.

Mira cómo cambia el patrón de atención con la máscara: toda la mitad superior derecha queda en cero.

Embeddings compartidos y la capa de salida

Al final del decoder, una capa lineal proyecta cada vector de $d_{\text{model}}$ al tamaño del vocabulario, y un softmax da la distribución sobre el siguiente token. El paper comparte la misma matriz de pesos entre los dos embeddings (fuente y objetivo, con un vocabulario compartido) y la proyección pre-softmax (Sección 3.4). Esto reduce parámetros y tiene sentido semántico: el embedding de un token y el vector que "predice" ese token viven en el mismo espacio.

Implementación completa

Aquí tienes un bloque encoder y uno decoder en PyTorch, fieles al paper (Post-LN), reutilizando la MultiHeadAttention de la lección anterior:

import torch
import torch.nn as nn

class FeedForward(nn.Module):
    def __init__(self, d_model=512, d_ff=2048, dropout=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )

    def forward(self, x):
        return self.net(x)


class EncoderLayer(nn.Module):
    def __init__(self, d_model=512, n_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ff = FeedForward(d_model, d_ff, dropout)
        self.norm1, self.norm2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x, src_mask=None):
        x = self.norm1(x + self.drop(self.attn(x, x, x, src_mask)))
        return self.norm2(x + self.drop(self.ff(x)))


class DecoderLayer(nn.Module):
    def __init__(self, d_model=512, n_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ff = FeedForward(d_model, d_ff, dropout)
        self.norms = nn.ModuleList([nn.LayerNorm(d_model) for _ in range(3)])
        self.drop = nn.Dropout(dropout)

    def forward(self, y, memory, tgt_mask, src_mask=None):
        y = self.norms[0](y + self.drop(self.self_attn(y, y, y, tgt_mask)))
        y = self.norms[1](y + self.drop(self.cross_attn(y, memory, memory, src_mask)))
        return self.norms[2](y + self.drop(self.ff(y)))


def causal_mask(n: int) -> torch.Tensor:
    # 1 = permitido, 0 = bloqueado; forma (1, 1, n, n) para broadcast sobre (B, h)
    return torch.tril(torch.ones(n, n)).view(1, 1, n, n)

Fíjate en la cross-attention: self.cross_attn(y, memory, memory). Las queries salen del decoder (y), las keys y values de la salida del encoder (memory). La src_mask sirve para ignorar los tokens de padding de la oración fuente.

Encoder-only, decoder-only

Una consecuencia histórica de esta arquitectura modular es que sus mitades resultaron útiles por separado:

• Solo encoder (BERT): atención bidireccional sin máscara causal; ideal para entender texto (clasificación, extracción).
• Solo decoder (GPT): self-attention causal, sin cross-attention; ideal para generar texto. Es la base de prácticamente todos los LLM conversacionales actuales.
• Encoder-decoder (T5, el original): sigue siendo natural para tareas de transformación de secuencia a secuencia.

Resumen

• Encoder: self-attention + FFN. Decoder: self-attention causal + cross-attention + FFN.
• Cada sub-capa usa residual + LayerNorm; el paper usa Post-LN, los modelos modernos Pre-LN.
• La FFN procesa cada posición por separado y concentra la mayoría de los parámetros.
• La máscara causal permite entrenar todas las posiciones en paralelo sin hacer trampa.