GenAI Course · Module 4 · Session 15

LSTM & GRU

⏱ 14 min read🎯 Intermediate🖼 2 diagrams🧪 1 live demo🐍 Python code

Plain RNNs forget after about 20 words. The LSTM (Long Short-Term Memory, 1997) adds a separate memory lane and three gates that decide what to keep, what to write and what to show. It powered Google Translate, Siri and most NLP until transformers arrived. Let us open it up and prove that the gates really keep gradients alive.

In this session

The cell state: a memory highway
Forget, input and output gates
One LSTM step in NumPy with real gate values
Measured: how the forget gate stops gradients vanishing
Bidirectional LSTMs and GRUs
Parameter counts: RNN vs GRU vs LSTM

01The LSTM cell

The LSTM’s cell state runs straight through, changed only by multiplication with the forget gate and addition of new info. That “highway” lets gradients survive long sequences.

Part	Formula	Job
Forget gate f	σ(W_f · [hₜ₋₁, xₜ] + b_f)	0 = forget the old memory, 1 = keep it
Input gate i	σ(W_i · [hₜ₋₁, xₜ] + b_i)	how much new information to write
Candidate g	tanh(W_g · [hₜ₋₁, xₜ] + b_g)	the new information itself
Cell state c	cₜ = f ⊙ cₜ₋₁ + i ⊙ g	long-term memory: keep some old, add some new
Output gate o	σ(W_o · [hₜ₋₁, xₜ] + b_o)	how much memory to reveal
Hidden state h	hₜ = o ⊙ tanh(cₜ)	short-term output for this step

🗂️

Like a bank clerk’s file. For each new document, the clerk decides what old notes to tear out (forget gate), what new notes to add (input gate), and what to tell the customer today (output gate). The file itself (cell state) travels from day to day.

lstm_step.py

import numpy as np
np.random.seed(1)
sig = lambda z: 1 / (1 + np.exp(-z))

X, H = 4, 3
W = {g: np.random.randn(X + H, H) * 0.5 for g in "fiog"}      # one weight matrix per gate
b = {g: np.zeros(H) for g in "fiog"}; b["f"] += 1.0             # common trick: start by remembering

def lstm_step(x, h, c):
    z = np.concatenate([x, h])
    f = sig(z @ W["f"] + b["f"])        # forget gate: how much old memory to keep
    i = sig(z @ W["i"] + b["i"])        # input gate: how much new info to write
    g = np.tanh(z @ W["g"] + b["g"])    # candidate new info
    o = sig(z @ W["o"] + b["o"])        # output gate: how much memory to show
    c = f * c + i * g                   # cell state: the long-term memory
    h = o * np.tanh(c)                  # hidden state: the short-term output
    return h, c, f, i, o

h = c = np.zeros(H)
for t in range(3):
    h, c, f, i, o = lstm_step(np.random.randn(X), h, c)
    print(f"t={t + 1} forget={f.round(2)} input={i.round(2)} output={o.round(2)}\n    cell={c.round(2)} hidden={h.round(2)}")

t=1 forget=[0.61 0.82 0.58] input=[0.59 0.52 0.43] output=[0.45 0.44 0.49] cell=[0.29 0.34 0.15] hidden=[0.12 0.14 0.07] t=2 forget=[0.81 0.75 0.41] input=[0.46 0.36 0.57] output=[0.42 0.64 0.81] cell=[0.09 0.31 0.5 ] hidden=[0.04 0.19 0.38] t=3 forget=[0.78 0.69 0.83] input=[0.33 0.49 0.51] output=[0.57 0.5 0.39] cell=[-0.12 -0.09 0.1 ] hidden=[-0.07 -0.05 0.04]

Every gate value is between 0 and 1 (sigmoid). The forget bias starts at +1, a standard trick so a new LSTM begins by remembering rather than forgetting.

02Why gates stop gradients vanishing

In a plain RNN the gradient is squeezed through tanh and the weights at every step. In an LSTM, the cell state is updated by multiplying by f and adding. If the network learns to keep f close to 1, the gradient travels back almost unchanged. We measure it, comparing a plain RNN, a new LSTM, and an LSTM whose forget gate is open (f ≈ 0.95, as a trained LSTM learns to do for important information):

bptt_vanish.py

import torch, torch.nn as nn
torch.set_num_threads(1)                    # same results on every run

def grad_to_first_token(cell, T):
    torch.manual_seed(0)
    x = torch.randn(1, T, 8, requires_grad=True)
    out, _ = cell(x)
    out[0, -1].sum().backward()                       # loss depends on the LAST output only
    return x.grad[0, 0].abs().mean().item()           # how much does the FIRST token still matter?

H = 32
torch.manual_seed(0); rnn = nn.RNN(8, H, batch_first=True)
torch.manual_seed(0); lstm = nn.LSTM(8, H, batch_first=True)
torch.manual_seed(0); lstm_open = nn.LSTM(8, H, batch_first=True)
with torch.no_grad():                                 # PyTorch gate order: input, forget, cell, output
    lstm_open.bias_ih_l0[H:2 * H].fill_(1.5); lstm_open.bias_hh_l0[H:2 * H].fill_(1.5)   # forget gate ≈ 0.95: "remember"

print("gradient reaching the first token:")
for T in [5, 20, 50, 100]:
    print(f"  length {T:3d}: RNN {grad_to_first_token(rnn, T):.0e}   LSTM (new) {grad_to_first_token(lstm, T):.0e}"
          f"   LSTM (forget gate open) {grad_to_first_token(lstm_open, T):.0e}")

gradient reaching the first token: length 5: RNN 5e-02 LSTM (new) 1e-02 LSTM (forget gate open) 9e-02 length 20: RNN 4e-06 LSTM (new) 6e-06 LSTM (forget gate open) 7e-02 length 50: RNN 3e-14 LSTM (new) 8e-12 LSTM (forget gate open) 2e-02 length 100: RNN 4e-28 LSTM (new) 2e-22 LSTM (forget gate open) 2e-02

Demo · Gradient reaching the first token

The key insight: with the forget gate open, the gradient after 100 tokens is still ~10⁻², while the plain RNN’s is ~10⁻²⁸. A fresh LSTM is no better than an RNN; it is the learnable forget gate that lets training discover what to remember.

03Bidirectional LSTMs and GRUs

A bidirectional LSTM (BiLSTM) runs one LSTM forwards and one backwards, exactly like the bidirectional RNN in Session 14. BiLSTMs were the standard for NER and classification before BERT.

The GRU (Gated Recurrent Unit, 2014) is a simpler cousin. It merges the cell and hidden state and uses two gates:

Same idea, different trade-off: GRU is lighter, LSTM has more control.

rnn_params.py

import torch.nn as nn
count = lambda m: sum(p.numel() for p in m.parameters())
inp, hid = 128, 256
for name, layer in [("RNN", nn.RNN(inp, hid)), ("GRU", nn.GRU(inp, hid)), ("LSTM", nn.LSTM(inp, hid)),
                    ("Bi-LSTM", nn.LSTM(inp, hid, bidirectional=True))]:
    print(f"{name:8s} {count(layer):>9,} parameters")

RNN 98,816 parameters GRU 296,448 parameters LSTM 395,264 parameters Bi-LSTM 790,528 parameters

Check: an RNN layer has 256 × (128 + 256) + 2 × 256 = 98,816 parameters. A GRU has exactly 3× that, an LSTM 4× (one block per gate plus the candidate), and a BiLSTM double the LSTM.

lstm_classifier.py

import torch
import torch.nn as nn

class SentimentLSTM(nn.Module):
    def __init__(self, vocab=20_000, emb=128, hidden=256, classes=3):
        super().__init__()
        self.emb = nn.Embedding(vocab, emb, padding_idx=0)
        self.lstm = nn.LSTM(emb, hidden, batch_first=True, bidirectional=True)
        self.head = nn.Linear(2 * hidden, classes)        # forward + backward states

    def forward(self, token_ids):
        _, (h, _) = self.lstm(self.emb(token_ids))
        return self.head(torch.cat([h[-2], h[-1]], dim=1))  # last forward + last backward state

Use…	When
GRU	smaller data, faster training, mobile/edge devices
LSTM	long sequences, when you need the extra control
Transformer (Module 5+)	almost all NLP today: parallel training and much longer memory

04Summary

The LSTM adds a cell state highway and three gates: forget, input, output.
cₜ = f ⊙ cₜ₋₁ + i ⊙ g: keep some old memory, add some new.
A forget gate near 1 keeps gradients alive for 100+ steps (measured: ~10⁻² vs ~10⁻²⁸).
GRU: 2 gates, no separate cell, ~25% fewer parameters.
LSTM = 4× and GRU = 3× an RNN’s parameters.

05Quick quiz

1. Which gate decides how much old memory to keep?

2. Why do gate values lie between 0 and 1?

3. How does the LSTM fight vanishing gradients?

4. How many gates does a GRU have?

5. An LSTM layer has about how many times the parameters of a plain RNN of the same size?

06References

Understanding LSTM Networks (Chris Olah)
Long Short-Term Memory (Hochreiter & Schmidhuber, 1997)
Learning Phrase Representations using RNN Encoder–Decoder (GRU) (Cho et al., 2014)
PyTorch nn.LSTM

Created by Ajeetkumar