Deep Learning

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

JrCodex·6 min read

Jr Codex Deep Learning Notes

Level: Advanced Prerequisites: Chapter 2: Recurrent Neural Networks Time to complete: ~30 minutes


Table of Contents

  1. The Problem to Solve
  2. The Core Idea: Gates
  3. The LSTM Cell, Piece by Piece
  4. Why LSTMs Fix the Vanishing Gradient Problem
  5. GRUs — a Simplified Alternative
  6. LSTM vs GRU vs Plain RNN
  7. LSTMs and GRUs in PyTorch
  8. 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_t
Data 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_t
GRU 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

  1. What specific structural difference lets an LSTM's cell state avoid the vanishing gradient problem that plain RNNs suffer from?
  2. What are the three gates in an LSTM, and what does each one control?
  3. 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