03 · Recurrent Layers & LSTMs 🔄

5 min read

Text is a sequence, and order matters. Recurrent networks read one token at a time while carrying a memory — and the LSTM is the design that made that memory actually work.

1. The 'Why'

The embedding layer (sub-module 02) gives you a sequence of word vectors, but a classifier that just averages them (as our toy model did) throws away order — and order is where much of language's meaning lives. "The movie was not good" and "was the movie not good" share the same words but differ in meaning; "not good" and "good" are near-opposites that a bag-of-words model can't tell apart. To capture meaning that depends on sequence, you need an architecture that processes tokens in order and remembers what it has seen. That's the recurrent neural network (RNN): it steps through the sequence one token at a time, maintaining a hidden state that acts as a running memory, updated at each step from the current input and the previous state.

The plain RNN is elegant but has a crippling flaw: over long sequences, the gradients that flow back through many time steps tend to shrink toward zero (the vanishing gradient problem), so the network effectively forgets anything more than a few steps back — it can't connect a pronoun to a noun mentioned a sentence earlier. The LSTM (Long Short-Term Memory) was invented to fix exactly this. It adds a separate cell state — a memory conveyor belt — and a set of learnable gates that decide what to remember, what to forget, and what to output at each step. This gating lets gradients flow across long spans, so LSTMs capture long-range dependencies that vanilla RNNs cannot. For years, nn.LSTM was the default for sequence tasks, and it remains a clear, powerful stepping stone to the Transformer.

2. Core Concepts

Hidden state & recurrence. An RNN processes a sequence step by step, updating a hidden state h at each token. The final hidden state summarizes the whole sequence; the per-step outputs give a representation at every position.

The vanishing gradient problem. Backpropagating through many time steps repeatedly multiplies small numbers, shrinking gradients toward zero. Plain RNNs therefore struggle to learn long-range dependencies.

LSTM gates & cell state. The LSTM adds a cell state plus input, forget, and output gates. The gates learn what information to keep or discard, letting memory (and gradients) persist over long sequences.

nn.LSTM I/O. With batch_first=True, input is (batch, seq_len, input_size). It returns output (per-step hidden states, (batch, seq_len, hidden_size)) and a (h_n, c_n) tuple (final hidden and cell states). Use output for per-token tasks, h_n for a whole-sequence summary.

Bidirectionality & stacking. bidirectional=True reads the sequence both ways (great for classification); num_layers>1 stacks LSTMs for more capacity.

3. Code in Action

A single LSTM layer

import torch
import torch.nn as nn

lstm = nn.LSTM(input_size=8, hidden_size=16, batch_first=True)  # 8-dim tokens -> 16-dim memory

x = torch.randn(2, 5, 8)          # (batch=2, seq_len=5, input_size=8) e.g. embeddings
output, (h_n, c_n) = lstm(x)
print(output.shape)               # torch.Size([2, 5, 16]) -> hidden state at EACH step
print(h_n.shape)                  # torch.Size([1, 2, 16]) -> final hidden state per sequence

An LSTM text classifier (embedding → LSTM → head)

import torch
import torch.nn as nn

class LSTMClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden, num_classes):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(embed_dim, hidden, batch_first=True)
        self.fc = nn.Linear(hidden, num_classes)

    def forward(self, ids):
        vecs = self.embed(ids)                 # (batch, seq, embed_dim)
        output, (h_n, c_n) = self.lstm(vecs)   # run the sequence through memory
        summary = h_n[-1]                      # final hidden state -> (batch, hidden)
        return self.fc(summary)                # -> class logits

model = LSTMClassifier(vocab_size=20, embed_dim=8, hidden=16, num_classes=2)
print(model(torch.tensor([[3, 4, 6, 0]])).shape)   # torch.Size([1, 2])

Bidirectional + stacked, for more power

import torch
import torch.nn as nn

bi_lstm = nn.LSTM(
    input_size=8, hidden_size=16,
    num_layers=2,            # stack two LSTM layers
    bidirectional=True,      # read the sequence forwards AND backwards
    batch_first=True,
)
out, (h_n, c_n) = bi_lstm(torch.randn(2, 5, 8))
print(out.shape)             # torch.Size([2, 5, 32]) -> hidden doubled by bidirection

4. Common Pitfalls

Pitfall 1 — Wrong input dimension order. nn.LSTM defaults to (seq_len, batch, features). If your data is batch-first (the common case after a DataLoader), you must pass batch_first=True, or the LSTM silently treats your batch axis as the time axis and learns nonsense.

Pitfall 2 — Feeding padded sequences without packing. Padding tokens still get processed and can pollute the final hidden state. For correct handling, use pack_padded_sequence/pad_packed_sequence so the LSTM skips pad steps — or at least pool using the mask from sub-module 01.

Pitfall 3 — Grabbing the wrong output tensor. output (all steps) and h_n (final step) serve different purposes. Using output[:, -1] on padded data grabs a pad position's state, not the true last token. Track real lengths, or use h_n from a packed sequence, to summarize correctly.

5. Further Reading & Watch List

⬅️ Prev: 02 · Embeddings · ➡️ Next: 04 · Intro to Transformers