LSTM: memoria con compuertas

La lección anterior, sobre el gradiente que se desvanece, dejó el diagnóstico cerrado: lo que hunde o dispara la señal de vuelta no es el valor de los pesos sino la forma del camino, y mientras cada paso hacia atrás cruce la misma matriz, el resultado será una potencia de la distancia. También dejó apuntada la salida —una vía por la que el estado avance sumando— y el nombre de la red que la lleva, la LSTM (long short-term memory). Nombrarla no es construirla: una memoria que sólo suma tiene un defecto tan grave como el que arregla, y es que nunca borra nada. Falta lo difícil, que es quién decide, paso a paso, qué se conserva, qué se escribe encima y qué de lo guardado sale a la luz.

Que son tres decisiones y no una se ve en una frase corriente. Toma Las cámaras que el ayuntamiento instaló el año pasado en la plaza no funcionan. Para conjugar funcionan hay que arrastrar el plural de cámaras doce posiciones, sin que ayuntamiento ni plaza —los dos en singular, los dos más cerca— lo pisen por el camino. Arrastrarlo intacto es una cosa; meterlo en la memoria cuando aparece cámaras y no en cada palabra que pasa es otra; y tenerlo guardado durante esas doce posiciones sin que gobierne lo que la red responde en ellas es una tercera. La RNN (recurrent neural network) vanilla hace las tres con el mismo gesto, porque sólo tiene un gesto: reescribir entero $\mathbf{h}_t$ en cada paso. Por eso no puede hacer ninguna bien.

Separar los tres trabajos empieza por separar los dos estados. Junto al $\mathbf{h}_t$ que ya conoces aparece un segundo vector, el estado de celda $\mathbf{c}_t \in \mathbb{R}^{d_h}$, que es la memoria propiamente dicha, y su forma de avanzar es lo único verdaderamente nuevo de esta lección. Todo lo demás de la LSTM existe para gobernar ese vector.

Tres compuertas, cada una por un fallo de la anterior

Empecemos por la vía aditiva sola y vayamos añadiendo lo que le falte. Llamemos candidato al vector que el paso $t$ propone guardar, una capa corriente a partir del token y del estado anterior:

y que la memoria lo vaya acumulando: $\mathbf{c}_t = \mathbf{c}_{t-1} + \tilde{\mathbf{c}}_t$. Para el gradiente esto es perfecto —la derivada de $\mathbf{c}_t$ respecto de $\mathbf{c}_{t-1}$ es la identidad, sin matriz que aplicar $d$ veces—, y para la memoria es inservible por dos motivos. Nunca olvida: lo que entró en la coordenada 3 leyendo cámaras sigue ahí trescientos tokens después, en un párrafo que habla de otra cosa. Y crece sin freno: cada término vive en $(-1, 1)$ y se suman $T$, así que una coordenada puede acabar valiendo $40$, donde el $\tanh$ que la lea está plano —la saturación que el explorador de activaciones de la lección sobre funciones de activaciónBloque 2 · Lección 2Funciones de activación y no linealidadEl escalón deja pasar el signo de la preactivación y tira lo demás, y su derivada vale cero allí donde existe. Sigmoide, tangente hiperbólica y ReLU devuelven la magnitud y admiten derivada; ninguna puede ser una recta, porque encadenar funciones afines vuelve a dar una función afín. ya te dejó ver—.

Los dos fallos son el mismo: la suma no tiene forma de decir no. Dásela con un factor por coordenada, entre $0$ y $1$, que la red calcule en cada paso a partir de lo que está leyendo. Ésa es la compuerta de olvido:

La activación es la sigmoide logística y no la $\tanh$ del candidato, por lo mismo que la capa de salida de la lección sobre la RNN vanillaBloque 3 · Lección 2La RNN vanillaUna sola matriz aplicada en todas las posiciones y un estado que viaja de un paso al siguiente. Con eso, ninguna de las formas de los pesos menciona la longitud del texto: la misma red lee dos tokens o cuarenta sin rellenar ni recortar nada, lo aprendido leyendo la posición 2 sirve tal cual en la 8, y dos frases con los mismos tokens en distinto orden dejan de ser la misma entrada.: hace falta una fracción de lo que había, y las fracciones van de $0$ a $1$. El $\odot$ pide leerse igual de despacio. Multiplicar coordenada a coordenada deja que la 3 se conserve entera mientras la 7 se borra en el mismo paso; una matriz mezclaría las dos, y mezclar es lo que menos le conviene a una memoria.

Queda el otro lado. Esta versión escribe $\tilde{\mathbf{c}}_t$ en cada paso, valga lo que valga el token: que y del escriben tanto como cámaras, y lo escrito sólo se deshace cerrando $\mathbf{f}_t$, que borra de paso lo que esa coordenada guardara de antes. Poner una condición a la escritura pide otro factor igual, la compuerta de entrada:

Y falta decir qué ve el resto de la red. Si el estado oculto fuera $\mathbf{h}_t = \tanh(\mathbf{c}_t)$ —toda la memoria, siempre—, el plural de cámaras gobernaría la respuesta en las doce posiciones del inciso, cuando lo que hace falta es que espere. Guardar algo para luego y no usarlo ahora son dos cosas distintas, así que hay una tercera compuerta, la de salida, y con ella el estado oculto es una vista parcial de la memoria:

Ahí está lo que la RNN vanilla no podía hacer: su $\mathbf{h}_t$ tenía dos trabajos a la vez, recordar y enseñar; la LSTM los reparte, y $\mathbf{c}_t$ recuerda mientras $\mathbf{h}_t$ enseña.

Con las cuatro piezas juntas —tres compuertas y un candidato— el paso completo queda así, con $\mathbf{h}_0 = \mathbf{c}_0 = \mathbf{0}$ por la misma razón que en la RNN:

Cuatro líneas idénticas y dos que no. Las de arriba son la recurrencia de la lección sobre la RNN vanillaBloque 3 · Lección 2La RNN vanillaUna sola matriz aplicada en todas las posiciones y un estado que viaja de un paso al siguiente. Con eso, ninguna de las formas de los pesos menciona la longitud del texto: la misma red lee dos tokens o cuarenta sin rellenar ni recortar nada, lo aprendido leyendo la posición 2 sirve tal cual en la 8, y dos frases con los mismos tokens en distinto orden dejan de ser la misma entrada. copiada cuatro veces, con $\mathbf{W}_{x\ast} \in \mathbb{R}^{d_h \times d_{\text{model}}}$, $\mathbf{W}_{h\ast} \in \mathbb{R}^{d_h \times d_h}$ y $\mathbf{b}_{\ast} \in \mathbb{R}^{d_h}$ en todas ellas, y ahí está el precio: la LSTM cuesta $4\left(d_h \cdot d_h + d_h \cdot d_{\text{model}} + d_h\right)$ pesos, cuatro veces la RNN vanilla. Lo que no cuesta es longitud —tampoco aquí aparece $T$ en ninguna forma—.

Ya está toda la maquinaria; muévela antes de derivarla. Recorre la secuencia y mira las tres compuertas abrirse y cerrarse en cada posición; luego sube el sesgo $\mathbf{b}_f$ de la de olvido y observa la última fila, lo que sobrevive de $\mathbf{c}_1$.

Por qué el gradiente sobrevive a la vía aditiva

Que la celda «suma» ya se ha dicho tres veces; ahora hay que cobrarlo. Mira la penúltima línea coordenada a coordenada, que es donde se ve lo que el $\odot$ significa:

La coordenada $j$ de la celda nueva depende de la coordenada $j$ de la vieja y de ninguna otra: por esta vía las coordenadas no se mezclan, porque no hay matriz que las mezcle. Derivando, y dejando quietas las compuertas —volveré sobre esto—, el camino directo de un paso al anterior vale

un número, no una matriz. Encadénalo desde el final hasta un paso lejano $k$, igual que la lección anterior encadenaba la recurrencia hacia atrás, y lo que llega a la celda de $k$ por esa vía es

Ponlo al lado del $\left(\gamma\,\sigma_{\max}\right)^{T-k}$ de la lección anterior y cuenta las diferencias, que son tres. La base ya no es una propiedad de una matriz compartida por todo, sino un número por coordenada y por paso que la red calcula a partir de lo que está leyendo. Por esta vía tampoco hay máscara del $\tanh$ que muerda en cada paso —ningún $\gamma \le 1$ que sólo pueda empeorar las cosas—, porque la celda cruza el paso con un producto y una suma y nada más. Y la tercera remata: $f_{t,j} = 1$ es alcanzable y además estable, porque una sigmoide saturada se queda en $1$ con pendiente casi nula, así que una coordenada que la red aprenda a mantener abierta multiplica por $1$ tantas veces como haga falta. El $\rho = 1$ de la RNN era un filo compartido por todas las coordenadas; esto es una decisión de cada una.

Dos concesiones cierran la cuenta, porque sin ellas esto se lee como una garantía. Si la red aprende $f_{t,j}$ alrededor de $0.7$, el producto se hunde igual que antes; la diferencia es que se hunde porque la red lo ha decidido y no porque la arquitectura lo imponga. Y la explosión sigue siendo posible por los caminos del recuadro de arriba, así que el recorte del gradiente de la lección anterior se sigue usando tal cual.

La LSTM, paso a paso, en NumPy

Dos celdas. La primera transcribe el recuadro de seis líneas y lo ejecuta sobre una secuencia corta; la segunda mide lo que sobrevive de un empujón dado al principio de una larga.

Los pesos me los da un generador con semilla fija y no están entrenados, así que las compuertas no significan nada todavía. Lo que sí he puesto a mano es el sesgo de la de olvido, y la celda lo corre con dos valores distintos: es la forma más corta de ver que $\mathbf{b}_f$ decide cuánto recuerda la red antes de aprender nada. Mira la última línea de cada bloque.

Las tres compuertas viven en $(0,1)$ y ninguna vale lo mismo en dos coordenadas, que es todo lo que la fórmula prometía. Compara ahora los dos bloques. Con $\mathbf{b}_f = 1$ la de olvido se mueve alrededor de $0.7$ y de $\mathbf{c}_1$ llega a $\mathbf{c}_6$ un $20\,\%$ escaso; con $\mathbf{b}_f = 3$ ronda $0.95$ y llega un $78\,\%$, con los mismos pesos y la misma frase. Cinco pasos bastan para que un sesgo separe una memoria de un colador, y de ahí la costumbre de inicializar $\mathbf{b}_f$ positivo: sale más barato que la red aprenda a olvidar que aprenda a recordar. La primera fila de $\mathbf{c}_t$, en cambio, es idéntica en los dos: en $t = 1$ la compuerta multiplica $\mathbf{c}_0 = \mathbf{0}$.

La segunda celda mide la vía de vuelta sin derivar nada: mueve el estado del primer paso una cantidad diminuta, deja correr treinta y nueve pasos más y compara. Es el sondeo numérico de la lección sobre BPTT (backpropagation through time)Bloque 3 · Lección 3Backpropagation through timeLa recurrencia comparte una matriz entre todas las posiciones, y eso reescribe su gradiente. Backpropagation through time es la pasada hacia atrás del bloque anterior con las capas tumbadas sobre el tiempo: un error por paso, transportado del final al principio por la misma W_hh, y un gradiente que —por reutilizar esa matriz en cada posición— es la suma de lo que cada paso le reclama. Se deriva completa y se comprueba, número a número, contra un sondeo numérico., aplicado a la distancia en vez de a un peso. Las compuertas de este juguete leen sólo el token —he desconectado el estado a mano— para medir la vía aditiva y nada más.

Cuatro números y tres lecturas. De la RNN llega $1.1 \times 10^{-3}$: $0.9^{39} \approx 0.016$ lo pone la matriz y el resto —un orden de magnitud— la máscara del $\tanh$, ese $\gamma \le 1$ que la lección anterior dijo que sólo podía empeorar las cosas. La segunda línea es la LSTM saliendo peor que la RNN: con $\mathbf{b}_f = 0$ la compuerta ronda $0.5$, el producto cae a $10^{-11}$ y la memoria es peor que la que venía a arreglar. La vía aditiva no regala nada, sólo hace alcanzable lo que en la RNN no lo era. Las dos últimas líneas son ese «alcanzable»: con la compuerta en $0.95$ sobrevive el $14\,\%$ y con ella en $0.993$, el $76\,\%$, a la distancia donde la RNN llevaba tres ceros tras el punto.

Comprueba tu intuición

Cinco preguntas: por qué la vía aditiva cambia el producto, qué hace cada compuerta en un caso concreto, cuánto cuesta la arquitectura, cómo se ve la diferencia en un número y qué es lo que la LSTM sigue sin garantizar.

for factor in [0.6, 0.99]:
    print(round(factor ** 40, 3))
 

Y el paso completo escrito por ti, que es la pieza sobre la que se monta todo lo que viene después.

La LSTM resuelve el problema de la lección anterior y cobra un precio que la cuenta de pesos sólo dice a medias. Ocho matrices, dos estados que arrastrar y tres compuertas que ajustar: cada pieza llegó aquí arreglando un fallo concreto, sí, pero ninguna con una prueba de que hiciera falta. Y de al menos una se puede sospechar. La compuerta de entrada decide cuánto se escribe justo cuando la de olvido acaba de decidir cuánto se conserva, y las dos preguntas son la misma vista desde los dos lados: si vas a guardar lo nuevo, es porque lo viejo te importa menos.

Alguien tiró de ese hilo. Fundir esas dos compuertas en una, borrar la de salida junto con el estado de celda entero y quedarse con tres cuartas partes de los pesos —conservando casi todo lo que esta lección acaba de comprar— es la siguiente lección, sobre la GRU (gated recurrent unit): la versión simplificada.