LSTM & GRU

The last lesson ended on a problem: a vanilla RNN's gradient shrinks geometrically as it flows back through time, so the network can't learn dependencies more than a couple of dozen steps long. This lesson covers the fix that carried sequence modeling for two decades — the LSTM and its lighter sibling the GRU — and puts both to work in PyTorch.

Why "I have a pen" isn't enough

Consider predicting the last word of: "I grew up in France, moved away for work, lived in three other countries, and after all these years I still speak fluent ___". The answer, French, depends on a word from 25 steps earlier. A vanilla RNN squashes its entire memory through a tanh at every step — relevant or not — so by the time "French" is needed, "France" has been overwritten dozens of times. What we need is a memory that is kept by default and only changed on purpose.

The conveyor belt

The LSTM (Long Short-Term Memory, 1997) adds a second track of state: the cell state c_t, alongside the usual hidden state h_t. Think of the cell state as a conveyor belt running the length of the sequence. At each step, information rides along largely untouched; the network can remove something from the belt or place something new on it, but only through small, learned, elementwise adjustments:

c_t = f_t · c_(t-1) + i_t · g_t

Read it as: keep a fraction f_t of the old memory, and add i_t worth of new content g_t. The update is additive, not a full rewrite — and if that reminds you of ResNet's skip connection, it should. Both create a path along which gradients flow without being squashed at every step, and both were invented to cure the same disease: signals dying in deep compositions. Watch the difference numerically:

The vanilla state collapses toward zero within a dozen steps; the gated cell still holds most of its signal after fifty. Gradients flowing backward enjoy the same protection.

Three gates: erase, write, reveal

Who decides what to keep and what to add? Gates — tiny learned networks, each a sigmoid over the current input x_t and previous hidden state h_(t-1), producing values between 0 and 1 that act as soft switches:

• Forget gate f_t — what to erase. Multiplies the old cell state elementwise. A value near 1 means "keep this memory slot", near 0 means "wipe it". Seeing a new subject in a sentence might trigger forgetting the old subject's gender.
• Input gate i_t — what to write. Controls how much of the freshly proposed content g_t (a tanh layer) is added onto the belt.
• Output gate o_t — what to reveal. The cell state is private. The hidden state that other layers actually see is a filtered view: h_t = o_t · tanh(c_t). The network can carry a memory for hundreds of steps without exposing it until it's needed.

Each gate has its own weights, all learned by backprop like everything else. That's the whole trick: memory management is not hard-coded — the network learns what's worth remembering, for how long, and when to use it.

GRU: the streamlined variant

The Gated Recurrent Unit (2014) asks: do we really need three gates and two states? It keeps just two gates and folds the cell state back into h_t:

• Update gate z_t merges forget and input into one decision — whatever fraction of memory you erase is exactly replaced by new content: h_t = (1 − z_t) · h_(t-1) + z_t · h̃_t.
• Reset gate r_t controls how much of the previous state is consulted when proposing new content.

The result has roughly 25% fewer parameters than an LSTM of the same hidden size, trains a bit faster, and performs comparably on most tasks.

Which one? Honest answer: it rarely matters much. Reasonable defaults — start with GRU for smaller datasets or when speed matters (fewer parameters, less overfitting); reach for LSTM on larger datasets and longer sequences, where its separate cell state sometimes gives it the edge. If the choice is decisive for your problem, you'll only find out by benchmarking both.

LSTM and GRU in PyTorch

Both are drop-in replacements for nn.RNN — same constructor, same batch_first, same output shapes. The one difference: nn.LSTM's second return value is a tuple (h_n, c_n) because of the extra cell state. Here is the sine next-step task from last lesson, with all three cells racing under identical conditions (Colab):

import torch
from torch import nn, optim
 
torch.manual_seed(0)
 
# data: sine wave -> (window, next value) pairs
t = torch.linspace(0, 60, 600)
series = torch.sin(t)
seq_len = 20
X = torch.stack([series[i:i + seq_len] for i in range(len(series) - seq_len)])
X, y = X.unsqueeze(-1), series[seq_len:].unsqueeze(-1)
 
CELLS = {"rnn": nn.RNN, "lstm": nn.LSTM, "gru": nn.GRU}
 
class SeqModel(nn.Module):
    def __init__(self, cell, hidden_size=32):
        super().__init__()
        self.rnn = CELLS[cell](1, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, 1)
 
    def forward(self, x):
        out, _ = self.rnn(x)          # for LSTM, "_" is the tuple (h_n, c_n)
        return self.fc(out[:, -1, :])
 
for cell in CELLS:
    torch.manual_seed(0)
    model = SeqModel(cell)
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.MSELoss()
    for epoch in range(200):
        loss = criterion(model(X), y)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        optimizer.zero_grad()
    print(f"{cell:>4} after 200 epochs: MSE = {loss.item():.6f}")

On this short, clean sequence all three learn — but the gated cells typically converge faster and to a lower loss, and the gap widens dramatically as seq_len grows (try 100). The exercise below makes you measure exactly that.

Gradient clipping

Gates solve vanishing gradients; exploding gradients get a blunter tool. Before each optimizer step, rescale the gradient vector if its norm exceeds a threshold:

loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
optimizer.zero_grad()

It's one line between backward() and step(), and it turns "loss went NaN at epoch 37" into a non-event. Clipping is near-universal practice when training recurrent networks.

Do RNNs still matter?

Honesty time: since ~2018, transformers ate NLP. Attention sees every position at once, trains in parallel instead of step by step, and scales to billions of parameters — no recurrent model competes on language benchmarks today. But recurrence is far from dead. RNNs process a stream with constant memory per step and no need to store a growing context window, which keeps them relevant for streaming and low-latency inference, wake-word detection and other tiny on-device models, and plenty of time-series work where a 50K-parameter GRU beats an over-parameterized transformer on 3,000 data points. The ideas you just learned — gating, additive state, learned forgetting — also live on inside modern state-space models. Learn the concepts; they keep resurfacing.

Next: putting sequence models to work — bidirectional and stacked RNNs, windowing and scaling real series, and a full forecasting pipeline.