Deep Sequence Modelling (RNN)

Lesson 4

LSTM & GRU

Gates that let recurrent networks remember longer than ten timesteps

In Lesson 3 we wrote the vanilla RNN recurrence:

\mathbf{h}_t = \tanh\!\left( \mathbf{W}_{xh}\mathbf{x}_t + \mathbf{W}_{hh}\mathbf{h}_{t-1} + \mathbf{b}_h \right)

The same matrix W_hh gets applied at every step. That single design choice is what gives the RNN its memory, and what destroys it. By the end of this lesson you will know exactly why a vanilla RNN forgets past ~10 timesteps, how the LSTM cell engineers a workaround, and how the GRU arrives at a simpler version of the same idea.

Why a vanilla RNN can't remember

To compute the gradient of the loss at step T with respect to the hidden state at an earlier step k, we apply the chain rule across every intervening step:

\frac{\partial \mathcal{L}_T}{\partial \mathbf{h}_k} \;=\; \frac{\partial \mathcal{L}_T}{\partial \mathbf{h}_T} \;\prod_{t=k+1}^{T} \frac{\partial \mathbf{h}_t}{\partial \mathbf{h}_{t-1}}

Each Jacobian inside the product expands to:

\frac{\partial \mathbf{h}_t}{\partial \mathbf{h}_{t-1}} \;=\; \mathbf{W}_{hh}^\top \cdot \mathrm{diag}\!\left(\tanh'(\mathbf{z}_t)\right)

Two terms multiply (T − k) times, and both are problems:

tanh′(z) ∈ (0, 1]

For any pre-activation away from zero, the derivative is < 1. Repeatedly multiplying values less than 1 drives the product toward zero exponentially.

W_hh raised to a power

If the largest singular value of W_hh is < 1, gradients vanish. If it's > 1, they explode. Hitting exactly 1 is unstable.

Empirically, this means a vanilla RNN can carry information across roughly 5–10 timesteps before the signal dies. That kills any task with longer-range dependencies, language modelling, music, long time series. We need an architecture whose recurrence Jacobian is not a repeated matrix multiplication.

The LSTM cell

The Long Short-Term Memory cell, Hochreiter & Schmidhuber, 1997, adds a second hidden quantity called the cell state c_t that flows through time on its own track. Three learnable gates control what gets written, kept, and read:

Gate	Equation	What it controls
Forget	f_t = σ(W_f[h_t−1, x_t] + b_f)	What of the prior cell state to erase
Input	i_t = σ(W_i[h_t−1, x_t] + b_i)	How much of the new candidate to write
Candidate	g_t = tanh(W_g[h_t−1, x_t] + b_g)	The new content being proposed
Output	o_t = σ(W_o[h_t−1, x_t] + b_o)	What of the cell state to expose as h_t

Here σ is the sigmoid (squashes to [0, 1], gate values), tanh squashes to [−1, 1] (signed content), and [h_t−1, x_t] denotes concatenation of the previous hidden state and the current input.

The state update combines the gates element-wise (denoted by ⊙):

\mathbf{c}_t = \mathbf{f}_t \odot \mathbf{c}_{t-1} \;+\; \mathbf{i}_t \odot \mathbf{g}_t

\mathbf{h}_t = \mathbf{o}_t \odot \tanh(\mathbf{c}_t)

Why this fixes vanishing gradients

Differentiate the cell-state update through time:

\frac{\partial \mathbf{c}_t}{\partial \mathbf{c}_{t-1}} \;=\; \mathbf{f}_t

The recurrence Jacobian is an elementwise product by the forget gate, not a repeated matrix multiplication. When f_t ≈ 1 the gradient flows across the cell-state highway nearly unchanged. Hochreiter called this the constant error carousel: a path where signal can travel for hundreds of timesteps without decay.

A practical consequence: initialising the forget-gate bias to 1 or 2 (Jozefowicz et al., 2015) biases the network toward remembering at the start of training. Most modern implementations do this by default.

Try it yourself. Three preset scenarios change only the forget-gate bias. A strong signal at t=1 followed by silence, watch whether the cell-state highway keeps the memory alive or wipes it:

Inside the LSTM cell · stage by stage

A strong signal arrives at t=1, then silence. Four operations make up one LSTM step: read inputs → compute four gates → update the cell-state highway → emit the hidden state. The forget gate alone decides whether the t=1 memory survives.

Stage:→→→

Timestep:

Stage 1 · Read inputs

The cell receives the current input x_t along with the previous hidden state h_{t−1} and cell state c_{t−1}.

x_1

h_0

c_0

x_1, h_0, c_0 → enter the cell

Cell-state magnitude across the sequence

Remember

0.57

t = 1

0.56

t = 2

0.56

t = 3

0.55

t = 4

t = 1 · stage 1/4Speed900ms

The GRU cell

The Gated Recurrent Unit, Cho et al., 2014, asks: do we really need a separate cell state, an output gate, and a separate input gate? Their answer is no. GRU collapses the LSTM into two gates and fuses the cell state back into the hidden state:

Gate	Equation	What it controls
Reset	r_t = σ(W_r[h_t−1, x_t] + b_r)	How much prior state contributes to the candidate
Update	z_t = σ(W_z[h_t−1, x_t] + b_z)	Interpolation weight between old h and new h

The candidate hidden state and the final state are:

\tilde{\mathbf{h}}_t = \tanh\!\left( \mathbf{W}\bigl[\mathbf{r}_t \odot \mathbf{h}_{t-1},\ \mathbf{x}_t\bigr] + \mathbf{b} \right)

\mathbf{h}_t = (1 - \mathbf{z}_t) \odot \mathbf{h}_{t-1} \;+\; \mathbf{z}_t \odot \tilde{\mathbf{h}}_t

The update gate z_t plays the same dual role that f_t and i_t played in the LSTM: when z_t ≈ 0, h_t ≈ h_t−1 (carry forward); when z_t ≈ 1, h_t ≈ h̃_t (overwrite). The same gradient-highway property holds, the recurrence Jacobian becomes an elementwise interpolation, not a repeated W^T product.

Parameter count: GRU vs LSTM

Per recurrent layer with hidden size H and input size I:

LSTM: 4 gate matrices → 4 · ((I + H) · H + H) parameters
GRU: 3 gate matrices → 3 · ((I + H) · H + H) parameters

GRU has ~25% fewer parameters and one fewer matmul per step. On smaller datasets that translates directly to faster convergence and less overfitting.

Empirical comparison

Chung et al. (2014) and the systematic search in Jozefowicz et al. (2015) ran LSTM vs GRU vs vanilla RNN across language modelling, speech, and music. Two robust findings:

LSTM and GRU both crush vanilla RNN

On any task with dependencies beyond ~10 steps, gating wins by a wide margin. This is no longer a research question.

LSTM vs GRU is task-dependent

No clean winner. LSTM tends to edge ahead on tasks demanding fine-grained memory control (music, long-context language modelling); GRU often matches or beats it on smaller datasets and shorter sequences with fewer parameters.

When to use which

GRU

Default for prototyping, small-data regimes, and time-series. Fewer params, fewer hyperparameters, often comparable accuracy.

LSTM

Pick when you have data + compute and the task needs long-range memory: speech, music, long-context language. The extra gate and cell state earn their keep.

Transformer

Pick when you can afford it. Attention is the parallelisable, longer-context replacement for both, covered in the AIAYN and GPT-2 courses.

The arc of the field: vanilla RNNs (1986) → LSTM (1997) → GRU (2014) → Transformer (2017). Each step traded away one constraint of the previous. LSTMs are no longer SOTA in 2026, but they are still everywhere in production for sequence problems where attention is overkill: streaming inference with bounded state, on-device models, modest-data time series, and as the recurrent component inside hybrid models.

A reference PyTorch snippet

Both cells are one line in PyTorch, but knowing what's underneath is what differentiates a candidate who has used nn.LSTM from one who can debug it.

Python

import torch
import torch.nn as nn
 
# LSTM: 4 gates per step, returns (output, (h_n, c_n))
lstm = nn.LSTM(input_size=64, hidden_size=128, num_layers=2, batch_first=True)
x = torch.randn(8, 50, 64)        # (batch, time, features)
out, (h_n, c_n) = lstm(x)
# out:  (8, 50, 128) , hidden state at every timestep
# h_n:  (2, 8, 128)  , final hidden state per layer
# c_n:  (2, 8, 128)  , final cell state per layer
 
# GRU: 3 gates per step, returns (output, h_n), no cell state
gru = nn.GRU(input_size=64, hidden_size=128, num_layers=2, batch_first=True)
out, h_n = gru(x)
# out:  (8, 50, 128)
# h_n:  (2, 8, 128)

A few things worth knowing:

batch_first=True swaps the default (time, batch, features) to (batch, time, features). The default is a footgun.
For variable-length sequences inside a batch, wrap inputs in nn.utils.rnn.pack_padded_sequence so padded steps don't pollute the hidden state.
Both modules accept an initial state. Defaulting to zeros is fine; for stateful inference (streaming) you carry (h_n, c_n) from the previous chunk.

What you can now do

Diagnose when a sequence model is failing because of vanishing gradients vs an architectural mismatch.
Read the LSTM and GRU equations and explain each gate's role.
Pick between LSTM and GRU based on dataset size, sequence length, and compute budget.
Recognise the constant error carousel, the structural reason gated cells learn longer dependencies than vanilla RNNs.

Check your understanding

1 / 4

Mathematically, why does a vanilla RNN's gradient vanish over long sequences?