Atención

La fórmula

Para $n$ tokens, se apilan sus consultas y claves (de $d_k$ números) y sus valores (de $d_v$) en matrices $Q,K\in\mathbb{R}^{n\times d_k}$ y $V\in\mathbb{R}^{n\times d_v}$:

Fila por fila: la puntuación del token $i$ sobre el $j$ es $s_{ij}=q_i\cdot k_j/\sqrt{d_k}$; el peso es $\alpha_{ij}=e^{s_{ij}}/\sum_l e^{s_{il}}$ (cada fila suma 1) y la salida es $o_i=\sum_j\alpha_{ij}\,v_j$. Es la «atención de producto punto escalado» de Vaswani et al. (2017), que la describen como una función que lleva una consulta y un conjunto de pares clave-valor a una salida: una suma ponderada de los valores.

En la escena, la consulta de «banco» es $q=(1,1)$ y la clave de «jardín» $k=(1,0.8)$: $q\cdot k=1\cdot1+1\cdot0.8=1.8$. Las puntuaciones $1.8,\,0.7,\,0,\,-1.6$ dan $e^{1.8}\approx6.0$, $e^{0.7}\approx2.0$, $e^{0}=1$ y $e^{-1.6}\approx0.2$; la suma es $\approx10$, y los pesos, $60\,\%$, $20\,\%$, $10\,\%$ y $2\,\%$. La escena no divide entre $\sqrt{2}$ para que las cuentas se sigan de cabeza; con la división, los pesos serían $46\,\%$, $21\,\%$, $13\,\%$ y $4\,\%$, y el vector nuevo, $(0.63,\,0.06)$ en vez de $(0.75,\,0.05)$: el mismo lado, menos marcado. Todos los números del ejemplo están escritos a mano (ver data.ts).

De dónde salen Q, K y V

Si $X$ es la matriz de vectores de entrada (embeddings; en un transformer, más la posición), $Q=XW^Q$, $K=XW^K$ y $V=XW^V$. Esas tres matrices son los parámetros que se aprenden; los pesos $\alpha_{ij}$ no: se calculan de nuevo para cada frase. En la autoatención, $Q$, $K$ y $V$ salen de la misma $X$. En la atención de Bahdanau et al. (2014), para traducir, las consultas venían del decodificador y las claves y valores del codificador, y la puntuación era una red pequeña, $e_{ij}=a(s_{i-1},h_j)$, en lugar de un producto punto; los pesos eran el mismo softmax y el resultado, la misma suma ponderada $c_i=\sum_j\alpha_{ij}h_j$. Los autores lo dicen así: el decodificador decide a qué partes de la frase de origen prestar atención.

Por qué $\sqrt{d_k}$

Si las componentes de $q$ y $k$ son independientes, con media 0 y varianza 1, el producto $q\cdot k=\sum_{i=1}^{d_k}q_ik_i$ tiene media 0 y varianza $d_k$ (nota 4 de Vaswani et al., 2017): su tamaño típico es $\sqrt{d_k}$, ocho con $d_k=64$. Dividir entre $\sqrt{d_k}$ devuelve la varianza a 1. Sin ello, el softmax se satura: si $\alpha=\text{softmax}(s)$, entonces $\partial\alpha_i/\partial s_j=\alpha_i(\delta_{ij}-\alpha_j)$, que tiende a 0 cuando un peso se acerca a 1 y los demás a 0. Es la «pendiente casi nula» de la narración. Con ocho claves al azar y $d_k=64$, el peso mayor promedia $0.89$ sin dividir y $0.36$ dividiendo (100 sorteos; el sorteo de la escena da $0.98$ y $0.35$).

Varias cabezas

Según Vaswani et al. (2017), varias cabezas permiten atender a la vez a información de distintos «subespacios» de representación; con una sola, el promedio lo impide. En su modelo base, $h=8$ y $d_k=d_v=512/8=64$, así que ocho cabezas cuestan casi lo mismo que una grande. BERT base tiene 12 capas de 12 cabezas (144) con 768 dimensiones, es decir, $768/12=64$ por cabeza (Devlin et al., 2018).

Los mapas reales de la lección

Se calcularon con real-attention.py (en la carpeta de la lección; corre en CPU): bert-base-uncased para el inglés y bert-base-multilingual-cased para el español, ambos con licencia Apache-2.0 y revisión fijada. Son los pesos de una sola cabeza tras softmax, sin promediar; las piezas de una palabra se suman. La cabeza de «banco» se eligió así: entre las cabezas en las que el mayor peso de «banco» cae en la palabra que decide su sentido en las dos frases (jardín y dinero), la de mayor suma de esos dos pesos. Solo 4 de 144 cumplen esa condición en español (capa 1, cabeza 4: 44 % y 65 %) y 6 de 144 en inglés (capa 3, cabeza 11: 93 % y 60 %). En cambio, el mayor peso de «banco» cae en una marca ([CLS], [SEP] o el punto) en 58 de 144 cabezas en español y en 113 de 144 en inglés. Las otras tres son las de mayor peso promedio en la palabra anterior, en la siguiente y en las marcas. Clark et al. (2019) describen estos mismos tipos de cabeza en BERT: las que miran a los delimitadores, las que miran a una posición fija y las que miran a toda la frase.

Costo

La matriz de pesos es $n\times n$: el cómputo por capa crece como $O(n^2\cdot d)$. Una capa recurrente cuesta $O(n\cdot d^2)$, pero necesita $O(n)$ pasos uno tras otro, y la información entre dos posiciones recorre un camino de longitud $O(n)$; en la autoatención, ambos son $O(1)$ (tabla 1 de Vaswani et al., 2017). Por eso la atención aprovecha bien una GPU, y por eso los contextos largos son caros.

En código

import numpy as np

def softmax(x):
    e = np.exp(x - x.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V, scale=True):
    d = K.shape[-1]
    scores = Q @ K.T                     # productos punto: qué tanto coincide cada consulta con cada clave
    if scale:
        scores = scores / np.sqrt(d)     # la pieza que la escena deja para el final
    weights = softmax(scores)            # cada fila suma 1
    return weights @ V, weights          # mezcla ponderada de los valores

# «Me senté en el banco de mi jardín»: los números de la escena (d = 2)
K = np.array([[-1.6, 0], [-0.3, 1], [-0.8, -0.8], [-0.8, -0.8],
              [0.5, -0.5], [-0.8, -0.8], [-0.8, -0.8], [1, 0.8]])
V = np.array([[0, 0], [0.5, 0], [0, 0], [0, 0],
              [0.5, 0.5], [0, 0], [0, 0], [1, 0]])     # (asiento, dinero)
q = np.array([[1.0, 1.0]])                          # la consulta de «banco»

out, w = attention(q, K, V, scale=False)
print(np.round(w, 2), np.round(out, 2))
out, w = attention(q, K, V)                          # con la división entre √2
print(np.round(w, 2), np.round(out, 2))
# [[0.02 0.2  0.02 0.02 0.1  0.02 0.02 0.6 ]] [[0.75 0.05]]
# [[0.04 0.21 0.04 0.04 0.13 0.04 0.04 0.46]] [[0.63 0.06]]

Y la comprobación de $\sqrt{d}$: la desviación típica de $q\cdot k$ crece como $\sqrt{d}$; al dividir, queda en 1.

rng = np.random.default_rng(0)
for d in (1, 2, 64, 512):
    q = rng.standard_normal((20000, d))
    k = rng.standard_normal((20000, d))
    s = (q * k).sum(axis=1)                         # 20 000 productos punto q·k
    print(d, round(s.std(), 1), round((s / np.sqrt(d)).std(), 2))
# 1 1.0 1.0
# 2 1.4 1.0
# 64 8.0 1.0
# 512 22.6 1.0

Lo que viene

Falta lo que la atención sola no hace: no sabe el orden de las palabras (si se barajan, cada token reparte lo mismo entre los mismos tokens), y al escribir texto cada token solo puede mirar hacia atrás (antes del softmax se pone $-\infty$ en las puntuaciones que miran al futuro). Con una codificación de posición, esa máscara, una red por token, conexiones residuales y normalización, se apilan bloques: el transformer, en la lección 22.