Scaled dot-product attention: Q, K, V y el misterioso √d_k

La ecuación que vale un paper

Todo el Transformer gira alrededor de una sola línea:

Vamos a desarmarla con calma, porque cada símbolo tiene una razón de ser.

Queries, keys y values: una búsqueda difusa

La metáfora más útil es la de una base de datos. En un diccionario de Python haces d[query] y obtienes el valor cuya clave coincide exactamente. La atención es una versión suave y diferenciable de esa búsqueda:

• Cada posición emite una query ($q$): "¿qué estoy buscando?"
• Cada posición expone una key ($k$): "esto es lo que ofrezco para que me encuentren".
• Cada posición guarda un value ($v$): "esto es lo que entrego si me eligen".

En vez de elegir una única clave, comparas la query contra todas las keys, obtienes un puntaje de similitud para cada una, conviertes esos puntajes en una distribución de probabilidad y devuelves el promedio ponderado de los values.

En self-attention, $Q$, $K$ y $V$ salen de la misma secuencia de entrada $X \in \mathbb{R}^{n \times d_{\text{model}}}$ mediante tres proyecciones lineales aprendidas:

con $W^Q, W^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ y $W^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$. Que sean proyecciones distintas importa: lo que una palabra busca no tiene por qué parecerse a lo que ofrece. Un verbo puede buscar su sujeto (query) mientras ofrece información de tiempo verbal (value).

Paso a paso con las formas de los tensores

Supongamos una secuencia de $n$ tokens.

1. Puntajes: $S = QK^\top$, con forma $n \times n$. La entrada $S_{ij} = q_i \cdot k_j$ mide cuánto le interesa la posición $j$ a la posición $i$.
2. Escalado: $S / \sqrt{d_k}$. Ya veremos por qué.
3. Softmax por fila: cada fila se convierte en una distribución que suma 1. La fila $i$ son los pesos de atención de la posición $i$ sobre todas las demás.
4. Mezcla: multiplicas por $V$ ($n \times d_v$) y obtienes una salida $n \times d_v$: cada posición es ahora una combinación convexa de los values.

Todo son multiplicaciones de matrices. No hay bucles sobre posiciones, no hay dependencias secuenciales: la GPU calcula las $n^2$ interacciones de una vez.

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

def scaled_dot_product_attention(q, k, v, mask=None):
    # q: (..., n_q, d_k)   k: (..., n_k, d_k)   v: (..., n_k, d_v)
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / math.sqrt(d_k)   # (..., n_q, n_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))
    weights = F.softmax(scores, dim=-1)                  # cada fila suma 1
    return weights @ v, weights                          # (..., n_q, d_v)

Softmax y temperatura

Antes de hablar del escalado, conviene tener intuición sobre softmax. Dado un vector de puntajes $z$:

convierte una lista de puntajes en porcentajes que suman 1: cada $e^{z_i}$ dividido por la suma de todos. El más alto se lleva la mayor parte

Softmax es sensible a la escala de sus entradas. Si multiplicas todos los puntajes por una constante grande, la distribución se vuelve casi one-hot: el máximo se lleva todo. Si los multiplicas por una constante chica, la distribución se aplana hacia uniforme. Juega con la temperatura en el widget y mira qué pasa con los pesos y con su entropía:

Dividir por $\sqrt{d_k}$ es exactamente ajustar esa temperatura. La pregunta es por qué ese valor.

El argumento de varianza

Supón que las componentes de $q$ y $k$ son variables aleatorias independientes con media 0 y varianza 1 (algo razonable con una buena inicialización y normalización). El producto punto es:

el producto punto: multiplicar componente a componente y sumar los $d_k$ resultados

Cada término $q_i k_i$ tiene media $\mathbb{E}[q_i]\mathbb{E}[k_i] = 0$ y varianza $\mathbb{E}[q_i^2]\mathbb{E}[k_i^2] = 1 \cdot 1 = 1$. Como los $d_k$ términos son independientes, las varianzas se suman:

en promedio da cero, pero su dispersión (la varianza) crece igual que $d_k$: con más dimensiones, puntajes más extremos

La desviación estándar del producto punto crece como $\sqrt{d_k}$. Con $d_k = 64$ (el valor del paper), los puntajes típicos andan en el orden de $\pm 8$; con dimensiones más grandes, peor. Puntajes así de grandes empujan al softmax a su régimen saturado: un elemento cerca de 1, el resto cerca de 0.

¿Y por qué es malo? Porque en la saturación los gradientes del softmax son minúsculos. La derivada $\partial\, \text{softmax}_i / \partial z_j = \text{softmax}_i(\delta_{ij} - \text{softmax}_j)$ se anula cuando todas las probabilidades son 0 o 1. El modelo deja de aprender a redistribuir su atención.

Dividiendo por $\sqrt{d_k}$:

dividir por $\sqrt{d_k}$ devuelve la dispersión a 1, sea cual sea la dimensión

Los puntajes vuelven a tener varianza unitaria, independientemente de la dimensión, y el softmax opera en su zona sana.

Por qué producto punto y no atención aditiva

El paper discute explícitamente la alternativa: la atención aditiva de Bahdanau, que usa una pequeña red feed-forward para calcular la compatibilidad. Ambas tienen complejidad teórica similar, pero el producto punto es mucho más rápido y eficiente en memoria en la práctica, porque se reduce a multiplicaciones de matrices altamente optimizadas.

Un detalle interesante: el paper observa que para $d_k$ chico ambas funcionan parecido, pero sin escalado la aditiva supera al producto punto con $d_k$ grande. El escalado es lo que permite quedarse con la opción rápida sin perder calidad.

Juega con la atención

En el widget de abajo puedes escribir una oración y ver la matriz de atención resultante. Cada fila es una query; cada columna, una key. El color indica el peso tras el softmax. Prueba cambiar la oración y fíjate cómo las filas siempre suman 1:

Una sutileza: la matriz no es simétrica

Como $W^Q \neq W^K$, en general $q_i \cdot k_j \neq q_j \cdot k_i$. Que "el" atienda mucho a "gato" no implica que "gato" atienda mucho a "el". Esta asimetría es deseable: las relaciones lingüísticas (modificador → modificado, pronombre → antecedente) son direccionales.

Resumen

• $Q$, $K$, $V$ son proyecciones lineales de la entrada que implementan una búsqueda suave.
• $QK^\top$ calcula todas las similitudes en paralelo; softmax las convierte en pesos.
• El factor $1/\sqrt{d_k}$ normaliza la varianza de los puntajes y evita la saturación del softmax.
• El costo es $O(n^2 d)$: cuadrático en la longitud de la secuencia.

En la próxima lección vamos a ver por qué una sola atención no alcanza y cómo el paper la multiplica en varias cabezas.