7 - Recurrent networks: the RNN and the LSTM, from scratch

Before the Transformer, sequences were processed with recurrent neural networks (RNNs). In chapter 7 Emma explains how they read a text token by token, carrying a hidden state from one token to the next, and why they hit a wall: in long sequences the information of the first tokens fades, because the gradient that should carry it back vanishes (or explodes). The LSTM, with its gates, was the first remedy. As Noah puts it, it is like a doctor who remembers clearly the last words of the consultation and ever more vaguely the first ones.

In this notebook we write an RNN and its training algorithm, backpropagation through time, and we measure that wall:

1. Import the libraries.
2. An RNN, step by step.
3. A task that needs memory.
4. Backpropagation through time, in NumPy.
5. Training the RNN.
6. The gradient that fades.
7. The LSTM and its gates.
8. How far can each network remember?

Step 1: Import the libraries

---

The RNN of steps 2 to 6 uses NumPy only. PyTorch checks our gradients in step 4 and trains the RNN and the LSTM of step 8.

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

Step 2: An RNN, step by step

---

The book's equations. The RNN starts with a hidden state $h_0$ of zeros and, for each token $x_t$, mixes it with the previous state and produces an output:

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

The same three matrices are used at every step: that is the parameter sharing across time of the book. We apply an RNN with random weights and a hidden state of 4 numbers to the book's sentence. Each token enters as a one-hot vector (a 1 in its position of the vocabulary), so $W_{xh}\,x_t$ is simply the row of $W_{xh}$ for that token.

sentence = ['The', 'patient', 'has', 'a', 'critical', 'condition']
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)        # one-hot input: the row of W_xh for this word
    print(f"h_{t} after {word!r:<12}", np.round(h, 3))

h_1 after 'The'        [ 0.063 -0.066  0.31   0.052]
h_2 after 'patient'    [-0.26   0.251  0.614  0.497]
h_3 after 'has'        [-0.628 -0.449 -0.081  0.533]
h_4 after 'a'          [-0.907 -0.24   0.038  0.34 ]
h_5 after 'critical'   [-0.63  -0.235  0.665  0.841]
h_6 after 'condition'  [-0.553  0.637  0.398  0.845]

Each hidden state mixes the new word with the previous state, so the last one depends on the whole sentence. With random weights these numbers mean nothing yet: training will adjust $W_{xh}$, $W_{hh}$ and $b_h$ so that the hidden state keeps what the task needs. Note also that every step needs the result of the previous one, so the steps cannot be computed at the same time.

Step 3: A task that needs memory

---

A simplified prescribing rule: an antibiotic from the penicillin family, amoxicillin, unless the patient is allergic to penicillin, in which case azithromycin. Each sequence starts with the allergy, allergy: penicillin or allergy: none, continues with $k$ words of clinical notes taken at random, and ends with prescribe:. The network must answer with the right drug, so it has to carry one piece of information, the allergy, across the $k$ words. By changing $k$ we control how far back it has to remember.

filler = ['the', 'patient', 'was', 'seen', 'today', 'in', 'clinic', 'with', 'a', 'mild', 'cough', 'and', 'no', 'fever',
          'she', 'reports', 'good', 'sleep', 'appetite', 'is', 'normal']
vocab = ['allergy:', 'penicillin', 'none', 'prescribe:', 'amoxicillin', 'azithromycin'] + filler
index = {w: i for i, w in enumerate(vocab)}

def make_sequences(n, k, rng):
    """n sequences with k words of notes between the allergy and the question. Returns token ids and the right drug."""
    allergic = rng.random(n) < 0.5
    X = np.array([[index['allergy:'], index['penicillin' if a else 'none']]
                  + [index[w] for w in rng.choice(filler, k)] + [index['prescribe:']] for a in allergic])
    y = np.where(allergic, index['azithromycin'], index['amoxicillin'])
    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"\nVocabulary: {len(vocab)} words")

allergy: none sleep is in clinic appetite a prescribe: -> amoxicillin
allergy: none in sleep in a fever and prescribe: -> amoxicillin
allergy: penicillin patient the appetite reports sleep and prescribe: -> azithromycin

Vocabulary: 27 words

Step 4: Backpropagation through time, in NumPy

---

The network reads the whole sequence and answers after the last token: $\hat{y} = \text{softmax}(W_{hy}\, h_T + b_y)$, with the cross-entropy as the loss. To train it we need the gradient of that loss with respect to the three shared matrices. Backpropagation through time (BPTT) unrolls the network into one layer per token and applies ordinary backpropagation to it, from the last step to the first:

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

where $1 - h_t^2$ is the derivative of tanh. Because the matrices are shared, their gradients are the sums of the contributions of every step. rnn_backward also records the size of $\partial L / \partial h_t$ at every step, which we will need in step 6. We check the result against PyTorch's automatic differentiation on the same weights.

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):                                # from the last token back to the first
        sizes.append(np.linalg.norm(d_h, axis=1).mean())
        d_a = d_h * (1 - states[t] ** 2)                     # back through tanh
        np.add.at(grads['W_xh'], X[:, t - 1], d_a)           # shared matrices: add up every step
        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                              # to the previous hidden state
    return grads, sizes[::-1]

# Check against PyTorch's autograd
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"Largest difference with PyTorch's gradients: {max(np.abs(grads[k] - tp[k].grad.numpy()).max() for k in p):.1e}")

Largest difference with PyTorch's gradients: 1.4e-16

The largest difference is 1.4e-16, the size of rounding errors in double precision: our backpropagation through time computes the same gradients as PyTorch.

Step 5: Training the RNN

---

Mini-batch gradient descent, with batches of 32 new sequences and $k = 5$ words between the allergy and the question. The book warns that in RNNs the gradients can also explode, so we add the usual protection, gradient clipping: if the size of the whole gradient is larger than 1, it is scaled down to 1 before the update.

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]          # gradient clipping at size 1
    if step % 500 == 0:
        accuracy = np.mean(rnn_forward(p, X_test)[0].argmax(axis=1) == y_test)
        print(f"step {step}: loss {-np.mean(np.log(probs[np.arange(len(y)), y])):.3f}, test accuracy {accuracy:.3f}")

step 500: loss 0.691, test accuracy 0.468

step 1000: loss 0.076, test accuracy 0.997

step 1500: loss 0.002, test accuracy 1.000

step 2000: loss 0.001, test accuracy 1.000

For the first 500 steps the loss stays at 0.691, the loss of guessing between two drugs ($\ln 2 = 0.693$), and the accuracy at 0.468. Then the network finds the allergy: at step 1,000 its accuracy is 0.997, and from step 1,500 it is 1.000. With 5 words in between, the RNN learns the task.

Step 6: The gradient that fades

---

The only way the network can learn to keep the allergy is through the gradient that travels from the answer back to the second token. At each step back, it is multiplied by $W_{hh}^\top$ and by the derivative of tanh, which is at most 1. We measure its size at every position of sequences with $k = 40$, for the untrained network and for the one trained in step 5.

X, y = make_sequences(200, 40, np.random.default_rng(5))
plt.figure(figsize=(9, 4))
for name, params in [('untrained', init_rnn(seed=0)), ('trained with 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} gradient size at the last token {sizes[-1]:.1e}, at the allergy (token 2) {sizes[1]:.1e}, "
          f"ratio {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('position in the sequence (2 = the allergy, 43 = the question)'); plt.ylabel('size of dL/dh_t')
plt.title('Backpropagation through time, k = 40'); plt.legend(); plt.grid(True); plt.show()

untrained            gradient size at the last token 4.0e-03, at the allergy (token 2) 3.1e-08, ratio 1.3e+05
trained with k = 5   gradient size at the last token 1.6e-05, at the allergy (token 2) 4.6e-21, ratio 3.5e+15

In the untrained network the gradient that reaches the allergy is 3.1e-08, against 4.0e-03 at the question: 1.3e+05 times smaller. On a logarithmic scale its curve is almost a straight line, which is what an exponential decay looks like. In the network trained with $k = 5$ the ratio is even larger, 3.5e+15. With 40 words in between, the loss has practically no way of telling the first steps what to keep: that is the vanishing gradient the book describes.

Step 7: The LSTM and its gates

---

The LSTM (Long Short-Term Memory) adds to the hidden state a second memory, the cell state $c_t$, and three gates, numbers between 0 and 1 computed with a sigmoid from $x_t$ and $h_{t-1}$, that decide what to forget, what to write and what to show:

$$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)$$

The key is the first term of $c_t$. Going back from $c_t$ to $c_{t-1}$ the gradient is multiplied only by the forget gate $f_t$, not by a matrix and a tanh. If the gate stays close to 1, the information, and its gradient, can cross many steps almost intact. PyTorch's nn.LSTM computes the four blocks with one matrix. Its initial biases are small random numbers, so the forget gate starts at about $\sigma(0) = 0.5$: the memory halves at every step until the network learns otherwise. A common practice is to start the forget bias at a positive value; we try 3, which opens the gate to $\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 orders the LSTM blocks as input, forget, cell, output
            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])                     # the answer after the last 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("PyTorch forget gate at the start: sigmoid(0) =", round(float(torch.sigmoid(torch.tensor(0.0))), 2),
      "| with bias 3: sigmoid(3) =", round(float(torch.sigmoid(torch.tensor(3.0))), 3))

PyTorch forget gate at the start: sigmoid(0) = 0.5 | with bias 3: sigmoid(3) = 0.953

Step 8: How far can each network remember?

---

We train three networks with exactly the same budget, 1,000 steps of Adam (a variant of gradient descent) with batches of 64 sequences and gradient clipping, for distances $k$ from 5 to 80 words, and measure the accuracy on 1,000 new sequences. Since half of the patients are allergic, 0.5 is the accuracy of guessing. The whole experiment takes a minute or two on a laptop.

distances = [5, 10, 20, 40, 80]
networks = {'RNN': ('RNN', None), 'LSTM, PyTorch initialization': ('LSTM', None), 'LSTM, forget bias 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 (words in between)':<30}" + "".join(f"{k:>7}" for k in distances))
for name, accuracies in results.items():
    print(f"{name:<30}" + "".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)   # the first two coincide
plt.axhline(0.5, color='gray', ls='--', label='guessing')
plt.xticks(distances, distances); plt.xlabel('words between the allergy and the question'); plt.ylabel('test accuracy')
plt.legend(); plt.grid(True); plt.show()

k (words in between)                5     10     20     40     80
RNN                             1.000  0.468  0.532  0.468  0.532
LSTM, PyTorch initialization    1.000  0.468  0.532  0.468  0.532
LSTM, forget bias 3             1.000  1.000  1.000  1.000  0.532

With 5 words in between, the three networks learn the task perfectly. From 10 words on, the RNN stays at 0.468 or 0.532, the share of one drug in the test set: it always answers the same thing, which is guessing. So does the LSTM with PyTorch's initialization, whose forget gate starts half closed. With the forget gate open from the start, the LSTM gets all 1,000 test sequences right up to $k = 40$. At 80 words it also falls back to guessing with this training budget. As the book says, the gates were a significant improvement: they move the wall much further, but they do not remove it.

Conclusion

---

• An RNN reads a sequence one token at a time, with the same weights at every step, and summarizes everything it has read in its hidden state.
• It is trained with backpropagation through time: the network is unrolled over the sequence and the gradients of the shared weights are added up over all the steps.
• On its way back, the gradient is multiplied at every step by the recurrent matrix and by the derivative of tanh, and it fades exponentially: what happened many tokens ago barely influences learning. Gradient clipping protects against the opposite problem, gradients that explode.
• The LSTM's cell state, controlled by the forget gate, gives the information and its gradient a path across time. It remembers much further, but not indefinitely, and it still has to read the tokens one after another.
• That sequential reading, which cannot be parallelized, is the wall the Transformer broke: with attention, every token looks directly at every other one. The notebook Attention and the Transformer builds it.