Quince cajas, cinco números

Descrito, el Transformer ya está entero. La lección anterior, sobre el encoder, el decoder y las máscaras, cerró la última operación que quedaba suelta, y con ella el modelo sabe leer una frase, escribir otra y no mirar lo que todavía no ha escrito. Construido no está. Lo que hay escrito son letras —$d_k$, $d_{\text{ff}}$, $h$, $N$—, y un programa no reserva memoria para una letra: hasta que cada una no valga un número, no hay ni una matriz que crear.

Los números del artículo son cinco, y ponerlos encima del dibujo es lo que cierra el bloque. Cinco y no quince: fijados ésos, todas las formas del modelo quedan determinadas —incluidas un par que parecen libres y no lo son—, y de ellos sale el recuento entero de parámetros y, con él, el reparto, que no está donde el dibujo sugiere. Antes hay que cerrar la caja que ninguna lección construyó, la que corona la columna del decoder, y su respuesta viene de mucho más atrás que el resto del bloque: es la tabla de embeddings de la lección sobre las representaciones densasBlock 1 · Lesson 6Representaciones densas: la idea centralDar a cada entrada del vocabulario unas pocas coordenadas reales, todas compartidas y ninguna reservada, rompe la equidistancia del one-hot y deja la tabla de la representación en una matriz que sí cabe. Los 15 millones de números que la llenan no los escribe nadie a mano: la hipótesis distribucional dice de dónde salen., leída al revés.

La figura del artículo está en inglés y este bloque la ha ido construyendo en español, así que lo primero es la correspondencia, rótulo por rótulo:

La caja de arriba lee la tabla del bloque 1 al revés

De lo que sale del último bloque del decoder, la fila $t$ es un vector de $d_{\text{model}}$ coordenadas, $\mathbf{x}^{\text{dec}}_t$, y lo que hace falta en esa posición es una probabilidad por cada entrada del vocabulario. El artículo dibuja dos cajas para ese salto: una matriz que lleva de $d_{\text{model}}$ a $\lvert V \rvert$ coordenadas y el reparto de siempre. Esa matriz podría ser suya, y no lo es:

donde $\mathbf{E} \in \mathbb{R}^{\lvert V \rvert \times d_{\text{model}}}$ es la tabla de embeddings del bloque 1, la misma que abajo del todo cambia cada token por su vector. La coordenada que le toca a la entrada $w$ dentro del softmax es $\mathbf{e}_w^{\top}\mathbf{x}^{\text{dec}}_t$: el producto escalar entre la fila de esa entrada y lo que el decoder lleva calculado. La entrada más probable es la que más apunta en la dirección que la posición pide.

Es la tabla leída en la otra dirección. Abajo se entra con una entrada y se sale con su vector; aquí se entra con un vector y se sale con una puntuación por entrada. Las $T_y$ posiciones salen a la vez del mismo producto, como todo en este bloque:

con el softmax por filas. Una fila por posición escrita, una columna por entrada del vocabulario, y ni un parámetro nuevo en toda la caja.

La misma tabla en tres sitios

El artículo va más lejos que eso y usa esa matriz tres veces: en el embedding de entrada, en el de salida y en la proyección al vocabulario. Eso exige un vocabulario común a los dos idiomas, y lo tiene: sus $37\,000$ entradas salen de un BPE (byte-pair encoding) entrenado sobre los dos textos a la vez, la técnica de la lección sobre la tokenizaciónBlock 1 · Lesson 2Tokenización: palabras, caracteres, subpalabrasAntes de decidir qué vector le corresponde a una unidad hay que decidir cuál es la unidad. Esa elección fija el tamaño del vocabulario, la longitud de las secuencias y qué significa que una palabra sea nueva.. Una tabla de $37\,000 \times 512$ son $18\,944\,000$ parámetros, y lo que se compra compartiéndola es no pagarla tres veces.

Queda un factor que 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. dejó fuera. Allí la entrada de la primera capa se escribió $\mathbf{x}_t = \mathbf{e}_{w_t} + \text{PE}_t$; el artículo multiplica el primer sumando:

Con los números delante se ve para qué. Todas las filas de $\text{PE}$ miden lo mismo, $\sqrt{256} = 16$ —es la cuenta $\text{PE}_{\text{pos}}^{\top}\text{PE}_{\text{pos}} = d_{\text{model}}/2$ de aquella lección—, mientras que una fila de $\mathbf{E}$ recién inicializada, con coordenadas pequeñas alrededor de cero, mide alrededor de $1$. Sumados así, el contenido queda debajo de la posición; multiplicar por $\sqrt{512} \approx 22.6$ los pone en el mismo orden de magnitud. El artículo escribe la multiplicación y no la justifica: ésta es la razón que se le suele dar, no una que él dé.

Cinco números fijan todas las formas

La tabla 3 del artículo tiene una fila por cada modelo que entrenaron. La primera, la que llaman base, es la que ha ido apareciendo suelta a lo largo del bloque:

Todo lo demás sale de esos cinco. El reparto 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. da $d_k = d_v = d_{\text{model}}/h = 64$, y con él $\mathbf{W}^Q_i$, $\mathbf{W}^K_i$ y $\mathbf{W}^V_i$ miden $512 \times 64$ y $\mathbf{W}^O$ mide $(h \cdot d_v) \times d_{\text{model}} = 512 \times 512$; el perceptrón por posiciones sube a $2\,048$ y vuelve a bajar; cada normalización aporta dos vectores de $512$. Lo que no está en la lista tampoco está en ninguna forma: ni $T$, ni el tamaño del batch. Esa tabla tiene además dos columnas que esta lección no toca, dropout y label smoothing: son del entrenamiento y no de la arquitectura, y quedan fuera del curso.

Un tercio del modelo es una tabla de búsqueda

Un bloque del encoder pesa $3\,150\,336$ parámetros —la cuenta 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., con sus sesgos y sus vectores de normalización dentro—. Uno del decoder lleva exactamente una atención y una normalización más:

El modelo entero son seis de cada, más la tabla que las dos columnas comparten:

Ahí está el reparto que el dibujo esconde. Son $18.9$ millones de encoder, $25.2$ de decoder y $18.9$ de tabla: un $30$, un $40$ y un $30$ por ciento. Lo que más pesa de una sola pieza no es ninguna de las quince cajas, es la tabla de búsqueda, que ella sola supera a las seis capas del encoder juntas —$18\,944\,000$ frente a $18\,902\,016$—. Y de ahí que el artículo la comparta: las dos copias que se ahorra suman $37\,888\,000$ parámetros, casi tanto como las doce capas.

Esa cuenta es una reconstrucción a partir de la figura, no la contabilidad del artículo, y conviene decir en qué se separan. El artículo declara $65$ millones y aquí salen $63.0$. Los dos que faltan caben enteros en dos detalles que la figura no muestra: el tamaño exacto del vocabulario —con $41\,000$ entradas la cifra cuadraría— y dónde pone cada implementación sus sesgos. La misma receta aplicada a la otra fila de la tabla 3, la que llaman big ($d_{\text{model}} = 1\,024$, $h = 16$, $d_{\text{ff}} = 4\,096$), da $214.2$ millones frente a los $213$ declarados.

El recuento en una función

Toda la sección anterior cabe en ocho líneas de aritmética con números enteros. Mira tres cosas: los tres porcentajes, la diferencia entre la tabla y el encoder entero, y qué le pasa al total si cada caja tuviera su propia tabla.

Los tres porcentajes salen $30.0$, $40.0$ y $30.0$, y el decoder pesa más que el encoder por lo único en lo que se diferencian: una subcapa de atención más, seis veces. La tabla le saca al encoder $41\,984$ parámetros, que sobre dieciocho millones es empatar. Y sin compartirla el modelo se iría a $100\,933\,632$, de $63$ millones a $101$ por una decisión que no cambia ni una operación del dibujo.

La última línea cambia los cinco argumentos por los de la fila big y da $214\,171\,648$. Lo que no vas a poder cambiar es la longitud del texto, porque $T$ no aparece en ninguna línea de la función: no es que se haya quedado fuera de la cuenta, es que no hay ninguna matriz cuya forma dependa de ella.

Comprueba tu intuición

Cinco preguntas: qué forma no fijan los cinco números, qué costaría no compartir la tabla, qué vas a encontrarte al abrir la figura, si un texto más largo pide más parámetros y qué sale por arriba del decoder.

import numpy as np
 
# E es la tabla compartida; X_dec sale del ultimo bloque
V, d_model, T_y = 37000, 512, 4
rng = np.random.default_rng(0)
E = rng.normal(size=(V, d_model))
X_dec = rng.normal(size=(T_y, d_model))
 
O = X_dec @ E.T
P = np.exp(O - O.max(axis=1, keepdims=True))
P /= P.sum(axis=1, keepdims=True)
print(P.shape, np.allclose(P.sum(axis=1), 1.0))
 

Ya no queda nada que describir. Las quince cajas tienen su operación, sus formas y sus números, y eso es justo lo que hace falta para escribir el modelo entero y ver qué contesta. Lo que no cabe es el tamaño: sesenta y tres millones de parámetros no se mueven dentro de un navegador, así que lo que se construye es la arquitectura con anchuras de juguete —las mismas fórmulas, dos capas en lugar de seis, dos cabezas en lugar de ocho— y una frase corta que se pueda seguir a ojo.

Eso es la lección siguiente, el proyecto: un Transformer desde cero. Casi todas las piezas están escritas ya, cada una en el desafío de su lección —el softmax por filas, la atención con su máscara, el layer norm, el perceptrón por posiciones, la codificación posicional—, y lo que queda es ensamblarlas y hacer una pasada hacia delante de principio a fin: de una frase de tokens a una probabilidad por entrada del vocabulario, con los mapas de atención a la vista por el camino. Entrenar de verdad pide una máquina que un navegador no tiene, y eso tiene su sitio más adelante en el bloque.