Construyendo un Transformer desde Cero

¿Por Qué Construirlo Nosotros Mismos?

A lo largo de los últimos ocho artículos hemos examinado cada componente de la arquitectura transformer de forma aislada: puntuaciones de atención, proyecciones Q/K/V, enmascaramiento causal, atención multi-cabezal, codificación posicional, conexiones residuales, redes feed-forward, y cómo estas piezas se combinan en encoders, decoders y modelos encoder-decoder. Cada concepto tenía sentido por sí solo, pero existe una brecha entre entender cada parte y ver cómo se conectan en un único modelo ejecutable. El objetivo de este artículo es cerrar esa brecha escribiendo un transformer decoder-only mínimo en PyTorch (aproximadamente 130 líneas de código), priorizando la claridad sobre el rendimiento.

No usaremos flash attention, KV caching, kernels fusionados, ni ninguna otra optimización. Cada línea se corresponde directamente con un concepto de un artículo anterior, y señalaremos esas conexiones a medida que avancemos. Al final, tendremos un modelo que se entrena en una tarea de juguete y produce resultados correctos, lo cual es una verificación útil de que nuestra comprensión de la teoría realmente se sostiene en la práctica.

¿Cómo Se Convierten los Tokens en Vectores?

Un transformer opera sobre vectores continuos, no sobre IDs de tokens discretos, así que el primer paso es convertir cada índice de token en un vector denso e inyectar información posicional. Cubrimos por qué la posición importa en el artículo 5: la auto-atención es equivariante a permutaciones, lo que significa que trata la entrada como un conjunto a menos que codifiquemos el orden explícitamente. El enfoque estándar de Vaswani et al. (2017) añade codificaciones posicionales sinusoidales a los token embeddings, aunque los embeddings posicionales aprendidos (como los usados en GPT-2) tienden a funcionar igual de bien para longitudes de contexto fijas. Usaremos embeddings aprendidos aquí porque el código es más simple.

La capa de embedding toma un tensor de IDs de tokens con forma $(B, T)$ (tamaño de batch $B$, longitud de secuencia $T$) y produce un tensor de forma $(B, T, d_{\text{model}})$. El embedding posicional hace lo mismo para los índices de posición $[0, 1, \ldots, T-1]$, y sumamos ambos elemento a elemento. La siguiente clase maneja ambos.

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class TokenEmbedding(nn.Module):
    """Token + positional embeddings (article 5)."""
    def __init__(self, vocab_size, d_model, max_seq_len):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_seq_len, d_model)

    def forward(self, x):
        B, T = x.shape
        tok = self.token_emb(x)                          # (B, T, d_model)
        pos = self.pos_emb(torch.arange(T, device=x.device))  # (T, d_model)
        return tok + pos                                  # (B, T, d_model)

Algo que vale la pena notar es que pos_emb se transmite (broadcast) a lo largo de la dimensión del batch. Cada secuencia en el batch recibe la misma codificación posicional, lo cual tiene sentido porque la posición 3 significa lo mismo independientemente de a qué oración pertenezca. La salida es un tensor $(B, T, d_{\text{model}})$ que lleva tanto información de contenido como de posición, listo para las capas de atención.

De Atención de Una Cabeza a Multi-Cabezal

Con las entradas embebidas en mano, necesitamos el mecanismo que permite a los tokens comunicarse entre sí. Construimos la intuición para esto en los artículos 2 a 4: cada token se proyecta en una query, key y value, las queries y keys producen puntuaciones de atención mediante un producto escalar escalado, enmascaramos las posiciones futuras para imponer causalidad, y los values ponderados por softmax se convierten en la salida. Comencemos con una sola cabeza de atención y luego extendamos a múltiples cabezas.

Recordemos la fórmula de atención por producto escalar escalado del artículo 2:

El escalado por $\sqrt{d_k}$ (donde $d_k$ es la dimensión de cada cabeza) evita que los productos escalares crezcan en magnitud a medida que aumenta la dimensión, lo que empujaría al softmax a regiones saturadas donde los gradientes se desvanecen. Sin él, el entrenamiento tiende a ser inestable para valores de $d_k$ superiores a aproximadamente 32. La máscara causal del artículo 3 establece las entradas triangulares superiores en $-\infty$ antes del softmax, asegurando que la posición $i$ solo pueda atender a las posiciones $j \leq i$.

Una sola cabeza captura un tipo de relación (quizás adyacencia sintáctica, quizás correferencia). El artículo 4 argumentó que queremos múltiples cabezas ejecutándose en paralelo, cada una con sus propias proyecciones Q/K/V operando sobre una porción de la dimensión del embedding, para que el modelo pueda atender a diferentes tipos de relaciones simultáneamente. Si tenemos $h$ cabezas y dimensión de embedding $d_{\text{model}}$, cada cabeza opera sobre $d_k = d_{\text{model}} / h$ dimensiones. En la práctica, implementamos esto proyectando Q, K y V a las $d_{\text{model}}$ dimensiones completas con una sola capa lineal, y luego reorganizando en $h$ cabezas. La siguiente clase implementa la auto-atención causal multi-cabezal completa.

class CausalSelfAttention(nn.Module):
    """Multi-head causal self-attention (articles 2-4)."""
    def __init__(self, d_model, n_heads, max_seq_len):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.qkv_proj = nn.Linear(d_model, 3 * d_model)  # single projection for Q, K, V
        self.out_proj = nn.Linear(d_model, d_model)

        # Precompute causal mask (article 3): lower-triangular = 1, upper = 0
        mask = torch.tril(torch.ones(max_seq_len, max_seq_len))
        self.register_buffer("mask", mask.unsqueeze(0).unsqueeze(0))  # (1, 1, T, T)

    def forward(self, x):
        B, T, C = x.shape
        # Project to Q, K, V and split into heads
        qkv = self.qkv_proj(x)                           # (B, T, 3*d_model)
        q, k, v = qkv.chunk(3, dim=-1)                   # each (B, T, d_model)

        # Reshape: (B, T, d_model) -> (B, n_heads, T, d_k)
        q = q.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        # Scaled dot-product attention (article 2)
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)  # (B, h, T, T)

        # Apply causal mask (article 3): set future positions to -inf
        scores = scores.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))

        weights = F.softmax(scores, dim=-1)               # (B, h, T, T)
        out = weights @ v                                  # (B, h, T, d_k)

        # Concatenate heads and project (article 4)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(out)                          # (B, T, d_model)

Hay algunos detalles que vale la pena destacar. Usamos una sola capa lineal qkv_proj para calcular Q, K y V en una sola multiplicación de matrices, y luego dividimos con chunk , lo cual es matemáticamente idéntico a tres proyecciones separadas pero más rápido porque emitimos un solo GEMM en lugar de tres. La máscara causal se registra como un buffer (no como un parámetro) para que se mueva automáticamente al dispositivo correcto y no sea actualizada por el optimizador. Y la out_proj final es la capa lineal aprendida que recombina las salidas concatenadas de las cabezas, como se describió en el artículo 4.

Ensamblando el Bloque Transformer y Apilándolo

Con la atención y los embeddings implementados, necesitamos dos ingredientes más de los artículos 6 y 7: conexiones residuales y la red feed-forward posicional (FFN). Un bloque transformer aplica atención, suma el resultado de vuelta a la entrada a través de una conexión residual, normaliza con layer norm, luego pasa por una FFN de dos capas con otra conexión residual y otra layer norm. Esta es la variante pre-norm (layer norm antes de cada sub-capa), que tiende a entrenar de forma más estable que el diseño post-norm original, y es lo que usan GPT-2 y la mayoría de los modelos decoder modernos.

La FFN del artículo 7 es un MLP simple de dos capas que expande la dimensión por un factor de 4, aplica una no linealidad y proyecta de vuelta hacia abajo. Usaremos GELU como la activación, siguiendo a GPT-2.

class FeedForward(nn.Module):
    """Position-wise feed-forward network (article 7)."""
    def __init__(self, d_model):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model),
        )

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


class TransformerBlock(nn.Module):
    """Pre-norm transformer block: LN -> attention -> residual -> LN -> FFN -> residual (article 6)."""
    def __init__(self, d_model, n_heads, max_seq_len):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads, max_seq_len)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = FeedForward(d_model)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))   # residual around attention
        x = x + self.ffn(self.ln2(x))    # residual around FFN
        return x

Cada conexión residual (el patrón x = x + ... ) cumple el propósito que discutimos en el artículo 6: permite que los gradientes fluyan directamente a través de la suma, evitando completamente los parámetros de la sub-capa si es necesario, lo que facilita mucho el entrenamiento de pilas profundas. Sin conexiones residuales, un transformer de 6 capas a menudo no converge en absoluto.

Ahora podemos apilar $N$ de estos bloques para formar el modelo completo. El modelo de lenguaje añade una layer norm final después del último bloque, luego una proyección lineal de $d_{\text{model}}$ al tamaño del vocabulario, produciendo un logit por token del vocabulario en cada posición de la secuencia. Durante el entrenamiento, calculamos la pérdida de entropía cruzada entre estos logits y los tokens objetivo desplazados (el token en la posición $i+1$ es la etiqueta para la posición $i$, porque estamos haciendo predicción del siguiente token como se describió en el artículo 8).

class DecoderTransformer(nn.Module):
    """Minimal decoder-only transformer language model."""
    def __init__(self, vocab_size, d_model, n_heads, n_layers, max_seq_len):
        super().__init__()
        self.embedding = TokenEmbedding(vocab_size, d_model, max_seq_len)
        self.blocks = nn.Sequential(
            *[TransformerBlock(d_model, n_heads, max_seq_len) for _ in range(n_layers)]
        )
        self.ln_final = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, idx, targets=None):
        x = self.embedding(idx)           # (B, T, d_model)
        x = self.blocks(x)                # (B, T, d_model)
        x = self.ln_final(x)              # (B, T, d_model)
        logits = self.head(x)             # (B, T, vocab_size)

        loss = None
        if targets is not None:
            # Flatten for cross-entropy: predictions vs next tokens
            B, T, V = logits.shape
            loss = F.cross_entropy(logits.view(B * T, V), targets.view(B * T))
        return logits, loss

Esa es toda la arquitectura. Contemos: TokenEmbedding tiene aproximadamente 10 líneas, CausalSelfAttention unas 30, FeedForward y TransformerBlock juntas unas 20, y DecoderTransformer otras 20. En aproximadamente 80 líneas de código del modelo, hemos implementado cada concepto de los artículos 1 a 8.

Antes de entrenar, es útil verificar que las formas son correctas ejecutando un pase forward con datos ficticios. El siguiente fragmento instancia un modelo pequeño y comprueba que la salida tiene las dimensiones esperadas.

# Shape verification with dummy data
vocab_size = 16
d_model = 64
n_heads = 4
n_layers = 2
max_seq_len = 32
batch_size = 2
seq_len = 10

model = DecoderTransformer(vocab_size, d_model, n_heads, n_layers, max_seq_len)
idx = torch.randint(0, vocab_size, (batch_size, seq_len))      # (2, 10)
targets = torch.randint(0, vocab_size, (batch_size, seq_len))   # (2, 10)

logits, loss = model(idx, targets)
print(f"Logits shape: {logits.shape}")   # expected: (2, 10, 16)
print(f"Loss: {loss.item():.4f}")        # expected: ~2.77 (≈ -ln(1/16))
# Logits shape: torch.Size([2, 10, 16])
# Loss: 2.8XXX (close to ln(16) ≈ 2.77, since the model is untrained)

La pérdida inicial debería estar cerca de $\ln(V)$ donde $V$ es el tamaño del vocabulario, porque un modelo sin entrenar asigna probabilidad aproximadamente uniforme a todos los tokens. Si vemos una pérdida de aproximadamente $\ln(16) \approx 2.77$ para nuestro vocabulario de 16 tokens, las formas y el cálculo de la pérdida son correctos y podemos proceder al entrenamiento.

Entrenamiento en una Tarea de Juguete

Un modelo de este tamaño (unos pocos miles de parámetros) no puede aprender lenguaje natural, pero puede aprender tareas algorítmicas simples que verifican que el mecanismo de atención funciona. Lo entrenaremos para invertir secuencias cortas: dada una entrada como [5, 3, 8, 0, SEP, ?, ?, ?, ?] , el modelo debería aprender a producir [0, 8, 3, 5] después del token separador. Esta tarea requiere que el modelo atienda a posiciones anteriores específicas (los patrones de atención deberían alinearse aproximadamente de forma inversa con la posición), lo que la convierte en una buena prueba de estrés para la atención causal.

Generaremos datos de entrenamiento al vuelo: secuencias aleatorias de longitud $L$, seguidas de un token separador, seguidas de la misma secuencia invertida. El modelo ve la concatenación completa como una sola secuencia y se entrena con el objetivo estándar de predicción del siguiente token. Solo calculamos la pérdida en la porción de salida (después del separador), ya que la porción de entrada no tiene un objetivo predecible.

# --- Toy task: learn to reverse a sequence ---
import torch

# Tokens 0..9 are data tokens, 10 = SEP, 11 = PAD (unused here)
VOCAB_SIZE = 12
SEP_TOKEN = 10
SEQ_LEN = 4       # length of sequence to reverse
TOTAL_LEN = 2 * SEQ_LEN + 1  # input + SEP + reversed output

def make_batch(batch_size):
    """Generate (input, target) pairs for the reversal task."""
    data = torch.randint(0, 10, (batch_size, SEQ_LEN))
    sep = torch.full((batch_size, 1), SEP_TOKEN)
    reversed_data = data.flip(1)
    full_seq = torch.cat([data, sep, reversed_data], dim=1)  # (B, 2*L+1)

    # Input is all tokens except the last; target is all tokens except the first
    x = full_seq[:, :-1]
    y = full_seq[:, 1:]
    return x, y

# Hyperparameters
D_MODEL = 64
N_HEADS = 4
N_LAYERS = 2
MAX_SEQ_LEN = TOTAL_LEN
LR = 3e-4
STEPS = 2000
BATCH_SIZE = 64

model = DecoderTransformer(VOCAB_SIZE, D_MODEL, N_HEADS, N_LAYERS, MAX_SEQ_LEN)
optimizer = torch.optim.AdamW(model.parameters(), lr=LR)

# Training loop
for step in range(STEPS):
    x, y = make_batch(BATCH_SIZE)
    logits, loss = model(x, y)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 400 == 0:
        print(f"Step {step:4d} | Loss: {loss.item():.4f}")

# --- Evaluate ---
model.eval()
with torch.no_grad():
    x_test, y_test = make_batch(5)
    logits, _ = model(x_test)
    preds = logits.argmax(dim=-1)  # greedy decoding

    for i in range(5):
        inp = x_test[i, :SEQ_LEN].tolist()
        expected = list(reversed(inp))
        got = preds[i, SEQ_LEN:].tolist()
        status = "PASS" if got == expected else "FAIL"
        print(f"Input: {inp} | Expected: {expected} | Got: {got} | {status}")

Con 2.000 pasos de entrenamiento y un tamaño de batch de 64, este modelo típicamente converge a una pérdida cercana a cero en la tarea de inversión en pocos minutos en CPU. La señal clave de que el entrenamiento está funcionando es la curva de pérdida: debería comenzar cerca de $\ln(12) \approx 2.48$ (uniforme sobre 12 tokens), descender rápidamente durante los primeros cientos de pasos a medida que el modelo aprende la estructura de la tarea, y aplanarse cerca de cero una vez que ha aprendido a invertir perfectamente.

Si el entrenamiento falla (la pérdida permanece alta u oscila), las causas más comunes son una tasa de aprendizaje demasiado alta (intenta bajarla a 1e-4), muy pocas capas o cabezas para que el modelo enrute la información correctamente, o un error en la máscara causal (que permitiría al modelo hacer trampa mirando tokens futuros durante el entrenamiento, y luego fallar en el momento de la prueba). Estos modos de fallo son instructivos en sí mismos, porque se corresponden directamente con los conceptos que hemos ido construyendo. La máscara causal impone la propiedad autoregresiva del artículo 3, las múltiples cabezas habilitan los patrones de atención paralelos del artículo 4, y las conexiones residuales del artículo 6 hacen posible que los gradientes fluyan a través de múltiples bloques apilados.

Ahora tenemos un transformer funcional, construido desde cero en aproximadamente 130 líneas. Cada clase corresponde a un concepto de esta serie: TokenEmbedding es el artículo 5, CausalSelfAttention son los artículos 2 a 4, TransformerBlock son los artículos 6 y 7, y DecoderTransformer une todo como en el artículo 8. El siguiente artículo pasa de la arquitectura al entrenamiento: ¿cómo pasamos de este pequeño modelo de juguete a un modelo de lenguaje grande que realmente entiende el lenguaje? La respuesta involucra pre-training a escala y fine-tuning con instrucciones.

Quiz

Pon a prueba tu comprensión de la implementación del transformer.