Long Short-Term Memory (LSTM)

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Vanilla RNNs and the Vanishing Gradient
5. The LSTM Architecture
6. Gate Equations in Detail
7. Gradient Flow Through the Cell State
8. The GRU Simplification
9. Bidirectional and Stacked RNNs
10. Common Pitfalls
11. Research Perspective
12. Summary
13. Exercises

Learning Objectives

1. Derive the vanilla RNN gradient and identify the vanishing/exploding gradient conditions.
2. Explain each LSTM gate and its role in memory control.
3. Prove that the cell state gradient has a multiplicative path through forget gates only.
4. Derive the GRU equations as a simplified LSTM with merged gates.
5. Analyze when LSTMs/GRUs are preferred over transformers.

Notation

• $\mathbf{x}_t \in \mathbb{R}^d$ — input at time $t$
• $\mathbf{h}_t \in \mathbb{R}^n$ — hidden state
• $\mathbf{c}_t \in \mathbb{R}^n$ — cell state (LSTM only)
• $\mathbf{f}_t, \mathbf{i}_t, \mathbf{o}_t \in (0,1)^n$ — forget, input, output gates
• $\tilde{\mathbf{c}}_t$ — candidate cell state
• $\odot$ — element-wise (Hadamard) product
• $\sigma(\cdot)$ — sigmoid function

Core Intuition

Vanilla RNNs struggle with long-range dependencies because gradients either vanish or explode when backpropagated through many time steps. The LSTM solves this by introducing a cell state $\mathbf{c}_t$ that flows through time with minimal transformation — analogous to a skip connection across time. Gates control what information to store, forget, and output, allowing the network to selectively maintain information over hundreds of time steps.

Vanilla RNNs and the Vanishing Gradient

Vanilla RNN:

Gradient through time. For loss at time $T$, the gradient w.r.t. $\mathbf{h}_t$ ($t < T$):

Vanishing: If $\lVert\text{diag}(\tanh')\mathbf{W}_h\rVert < 1$, the product $\to 0$ exponentially in $(T-t)$.

Exploding: If $\lVert\text{diag}(\tanh')\mathbf{W}_h\rVert > 1$, the product $\to \infty$.

Since $\tanh'(z) \in (0, 1]$ and is typically $< 1$, vanishing dominates for long sequences.

The LSTM Architecture

The LSTM (Hochreiter & Schmidhuber, 1997) introduces:

1. A cell state $\mathbf{c}_t$ — a memory highway with additive updates.
2. Three gates — sigmoid-activated controllers that modulate information flow.

The cell state update is additive (like a residual connection through time):

Gate Equations in Detail

Forget gate (what to erase from memory):

Input gate (what new information to store):

Candidate cell state (proposed new content):

Cell state update:

Output gate (what to expose from memory):

Hidden state:

Gradient Flow Through the Cell State

The key advantage: the gradient through the cell state from $t$ to $t-1$:

Over $T-t$ steps:

Key insight: The forget gate $\mathbf{f}_k \in (0,1)^n$ is learned. If the network sets $f_k \approx 1$ for some dimensions, the gradient flows nearly unattenuated through those dimensions, enabling long-range memory. Compare with the vanilla RNN where the "forget factor" $\tanh'(z)\mathbf{W}_h$ is not directly controllable.

Analogy: The cell state is a "conveyor belt" — information travels forward with the forget gate controlling which items stay on the belt.

The GRU Simplification

The Gated Recurrent Unit (Cho et al., 2014) merges the forget and input gates and removes the cell state:

Update gate: $\mathbf{z}_t = \sigma(\mathbf{W}_z[\mathbf{h}_{t-1}, \mathbf{x}_t])$

Reset gate: $\mathbf{r}_t = \sigma(\mathbf{W}_r[\mathbf{h}_{t-1}, \mathbf{x}_t])$

Candidate: $\tilde{\mathbf{h}}_t = \tanh(\mathbf{W}[\mathbf{r}_t \odot \mathbf{h}_{t-1}, \mathbf{x}_t])$

State update:

Comparison to LSTM:

• Fewer parameters (2 gates instead of 3, no separate cell state).
• The update $\mathbf{z}_t$ plays dual role: $\mathbf{f}_t = 1 - \mathbf{z}_t$ and $\mathbf{i}_t = \mathbf{z}_t$.
• Empirically similar performance for most tasks.

Bidirectional and Stacked RNNs

Bidirectional: Run two RNNs — one forward, one backward — and concatenate:

Captures both past and future context (not applicable for autoregressive generation).

Stacked (deep): Multiple RNN layers, each taking the previous layer's hidden states as input. Depth increases representational power.

Common Pitfalls

Pitfall 1. Initializing forget gate biases to zero. This means $f_t \approx 0.5$ initially — the network starts by forgetting half the cell state. Best practice: initialize $\mathbf{b}_f$ to 1 or 2 so $f_t \approx 1$ initially (remember by default).

Pitfall 2. Not using gradient clipping. Even with LSTMs, gradients can still explode through the hidden-to-hidden paths in the gates. Clip gradient norms to a maximum value.

Pitfall 3. Using LSTMs where transformers suffice. For tasks with fixed-length inputs and parallelizable training, transformers are faster and often more accurate. LSTMs remain relevant for online/streaming settings.

Research Perspective

The LSTM (1997) was one of the most impactful architectural innovations in deep learning, enabling breakthroughs in speech recognition, machine translation, and language modeling before transformers. The gating mechanism inspired many subsequent architectures. Today, state-space models (Mamba, S4) attempt to combine the LSTM's sequential efficiency with the transformer's parallelism.

Summary

• Vanilla RNNs suffer from vanishing gradients: $\prod_k \tanh'(z_k)\mathbf{W}_h \to 0$.
• LSTM adds a cell state with additive updates: gradient flows through $\prod_k \mathbf{f}_k$.
• Three gates (forget, input, output) control memory read/write/erase.
• GRU simplifies to two gates and no separate cell state.
• Forget gate bias initialization to 1 is critical for long-range dependencies.
• LSTMs are the temporal analog of residual connections.

Exercises

Exercise 1. For a vanilla RNN with $\mathbf{W}_h = 0.9\mathbf{I}$ and $\tanh$ activation, compute the gradient magnitude after 50 time steps.

Exercise 2. Derive equation (11) by recursive application of the chain rule through the cell state.

Exercise 3. Show that with $\mathbf{f}_t = \mathbf{1}$ and $\mathbf{i}_t = \mathbf{0}$, the LSTM perfectly preserves the cell state (infinite memory).

Exercise 4. Prove that the GRU update (equation 12) is a convex combination of $\mathbf{h}_{t-1}$ and $\tilde{\mathbf{h}}_t$.

Exercise 5. Count the total number of parameters in a single LSTM layer with input dimension $d$ and hidden dimension $n$.