Atención de Bahdanau (aditiva)

Un modelo con atención necesita, en cada paso de la salida, un peso por cada posición de la entrada. La lección anterior, sobre la idea de atención, dejó montado todo lo que rodea a esos pesos: salen de un softmax que recorre las posiciones de la fuente, sirven para mezclar los estados del encoder en un vector de contexto propio de cada paso, y ese vector entra en la recurrencia del decoder con su matriz $\mathbf{W}_{ch}$. Lo que no dejó escrito es lo que los produce. Antes del softmax hay una función, $a$, que recibe una pareja (dónde está el decoder y qué hay en una posición de la entrada) y devuelve una puntuación, y de ella sólo se pidió que fuese un número y que se pudiera derivar. Esta lección la escribe, con la forma que le dio Bahdanau en 2014.

La pregunta es pequeña y concreta: ¿qué función se traga dos vectores y devuelve un número que signifique «estos dos se corresponden»? La respuesta de Bahdanau es una red neuronal diminuta, con una sola capa oculta. Pero lo que más enseña no es esa red, sino la restricción que decide su forma: la puntuación tiene que ser función de la pareja. Una función que opine por separado sobre cada vector y sume las dos opiniones no puntúa una correspondencia: puntúa dos cosas sueltas, y más abajo se ve que eso destruye el mecanismo entero.

Vuelve al mapa de la lección anterior y mira tres filas. La de i y la de read se llevan casi todo de la misma posición, la de leí, con pesos de $0.74$ y $0.81$. La de a, la siguiente, deja esa misma posición en $0.05$. El estado $\bar{\mathbf{h}}_2$ no ha cambiado entre esas tres filas (el encoder lo calculó una vez y ahí sigue), y sin embargo recibe dos puntuaciones altas y una baja. Lo único distinto es con qué estado del decoder se le comparaba cada vez.

Una red pequeña que puntúa parejas

La puntuación aditiva proyecta los dos estados a un espacio propio, de $d_a$ coordenadas, y los suma dentro de él:

con $\mathbf{W}_{sa}, \mathbf{W}_{ha} \in \mathbb{R}^{d_a \times d_h}$, una por lado. Las dos tienen $d_h$ columnas porque en este bloque el encoder y el decoder miden lo mismo, pero la suma no lo exige. Sólo pide que las dos proyecciones caigan en el mismo espacio, es decir, que las dos matrices tengan $d_a$ filas; las columnas de cada una son cosa de su lado. Si los estados del decoder tuvieran $d_s$ coordenadas, $\mathbf{W}_{sa}$ sería de $d_a \times d_s$ y el resto de la puntuación quedaría igual. En el artículo de Bahdanau, de hecho, no coinciden: su encoder lee la frase en los dos sentidos y concatena los dos estados, así que mide el doble que el decoder.

Esa suma es la preactivación de una capa oculta, y lo que sigue es lo que ya hace cualquier capa: aplastarla y leerla.

con $\mathbf{v}_a \in \mathbb{R}^{d_a}$. Es un perceptrón multicapa (multilayer perceptron, MLP) del bloque 2 en su tamaño mínimo: una capa oculta de $d_a$ unidades con $\tanh$, y una salida lineal de una sola neurona. Esa neurona es la de la lección sobre la neuronaBloque 2 · Lección 1La neurona artificialUna suma ponderada y un umbral bastan para que un solo objeto convierta un vector en una decisión. Sus pesos son la perpendicular de una frontera que parte el espacio en dos, y el número que calcula antes de decidir mide a qué distancia de esa frontera ha caído la entrada., $z = \mathbf{w}^{\top}\mathbf{x} + b$: su entrada es la capa oculta, $\tanh(\mathbf{p}_{ij})$, y sus pesos son $\mathbf{v}_a$, uno por unidad oculta. Llamemos a $\mathbf{v}_a$ el vector de lectura, porque convierte los $d_a$ números de la capa oculta en el único número que es la puntuación.

Ninguna de las dos capas lleva sesgo, y por razones distintas. El de la salida sobra porque sumaría el mismo número a toda la fila, y una constante por fila no sobrevive al softmax. El de la capa oculta no sobraría: dentro de la $\tanh$ desplaza el punto donde trabaja cada unidad, más cerca o más lejos de la saturación, y eso sí cambia las puntuaciones. Bahdanau no lo pone y muchas implementaciones sí; aquí sigo la fórmula del artículo.

Cuenta los pesos: $2 d_a d_h + d_a$. No aparece $T_x$, no aparece $T_y$ y no aparece la pareja. Hay un solo juego de pesos y con él se puntúan las $T_x \cdot T_y$ parejas de esta frase y las de todas las demás, que es la misma economía de la recurrencia del bloque anterior (un $\mathbf{W}_{hh}$ para todas las posiciones) aplicada a una rejilla en vez de a una fila.

Sin la no linealidad, todos los pasos mirarían al mismo sitio

Quita la $\tanh$ y mira qué queda. El vector de lectura entra en la suma y la reparte:

dos números que se suman y que no se han mirado el uno al otro. El primero, $\mathbf{v}_a^{\top}\mathbf{W}_{sa}\mathbf{s}_{i-1}$, depende del paso $i$ pero no de la posición $j$. En la fila $i$, es decir, en las $T_x$ puntuaciones $e_{i1}, \dots, e_{iT_x}$ que el softmax del paso $i$ convierte en pesos, vale lo mismo en todas: es una constante sumada a la fila entera. Y una constante así no cambia el softmax. Es la invariancia al desplazamiento de la lección sobre funciones de pérdidaBloque 2 · Lección 5Funciones de pérdida: MSE y entropía cruzadaUna pérdida resume en un número lo mal que lo hace la red, y entre dos candidatas decide su derivada. El error cuadrático medio entrega su corrección más pequeña justo donde la red se equivoca con más seguridad; la entropía cruzada, que sale de leer la salida como probabilidad, deja esa derivada en la diferencia entre lo predicho y la etiqueta., la misma que permite restar el máximo antes de exponenciar, y aquí hace que ese primer número se cancele arriba y abajo:

El índice $i$ ha desaparecido del lado derecho. Los $T_y$ pasos reciben la misma fila de pesos, y con ella el mismo vector de contexto, de modo que el modelo vuelve a ser el resumen fijo de las lecciones anteriores: una alineación aprendida, sí, pero congelada, sólo que congelada donde diga el entrenamiento en vez de en el último token. La $\tanh$ es lo que impide esa separación. Aplasta la suma antes de leerla, y una función no lineal de una suma no se deja repartir entre sus sumandos, así que la puntuación no puede volver a partirse en una opinión sobre el decoder más otra sobre el encoder.

Por dónde le llega el gradiente

Ahora hay parámetros que entrenar, y la pregunta es si la pérdida los alcanza. La cadena empieza en la mezcla, donde $\alpha_{ij}$ multiplica a un vector conocido:

y sigue por el softmax, que reparte cada puntuación entre todos los pesos de su fila. Del otro lado sale la derivada respecto de una puntuación:

Léela despacio, porque dice cómo aprende a mirar. Tiene dos partes, y la primera es lo que le pasa a la mezcla cuando la puntuación sube. Si $e_{ij}$ crece un poco, la posición $j$ gana peso a costa de las demás y $\mathbf{c}_i$ se desplaza hacia $\bar{\mathbf{h}}_j$. En concreto, $\partial \mathbf{c}_i / \partial e_{ij} = \alpha_{ij}\left(\bar{\mathbf{h}}_j - \mathbf{c}_i\right)$, que es la misma cuenta del desplegable con $\bar{\mathbf{h}}_k$ en el sitio de $\partial \ell / \partial \alpha_{ik}$. La segunda parte es $\nabla_{\mathbf{c}_i}\ell$, cuyo opuesto apunta hacia donde la pérdida baja: hacia donde el descenso de gradiente querría llevar la mezcla. La derivada es el producto escalar de las dos, así que mide si acercar $\mathbf{c}_i$ a $\bar{\mathbf{h}}_j$ lo lleva hacia donde conviene.

De ahí sale lo que el entrenamiento hace con la puntuación. Si acercar la mezcla a $\bar{\mathbf{h}}_j$ hace bajar la pérdida, la puntuación sube; si la hace subir, la puntuación baja; y si $\bar{\mathbf{h}}_j$ coincide con $\mathbf{c}_i$, subir la puntuación no mueve la mezcla, así que no recibe corrección ninguna. Una puntuación, por tanto, no se corrige según lo que valga su estado por sí solo, sino según lo que aporta frente a la mezcla que la fila ya construyó. Tiene sentido que sea una comparación y no una nota: los pesos suman $1$, así que subir uno es bajar los demás, y no hay forma de premiar a una posición sin decir a costa de cuál.

De ahí para atrás todo es cadena conocida. La lectura devuelve $\tanh(\mathbf{p}_{ij})$, la $\tanh$ aporta su derivada $1 - \tanh^{2}$, y cada matriz recoge el producto exterior con el vector que multiplicaba:

Las $T_x \cdot T_y$ parejas suman sobre los mismos tres arrays, que es lo que significa compartir pesos.

Y hay una consecuencia que no está en los parámetros. Fíjate en dónde aparece $\bar{\mathbf{h}}_j$: en la mezcla, como uno de los sumandos, y en las $T_y$ puntuaciones, como argumento de $a$. Son dos sitios, así que el gradiente le llega por dos rutas y hay que sumarlas:

A esas dos se les suma la que ya traía el bloque anterior, la que baja desde $\bar{\mathbf{h}}_{j+1}$ por la recurrencia del encoder, y ahí está lo interesante: las dos de arriba no pasan por ella. Van del paso $i$ del decoder a la posición $j$ de la entrada sin recorrer el trecho de recurrencia que separa esa posición del final de la frase, así que no multiplican por $\mathbf{W}^{\text{enc}}_{hh}$ ni una sola vez y su longitud no crece cuando $j$ se aleja. El producto de matrices que la lección sobre el gradiente desvanecidoBloque 3 · Lección 4El gradiente que se desvaneceMultiplicar por la misma matriz recurrente en cada paso vuelve el gradiente exponencial en la distancia. Lo que garantiza que se desvanezca es γ·σ_max < 1, la mayor pendiente del tanh por el mayor valor singular de W_hh; el radio espectral solo marca el ritmo de la matriz sola. El recorte del gradiente frena la explosión pero no rescata lo desvanecido, y eso es lo que obligará a cambiar de arquitectura. señalaba como causa del problema no aparece en estas dos rutas.

La puntuación, en NumPy

La primera celda calcula la rejilla entera. Guarda los estados por filas (una fila por paso, una fila por posición) y por eso las matrices aparecen transpuestas respecto de las ecuaciones. Mira dos cosas: cuánto se parece a un mapa el reparto que sale de unos pesos recién inicializados, y qué pasa al pedir la versión sin $\tanh$.

El mapa recién inicializado no dice nada: el mayor peso de cada fila ronda $0.18$ cuando el reparto uniforme sería $0.167$, así que las ocho filas son prácticamente la media de la frase. La alineación interpretable del explorable no viene de la arquitectura, viene del entrenamiento; lo que la arquitectura garantiza es que haya por dónde entrenarla. Y lo que sale tras la línea en blanco es la sección anterior comprobada a mano: sin la $\tanh$, la fila $1$ y la fila $8$ salen iguales hasta el último decimal.

La segunda celda persigue las dos rutas por separado. Necesita una pérdida y usa la más barata posible (una cuyo gradiente respecto de cada $\mathbf{c}_i$ es un vector dado) para que lo único que se esté midiendo sea el mecanismo.

La última línea es la comprobación: sumadas, las dos rutas reproducen el gradiente numérico hasta el orden de $10^{-10}$, y por separado cada una se equivoca en la primera cifra. Las normas dicen algo más, y es honesto mirarlo: con los pesos sin entrenar la ruta del valor es once veces mayor que la de la puntuación. Al principio casi todo el gradiente llega por donde llegaría aunque la alineación fuese uniforme, y es la ruta pequeña, la que pasa por $a$, la única que puede sacar al mapa de esa uniformidad. Cambia la escala de va a $2.0$ y ejecuta otra vez las dos celdas: las puntuaciones se separan, los pesos se afilan y las dos normas se acercan.

Comprueba tu intuición

Cinco preguntas: qué se rompe sin la $\tanh$, cuánto pesa la función $a$, por dónde le llega el gradiente a un estado del encoder, qué dice la derivada de una puntuación, y qué parte del cálculo se puede sacar del bucle.

Queda una cuenta pendiente de la última pregunta. Las dos proyecciones se sacan del bucle porque cada una depende de un solo índice, pero entre la suma y el vector de lectura hay una $\tanh$, y eso obliga a construir la rejilla de $T_x \cdot T_y$ vectores de $d_a$ coordenadas de una en una. Una frase de sesenta tokens traducida a sesenta pide $3\,600$ de esos vectores, y ninguna de las operaciones que los producen es un producto de matrices grande: son sumas y $\tanh$ sobre un bloque de tres dimensiones. El hardware que entrena estas redes hace una cosa mucho mejor que las demás, y no es ésa.

Hay otra forma de puntuar una pareja que evita el problema entero, y su fórmula cabe en cinco símbolos: la atención de Luong, multiplicativa, que es la lección siguiente. En vez de proyectar, sumar y aplastar, multiplica los dos estados a través de una matriz, con lo que la rejilla completa de puntuaciones pasa a ser un producto de matrices. La lección compara las dos puntuaciones (qué cuesta cada una, qué puede hacer cada una) y ahí empieza a decidirse cómo será el bloque 5.