Varias cabezas, varios repartos

La lección anterior justificó el $\sqrt{d_k}$ de la fórmula de la atención contando cuánto crece un producto escalar cuando se le añaden coordenadas, y al cerrar esa cuenta dejó dos cosas señaladas. La primera es de forma: una capa de auto-atención no le da a cada posición más que una fila de $\mathbf{A}$, fabricada por un único par de matrices, $\mathbf{W}^Q$ y $\mathbf{W}^K$ —una manera de mirar la frase, la misma para todas las posiciones y para todas las relaciones que haya que encontrar entre ellas—. La segunda es más pequeña y sale del mismo sitio: $d_k$ entra en la fórmula sin que nada diga cuánto debería valer.

Mira lo que esa única fila le pide a una frase como el ratón pequeño que persiguen los gatos duerme. duerme tiene que encontrar a ratón por encima de la subordinada entera, y persiguen, tres posiciones antes, tiene que encontrar a gatos, que le queda detrás: dos verbos, dos sujetos, en direcciones opuestas. Ninguna relación sola los resuelve —quién es sustantivo no distingue ratón de gatos, y quién concuerda en número no distingue un sujeto de un determinante—, así que la única fila que hay tiene que llevar las dos a la vez. Y puede: las puntuaciones se suman antes del softmax, y la suma decide. Esta lección trata de lo que esa suma cuesta, y de la alternativa, que consiste en no hacerla.

Las rejillas de abajo son esa misma frase vista por cuatro reglas distintas, una por cabeza, y en el quinto botón por la capa de la lección anterior, que lleva las cuatro sumadas en una sola fila. Ponte en la fila de duerme y lee el panel de abajo entero antes de tocar nada: la primera cabeza empata ratón con gatos a $0.30$; la segunda descarta gatos —le deja $0.01$— pero reparte su peso mayor entre tres singulares y ninguno de los tres es el sujeto; dos no dicen nada de esa fila; y la capa de una sola cabeza contesta ratón. Pásate luego a la fila de persiguen: la primera cabeza repite su fila anterior número a número, porque no sabe distinguir un verbo de otro; la segunda se da la vuelta entera; y la respuesta de la capa de una cabeza cambia a gatos.

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 2Auto-atención: la secuencia se mira a sí mismaSin decoder no hay quien fabrique las consultas, y la salida es que las fabrique la propia frase. Pero usar los vectores tal cual no vale: con las tres listas iguales, 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 pesos 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 2Auto-atención: la secuencia se mira a sí mismaSin decoder no hay quien fabrique las consultas, y la salida es que las fabrique la propia frase. Pero usar los vectores tal cual no vale: con las tres listas iguales, 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 pesos 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.

Dónde se normaliza

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. Dos coordenadas que midan dos cosas distintas llegan sumadas a la normalización, y lo que salga de esa suma es lo único que la fila tiene para repartir. Un término grande se lleva la fila; lo que otra coordenada tenía que decir de esa misma pareja desaparece, y no queda dónde ir a buscarlo.

El explorable de arriba es ese caso con $d_k = 4$ y una regla por coordenada, así que la cuenta se puede seguir número a número. Vuelve a las llaves del coche están ahí, la frase de la lección sobre la auto-atenciónBlock 5 · Lesson 2Auto-atención: la secuencia se mira a sí mismaSin decoder no hay quien fabrique las consultas, y la salida es que las fabrique la propia frase. Pero usar los vectores tal cual no vale: con las tres listas iguales, 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 pesos donde la longitud de la frase no aparece., y a su fila de llaves. La regla que lleva un sustantivo hacia su determinante y sus modificadores, ella sola, pone en esa fila su peso mayor sobre ahí, con $0.40$, y deja las en $0.23$. Añádele la de concordancia, que carga sobre las y están por ser los plurales de la frase, y la misma fila cambia de ganador: las sube a $0.37$ y ahí baja a $0.17$. Ninguna de las dos reglas ha cambiado un solo número. Lo que ha pasado es que se sumaron antes de normalizar.

Ahí está la operación entera. Separar las coordenadas en $h$ grupos y darle a cada grupo su propio softmax es normalizar antes de combinar, en vez de combinar antes de normalizar. 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 2Auto-atención: la secuencia se mira a sí mismaSin decoder no hay quien fabrique las consultas, y la salida es que las fabrique la propia frase. Pero usar los vectores tal cual no vale: con las tres listas iguales, 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 pesos 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 1Adió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 6Apilar 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 ningún peso, 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 pesos distintos, 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 pesos, 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.