Attention y el bloque transformer, derivados de una media

De 1997 a aproximadamente 2017, el resumen era un estado recurrente: mantener un vector $\mathbf{h}$ y actualizarlo en cada token, $\mathbf{h}_t = f(\mathbf{h}_{t-1}, \mathbf{x}_t)$. Tamaño fijo, entrada variable, exactamente la forma adecuada.

Fallaba de tres maneras, y la arquitectura de este capítulo responde a las tres. Hacer backpropagation a través de $T$ pasos multiplica $T$ jacobianos, así que el gradient se desvanece o explota: la enfermedad que el Capítulo 5 midió dentro de un único nodo $\tanh$. La LSTM¹ se diseñó precisamente contra eso y empujó el rango utilizable de decenas de pasos a cientos, sin cambiar el hecho de que la información del token 5 llega al token 500 solo si sobrevive a 495 actualizaciones secuenciales. Toda la fuente tenía que caber en un vector: en traducción sequence-to-sequence², un codificador comprime la entrada en su estado final. Bahdanau, Cho y Bengio pusieron nombre a ese cuello de botella y lo arreglaron en 2014, tres años antes del transformer, dejando que el decodificador tomara una suma ponderada de todos los estados del codificador con pesos que calculaba él mismo.³ Todo lo que sigue es esa idea, aplicada por una secuencia a sí misma, con la recurrencia eliminada. Y la actualización es secuencial por construcción: $\mathbf{h}_t$ necesita $\mathbf{h}_{t-1}$, y una GPU con diez mil núcleos no puede hacer nada con eso. La arquitectura que ganó no es obviamente más lista; es aquella cuyo paso caro es una multiplicación de matrices.

El otro sesgo inductivo clásico, la convolución —deslizar un filtro pequeño por toda la entrada, de modo que una característica detectada en cualquier sitio se detecte en todas partes— tampoco se construye aquí; es casi exactamente lo correcto para imágenes y queda delegado a un curso de visión. Ni la recurrencia ni la convolución reaparecen después de esta página, por eso ninguna recibe un capítulo: el Capítulo 1 prometió que las omisiones se declararían, no se esconderían.

La función más obvia que toma un número variable de vectores y devuelve un vector es la media:

Cualquier número de entradas, tamaño de salida fijo, diferenciable, gratis. Una tabla de embedding más esta media más una capa lineal hacia el vocabulario es un modelo de lenguaje completo en quince líneas. También es terrible, y la forma en que es terrible es toda la derivación.

El corpus de abajo es un megabyte de Shakespeare, 1.115.394 caracteres, pasado por un tokenizador BPE a nivel de bytes del tipo construido en el Capítulo 7 con un vocabulario de 1024: 459.760 tokens a 2,43 caracteres cada uno, dividido 90/10. Todos los modelos tienen anchura 128, ven 128 tokens y entrenan durante 3000 pasos de AdamW a $10^{-3}$ con un batch de 64. La perplejidad se mide en la partición reservada.⁴

Lee la segunda fila dos veces. Promediar el contexto no ayuda un poco; hace que el modelo sea cuatro veces peor que ignorar el contexto por completo. Dos razones, ambas demostrables en vez de empíricas.

La media no puede ver el orden. La suma conmuta, así que barajar la context window deja el resumen intacto; no aproximadamente:

A = torch.tril(torch.ones(T, T))
A = A / A.sum(1, keepdim=True)          # rows of the averaging matrix
y = x[torch.randperm(T)]                # the same tokens, shuffled
print((A[-1] @ x - A[-1] @ y).abs().max().item())

2.9802322387695312e-08

Ruido de coma flotante en una suma reordenada: los dos resúmenes son el mismo vector. Un modelo cuya única vista del contexto es una media no puede distinguir the dog bit the man de the man bit the dog. La tercera fila demuestra que esto no se arregla añadiendo posiciones a las entradas: un embedding de posición aprendido en cada token antes de promediar compró 2,14 puntos de 188. Las posiciones entran en la suma, y la suma las olvida.

Y la media ahoga el presente. En la posición 100, el token actual es una centésima parte del resumen. Eso tiene una reparación barata que ya conoces: conservar el token y sumarle el resumen: una conexión residual, del Capítulo 6, y la cuarta fila muestra lo que hace. Con la dilución reparada, la media uniforme no aporta nada en absoluto: 60,45 frente a una línea base de 59,71. Todos los tokens están ahí, ponderados por igual, y ponderar por igual equivale a no tener información.

El problema no es promediar. Son los pesos.

Promediar sobre un prefijo creciente parece un bucle. Es una multiplicación por una matriz triangular inferior cuyas filas suman uno; y también es, exactamente, una softmax:

loop = torch.stack([x[:t + 1].mean(0) for t in range(T)])   # the obvious version

A = torch.tril(torch.ones(T, T))
A = A / A.sum(1, keepdim=True)
mat = A @ x                                                  # the same thing

S = torch.zeros(T, T).masked_fill(torch.tril(torch.ones(T, T)) == 0, float("-inf"))
soft = F.softmax(S, dim=-1) @ x                              # and the same thing again

loop vs matmul   max |diff| = 5.960464477539063e-08
loop vs softmax  max |diff| = 5.960464477539063e-08

the averaging matrix A (rows sum to 1, upper triangle is zero):
  1.000 0.000 0.000 0.000 0.000 0.000
  0.500 0.500 0.000 0.000 0.000 0.000
  0.333 0.333 0.333 0.000 0.000 0.000
  0.250 0.250 0.250 0.250 0.000 0.000
  0.200 0.200 0.200 0.200 0.200 0.000
  0.167 0.167 0.167 0.167 0.167 0.167

Tres componentes con nombre de un transformer están ya en pantalla. El triángulo es la máscara causal, impuesta por el objetivo: si la posición $t$ pudiera ver la posición $t{+}1$, la respuesta estaría en la entrada: la fuga que el Capítulo 6 te dijo que auditaras, pero dentro de la arquitectura. La softmax es cómo se implementa la máscara: poner las entradas prohibidas a $-\infty$ las envía exactamente a cero y normaliza lo que queda, así que enmascarar y normalizar son una sola operación. (Usa $-\infty$, no -1e9: es el valor que el enmascaramiento significa, sobrevive a una conversión a float16 como $-\infty$ y te evita decidir si la constante que elegiste es lo bastante grande para el rango en el que estás, que es la caja de coma flotante del Capítulo 2 haciendo una pregunta que no tienes que responder.) Y los scores son el parámetro libre. La media uniforme es lo que obtienes cuando todos los scores permitidos son el mismo número; pon ahí cualquier número y la softmax los convierte en pesos válidos.

El resto de este capítulo es una pregunta: ¿de dónde salen esos números?

No pueden ser parámetros simples. Una matriz $T \times T$ aprendida sería idéntica para cada frase: podría codificar “mira cuatro tokens atrás”, pero nunca “mira el sustantivo al que se refiere este pronombre”. El peso que enlaza la posición $t$ con la posición $i$ debe depender de lo que hay en ambas posiciones, porque la relevancia es una relación, no una propiedad: la palabra it no es intrínsecamente relevante; es relevante para algo.

La función más barata de dos vectores que devuelve un número es el producto escalar del Capítulo 1. Puntúa la posición $i$ para la posición $t$ como $\mathbf{x}_t \cdot \mathbf{x}_i$ y el mecanismo funciona, mal, de dos maneras que obligan a todo lo demás. El producto escalar de un vector consigo mismo es su norma al cuadrado, así que cada token atendería sobre todo a sí mismo. Y la relación sería simétrica: si it atiende mucho a animal, entonces animal atiende mucho a it, lo cual es falso en el lenguaje, donde un adjetivo necesita su sustantivo mucho más de lo que el sustantivo necesita el adjetivo.

Así que dale a cada token dos roles, como dos mapas lineales aprendidos de él: lo que esta posición está buscando, $\mathbf{q}_t = W_q\mathbf{x}_t$, la query; y lo que ofrece para que lo encuentren, $\mathbf{k}_i = W_k\mathbf{x}_i$, la key. Puntúa $\mathbf{q}_t \cdot \mathbf{k}_i$ y la simetría desaparece, porque $W_q \neq W_k$: un token puede anunciar una cosa y buscar otra.

Aún queda una cosa mal. La suma ponderada era sobre los propios $\mathbf{x}_i$, lo que obliga a que lo que se copia sea lo que se empareja. Emparejar quiere las características que identifican un token; copiar quiere las características útiles aguas abajo. Así que aprende un tercer mapa, $\mathbf{v}_i = W_v\mathbf{x}_i$, el value, y suma esos.

La fórmula es ahora contabilidad:

con $M$ como máscara causal, cero en y por debajo de la diagonal y $-\infty$ por encima. En código son treinta líneas, veinte de ellas formas:

class Head(nn.Module):
    """One head of causal self-attention."""

    def __init__(self, d_model, d_head, block):
        super().__init__()
        self.q = nn.Linear(d_model, d_head, bias=False)      
        self.k = nn.Linear(d_model, d_head, bias=False)      
        self.v = nn.Linear(d_model, d_head, bias=False)      
        self.d_head = d_head
        self.register_buffer("mask", torch.tril(torch.ones(block, block)).bool())

    def forward(self, x):
        T = x.shape[1]
        q, k, v = self.q(x), self.k(x), self.v(x)
        s = q @ k.transpose(-2, -1) / math.sqrt(self.d_head)          
        s = s.masked_fill(~self.mask[:T, :T], float("-inf"))          
        w = F.softmax(s, dim=-1)                                      
        return w @ v                                                  

Puntuar, enmascarar, normalizar, mezclar. Todo lo demás es una proyección.

Casi toda explicación de $\sqrt{d_k}$ dice “para evitar que la softmax se sature”, lo cual es cierto y no explica nada. El argumento son dos líneas de la varianza del Capítulo 2. Si las entradas de $\mathbf{q}$ y $\mathbf{k}$ son independientes, con media cero y varianza uno, cada producto $q_j k_j$ tiene varianza uno, y las varianzas de cosas independientes se suman:

Así que los scores tienen desviación típica $\sqrt{d_k}$. Medido sobre veinte mil pares aleatorios:

     d     Var(q.k)         std   sqrt(d)
     4        3.975       1.994     2.000
    16       16.071       4.009     4.000
    64       64.249       8.016     8.000
   256      253.065      15.908    16.000
  1024     1015.562      31.868    32.000

Por qué importa: la softmax es sensible a la escala de una forma en que una capa lineal no lo es. Duplicar la entrada de una capa lineal duplica su salida; multiplicar los scores por diez antes de una softmax convierte una mezcla suave en una elección dura. Una fila de 64 scores, con y sin la división:

“Tokens efectivos” es el exponencial de la entropía: sobre cuántas posiciones promedia realmente la fila. Sin dividir, en $d_k = 256$, una head recién inicializada atiende exactamente a un token de 64, elegido solo por el sorteo aleatorio.

Eso es malo hacia delante y peor hacia atrás, con una forma que el Capítulo 5 ya midió en una $\tanh$. Una softmax comprometida con una entrada tiene casi derivada cero: la diagonal de su jacobiano es $w_i(1-w_i)$, cero en ambos extremos. Sobre dos mil filas aleatorias:

En $d_k = 1024$, siete filas de cada diez están congeladas antes de empezar el entrenamiento, y una head que empieza congelada no puede aprender dónde mirar. Dividida, la cantidad se mantiene plana en 0,96 en cualquier anchura y nada se satura.

Ahora la parte que nadie publica: ¿cambia la perplejidad final? Elimina la división y entrena, con cuatro anchuras de head:

Las dos primeras filas salen del presupuesto de 3000 pasos anterior; las tres últimas son una ejecución más corta —1500 pasos, batch de 32, una head, sin normalización antes de las proyecciones— con ambas variantes bajo ajustes idénticos.

En $d_k = 32$ la división no vale nada y la ejecución sin ella queda ligerísimamente por delante. Eso no es una licencia para quitarla, porque en 256 vale 18 puntos de perplejidad y en 1024 vale 17. El mecanismo se ve en los propios scores:

La head sin dividir no se recupera. Se desboca: la desviación típica de sus scores pasa de 21 en la inicialización a 5147, la entropía de attention cae a cero, y el 99,9 % de las filas pone más de 0,99 de su peso en un único token. Una vez que una head es un selector duro, su gradient es casi cero y nada la trae de vuelta, así que el colapso es estable. La head dividida se queda en una desviación típica de score de 3,44 tras el mismo entrenamiento, que es una mezcla suave que todavía puede cambiarse.

Vaswani et al. dicen exactamente esto y nada más: sospechan que los productos “crecen en magnitud para valores grandes de $d_k$” y dividen.⁵ La palabra grandes soporta el peso, y las tablas dicen dónde empieza lo grande: nada en 32, todo en 256.

Una head es una fila de softmax por posición, así que contiene una respuesta a “qué es relevante aquí”. Predecir la palabra después de the en the animal that crossed the wet street necesita a la vez el hueco sintáctico, el sujeto y el token anterior, y una única distribución de probabilidad no puede concentrarse en tres lugares. Así que ejecuta varias heads en paralelo, cada una de anchura $d_{\text{model}}/h$, concatena y mezcla con una matriz más $W_o$: has particionado la anchura, no la has aumentado.

Attention también hace exactamente una cosa: mueve información entre posiciones. Cada operación del código anterior es lineal a lo largo del eje de características, y el Capítulo 5 demostró qué es una pila de mapas lineales. Así que cada bloque también lleva una pequeña MLP aplicada a cada posición de forma independiente, que expande la anchura por cuatro y vuelve, con una GELU en medio. Merece la pena memorizar la división del trabajo: attention mezcla entre posiciones; la red feed-forward calcula dentro de una posición.

La escalera completa, cada fila añadiendo una pieza a la fila de arriba:

Los pesos aprendidos vencen a los uniformes por 14 puntos de perplejidad, que es todo el argumento de este capítulo en una fila. Cuatro heads compran otros 3 por 16.512 parámetros extra. Y la misma head vale 9 puntos más sumada que reemplazando: attention trae información hacia dentro, no decide qué es una posición.

Ahora dónde se sientan realmente los parámetros, lo que sorprende a quienes solo han visto el diagrama:

Dos tercios de cada bloque transformer son la red feed-forward, a cualquier escala, porque attention tiene cuatro matrices $d \times d$ y la MLP tiene el equivalente a ocho. Sea lo que sea que sepa un modelo, la mayoría de los parámetros que lo sostienen están en la MLP por posición.

LayerNorm se construyó y midió en el Capítulo 6, y este capítulo la usa tal y como quedó allí; las conexiones residuales se nombraron y ablacionaron allí, y se construyen aquí. Las filas de “sumada, no reemplazando” de arriba son conexiones residuales, que valen 188 puntos de perplejidad para la media y 9 para una head. LayerNorm⁷ normaliza cada ejemplo a través de sus características, y el Capítulo 6 dio las razones por las que ella, y no BatchNorm, sobrevivió aquí: sin dependencia del batch, sin estadísticas acumuladas, idéntica en entrenamiento e inferencia, indiferente a la longitud de secuencia. Cada una de esas razones se convierte en requisito cuando generas un token cada vez para un usuario, que es donde acaba el Capítulo 13. Cuesta 768 parámetros y compra 1,8 puntos de perplejidad.

class Block(nn.Module):
    def forward(self, x):
        x = x + self.att(self.ln1(x))     
        x = x + self.ff(self.ln2(x))      
        return x

Mira dónde se coloca la normalización: en la entrada de cada subcapa, con la ruta residual de entrada a salida nunca normalizada. Eso es pre-norm. El artículo de 2017 hace lo contrario, x = LayerNorm(x + Att(x)): post-norm, que pone una LayerNorm en la propia ruta residual.

Xiong et al. explicaron la diferencia a través del gradient en la inicialización, que en una red post-norm está mal escalado con la profundidad: la razón por la que el transformer original necesitaba un warmup de learning rate para entrenar siquiera.⁸ Doce bloques, 1000 pasos, learning rate $3 \times 10^{-3}$:

gradient norm per block at initialisation, before any step
  pre-norm    block 1 0.0498 ... block 12 0.0657   ratio last/first  1.32
  post-norm   block 1 0.0977 ... block 12 0.1613   ratio last/first  1.65

  pre-norm,  no warmup          perplexity   37.82
  pre-norm,  200-step warmup    perplexity   37.62
  post-norm, no warmup          perplexity  308.05
  post-norm, 200-step warmup    perplexity   37.88

Post-norm sin warmup es ocho veces peor, y post-norm con warmup iguala exactamente a pre-norm. Warmup no es aquí una buena práctica general; es un parche para una disposición concreta de la normalización, y mover la LayerNorm elimina su necesidad. Por eso prácticamente todos los modelos desde 2019 son pre-norm, y por eso el diagrama de 2017 debe leerse como historia, no como especificación.

Elimina los embeddings de posición y el modelo sigue entrenando; simplemente no puede saber dónde está nada, y eso es una simetría, no un fallo de entrenamiento. Nada en el score de attention menciona las propias $t$ o $i$, así que permutar la entrada permuta la salida: self-attention es equivariante a permutaciones. Es la ceguera al orden de la media con un disfraz mejor: la máscara causal restaura parte del orden, ya que cada posición ve un prefijo distinto, pero dentro de un prefijo todos los ordenamientos son iguales.

Cuatro formas de inyectar posición, entrenadas en ventanas de 64 tokens y evaluadas en 64, 128 y 256, más allá de cualquier longitud que vieran:

Embeddings absolutos aprendidos —un vector por posición, sumado al token— ganan en la longitud entrenada y luego caen por un precipicio, porque la posición 100 nunca estuvo en un batch y su embedding sigue siendo el vector aleatorio con el que empezó. Sinusoides, la elección original, se calculan en vez de aprenderse, a partir de senos y cosenos con frecuencias espaciadas geométricamente; el artículo de 2017 esperaba que eso extrapolara, y la tabla dice que no: la función está definida en la posición 200, pero el modelo nunca aprendió a leerla allí. RoPE⁹ no añade nada y en su lugar rota query y key por un ángulo proporcional a la posición, en cortes bidimensionales; como rotar ambos lados de un producto escalar por igual lo deja sin cambios, el score acaba dependiendo solo de $t - i$, así que la posición se vuelve relativa gratis y no hay una tabla que se agote. Se degrada, pero se degrada. ALiBi¹⁰ es el resultado más simple y más extraño aquí: una penalización lineal sobre el score proporcional a la distancia, con una pendiente distinta por head. Su perplejidad mejora cuando la context window crece más allá de la longitud de entrenamiento, de 44,95 a 42,49, porque la penalización está definida a cualquier distancia y cada head sigue haciendo lo que se entrenó para hacer.

La lección sobrevive a la tabla: una arquitectura que no puede representar algo es un problema distinto de una que nunca aprendió ese rango, y el segundo es el que muerde. También es la maquinaria detrás de cada anuncio de “hemos ampliado el contexto a 128K”: casi siempre son reescalados de una codificación rotatoria, y por eso el Capítulo 16 dice que el límite de contexto se mueve en vez de desaparecer.

Dropout se hereda de la misma forma: aparece en los pesos de attention después de la softmax, en la salida de cada subcapa antes de la suma residual y en la suma de embedding, haciendo exactamente lo que describió el Capítulo 6. En grandes ejecuciones de pretraining suele ponerse a cero, porque un modelo que ve cada token una sola vez no está en posición de sobreajustar.

Dos tensores de la capa tienen forma $n \times n$, donde $n$ es el número de tokens: los scores y los pesos después de la softmax. Todo lo demás —cada proyección, toda la MLP— es lineal en $n$.

Una capa de attention, 512 de anchura, 8 heads, batch de uno, float32, en una GPU de portátil. Lee las dos columnas de milisegundos solo por sus proporciones: son tiempo real en una tarjeta de portátil de 8 GB que baja de 1.785 MHz a menos de 300 MHz cuando se calienta, así que una ejecución en frío del mismo código vuelve de siete a diez veces más rápido y una ocupada aún más lenta. Las columnas de megabytes son recuentos de bytes del asignador y no se mueven.

  tokens   ms total    ms x4   ms projections   attn matrix MB    peak MB    MB x4
     128      2.246        -            1.324              0.5       14.6        -
     256      2.855     1.27            2.113              2.0       19.2     1.31
     512      5.761     2.02            3.105              8.0       34.4     1.79
    1024     16.414     2.85            4.008             32.0       89.1     2.59
    2048     51.573     3.14            9.989            128.0      296.1     3.32
    4096    225.432     4.37           20.176            512.0     1100.1     3.72
    8192    832.838     3.69           40.106           2048.0     4300.1     3.91
   16384   OUT OF MEMORY                                 8192.0

fitted exponent (log-log slope, last four rows):  time ~ n^1.91   memory ~ n^1.87

Las columnas x4 son la proporción respecto a la fila superior, y duplicar $n$ converge exactamente a 4 tanto para tiempo como para memoria: 3,91 en el último paso frente a un 4 teórico. La columna de proyecciones es el control: 4,0 ms en 1024 tokens frente a 40,1 ms en 8192, un factor de diez para un factor de ocho. Lineal, como se prometió.

Luego la última fila. Una capa de attention, una secuencia, sin modelo alrededor, se queda sin memoria en una GPU de 8 GB con 16.384 tokens: solo la matriz de scores serían 8 GB, al ser 8 heads por 16.384 por 16.384 por 4 bytes. No el modelo; un tensor intermedio en una capa.

Ese es el hecho físico debajo de tres capítulos posteriores. Es por lo que una context window tiene un límite, que el Capítulo 16 convierte en precio. Es por lo que existe FlashAttention, que calcula el mismo resultado en teselas sin almacenar nunca la matriz: una optimización de memoria antes que de velocidad.¹¹ Y es la aritmética detrás del precio de un prompt largo, que el Capítulo 24 paga en un bucle de agent; un asunto separado del otro hallazgo de ese capítulo, que un modelo también usa peor un contexto largo, cosa que mide y se niega a achacar a esta fórmula.

El artículo de 2017 describe un encoder-decoder: una pila que lee la fuente con attention sin máscara, una segunda que genera el objetivo de forma causal, y un tercer tipo de attention en medio donde las queries del decodificador se encuentran con las keys del codificador. Eso es correcto para traducción, donde entrada y salida son dos secuencias.

Lo que ganó fue la mitad decoder-only: una pila, causal de principio a fin, entrada y salida en la misma secuencia; y la razón no es la elegancia. “Predecir el siguiente token” funciona sobre cualquier texto, así que el conjunto de entrenamiento es internet en vez de un corpus paralelo, y todo se convierte en esa única tarea: una traducción es un documento que contiene fuente y luego objetivo, una pregunta y su respuesta son un documento, una conversación con una tool call en medio es un documento. El Capítulo 11 trata de cómo se fabrica esa última. Los codificadores no desaparecieron: uno ve toda la entrada de una vez, que es lo que quieres cuando el trabajo es representar un texto en vez de continuarlo, y por eso los embeddings de retrieval del Capítulo 19 vienen de codificadores y no del modelo que chatea.

Con el bloque definido, el tamaño del modelo es aritmética. Por bloque, con anchura $d$ y una expansión por cuatro: $4d^2 + 4d$ para $W_q, W_k, W_v, W_o$ con sesgos en las cuatro, como los tiene GPT-2; la tabla anterior deja fuera el sesgo en tres de ellas, de ahí 2304 menos por bloque en $d = 768$; $8d^2 + 5d$ para la MLP; $4d$ para dos LayerNorms: $12d^2 + 13d$, más una tabla de tokens de $V \times d$ y, para posiciones absolutas, $n_{\text{ctx}} \times d$. Para la forma de GPT-2 small —$d = 768$, 12 bloques, un vocabulario de 50.257, un contexto de 1024, la capa de salida compartiendo pesos con el embedding—:

  token embeddings     50,257 x 768 = 38,597,376
  position embeddings   1,024 x 768 =    786,432
  one block                             7,087,872
  12 blocks                            85,054,464
  final LayerNorm         2 x 768 =        1,536
  total (weights tied)                124,439,808

Que es el tamaño publicado de ese modelo. La fórmula no es una aproximación; es el modelo. Observa también que casi un tercio de un modelo pequeño es la tabla de embedding, por eso el tamaño del vocabulario es una decisión arquitectónica y no de preprocesamiento: el intercambio que planteó el Capítulo 7.

La perplejidad es un número sobre un corpus. Lo que hace una head es otra pregunta, y un modelo entrenado en un megabyte de Shakespeare es el instrumento equivocado para ello: lo honesto sobre el mapa de attention de un modelo de 500.000 parámetros es decir que en su mayoría no es interpretable. Así que: un lenguaje donde la pregunta tiene una respuesta correcta.

La ilustración clásica es the animal did not cross the street because it was too tired, donde it es el animal, frente a …because it was too wet, donde una palabra mueve el referente a la calle. Son esquemas de Winograd¹⁴: pares de frases idénticas salvo por una palabra, donde esa palabra decide a qué se refiere un pronombre.

También son resolubles haciendo trampas, que es la parte que los tutoriales omiten. Si los dos candidatos son un animal y un lugar, tired y wet identifican el referente por categoría, y un modelo que solo sabe qué palabras están presentes acierta sin saber nada de orden. Medido en esa versión de la tarea, con pares animal/lugar reservados:

uniform causal average           held-out referent accuracy 100.0 %
one transformer block            held-out referent accuracy  91.7 %

La bolsa de palabras vence al transformer. Cualquier demostración construida sobre esa frase no prueba nada sobre attention.

Así que cierra el agujero: toma ambos candidatos de un único conjunto de dieciséis sustantivos, cualquiera de los cuales puede aparecer en cualquiera de los dos huecos, y divide los adjetivos por rol en vez de por categoría: cuatro hacen que it sea quien cruza (tired, scared, slow, weak), cuatro hacen que sea lo cruzado (wet, wide, busy, steep).

the {x} did not cross the {y} because it was too {adj} , so the {ref} waited .

Entrena como un predictor ordinario del siguiente token, puntúa una posición —la palabra después de so the— y construye el conjunto reservado a partir de pares de sustantivos cuyo orden invertido estuvo en entrenamiento, de modo que cualquier cosa que sepa qué dos sustantivos están presentes pero no cuál vino primero debe responder al revés.

El azar entre los dos sustantivos presentes es 50 %. La media uniforme cae en 27,9 % y responde con el sustantivo equivocado del par exactamente la mitad de las veces: la firma de algo que sabe qué palabras hay y nada sobre su orden, tal y como predijo la prueba de barajado tres secciones atrás.

Ahora el mapa: la attention en la posición que debe nombrar el referente, promediada sobre las cuatro heads de cada bloque, para las dos frases que difieren en una palabra. Una media uniforme pondría 0,067 en cada uno de los quince tokens visibles.

the animal did not cross the street because it was too tired , so the animal waited .
  blk 1  the:0.00 animal:0.70 did:0.00 not:0.00 cross:0.00 the:0.00 street:0.06
         because:0.00 it:0.00 was:0.00 too:0.00 tired:0.00 ,:0.05 so:0.00 the:0.19
  blk 2  the:0.00 animal:0.00 did:0.00 not:0.00 cross:0.00 the:0.00 street:0.00
         because:0.00 it:0.00 was:0.00 too:0.00 tired:1.00 ,:0.00 so:0.00 the:0.00

the animal did not cross the street because it was too wet , so the street waited .
  blk 1  the:0.00 animal:0.70 did:0.00 not:0.00 cross:0.00 the:0.00 street:0.06
         because:0.00 it:0.00 was:0.00 too:0.00   wet:0.00 ,:0.05 so:0.00 the:0.19
  blk 2  the:0.00 animal:0.00 did:0.00 not:0.00 cross:0.03 the:0.00 street:0.49
         because:0.00 it:0.00 was:0.00 too:0.20   wet:0.03 ,:0.00 so:0.00 the:0.25

El bloque 1 es idéntico en ambas frases: 0,70 sobre el primer sustantivo, sea cual sea el adjetivo. Eso no es un fallo sino una prueba: en la primera capa, la query de una posición es función del propio token y del índice de esa posición, y the en la posición 14 es el mismo token en ambas frases. Una head de primera capa no puede condicionar sobre una palabra que aún no ha traído. Así que el bloque 1 hace lo único útil disponible y arrastra el primer sustantivo hacia delante.

El bloque 2 es donde las frases se separan, y la misma fila a través de los ocho adjetivos muestra la regla que encontró el modelo:

Para un adjetivo de quien cruza, el segundo bloque gasta todo su peso en el adjetivo, porque la respuesta ya está en el flujo residual —el bloque 1 la puso ahí— y solo necesita confirmación. Para un adjetivo de lo cruzado, va a buscar el otro sustantivo. Eso es un circuito de dos saltos: una head mueve un candidato hacia delante, una head de una capa posterior lee un token que decide si conservarlo. La composición entre capas es el mecanismo, y por eso un bloque llegó al 92,7 % y dos llegaron al 100 %.

También es la forma del circuito mejor documentado en modelos reales. Induction heads —una head de token anterior que alimenta a una head en la capa siguiente que completa el patrón [A][B] … [A] → [B]— son lo que el trabajo de interpretabilidad de Anthropic identifica detrás de una gran parte del in-context learning, y se forman en un momento identificable durante el pretraining. Este capítulo no intenta ese análisis: queda delegado, con ambos artículos en las referencias, porque leer circuitos de un modelo real es un campo de investigación y no una sección.

Por último, la implementación. Las treinta líneas de arriba, con sus pesos copiados de los propios de PyTorch:

ours vs nn.MultiheadAttention           max |diff| = 1.7881393432617188e-07
ours vs F.scaled_dot_product_attention  max |diff| = 1.7881393432617188e-07

$1.8 \times 10^{-7}$ en salidas cuya magnitud media es 0,159: la misma aritmética en un orden distinto, con precisión float32.

Tienes la arquitectura de la que parte cada modelo del resto de este curso, y es más pequeña que su reputación: una media ponderada cuyos pesos se aprenden, una MLP por posición que sostiene dos tercios de los parámetros, dos normalizaciones y dos sumas, apiladas.

Lo que no tienes es un modelo que sepa algo, y apilar no lo arreglará por sí solo. Dos bloques en este corpus alcanzan una perplejidad de entrenamiento de 14,49 y una perplejidad de validación de 40,57, frente a 18,77 y 38,07 de un bloque: más capacidad, mejor en lo que ha visto, peor en lo que no ha visto, que es la tabla del Capítulo 6 con un transformer dentro. La distancia entre este modelo y aquellos con los que hablan los capítulos 14 a 30 no es arquitectónica. Es el mismo bloque, repetido más veces, sobre muchísimo más texto.

Eso lo convierte en un problema de contabilidad, y la contabilidad es más extraña de lo que parece. ¿Cuánto texto, y de dónde lo saca alguien? ¿Cuánta aritmética, y cómo la estimas antes de gastar el dinero? Dado un presupuesto fijo, ¿es mejor hacer el modelo más grande o enseñarle más datos? ¿Y hay una respuesta correcta, o solo una moda? El Capítulo 10 responde a las tres por medición, y pone precio a la forma útil más barata de la pregunta: ¿cuánto cuesta, hoy, entrenar desde cero un modelo como GPT-2?

Tres explicaciones de este material son mejores que esta en aquello para lo que sirven, y este capítulo está escrito para leerse junto a ellas. The Illustrated Transformer, de Jay Alammar, es la mejor imagen del flujo de datos que se ha dibujado. The Annotated Transformer, de Harvard NLP, es el artículo de 2017 con código ejecutable intercalado línea a línea. Let's build GPT: from scratch, in code, spelled out, de Andrej Karpathy, construye el mismo modelo en directo en dos horas, y la escalera de ablaciones de arriba es la misma columna vertebral medida en otro corpus. Para la cuestión de interpretabilidad que este capítulo solo roza, las fuentes primarias son Elhage et al., A Mathematical Framework for Transformer Circuits (2021), y Olsson et al., In-context Learning and Induction Heads (2022), ambos del grupo de interpretabilidad de Anthropic.