Dentro del Transformer

El artículo de 2017 "Attention Is All You Need" de Vaswani et al. es probablemente el paper de machine learning más trascendental de la última década. Introdujo el Transformer, que descartó los modelos de secuencia recurrentes y convolucionales en favor de un mecanismo construido enteramente sobre la atención. Cada modelo de lenguaje importante desde entonces —GPT, BERT, PaLM, LLaMA, Claude— se asienta sobre esta base.

Si trabajas en IA, entender el Transformer no es opcional. Así que lo explicaré pieza por pieza.

El problema de la recurrencia

Antes de los Transformers, el modelado de secuencias se basaba en RNNs y LSTMs. Estas procesan los tokens uno a la vez, manteniendo un estado oculto (hidden state) que acumula información de izquierda a derecha. Ese diseño secuencial presenta dos problemas reales:

El entrenamiento no se puede paralelizar. Cada paso de tiempo depende del anterior, por lo que no puedes distribuir el trabajo entre los núcleos de la GPU de ninguna manera útil.
Las dependencias de largo alcance decaen. Incluso las LSTMs, diseñadas específicamente para mantener viva la información en secuencias largas, tienen dificultades cuando el contexto relevante está cientos de tokens atrás. El estado oculto es un cuello de botella.

El Transformer elimina la recurrencia por completo. Cada token puede atender a cualquier otro directamente y en paralelo.

La arquitectura a simple vista

El Transformer original es de tipo encoder-decoder. El encoder procesa la secuencia de entrada y produce un conjunto de representaciones. El decoder genera la salida token por token, atendiendo tanto a la salida del encoder como a sus propios tokens generados anteriormente.

Para modelos de lenguaje como GPT, solo se utiliza el decoder. Para modelos tipo BERT, solo el encoder. El mecanismo central —self-attention— es el mismo en ambos.

Self-attention: el mecanismo central

El self-attention permite que cada token en una secuencia observe a todos los demás tokens y decida cuánta "atención" prestarles. Así es como funciona.

Dada una secuencia de entrada de $n$ tokens, cada uno representado como un vector de $d$ dimensiones, computamos tres matrices:

Q (Queries): ¿qué está buscando este token?
K (Keys): ¿qué ofrece este token?
V (Values): ¿qué información contiene este token?

Cada una se obtiene multiplicando la entrada por una matriz de pesos aprendida:

Q = XW_Q \qquad K = XW_K \qquad V = XW_V

con $Q, K \in \mathbb{{R}}^{{n \times d_k}}$ y $V \in \mathbb{{R}}^{{n \times d_v}}$ . La atención se calcula como:

\mathrm{{Attention}}(Q, K, V) = \mathrm{{softmax}}\!\left(\frac{{QK^\top}}{{\sqrt{{d_k}}}}\right)V

Analizando esto:

$QK^\top$ produce una matriz $n \times n$ donde la entrada $(i, j)$ es el producto punto entre la query $i$ y la key $j$. Esto mide cuánto debe atender el token $i$ al token $j$.
El $1/\sqrt{{d_k}}$ es un factor de escala. Sin él, los productos punto crecen a medida que $d_k$ aumenta, empujando a la función softmax hacia regiones con gradientes que se desvanecen. La raíz cuadrada mantiene la varianza estable.
softmax normaliza cada fila para que sume 1, de modo que las puntuaciones brutas se conviertan en una distribución de probabilidad sobre los tokens.
La multiplicación final por $V$ genera la salida: la representación de cada token es una suma ponderada de todos los vectores de valor, ponderada por la atención.

En código:

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

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'))

    weights = F.softmax(scores, dim=-1)
    return torch.matmul(weights, V), weights

El parámetro mask aparece en el decoder, donde no quieres que los tokens atiendan a posiciones futuras. Eso es el causal masking: el token $i$ solo puede ver los tokens del $0$ al $i$.

Multi-head attention

Un solo head de atención solo puede capturar un tipo de relación a la vez; por ejemplo, la "concordancia sujeto-verbo" o la "resolución de correferencia". Para capturar varias a la vez, el Transformer utiliza multi-head attention.

En lugar de una sola operación de atención con Q, K, V de $d$ dimensiones, se dividen en $h$ heads, cada uno con $d/h$ dimensiones. Cada head calcula la atención de forma independiente, y las salidas se concatenan y proyectan:

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_Q = torch.nn.Linear(d_model, d_model)
        self.W_K = torch.nn.Linear(d_model, d_model)
        self.W_V = torch.nn.Linear(d_model, d_model)
        self.W_O = torch.nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.size()

        # Project and reshape into (batch, heads, seq_len, d_k)
        Q = self.W_Q(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_K(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_V(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # Attention per head
        attn_output, _ = scaled_dot_product_attention(Q, K, V, mask)

        # Concatenate heads and project
        attn_output = attn_output.transpose(1, 2).contiguous().view(
            batch_size, seq_len, -1
        )
        return self.W_O(attn_output)

En la práctica, diferentes heads aprenden a atender a diferentes cosas. Los trabajos de interpretabilidad han demostrado que algunos heads se especializan en relaciones sintácticas (atendiendo al token anterior o al inicio de la oración), mientras que otros captan relaciones semánticas (atendiendo al sujeto de una cláusula relativa, por ejemplo).

Positional encoding

El self-attention es invariante a las permutaciones. Si barajas los tokens de entrada, los puntajes de atención cambian, pero el mecanismo en sí no tiene una noción integrada de que la posición importa. Dado que el orden de las palabras es fundamental para el lenguaje, el Transformer añade positional encodings a los embeddings de entrada.

El paper original utiliza funciones sinusoidales:

PE_{{(pos,\,2i)}} = \sin\!\left(\frac{{pos}}{{10000^{{2i/d_{{model}}}}}}\right) \qquad PE_{{(pos,\,2i+1)}} = \cos\!\left(\frac{{pos}}{{10000^{{2i/d_{{model}}}}}}\right)

Cada posición recibe un patrón único de senos y cosenos a diferentes frecuencias. La intuición es que cada dimensión oscila a un ritmo diferente, por lo que cada posición obtiene una "huella dactilar" distintiva. Una propiedad útil es que el encoding para la posición $p+k$ puede escribirse como una transformación lineal del encoding para la posición $p$, lo que puede ayudar al modelo a aprender posiciones relativas.

Los modelos modernos han evolucionado. BERT y GPT-2 utilizan embeddings de posición aprendidos. LLaMA y muchos LLMs actuales utilizan rotary position embeddings (RoPE), que codifican la posición rotando los vectores de query y key. RoPE decae naturalmente la atención hacia tokens distantes y generaliza mejor a longitudes de secuencia que el modelo no vio durante el entrenamiento.

El bloque Transformer completo

Un solo bloque de Transformer tiene:

Multi-head self-attention
Conexión residual + layer normalization
Feed-forward network (dos capas lineales con una no linealidad)
Conexión residual + layer normalization

\text{{h}} = \mathrm{{LayerNorm}}\big(x + \mathrm{{MultiHeadAttention}}(x)\big)

\text{{output}} = \mathrm{{LayerNorm}}\big(\text{{h}} + \mathrm{{FFN}}(\text{{h}})\big)

La red feed-forward (FFN) se ejecuta de forma independiente en cada posición. Típicamente expande la dimensión 4 veces, aplica GELU o ReLU, y luego proyecta de vuelta. Investigaciones recientes sugieren que las capas FFN actúan como memorias de clave-valor (key-value memories), almacenando conocimiento fáctico aprendido durante el entrenamiento; razón por la cual las técnicas de edición de modelos suelen dirigirse a ellas.

Las conexiones residuales no son opcionales. Sin ellas, los gradientes se desvanecen en modelos profundos. El flujo residual actúa como una autopista que lleva información de las capas iniciales a las posteriores, y cada capa añade su propia contribución.

Encoder-decoder vs decoder-only

El Transformer original tiene ambos bloques. El encoder procesa la entrada completa con atención bidireccional: cada token ve a todos los demás. El decoder genera la salida de forma autorregresiva con causal masking, además de capas de cross-attention que atienden a la salida del encoder.

El campo ha convergido mayoritariamente en modelos decoder-only para la generación. GPT, LLaMA y la mayoría de los LLMs modernos son decoder-only con causal masking. BERT y sus variantes son encoder-only con atención bidireccional, lo que los hace fuertes para clasificación y extracción, pero inutilizables para generación. La estructura encoder-decoder todavía aparece en T5 y en traducción automática, donde la entrada y la salida son secuencias claramente distintas.

Por qué ganaron los Transformers

El dominio del Transformer se debe a tres propiedades.

Paralelismo. Cada token se procesa al mismo tiempo durante el entrenamiento. En el hardware de GPU moderno, eso se traduce en aceleraciones masivas sobre los modelos recurrentes. Entrenar GPT-3 con LSTMs habría sido computacionalmente inviable.

Contexto flexible. El self-attention le da a cada token un camino directo a cualquier otro token, independientemente de la distancia. Una LSTM necesita que la información sobreviva al paso por cada estado oculto intermedio. Un Transformer solo necesita que el peso de atención sea distinto de cero.

Escalabilidad. Los Transformers escalan de manera predecible con el cómputo, los datos y los parámetros. Las leyes de escalado (scaling laws de Kaplan et al., Hoffmann et al.) se descubrieron específicamente para esta arquitectura. Más parámetros, más datos, más cómputo, y obtienes de manera confiable un mejor modelo.

La gran limitación es el costo cuadrático del self-attention: $O(n^2)$ en la longitud de la secuencia. Una secuencia de 100,000 tokens tiene una matriz de atención con 10,000 millones de entradas. Esto ha impulsado mucha investigación en variantes eficientes: sparse attention, linear attention, FlashAttention (que es en realidad una optimización de acceso a memoria) y enfoques de ventana deslizante (sliding window). Pero el mecanismo central sigue siendo el scaled dot-product attention del paper original.

Implicaciones prácticas

Si estás construyendo sistemas sobre LLMs, entender el Transformer te ayuda a razonar sobre cómo se comportan. Las dependencias de largo alcance deberían funcionar en teoría, pero en la práctica la atención se diluye en contextos muy largos. El conocimiento del modelo reside en las capas FFN y puede ser localizado, que es como funciona la edición de modelos. Diferentes attention heads capturan diferentes fenómenos lingüísticos, lo cual puedes explotar para la interpretabilidad.

El Transformer no es solo una arquitectura. Es el sustrato computacional sobre el que corre la era actual de la IA.