LSTM & GRU:
Teaching Networks to Remember

The vanishing gradient is not a numerical accident.
It is a theorem: multiply any matrix with spectral radius < 1 by itself
enough times and the result approaches zero.

The LSTM's answer was architectural, not numerical:
stop multiplying. Give memory a path that adds instead.

01 — A Precise Statement of the Problem

Before building the solution, let's state the problem with full precision. From Day 20, the gradient of the loss at step T with respect to the hidden state at step 1 is:

The root cause is the multiplicative recurrence: to update the hidden state, we always multiply by W_h. There is no additive path — every gradient step must pass through this multiplication. For the RNN, the only connection between h_1 and h_T is through this chain of multiplications.

The architectural insight behind both LSTM and GRU is the same:

02 — LSTM: Architecture from First Principles

The Long Short-Term Memory was introduced by Hochreiter & Schmidhuber in 1997 — one of the most cited papers in deep learning history. It introduces two running states instead of one:

Three gates control the flow of information. Each gate produces a vector of values between 0 and 1 (via sigmoid), acting as a soft switch: 0 = block completely, 1 = pass through completely, 0.7 = pass through 70%.

The Four Equations

The Gradient Highway: Mathematical Proof

Let's verify exactly why the cell state fixes the vanishing gradient. The gradient of C_T with respect to C_1:

The forget gate is learned, not fixed. When the network encounters information it wants to remember for a long time — a subject at the start of a sentence that needs to match a pronoun 40 words later — it learns to set f_t ≈ 1 for those dimensions of C across those steps. The gradient path stays open exactly where it needs to be.

03 — LSTM Cell Diagram

04 — Worked Forward Pass: "I love NLP"

Let's trace one LSTM step manually. hidden_size = 1 (scalar for clarity). All weights = 0.5, all biases = 0. Previous state: h_0 = 0, C_0 = 0. Input: x_1 = 1 ("I").

── Step 1: x_1="I"=1, h_0=0, C_0=0 ────────────────────────────
z = [h_0 ; x_1] = [0, 1]    (concat)

f_1 = σ(0.5×0 + 0.5×1) = σ(0.5) = 0.622    ← forget gate
i_1 = σ(0.5×0 + 0.5×1) = σ(0.5) = 0.622    ← input gate
g_1 = tanh(0.5×0 + 0.5×1) = tanh(0.5) = 0.462  ← candidate
o_1 = σ(0.5×0 + 0.5×1) = σ(0.5) = 0.622    ← output gate

C_1 = f_1 × C_0 + i_1 × g_1 = 0.622×0 + 0.622×0.462 = 0.287
h_1 = o_1 × tanh(C_1) = 0.622 × tanh(0.287) = 0.622 × 0.280 = 0.174

── Step 2: x_2="love"=1, h_1=0.174, C_1=0.287 ─────────────
z = [0.174, 1]

f_2 = σ(0.5×0.174 + 0.5×1) = σ(0.587) = 0.643
i_2 = σ(0.5×0.174 + 0.5×1) = σ(0.587) = 0.643
g_2 = tanh(0.587) = 0.528
o_2 = σ(0.587) = 0.643

C_2 = 0.643×0.287 + 0.643×0.528 = 0.184 + 0.340 = 0.524
     ↑ C_1 still alive in C_2 — memory preserved
h_2 = 0.643 × tanh(0.524) = 0.643 × 0.481 = 0.309

05 — What Each Gate Learns in Practice

After training on real data, the four gates develop interpretable specialisations. Here's what mechanistic interpretability research and practical experience have found:

06 — LSTM Variants Worth Knowing

Peephole Connections (Gers & Schmidhuber, 2000)

In the standard LSTM, gates only see h_{t-1} and x_t. Peephole connections let the gates also see the cell state C_{t-1} directly:

Peepholes let gates "know" what's already stored in memory before deciding whether to forget or write. Useful for tasks requiring precise timing — counting exact beats in music generation, or matching opening/closing brackets in code. In practice, the improvement is small and peepholes are rarely used today.

Coupled Forget and Input Gates

A simpler variant that ties f_t and i_t together: i_t = 1 − f_t. If you forget 60% of the old state, you write 40% of the new state. This halves the number of gate parameters and works well on many tasks. The GRU takes this coupling idea further.

Bidirectional LSTM (BiLSTM)

Two LSTMs run in opposite directions over the same sequence — one left-to-right (forward), one right-to-left (backward). Their hidden states are concatenated at each timestep: h_t = [h_t_fwd ; h_t_bwd]. Every token now has access to both past and future context. BiLSTMs were state-of-the-art for NER, POS tagging, and coreference resolution from ~2016–2018, before being replaced by BERT.

Stacked (Deep) LSTM

Multiple LSTM layers stacked vertically — the h_t sequence from layer k becomes the x_t input sequence for layer k+1. GPT-2's predecessor used a 2-layer LSTM for language modelling. Google's Neural Machine Translation (GNMT) used an 8-layer encoder-decoder LSTM in 2016. Depth provides hierarchical representation: lower layers capture local syntax, higher layers capture longer-range semantics.

07 — GRU: The Streamlined Alternative

The Gated Recurrent Unit (Cho et al., 2014) was introduced in the same paper as the seq2seq architecture. Cho's observation: the LSTM has three gates and two states (h and C). This is more complexity than necessary. Can you get the same gradient flow benefits with fewer parameters?

The GRU collapses the LSTM's four equations into two gates and one state:

GRU Cell Diagram

How GRU Achieves the Gradient Highway

The GRU's gradient path is through the update gate equation. When z_t ≈ 0:

GRU vs LSTM: The Key Differences

08 — Gradient Flow: RNN vs LSTM vs GRU

The gradient bars tell the story clearly. The RNN gradient at step T−50 is near zero — weights at early timesteps receive no meaningful update. LSTM and GRU both maintain meaningful gradient signal even at step T−50, because their gradient paths bypass the multiplicative W_h chain.

09 — Code: LSTM and GRU from Scratch

Complete implementations of both architectures — every gate written explicitly, a stacked bidirectional variant, and a training comparison against vanilla RNN on a long-range dependency task.

import torch
import torch.nn as nn
import math

# ─────────────────────────────────────────────────────────────────────────
# 1. LSTM Cell — all four equations explicit, every gate visible
# ─────────────────────────────────────────────────────────────────────────
class LSTMCellManual(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size
        # One combined weight matrix for all 4 gates — efficient but readable when split
        # W: (4*hidden, input+hidden)  —  [forget | input | candidate | output]
        self.W = nn.Linear(input_size + hidden_size, 4 * hidden_size)
        self._init_weights()

    def _init_weights(self):
        # Orthogonal init for the recurrent part, uniform for input part
        nn.init.orthogonal_(self.W.weight[:, :self.hidden_size])
        nn.init.xavier_uniform_(self.W.weight[:, self.hidden_size:])
        # Forget gate bias to 1.0 — encourages remembering at initialisation
        nn.init.ones_(self.W.bias[:self.hidden_size])

    def forward(self, x_t, h_prev, c_prev):
        """
        x_t:    (batch, input_size)
        h_prev: (batch, hidden_size)
        c_prev: (batch, hidden_size)
        returns: h_t, c_t — each (batch, hidden_size)
        """
        combined = torch.cat([h_prev, x_t], dim=1)   # (batch, hidden+input)
        gates    = self.W(combined)                 # (batch, 4*hidden)

        # Split into 4 gates along dim=1
        f, i, g, o = gates.chunk(4, dim=1)

        f = torch.sigmoid(f)       # forget gate  → (0, 1)
        i = torch.sigmoid(i)       # input gate   → (0, 1)
        g = torch.tanh(g)          # candidate    → (-1, 1)
        o = torch.sigmoid(o)       # output gate  → (0, 1)

        c_t = f * c_prev + i * g   # cell update — the gradient highway (addition!)
        h_t = o * torch.tanh(c_t)  # hidden state — filtered cell
        return h_t, c_t


# ─────────────────────────────────────────────────────────────────────────
# 2. GRU Cell — two gates, one state
# ─────────────────────────────────────────────────────────────────────────
class GRUCellManual(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size
        # Reset + update gates share one projection from [h, x]
        self.W_rz = nn.Linear(input_size + hidden_size, 2 * hidden_size)
        # Candidate uses reset-gated h_{t-1}, so separate projection
        self.W_n_h = nn.Linear(hidden_size, hidden_size, bias=False)
        self.W_n_x = nn.Linear(input_size, hidden_size)

    def forward(self, x_t, h_prev):
        """
        x_t:    (batch, input_size)
        h_prev: (batch, hidden_size)
        returns: h_t (batch, hidden_size)
        """
        combined = torch.cat([h_prev, x_t], dim=1)
        rz       = torch.sigmoid(self.W_rz(combined))   # (batch, 2*hidden)
        r, z     = rz.chunk(2, dim=1)                    # reset, update gates

        # Candidate: reset gate filters how much h_{t-1} contributes
        n = torch.tanh(self.W_n_h(r * h_prev) + self.W_n_x(x_t))

        # Interpolate: z=1 means full update, z=0 means carry forward
        h_t = (1 - z) * h_prev + z * n
        return h_t


# ─────────────────────────────────────────────────────────────────────────
# 3. Full sequence processors (many-to-many and many-to-one)
# ─────────────────────────────────────────────────────────────────────────
class LSTMSequence(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=1):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers  = num_layers
        # Stack of LSTM cells — each layer takes previous layer's h sequence as input
        self.cells = nn.ModuleList([
            LSTMCellManual(input_size if l == 0 else hidden_size, hidden_size)
            for l in range(num_layers)
        ])
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        """x: (batch, seq_len, input_size)"""
        B, T, _ = x.shape
        # Initialise all layers' h and c to zeros
        h = [torch.zeros(B, self.hidden_size) for _ in range(self.num_layers)]
        c = [torch.zeros(B, self.hidden_size) for _ in range(self.num_layers)]

        for t in range(T):
            layer_input = x[:, t, :]
            for l, cell in enumerate(self.cells):
                h[l], c[l] = cell(layer_input, h[l], c[l])
                layer_input = h[l]   # next layer reads this layer's h

        return self.fc(h[-1])   # classify from final layer's final hidden state


class GRUSequence(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.hidden_size = hidden_size
        self.cell = GRUCellManual(input_size, hidden_size)
        self.fc   = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        B, T, _ = x.shape
        h = torch.zeros(B, self.hidden_size)
        for t in range(T):
            h = self.cell(x[:, t, :], h)
        return self.fc(h)


# ─────────────────────────────────────────────────────────────────────────
# 4. Long-range dependency task — first token determines label
#    The model must remember token[0] across the full sequence
# ─────────────────────────────────────────────────────────────────────────
def make_long_range_batch(batch_size, seq_len, input_size):
    """Label = sign of first token's first dimension. Model must remember it."""
    x = torch.randn(batch_size, seq_len, input_size)
    y = (x[:, 0, 0] > 0).long()   # label from FIRST token only
    return x, y

def train_compare(seq_len=40, steps=300):
    """Train RNN, LSTM, GRU on a long-range task and compare final accuracy."""
    input_sz, hidden_sz = 8, 32

    from rnn_day20 import SimpleRNN   # conceptual import — use nn.RNN below instead
    models = {
        "RNN ":  nn.Sequential(nn.RNN(input_sz, hidden_sz, batch_first=True)),
        "LSTM": LSTMSequence(input_sz, hidden_sz, output_size=2),
        "GRU ": GRUSequence(input_sz, hidden_sz, output_size=2),
    }

    results = {}
    for name, model in models.items():
        if name.strip() == "RNN":   # wrap nn.RNN output
            continue
        optim = torch.optim.Adam(model.parameters(), lr=1e-3)
        correct = 0
        for step in range(steps):
            x, y = make_long_range_batch(64, seq_len, input_sz)
            optim.zero_grad()
            out  = model(x)
            loss = nn.CrossEntropyLoss()(out, y)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optim.step()
            if step == steps - 1:
                correct = (out.argmax(1) == y).float().mean().item()
        results[name] = correct
    return results


# ─────────────────────────────────────────────────────────────────────────
# 5. Sanity checks — shapes, gradient norms, parameter counts
# ─────────────────────────────────────────────────────────────────────────
batch, seq_len, input_sz, hidden_sz = 4, 20, 8, 16
x = torch.randn(batch, seq_len, input_sz)
y = torch.randint(0, 2, (batch,))

lstm_model = LSTMSequence(input_sz, hidden_sz, output_size=2, num_layers=2)
gru_model  = GRUSequence(input_sz, hidden_sz, output_size=2)
pt_lstm    = nn.LSTM(input_sz, hidden_sz, batch_first=True)
pt_gru     = nn.GRU(input_sz, hidden_sz, batch_first=True)

# Forward pass
out_lstm = lstm_model(x)
out_gru  = gru_model(x)
pt_out_lstm, (h_n, c_n) = pt_lstm(x)
pt_out_gru,   h_n_gru   = pt_gru(x)

print(f"Manual LSTM output:  {out_lstm.shape}")
print(f"Manual GRU  output:  {out_gru.shape}")
print(f"PyTorch LSTM all h:  {pt_out_lstm.shape} | final h: {h_n.shape} | final c: {c_n.shape}")
print(f"PyTorch GRU  all h:  {pt_out_gru.shape}  | final h: {h_n_gru.shape}")

# Parameter counts
lstm_params = sum(p.numel() for p in lstm_model.parameters())
gru_params  = sum(p.numel() for p in gru_model.parameters())
print(f"\nParameter counts (2-layer LSTM vs 1-layer GRU, h={hidden_sz}, x={input_sz}):")
print(f"  Manual 2-layer LSTM: {lstm_params:,}")
print(f"  Manual 1-layer GRU:  {gru_params:,}")

# Gradient norm after backprop — LSTM should maintain better gradient to early steps
loss = nn.CrossEntropyLoss()(out_lstm, y)
loss.backward()
grad_norm = sum(p.grad.norm()**2 for p in lstm_model.parameters() if p.grad is not None)**0.5
print(f"\nLSTM gradient norm (seq_len={seq_len}): {grad_norm:.4f}")
print(f"Forget gate bias init: {lstm_model.cells[0].W.bias[0].item():.1f} (should be 1.0)")

# PyTorch built-in LSTM: packed sequence for variable-length inputs
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
lengths = torch.tensor([20, 15, 10, 8])   # variable-length sequences in batch
packed  = pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=True)
packed_out, _ = pt_lstm(packed)
unpacked, _   = pad_packed_sequence(packed_out, batch_first=True)
print(f"\nPacked sequence output: {unpacked.shape} (no wasted compute on padding)")

10 — Where LSTM and GRU Still Win in 2026

Transformers dominate NLP. But LSTM and GRU are not dead — they remain the right tool in specific, well-defined scenarios:

11 — The Remaining Bottleneck: Why Even LSTM Isn't Enough for Translation

LSTM and GRU solve the vanishing gradient. But there's a second problem they don't solve — the information bottleneck in seq2seq models.

Recall the seq2seq architecture from Day 20: an encoder reads the full source sentence and compresses it into a single context vector (the final hidden state h_T). A decoder then generates the target sentence conditioned on this one vector.

For short sentences, this works well. For longer sentences (20+ words), the fixed-size context vector becomes a bottleneck — it simply cannot losslessly encode everything. Google's GNMT (2016) using 8-layer stacked LSTMs achieved impressive results but still degraded on long sentences.

The fix — introduced by Bahdanau et al. in 2015 — was attention: instead of compressing everything into one vector, let the decoder dynamically look at all encoder hidden states at each decoding step, and weight them by relevance. This is the direct ancestor of the transformer's self-attention. Day 24 covers it in full.

12 — A Thought to Sit With

13 — Coming Up Next

Tomorrow — Day 22: Tokenization. Before any model — RNN, LSTM, or transformer — can process text, that text must be converted into tokens and then into integer IDs. How you split text into tokens is not obvious, and the choice has profound effects on model vocabulary size, multilingual capability, handling of rare words, and even mathematical reasoning ability. We cover whitespace tokenisation and its failures, then BPE (used by GPT), WordPiece (used by BERT), and SentencePiece (used by LLaMA and multilingual models) — with full implementations.

14 — Go Deeper

▶ Illustrated Guide to LSTM and GRU — Michael Phi
YouTube · 13 min · Clear animated gate-by-gate walkthrough. Best visual complement to the equations in this post.

🔗 Understanding LSTMs — Christopher Olah (2015)
The most-read blog post on LSTMs ever written. The gate diagrams are widely reproduced. Read this after this post for a visual reinforcement of everything we derived.

📄 Long Short-Term Memory — Hochreiter & Schmidhuber (1997)
The original LSTM paper. Section 2 (constant error carousel) is the key intuition — the authors' way of explaining why additive cell state prevents vanishing gradients. 27 years old, still worth reading.

📄 Learning Phrase Representations Using RNN Encoder-Decoder — Cho et al. (2014)
The GRU paper. Also the seq2seq paper. Section 3 introduces the GRU equations; Section 4 is the encoder-decoder architecture. Both in 8 pages.

📄 Neural Machine Translation by Jointly Learning to Align and Translate — Bahdanau et al. (2015)
The attention paper. The direct predecessor of transformer attention. Figure 3 (alignment visualisation) is one of the most beautiful illustrations in NLP — a heat map showing which source words the decoder attends to when generating each target word.