La raíz que faltaba en la atención

La auto-atención quedó montada del todo salvo por un símbolo. Las tres proyecciones $\mathbf{Q}$, $\mathbf{K}$ y $\mathbf{V}$ tienen su argumento, la matriz de pesos $\mathbf{A}$ es cuadrada ($T \times T$) porque la frase se mira a sí misma, y el softmax por filas reparte una unidad entre las posiciones de cada fila. Pero la $\sqrt{d_k}$ del denominador se escribió en la fórmula sin explicar por qué está ahí, y sigue sin explicación. Es el único préstamo que sigue abierto, y esta lección lo salda.

La pregunta tiene más filo del que parece. ¿Por qué una raíz, y no el propio $d_k$? ¿Y por qué de $d_k$, y no de $d_{\text{model}}$ o de la longitud de la frase? Las dos se contestan de una vez, porque las dos dependen de averiguar cuánto se aleja del cero una suma cuando se le van añadiendo sumandos. Para hacerlo no hace falta teoría de la probabilidad compleja, sólo promediar, y la cuenta empieza promediando sobre una lista finita que se puede recorrer entera.

Toma la versión más cruda posible de una consulta (query) y una clave (key), es decir, de una fila de $\mathbf{Q}$ y una fila de $\mathbf{K}$: todas sus coordenadas valen $+1$ o $-1$. La puntuación es entonces una cuenta de coincidencias, porque cada término aporta $+1$ si las dos coordenadas comparten signo y $-1$ si no. Con cuatro coordenadas, tres coincidencias y una discrepancia dan $3 - 1 = 2$. Con cuatrocientas, la misma proporción (trescientas contra cien) da $200$.

Y ahí está lo que hay que ver, que no es que los términos se hayan hecho mayores: siguen valiendo $\pm 1$. Es que hay cien veces más. Una desviación pequeña respecto de la mitad exacta, repetida sobre más sumandos, se convierte en un número grande, y ese número se lo va a comer una exponencial.

Con signos al azar la proporción no se queda fija en tres contra uno, y la puntuación crece menos: su tamaño típico va con $\sqrt{d_k}$, como deduce la sección siguiente. La figura pone una misma fila de seis puntuaciones a ese tamaño en tres anchuras y prueba tres divisores antes del softmax: $1$ (es decir, no dividir), $\sqrt{d_k}$ y el propio $d_k$. Cada grupo de barras es el reparto de la fila, un peso por posición, y el número de encima es lo que se lleva la posición que gana. Busca la fila que no cambia de una columna a otra.

Una puntuación crece con la raíz de la anchura

Cada casilla de $\mathbf{Q}\mathbf{K}^{\top}$ es el producto escalar de una fila de $\mathbf{Q}$ (la consulta de una posición) con una fila de $\mathbf{K}$ (la clave de otra). Las dos viven en $\mathbb{R}^{d_k}$, y las escribo $\mathbf{q}$ y $\mathbf{k}$ sin índice de posición porque la cuenta es la misma para cualquier pareja. La puntuación es:

una suma de $d_k$ términos. De qué tamaño sale depende de qué tamaño entra, así que fijemos primero el caso exacto: las coordenadas valen $+1$ o $-1$, como en la cuenta de coincidencias de arriba. Llamemos $s_m = q_m k_m$ al término $m$-ésimo, que vale entonces $+1$ o $-1$, y $e = \sum_m s_m$.

Elevar al cuadrado separa la diagonal del resto:

porque $s_m^{2} = 1$ sea cual sea el signo. El primer sumando ya no depende de nada. El segundo es el que hay que promediar, y promediar aquí quiere decir algo perfectamente finito: recorrer las $2^{d_k}$ combinaciones de signos que puede tomar $\mathbf{q}$, con $\mathbf{k}$ fija, y hacer la media. La escribiremos $\mathbb{E}\left[\cdot\right]$, la media de una cantidad sobre el azar que se le supone. Los términos cruzados se van todos.

De modo que del cuadrado sobrevive el primer sumando y nada más:

Con $\mathbb{E}\left[e\right] = 0$, eso es cuanto hace falta. La puntuación se reparte alrededor del cero, y lo que mide su alejamiento típico (la desviación típica) es la raíz de ese promedio:

que es la respuesta a la mitad de la pregunta. Dividir entre $\sqrt{d_k}$ deja $\mathbb{E}\left[\left(e/\sqrt{d_k}\right)^{2}\right] = 1$, y en ese $1$ no aparece $d_k$: la misma escala con consultas de ocho coordenadas y con consultas de mil.

Con el divisor puesto, la operación entera (producto, división, softmax y mezcla) pasa a tener nombre propio: es el producto interno escalado, que es como el artículo escribe scaled dot-product. El producto interno es el producto escalar de siempre, el de la línea de arriba; lo escalado es esta división y nada más.

Los signos estaban ahí para que la cuenta fuese exacta, pero fíjate en qué propiedades le hicieron falta de verdad. Dos: que $\mathbb{E}\left[s_m\right] = 0$ y que $\mathbb{E}\left[s_m^{2}\right] = 1$. Cualquier inicialización centrada cuyas coordenadas sean independientes y de tamaño típico $1$ las cumple, y con ellas la misma cancelación deja el mismo $\mathbb{E}\left[e^{2}\right] = d_k$. Las gaussianas de rng.normal que el curso lleva usando desde la lección sobre implementar el MLP (multilayer perceptron)Bloque 2 · Lección 9Implementar un MLP desde ceroLas cuatro piezas del entrenamiento están desarrolladas y ninguna celda las ha ejecutado seguidas. Montarlas cuesta apilar la vuelta por filas (con el 1/B que trae la media), y tomar la decisión que ninguna fórmula toma: de dónde parten los pesos. Todos iguales se quedan iguales para siempre, y por eso la inicialización es al azar. son una de esas inicializaciones, y la primera celda lo mide con ellas.

Conviene decir en voz alta de qué habla esa hipótesis, porque no habla del modelo que te encontrarás. Habla de una red recién inicializada: nada garantiza que después de miles de pasos las coordenadas de $\mathbf{q}$ sigan siendo independientes ni midiendo $1$. Por eso el divisor se fija al diseñar la capa y no se mide sobre la marcha: arregla el arranque, que es donde el problema detiene el entrenamiento antes de que llegue a empezar.

Queda la parte que suele darse por sabida: por qué la raíz y no $d_k$. El divisor tiene que quitarle a la puntuación exactamente lo que la anchura le puso, y lo que le puso es un factor $\sqrt{d_k}$. Dividir entre $d_k$ dejaría las puntuaciones en un tamaño típico de $1/\sqrt{d_k}$, cada vez más cerca del cero cuanto más ancho el modelo; y unas puntuaciones que tienden a cero dan un softmax que tiende al reparto uniforme, $1/T$ en cada casilla, con lo que la fila deja de elegir nada. Son dos fallos, uno a cada lado, y la raíz es el único exponente que no cae en ninguno.

Una fila cerrada no devuelve gradiente

Si las puntuaciones miden $\sqrt{d_k}$, las diferencias entre las de una misma fila miden eso también, y lo que el softmax mira son las diferencias:

con $m$ recorriendo las posiciones de la fila $i$. Dos puntuaciones separadas por $8$ (una distancia corriente cuando el tamaño típico es $8$, que es lo que da $d_k = 64$) quedan a un factor $\exp(8) \approx 2\,981$ la una de la otra, y con eso el reparto se ha terminado: una posición se lleva casi la unidad entera.

Que una fila elija fuerte no suena a defecto. El defecto está en la derivada.

Deriva $\alpha_{ij}$ respecto de las puntuaciones de su propia fila. Es la misma cuenta que la lección sobre las funciones de pérdidaBloque 2 · Lección 5Funciones de pérdida: MSE y entropía cruzadaUna pérdida resume en un número lo mal que lo hace la red, y entre dos candidatas decide su derivada. El error cuadrático medio entrega su corrección más pequeña justo donde la red se equivoca con más seguridad; la entropía cruzada, que sale de leer la salida como probabilidad, deja esa derivada en la diferencia entre lo predicho y la etiqueta. hizo con el softmax de la capa de salida, y salen los mismos dos casos según el índice coincida o no:

Pon ahora la fila cerrada, con un peso en casi $1$ y el resto en casi $0$. El caso de la diagonal da $1$ por $0$; el cruzado, en cualquier casilla, arrastra al menos un factor casi nulo. Las dos derivadas se apagan a la vez, y con ellas la fila entera.

Lo que eso cuesta se ve mirando qué hay antes de las puntuaciones. Sólo $\mathbf{W}^Q$ y $\mathbf{W}^K$: las puntuaciones salen de $\mathbf{Q}\mathbf{K}^{\top}$ y no entra nada más. Si el gradiente no atraviesa el softmax, esas dos matrices no reciben nada por esa fila, y son precisamente las que tendrían que aprender a quién debe mirar cada posición. La capa sigue calculando (una salida saturada es una salida perfectamente definida); lo que ha dejado de poder es cambiar de opinión.

Merece la pena pararse en por qué el bloque 2 no se topó con esto, porque la diferencia es de una línea. Allí el softmax iba seguido de la entropía cruzada, y el $-\log$ de la pérdida aporta el factor inverso exacto: las dos juntas dieron $\nabla_{\mathbf{z}}\,\ell = \hat{\mathbf{y}} - \mathbf{y}$, que vale $\hat{y}_c - 1$ en la clase correcta y no se anula por muy segura que esté la red. Aquí detrás del softmax no hay ningún logaritmo, sólo la multiplicación por $\mathbf{V}$, la matriz de valores (values). La jacobiana no se cancela contra nada: se queda entera, ceros incluidos.

Y así reaparecen dos cosas que el curso dio por cerradas en bloques distintos. La curva $\alpha\left(1 - \alpha\right)$ es la misma que la derivada de la sigmoide, con su máximo de $0.25$ en la mitad y sus dos colas planas: la saturación de la lección sobre las funciones de activaciónBloque 2 · Lección 2Funciones de activación y no linealidadEl escalón deja pasar el signo de la preactivación y tira lo demás, y su derivada vale cero allí donde existe. Sigmoide, tangente hiperbólica y ReLU devuelven la magnitud y admiten derivada; ninguna puede ser una recta, porque encadenar funciones afines vuelve a dar una función afín., en una capa donde no hay ninguna sigmoide. Y un gradiente que llega apagado a las matrices que tenían que aprender es el bloque 3 entero, en una arquitectura sin una sola recurrencia. Lo que las une no era la sigmoide ni el estado recurrente: era que en las dos había una función aplanándose.

Una fila entera se resume en un número, y sale de sumar el caso de la diagonal sobre las $T$ posiciones:

donde la última igualdad usa que la fila suma $1$. Con la fila repartida por igual vale $1 - 1/T$; con la fila entera en una posición vale $0$. Es la suma de la diagonal de la jacobiana, y mide lo mismo por los dos lados: cuánto reparte esa fila, y cuánto le queda por aprender.

El divisor, medido en NumPy

La primera celda saca la cuenta del caso de los signos y la comprueba donde de verdad se aplica: coordenadas gaussianas, cinco anchuras, la columna del promedio medido contra la columna de $\sqrt{d_k}$.

Las dos columnas centrales coinciden hasta donde llega la precisión de medir con ocho mil parejas, y la de la derecha se queda clavada en $1$: dividir deja el tamaño donde estaba, sea cual sea la anchura. La columna de la media es pequeña frente al tamaño típico de su fila (a $d_k = 1\,024$, $0.236$ contra $32.03$), que es lo que dice $\mathbb{E}\left[e\right] = 0$ cuando se mide en vez de demostrarse.

La segunda celda pasa de las puntuaciones a lo que sale de ellas. Mira sobre todo la última columna, que es la que decide si la capa puede aprender.

Léela por parejas de líneas. Sin dividir, el peso mayor sube de $0.670$ a $0.953$ y el reparto $1 - \sum_j \alpha_{ij}^{2}$ se hunde de $0.445$ a $0.066$ (el máximo posible con ocho posiciones es $1 - 1/8 = 0.875$), mientras la mayor casilla de la diagonal de la jacobiana cae de $0.1755$ a $0.0321$: más de cinco veces menos gradiente por el mero hecho de haber ensanchado el modelo. Dividiendo, las tres columnas dejan de moverse: $0.36$ de peso mayor, $0.77$ de reparto y $0.216$ de derivada en las tres anchuras, incluida la de $512$. Ésa es la propiedad que se compra con el divisor, y no es que los números sean mejores, es que dejan de depender de $d_k$.

Comprueba tu intuición

Cinco preguntas: cuánto mide una puntuación, por qué la raíz y no otra cosa, qué pierde una fila cerrada, por qué el bloque 2 se libró de esto y qué es lo que dividir no toca.

Con el divisor puesto, la escala de las puntuaciones deja de depender de la anchura, y el reparto y el gradiente dejan de depender de ella con la misma cuenta. Pero mira lo que el divisor da por hecho: recibe $d_k$ ya decidido. Nada de lo que hemos contado dice cuánto debería medir una consulta (sólo cómo compensar lo que mida), y hasta aquí $d_k$ es un número que apareció porque las proyecciones de la lección anterior tenían que salir a alguna anchura.

Hay una segunda cosa que esta capa no sabe hacer, y apunta al mismo sitio. Una fila de $\mathbf{A}$ es un reparto y sólo uno: en las llaves del coche están ahí, la posición de están tiene que decidir con un único juego de pesos si mira a llaves, que es con quien concuerda, o a del coche, que es lo que la separa de ella; y con una fila no puede hacer las dos cosas. La respuesta a las dos preguntas es la misma y ocupa la lección siguiente, sobre la atención multi-head: varias parejas de proyecciones trabajando en paralelo, cada una con su propio reparto y su propia anchura, repartida entre todas a partir de $d_{\text{model}}$. Ahí $d_k$ deja de ser un número dado y pasa a ser una decisión con un motivo.