Proyecto: un Transformer desde cero

Ocho lecciones han desmontado Attention is All You Need pieza a pieza y en ninguna ha llegado a ejecutarse nada. La anterior, sobre los cinco números que fijan todas las formas del modelo, era la última que quedaba por poner: con ella el Transformer está entero sobre el papel y sigue sin existir como programa. La distancia entre las dos cosas es corta —unas sesenta líneas de NumPy— y conviene recorrerla una vez.

Lo que salga por arriba no significará nada, y vale más decirlo antes de ejecutar que después: los pesos son los que devuelve el generador de números aleatorios y nadie los ha corregido jamás. Con me gusta la casa roja a la entrada, este modelo propondrá lo primero que le salga, y ajustarlo para que deje de hacerlo queda fuera de esta lección. Aun así hay tres cosas comprobables con pesos cualesquiera, y son exactamente las que dicen si la arquitectura está bien montada: que las formas encajan de un extremo al otro, que ninguna posición del decoder ve lo que todavía no se ha escrito, y que el recuento de parámetros es el de la lección anterior con otros números dentro.

Montar quince cajas suena peor de lo que es, y la razón cabe en una frase: salvo en los dos extremos, todo lo que hay dentro recibe una matriz de $T$ filas y $d_{\text{model}}$ columnas y devuelve otra de la misma forma. La atención multi-head baja a $d_v$ por cabeza y $\mathbf{W}^O$ la sube de vuelta; el perceptrón por posiciones sube a $d_{\text{ff}}$ y $\mathbf{W}_2$ la baja; el layer norm no toca ninguna forma. No es una casualidad afortunada: la conexión residual de la lección sobre los residuales y el layer normBlock 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. lo exige, porque suma casilla a casilla lo que entró con lo que salió, y dos matrices de formas distintas no se suman.

De modo que el programa entero cambia de forma dos veces, y las dos están en los extremos: abajo, donde una lista de enteros se convierte en vectores, y arriba, donde cada vector se convierte en una probabilidad por entrada del vocabulario. Entre medias no hay nada que ajustar.

Todo lo de en medio conserva la forma

El modelo de esta lección es el del artículo con las anchuras divididas por dieciséis: $N = 2$ bloques por columna, $d_{\text{model}} = 32$, $h = 2$ cabezas —y por tanto $d_k = d_v = 16$—, $d_{\text{ff}} = 128$, y un vocabulario de $\lvert V \rvert = 12$ entradas compartido por los dos idiomas, como en el artículo. La frase es me gusta la casa roja, cinco tokens, y su traducción i like the red house, cinco más <EOS>, el símbolo de fin de secuencia: $T_x = 5$ y $T_y = 6$.

Por abajo entra la caja de la lección sobre la codificación posicionalBlock 5 · Lesson 5Cada posición, un puñado de relojesLa capa de atención está terminada y sigue leyendo la frase como un montón de vectores sin orden, así que lo que hay que arreglar es la entrada: cada fila tiene que llegar sabiendo dónde estaba. Meter el número de la posición rompe la escala, dividirlo entre la longitud de la frase hace que la misma coordenada signifique cosas distintas en textos distintos, y aprender un vector por posición deja la 7 y la 8 sin relación ninguna. Los senos y cosenos del artículo resuelven las tres: acotados por construcción, con una longitud de onda por pareja de coordenadas, y con una propiedad que se demuestra en tres líneas —moverse k posiciones es girar un ángulo fijo dentro de cada pareja, con la misma matriz en toda la frase, y el producto escalar entre dos posiciones sólo depende de la distancia que las separa., con el factor que la lección anterior le añadió:

la misma operación en las dos columnas, con la misma tabla $\mathbf{E}$. Apiladas por filas dan $\mathbf{X}^{\text{enc}} \in \mathbb{R}^{5 \times 32}$ y $\mathbf{X}^{\text{dec}} \in \mathbb{R}^{6 \times 32}$. A partir de ahí sólo hay una línea, la de la lección sobre los residuales y el layer normBlock 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.:

aplicada cinco veces con cinco cosas distintas en el hueco. Ésa es toda la arquitectura:

$\text{FFN}$ es como el artículo llama al perceptrón por posiciones: el mismo objeto con otro nombre. Y por arriba sale la tabla leída al revés, que es la caja que cerró la lección anterior:

con el softmax por filas. Ninguna de esas líneas cambia el número de columnas de lo que recibe, y de ahí sale la propiedad que hace corto el programa: $N$ no aparece en ninguna forma, sólo en cuántas veces da la vuelta un bucle. Poner seis bloques en lugar de dos no toca una sola matriz.

Por qué la máscara va en todos los bloques y no sólo en el primero

La lección sobre el encoder, el decoder y las máscarasBlock 5 · Lesson 7Escribir no es leer: el decoder y su máscaraTodo lo que lleva construido el bloque supone la frase entera delante y las T posiciones calculadas a la vez, y hay media arquitectura para la que eso es falso: un modelo que escribe no tiene el resultado, lo va produciendo. Entrenarlo de una sola pasada pone la respuesta al lado de la posición que tiene que adivinarla, y la solución no es un bucle sino un sumando: menos infinito en la puntuación, antes del softmax, para que la exponencial lo borre del numerador y del denominador a la vez. Eso deja cada fila quedándose con su propio principio y volviéndolo a normalizar, así que la primera posición se lleva la unidad entera y ninguna fila deja de sumar 1. Con eso montado, la segunda subcapa del decoder es la atención del bloque anterior sin un solo cambio: consultas de lo que se escribe, claves y valores de lo que se leyó, mapa rectangular y ninguna máscara. dejó demostrado que una llamada enmascarada respeta el pasado: en su fila $t$ no entra ninguna columna a la derecha de $t$. Lo que hace falta al montar el modelo es otra cosa, y es lo único que este ensamblaje añade a lo ya escrito: que la propiedad sobreviva a apilar. Entre lo que entra al decoder y lo que sale hay seis subcapas, tres por bloque, y sólo dos llevan máscara.

Llamemos respetar el pasado a esto: la fila $t$ de lo que sale depende únicamente de las filas $1$ a $t$ de lo que entró. Tres observaciones bastan.

La primera es que las cajas que trabajan fila a fila lo respetan sin hacer nada. El perceptrón por posiciones aplica los mismos pesos a cada fila por separado, el layer norm normaliza cada fila con su propia media y su propia desviación típica, y la suma residual suma la fila $t$ con la fila $t$. Es la cuenta 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.: de las quince cajas del dibujo, sólo las tres de atención miran a otra posición.

La segunda es que la auto-atención enmascarada lo respeta por construcción, que es de lo que trataba $\mathbf{M}$: su fila $t$ de pesos vale cero de la columna $t+1$ en adelante, así que lo que sale en $t$ mezcla valores de filas que no pasan de $t$. Y la tercera es que la atención encoder-decoder lo respeta por un motivo distinto, y conviene separarlo: sus claves y sus valores no son del decoder, son del encoder. Mezcla las cinco posiciones de la fuente enteras —y debe hacerlo, la fuente está leída desde el principio—; de la salida sólo usa la consulta de su propia fila.

Respetar el pasado se conserva al componer, y con eso está hecho: con $\mathbf{Z} = f(\mathbf{X})$, la fila $t$ de $g(\mathbf{Z})$ depende sólo de las filas $\leq t$ de $\mathbf{Z}$, y cada una de ésas sólo de las filas $\leq t$ de $\mathbf{X}$. Encadena las seis subcapas y sale la propiedad entera:

que es la barra vertical de $P\left(y_t \mid y_{<t}, x_{1:T_x}\right)$ diciendo la verdad.

La cadena tiene la fuerza de su eslabón más débil, y ahí está el motivo de que el artículo ponga máscara en la primera subcapa de cada uno de sus seis bloques. Quítasela al segundo bloque de este modelo y su auto-atención mezclará, en la fila $3$, la fila $6$ de lo que le entregó el primero; todo lo que venga por encima hereda la mezcla, y la barra vertical pasa a ser mentira. La última celda lo mide con números.

El modelo entero en cuatro celdas

Cuatro celdas, en este orden: la frase y su entrada, las piezas, el modelo, y lo que se puede comprobar sin haber entrenado nada. Cada una necesita las anteriores, así que ejecútalas seguidas.

La primera fija el vocabulario compartido, convierte las dos frases en índices y construye lo que entra por abajo en las dos columnas. Mira la tercera línea que imprime: lo que entra al decoder no es la traducción, es la traducción corrida una posición, con <GO> —el símbolo de arranque— en la primera y sin el <EOS> del final.

La segunda no trae nada nuevo, y ése es el hallazgo: son seis funciones y las seis son el cuerpo de un desafío que ya resolviste. softmax_filas viene 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.; atencion es el de la lección sobre las máscarasBlock 5 · Lesson 7Escribir no es leer: el decoder y su máscaraTodo lo que lleva construido el bloque supone la frase entera delante y las T posiciones calculadas a la vez, y hay media arquitectura para la que eso es falso: un modelo que escribe no tiene el resultado, lo va produciendo. Entrenarlo de una sola pasada pone la respuesta al lado de la posición que tiene que adivinarla, y la solución no es un bucle sino un sumando: menos infinito en la puntuación, antes del softmax, para que la exponencial lo borre del numerador y del denominador a la vez. Eso deja cada fila quedándose con su propio principio y volviéndolo a normalizar, así que la primera posición se lleva la unidad entera y ninguna fila deja de sumar 1. Con eso montado, la segunda subcapa del decoder es la atención del bloque anterior sin un solo cambio: consultas de lo que se escribe, claves y valores de lo que se leyó, mapa rectangular y ninguna máscara., con su máscara opcional; multi_head reparte las columnas entre las cabezas como el 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.; layer_norm y subcapa son los de la lección sobre los residuales y el layer normBlock 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.. Lo único escrito para hoy es ffn, que son dos líneas.

La tercera saca los pesos del generador y encadena las subcapas en el orden de la tabla de arriba. mascaras es un interruptor por bloque y viene puesto en los dos; la última celda lo usará. De lo que imprime, mira la forma de la salida, que las seis filas suman $1$, y la pérdida por posición al lado de $\ln 12$.

La pérdida sale $3.292$ y $\ln 12 = 2.485$. Un modelo recién inicializado no está en el reparto uniforme: está peor. Repartir por igual es lo mejor que puede hacer quien no sabe nada, y éste ni siquiera reparte por igual —reparte desigual y sin motivo, y paga la diferencia—. Entrenar consistiría, para empezar, en bajar hasta esa línea; cruzarla es lo que se parecería a traducir.

La cuarta comprueba lo comprobable: los seis mapas con su forma, la primera fila de la salida, y la tabla que mide la sección anterior —el mismo modelo con la máscara puesta en los dos bloques, en uno, en el otro y en ninguno, mirando cuánto se mueven las cuatro primeras filas de la salida al reescribir el final de lo que entra—.

Los mapas cuadran con la tabla: dos por capa de atención —una por cabeza—, $5 \times 5$ los del encoder, $6 \times 6$ los enmascarados y $6 \times 5$ los cruzados, el único rectangular del modelo, y lo es porque las dos frases no miden lo mismo. En las seis posiciones lo más probable es la misma entrada, the, con un $0.589$ en la primera: no es que sepa inglés, es la fila de $\mathbf{E}$ que más apunta en la dirección que le sale al decoder, decidida por el mismo generador que puso los pesos.

La tabla de las máscaras es la afirmación de la sección anterior con números. Sólo la primera fila da $0$, y da $0$ exacto; enmascarar un bloque y no el otro deja $0.0768$ o $0.1651$ según cuál se quite, y no enmascarar ninguno, $0.194$. Un número pequeño no vale: o la fila anterior no se mueve, o el modelo ha leído lo que todavía no había escrito.

Y los $59\,008$ parámetros son la fórmula de la lección anterior con otros cinco números dentro —$12\,576$ por bloque de encoder, $16\,736$ por bloque de decoder y la tabla compartida—:

El modelo base del artículo tiene mil sesenta y ocho veces más pesos que éste y ni una caja más.

Comprueba tu intuición

Tres preguntas: qué demuestra un forward pass con pesos sin entrenar, si la máscara del primer bloque le alcanza al decoder entero, y cuántos parámetros guarda una capa de atención.

Queda una cosa que este modelo no ha hecho ni puede hacer de una pasada. Las seis filas salen a la vez porque la salida correcta está sobre la mesa —el teacher forcing que su lecciónBlock 5 · Lesson 7Escribir no es leer: el decoder y su máscaraTodo lo que lleva construido el bloque supone la frase entera delante y las T posiciones calculadas a la vez, y hay media arquitectura para la que eso es falso: un modelo que escribe no tiene el resultado, lo va produciendo. Entrenarlo de una sola pasada pone la respuesta al lado de la posición que tiene que adivinarla, y la solución no es un bucle sino un sumando: menos infinito en la puntuación, antes del softmax, para que la exponencial lo borre del numerador y del denominador a la vez. Eso deja cada fila quedándose con su propio principio y volviéndolo a normalizar, así que la primera posición se lleva la unidad entera y ninguna fila deja de sumar 1. Con eso montado, la segunda subcapa del decoder es la atención del bloque anterior sin un solo cambio: consultas de lo que se escribe, claves y valores de lo que se leyó, mapa rectangular y ninguna máscara. puso al lado de la máscara—, y al escribir de verdad no la hay: el decoder arranca con <GO> solo, elige un token, lo realimenta y vuelve a llamarse. Ese bucle es el desafío, escrito contra una función paso que hace de modelo.

Lo que has montado es el modelo del artículo, y eso incluye una decisión que en 2017 no se discutía: dos columnas, porque la tarea era traducir. La fuente entera por un lado, lo escrito por el otro, y una subcapa que las junta. Nada de lo que has escrito obliga a esa forma —las dos columnas comparten la tubería, las piezas y hasta la tabla—, y en cuanto la tarea deja de ser traducir, la pregunta de cuántas columnas hacen falta se abre.

Las dos respuestas que se dieron son la lección siguiente, sobre BERT y GPT: una se quedó con la columna de la izquierda y con ninguna máscara, la otra con la de la derecha y con todas, y las dos dejaron de traducir para entrenarse contra textos que no tienen traducción al lado. La aritmética que las mueve es la que acabas de escribir, sin una operación nueva; lo que cambia es qué se les pide predecir.