El transformer completo

El pase completo, ecuación por ecuación

Para una secuencia de $T$ tokens con números $t_1,\dots,t_T$, una tabla de embeddings $E$ (una fila de $d$ números por palabra del vocabulario) y $L$ bloques, en la variante «con normalización previa» de GPT-2, que es la de nuestro juguete ($d=8$, $L=2$):

y, en cada bloque, con $\mathrm{LN}$ la normalización de capa:

Esas dos sumas son las conexiones residuales: cada parte suma su aporte a $x$. Al salir del último bloque, la lista del último token da los puntajes (en inglés, logits) y softmax las probabilidades:

«Comparar con cada palabra del vocabulario» es exactamente esto: un producto punto con cada fila de $E$, la misma tabla de la entrada. Reutilizarla ahorra $V\cdot d$ parámetros; ya lo hacía el artículo de 2017 (sección 3.4) y la técnica se estudia en Press y Wolf (2017). Al entrenar se usan las $T$ posiciones a la vez: cada una predice el token siguiente.

La atención con máscara

Con $Q$, $K$ y $V$ como en la lección 21 y la máscara $M_{ij}=0$ si $j\le i$ y $-\infty$ si $j>i$:

Como $e^{-\infty}=0$, el futuro recibe peso cero. Así lo describe el artículo original para el decodificador: se pone en $-\infty$ toda entrada de softmax que corresponda a una conexión «ilegal», para preservar la propiedad autorregresiva (sección 3.2.3). Nuestro juguete tiene dos cabezas de $d_h=4$ números por bloque; la tabla de la escena es una cabeza del segundo bloque.

El experimento. Entrenamos dos modelos idénticos (misma semilla, mismos datos, 4 000 pasos) y a uno le quitamos la máscara. Con frases nuevas, el modelo con máscara tiene una pérdida de 0.527 nats por palabra, casi la entropía de la gramática (0.526, el mínimo posible). El modelo sin máscara baja su pérdida de entrenamiento a 0.059 copiando, pero usado como se usa al escribir (con máscara, sin futuro) da 3.088, peor que repartir por igual entre doce palabras: $\ln 12 = 2.485$.

La red de cada bloque

$\mathrm{MLP}(x)=W_2\,\mathrm{GELU}(W_1x+b_1)+b_2$, con $W_1$ de $d\times 4d$ y $W_2$ de $4d\times d$, la misma para todas las posiciones del bloque («aplicada a cada posición por separado y de forma idéntica», sección 3.3 del artículo). El artículo original usaba ReLU; GPT-1 y GPT-2 usan GELU (Hendrycks y Gimpel, 2016), una ReLU suave: $\mathrm{GELU}(x)=x\,\Phi(x)$, con $\Phi$ la distribución normal acumulada.

Normalización de capa

Sobre los $d$ números de una lista (Ba, Kiros y Hinton, 2016):

$\gamma$ y $\beta$ son los «dos ajustes que también se aprenden» (ganancia y sesgo, $2d$ parámetros por normalización). El artículo de 2017 normalizaba después de sumar, $\mathrm{LN}(x+\mathrm{Subcapa}(x))$; GPT-2 la movió a la entrada de cada parte y añadió una más tras el último bloque. Xiong y colegas (2020) muestran que, con la normalización previa, los gradientes están bien comportados desde el inicio y se puede entrenar sin la fase de calentamiento de la tasa de aprendizaje. En la escena, la lista de «come» al entrar al segundo bloque tiene promedio 1.10 y desviación estándar 0.78.

Codificación de posición

La del artículo de 2017 (sección 3.5), con $pos$ el lugar e $i$ el índice del par de números:

Las longitudes de onda van de $2\pi$ a $10000\cdot2\pi$. Los autores probaron también posiciones aprendidas y obtuvieron resultados casi idénticos; GPT-1 aprende las suyas. «La atención sola no sabe el orden» es exacto para la atención sin máscara: si se permutan los tokens, sus salidas se permutan igual. Con máscara causal la cosa es más sutil: Haviv y colegas (2022) muestran que un modelo de lenguaje causal sin codificación de posición aprende una noción implícita de posición, gracias a la propia máscara. Aun así, casi todos los modelos usan alguna. En la escena, el cálculo «sin posición» usa una cabeza real de nuestro modelo sobre los embeddings solos y sin máscara.

Contar parámetros

Con $V$ palabras, $d$ números por lista, $L$ bloques, $ctx$ posiciones aprendidas y una red de $d\to4d\to d$:

Parte	Parámetros	Juguete ($V{=}12$, $d{=}8$, $L{=}2$)
Embeddings (y posiciones aprendidas)	$V d + ctx\,d$	96 (posición con ondas: 0)
Atención	$L\,(4d^2+4d)$	576
Red	$L\,(8d^2+5d)$	1 104
Normalizaciones	$(4L+2)\,d$	80
Total	$\approx 12\,L\,d^2$ si $d$ es grande	1 856

La aproximación $N\approx 12\,L\,d^2$ es la de Kaplan y colegas (2020, ecuación 2.1), sin contar embeddings. Con el vocabulario de GPT-2 (50 257) y $ctx=2048$, la tabla da 125 226 240 para 12 bloques con $d=768$ y 174 604 259 328 para 96 bloques con $d=12\,288$; el artículo de GPT-3 (tabla 2.1) lista 125 millones para «GPT-3 Small» y 175.0 mil millones para GPT-3. El reparto por partes de la escena es nuestro cálculo con esta tabla. Para GPT-2 pequeño (contexto 1 024) da 124 439 808, lo mismo que cuenta la biblioteca Hugging Face Transformers (comprobado en train_tiny.py); el artículo de GPT-2 decía 117 millones.

En código

import numpy as np

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

def causal_attention(Q, K, V):
    T, d = Q.shape
    scores = Q @ K.T / np.sqrt(d)
    scores[np.triu_indices(T, k=1)] = -np.inf      # máscara causal: nada de mirar al futuro
    w = softmax(scores)
    return w, w @ V

def layer_norm(x, g, b, eps=1e-5):
    return g * (x - x.mean(-1, keepdims=True)) / np.sqrt(x.var(-1, keepdims=True) + eps) + b

def n_params(d, L, V, ctx, ff=None):
    ff = ff or 4 * d
    return V*d + ctx*d + L*(4*d*d + 4*d) + L*(2*d*ff + ff + d) + (4*L + 2)*d

rng = np.random.default_rng(0)
Q = K = V = rng.normal(size=(3, 4))
w, out = causal_attention(Q, K, V)
print(np.round(w, 2))
x = rng.normal(2.0, 3.0, size=8)
y = layer_norm(x, np.ones(8), np.zeros(8))
print(round(y.mean(), 3), round(y.std(), 3))
print(n_params(8, 2, 12, 0), n_params(12288, 96, 50257, 2048))
# [[1.   0.   0.  ]
#  [0.25 0.75 0.  ]
#  [0.17 0.13 0.7 ]]
# 0.0 1.0
# 1856 174604259328

El modelo de la lección está completo en model.ts y se entrena con train_tiny.py (PyTorch, en CPU, menos de un minuto): ambos viven junto a esta lección, y verify-forward.mts comprueba que el pase en TypeScript coincide con PyTorch (diferencia máxima de 7·10⁻⁶).

Lo que viene

Un modelo con esta arquitectura empieza con parámetros al azar. La lección 23 muestra qué aprende al entrenarlo para predecir la siguiente palabra con texto a escala de internet, y la 24, cuánto cuesta. Cómo se sirve (la caché KV, que evita recalcular los tokens anteriores) llega en la lección 26.