Long Short-Term Memory (LSTM) cells

The Vanishing Gradient Problem (Why We Need LSTMs)

In a standard RNN, hidden state $h_t = \tanh(W_{hh} h_{t-1} + W_{xh} x_t + b_h)$. During backpropagation through time (BPTT), we compute:

$\frac{\partial L}{\partial h_{t-k}} = \frac{\partial L}{\partial h_t} \prod_{i=1}^{k} \frac{\partial h_{t-i+1}}{\partial h_{t-i}}$

Why this step? The chain rule forces us to multiply Jacobians across all time steps. Each Jacobian $\frac{\partial h_{t-i+1}}{\partial h_{t-i}}$ contains $W_{hh}$ and the derivative of $\tanh$, which is $\leq 1$.

If $||W_{hh}|| < 1$, the product $\prod_{i=1}^{k} W_{hh}$ shrinks exponentially as $k$ grows (vanishing gradient). If $||W_{hh}|| > 1$, it explodes. Either way, learning long-term dependencies becomes impossible.

The LSTM fix: The cell state $C_t$ flows through time with additive updates, not multiplicative ones. This creates gradient highways that preserve information.

LSTM Architecture: The Four Gates Derivation

An LSTM has:

• Input $x_t$ (current input vector)
• Previous hidden state $h_{t-1}$
• Previous cell state $C_{t-1}$
• Three gates: forget, input, output
• Outputs: new cell state $C_t$ and hidden state $h_t$

Step 1: Forget Gate (What to Discard)

$f_t = \sigma(W_f \cdot [h_{t-1}, x_t] + b_f)$

What: A vector in $[0, 1]^{n_c}$ where $n_c$ is cell state dimension.

Why sigmoid? We want values between 0 (completely forget) and 1 (completely remember). Sigmoid $\sigma(z) = \frac{1}{1 + e^{-z}}$ outputs this range.

How it works: $f_t$ is element-wise multiplied with $C_{t-1}$. If $f_t[i] = 0$, dimension $i$ of the cell state is erased. If $f_t[i] = 1$, it's preserved.

Notation: $[h_{t-1}, x_t]$ means concatenating the two vectors: if $h_{t-1} \in \mathbb{R}^{n_h}$ and $x_t \in \mathbb{R}^{n_x}$, then $[h_{t-1}, x_t] \in \mathbb{R}^{n_h + n_x}$. The weight matrix $W_f \in \mathbb{R}^{n_c \times (n_h + n_x)}$.

Step 2: Input Gate (What New Info to Store)

$i_t = \sigma(W_i \cdot [h_{t-1}, x_t] + b_i)$ $\tilde{C}_t = \tanh(W_C \cdot [h_{t-1}, x_t] + b_C)$

What: $i_t$ is a gate vector in $[0, 1]^{n_c}$; $\tilde{C}_t$ is a candidate cell state in $[-1, 1]^{n_c}$.

Why two parts? We separate what to add ($\tilde{C}_t$) from how much to add ($i_t$). The candidate $\tilde{C}_t$ is computed using $\tanh$ (outputs $[-1,1]$) because cell states need signed values to represent both positive and negative information.

Why this step? Without the gate $i_t$, we'd always add the full candidate. The gate lets the network learn when new information is relevant (gate open) vs. when to ignore it (gate closed).

Step 3: Update Cell State (Combine Forget + Input)

$C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t$

Notation: $\odot$ is element-wise (Hadamard) product.

Derivation from scratch:

1. Start with old cell state $C_{t-1}$
2. Multiply by forget gate: $f_t \odot C_{t-1}$ (keeps relevant memories)
3. Compute what to add: $i_t \odot \tilde{C}_t$ (new relevant info)
4. Add them: $C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t$

Why this is brilliant: This update is additive, not multiplicative! During backpropagation:

$\frac{\partial C_t}{\partial C_{t-1}} = f_t$

The gradient flows through $f_t$ without repeated matrix multiplications. If $f_t \approx 1$ (forget gate open), gradients flow unchanged. This is the gradient highway that solves vanishing gradients.

Step 4: Output Gate (What to Expose)

$o_t = \sigma(W_o \cdot [h_{t-1}, x_t] + b_o)$ $h_t = o_t \odot \tanh(C_t)$

What: $o_t$ controls which parts of the cell state to expose as the hidden state $h_t$.

Why $\tanh(C_t)$? Cell states can grow large over time. $\tanh$ squashes to $[-1, 1]$ for stability.

Why not just $h_t = C_t$? The cell state is the internal memory (can be large, unbounded). The hidden state is the external representation we pass to the next layer—we want it controlled and squashed.

Why this step? The output gate lets the network decide: "I have this memory, but do I want to expose it now?" For example, storing grammatical gender but only outputing it when generating a pronoun.

Complete LSTM Forward Pass

Given inputs $x_1, x_2, \ldots, x_T$, initialize $h_0 = 0$, $C_0 = 0$. For $t = 1$ to $T$:

Parameter count: If $n_x$ = input dim, $n_h$ = hidden dim, $n_c$ = cell dim (usually $n_c = n_h$):

• Each gate has weight matrix $\mathbb{R}^{n_h \times (n_h + n_x)}$ and bias $\mathbb{R}^{n_h}$
• 4 sets of parameters (forget, input, candidate, output)
• Total: $4 \times [n_h \times (n_h + n_x) + n_h]$ parameters

Comparison to vanilla RNN:

• RNN: gradient multiplies by $W_{hh} \cdot \text{diag}(\tanh'(h_{t-1}))$ at each step → exponential decay
• LSTM: gradient multiplies by $f_t$ (values near 1) → linear flow, no exponential decay

The cell state acts as a gradient superhighway—if forget gates stay open ($f_t \approx 1$), gradients flow back hundreds of steps with minimal degradation.

Why LSTMs help: The word "not" appears early but affects "bad" later. An RNN might forget "not" by the time it sees "bad".

LSTM behavior:

• When processing "not", input gate $i_t$ opens, storing negation in $C_t$
• Forget gate $f_t$ stays high through "was" (don't forget negation)
• When "bad" arrives, the cell state still contains "negation", so the network outputs positive

Numerical walkthrough (simplified, 1D):

1. $t=1$ ("not"): $i_1 = 0.9$, $\tilde{C}_1 = 0.8$ → $C_1 = 0.72$ (store negation)
2. $t=2$ ("bad"): $f_2 = 0.95$, $C_2 = 0.95 \times 0.72 + \ldots = 0.68$ (negation preserved)
3. Output gate uses $C_2$ to predict positive sentiment

Why this step? The additive update $C_2 = f_2 \odot C_1 + \ldots$ means $C_2$ directly depends on $C_1$ with minimal information loss, unlike RNN's $h_2 = \tanh(W h_1 + \ldots)$ which can erase $h_1$.

Sequence: $x_1=[, x_2=[, x_3=[, x_4=[ x_5=], \ldots$

LSTM solution:

• Learn $W_C, W_i$ such that $\tilde{C}_t = +1$ when $x_t = [$ and $i_t = 1$
• Learn $f_t = 1$ always (never forget)
• Cell state: $C_1 = 1, C_2 = 2, C_3 = 3, C_4 = 4, \ldots$
• When seeing"]", learn to decrement

Why this works: $C_t = 1 \cdot C_{t-1} + 1 \cdot (+1) = C_{t-1} + 1$

The cell state is literally counting! The additive update enables this. A vanilla RNN with $h_t = \tanh(W h_{t-1} + \ldots)$ can't maintain an exact count because $\tanh$ bounds values to $[-1, 1]$.

Why it feels right: The name "forget gate" sounds like a switch.

The fix: Gates are continuous multipliers. $f_t = 0.6$ means every dimension of $C_{t-1}$ gets scaled by 0.6—it's partial forgetting, not random erasure. The network learns smooth interpolations, not hard decisions. This continuity is why LSTMs can be trained with gradient descent.

Steel-man: This confusion arises because we want interpretability, and binary decisions are easier to reason about. But soft gates give the network flexibility to learn nuanced timing.

Why it feels right: Both carry information forward time.

The fix:

• $C_t$ is the unfiltered long-term memory (can be large, unrestricted)
• $h_t = o_t \odot \tanh(C_t)$ is the filtered output (bounded, task-relevant)

Example: When reading "John went to the store. He bought milk", the cell state might store [John, male, past-tense, store-context]. But when generating the next word after "He", the hidden state only exposes [male] (controlled by output gate) because that's what's needed for pronoun consistency.

Why this design? Separating storage ($C_t$) from exposure ($h_t$) lets the network keep rich internal state without forcing every detail into the output.

Why it feels right: We said LSTMs solve vanishing gradients and preserve long-term dependencies.

The fix: LSTMs can learn long dependencies if the data supports it, but they don't memorize perfectly. The forget gate $f_t$ is learned—if the training data doesn't require long-term memory, the network learns to forget. Also, cell states can still saturate or suffer from gradient issues with extremely long sequences (100k+ steps).

Practical limit: LSTMs handle dependencies of ~100-200 steps well. For longer, use attention mechanisms or Transformers.

Backpropagation Through Time in LSTMs

The key advantage: gradients flow through the cell state path with minimal degradation.

Gradient flow equation: $\frac{\partial L}{\partial C_{t-k}} = \frac{\partial L}{\partial C_t} \prod_{i=t-k+1}^{t} f_i$

Why this matters: If forget gates average $f_i \approx 0.9$, after 100 steps: $\prod_{i=1}^{100} 0.9 = 0.9^{100} \approx 2.7 \times 10^{-5}$

Still small, but much better than vanilla RNN where: $\prod_{i=1}^{100} ||W_{hh}|| \cdot |\tanh'|| < 0.25^{100} \approx 10^{-60}$

Why this step? The forget gate values are learned, not fixed. The network can keep $f_t$ close to 1 for important memories, making $\prod f_i$ much larger in practice.

An LSTM is like having a notebook with smart sticky notes:

• The cell state ($C_t$) is your notebook where you write important facts
• The forget gate is a smart eraser—it only erases stuff you don't need anymore
• The input gate is a filter—it decides if new info is important enough to write down
• The output gate is a privacy filter—you have the memory, but you only share what's relevant right now

So if the story is "The wizard, who wore a purple hat, cast a spell", the LSTM:

1. Writes "wizard, purple hat" in the notebook (input gate lets it in)
2. Keeps "wizard" the whole time (forget gate says "don't erase")
3. When the story asks "Who cast the spell?", the output gate shows "wizard" from the notebook

The notebook (cell state) has a straight path through time—info can flow back without getting corrupted. That's why LSTMs remember better than regular RNNs!

Order: FIO-C-O, where C (cell state) is the result of combining F and I, and O comes last.

Connections

• 3.5.01-Recurrent-Neural-Networks-(RNN)-basics: LSTMs extend RNNs with gated memory
• 3.5.02-Backpropagation-Through-Time-(BPTT): Why vanilla BPTT fails, LSTMs fix it
• 3.5.03-Vanishing-and-Exploding-Gradientsin-RNNs: The core problem LSTMs solve
• 3.5.05-Gated-Recurrent-Units-(GRU): Simplified alternative with2 gates instead of 3
• 3.5.06-Bidirectional-RNNs: LSTMs can be made bidirectional
• 3.6.01-Attention-Mechanisms: Modern alternative that handles even longer dependencies
• 4.2.01-Transformers: Replaced LSTMs for many NLP tasks

#flashcards/ai-ml

What problem do LSTMs solve that vanilla RNNs cannot? :: Vanishing gradients during long sequences—gradients shrink exponentially in RNNs but flow additively through LSTM cell states, preserving long-term dependencies.

Concept Map

Hinglish (regional understanding)

Hinglish (regional understanding)