Cada posición, un puñado de relojes

Queda una casilla por rellenar y no está en la capa, sino en lo que se le entrega. La atención multi-head cerró todo lo demás —$h$ maneras de mirar la frase, concatenadas y proyectadas de vuelta al precio de una—, y dejó dicho que ninguna cantidad de cabezas consigue que $\mathbf{X}$ tenga una fila primera y una última: eso hay que escribirlo en cada fila antes de que la capa la lea. De qué se escribe ahí trata esta lección.

Las respuestas fáciles fallan de tres maneras, y conviene mirarlas porque entre las tres describen lo que hace falta. Escribir el número de la posición en una coordenada rompe la escala: en la posición $300$ esa coordenada vale $300$ mientras el resto del vector se mueve entre $-1$ y $1$, y se lleva por delante cualquier producto escalar en el que entre. Dividirlo entre la longitud de la frase, $\text{pos}/T$, arregla la escala y estropea otra cosa: el $0.5$ de una frase de ocho tokens señalaría el cuarto token y el de un texto de trescientos señalaría el ciento cincuenta, y sería la misma coordenada con el mismo valor. Y aprender un vector por posición, una tabla como la de los embeddings, funciona —hay modelos que lo hacen— a cambio de que la tabla tenga exactamente tantas filas como posiciones se hayan entrenado, y de que la posición $7$ y la $8$ sean dos filas sin ninguna relación entre ellas: nada dice que sean vecinas.

Ese último defecto es el que apunta a lo que se quiere de verdad. Piensa en un cuentakilómetros: la rueda de la derecha gira deprisa y distingue perfectamente dos posiciones seguidas, pero se repite cada diez; la de la izquierda casi no se mueve, así que no distingue vecinas, pero sí el principio del recorrido del final. Ninguna rueda sola sirve y todas juntas dan un número que no se repite —y donde, además, dos lecturas separadas por la misma cantidad se parecen igual en cualquier tramo, porque cada rueda ha avanzado lo mismo—. Los senos y cosenos son esa idea sin escalones: una pareja de coordenadas por «rueda», cada pareja girando a su propia velocidad.

Una escalera de longitudes de onda

El artículo define la codificación posicional entrada por entrada, con $\text{pos}$ como índice de la posición —es el $t$ del resto del curso, escrito con la letra del artículo— y con $\text{PE}$ de positional encoding:

para $i = 0, 1, \dots, d_{\text{model}}/2 - 1$. Léelo despacio, porque la estructura está en los índices: las coordenadas van por parejas, la pareja $i$ ocupa las coordenadas $2i$ y $2i+1$, y las dos comparten el mismo argumento. Llamemos $\omega_i$ a ese divisor invertido, que es lo que la pareja gira por cada posición que se avanza:

A la derecha está la longitud de onda: cuántas posiciones tarda esa pareja en volver al mismo sitio. La pareja $0$ tiene $\omega_0 = 1$ y vuelve cada $2\pi \approx 6.28$ posiciones; la última se queda justo por debajo de $10\,000 \cdot 2\pi$, unas $63\,000$ posiciones, que sobre cualquier texto es una franja que apenas se mueve. Entre las dos, las longitudes de onda crecen multiplicándose siempre por el mismo factor: son una progresión geométrica, y eso es la escalera que el explorable dibuja de izquierda a derecha.

Lo que sale de ahí es una matriz $\text{PE} \in \mathbb{R}^{T \times d_{\text{model}}}$ cuya fila $\text{pos}$ es $\text{PE}_{\text{pos}}^{\top}$, exactamente la forma de la $\mathbf{X}$ de las tres lecciones anteriores. Y se junta con ella sumando:

Fíjate en que la primera pega ha desaparecido sola: seno y coseno viven en $[-1, 1]$ pase lo que pase con $T$, así que la codificación de la posición $3$ y la de la posición $3\,000$ son del mismo tamaño. La segunda también, y por una razón distinta: $T$ no aparece en ninguna parte de la fórmula, de modo que la posición $5$ se codifica igual en una frase corta que en un texto largo.

Sumar en lugar de concatenar es una decisión, y merece decirse en voz alta lo que cuesta. Concatenar mantendría separadas por construcción las coordenadas del significado y las de la posición, y se pagaría en anchura: reservar $64$ de las $512$ para la posición deja $448$ para todo lo demás. Sumando no se paga anchura ninguna, y a cambio las dos señales comparten las mismas coordenadas y nada garantiza que se puedan volver a separar. Lo que hay a favor es que $\mathbf{W}^Q$ y $\mathbf{W}^K$ se aprenden y pueden quedarse con unas direcciones y no con otras, y que la codificación tiene una estructura muy suya que un embedding entrenado no tiene por qué imitar. Plausible, entonces; no demostrado.

Desplazarse una distancia fija es girar un ángulo fijo

De las tres pegas del principio queda la tercera, que era la que de verdad pedía algo: que dos posiciones vecinas se parezcan, y que se parezcan igual estén donde estén. Aquí se puede demostrar, y ocupa tres líneas.

Fija un desplazamiento $k$ y quédate con una sola pareja. Las fórmulas de la suma de ángulos dicen qué le pasa a esa pareja al avanzar $k$ posiciones:

—y las he escrito con los factores en ese orden a propósito, porque así las dos líneas son ya un producto de matrices. Las coordenadas nuevas son una combinación de las dos viejas, y los coeficientes de esa combinación son cuatro números que sólo miran a $\omega_i$ y a $k$:

Mira lo que no está en esa matriz: $\text{pos}$. La posición de partida se ha quedado entera dentro del vector de la derecha, y el desarrollo la ha dejado ahí. Como cada pareja se transforma sin tocar a las demás, apilar las $d_{\text{model}}/2$ matrices de $2 \times 2$ a lo largo de la diagonal da una sola matriz para el vector entero:

Esa es la propiedad, y conviene leer bien lo que dice: hay una matriz por cada distancia $k$, y esa matriz vale para todas las posiciones de la frase. Avanzar tres posiciones es aplicar $\mathbf{M}_3$, se avance desde la primera palabra o desde la última. Cada bloque es además un giro —sus dos columnas son perpendiculares y de longitud $1$—, así que $\mathbf{M}_k$ no estira ni encoge nada: todas las filas de $\text{PE}$ miden lo mismo.

Lo que de todo esto puede ver una puntuación

$\mathbf{M}_k$ es una propiedad de la codificación, y la capa de atención no mira matrices: mira productos escalares. Merece la pena hacer uno y ver qué queda. Con $\cos(A - B) = \cos A \cos B + \sin A \sin B$ aplicado pareja a pareja:

$\text{pos}$ ha vuelto a desaparecer. Dos posiciones separadas por $k$ puntúan entre ellas el mismo número al principio de la frase y en la mitad, y ese número sólo depende de la distancia. Con $k = 0$ cada término vale $1$ y la suma da $d_{\text{model}}/2$, que es lo máximo: ninguna posición puntúa contra otra tanto como contra sí misma.

Y ahora la concesión, porque esto es menos de lo que parece. La capa no calcula $\text{PE}_{\text{pos}}^{\top}\text{PE}_{\text{pos}+k}$: calcula $\mathbf{q}^{\top}\mathbf{k}$ con las dos proyectadas desde $\mathbf{e}_{w_t} + \text{PE}_t$, así que al desarrollar el producto salen cuatro términos y sólo uno de ellos es puramente posicional. Lo que la forma sinusoidal garantiza es que ese término existe y que depende sólo de la distancia; una cabeza que quiera aprender «el determinante que tengo justo delante» —el la de la casa amarilla— tiene con qué construirlo, porque la señal de «una posición atrás» es la misma en toda la frase. Si lo construye o no lo decide el entrenamiento, y no esta lección.

La escalera en NumPy

La primera celda monta $\text{PE}$ con dos rodajas y sin un solo bucle, y después imprime la longitud de onda de unas cuantas parejas.

La fila $0$ sale $(0, 1, 0, 1, \dots)$, que es lo que tiene que salir: en la posición $0$ todos los ángulos valen $0$. Las tres primeras filas de la tabla son la escalera: la pareja $0$ vuelve cada $6.3$ posiciones, la $1$ cada $11.2$, la $2$ cada $19.9$ —cada una $1.78$ veces la anterior— y la $15$, la última de este $d_{\text{model}}$, cada $35\,332.9$, que sobre $24$ posiciones es una recta. La última línea es la observación del final de la sección anterior comprobada: todas las filas miden $4$, que es $\sqrt{32/2}$.

La segunda celda comprueba las dos propiedades, que es donde está el contenido de la lección.

La primera línea es $\mathbf{M}_k$ movida a través de las veinticuatro posiciones de una vez, y el error máximo se queda en el orden de $10^{-15}$: la misma matriz, todas las filas. La segunda es la otra propiedad —cuatro posiciones de partida y cuatro veces $12.272398$, que es exactamente la suma de cosenos— y la tabla del final dice cómo se comporta esa cantidad al alejarse: $16$ contra sí misma, $15.314$ a un paso, $12.272$ a tres, y luego baja despacio y sin ser monótona, porque es una suma de cosenos y no una distancia.

Comprueba tu intuición

Cinco preguntas: por qué no vale el número de la posición, de qué depende $\mathbf{M}_k$, cuánto puntúa una posición contra sí misma, qué se rompe al quitar los cosenos, y qué garantiza sumar en vez de concatenar.

Con esto la entrada está completa: cada fila de $\mathbf{X}$ trae lo que dice el token y dónde estaba, y la capa de atención de las cuatro lecciones anteriores ya se puede calcular entera sobre ella. Lo que queda es apilarla, que es lo que convierte una capa en un modelo. Y ahí aparece algo que esta lección deja pendiente sin decirlo: la posición se suma una vez, abajo del todo, y tiene que llegar viva hasta la última capa.

Lo que el artículo dibuja alrededor de la atención —una conexión que se salta la capa entera, una normalización, un perceptrón por posiciones— parece fontanería alrededor de lo importante y no lo es. Es lo que decide si una señal metida abajo sigue ahí arriba, y el bloque 3 ya midió por qué al seguir un gradiente a lo largo de una cadena de productos y verlo desvanecerse. Eso es la lección siguiente, sobre el bloque completo con residuales y layer norm.