LSTM — Module 8.1 | Codextroop

How This Lesson Fits the Module & Engineering Practice

By this point in the module, the main weakness of a simple RNN is clear: it struggles to preserve useful signal across long sequences because training through time is unstable. Long Short-Term Memory (LSTM) networks were designed to address that weakness by giving the model explicit control over what to keep, forget, and expose.

LSTM became one of the most influential sequence architectures in practical machine learning, powering forecasting, speech systems, and early high-performing NLP pipelines before the later rise of transformer-based models in Volume 10 and beyond.

Learning Objectives

By the end of this lesson, students should be able to:

Explain why LSTM was introduced as an improvement over simple RNNs.
Describe the cell state and the input, forget, and output gates conceptually.
Interpret the difference between hidden state and cell state in an LSTM.
Build and read an nn.LSTM model in PyTorch.
Identify common sequence tasks where LSTM is more suitable than a basic RNN.
Understand the tradeoff between expressive gating and increased complexity.

Definition

An LSTM is a gated recurrent architecture that introduces a persistent cell state and learned gates to regulate what information is written, retained, and exposed over time.

Why LSTM Was Needed

Simple RNNs can represent sequential dependence in principle, but in practice they struggle to learn long-range dependencies because gradients either vanish or explode during BPTT. LSTM addresses this by creating a more stable memory path called the cell state.

Instead of repeatedly overwriting one hidden representation in an uncontrolled way, LSTM learns gates that decide how much prior memory to keep, how much new information to write, and how much of the current state to reveal to downstream computation.

The Main Components

Forget gate

Decides what old information to erase.
Protects against irrelevant carryover.
Critical for state maintenance.

Input gate

Controls what new content enters memory.
Filters the current input.
Prevents noisy overwrites.

Output gate

Controls what part of memory is exposed.
Shapes the hidden state.
Supports task-specific readout.

Hidden State Versus Cell State

State	Role	Intuition
Hidden state `h_t`	Exposed working representation	What the model currently shows outwardly
Cell state `c_t`	Longer-term memory channel	What the model tries to preserve internally

This separation is the heart of LSTM. The cell state creates a more direct path for information and gradients to travel across time, while the hidden state acts more like the current visible summary.

Information Flow Through an LSTM Step

1. Forget

Decide which parts of c_(t-1) should remain.

2. Write

Decide which new candidate information should enter memory.

3. Update memory

Produce the new cell state c_t.

4. Expose

Use the output gate to form the hidden state h_t.

PyTorch Example: Sequence Classification with LSTM

import torch
from torch import nn

class LSTMClassifier(nn.Module):
    def __init__(self, input_size=20, hidden_size=32, num_classes=4):
        super().__init__()
        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size, batch_first=True)
        self.head = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        output, (h_n, c_n) = self.lstm(x)
        last_hidden = output[:, -1, :]
        return self.head(last_hidden)

model = LSTMClassifier()
x = torch.randn(12, 15, 20)
logits = model(x)
print(logits.shape)  # torch.Size([12, 4])

Understanding the Returned States

Unlike nn.RNN, an LSTM returns both the final hidden state and the final cell state:

output: hidden representations for every step.
h_n: final hidden state.
c_n: final cell state.

Tensor	Typical shape with one layer	Meaning
`output`	`(B, T, H)`	Per-step hidden outputs
`h_n`	`(1, B, H)`	Final hidden summary
`c_n`	`(1, B, H)`	Final memory state

Where LSTM Helps Most

Strong use cases

Longer temporal dependencies.
Noisy sequential signals.
Forecasting or text tasks needing persistent context.

Tradeoffs

More parameters than simple RNNs.
More complex internal dynamics.
Still sequential and harder to parallelize than later architectures.

Common Mistake

Do not assume LSTM “stores everything forever.” It learns to preserve some information, not all information. If the task, data, or training setup is poor, LSTM can still forget useful context.

Misconception

“The gates are hand-written rules.” They are not. The forget, input, and output gates are learned differentiable functions trained from data during BPTT.

LSTM as a Bridge to Modern Sequence Modeling

LSTM was a major milestone because it made longer dependencies much more learnable than in plain RNNs. It powered many applications before newer sequence models became dominant. Understanding LSTM is still valuable because it clarifies what later architectures are trying to replace, preserve, or improve.

The next lecture, GRU, shows a streamlined gated alternative. From there, Volume 09 will move toward language-first sequence pipelines, beginning with tokenization.

Knowledge Check

Short Answer: Why was LSTM introduced? Answer: To improve recurrent learning on longer sequences by addressing vanishing-gradient and memory problems in simple RNNs.
True/False: LSTM includes both a hidden state and a cell state. Answer: True.
Multiple Choice: Which gate controls what old information should be kept or removed? (a) forget gate, (b) pooling gate, (c) flatten gate. Answer: (a).
Short Answer: What is the main purpose of the cell state? Answer: To provide a more stable memory path across time.
True/False: LSTM completely eliminates all recurrent training difficulties. Answer: False.
Multiple Choice: In PyTorch, nn.LSTM returns: (a) only logits, (b) output and (h_n, c_n), (c) only c_n. Answer: (b).
Short Answer: Why is LSTM usually better than a simple RNN for long dependencies? Answer: Because its gated memory path helps preserve information and gradients over longer spans.
True/False: LSTM gates are learned from data. Answer: True.
Multiple Choice: Which next lecture covers a simpler gated recurrent alternative? (a) GRU, (b) CNN, (c) Quantization. Answer: (a).
Short Answer: What is one cost of using LSTM instead of a simple RNN? Answer: More parameters and more computational complexity.

Key Takeaways

LSTM improves recurrent modeling by separating exposed hidden state from longer-term cell state.
Its gates regulate forgetting, writing, and exposing information.
PyTorch nn.LSTM returns both hidden and cell state summaries.
LSTM is more robust than a simple RNN on long dependencies, though it is more complex.
Next, GRU introduces a simpler gated recurrent alternative.

Trainer’s Guide

Hands-on idea: Have learners compare nn.RNN and nn.LSTM signatures in PyTorch, then explain why LSTM returns two state tensors instead of one.

Discussion prompt: Ask which kinds of application requirements justify the extra complexity of LSTM over a basic RNN.

Recap: LSTM adds gated memory and a cell state to help recurrent models preserve important context over longer time spans. Continue with GRU.