Sequence Models
LSTM Networks — Long Short-Term Memory for Sequential Data
LSTM networks solve the vanishing gradient problem of vanilla RNNs by using gating mechanisms to control information flow across time steps. The cell state acts as an information highway.
- Three Gates Control Flow — Forget, input, and output gates decide what to discard, store, and output
- Cell State Prevents Vanishing — Linear addition preserves gradient magnitude across hundreds of time steps
- Bidirectional Uses Context — Processing in both directions captures past and future information
"LSTMs gave RNNs a long-term memory that doesn't fade."
LSTM Networks — Gates, Cell State and Bidirectional Architectures
What You'll Learn
- 1Write the four LSTM gate equations and the cell state update from memory
- 2Explain why the additive cell update keeps the gradient close to the identity
- 3Count LSTM parameters and compare them against a GRU of the same hidden size
- 4Choose between unidirectional, bidirectional and stacked LSTMs for a given task
LSTM (Long Short-Term Memory) networks solve the vanishing gradient problem of vanilla RNNs by using gating mechanisms to control information flow across time steps. Where a vanilla RNN multiplies the gradient by at every step, an LSTM adds a cell state whose update is a sum, so the gradient has an unattenuated path back through time.
See our RNN Deep Dive tutorial for the fundamentals of vanilla RNNs and why LSTMs were invented.
Why LSTMs?
LSTM Cell Architecture
The Three Gates
Forget Gate
Input Gate
Cell State Update
Output Gate
Complete LSTM Equations
All four gates read the same concatenation of size , and each one needs its own matrix plus a bias. That is why an LSTM layer costs a full third more parameters than a GRU of the same width:
How this diagram works: This diagram shows the internal architecture of a single LSTM cell. The cell state (blue bar at top) acts as an information highway, flowing across time steps with only linear interactions (multiplication and addition). The three gates — forget (red), input (green), and output (purple) — are sigmoid neural networks that output values between 0 and 1, acting as soft switches. The forget gate controls what to discard from the cell state, the input gate controls what new information to store, and the output gate controls what to reveal. The critical insight is that the cell state update is linear: . Because gradients flow through addition (not matrix multiplication), they are preserved across many time steps, solving the vanishing gradient problem that plagues vanilla RNNs.
Why LSTMs Solve Vanishing Gradients
The contrast with a vanilla RNN is exact: that product carries at every step and decays like , while the LSTM product decays like the average forget gate value raised to . With the bias initialized to 1.0 the gates start at , and training pushes them above 1 where the product stops shrinking at all.
Bidirectional LSTM
Stacked LSTM
Two bidirectional layers at over a 512 dimensional input therefore hold about 8.4 million parameters before a single embedding row, which is why stacking is the first thing to shrink when a small dataset starts to overfit.
PyTorch Implementation
LSTM vs. GRU Comparison
| Feature | LSTM | GRU |
|---|---|---|
| Gates | 3 (forget, input, output) | 2 (update, reset) |
| Cell state | Yes | No |
| Parameters | More () | Fewer () |
| Speed | Slower | Faster |
| Performance | Similar | Similar |
| Memory | Better for very long sequences | Competitive |
Summary
Practice Exercises
-
Mathematical: Derive the gradient for the LSTM cell state update. Why does this prevent vanishing gradients?
-
Coding: Implement an LSTM cell from scratch using only
torch.tensoroperations (nonn.LSTM). Verify that your implementation matches PyTorch's output. -
Experiment: Compare LSTM vs. GRU vs. Transformer on a long-range dependency task (e.g., copying the first element of a 500-element sequence). Which handles long sequences best?
-
Application: Build a bidirectional LSTM for sentiment analysis on IMDB reviews. Compare with a Transformer-based model. What are the tradeoffs?
-
Research: Read the original LSTM paper (Hochreiter and Schmidhuber, 1997). What was the original motivation? How has the architecture evolved?
What to Learn Next
-> GRU Networks Gated Recurrent Unit — a simpler alternative to LSTM with similar performance.
-> RNN Deep Dive Vanilla RNNs, BPTT, and why vanishing gradients motivated the invention of LSTM and GRU.
-> Sequence to Sequence Encoder-decoder architectures for translation, summarization, and text generation.
-> Attention Mechanisms Self-attention, multi-head attention, and how Transformers replaced RNNs for sequence modeling.
-> Vision Transformers ViT, DeiT, Swin — how Transformers are challenging CNN dominance in computer vision.
-> DL Systems Design Designing production ML systems — data pipelines, training infrastructure, and serving at scale.
Common Mistakes
Leaving the forget gate bias at zero
Fix: Set b_f to 1.0 or 2.0 so the gate opens at sigma(1) = 0.73 and the cell keeps memory instead of shedding half of it each step.
Running a bidirectional LSTM on streaming data
Fix: A bidirectional model reads the future, so use a unidirectional LSTM online and keep the bidirectional one for offline tagging.
Stacking four layers on a small dataset
Fix: Each layer multiplies the budget by 4(n^2 + n d + n); add depth only after training loss has plateaued with two.
Best Practices
- Clip the global gradient norm at 1.0 and pack padded sequences so the final hidden state comes from a real token
- Classify from the last layer's hidden state, concatenating the two directions when the model is bidirectional
- Benchmark against a GRU of the same width before adding depth: on short sequences it usually lands within 0.5 F1 and trains faster
Key Takeaways
- Forget, input and output gates all read [h_{t-1}, x_t]; the first two write the cell state and the third filters it
- The cell update is additive, so d c_t over d c_{t-1} stays close to diag(f_t) and gradients survive hundreds of steps
- One LSTM layer costs 4(n^2 + n d + n) parameters, exactly a third more than a GRU at the same width
- Bidirectional doubles parameters and needs the full sequence, while stacking multiplies the budget by the layer count
Frequently Asked Questions
Why does adding a cell state fix the vanishing gradient?
In a vanilla RNN the gradient multiplies at every step, so its norm decays like the spectral norm raised to the th power. The LSTM cell state is updated by addition, and its gradient carries one diagonal factor per step instead of a full matrix. When the forget gate sits near 1 the product of those diagonals stays near the identity, so a signal from 500 steps back still arrives with most of its magnitude.
Is the forget gate bias initialization worth the trouble?
Yes. With bias 0 the gate starts at , so the cell halves its memory every step at the start of training and the gradient must fight that decay for thousands of updates. With bias 1.0 the gate opens at , and with bias 2.0 it reaches 0.88. Most frameworks now ship a positive forget bias for exactly this reason.
When should you use a bidirectional LSTM?
Only when the whole sequence is available before you predict. Named entity recognition and sentiment tagging read the full sentence, so both directions help and the doubled cost buys real accuracy. A streaming forecaster or an autoregressive decoder cannot see the future, and running it bidirectionally would leak the answer into its own input.
Interview Questions
Q1: Give the parameter count of a single LSTM layer.
Each of the four gates needs an n x (n + d) matrix plus n biases, so P = 4 (n^2 + n d + n). At n = d = 512 that is 4 times (262,144 + 262,144 + 512) = 2,099,200 parameters per direction per layer.
Q2: Why does the LSTM not guarantee long range memory?
The cell gradient still passes through one diagonal factor per step, so a forget gate drifting to 0.6 over 500 steps multiplies out to about 1e-118. The architecture makes long memory possible by keeping gates near 1, it does not enforce it.
Q3: How do you set the forget bias in a trained PyTorch LSTM?
Iterate over model.lstm.named_parameters, find the bias vector, which holds four blocks of n values in the order input, forget, output, cell, and fill the forget slice with 1.0. That puts sigma(1) = 0.73 at step zero.