Las matemáticas de la atención paso a paso: de la matriz de entrada X al softmax y V

Cada vez que un modelo de lenguaje produce un token, se ejecuta una y otra vez el mismo cálculo: la atención. No es una forma de hablar. En cada pasada hacia delante, para cada token y en cada capa, se ejecuta la misma aritmética.

La mayoría de las explicaciones se quedan en la intuición de «consultas, claves y valores», y eso deja las matemáticas reales como una caja negra. Este artículo abre la caja. Explica qué es de verdad la matriz de entrada, de dónde salen Q, K y V, un ejemplo numérico lo bastante pequeño como para ejecutarlo tú mismo y por qué se llama self-attention. Después debería ser posible seguir a mano una sola cabeza de atención.

El problema que resuelve la atención

Toma esta frase: «The animal didn't cross the street because it was tired» («El animal no cruzó la calle porque estaba cansado»).

Cuando el modelo llega al token «it», tiene que averiguar a qué se refiere la palabra. ¿Al animal? ¿A la calle? La palabra sola no dice nada, así que la respuesta tiene que salir de la secuencia que la rodea.

La atención lo resuelve dejando que cada token asigne a cada uno de los demás un peso, un número que indica cuánto debe fijarse en esa posición mientras construye su representación actualizada. Para «it», los pesos podrían salir así:

"it" → "animal"  : 0.65
"it" → "street"  : 0.20
"it" → "tired"   : 0.10
"it" → others    : 0.05

Nadie escribe esos números a mano. Salen de un cálculo aprendido construido sobre tres vectores, Query, Key y Value, y la fórmula entera cabe en una línea:

Attention(Q, K, V) = softmax( QKᵀ / √dₖ ) · V

Todo lo que sigue no hace más que desmenuzar esa línea.

Paso 0: la matriz de entrada X

Antes de que puedan existir Q, K o V, tiene que haber una matriz de entrada, X.

Cada token de «The cat sat» se asigna a un vector de embedding aprendido. Con un tamaño de embedding de d_model = 4, queda así:

The → [0.2, 0.5, 0.1, 0.7]
cat → [0.8, 0.3, 0.6, 0.2]
sat → [0.4, 0.9, 0.5, 0.1]

Apila los vectores fila a fila y obtienes X, cuya forma es [n_tokens, d_model].

Hay un detalle fácil de pasar por alto: X no es constante entre capas. Cada capa de atención toma la salida de la capa anterior como su propia X. En la primera capa, X es solo el token en bruto más la información de posición. Más arriba en la pila, por ejemplo en la capa 12 de GPT-3, que tiene 96 capas, X ya codifica un significado contextual profundo construido a lo largo de muchas rondas de atención.

Paso 1: calcular Q, K y V

Nadie asigna consultas, claves y valores a mano. Los tres se calculan a partir de X mediante tres matrices de pesos aprendidas por separado:

Q = X · WQ
K = X · WK
V = X · WV

WQ, WK y WV son matrices de parámetros que quedan congeladas al terminar el entrenamiento. Las formas encajan así:

X.shape  = [n, d_model]
WQ.shape = [d_model, d_k]   → Q.shape = [n, d_k]
WK.shape = [d_model, d_k]   → K.shape = [n, d_k]
WV.shape = [d_model, d_v]   → V.shape = [n, d_v]

La distinción importa, así que aquí va en una tabla:

Símbolo	Qué es	¿Depende de la entrada?
WQ, WK, WV	Matrices de proyección aprendidas	No, fijas tras el entrenamiento
Q, K, V	Activaciones calculadas	Sí, siguen a la X actual

Por eso «bank» en «river bank» (la orilla del río) y «bank» en «deposit at the bank» (ingresar en el banco) acaban con patrones de atención totalmente distintos. El contexto produce una X diferente, que pasa por las mismas matrices de pesos y sale convertida en vectores Q, K y V distintos.

Un buscador sirve como modelo mental práctico:

• Q es lo que busca el token.
• K es lo que cada token anuncia de sí mismo.
• V es la información que se devuelve cuando hay coincidencia.

Q y K deciden adónde mirar. V decide qué sale realmente.

Paso 2: el producto escalar QKᵀ

Toma Q y multiplícala por K traspuesta. Eso da una matriz cuadrada [n, n] de puntuaciones en bruto, donde la entrada de la fila i, columna j, indica con qué fuerza atiende el token i al token j. Antes de rellenar números, tiene este aspecto:

Keys
       cat   eats  fish
     ┌──────┬──────┬──────┐
Q cat│  ?   │  ?   │  ?   │
  eat│  ?   │  ?   │  ?   │
 fish│  ?   │  ?   │  ?   │
     └──────┴──────┴──────┘

Un producto escalar mayor significa que los dos tokens son más relevantes el uno para el otro.

Paso 3: escalar por √dₖ

Cuando d_k es grande, los productos escalares en bruto se disparan. Con d_k = 512, puntuaciones como [180, 210, 95] son normales. Pásalas por el softmax y la salida es [0.00, 1.00, 0.00], una distribución completamente colapsada. Entonces los gradientes se desvanecen y el entrenamiento se estanca.

Dividir entre √d_k mantiene las puntuaciones en un rango razonable. Esa división es la que le da el nombre de scaled dot-product attention.

Paso 4: el softmax convierte puntuaciones en pesos

El softmax se aplica fila a fila, y cada fila se convierte en una distribución de probabilidad sobre todos los tokens. Este es el cálculo completo para las puntuaciones [3.2, 4.0, 1.0, 6.2, 8.0]:

exp(3.2) =     24.53
exp(4.0) =     54.60
exp(1.0) =      2.72
exp(6.2) =    492.75
exp(8.0) =  2,980.96
         ──────────
Σ        =  3,555.56

Token 1:   24.53 / 3555.56 →  0.69%
Token 2:   54.60 / 3555.56 →  1.54%
Token 3:    2.72 / 3555.56 →  0.08%
Token 4:  492.75 / 3555.56 → 13.86%
Token 5: 2980.96 / 3555.56 → 83.83%  ← winner

El token 5 se lleva el 83,83 % aunque solo supera al token 4 por 1,8 puntos. Esa amplificación no lineal es lo que permite al modelo comprometerse con firmeza con los tokens más relevantes.

Consejo de estabilidad: resta primero el máximo de la fila, como en np.exp(x - x.max()), y solo después divide. Un float de 64 bits no puede representar exp(800), pero tras el desplazamiento todos los exponentes quedan en zona segura y las probabilidades resultantes son matemáticamente idénticas.

Paso 5: la suma ponderada con V

El último paso multiplica los pesos por V y lo suma todo. Para «it», eso significa:

output("it") =
    0.65 × V("animal")
  + 0.20 × V("street")
  + 0.10 × V("tired")
  + 0.05 × V(rest)

Lo que sale es un vector nuevo para «it» que contiene sobre todo información de «animal», un poco de «street» y muy poco del resto. Esta representación contextualizada es lo que pasa a la siguiente capa.

Un ejemplo completo en NumPy

La frase aquí es «cat eats fish» con d_model = 2:

import numpy as np

X  = np.array([[1, 2], [3, 1], [2, 4]], dtype=float)
WQ = np.array([[1, 0], [0, 1]], dtype=float)  # identity
WK = np.array([[1, 0], [0, 1]], dtype=float)  # identity
WV = np.array([[0, 1], [1, 0]], dtype=float)  # swaps dimensions

Q, K, V = X @ WQ, X @ WK, X @ WV

# Raw scores for "fish" (Q=[2,4]) against each key:
# fish→cat:  [2,4]·[1,2] = 10
# fish→eats: [2,4]·[3,1] = 10
# fish→fish: [2,4]·[2,4] = 20  ← self-attention dominates

scores = Q @ K.T / np.sqrt(2)
e = np.exp(scores - scores.max(-1, keepdims=True))
weights = e / e.sum(-1, keepdims=True)  # softmax row-wise

output = weights @ V
# output[fish] ≈ [3.997, 2.001] ≈ V(fish)
# fish self-attends at ~99.8% in this toy example

En un modelo entrenado de verdad, WQ y WK romperían esta simetría y darían una atención con sentido entre tokens. Las matrices identidad solo sirven para que la aritmética siga siendo legible.

Por qué se llama self-attention

Q, K y V salen todas de la misma matriz de origen, X. Cada token produce una consulta, una clave y un valor a partir de su propia representación, y luego cada token atiende a todos los demás dentro de esa misma secuencia.

Eso distingue la self-attention, la que usan GPT y LLaMA, de la cross-attention, la de los modelos codificador-decodificador como T5, donde Q sale de una secuencia y K y V de otra.

Notas prácticas de ingeniería

Caché KV. Recalcular K y V para todos los tokens anteriores en cada paso de generación costaría O(n²). Una caché KV guarda todos los vectores K y V previos y los reutiliza, lo que reduce el coste por paso a O(n) a cambio de una memoria que crece con la longitud de la secuencia.

Dimensión de la cabeza. La configuración habitual es d_k = d_model / num_heads. GPT-2 tiene 768 / 12 = 64 y GPT-3 tiene 12288 / 96 = 128.

Grouped Query Attention (GQA). LLaMA-2, Mistral y Gemma permiten que varias cabezas Query compartan una sola cabeza KV. Eso recorta mucho la memoria de la caché KV y es una de las grandes razones por las que estos modelos aguantan contextos largos en hardware de consumo.

La atención en una frase

No hay magia. La atención es una tabla de búsqueda aprendida, diferenciable y con pesos suaves que se recalcula desde cero para cada entrada y produce representaciones que llevan toda la estructura relacional del texto.