Construye un Transformer desde cero en PyTorch: atención, codificador y decodificador

Los Transformers demostraron que el modelado de secuencias puede construirse por completo sobre la atención, sin recurrencia y sin convoluciones. Los autores del artículo original del Transformer propusieron un diseño codificador-decodificador basado en self-attention y mostraron que se entrena muy bien en paralelo y rinde con fuerza en traducción automática.

Cuando los Transformers solo aparecen a través de bibliotecas de alto nivel, es fácil tratarlos como magia. Escribir uno mismo es de las mejores formas de ver qué ocurre realmente dentro. Lo que sigue es una versión compacta en PyTorch de las partes centrales: empieza por cómo se codifican las posiciones, pasa por los dos tipos de atención y el bloque feed-forward, sigue con las capas del codificador y termina con un esbozo del decodificador. Buenos compañeros de lectura son el propio artículo, The Annotated Transformer de Harvard NLP y la referencia de torch.nn.

Por qué escribirlo a mano

El artículo es claro, pero convertirlo en código que funcione implica muchas decisiones pequeñas: formas de los tensores, enmascarado, normalización, conexiones residuales y cómo encajan los módulos entre sí. The Annotated Transformer se hizo famoso porque convierte el artículo en una implementación didáctica línea a línea, lo que facilita mucho razonar sobre el diseño.

Construirlo tú mismo también hace tangibles tres ideas definitorias:

1. La atención asume el trabajo que antes hacía la recurrencia, que es relacionar los tokens entre sí.
2. La posición tiene que inyectarse de forma explícita, porque la atención por sí sola no tiene noción del orden de las palabras.
3. Los bloques del Transformer son modulares, algo que encaja muy bien con torch.nn.Module, ModuleList, las capas lineales, las capas de normalización y el dropout.

Una imagen mental sencilla

El modelo clásico tiene dos pilas. La pila del codificador aplica self-attention y una red feed-forward en cada capa. La pila del decodificador hace lo mismo, además de self-attention enmascarada y atención sobre la salida del codificador.

Cada subcapa va envuelta en una conexión residual y layer normalization. La normalización actúa sobre las últimas dimensiones de características y puede usar parámetros aprendibles de escala y desplazamiento. Visto así, un bloque es menos misterioso de lo que suena. Son sobre todo proyecciones lineales, atención, normalización y caminos residuales.

Paso 1: embeddings y codificación posicional

Los Transformers trabajan con vectores, así que los IDs de tokens se convierten primero en embeddings densos. Después se suma una codificación posicional para que el modelo distinga «cat chased dog» de «dog chased cat». El artículo usó codificaciones sinusoidales para ello.

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

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # shape: (1, max_len, d_model)

        self.register_buffer("pe", pe)

    def forward(self, x):
        seq_len = x.size(1)
        return x + self.pe[:, :seq_len]

La etapa de entrada combina un embedding con un factor de escala:

class InputEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.scale = math.sqrt(d_model)

    def forward(self, x):
        return

Multiplicar por sqrt(d_model) sigue el diseño original y mantiene la magnitud de los embeddings en proporción con el resto de la red.

Paso 2: atención de producto escalar escalado

Es el núcleo del modelo. Dadas las consultas Q, las claves K y los valores V, la atención multiplica Q por la traspuesta de K, divide por la raíz cuadrada de la dimensión de las claves, aplica softmax y usa los pesos resultantes para mezclar V.

La división por sqrt(d_k) importa. Los productos escalares crecen con la dimensión y, sin escalado, el softmax se volvería demasiado puntiagudo.

def scaled_dot_product_attention(q, k, v, mask=None):
    d_k = q.size(-1)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))

    attn_weights = F.softmax(scores, dim=-1)
    output = torch.matmul(attn_weights, v)
    return output, attn_weights

Esta pequeña función permite que cada token decida dinámicamente qué otros tokens merecen atención. En self-attention, las consultas, las claves y los valores salen de la misma secuencia. En la atención codificador-decodificador, el decodificador lanza sus consultas sobre las salidas del codificador.

Paso 3: atención multi-head

Un gran patrón de atención se sustituye por varios más pequeños. Las entradas se proyectan en varios subespacios, la atención se ejecuta en paralelo en cada uno y los resultados se concatenan y se proyectan de vuelta a la dimensión del modelo.

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)

    def split_heads(self, x):
        batch_size, seq_len, d_model = x.size()
        x = x.view(batch_size, seq_len, self.num_heads, self.d_k)
        return x.transpose(1, 2)  # (batch, heads, seq, d_k)

    def combine_heads(self, x):
        batch_size, heads, seq_len, d_k = x.size()
        x = x.transpose(1, 2).contiguous()
        return x.view(batch_size, seq_len, heads * d_k)

    def forward(self, q, k, v, mask=None):
        q = self.split_heads(self.w_q(q))
        k = self.split_heads(self.w_k(k))
        v = self.split_heads(self.w_v(v))

        attn_output, attn_weights = scaled_dot_product_attention(q, k, v, mask)
        output = self.combine_heads(attn_output)
        return self.w_o(output), attn_weights

El sistema de módulos de PyTorch encaja bien aquí, porque cada proyección aprendida es simplemente una capa lineal dentro de un nn.Module.

Paso 4: red feed-forward y layer norm

Tras la atención, cada posición pasa por la misma red feed-forward. El diseño original usa dos capas lineales con una no linealidad entre ellas. La salida se une al camino residual y se normaliza.

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff, 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)

LayerNorm de PyTorch normaliza sobre las dimensiones finales de características y puede aprender valores de escala y sesgo por elemento, que es como se usa la normalización en las pilas de un Transformer. Con esas piezas listas, una capa del codificador se monta fácilmente:

class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = FeedForward(d_model, d_ff, dropout)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        attn_out, _ = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_out))

        ffn_out = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_out))
        return x

Paso 5: la pila del codificador

El modelo original apila varias capas de codificador idénticas. En PyTorch eso suele hacerse con nn.ModuleList, que guarda submódulos y además los registra para la optimización y los checkpoints.

class Encoder(nn.Module):
    def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, dropout=0.1):
        super().__init__()
        self.embed = InputEmbedding(vocab_size, d_model)
        self.pos = PositionalEncoding(d_model)
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        self.dropout = nn.Dropout(dropout)

    def forward(self, src, mask=None):
        x = self.embed(src)
        x = self.pos(x)
        x = self.dropout(x)

        for layer in self.layers:
            x = layer(x, mask)

        return x

En este punto ya hay un codificador utilizable para tareas como clasificación o aprendizaje de representaciones de secuencias. El modelo codificador-decodificador completo necesita después una pila de decodificador y una proyección final sobre el vocabulario, una disposición que describe el artículo y que replica la implementación de Harvard.

Paso 6: el decodificador

Una capa de decodificador completa introduce dos ideas adicionales:

1. Self-attention enmascarada, que impide que un token mire los tokens posteriores durante el entrenamiento.
2. Cross-attention, en la que el decodificador atiende a las salidas del codificador.

En pseudocódigo aproximado, una capa de decodificador funciona así:

# Pseudocode structure
x = norm1(x + masked_self_attention(x, x, x))
x = norm2(x + cross_attention(x, encoder_output, encoder_output))
x = norm3(x + feed_forward(x))

Una vez que el codificador tiene sentido, el decodificador es casi lo mismo otra vez, con un bloque de atención más y una máscara causal.

Comprobar las formas sobre la marcha

Un hábito que ahorra horas es pasar un lote aleatorio por cada pieza en cuanto está escrita. Mete un tensor de forma (batch, sequence, model dimension) en la codificación posicional y la forma no debería cambiar. Pásalo por la atención multi-head y la salida también debería conservar esa forma, mientras que los pesos de atención vuelven como (batch, heads, sequence, sequence). El bloque feed-forward ensancha internamente la última dimensión y la devuelve al tamaño del modelo, así que la forma de salida vuelve a coincidir con la de entrada.

Si alguna de esas comprobaciones falla, el problema casi siempre está en cómo se dividen o se juntan las cabezas. Detectarlo en esa fase es mucho más barato que descubrirlo tras un entrenamiento largo que produce basura en silencio.

Otra prueba útil es un modelo deliberadamente diminuto. Dos capas, una dimensión de modelo pequeña y un conjunto de datos de juguete, como copiar una secuencia corta, bastan para ver si la pérdida baja de verdad. Un modelo que no aprende a copiar tiene un bug en algún sitio, y uno que sí puede ya está haciendo bien la parte difícil.

Bugs frecuentes

Cuando un Transformer escrito desde cero se niega a entrenar, la causa suele ser una de estas:

• formas de tensores que se estropean al dividir o juntar las cabezas
• errores de broadcasting en las máscaras, sobre todo en la self-attention del decodificador
• codificación posicional olvidada, lo que deja a la atención ciega al orden
• normalización colocada en un sitio distinto al de la variante con la que se compara
• capas guardadas en una lista normal de Python en lugar de una ModuleList, de modo que nunca se registran

La arquitectura es elegante. Los problemas suelen esconderse en la fontanería.

Qué llevarse

Ningún componente es abrumadoramente complicado. La fuerza viene de la combinación, que sigue siendo sencilla, se paraleliza bien y expresa mucho. Los ingredientes son embeddings aprendidos, señales de posición, self-attention, capas feed-forward por posición, atajos residuales y normalización. Esa receta se convirtió en la base de los modelos de secuencias modernos.

Para entenderlo de verdad, ve más allá de los diagramas. Construye uno, imprime las formas de los tensores, observa cómo se desplazan las puntuaciones de atención y rompe las máscaras a propósito antes de arreglarlas. Después de eso, el modelo deja de parecer magia y empieza a parecer ingeniería.