Apilar sin perder lo de abajo

A la atención ya no le falta nada: tiene sus tres proyecciones, su divisor, sus cabezas y, desde la lección anterior sobre la codificación posicional, unas filas que saben dónde están. Lo que no existe todavía es una capa que se pueda poner encima de otra. Y el artículo pone seis, con lo que sale de cada una haciendo de entrada de la siguiente.

Repetir seis veces una operación que ya funciona suena a no hacer nada, y es justo donde se estrelló el bloque 3. Allí lo que se repetía era un paso de la recurrencia, y componer una función consigo misma multiplica derivadas: un producto de muchos factores parecidos se va a cero o se dispara, y lo que venía de abajo no llega arriba. La profundidad de un Transformer no la pone la frase —la ponen las capas, y son seis se alargue el texto lo que se alargue—, pero seis productos siguen siendo seis productos. Esta lección es lo que hay que ponerle alrededor a una subcapa para que esa cuenta salga distinta, y son dos cosas.

Vuelve al diagrama de la lección sobre el adiós a la recurrenciaBloque 5 · Lección 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. y mira las cajas que dejó sin explicar. Cinco de las quince se llaman «Suma y layer norm», y de cada una sale una línea que baja por el lado de fuera de su columna y vuelve a entrar por debajo de la subcapa que esa caja cierra. Pincha una: esa línea es la mitad de lo que la caja suma.

Ninguna subcapa reemplaza lo que recibe

La conexión residual es la línea del dibujo escrita en símbolos, y no tiene más que esto:

donde $\text{Sublayer}$ es un hueco —como el $\varphi$ del bloque 2 o el $a$ del bloque anterior— que ocupa la atención multi-head en la primera subcapa y el perceptrón por posiciones en la segunda. Lo interesante está en su derivada. Pon las $T$ filas de $\mathbf{X}$ una debajo de otra, como en la lección sobre por qué falla el MLP (multilayer perceptron) en secuenciasBloque 3 · Lección 1Por qué el MLP falla con secuenciasPoner los vectores de los tokens en fila mete el orden dentro de la entrada, y cobra dos cosas por ello: una longitud fija que hay que truncar o rellenar, y una primera capa partida en tantos bloques como posiciones, cada uno ajustado sólo por lo que cae en la suya. Una red que aprende «no» en las cuatro primeras posiciones acierta el 0.500 cuando «no» aparece en la quinta., y llama $\mathbf{x} \in \mathbb{R}^{T \cdot d_{\text{model}}}$ a esa columna. Derivar una suma es derivar cada sumando:

Ahí está todo. El jacobiano de la subcapa depende de sus pesos y puede valer lo que sea, incluso cero; el otro sumando es la identidad y no depende de nada. Y como la regla de la cadena de su lecciónBloque 2 · Lección 7La regla de la cadena, en serioUna capa oculta hace que un mismo peso llegue a la pérdida por varios caminos a la vez, y la regla de la cadena de una variable no cubre ese caso. Dentro de un camino las derivadas se multiplican; entre caminos se suman. Apilar esas sumas da la jacobiana, y con ella derivar una composición vuelve a ser un producto de matrices. convierte una composición en un producto de jacobianos, apilar $N$ bloques multiplica $N$ paréntesis así, cada uno con su identidad dentro.

En la lección sobre el gradiente que se desvaneceBloque 3 · Lección 4El gradiente que se desvaneceMultiplicar por la misma matriz recurrente en cada paso vuelve el gradiente exponencial en la distancia: el radio espectral de W_hh decide si se desvanece o explota, la saturación del tanh solo agrava lo primero, y el recorte del gradiente frena la explosión pero no rescata lo desvanecido —que es lo que obligará a cambiar de arquitectura—. no había paréntesis: cada factor era un producto pelado y el único camino entre el paso $1$ y el paso $T$ los atravesaba todos. La LSTM (long short-term memory) respondió abriendo una vía aditiva al lado. La conexión residual es esa misma idea con la profundidad en el sitio del tiempo: aquí lo que se apila no son posiciones, son capas.

Con eso queda pagada la deuda de la lección anterior. La codificación posicional se suma una vez, abajo del todo, y ninguna caja vuelve a recordar después dónde estaba cada token: el sumando que vale $\mathbf{I}$ es lo que la lleva hasta arriba.

Una sola anchura de abajo arriba

La suma trae una restricción escondida: $\mathbf{X}$ y $\text{Sublayer}(\mathbf{X})$ se suman casilla a casilla, así que la subcapa tiene prohibido devolver otra forma que no sea $T \times d_{\text{model}}$. Ninguna de las dos la respeta por dentro y las dos hacen el viaje de vuelta: la atención multi-head baja a $d_v$ por cabeza y $\mathbf{W}^O$ la devuelve —para esto lo dejó señalado la lección sobre la atención multi-headBloque 5 · Lección 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.—, y el perceptrón por posiciones sube a $d_{\text{ff}}$ y vuelve a bajar.

Eso convierte a $d_{\text{model}}$ en el único número que hay que elegir para todo el modelo. En el perceptrón multicapa del bloque 2 cada capa tenía su anchura $d_l$ y podía estrecharse a voluntad; aquí, si una capa cambiara de anchura, la siguiente no tendría a qué sumarle nada. Los $512$ del artículo son los mismos en la primera capa y en la sexta, y esa rigidez es el precio de apilar.

Cada fila se normaliza con sus propias coordenadas

La suma crea un problema ella sola: si cada subcapa añade sin quitar, lo que circula va creciendo capa a capa, y la lección sobre el producto interno escaladoBloque 5 · Lección 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. ya midió lo que le pasa a un softmax cuando le llegan números grandes. Hay que devolver la escala a su sitio después de cada suma, y eso es el layer norm.

Fijemos la fila de la posición $t$, con sus $d_{\text{model}}$ coordenadas $x_{t1}, \dots, x_{t\,d_{\text{model}}}$. Su media y su desviación típica son las de esos números y de ningún otro:

Con ellas se centra y se escala la fila, y después entran dos vectores que se aprenden, $\boldsymbol{\gamma}$ y $\boldsymbol{\beta}$, los dos de $d_{\text{model}}$ coordenadas:

$\varepsilon$ es una constante pequeña —$10^{-5}$ en casi todas las implementaciones— que evita dividir entre cero si una fila tiene sus coordenadas iguales. Los otros dos parecen deshacer el trabajo que se acaba de hacer, y en parte lo hacen: dejar cada fila con media $0$ y desviación $1$ es una condición que nadie ha pedido, y $\boldsymbol{\gamma}$ y $\boldsymbol{\beta}$ devuelven esa libertad convertida en parámetros. Del todo no pueden deshacerla: $\mu_t$ y $s_t$ cambian de fila en fila, y estos dos vectores son los mismos en todas.

Por qué la fila y no el batch

Sumar a lo ancho de la fila es una decisión, y la alternativa es la contraria: normalizar cada coordenada a lo largo de la columna, que es la normalización por batch. Sobre secuencias falla por tres sitios. Lo que saliera en la posición $t$ dependería de qué otras frases hubieran caído en el mismo batch; dependería de $T$ y de lo que hubiera en las demás posiciones, incluidas las de padding, que no dicen nada; y habría que guardar estadísticas del entrenamiento para responder después a una frase suelta.

Con la fila no aparece ninguna de las tres, porque $\mu_t$ y $s_t$ salen de la fila $t$ y de nada más. Guarda la tercera: la caja «Suma y layer norm» no mezcla posiciones, y por eso la afirmación de la lección sobre el adiós a la recurrenciaBloque 5 · Lección 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. —tres de las quince cajas miran más allá de su propia fila— sigue en pie ahora que sabemos qué hacen las otras doce.

La subcapa que no mira a ninguna otra posición

Falta la segunda subcapa: la que el artículo llama FFN (feed-forward network) y esta lección llama el perceptrón por posiciones —los dos nombres son la misma caja—. Por dentro es el perceptrón multicapa del bloque 2 con dos capas y ReLU en medio, aplicado a una fila cada vez:

con $\mathbf{W}_1 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ff}}}$ y $\mathbf{W}_2 \in \mathbb{R}^{d_{\text{ff}} \times d_{\text{model}}}$, y con el subíndice nombrando cuál de las dos capas es. $\mathbf{1}_T$ es la columna de unos del bloque 2, que aquí repite el sesgo una vez por posición en lugar de una vez por ejemplo, y la ReLU es el $\max(0, \cdot)$ del artículo.

Las dos matrices son las mismas para todas las posiciones: la fila $3$ y la fila $300$ pasan por los mismos pesos, así que la subcapa no puede enterarse de cuál está transformando. Eso es lo que significa «por posiciones» y lo que la deja fuera de las cajas que mezclan. Su papel sale de ahí por descarte: la atención no transforma, mezcla, y éste es el único sitio donde el contenido de una posición pasa por una no linealidad propia.

El reparto no es el que sugiere el dibujo. Con $d_{\text{model}} = 512$ y $d_{\text{ff}} = 2\,048$:

A la izquierda las dos matrices del perceptrón por posiciones, a la derecha las cuatro de la atención sumadas: la caja que parece fontanería lleva el doble de pesos que la que le da nombre a la arquitectura. Ninguno de los dos números menciona a $T$.

El bloque completo, en dos líneas

Con las tres piezas puestas, un bloque del encoder es la misma línea dos veces:

con $\text{MultiHead}$ en el hueco la primera vez y $\text{FFN}$ la segunda, y con la salida de la primera haciendo de $\mathbf{X}$ de la segunda. Sesgos y vectores de normalización incluidos, son $3\,150\,336$ parámetros por bloque, y el encoder del artículo apila $N = 6$: unos $18.9$ millones, ninguno de ellos dependiente de la longitud del texto.

Y ahora la concesión. El artículo normaliza después de sumar, y eso deja el layer norm encima del camino de la identidad: el sumando que vale $\mathbf{I}$ está dentro del paréntesis, pero lo que sale de la caja ha pasado además por una normalización, así que el camino limpio de la primera sección no es tan limpio. Casi todas las implementaciones de hoy escriben $\mathbf{X} + \text{Sublayer}\left(\text{LayerNorm}(\mathbf{X})\right)$ —normalizar lo que entra en la subcapa y dejar la suma sin tocar—, que es lo que se conoce como pre-norm. Cuál entrena mejor queda fuera de este curso; que son dos operaciones distintas, no.

Las dos piezas en NumPy

La primera celda es el layer norm sobre cinco posiciones. Mira tres cosas: la media y la desviación típica de cada fila, qué devuelven $\boldsymbol{\gamma}$ y $\boldsymbol{\beta}$, y qué le pasa a la fila $0$ cuando se toca la fila $4$.

Las cinco filas entran con medias entre $0.98$ y $2.89$ y salen con media $0$ y desviación $1$; con $\boldsymbol{\gamma} = 4$ y $\boldsymbol{\beta} = -1$ salen con media $-1$ y desviación $4$, que es lo que esos dos vectores significan. Las dos últimas líneas son la sección anterior comprobada: tocar la fila $4$ mueve la fila $0$ en $0.0$ por filas y en $1.889$ por columnas.

La segunda mide lo que hace la suma, y para verla sola le quito al bloque el layer norm y la atención: quedan $12$ perceptrones por posiciones, uno encima de otro, con y sin la línea que los rodea. La escala de los pesos la he puesto baja a propósito, $0.15$: al arrancar, una subcapa es una corrección pequeña.

Cada número es una derivada medida por sondeo numérico: cuánto se mueve la salida entera al empujar una sola coordenada de la entrada. Sin la suma cae de $0.109$ con una capa a $5.182 \times 10^{-7}$ con doce —siete órdenes de magnitud, y con doce capas, no con doce mil—, porque es un producto de doce factores pequeños. Con la suma se queda entre $1.00$ y $1.80$ todo el camino: el sumando que no atraviesa ningún peso vale $1$, y los demás sólo lo corrigen.

Comprueba tu intuición

Cinco preguntas: qué añade la suma a la derivada, qué le exige a la subcapa, qué números entran al normalizar una fila, dónde están los parámetros y qué cambia al mover el layer norm de sitio.

El bloque está entero y es apilable, y con eso está terminada una de las dos columnas del dibujo. Vuelve a mirarlo: la de la derecha tiene tres subcapas donde la izquierda tiene dos, la de en medio recibe una flecha que cruza desde la otra columna, y la de abajo se llama «auto-atención enmascarada», que es una palabra que no ha salido en estas seis lecciones.

Las tres cosas son la misma cosa, y es la que separa leer un texto de escribirlo. Un encoder puede mirar la frase entera porque la tiene delante; un decoder produce la salida token a token, y al escribir el cuarto no puede mirar el quinto, que todavía no existe. Eso es la lección siguiente, sobre el encoder, el decoder y las máscaras, y lo bueno es que no toca ni una línea de lo que llevamos: tachar el futuro se hace sumándole a la rejilla de puntuaciones una matriz de ceros y de menos infinitos antes del softmax.