Escribir no es leer: el decoder y su máscara

La suma y el layer norm de la lección anterior, sobre los residuales, envuelven lo que se les ponga dentro y no preguntan qué es, así que poner una capa encima de otra dejó de ser el problema. El que queda no es de profundidad. Todo lo que este bloque ha construido da por hecho que la frase está entera delante y que las $T$ posiciones se calculan a la vez, y hay media arquitectura para la que eso es falso: un modelo que traduce no tiene la traducción delante, la escribe, un token detrás de otro. Y al entrenarlo calculándolo todo de golpe —que es lo que se ganó quitando la recurrencia— la respuesta acaba sentada al lado de la posición que tiene que adivinarla.

Míralo en el gato bebe leche, que es lo que el modelo tendría que aprender a escribir. Si las cuatro posiciones se calculan a la vez y cada una mira a todas —que es exactamente lo que hace 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.—, la posición que debe producir bebe está mirando leche mientras la produce. Adivinar con la respuesta delante sale muy bien: la pérdida baja a cero durante el entrenamiento. Al generar de verdad no hay nada a la derecha, porque todavía no se ha escrito, y el modelo no ha aprendido a hacer otra cosa. Esta lección prohíbe esa mirada sin renunciar a calcular las $T$ posiciones a la vez, y de paso monta la subcapa que une las dos columnas del artículo.

Vuelve a la rejilla de la auto-atención y enciende la máscara. Es la misma frase, con las mismas proyecciones y la misma cuenta; lo único que cambia es que desaparece todo lo que hay a la derecha de la diagonal y que las casillas que quedan se reparten entre ellas la unidad que había que repartir. Empieza por la primera fila, que es el caso extremo: no tiene a nadie a quien mirar más que a sí misma.

Tachar el futuro dentro del softmax

La máscara es una matriz del tamaño de la rejilla, y sus casillas sólo toman dos valores:

Con ella, la llamada a la atención cambia en un sumando y en nada más:

con $\mathbf{M} \in \mathbb{R}^{T \times T}$ y el softmax por filas, como siempre. Nada de lo de dentro se toca: las tres proyecciones son las mismas, el divisor de la lección sobre el producto interno escaladoBlock 5 · Lesson 3La raíz que faltaba en la atenciónDe la fórmula de la atención queda un solo símbolo sin justificar, y no se justifica midiendo: se cuenta. Una puntuación es una suma de d_k términos, y lo que crece al añadir sumandos no es su tamaño sino cuánto se aleja del cero la suma, que va con la raíz de d_k. Lo que ese tamaño le hace al softmax es lo que el curso ya vio dos veces con otra ropa: la fila se cierra sobre una posición, las dos derivadas del softmax se apagan a la vez y las dos únicas matrices que hay antes de las puntuaciones dejan de recibir gradiente. Dividir entre la raíz devuelve las puntuaciones a una escala que no menciona la anchura del modelo, y con ellas el reparto y el gradiente. es el mismo, y las cabezas de la lección sobre la atención multi-headBlock 5 · Lesson 4Varias cabezas, varios repartosUna posición tiene una fila y en esa fila cabe una sola manera de mirar la frase. La atención multi-head monta h atenciones en paralelo sobre la misma X, cada una con sus tres proyecciones y su propio mapa, y las junta concatenando y proyectando. Lo que cambia no son las puntuaciones sino dónde se normaliza: cada grupo de coordenadas con su propio softmax, en vez de todas sumadas antes de uno solo, que es lo que deja a una regla tapar a otra sin dejar rastro. Y de repartir la anchura, d_k = d_model/h, sale que el coste no depende de h: los mismos parámetros y las mismas multiplicaciones con una cabeza que con dieciséis. son las mismas —cada una recibe la misma $\mathbf{M}$—.

Por qué $-\infty$ y no otra cosa se ve escribiendo el peso entero:

donde $e_{ij}$ es la puntuación ya dividida por $\sqrt{d_k}$. Como $\exp(-\infty) = 0$, una casilla tachada desaparece del numerador y de su término del denominador, las dos cosas a la vez. Eso es lo que hace que no sobre nada: la unidad que el softmax reparte se la reparten enteramente las casillas que quedan.

Las dos alternativas que se le ocurren a cualquiera fallan justo ahí. Poner un $0$ en la puntuación no tacha nada: $\exp(0) = 1$ es un peso perfectamente normal, y en la celda de más abajo la fila $0$ acaba dándole al futuro el $0.525$ de su peso. Y calcular el softmax como siempre y poner a cero los pesos después deja el denominador como estaba, así que la fila suma lo que sumara su principio —$0.085$, $0.291$, $0.871$— y lo que sale ya no es una mezcla sino una mezcla encogida, con un factor distinto en cada fila.

Enmascarar es renormalizar el principio de la fila

Hay una manera de leer la máscara que la deja en su sitio, y sale de la fórmula de arriba en dos pasos. Para $j \leq i$ el sumando $M_{ij}$ vale $0$, de modo que el numerador es el de siempre y lo único que cambia es el denominador, que se queda con los términos hasta $i$. Divide arriba y abajo por la suma completa $\sum_{m=1}^{T}\exp(e_{im})$ y aparecen los pesos sin máscara:

Es decir: coge la fila que 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. calculaba, tira la parte de la derecha y divide lo que queda entre lo que sumaba. De ahí salen tres cosas, y ninguna necesita otra cuenta. El orden dentro de lo que sobrevive no cambia, porque todas esas casillas se dividen entre el mismo número: la máscara no altera preferencias, retira candidatos. La fila $1$ se queda con un solo candidato y por tanto con $\alpha_{11} = 1$, la sepa el modelo o no. Y la máscara no es una segunda operación al lado del softmax, es el mismo softmax sobre menos casillas.

Una pasada en lugar de T

Falta decir qué entra en el decoder, porque hasta ahora hemos hablado de lo que no puede mirar y no de lo que tiene delante. Lo que entra es la salida correcta corrida una posición: en la fila $i$ va el token $y_{i-1}$, el que el modelo debería haber escrito en el paso anterior, y en la fila $1$ va <GO>, el símbolo de arranque, porque antes de la primera no hay ninguna. Así la fila $i$ tiene delante $y_{<i}$ y tiene que producir $y_i$, que es precisamente la distribución que el bloque anterior escribía

una por fila. Eso es el teacher forcing de la lección sobre seq2seqBlock 3 · Lesson 8Seq2seq: de secuencia a secuenciaDos RNN encadenadas: un encoder que lee la secuencia de entrada y la comprime en un solo vector de contexto, y un decoder que escribe una secuencia nueva a partir de ese vector. Entrena de punta a punta —el encoder aprende sólo por el gradiente que el decoder le devuelve en c— sobre una tarea de juguete, invertir tiras de símbolos, que se resuelve en el navegador. Funciona con las cortas y se degrada con las largas: toda la entrada, sea cual sea su longitud, tiene que caber en los mismos d_h números. Ese cuello de botella es la puerta al bloque 4., sin ningún cambio en la idea: se le da al modelo el token verdadero en lugar del que él escribió. Lo que cambia es el precio. Allí las $T_y$ filas salían de $T_y$ pasos de una recurrencia; aquí salen de un producto de matrices, y la máscara es lo que hace que esa barra vertical diga la verdad —sin ella, la fila $i$ estaría condicionando también con $y_{i}$, $y_{i+1}$ y todo lo demás—.

Y ahora las dos concesiones, que van juntas. La máscara hace paralelo el entrenamiento, no la generación: al escribir de verdad no hay ninguna salida correcta que correr una posición, la fila $i$ necesita el token que produjo la pasada anterior, y son $T_y$ pasadas, una por token. Lo otro es que el triángulo no ahorra ni una multiplicación: las $T^{2}$ puntuaciones se calculan enteras y después se tachan más de la mitad. La máscara es una prohibición, no una optimización.

La subcapa que mira a la otra columna

El decoder lleva tres subcapas donde el encoder lleva dos, y la de en medio es la única caja de todo el artículo que junta las dos columnas. Es la atención del bloque anterior, entera, con las tres listas leídas de dos sitios distintos:

donde $\mathbf{X}^{\text{dec}} \in \mathbb{R}^{T_y \times d_{\text{model}}}$ es lo que se lleva escrito y $\mathbf{X}^{\text{enc}} \in \mathbb{R}^{T_x \times d_{\text{model}}}$ es lo que el encoder dejó leído. Los superíndices de $\mathbf{X}$ dicen de qué columna viene la matriz; los de $\mathbf{W}$ siguen diciendo qué papel proyecta cada una, como desde 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.. Ninguno de los dos cuenta capas.

La forma del mapa cae sola y es la que delata de dónde viene esta subcapa:

rectangular, una fila por posición de la salida y una columna por posición de la entrada. Es la rejilla del bloque anterior, la de la lección sobre la atención multiplicativaBlock 4 · Lesson 5Atención de Luong (multiplicativa)La misma casilla de la lección anterior, llena de otra forma: en vez de proyectar los dos estados y aplastar su suma, se multiplican a través de una sola matriz. Sale una puntuación sin capa oculta, con casi los mismos pesos y un gradiente más corto, y con una propiedad que la aditiva no tiene: la rejilla entera de puntuaciones es un producto de matrices. Aquí se comparan las dos por coste —operaciones, memoria intermedia y forma del cálculo— y se dice qué impide todavía calcular la rejilla de una sola vez., con los mismos dos ejes y la misma fila sumando $1$. Lo que ha cambiado desde entonces son tres cosas y ninguna es la operación: la puntuación es el producto interno escalado en lugar del modelo de alineación, hay $h$ cabezas en paralelo en lugar de una, y la consulta ya no es el estado de una red recurrente sino una fila de una matriz que existe entera.

Aquí no hay máscara, y merece la pena decir por qué, porque suena raro después de la sección anterior. Lo que se prohíbe es leer lo que todavía no se ha escrito, y la entrada no se está escribiendo: se leyó entera antes de empezar. La posición $T_x$ de la fuente está disponible desde el primer token de la salida igual que la posición $1$. (Las implementaciones de verdad sí ponen otra máscara aquí, la del padding, para que las posiciones de relleno de un batch no cuenten; es el mismo mecanismo aplicado a otro problema y queda fuera de esta lección.)

Con esto, las tres llamadas a la atención del artículo se pueden poner en una tabla, y son las tres cajas que se encendían en el diagrama de 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.:

Una sola función, tres juegos de argumentos. Y las tres subcapas del decoder van envueltas en la línea de la lección anterior, sin ninguna excepción: se suma lo que entró y se normaliza el resultado, tres veces.

Las dos atenciones en NumPy

La primera celda es la máscara sola, sobre una rejilla de puntuaciones cuyos números he puesto yo a mano para que las cuentas se puedan seguir a ojo. Mira tres cosas: la fila $0$, la columna de las sumas, y qué pasa con las dos maneras de equivocarse.

Las tres primeras líneas son el argumento entero. Con $-\infty$ la fila $0$ sale [1. 0. 0. 0.] y las cuatro sumas valen $1$. Con un $0$ en la puntuación las sumas también valen $1$, y ahí está la trampa: la fila reparte bien, pero reparte entre gente que no debería estar —el $0.525$ que se van tres columnas del futuro—. Con los pesos tachados después, las sumas caen a $0.085$, $0.291$, $0.871$ y $1$: cuatro filas escaladas de cuatro maneras distintas, que es lo peor de los tres resultados aunque sea el que parece más inofensivo.

La segunda celda monta las dos atenciones del decoder y comprueba lo único que la máscara promete.

Los dos mapas tienen la forma de la tabla —$5 \times 5$ el propio, $5 \times 7$ el cruzado— y las diez filas suman $1$. De las $25$ casillas del primero quedan vivas $15$, que es la cuenta $T(T+1)/2$. Y las dos últimas líneas son la propiedad por la que existe todo esto: reescribir las posiciones $4$ y $5$ del decoder mueve lo que sale en las tres primeras exactamente $0.0$, mientras que la cuarta —que es una de las que he reescrito— sí se mueve. Cambia causal=True por causal=False y ejecuta otra vez: ese $0.0$ deja de serlo, y eso es un modelo que ha leído lo que todavía no ha escrito.

Comprueba tu intuición

Cinco preguntas: por qué $-\infty$ y no un cero, cuántas casillas sobreviven, qué lleva máscara y qué no, qué se paraleliza de verdad, y qué sale en la primera fila.

import numpy as np
 
E = np.array([[2.0, 1.0, 0.5, 3.0],
              [0.0, 2.0, 1.0, 1.0],
              [1.0, 0.0, 2.0, 0.5],
              [0.5, 1.0, 0.0, 2.0]])
M = np.triu(np.full((4, 4), -np.inf), k=1)
Z = np.exp(E + M - (E + M).max(axis=1, keepdims=True))
A = Z / Z.sum(axis=1, keepdims=True)
print(np.round(A[0], 2))
 

Ya no queda ninguna caja del dibujo sin una lección detrás, salvo la de arriba del todo, que convierte cada fila en una probabilidad por entrada del vocabulario. Lo que no ha aparecido nunca en estas siete lecciones son los números: cuántos bloques se apilan, cuánto miden, cuántas cabezas hay, qué tablas comparten pesos y cuántos parámetros suma todo junto. Cada pieza se ha construido con las anchuras que le hacían falta al argumento, y el artículo eligió unas concretas.

Ésa es la lección siguiente, sobre la arquitectura completa del artículo: la figura 1 caja por caja, con sus tamaños, puesta al lado de lo que llevamos construido. Es la lección que convierte Attention is All You Need en un texto que puedes abrir y leer de principio a fin, que era la promesa con la que empezó el bloque.