Sequence Models
LSTMs & GRUs
Chapter 2, Section 5 established that plain RNNs struggle to retain information across long sequences, due to vanishing gradients compounding through Backpropag
Jr Codex Deep Learning Notes
Level: Advanced Prerequisites: Chapter 2: Recurrent Neural Networks Time to complete: ~30 minutes
Table of Contents
- The Problem to Solve
- The Core Idea: Gates
- The LSTM Cell, Piece by Piece
- Why LSTMs Fix the Vanishing Gradient Problem
- GRUs — a Simplified Alternative
- LSTM vs GRU vs Plain RNN
- LSTMs and GRUs in PyTorch
- Summary & Next Steps
1. The Problem to Solve
Chapter 2, Section 5 established that plain RNNs struggle to retain information across long sequences, due to vanishing gradients compounding through Backpropagation Through Time. Long Short-Term Memory (LSTM) networks, introduced in 1997 but not widely adopted until the 2010s, were specifically designed to solve this.
2. The Core Idea: Gates
An LSTM's key innovation is adding a separate cell state — a dedicated memory pathway — alongside the hidden state, along with gates: small learned neural networks (each just a dense layer with a sigmoid activation, Module 2, Chapter 2) that control how information flows into, out of, and through this memory.
Why "Gates"
─────────────────────────────────────────
A sigmoid output near 0 means "let almost NOTHING through";
near 1 means "let almost EVERYTHING through" — the gate
acts like a learned, CONTINUOUS on/off switch, applied via
element-wise multiplication with the corresponding
information stream.
─────────────────────────────────────────
3. The LSTM Cell, Piece by Piece
The Three Gates and the Cell State
─────────────────────────────────────────
Forget gate (f_t): decides what to DISCARD from the
PREVIOUS cell state — "how much of
the old memory is still relevant?"
Input gate (i_t): decides what NEW information to
ADD to the cell state — "how
much of the current input is
worth remembering?"
Output gate (o_t): decides what part of the
(updated) cell state to
EXPOSE as the hidden state
— "what's relevant to
output RIGHT NOW?"
Cell state (C_t): the actual long-term
MEMORY — updated by
combining the forget
and input gates
─────────────────────────────────────────
import numpy as np
def sigmoid(z):
return 1 / (1 + np.exp(-z))
def lstm_cell(x_t, h_prev, C_prev, params):
combined = np.concatenate([h_prev, x_t]) # combine hidden state and current input
f_t = sigmoid(params["W_f"] @ combined + params["b_f"]) # forget gate
i_t = sigmoid(params["W_i"] @ combined + params["b_i"]) # input gate
C_candidate = np.tanh(params["W_C"] @ combined + params["b_C"]) # candidate new memory
o_t = sigmoid(params["W_o"] @ combined + params["b_o"]) # output gate
C_t = f_t * C_prev + i_t * C_candidate # UPDATE cell state: forget old + add new
h_t = o_t * np.tanh(C_t) # decide what to EXPOSE as the hidden state
return h_t, C_tData Flow Through One LSTM Step
─────────────────────────────────────────
C(t-1) ──►[× f_t]──►[+]──────────────► C(t)
▲
[i_t × C_candidate]
h(t-1), x(t) ──► compute f_t, i_t, C_candidate, o_t
C(t) ──►[tanh]──►[× o_t]──► h(t)
─────────────────────────────────────────
4. Why LSTMs Fix the Vanishing Gradient Problem
The Key Structural Difference From a Plain RNN
─────────────────────────────────────────
A plain RNN's hidden state is recomputed ENTIRELY at every
step (Chapter 2, Section 2) — repeatedly passing through
tanh and a weight matrix multiplication, which is exactly
what causes gradients to shrink or grow exponentially
(Chapter 2, Section 5).
An LSTM's CELL STATE update (C_t = f_t * C_prev + ...) is
ADDITIVE, not purely multiplicative through a weight matrix
— when the forget gate f_t is close to 1, the cell state
can carry gradient information across MANY time steps with
MUCH LESS decay, since it's not repeatedly squashed through
tanh and re-multiplied by a weight matrix at every step the
way a plain RNN's hidden state is.
─────────────────────────────────────────
This is conceptually similar to Module 4, Chapter 2's ResNet residual connections — both provide a more direct path for gradients to flow across many steps (time steps here, layers there), sidestepping the repeated multiplicative shrinkage that causes vanishing gradients.
5. GRUs — a Simplified Alternative
The Gated Recurrent Unit (GRU), introduced in 2014, simplifies the LSTM by combining the forget and input gates into a single update gate, and merging the cell state and hidden state into one.
def gru_cell(x_t, h_prev, params):
combined = np.concatenate([h_prev, x_t])
z_t = sigmoid(params["W_z"] @ combined + params["b_z"]) # update gate (combines forget + input)
r_t = sigmoid(params["W_r"] @ combined + params["b_r"]) # reset gate
combined_reset = np.concatenate([r_t * h_prev, x_t])
h_candidate = np.tanh(params["W_h"] @ combined_reset + params["b_h"])
h_t = (1 - z_t) * h_prev + z_t * h_candidate # blend OLD hidden state and NEW candidate
return h_tGRU vs LSTM — Structural Difference
─────────────────────────────────────────
LSTM: 3 gates (forget, input, output), SEPARATE cell
state and hidden state — more expressive,
more parameters
GRU: 2 gates (update, reset), a SINGLE combined
state — fewer parameters, often trains
faster, with COMPARABLE performance on many
tasks
─────────────────────────────────────────
6. LSTM vs GRU vs Plain RNN
Practical Guidance
─────────────────────────────────────────
Plain RNN: rarely used in practice today, except as a
teaching tool (Chapter 2) — vanishing
gradients make it unreliable for anything
beyond very short sequences
LSTM: the traditional default for tasks needing
to capture longer-range dependencies —
more parameters, slightly slower to
train, often marginally more accurate
GRU: a common alternative when training
speed or a smaller parameter count
matters — frequently performs
comparably to LSTM despite being
simpler
Modern practice: for the LONGEST-range dependencies
and best overall performance,
both are increasingly superseded
by Transformers (Module 6) —
this is exactly Chapter 5's
closing topic
─────────────────────────────────────────
7. LSTMs and GRUs in PyTorch
import torch
import torch.nn as nn
lstm = nn.LSTM(input_size=10, hidden_size=20, batch_first=True)
gru = nn.GRU(input_size=10, hidden_size=20, batch_first=True)
sequence = torch.randn(1, 5, 10) # (batch, time_steps, features)
lstm_output, (h_n, c_n) = lstm(sequence) # LSTM returns hidden state AND cell state
print(f"LSTM output: {lstm_output.shape}, hidden: {h_n.shape}, cell: {c_n.shape}")
gru_output, h_n_gru = gru(sequence) # GRU returns ONLY hidden state (Section 5's simplification)
print(f"GRU output: {gru_output.shape}, hidden: {h_n_gru.shape}")
# A stacked (multi-layer) LSTM — common for more complex sequence tasks
stacked_lstm = nn.LSTM(input_size=10, hidden_size=20, num_layers=2, batch_first=True, dropout=0.3)8. Summary & Next Steps
Key Takeaways
- LSTMs introduce a separate cell state and three gates (forget, input, output) that control information flow, specifically to address plain RNNs' vanishing gradient problem.
- The cell state's additive update rule lets gradients flow across many time steps with far less decay than a plain RNN's fully multiplicative hidden state update.
- GRUs simplify the LSTM design (2 gates, one combined state), often training faster with comparable performance on many tasks.
- Both LSTMs and GRUs meaningfully outperform plain RNNs on long sequences, but are increasingly superseded by Transformers (Module 6) for the longest-range dependencies.
Concept Check
- What specific structural difference lets an LSTM's cell state avoid the vanishing gradient problem that plain RNNs suffer from?
- What are the three gates in an LSTM, and what does each one control?
- What's the main tradeoff between choosing an LSTM versus a GRU for a given task?
Next Chapter
→ Chapter 4: Sequence-to-Sequence & Encoder-Decoder Models
Jr Codex — 1-on-1 Personalized Coaching | Back to Module Index | Back to DL Index