11 LSTM & GRU

TL;DR

The LSTM (1997) fixes vanishing gradients by routing a cell state \(c_t\) through the network on a protected highway — gating mechanisms decide what to forget, what to write, and what to read, keeping the gradient path near-constant. The GRU (2014) achieves comparable results with two gates instead of three. Both were the dominant sequence model for nearly two decades, powering speech recognition, translation, and language modeling until the Transformer.

Depends on: Vanilla RNN & BPTT

11.1 Why this matters

The vanilla RNN’s Jacobian product \(\prod_i \mathbf{W}_h^\top \operatorname{diag}(1 - h_i^2)\) shrinks exponentially with sequence length. Hochreiter (Hochreiter, 1991) diagnosed this in 1991; Hochreiter & Schmidhuber (Hochreiter & Schmidhuber, 1997) solved it in 1997 with the LSTM. The key insight: replace the single scalar hidden state with a cell state \(c_t\) that travels through time with only additive updates — no repeated multiplication by the same weight matrix.

11.2 The mechanism

11.2.1 LSTM

Let \([h_{t-1}; x_t]\) denote the concatenation of the previous hidden state and current input. The LSTM computes four quantities at each step:

\[ \begin{aligned} f_t &= \sigma\!\bigl(\mathbf{W}_f [h_{t-1}; x_t] + \mathbf{b}_f\bigr) & &\text{(forget gate — how much of } c_{t-1} \text{ to keep)}\\ i_t &= \sigma\!\bigl(\mathbf{W}_i [h_{t-1}; x_t] + \mathbf{b}_i\bigr) & &\text{(input gate — how much new info to write)}\\ \tilde{c}_t &= \tanh\!\bigl(\mathbf{W}_c [h_{t-1}; x_t] + \mathbf{b}_c\bigr) & &\text{(candidate cell update)}\\ o_t &= \sigma\!\bigl(\mathbf{W}_o [h_{t-1}; x_t] + \mathbf{b}_o\bigr) & &\text{(output gate — how much cell to expose)} \end{aligned} \]

The cell state and hidden state update:

\[ c_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t, \qquad h_t = o_t \odot \tanh(c_t). \]

The critical term is \(c_t = f_t \odot c_{t-1} + \ldots\): the cell state accumulates via addition, not matrix multiplication. Gradients flow backward through \(c_t\) as \(\partial c_t / \partial c_{t-1} = f_t\) — a scalar gate, not a full Jacobian product. As long as \(f_t \approx 1\), the gradient is preserved across many steps.

11.2.2 GRU: simplified gating

Cho et al. (Cho et al., 2014) proposed the GRU in 2014 as a lighter alternative. Two gates replace three; the separate cell state is eliminated:

\[ \begin{aligned} z_t &= \sigma\!\bigl(\mathbf{W}_z [h_{t-1}; x_t]\bigr) & &\text{(update gate — interpolation weight)}\\ r_t &= \sigma\!\bigl(\mathbf{W}_r [h_{t-1}; x_t]\bigr) & &\text{(reset gate — how much past to expose)}\\ \tilde{h}_t &= \tanh\!\bigl(\mathbf{W}[r_t \odot h_{t-1}; x_t]\bigr) & &\text{(candidate hidden state)}\\ h_t &= (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t. & & \end{aligned} \]

The update gate \(z_t\) linearly interpolates between keeping the old hidden state and replacing it with the candidate — a clean analogue of the LSTM forget + input gates combined.

LSTM vs GRU: which to use?

Empirically, the two perform comparably on most tasks. GRU trains faster (fewer parameters), LSTM gives finer-grained control over memory. The Transformer later made the choice moot for most applications — but both remain relevant in low-latency, streaming, and on-device settings where sequence-parallel architectures are too large.

11.2.3 Minimal sketch

The sketch implements both cells in NumPy and runs them on the same input to compare hidden states.

import numpy as np

rng = np.random.default_rng(0)
d_x, d_h = 4, 8

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

# LSTM weights (concat [h; x] → 4 gates)
W_lstm = rng.normal(scale=0.1, size=(4 * d_h, d_h + d_x))
b_lstm = np.zeros(4 * d_h)

def lstm_step(h, c, x):
    gates = W_lstm @ np.concatenate([h, x]) + b_lstm
    f, i, g, o = np.split(gates, 4)
    c = sigmoid(f) * c + sigmoid(i) * np.tanh(g)
    h = sigmoid(o) * np.tanh(c)
    return h, c

# GRU weights
W_z = rng.normal(scale=0.1, size=(d_h, d_h + d_x))
W_r = rng.normal(scale=0.1, size=(d_h, d_h + d_x))
W_h = rng.normal(scale=0.1, size=(d_h, d_h + d_x))

def gru_step(h, x):
    hx = np.concatenate([h, x])
    z  = sigmoid(W_z @ hx)
    r  = sigmoid(W_r @ hx)
    h_ = np.tanh(W_h @ np.concatenate([r * h, x]))
    return (1 - z) * h + z * h_

x = rng.normal(size=(5, d_x))
h_l, c_l = np.zeros(d_h), np.zeros(d_h)
h_g = np.zeros(d_h)

for t in range(5):
    h_l, c_l = lstm_step(h_l, c_l, x[t])
    h_g      = gru_step(h_g, x[t])

print("LSTM h_T norm:", np.linalg.norm(h_l).round(4))
print("GRU  h_T norm:", np.linalg.norm(h_g).round(4))

LSTM h_T norm: 0.0742
GRU  h_T norm: 0.1164

BiLSTM: reading in both directions

A bidirectional LSTM (BiLSTM) runs two LSTM cells over the sequence — one forward, one backward — and concatenates their hidden states: \(h_t^{\text{bi}} = [\overrightarrow{h}_t; \overleftarrow{h}_t]\). This gives each position context from both past and future, which is why BiLSTMs underpin BERT’s encoder architecture and most pre-Transformer NER and parsing systems.

11.3 Application & impact

Concept here	What it became	Where
Cell state highway	Residual stream	Transformer (residual connections, 2017)
Forget gate	Decay / forgetting mechanism	S4, Mamba (2022–2023)
Gating (sigmoid × tanh)	SwiGLU, GLU activations	LLaMA, Mistral
LSTM encoder	BERT encoder backbone	Pre-Transformer NLP (2014–2018)
GRU	Lighter recurrent baseline	Still used in low-latency streaming

Speech recognition (Google, 2015): replaced HMM-GMM pipelines with LSTM stacks.
Neural MT (2014–2016): Seq2Seq with LSTMs was the state of the art before attention.
Language models (ELMo, 2018): BiLSTM over large text corpora produced the first strong contextualised word representations before BERT.

Key takeaway

The LSTM’s cell state turns the Jacobian product \(\prod \mathbf{W}_h^\top\) into a scalar product \(\prod f_t\) — bounded between 0 and 1, but controllable. The gradient highway is the architectural solution to vanishing gradients; gating is its implementation. Everything downstream — attention, residual connections, the Transformer — is in some sense a variation on this idea.

→ Next: Seq2Seq RNN Encoder-Decoder