7 - Redes recurrentes: la RNN y la LSTM, desde cero

Antes del Transformer, las secuencias se procesaban con redes neuronales recurrentes (RNN). En el capítulo 7, Ana explica cómo leen un texto token a token, arrastrando un estado oculto de un token al siguiente, y por qué chocaron con un muro: en las secuencias largas, la información de los primeros tokens se desvanece, porque el gradiente que debería llevarla hacia atrás se desvanece (o explota). La LSTM, con sus puertas, fue el primer remedio. Como dice Luis, es como un médico que recuerda con claridad las últimas palabras de la consulta y cada vez más vagamente las primeras.

En este cuaderno escribimos una RNN y su algoritmo de entrenamiento, la retropropagación a través del tiempo, y medimos ese muro:

1. Importar las bibliotecas.
2. Una RNN, paso a paso.
3. Una tarea que necesita memoria.
4. La retropropagación a través del tiempo, en NumPy.
5. Entrenar la RNN.
6. El gradiente que se desvanece.
7. La LSTM y sus puertas.
8. ¿Hasta dónde recuerda cada red?

Paso 1: Importar las bibliotecas

---

La RNN de los pasos 2 a 6 solo usa NumPy. PyTorch comprueba nuestros gradientes en el paso 4 y entrena la RNN y la LSTM del paso 8.

import numpy as np
import matplotlib.pyplot as plt
import torch

Paso 2: Una RNN, paso a paso

---

Las ecuaciones del libro. La RNN empieza con un estado oculto $h_0$ de ceros y, para cada token $x_t$, lo mezcla con el estado anterior y produce una salida:

$$h_t = \tanh(W_{xh}\, x_t + W_{hh}\, h_{t-1} + b_h) \qquad\qquad y_t = W_{hy}\, h_t + b_y$$

Se usan las mismas tres matrices en todos los pasos: es la compartición de parámetros a través del tiempo del libro. Aplicamos una RNN con pesos aleatorios y un estado oculto de 4 números a la frase del libro. Cada token entra como un vector one-hot (un 1 en su posición del vocabulario), así que $W_{xh}\,x_t$ es simplemente la fila de $W_{xh}$ de ese token.

sentence = ['La', 'paciente', 'tiene', 'fiebre', 'alta']
rng = np.random.default_rng(0)
W_xh, W_hh, b_h = rng.normal(0, 0.5, (len(sentence), 4)), rng.normal(0, 0.5, (4, 4)), np.zeros(4)

h = np.zeros(4)                                      # h_0
for t, word in enumerate(sentence, start=1):
    h = np.tanh(W_xh[t - 1] + h @ W_hh + b_h)        # entrada one-hot: la fila de W_xh de esta palabra
    print(f"h_{t} tras {word!r:<12}", np.round(h, 3))

h_1 tras 'La'         [ 0.063 -0.066  0.31   0.052]
h_2 tras 'paciente'   [-0.36   0.263  0.466  0.456]
h_3 tras 'tiene'      [-0.339 -0.599 -0.443 -0.131]
h_4 tras 'fiebre'     [-0.862 -0.425 -0.078 -0.126]
h_5 tras 'alta'       [-0.363 -0.669  0.589  0.501]

Cada estado oculto mezcla la palabra nueva con el estado anterior, así que el último depende de toda la frase. Con pesos aleatorios, estos números todavía no significan nada: el entrenamiento ajustará $W_{xh}$, $W_{hh}$ y $b_h$ para que el estado oculto guarde lo que necesita la tarea. Fíjate además en que cada paso necesita el resultado del anterior, así que los pasos no se pueden calcular a la vez.

Paso 3: Una tarea que necesita memoria

---

Una regla de prescripción simplificada: un antibiótico de la familia de la penicilina, la amoxicilina, salvo que el paciente sea alérgico a la penicilina, en cuyo caso azitromicina. Cada secuencia empieza con la alergia, alergia: penicilina o alergia: ninguna, sigue con $k$ palabras de notas clínicas tomadas al azar y termina con prescribir:. La red debe responder con el fármaco correcto, así que tiene que llevar un dato, la alergia, a través de las $k$ palabras. Cambiando $k$ controlamos hasta dónde tiene que recordar.

filler = ['la', 'paciente', 'acude', 'hoy', 'a', 'consulta', 'con', 'tos', 'leve', 'y', 'sin', 'fiebre', 'refiere',
          'buen', 'sueño', 'el', 'apetito', 'es', 'normal', 'duerme', 'bien']
vocab = ['alergia:', 'penicilina', 'ninguna', 'prescribir:', 'amoxicilina', 'azitromicina'] + filler
index = {w: i for i, w in enumerate(vocab)}

def make_sequences(n, k, rng):
    """n secuencias con k palabras de notas entre la alergia y la pregunta. Devuelve los ids de los tokens y el fármaco correcto."""
    allergic = rng.random(n) < 0.5
    X = np.array([[index['alergia:'], index['penicilina' if a else 'ninguna']]
                  + [index[w] for w in rng.choice(filler, k)] + [index['prescribir:']] for a in allergic])
    y = np.where(allergic, index['azitromicina'], index['amoxicilina'])
    return X, y

X, y = make_sequences(3, 6, np.random.default_rng(1))
for seq, target in zip(X, y):
    print(" ".join(vocab[i] for i in seq), "->", vocab[target])
print(f"\nVocabulario: {len(vocab)} palabras")

alergia: ninguna es duerme consulta con normal leve prescribir: -> amoxicilina
alergia: ninguna consulta es consulta leve buen fiebre prescribir: -> amoxicilina
alergia: penicilina paciente la normal el es fiebre prescribir: -> azitromicina

Vocabulario: 27 palabras

Paso 4: La retropropagación a través del tiempo, en NumPy

---

La red lee toda la secuencia y responde tras el último token: $\hat{y} = \text{softmax}(W_{hy}\, h_T + b_y)$, con la entropía cruzada como pérdida. Para entrenarla necesitamos el gradiente de esa pérdida respecto a las tres matrices compartidas. La retropropagación a través del tiempo (BPTT) despliega la red en una capa por token y le aplica la retropropagación de siempre, del último paso al primero:

$$\frac{\partial L}{\partial h_{t-1}} = \Big(\frac{\partial L}{\partial h_t} \odot (1 - h_t^2)\Big)\, W_{hh}^\top$$

donde $1 - h_t^2$ es la derivada de tanh. Como las matrices son compartidas, sus gradientes son las sumas de las contribuciones de todos los pasos. rnn_backward también anota el tamaño de $\partial L / \partial h_t$ en cada paso, que necesitaremos en el paso 6. Comprobamos el resultado con la diferenciación automática de PyTorch sobre los mismos pesos.

V, H = len(vocab), 16

def init_rnn(seed):
    r = np.random.default_rng(seed)
    return {'W_xh': r.normal(0, 0.5, (V, H)), 'W_hh': r.normal(0, 1 / np.sqrt(H), (H, H)), 'b_h': np.zeros(H),
            'W_hy': r.normal(0, 1 / np.sqrt(H), (H, V)), 'b_y': np.zeros(V)}

def rnn_forward(p, X):
    h = np.zeros((len(X), H))
    states = [h]
    for t in range(X.shape[1]):
        h = np.tanh(p['W_xh'][X[:, t]] + h @ p['W_hh'] + p['b_h'])
        states.append(h)
    scores = h @ p['W_hy'] + p['b_y']
    scores -= scores.max(axis=1, keepdims=True)
    probs = np.exp(scores) / np.exp(scores).sum(axis=1, keepdims=True)
    return probs, states

def rnn_backward(p, X, y, probs, states):
    n, T = X.shape
    grads = {key: np.zeros_like(value) for key, value in p.items()}
    d_scores = probs.copy()
    d_scores[np.arange(n), y] -= 1
    d_scores /= n
    grads['W_hy'], grads['b_y'] = states[-1].T @ d_scores, d_scores.sum(axis=0)
    d_h = d_scores @ p['W_hy'].T
    sizes = []
    for t in range(T, 0, -1):                                # del último token hacia atrás hasta el primero
        sizes.append(np.linalg.norm(d_h, axis=1).mean())
        d_a = d_h * (1 - states[t] ** 2)                     # hacia atrás por tanh
        np.add.at(grads['W_xh'], X[:, t - 1], d_a)           # matrices compartidas: se suman todos los pasos
        grads['W_hh'] += states[t - 1].T @ d_a
        grads['b_h'] += d_a.sum(axis=0)
        d_h = d_a @ p['W_hh'].T                              # al estado oculto anterior
    return grads, sizes[::-1]

# Comprobación con la diferenciación automática de PyTorch
p = init_rnn(seed=1)
X, y = make_sequences(8, 6, np.random.default_rng(3))
probs, states = rnn_forward(p, X)
grads, _ = rnn_backward(p, X, y, probs, states)
tp = {key: torch.tensor(value, requires_grad=True) for key, value in p.items()}
h = torch.zeros(len(X), H, dtype=torch.float64)
for t in range(X.shape[1]):
    h = torch.tanh(tp['W_xh'][X[:, t]] + h @ tp['W_hh'] + tp['b_h'])
torch.nn.functional.cross_entropy(h @ tp['W_hy'] + tp['b_y'], torch.tensor(y)).backward()
print(f"Mayor diferencia con los gradientes de PyTorch: {max(np.abs(grads[k] - tp[k].grad.numpy()).max() for k in p):.1e}")

Mayor diferencia con los gradientes de PyTorch: 1.4e-16

La mayor diferencia es 1,4e-16, del tamaño de los errores de redondeo en doble precisión: nuestra retropropagación a través del tiempo calcula los mismos gradientes que PyTorch.

Paso 5: Entrenar la RNN

---

Descenso del gradiente por mini-lotes, con lotes de 32 secuencias nuevas y $k = 5$ palabras entre la alergia y la pregunta. El libro advierte de que en las RNN los gradientes también pueden explotar, así que añadimos la protección habitual, el recorte del gradiente (gradient clipping): si el tamaño del gradiente completo es mayor que 1, se reduce a 1 antes de la actualización.

p = init_rnn(seed=0)
train_rng = np.random.default_rng(0)
X_test, y_test = make_sequences(1000, 5, np.random.default_rng(99))
for step in range(1, 2001):
    X, y = make_sequences(32, 5, train_rng)
    probs, states = rnn_forward(p, X)
    grads, _ = rnn_backward(p, X, y, probs, states)
    size = np.sqrt(sum((g ** 2).sum() for g in grads.values()))
    for key in p:
        p[key] -= 0.1 * min(1.0, 1.0 / size) * grads[key]          # recorte del gradiente a tamaño 1
    if step % 500 == 0:
        accuracy = np.mean(rnn_forward(p, X_test)[0].argmax(axis=1) == y_test)
        print(f"paso {step}: pérdida {-np.mean(np.log(probs[np.arange(len(y)), y])):.3f}, exactitud en prueba {accuracy:.3f}")

paso 500: pérdida 0.691, exactitud en prueba 0.468

paso 1000: pérdida 0.076, exactitud en prueba 0.997

paso 1500: pérdida 0.002, exactitud en prueba 1.000

paso 2000: pérdida 0.001, exactitud en prueba 1.000

Durante los primeros 500 pasos, la pérdida se queda en 0,691, la pérdida de adivinar entre dos fármacos ($\ln 2 = 0{,}693$), y la exactitud en 0,468. Después, la red encuentra la alergia: en el paso 1.000 su exactitud es 0,997, y a partir del paso 1.500, 1,000. Con 5 palabras en medio, la RNN aprende la tarea.

Paso 6: El gradiente que se desvanece

---

La única forma de que la red aprenda a guardar la alergia es el gradiente que viaja desde la respuesta hasta el segundo token. En cada paso hacia atrás se multiplica por $W_{hh}^\top$ y por la derivada de tanh, que como mucho vale 1. Medimos su tamaño en cada posición de secuencias con $k = 40$, para la red sin entrenar y para la entrenada en el paso 5.

X, y = make_sequences(200, 40, np.random.default_rng(5))
plt.figure(figsize=(9, 4))
for name, params in [('sin entrenar', init_rnn(seed=0)), ('entrenada con k = 5', p)]:
    probs, states = rnn_forward(params, X)
    _, sizes = rnn_backward(params, X, y, probs, states)
    sizes = np.array(sizes)
    print(f"{name:<20} tamaño del gradiente en el último token {sizes[-1]:.1e}, en la alergia (token 2) {sizes[1]:.1e}, "
          f"cociente {sizes[-1] / sizes[1]:.1e}")
    plt.semilogy(range(1, len(sizes) + 1), sizes, 'o-', ms=3, label=name)
plt.axvline(2, color='gray', ls='--')
plt.xlabel('posición en la secuencia (2 = la alergia, 43 = la pregunta)'); plt.ylabel('tamaño de dL/dh_t')
plt.title('Retropropagación a través del tiempo, k = 40'); plt.legend(); plt.grid(True); plt.show()

sin entrenar         tamaño del gradiente en el último token 4.0e-03, en la alergia (token 2) 3.1e-08, cociente 1.3e+05
entrenada con k = 5  tamaño del gradiente en el último token 1.6e-05, en la alergia (token 2) 4.6e-21, cociente 3.5e+15

En la red sin entrenar, el gradiente que llega a la alergia es 3,1e-08, frente a 4,0e-03 en la pregunta: 1,3e+05 veces más pequeño. En escala logarítmica su curva es casi una línea recta, que es el aspecto de un decaimiento exponencial. En la red entrenada con $k = 5$ el cociente es todavía mayor, 3,5e+15. Con 40 palabras en medio, la pérdida prácticamente no tiene forma de decir a los primeros pasos qué deben guardar: es el gradiente que se desvanece del que habla el libro.

Paso 7: La LSTM y sus puertas

---

La LSTM (Long Short-Term Memory) añade al estado oculto una segunda memoria, el estado de celda $c_t$, y tres puertas, números entre 0 y 1 calculados con una sigmoide a partir de $x_t$ y $h_{t-1}$, que deciden qué olvidar, qué escribir y qué mostrar:

$$f_t = \sigma(W_f [h_{t-1}, x_t] + b_f) \qquad i_t = \sigma(W_i [h_{t-1}, x_t] + b_i) \qquad o_t = \sigma(W_o [h_{t-1}, x_t] + b_o)$$

$$c_t = f_t \odot c_{t-1} + i_t \odot \tanh(W_c [h_{t-1}, x_t] + b_c) \qquad\qquad h_t = o_t \odot \tanh(c_t)$$

La clave es el primer término de $c_t$. Al volver de $c_t$ a $c_{t-1}$, el gradiente solo se multiplica por la puerta de olvido $f_t$, no por una matriz y una tanh. Si la puerta se mantiene cerca de 1, la información, y su gradiente, pueden cruzar muchos pasos casi intactos. nn.LSTM de PyTorch calcula los cuatro bloques con una sola matriz. Sus sesgos iniciales son pequeños números aleatorios, así que la puerta de olvido empieza en torno a $\sigma(0) = 0{,}5$: la memoria se reduce a la mitad en cada paso hasta que la red aprende otra cosa. Una práctica habitual es empezar el sesgo de olvido en un valor positivo; probamos con 3, que abre la puerta a $\sigma(3) = 0{,}95$.

class SequenceNet(torch.nn.Module):
    def __init__(self, kind, hidden=32, forget_bias=None):
        super().__init__()
        self.embed = torch.nn.Embedding(len(vocab), 16)
        self.rnn = {'RNN': torch.nn.RNN, 'LSTM': torch.nn.LSTM}[kind](16, hidden, batch_first=True)
        self.out = torch.nn.Linear(hidden, len(vocab))
        if forget_bias is not None:                        # PyTorch ordena los bloques de la LSTM así: entrada, olvido, celda, salida
            with torch.no_grad():
                self.rnn.bias_ih_l0[hidden:2 * hidden].fill_(forget_bias)
                self.rnn.bias_hh_l0[hidden:2 * hidden].fill_(0.0)

    def forward(self, x):
        states, _ = self.rnn(self.embed(x))
        return self.out(states[:, -1])                     # la respuesta tras el último token

def train_and_test(kind, k, forget_bias=None, steps=1000, seed=0):
    torch.manual_seed(seed)
    data_rng = np.random.default_rng(seed)
    net = SequenceNet(kind, forget_bias=forget_bias)
    optimizer = torch.optim.Adam(net.parameters(), lr=0.01)
    for _ in range(steps):
        X, y = make_sequences(64, k, data_rng)
        loss = torch.nn.functional.cross_entropy(net(torch.tensor(X)), torch.tensor(y))
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
        optimizer.step()
    X, y = make_sequences(1000, k, np.random.default_rng(99))
    with torch.no_grad():
        return (net(torch.tensor(X)).argmax(dim=1).numpy() == y).mean()

print("Puerta de olvido de PyTorch al principio: sigmoid(0) =", round(float(torch.sigmoid(torch.tensor(0.0))), 2),
      "| con sesgo 3: sigmoid(3) =", round(float(torch.sigmoid(torch.tensor(3.0))), 3))

Puerta de olvido de PyTorch al principio: sigmoid(0) = 0.5 | con sesgo 3: sigmoid(3) = 0.953

Paso 8: ¿Hasta dónde recuerda cada red?

---

Entrenamos tres redes exactamente con el mismo presupuesto, 1.000 pasos de Adam (una variante del descenso del gradiente) con lotes de 64 secuencias y recorte del gradiente, para distancias $k$ de 5 a 80 palabras, y medimos la exactitud con 1.000 secuencias nuevas. Como la mitad de los pacientes son alérgicos, 0,5 es la exactitud de adivinar al azar. Todo el experimento tarda uno o dos minutos en un portátil.

distances = [5, 10, 20, 40, 80]
networks = {'RNN': ('RNN', None), 'LSTM, inicialización de PyTorch': ('LSTM', None), 'LSTM, sesgo de olvido 3': ('LSTM', 3.0)}
results = {name: [train_and_test(kind, k, forget_bias=fb) for k in distances] for name, (kind, fb) in networks.items()}

print(f"{'k (palabras en medio)':<33}" + "".join(f"{k:>7}" for k in distances))
for name, accuracies in results.items():
    print(f"{name:<33}" + "".join(f"{a:>7.3f}" for a in accuracies))

plt.figure(figsize=(8, 4))
for (name, accuracies), style in zip(results.items(), ['o-', 's--', '^-']):
    plt.semilogx(distances, accuracies, style, ms=9 if style == 'o-' else 6, label=name)   # las dos primeras coinciden
plt.axhline(0.5, color='gray', ls='--', label='azar')
plt.xticks(distances, distances); plt.xlabel('palabras entre la alergia y la pregunta'); plt.ylabel('exactitud en prueba')
plt.legend(); plt.grid(True); plt.show()

k (palabras en medio)                  5     10     20     40     80
RNN                                1.000  0.468  0.532  0.468  0.532
LSTM, inicialización de PyTorch    1.000  0.468  0.532  0.468  0.532
LSTM, sesgo de olvido 3            1.000  1.000  1.000  1.000  0.532

Con 5 palabras en medio, las tres redes aprenden la tarea a la perfección. A partir de 10 palabras, la RNN se queda en 0,468 o 0,532, la proporción de uno de los fármacos en el conjunto de prueba: responde siempre lo mismo, que es adivinar. Igual le ocurre a la LSTM con la inicialización de PyTorch, cuya puerta de olvido empieza medio cerrada. Con la puerta de olvido abierta desde el principio, la LSTM acierta las 1.000 secuencias de prueba hasta $k = 40$. Con 80 palabras también vuelve a adivinar con este presupuesto de entrenamiento. Como dice el libro, las puertas fueron una mejora significativa: alejan mucho el muro, pero no lo eliminan.

Conclusión

---

• Una RNN lee una secuencia token a token, con los mismos pesos en todos los pasos, y resume todo lo que ha leído en su estado oculto.
• Se entrena con la retropropagación a través del tiempo: la red se despliega sobre la secuencia y los gradientes de los pesos compartidos se suman en todos los pasos.
• En su camino hacia atrás, el gradiente se multiplica en cada paso por la matriz recurrente y por la derivada de tanh, y se desvanece de forma exponencial: lo que ocurrió muchos tokens atrás apenas influye en el aprendizaje. El recorte del gradiente protege del problema contrario, los gradientes que explotan.
• El estado de celda de la LSTM, controlado por la puerta de olvido, da a la información y a su gradiente un camino a través del tiempo. Recuerda mucho más lejos, pero no indefinidamente, y sigue teniendo que leer los tokens uno detrás de otro.
• Esa lectura secuencial, que no se puede paralelizar, es el muro que rompió el Transformer: con la atención, cada token mira directamente a todos los demás. Lo construye el cuaderno La atención y el Transformer.