Varias cabezas, varios repartos

La lección anterior dejó dos cosas abiertas. Una capa de auto-atención le da a cada posición una sola fila de $\mathbf{A}$, fabricada por un único par de matrices, $\mathbf{W}^Q$ y $\mathbf{W}^K$; y $d_k$ entra en la fórmula sin que nada diga cuánto debería valer.

Mira la frase el ratón pequeño que persiguen los gatos duerme. duerme tiene que encontrar a su sujeto, ratón, y no le basta un criterio: «¿es un sustantivo?» no distingue ratón de gatos, y «¿concuerda conmigo en número?» no distingue un sujeto de un determinante. Llamaré regla a cada criterio de ésos. La fila de duerme necesita las dos a la vez, y las combina sumándolas: cada regla le da a cada palabra una puntuación, el producto escalar las suma y el softmax reparte la fila según el total. Gana la palabra que cumple las dos.

Sumar sirve cuando se busca una sola palabra que cumpla todas las reglas, pero una posición puede necesitar dos palabras distintas. persiguen necesita a quien persigue, gatos, y a quien es perseguido, ratón, y su fila es un único reparto que suma $1$: o se queda con una y pierde la otra, o mezcla las dos. Además, una vez sumadas, las reglas ya no se pueden separar: la fila no dice cuánto puso cada una.

La alternativa es no sumar antes de repartir: separar las reglas en grupos, cada uno con su propia atención y su propio softmax, y juntar después lo que lee cada uno. Cada una de esas atenciones es una cabeza, y la capa que monta varias es la atención multi-head, de la que trata esta lección.

El explorable de abajo pone una sola regla en cada cabeza: la 1 es la del sustantivo, la 2 la concordancia, y la 3 y la 4 relacionan sustantivos con sus determinantes y modificadores. El quinto botón no junta las cabezas: es la capa de la lección sobre la auto-atenciónBlock 5 · Lesson 2 · In SpanishAuto-atención: la secuencia se mira a sí mismaSin encoder ni decoder recurrentes no hay quien fabrique las consultas, las claves ni los valores, y la salida es que los fabrique la propia frase. Pero usar sus vectores tal cual en los tres papeles no vale: lo que más se parece a un vector es él mismo, la diagonal gana todas las filas y la capa devuelve lo que le entró, sin un solo parámetro que entrenar. Lo que arregla eso son tres proyecciones de la misma X (una para preguntar, otra para ser preguntado, otra para mezclarse), y con ellas la rejilla se vuelve cuadrada, deja de ser simétrica y cuesta un número de parámetros donde la longitud de la frase no aparece., con las cuatro reglas sumadas antes de un único softmax. La fila es quien pregunta y la columna quien contesta, y el panel de debajo resume, para la fila que elijas, a qué palabra tira cada mapa.

En la fila de duerme, la cabeza 1 empata ratón con gatos, a $0.30$ cada uno; la 2 prefiere tres singulares que no son el sujeto, a $0.26$; y la 3 y la 4 no dicen nada. Ninguna cabeza encuentra el sujeto, porque las dos reglas que hacían falta han ido a cabezas distintas. La suma sí lo encuentra, ratón con $0.31$, pero a cambio borra lo que dijo cada regla.

En la fila de persiguen, que necesitaba a gatos y a ratón, la suma se queda con gatos, con $0.36$, y a ratón le deja $0.07$. La cabeza 1, en cambio, le sigue dando $0.30$, así que con varias cabezas lo que se lee de ratón llega a la salida. Esa cabeza no sabe que ratón es el perseguido, porque ninguna de estas reglas busca objetos y ella lee los dos sustantivos por igual; lo que la separación evita es que la suma lo borre. Por eso una cabeza de verdad lleva varias reglas: dentro de ella se suman las que tienen que cumplirse a la vez, y en cabezas distintas pueden acabar las que buscan cosas distintas.

La misma matriz leída por varias atenciones a la vez

Cada cabeza es la capa de la lección sobre la auto-atenciónBlock 5 · Lesson 2 · In SpanishAuto-atención: la secuencia se mira a sí mismaSin encoder ni decoder recurrentes no hay quien fabrique las consultas, las claves ni los valores, y la salida es que los fabrique la propia frase. Pero usar sus vectores tal cual en los tres papeles no vale: lo que más se parece a un vector es él mismo, la diagonal gana todas las filas y la capa devuelve lo que le entró, sin un solo parámetro que entrenar. Lo que arregla eso son tres proyecciones de la misma X (una para preguntar, otra para ser preguntado, otra para mezclarse), y con ellas la rejilla se vuelve cuadrada, deja de ser simétrica y cuesta un número de parámetros donde la longitud de la frase no aparece., entera: sus tres proyecciones sobre la misma matriz y su propio mapa. Lo único que hay que añadir es el subíndice que allí no hizo falta escribir porque sólo había una.

con $i = 1, \dots, h$, y con $\mathbf{W}^Q_i, \mathbf{W}^K_i \in \mathbb{R}^{d_{\text{model}} \times d_k}$ y $\mathbf{W}^V_i \in \mathbb{R}^{d_{\text{model}} \times d_v}$. Fíjate en que las $h$ leen la misma $\mathbf{X}$: no se reparten la frase, ni se reparten las coordenadas de entrada. Cada una saca de ella sus propias consultas (queries), claves (keys) y valores (values), y con ellas su propio mapa $\mathbf{A}_i \in \mathbb{R}^{T \times T}$, cuadrado y con las filas sumando $1$ igual que el de la lección sobre la auto-atenciónBlock 5 · Lesson 2 · In SpanishAuto-atención: la secuencia se mira a sí mismaSin encoder ni decoder recurrentes no hay quien fabrique las consultas, las claves ni los valores, y la salida es que los fabrique la propia frase. Pero usar sus vectores tal cual en los tres papeles no vale: lo que más se parece a un vector es él mismo, la diagonal gana todas las filas y la capa devuelve lo que le entró, sin un solo parámetro que entrenar. Lo que arregla eso son tres proyecciones de la misma X (una para preguntar, otra para ser preguntado, otra para mezclarse), y con ellas la rejilla se vuelve cuadrada, deja de ser simétrica y cuesta un número de parámetros donde la longitud de la frase no aparece..

Quedan entonces $h$ salidas de $T \times d_v$ y la capa tiene que devolver una. Se ponen una al lado de otra y se proyectan:

donde $\text{Concat}$ pega las $h$ matrices por columnas (una fila por posición, con los $h$ trozos seguidos dentro de ella, de modo que el resultado es $T \times \left(h \cdot d_v\right)$), y $\mathbf{W}^O \in \mathbb{R}^{\left(h \cdot d_v\right) \times d_{\text{model}}}$ es la única matriz de la capa que no pertenece a ninguna cabeza. Su superíndice es una etiqueta de papel, como los de las otras tres.

Sumar antes o después del softmax

La puntuación de una cabeza ya era una suma, y la lección anterior la escribió así precisamente para poder contarla:

donde $q_{im}$ es la coordenada $m$ de la consulta de la posición $i$. Cada coordenada de la comparación aporta su término, y el softmax no ve los términos: ve el total.

Las reglas de la introducción son esos términos. El explorable de arriba está hecho así a propósito: cada regla es una columna de $\mathbf{W}^Q$ y otra de $\mathbf{W}^K$, de modo que la coordenada $m$ de la consulta por la coordenada $m$ de la clave es la puntuación de la regla $m$. La capa de una sola cabeza lleva las cuatro columnas, así que compara en $d_k = 4$ coordenadas (y divide su suma entre $\sqrt{4} = 2$); cada una de las cuatro cabezas se queda con una sola columna, $d_k = 1$, y su puntuación es una regla y nada más. En una capa entrenada nadie escribe las reglas y una coordenada no tiene por qué significar nada por sí sola, pero la cuenta es la misma: términos que llegan sumados al softmax.

Sumar dentro del softmax tiene una consecuencia que cabe en una línea. Si $a_j$ y $b_j$ son lo que dos reglas le dan a la palabra $j$, la exponencial convierte su suma en un producto:

Repartir el total es multiplicar, palabra a palabra, lo que pondría cada regla por su cuenta, y reescalar para que la fila vuelva a sumar $1$. Un producto funciona como un «y»: si una regla puntúa bajo a una palabra, ésta se queda sin peso, diga lo que diga la otra. Es lo que le pasa a ratón en la fila de persiguen: la concordancia lo hunde, porque no coincide en número, y la regla del sustantivo, que lo puntuaba igual que a gatos, no puede sacarlo de ahí. Y como lo que sale es un solo reparto, ya no se puede saber cuánto puso cada regla.

Ahí está la operación entera. Separar las coordenadas en $h$ grupos y darle a cada grupo su propio softmax es repartir antes de combinar, en vez de combinar antes de repartir. Salen $h$ distribuciones en lugar de una, cada una con su denominador, y llegan hasta la salida sin tocarse; lo que las combina es $\mathbf{W}^O$, que es una matriz que se aprende y que las recibe ya repartidas. La suma que la capa de una cabeza hacía dentro del softmax era irreversible; hecha fuera, todavía se puede decidir cuánto pesa cada parte.

El reparto de la anchura

Queda elegir $d_k$, que era la otra pregunta abierta, y se elige contando. Cada cabeza cuesta lo que costaba la capa entera de la lección sobre la auto-atenciónBlock 5 · Lesson 2 · In SpanishAuto-atención: la secuencia se mira a sí mismaSin encoder ni decoder recurrentes no hay quien fabrique las consultas, las claves ni los valores, y la salida es que los fabrique la propia frase. Pero usar sus vectores tal cual en los tres papeles no vale: lo que más se parece a un vector es él mismo, la diagonal gana todas las filas y la capa devuelve lo que le entró, sin un solo parámetro que entrenar. Lo que arregla eso son tres proyecciones de la misma X (una para preguntar, otra para ser preguntado, otra para mezclarse), y con ellas la rejilla se vuelve cuadrada, deja de ser simétrica y cuesta un número de parámetros donde la longitud de la frase no aparece.:

hay $h$ de ellas, y encima está $\mathbf{W}^O$ con $h \cdot d_v \cdot d_{\text{model}}$ más, de modo que la capa completa pide

Asigna ahora a $d_k$ y a $d_v$ el valor $d_{\text{model}}/h$ (repartir la anchura del modelo entre las cabezas, en vez de darle a cada una la anchura entera) y mira lo que ocurre:

La $h$ se ha ido. Ocho cabezas cuestan lo mismo que una, y dieciséis lo mismo que ocho. Las multiplicaciones se comportan igual: la lección sobre el adiós a la recurrenciaBlock 5 · Lesson 1 · In SpanishAdiós a la recurrenciaLa atención ya no necesita a la recurrencia, pero la máquina que la rodea sigue siendo recurrente, y eso fija dos cosas que ningún ordenador arregla: una frase de T tokens obliga a T pasos que no se pueden solapar, y lo que la posición 1 le aporta a la posición T pasa por T menos 1 productos por la misma matriz. Quitar la recurrencia deja las dos en un número que no crece con la frase, a cambio de una rejilla de T por T puntuaciones que sale más barata que la recurrencia mientras la frase mida menos que el modelo. Y a cambio de otra cosa que no se paga en multiplicaciones: el orden. midió una atención en $T^{2}\left(d_k + d_v\right)$, y $h$ de ellas dan $h \cdot T^{2}\left(d_k + d_v\right) = 2T^{2}d_{\text{model}}$, que es exactamente lo que costaba una sola de anchura $d_{\text{model}}$.

De modo que $h$ se elige por lo que hace y no por lo que cuesta, y $d_k$ deja de ser un número suelto: es lo que queda de $d_{\text{model}}$ al partirlo. Con las cifras del artículo, $d_{\text{model}} = 512$ y $h = 8$, sale $d_k = 64$ (el número con el que la lección anterior hizo sus cuentas) y la capa entera pesa $4 \cdot 512^{2} = 1\,048\,576$ parámetros.

Lo que se paga conviene decirlo, porque no es cero. Cada cabeza compara en $64$ coordenadas donde una sola habría comparado en $512$: no son ocho comparaciones tan finas como la que había, son ocho comparaciones ocho veces más estrechas. El trato es ése, y es una apuesta sobre el lenguaje (que compensan varias relaciones gruesas frente a una sola muy fina) y no un teorema.

Concatenar y proyectar, no promediar

La alternativa evidente a concatenar es promediar: dale a cada cabeza la anchura entera, $d_v = d_{\text{model}}$, y devuelve $\frac{1}{h}\sum_i \text{head}_i$, que ya tiene la forma correcta sin necesidad de ninguna matriz más. Merece la pena verla como lo que es, un caso particular de lo que la capa ya hace. Toma la $\mathbf{W}^O$ que apila $h$ copias de $\frac{1}{h}\mathbf{I}$:

Concatenar con una $\mathbf{W}^O$ que se aprende puede, entonces, dar el promedio y cualquier otra cosa; promediar es fijar de antemano un punto de todo eso. Y fijarlo cuesta algo muy concreto: obliga a las $h$ cabezas a escribir en las mismas $d_{\text{model}}$ coordenadas, de manera que dos cabezas que han encontrado dos cosas distintas pueden cancelarse antes de que nadie las lea. Concatenadas, cada una tiene sus $d_v$ coordenadas para ella sola, y $\mathbf{W}^O$ decide después, coordenada de salida por coordenada de salida, a cuál hace caso.

Esa matriz hace además una segunda cosa que aquí sólo queda señalada: devuelve la salida a $d_{\text{model}}$, la anchura con la que entró. Eso es lo que permite apilar estas capas una sobre otra, y quien lo necesita de verdad es una lección más adelante, sobre el bloque completo con residuales y layer normComing up · Block 5 · Lesson 6 · In SpanishApilar sin perder lo de abajoLa capa de atención ya está completa, y el artículo pone seis una encima de otra. Componer seis funciones multiplica seis derivadas, que es exactamente la cuenta con la que se estrelló el bloque 3, así que apilar no es repetir: hace falta algo alrededor de cada subcapa. Son dos cosas y las dos caben en una línea. La suma —la subcapa devuelve lo que ha calculado más lo que le entró— abre un camino de la salida a la entrada que no atraviesa ninguna matriz, y de paso obliga a que d_model sea un solo número en todo el modelo, que es para lo que estaba la proyección de salida de multi-head. El layer norm centra cada fila con sus propias coordenadas, sin mirar a las demás posiciones ni a las demás frases del batch. Y entre las dos subcapas del bloque está el perceptrón por posiciones, que parece fontanería y lleva el doble de parámetros que la atención..

Falta decir lo que esta construcción no hace, porque los dibujos de los artículos invitan a creer lo contrario. Nada en la pérdida le asigna un papel a una cabeza: no existe ningún término que diga «la cabeza 2 se ocupa de la concordancia». Las cabezas hacen cosas distintas porque empiezan distintas, y por nada más. Los nombres del explorable de arriba se los he puesto yo, escribiendo las cuatro reglas a mano; en una capa entrenada no los pone nadie y hay que salir a buscarlos después, cuando están. Ni siquiera el orden significa nada: intercambia dos cabezas junto con sus dos bloques de filas de $\mathbf{W}^O$ y no cambia un número de la salida, así que «la cabeza 3» no nombra lo mismo en dos entrenamientos del mismo modelo.

Las cabezas en NumPy

La primera celda monta la capa con las proyecciones enteras y un bloque de columnas por cabeza, que es como se implementa de verdad, y después pasa la cuenta de parámetros por cinco valores de $h$.

Las dos filas impresas son la misma posición vista por dos cabezas, y no se parecen en nada: la primera pone $0.591$ sobre la posición $1$ y la segunda $0.531$ sobre la $5$, con la que la primera gastaba $0.005$. Aquí las cabezas se han separado sólo porque rng les dio matrices distintas, que es todo lo que las separa también en un modelo entrenado. Y la tabla es la cancelación de arriba comprobada: las tres columnas de la derecha no se mueven de $786\,432$, $262\,144$ y $1\,048\,576$ mientras $d_k$ baja de $512$ a $32$.

La segunda celda mide la concesión. Si las $h$ cabezas empiezan con los mismos parámetros, no hay $h$ cabezas.

Con los cuatro bloques de columnas iguales las cuatro distancias salen exactamente $0$: los cuatro mapas son el mismo, y $\mathbf{W}^O$ recibe cuatro copias de un mismo vector por posición. Mover un solo bloque separa a esa cabeza de las otras tres ($0.458$ en las tres casillas de la primera línea) y deja a esas tres idénticas entre ellas, que es lo que dicen los tres ceros de la segunda. La asimetría no aparece sola: entra por la inicialización, y en un entrenamiento se mantiene porque las cabezas y los bloques de filas de $\mathbf{W}^O$ empiezan todos distintos. Con todo simétrico, el gradiente no tendría por dónde romperlo.

Comprueba tu intuición

Cinco preguntas: qué separa a $h$ cabezas de una sola, cuánto pesa la capa, de dónde sale $d_k$, qué calcula la capa con las cabezas iguales, y si promediar era una opción.

Con $h$ cabezas, una posición tiene $h$ respuestas donde tenía una, y $\mathbf{W}^O$ decide qué hacer con ellas. Pero mira de dónde salen las $h$: de la misma $\mathbf{X}$, con la misma fórmula, y en $e_{ij}$ siguen sin aparecer $i$ ni $j$ por ningún sitio. Multiplicar las cabezas multiplica qué se compara, y no añade ni una coordenada de dónde. Ninguna de las $h$ puede aprender «el sustantivo que tengo a la izquierda», por muchas que pongas, porque en lo que la capa lee no hay izquierda.

La pieza que falta no está en la capa: está en la entrada, y consiste en que dos posiciones distintas le lleguen a $\mathbf{X}$ como dos filas distintas aunque el token sea el mismo. Eso es la lección siguiente, sobre la codificación posicional, y lo que la hace interesante no es que haga falta (eso ya se sabía desde la primera lección de este bloque), sino la forma concreta que elige el artículo: unos senos y unos cosenos con los que la distancia entre dos posiciones se puede escribir como una operación lineal sobre lo que se ha añadido. Esa propiedad se enseña, no se afirma, y ocupa la lección entera.