DL S033 · LSTM and GRU — Gating the Gradient Highway

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 33 / 80 · Module M06 · Enriched pass 2026-07-27

The story

0 · The paper that took ten years to be believed

Vienna, 1997. Sepp Hochreiter and Jürgen Schmidhuber publish Long Short-Term Memory in Neural Computation, volume 9, issue 8. The paper is 32 dense pages. It contains a strange-looking cell diagram with something called a "constant error carousel", three multiplicative gates, and an argument (in section 3.2) that gradient flow through this cell is uniform across time — no vanishing, no exploding.

Almost nobody believed them. The paper accumulated a slow trickle of citations through the late 1990s and early 2000s. Neural nets themselves were in a research winter; the ImageNet moment was still fifteen years away. Even people who were doing sequence modelling preferred simpler RNNs, or Hidden Markov Models, or n-grams.

Then two things happened at once. In 2013 Alex Graves used stacked LSTMs to smash the state of the art on handwriting recognition and TIMIT speech, then wrote a landmark technical report generating handwriting with an LSTM one stroke at a time — the outputs went viral. In 2014 Google Translate switched to an encoder-decoder LSTM (Sutskever, Vinyals, Le Sequence to Sequence Learning) and delivered a jump in translation quality that stunned the industry. By 2015, LSTMs were the default sequence model in every deep-learning textbook. Hochreiter's 1997 paper — snubbed for over a decade — had, at long last, taken over.

The trick is beautiful and, once you see it, obvious. The plain RNN destroys gradient by multiplying the hidden state by $W_{hh}$ at every step. The LSTM has a separate cell state that is added to rather than multiplied by the previous value. Addition preserves gradient. Multiplication (usually) does not. Everything else — the three gates, the tanh activations, the peephole connections — is engineering around that one idea.

Today you build the LSTM. You feel why the cell state is a gradient highway. You meet its 2014 cousin the GRU (simpler, similar performance, one fewer gate). And you'll see the same architectural motif — additive skip path guarded by a sigmoid gate — repeated verbatim in Highway Networks (2015), ResNets (2015), and every transformer residual connection ever written. Once you spot the pattern, you spot it everywhere.

1 · What you'll build today

One sentence: You will implement an LSTM cell in ~30 lines of PyTorch, compare its gradient flow to the plain RNN from S032, and internalise why "additive cell state + sigmoid gates" is the architectural motif that makes deep sequence models trainable.

By the end you will:

• Draw the LSTM cell diagram from memory (input, forget, output, candidate — four learned linear maps).
• Explain, with equations and pictures, why $\partial c_t / \partial c_{t-1}$ is (approximately) the identity matrix scaled by the forget gate, and why that makes gradient flow across hundreds of timesteps.
• Implement LSTMCell and GRUCell from scratch in PyTorch.
• Trace the intellectual lineage: LSTM (1997) → forget-gate improvement (Gers et al. 2000) → Highway Networks (2015) → ResNets (2015) → transformer residual (2017) → modern xLSTM revival (2024).
• Choose sensibly between LSTM, GRU, and transformer for a given task.

2 · The core idea — a gradient highway

The plain RNN update was:

Gradient flowing backward through this update gets multiplied by $\operatorname{diag}(1 - \tanh^2) \cdot W_{hh}$ at every step. That product goes to zero (or infinity) fast.

The LSTM's central trick: introduce a second hidden vector called the cell state $c_t$, and update it via addition rather than multiplication.

where $\odot$ is elementwise product, $f_t$ (forget gate) and $i_t$ (input gate) are sigmoid vectors in $(0, 1)^H$, and $\tilde{c}_t$ is a candidate update from the current input. The Jacobian of this update:

That is a diagonal matrix, with entries in $(0, 1)$. If the forget gate is close to 1 (the model wants to remember), the Jacobian is approximately the identity matrix. Multiply the identity by itself a hundred times and you get the identity — no vanishing, no exploding. Gradient flows across time as if along a highway.

2.1 Why this simple observation was a breakthrough

Hochreiter's 1991 thesis had already identified the disease. Several groups in the mid-1990s tried architectural workarounds — Bengio et al. 1994, Lin et al. 1995 — that limited how much gradient could vanish, but none preserved it fully. The LSTM's insight was structural, not incremental: do not try to prevent multiplication from destroying gradient; add an additive path where gradient does not go through multiplication at all. That reframing — "route around the problem instead of shrinking it" — is a recurring pattern in deep learning. It later reappeared as the ResNet trick (2015): don't try to make deep networks trainable by clever initialisation; add a skip connection so gradient bypasses the depth entirely.

3 · The four gates — with pictures then equations

An LSTM cell computes four things at each timestep, all from the same inputs (x_t, h_{t-1}):

1. Forget gate $f_t = \sigma(W_f [x_t, h_{t-1}] + b_f) \in (0,1)^H$ — "what fraction of each cell-state dimension should we erase?"
2. Input gate $i_t = \sigma(W_i [x_t, h_{t-1}] + b_i) \in (0,1)^H$ — "what fraction of each candidate dimension should we admit?"
3. Candidate cell state $\tilde{c}_t = \tanh(W_c [x_t, h_{t-1}] + b_c) \in (-1,1)^H$ — "if we did write to the cell, what would we write?"
4. Output gate $o_t = \sigma(W_o [x_t, h_{t-1}] + b_o) \in (0,1)^H$ — "what fraction of each cell-state dimension should be exposed as the hidden state right now?"

Then the state updates:

The gate values are computed from both the current input and the previous hidden state, so the model can decide "given what I just saw and what I already know, forget dimensions 3–7 of my memory and write new stuff to dimensions 12–20".

3.1 A worked example on 2-dimensional state

Set $H = 2$ so we can inspect every number. Suppose at timestep $t$:

• $c_{t-1} = [0.8, -0.3]$ (some accumulated memory)
• $f_t = [0.95, 0.02]$ (keep dim 0 mostly, aggressively forget dim 1)
• $i_t = [0.10, 0.80]$ (write a little to dim 0, a lot to dim 1)
• $\tilde{c}_t = [0.5, 0.9]$

Then:

Notice how dim 1 of $c$ has been rewritten (-0.3 → 0.71) while dim 0 has been slightly updated (0.8 → 0.81). The gates gave the model fine-grained control over which memory slots to preserve vs overwrite. Try this on paper for one or two more timesteps until it clicks.

3.2 Peephole connections (a footnote)

Gers, Schraudolph, Schmidhuber (2002) proposed peephole connections — let the gates see $c_{t-1}$ directly, not just via $h_{t-1}$:

Marginal improvement, mostly on very precise timing tasks. Rarely used in modern PyTorch code. Mention it for completeness; skip in your own implementations.

4 · Implement LSTMCell from scratch

The four-gate LSTM in PyTorch, written out explicitly:

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class MyLSTMCell(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        # One big linear map from [x, h] to [f, i, ĉ, o] concatenated.
        # This is exactly what PyTorch's built-in LSTMCell does — bundling the four
        # gate projections into a single (4H) matmul is a huge GPU efficiency win.
        self.linear = nn.Linear(input_dim + hidden_dim, 4 * hidden_dim)
        self.H = hidden_dim
        # Initialise forget-gate bias to 1 — Jozefowicz et al. 2015 trick.
        with torch.no_grad():
            self.linear.bias[hidden_dim:2*hidden_dim].fill_(1.0)
 
    def forward(self, x, state):
        h_prev, c_prev = state                          # (B, H) each
        cat = torch.cat([x, h_prev], dim=-1)            # (B, D+H)
        gates = self.linear(cat)                        # (B, 4H)
        f, i, g, o = gates.chunk(4, dim=-1)             # each (B, H)
 
        f = torch.sigmoid(f)                            # forget
        i = torch.sigmoid(i)                            # input
        g = torch.tanh(g)                               # candidate (ĉ)
        o = torch.sigmoid(o)                            # output
 
        c = f * c_prev + i * g                          # additive cell update
        h = o * torch.tanh(c)
        return h, (h, c)

Three implementation points that matter:

1. Bundled 4H projection. Instead of four separate nn.Linear(D+H, H) layers, we do one nn.Linear(D+H, 4H) and chunk the output. This is 4× fewer matmul launches on GPU. PyTorch's built-in nn.LSTMCell does the same. Roughly 2× faster in practice.
2. Forget-gate bias init to 1. Jozefowicz, Zaremba, Sutskever 2015, An Empirical Exploration of RNN Architectures showed that initialising the forget-gate bias to 1 (so the initial forget gate is $\sigma(1) \approx 0.73$, closer to "keep") speeds up training substantially and often improves final performance. This is now standard.
3. torch.sigmoid and torch.tanh are numerically safer than F.sigmoid / F.tanh (deprecated aliases). Use the tensor-methods x.sigmoid() / x.tanh() if you prefer.

4.1 Unrolling over a sequence

def lstm_forward(cell, xs, h0=None, c0=None):
    B, T, D = xs.shape
    H = cell.H
    if h0 is None: h0 = xs.new_zeros(B, H)
    if c0 is None: c0 = xs.new_zeros(B, H)
    state = (h0, c0)
    outputs = []
    for t in range(T):
        h, state = cell(xs[:, t], state)
        outputs.append(h)
    return torch.stack(outputs, dim=1), state

Now swap this into the character-RNN training loop from S032 and retrain. You'll see the loss drop from ~2.0 (plain RNN final) to ~1.5 in the same number of steps — a large improvement for a small code change. Generations become noticeably more coherent, with fewer of the "Wille the his stall" gibberish patterns.

5 · Why the gradient really flows — a detailed derivation

Let's verify Hochreiter's claim that gradient through the LSTM cell state does not vanish. Compute $\partial c_t / \partial c_{t-1}$ from the update:

Take the total derivative w.r.t. $c_{t-1}$:

Without peephole connections, $f_t, i_t, \tilde{c}_t$ depend on $c_{t-1}$ only through $h_{t-1}$, and $h_{t-1}$ depends on $c_{t-1}$ through $h_{t-1} = o_{t-1} \odot \tanh(c_{t-1})$. So the indirect terms exist but are heavily attenuated by the tanh (which saturates) and by the previous o gate.

The direct path — the highway — is $\operatorname{diag}(f_t)$. When the forget gate is close to 1 (the model is keeping memory), this is essentially the identity. Multiply the identity by itself $T$ times, you get the identity. Signal survives.

By contrast, in a plain RNN the only path is:

There is no highway. Every gradient signal must traverse this multiplicative Jacobian, and its spectral norm is almost never 1.

5.1 Measure it yourself

Take the gradient-norm experiment from S032 §5.1 and swap nn.RNN for nn.LSTM. Same T=100, same hidden dim, same seed. You'll see the gradient norm at $t=0$ decay from ~1e-11 (plain RNN) to something like ~1e-2 (LSTM) — nine orders of magnitude improvement, entirely from the architectural change. This is the empirical validation of Section 2's theory.

6 · The GRU — same idea, one fewer gate

Cho, Merriënboer, Bahdanau, Bengio (2014), Learning Phrase Representations using RNN Encoder-Decoder, introduced the Gated Recurrent Unit as a leaner LSTM. Two gates instead of four. No separate cell state.

The update rule is a convex combination of the previous hidden state and a candidate — no separate cell state, no output gate. The reset gate controls how much of $h_{t-1}$ is exposed to the candidate computation; the update gate controls how much of $h_{t-1}$ is preserved in the new hidden state.

Notice the Jacobian $\partial h_t / \partial h_{t-1}$ has a direct-path term $\operatorname{diag}(1 - z_t)$. Same trick as the LSTM's cell state, just merged into the hidden state itself.

6.1 GRU in code

class MyGRUCell(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.linear_rz = nn.Linear(input_dim + hidden_dim, 2 * hidden_dim)
        self.linear_h  = nn.Linear(input_dim + hidden_dim, hidden_dim)
        self.H = hidden_dim
 
    def forward(self, x, h_prev):
        cat = torch.cat([x, h_prev], dim=-1)
        r, z = self.linear_rz(cat).chunk(2, dim=-1)
        r, z = torch.sigmoid(r), torch.sigmoid(z)
        cat2 = torch.cat([x, r * h_prev], dim=-1)
        h_tilde = torch.tanh(self.linear_h(cat2))
        h = (1 - z) * h_prev + z * h_tilde
        return h

Parameter count. For hidden dim $H$ and input dim $D$:

• LSTM: $4H(D + H + 1)$ parameters per cell.
• GRU: $3H(D + H + 1)$ parameters per cell.

Roughly 25% fewer parameters, so faster to train and inference for the same $H$.

6.2 LSTM vs GRU — the empirical answer

Chung et al. 2014 benchmarked them on polyphonic music, speech signal modelling, and machine translation. Result: neither consistently wins. GRUs slightly ahead on smaller datasets and smaller models; LSTMs slightly ahead on larger datasets when you have more training compute. On language modelling in particular the gap is essentially noise.

Practical guidance for 2025 if you're actually using recurrent layers:

• Default to GRU if you don't have strong reasons to choose. Fewer params, easier to tune, essentially same quality.
• Use LSTM if you're building on top of pretrained models (many older speech and OCR systems ship LSTM weights).
• Prefer transformers or SSMs for anything language-model-shaped. GRUs and LSTMs are still useful for streaming inference and edge deployment.

7 · The pattern — additive skip + sigmoid gate — is everywhere

Once you see "additive update + sigmoid gate", you spot it in every modern architecture.

7.1 Highway Networks (Srivastava, Greff, Schmidhuber 2015)

For a feedforward deep network, Highway Networks proposed:

where $T$ is a sigmoid "transform gate". Compare to the GRU update rule — it is literally the same equation, applied across depth instead of across time. Highway Networks allowed training of 100+ layer feedforward networks in mid-2015, months before ResNets.

7.2 ResNets (He, Zhang, Ren, Sun 2015, ImageNet)

The most cited computer vision paper of the decade did something simpler:

No gate. Just addition. ResNet-152 hit human-level accuracy on ImageNet and won ILSVRC 2015. This is the LSTM's cell-state highway, stripped down to its essence. Every deep network you train today has residual connections, and every one of them exists because Hochreiter figured out in 1997 that addition preserves gradient and multiplication destroys it.

7.3 Transformer residual connections (Vaswani et al. 2017)

Every transformer block is:

Two additive skip connections per block, one around attention and one around the MLP. Without them, transformers of depth > 6 or so fail to train. With them, we scale to 100+ layers routinely. Same insight, third generation.

8 · When to use LSTM/GRU in 2025

Use them when:

• Streaming or low-latency inference matters (real-time speech, keystroke prediction, sensor fusion). One recurrent step is $O(H^2)$ compute and $O(H)$ memory — constant regardless of history length. Transformers require caching all prior tokens.
• Sequence lengths are enormous (100k+) and attention is unaffordable. Though SSMs (Mamba-2, RWKV-7) are now usually a better choice here.
• You're deploying to edge devices (phones, watches, MCUs) where the model must fit in tens of KB. A 1-layer GRU with $H = 64$ has ~25k parameters. A 1-layer transformer of comparable expressiveness is roughly 10× larger.
• You need strong inductive bias for temporal ordering. Vanilla transformers rely on positional encodings; RNNs have order baked into the architecture.

Avoid them when:

• You need to train on lots of parallel hardware. RNNs are sequential-by-construction and leave 99% of GPU/TPU cores idle during training. Transformers and structured SSMs parallelise across sequence length.
• You need long-range dense information access (few-shot in-context learning, long-form reasoning). Attention lets any position directly reference any other; the recurrent bottleneck of a fixed-size $h$ vector cannot compete.
• Your dataset is large (>10B tokens) and quality is paramount. Attention scales better with data at this regime.

9 · War stories

10 · Try it yourself

11 · Recap in five sentences

The LSTM's cell state — updated as $c_t = f_t \odot c_{t-1} + i_t \odot \tilde c_t$ — creates an additive gradient highway whose Jacobian is $\operatorname{diag}(f_t)$, close to identity whenever the forget gate is near 1, which is how it survives hundreds of timesteps without vanishing. The four gates (forget, input, candidate, output) give the model fine-grained control over what to erase, what to write, and what to expose. GRUs merge cell state and hidden state and use two gates instead of four; empirically the two architectures perform comparably. The "additive skip + sigmoid gate" pattern reappears verbatim in Highway Networks, ResNets, and transformer residual connections — Hochreiter's 1997 insight is quietly present in every modern deep network. In 2024 xLSTM brought Hochreiter himself back with an evolved recurrent architecture competitive with billion-scale transformers, closing a thirty-three-year arc that started with his diploma thesis.

🧠 Retention scaffold

One-line summary (write it in your own words): ________________________________________________

Spaced review: re-read §2 (the highway idea) and §5 (Jacobian derivation) in 24 hours. Revisit the full session on day 7.

Next session (S034): we stack two LSTMs into an encoder–decoder (seq2seq) and translate short English sentences to French, reproducing the Sutskever 2014 result that put LSTMs on the map. You'll also feel the bottleneck (a single fixed-size context vector) that motivates attention in S035.

Sticky note (keep on your desk): Addition preserves gradient, multiplication destroys it. Every skip connection you've ever seen — LSTM cell state, ResNet, transformer residual — is one line of the same solution to Hochreiter's 1991 problem.

Further reading

• Original. Hochreiter & Schmidhuber 1997, Long Short-Term Memory.
• Forget gate. Gers, Schmidhuber, Cummins 2000, Learning to Forget: Continual Prediction with LSTM.
• Empirical comparison. Chung et al. 2014, arXiv:1412.3555; Jozefowicz, Zaremba, Sutskever 2015, An Empirical Exploration of RNN Architectures.
• The definitive pedagogy. Olah 2015, Understanding LSTM Networks.
• Skip-connection family. Srivastava et al. 2015 Highway Networks; He et al. 2015 ResNets; Vaswani et al. 2017 Transformers.
• 2024 revival. Beck et al. 2024, xLSTM; Peng et al. 2023, RWKV; Gu & Dao 2023, Mamba.

Previous: ← DL S032 · Next: DL S034 →