Transformers: Atención: cada palabra mira a todas las demás

Lee esta frase: «La gata que perseguía al ratón estaba muy cansada». Para saber que la última palabra va en femenino hay que recordar «gata», siete palabras atrás, y no confundirla con «ratón», que está más cerca. Un modelo de $n$ -gramas que solo mira las dos o tres palabras anteriores no puede hacerlo. Capturar dependencias lejanas fue durante décadas el gran problema del modelado del lenguaje.

Antes de la atención: leer de izquierda a derecha

Las redes recurrentes (Elman, 1990) leen el texto palabra a palabra y van actualizando un vector de estado, una especie de memoria resumida: $h_{t} = σ (W h_{t - 1} + U x_{t})$ . En la práctica olvidaban rápido, por el gradiente que se desvanece a lo largo de la cadena. Las LSTM de Sepp Hochreiter y Jürgen Schmidhuber (1997) añadieron puertas que deciden qué guardar y qué olvidar, y dominaron la traducción automática y el reconocimiento de voz durante veinte años.

Pero seguían siendo secuenciales. Para traducir, un codificador comprimía toda la frase de entrada en un único vector, y el decodificador generaba la traducción a partir de él (Sutskever, Vinyals y Le, 2014). Una frase larga no cabe bien en un vector. En 2014, Dzmitry Bahdanau, Kyunghyun Cho y Yoshua Bengio propusieron dejar que el decodificador, en cada paso, mirara todas las palabras de la entrada y eligiera en cuáles fijarse. Lo llamaron atención.

Atención: una búsqueda suave en un diccionario

Piensa en un diccionario de Python: buscas una clave y obtienes su valor. La atención es una versión «suave» y diferenciable de esa búsqueda. Cada palabra (cada token) produce tres vectores mediante tres matrices aprendidas, a partir de su embedding $𝐱$ :

una consulta $𝐪 = W_{Q} 𝐱$ : «¿qué estoy buscando?»;
una clave $𝐤 = W_{K} 𝐱$ : «¿qué ofrezco?»;
un valor $𝐯 = W_{V} 𝐱$ : «esta es mi información si me eliges».

El parecido entre la consulta de la palabra $i$ y la clave de la palabra $j$ es un producto escalar. Un softmax convierte esas puntuaciones en pesos que suman 1, y la salida es la media de los valores ponderada por esos pesos. Con todas las palabras a la vez, en forma matricial:

Atención (Q, K, V) = softmax (\frac{Q K^{⊤}}{\sqrt{d_{k}}}) V, softmax (𝐬)_{j} = \frac{e^{s_{j}}}{\sum_{l} e^{s_{l}}} .

Tres «cabezas» de atención sobre la misma frase. Haz clic en una palabra para ver a qué palabras atiende; la matriz muestra todas las filas a la vez. Las matrices

W_{Q}

W_{K}

están diseñadas a mano sobre rasgos simples (sustantivo, género, posición, función de sujeto) para que se entienda qué hace cada cabeza; en un modelo real se aprenden y los rasgos los construyen las capas anteriores. Activa la máscara causal para ver la atención de un modelo generativo, que no puede mirar el futuro.

Fíjate en la cabeza de concordancia: «cansada» atiende a «gata» y no a «ratón», aunque «ratón» esté más cerca. La distancia no importa. Cualquier palabra puede consultar directamente a cualquier otra, en un solo paso.

¿Por qué dividir entre $\sqrt{d_{k}}$ ?

El factor de la fórmula no es decorativo. Sale de un cálculo de varianza.

Proposición (escala del producto escalar)

Si las componentes de $𝐪, 𝐤 \in ℝ^{d_{k}}$ son independientes, con media 0 y varianza 1, entonces

𝔼 [𝐪 \cdot 𝐤] = 0, Var (𝐪 \cdot 𝐤) = d_{k} .

Demostración

$𝐪 \cdot 𝐤 = \sum_{i = 1}^{d_{k}} q_{i} k_{i}$ . Cada término tiene media $𝔼 [q_{i}] 𝔼 [k_{i}] = 0$ y varianza $𝔼 [q_{i}^{2}] 𝔼 [k_{i}^{2}] = 1$ , y los términos son independientes, así que las varianzas se suman.

Con $d_{k} = 128$ , las puntuaciones tendrían una desviación típica de más de 11. El softmax de números tan dispares es prácticamente un «uno y todo ceros», y en esa zona su gradiente es casi nulo: la red dejaría de aprender. Dividir entre $\sqrt{d_{k}}$ devuelve la varianza a 1. Puedes comprobarlo en la figura desactivando la división.

El bloque Transformer

Un Transformer apila decenas de bloques iguales. Cada uno tiene dos partes:

Atención multicabeza. Varias atenciones en paralelo, cada una con sus propias $W_{Q}, W_{K}, W_{V}$ , que pueden especializarse en relaciones distintas (sintaxis, posición, correferencia…). Sus salidas se concatenan y se mezclan con otra matriz.
Una red neuronal de dos capas aplicada a cada posición por separado. Se cree que en ella se almacena buena parte del «conocimiento» factual del modelo.

Alrededor de cada parte hay una conexión residual, $𝐱 \leftarrow 𝐱 + Capa (𝐱)$ , y una normalización. La atención por sí sola no distingue el orden de las palabras: «el perro muerde al hombre» y «el hombre muerde al perro» darían lo mismo. Por eso se suma a cada embedding una codificación de posición. El artículo original usaba senos y cosenos de distintas frecuencias,

{PE}_{(p, 2 i)} = \sin (\frac{p}{10000^{2 i / d}}), {PE}_{(p, 2 i + 1)} = \cos (\frac{p}{10000^{2 i / d}}),

elegidos porque desplazar la posición equivale a una rotación, una transformación lineal. La cabeza de «palabra anterior» de la figura aprovecha justamente esa propiedad.

Generar texto: la máscara causal

Un modelo como GPT es un Transformer «solo decodificador»: en cada posición predice el token siguiente. Para que no haga trampa durante el entrenamiento, se impide que cada posición mire a las posteriores poniendo $- \infty$ en esas puntuaciones antes del softmax. Gracias a esta máscara causal, una sola pasada sobre un texto de $T$ tokens da $T$ predicciones de entrenamiento a la vez.

Aquí está la razón práctica del éxito del Transformer: a diferencia de las redes recurrentes, procesa todas las posiciones en paralelo, con multiplicaciones de matrices que las GPU ejecutan muy deprisa. Tiene un precio: la matriz de atención tiene $T \times T$ entradas, así que el coste crece con el cuadrado de la longitud del contexto. Buena parte de la investigación actual (FlashAttention, atención dispersa, modelos de espacio de estados) intenta abaratar ese término.

Lo que sabemos en teoría

Teorema (aproximación universal de secuencias; Yun et al., 2020)

Los Transformers con codificación de posición pueden aproximar con precisión arbitraria cualquier función continua de secuencias de longitud fija en secuencias, sobre un dominio compacto.

Es la versión, para secuencias, del teorema de aproximación universal. Pérez, Barceló y Marinković (2021) demostraron además que, con precisión aritmética ilimitada, los Transformers son Turing-completos: se cierra el círculo con el primer capítulo. Más intrigante aún es el aprendizaje en contexto: se ha demostrado que la atención puede implementar, dentro de una sola pasada, algoritmos como un paso de descenso de gradiente sobre ejemplos dados en el propio texto (von Oswald et al., 2023). Es una pista de cómo un LLM «aprende» de los ejemplos de una pregunta sin cambiar sus pesos.

Ya tenemos todas las piezas. En el último capítulo las juntamos a escala: billones de tokens, cientos de miles de millones de parámetros y un objetivo tan simple como predecir la palabra siguiente.

Referencias

J. L. Elman (1990). «Finding Structure in Time». Cognitive Science, 14(2).
S. Hochreiter y J. Schmidhuber (1997). «Long Short-Term Memory». Neural Computation, 9(8).
I. Sutskever, O. Vinyals y Q. V. Le (2014). «Sequence to Sequence Learning with Neural Networks». NeurIPS.
D. Bahdanau, K. Cho y Y. Bengio (2015). «Neural Machine Translation by Jointly Learning to Align and Translate». ICLR (arXiv 2014).
A. Vaswani et al. (2017). «Attention Is All You Need». NeurIPS.
C. Yun, S. Bhojanapalli, A. S. Rawat, S. J. Reddi y S. Kumar (2020). «Are Transformers universal approximators of sequence-to-sequence functions?». ICLR.
J. Pérez, P. Barceló y J. Marinković (2021). «Attention is Turing Complete». JMLR, 22.
J. Su et al. (2021). «RoFormer: Enhanced Transformer with Rotary Position Embedding». arXiv:2104.09864.
J. von Oswald et al. (2023). «Transformers Learn In-Context by Gradient Descent». ICML.

§Antes de la atención: leer de izquierda a derecha

§Atención: una búsqueda suave en un diccionario

§¿Por qué dividir entre dk?

§El bloque Transformer

§Generar texto: la máscara causal

§Lo que sabemos en teoría

§Referencias