RNNs, LSTMs & GRUs

By the end of this lesson you will be able to explain how a recurrent neural network (RNN) reads a sequence one step at a time, see in numbers why a plain RNN forgets long-ago inputs, and understand how the gates in LSTM and GRU cells fix that. You will finish by training a small LSTM in PyTorch to predict a wave.

Many kinds of data are sequences, where order matters: the words of a sentence, the notes of a song, daily sales at a shop in Pune, readings from a heart-rate sensor. "Dog bites man" and "man bites dog" use the same words with very different meanings. The networks so far take a fixed-size input all at once. Recurrent networks take one item at a time and keep a memory.

The idea: a loop with memory

An RNN reads the sequence one step at a time. At each step it combines two things:

1. the current input x_t (for example, the current word's vector);
2. its hidden state h_{t-1}, a small vector that summarises everything it has read so far.

It mixes them with weights and a tanh to produce the new hidden state:

h_t = tanh(x_t @ W_x + h_{t-1} @ W_h + b)

The same weights W_x and W_h are used at every step, so the network can read sequences of any length.

Think of a cricket commentator on the radio. Ball after ball, they update a mental summary: the score, who is batting, how the bowler is doing. They do not replay the whole match before every sentence; they carry the summary forward. The hidden state is that running summary.

Step 1: run an RNN cell over a sequence. The task is just to watch the shapes and the memory. The input is a sequence of 6 steps with 2 features each.

import numpy as np

rng = np.random.default_rng(1)
n_in, n_hidden = 2, 4
W_x = rng.normal(0, 0.5, (n_in, n_hidden))
W_h = rng.normal(0, 0.5, (n_hidden, n_hidden))
b = np.zeros(n_hidden)

sequence = rng.normal(size=(6, n_in))       # 6 time steps, 2 features each
h = np.zeros(n_hidden)                       # empty memory at the start
states = []
for t, x_t in enumerate(sequence):
    h = np.tanh(x_t @ W_x + h @ W_h + b)     # same weights at every step
    states.append(h)
    print(f"step {t}: input {x_t.round(2)} -> hidden {h.round(3)}")

print("all hidden states:", np.array(states).shape, " final state:", h.shape)

step 0: input [-2.71 -1.89] -> hidden [-0.868 -0.911  0.059  0.839]
step 1: input [-0.17 -0.42] -> hidden [-0.039 -0.331  0.67  -0.104]
step 2: input [0.21 0.22] -> hidden [ 0.257  0.074 -0.267 -0.313]
step 3: input [ 2.12 -1.11] -> hidden [-0.124  0.627  0.49  -0.94 ]
step 4: input [-0.38  2.04] -> hidden [ 0.546  0.282 -0.916  0.428]
step 5: input [0.65 0.66] -> hidden [0.372 0.497 0.465 0.326]
all hidden states: (6, 4)  final state: (4,)

The network produced one hidden state per step, shape (6, 4). For a whole-sequence task such as "is this review positive?", you use the final state. For a per-step task such as tagging each word, you use all of them.

Training through time, and why plain RNNs forget

To train an RNN, you unroll the loop: step 1, step 2, ... step T become a very deep network whose layers share weights. Backpropagation then runs back through every step. This is called backpropagation through time.

The problem: the gradient that travels from the last step back to the first is multiplied by the RNN's slope at every step in between. If those factors are mostly below 1, the gradient shrinks exponentially, just like in a deep plain network (CNN Architectures showed the same effect). A plain RNN therefore struggles to learn that something at step 3 matters at step 50.

Step 2: measure how fast the memory fades. We compute how much a small change in the first hidden state still affects the hidden state after T steps.

import numpy as np

rng = np.random.default_rng(0)
n_hidden = 16
W_h = rng.normal(0, 1, (n_hidden, n_hidden)) / np.sqrt(n_hidden) * 0.9
h = rng.normal(0, 0.5, n_hidden)
J = np.eye(n_hidden)                          # d h_t / d h_0, starts as identity

for t in range(1, 51):
    x_part = rng.normal(0, 0.5, n_hidden)     # some input at every step
    h = np.tanh(h @ W_h + x_part)
    J = J @ W_h @ np.diag(1 - h ** 2)         # chain rule through one step
    if t in (1, 5, 10, 20, 50):
        print(f"after {t:>2} steps: influence of the first state = {np.linalg.norm(J):.2e}")

after  1 steps: influence of the first state = 2.53e+00
after  5 steps: influence of the first state = 5.75e-01
after 10 steps: influence of the first state = 1.97e-01
after 20 steps: influence of the first state = 1.30e-02
after 50 steps: influence of the first state = 3.36e-05

The influence of the first state falls by many orders of magnitude within a few dozen steps. Gradients through this path are just as tiny, so the network cannot learn long-range patterns. With larger weights, the opposite can happen and gradients explode; a common guard is gradient clipping, which caps the gradient's size before each update (torch.nn.utils.clip_grad_norm_ in PyTorch).

LSTM: a memory with gates

The LSTM (long short-term memory, Hochreiter and Schmidhuber, 1997) adds a second memory, the cell state c, and three gates. A gate is a set of sigmoid outputs between 0 and 1 that act like dimmer switches:

• the forget gate f decides how much of the old cell state to keep;
• the input gate i decides how much new information to write;
• the output gate o decides how much of the cell to show as the hidden state.

The updates are:

1. c_t = f * c_{t-1} + i * g, where g is a tanh "candidate" of new content;
2. h_t = o * tanh(c_t).

Think of a student's diary written in pencil. Each evening she can erase some old notes (forget gate), write new ones (input gate), and decide what to tell her friend about (output gate). The key detail is the first update: the cell state changes by adding, not by repeated multiplication through a squashing function. When the forget gate is near 1, the old memory flows forward almost unchanged, and so does the gradient. It is the same trick as ResNet's skip connection, applied through time.

Step 3: one LSTM step in NumPy.

import numpy as np

def sigmoid(z):
    return 1 / (1 + np.exp(-z))

rng = np.random.default_rng(2)
n_in, n_hidden = 3, 4
W = rng.normal(0, 0.5, (n_in + n_hidden, 4 * n_hidden))   # all four parts in one matrix
b = np.zeros(4 * n_hidden)
b[n_hidden:2 * n_hidden] = 1.0           # a common trick: start the forget gate open

x_t = np.array([0.5, -1.0, 0.3])
h_prev = np.zeros(n_hidden)
c_prev = np.array([0.8, -0.5, 0.1, 0.0])  # something remembered earlier

z = np.concatenate([x_t, h_prev]) @ W + b
i, f, g, o = np.split(z, 4)
i, f, o, g = sigmoid(i), sigmoid(f), sigmoid(o), np.tanh(g)
c = f * c_prev + i * g
h = o * np.tanh(c)

print("input gate: ", i.round(2))
print("forget gate:", f.round(2))
print("output gate:", o.round(2))
print("new cell:   ", c.round(3))
print("new hidden: ", h.round(3))
print("LSTM weights:", W.size + b.size, " vs plain RNN:", (n_in + n_hidden) * n_hidden + n_hidden)

input gate:  [0.48 0.57 0.47 0.47]
forget gate: [0.73 0.79 0.64 0.72]
output gate: [0.67 0.67 0.42 0.53]
new cell:    [ 0.85  -0.54  -0.269  0.245]
new hidden:  [ 0.461 -0.329 -0.109  0.127]
LSTM weights: 128  vs plain RNN: 32

The forget gate values sit well above 0.5 because we started its bias at 1, so most of the old cell state is kept. An LSTM has four sets of weights (input, forget, candidate, output), so it has about four times the parameters of a plain RNN of the same size.

GRU: a lighter cousin

The GRU (gated recurrent unit, Cho and colleagues, 2014) merges the cell and hidden state and uses two gates: an update gate that blends old and new memory, and a reset gate that decides how much of the past to use when proposing new content. It has three sets of weights instead of four, trains a little faster, and often performs about as well as an LSTM. Try both on your problem.

Recurrent layers in PyTorch

PyTorch provides nn.RNN, nn.LSTM and nn.GRU. With batch_first=True, inputs have shape (batch, time steps, features). Each returns the output at every step and the final state. The LSTM's final state is a pair (h_n, c_n).

import torch
from torch import nn

torch.manual_seed(0)
x = torch.randn(8, 20, 3)                    # 8 sequences, 20 steps, 3 features
for layer in [nn.RNN(3, 16, batch_first=True), nn.GRU(3, 16, batch_first=True),
              nn.LSTM(3, 16, batch_first=True)]:
    output, state = layer(x)
    h_n = state[0] if isinstance(state, tuple) else state
    params = sum(p.numel() for p in layer.parameters())
    print(f"{layer.__class__.__name__:<5} output {tuple(output.shape)}  final h {tuple(h_n.shape)}  parameters {params}")

RNN   output (8, 20, 16)  final h (1, 8, 16)  parameters 336
GRU   output (8, 20, 16)  final h (1, 8, 16)  parameters 1008
LSTM  output (8, 20, 16)  final h (1, 8, 16)  parameters 1344

The parameter counts show the 1 : 3 : 4 ratio of RNN, GRU and LSTM. (PyTorch keeps two bias vectors per gate group, which is why the numbers are slightly higher than our NumPy count.) The final hidden state has a first dimension of 1, for the number of layers. Setting num_layers=2 stacks two recurrent layers, and bidirectional=True adds a second pass that reads the sequence backwards, which helps when the whole sequence is available at once.

Step 4: train an LSTM to predict a wave

Let us train a small LSTM to look at 30 points of a noisy sine wave and predict the next point. Think of it as forecasting tomorrow's temperature in Lucknow from the last month. This is plain code with its real output.

import math
import torch
from torch import nn

torch.manual_seed(0)
t = torch.arange(0, 400, 0.1)
wave = torch.sin(t) + 0.1 * torch.randn_like(t)
window = 30
X = torch.stack([wave[i:i + window] for i in range(len(wave) - window)]).unsqueeze(-1)
y = wave[window:].unsqueeze(-1)
X_train, y_train, X_test, y_test = X[:3000], y[:3000], X[3000:], y[3000:]

class Forecaster(nn.Module):
    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(1, 32, batch_first=True)
        self.head = nn.Linear(32, 1)

    def forward(self, x):
        output, (h_n, c_n) = self.lstm(x)
        return self.head(output[:, -1])          # use the last time step

model = Forecaster()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(1, 6):
    perm = torch.randperm(len(X_train))
    for i in range(0, len(X_train), 64):
        idx = perm[i:i + 64]
        loss = nn.functional.mse_loss(model(X_train[idx]), y_train[idx])
        optimizer.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)   # guard against explosions
        optimizer.step()
    with torch.no_grad():
        test_mse = nn.functional.mse_loss(model(X_test), y_test).item()
    print(f"epoch {epoch}: test MSE {test_mse:.4f}")
print("noise variance (best possible MSE):", round(0.1 ** 2, 4))

epoch 1: test MSE 0.0156
epoch 2: test MSE 0.0163
epoch 3: test MSE 0.0134
epoch 4: test MSE 0.0122
epoch 5: test MSE 0.0126
noise variance (best possible MSE): 0.01

The test error quickly falls close to the noise level of 0.01, which is the best any model can do because the noise itself cannot be predicted. For real time-series work, also compare against simple baselines such as "tomorrow equals today", as described in Time Series Forecasting with ML.

Where RNNs stand today

For text, transformers have largely replaced RNNs, because they process all positions in parallel and handle long-range links better (see Attention & Transformers). RNNs, LSTMs and GRUs are still useful for streaming data that arrives one step at a time, small on-device models, and many time-series problems.

Common mistakes

• Forgetting batch_first=True. Without it, PyTorch expects (time steps, batch, features).
• Using the whole output when you meant the last step. For one prediction per sequence, use output[:, -1] or h_n[-1].
• No gradient clipping. Recurrent networks can have sudden exploding gradients; clipping is cheap insurance.
• Shuffling time steps inside a sequence. Shuffle whole sequences (windows) between batches, never the order within one.
• Testing on random windows from the training period. For forecasting, test on a later period, as we did.

Key points

• An RNN reads a sequence step by step, updating a hidden state with the same weights at every step.
• Backpropagation through time multiplies many slopes, so plain RNNs forget long-range information.
• LSTMs add a cell state and forget, input and output gates; additive updates keep memory and gradients flowing.
• GRUs use two gates and fewer parameters, and often work as well as LSTMs.
• In PyTorch use nn.LSTM(..., batch_first=True), take the last step for sequence-level predictions, and clip gradients.

from torch import nn

def predicted(groups, n_in, n_h):
    return groups * (n_h * (n_in + n_h) + 2 * n_h)

for layer, groups in [(nn.RNN(10, 64), 1), (nn.GRU(10, 64), 3), (nn.LSTM(10, 64), 4)]:
    actual = sum(p.numel() for p in layer.parameters())
    print(f"{layer.__class__.__name__:<5} predicted {predicted(groups, 10, 64):>6,}  actual {actual:>6,}")

RNN   predicted  4,864  actual  4,864
GRU   predicted 14,592  actual 14,592
LSTM  predicted 19,456  actual 19,456