LSTMs and GRUs

In the previous chapter, RNNs gave neural networks memory. A hidden state carried information forward from one time step to the next.

But that memory was fragile. As sequences became longer, earlier information could gradually disappear.

Can we build a recurrent network that learns what to remember and what to forget?

That is the purpose of LSTMs and GRUs. They are gated recurrent networks: instead of blindly overwriting memory at every step, they learn how information should flow.

The Pain Of A Vanilla RNN

Consider this sentence:

The word not changes the meaning, but many time steps pass before the final object appears. A vanilla RNN has to preserve that information through repeated hidden-state updates.

Every step rewrites the hidden state. Important information can be diluted, overwritten, or lost.

A vanilla RNN is like having one whiteboard. At every time step, you erase it and write a new summary. After many updates, old details are unlikely to survive exactly.

LSTMs add a better memory system. Instead of one constantly overwritten hidden state, an LSTM has a long-term cell state plus gates that control memory.

Hidden State Versus Cell State

The cell state is often described as a conveyor belt. Information can flow along it while gates decide what to erase, write, or reveal.

The Three LSTM Questions

At each time step, the LSTM asks:

• What old information should I forget?
• What new information should I store?
• What part of memory should I output right now?

The gates are not decorative. They are answers to the memory-management problem.

Vanishing Gradients Through Time

During backpropagation, gradients in an RNN must travel backward through many repeated time steps. If each step multiplies the gradient by a number less than 1, the signal shrinks quickly:

Once the gradient becomes tiny, the network struggles to learn that an early word mattered for a much later prediction.

Forget Gate

The forget gate decides how much old memory to keep:

Since the sigmoid outputs values between 0 and 1, the gate acts like a soft valve. A value near 0 erases memory. A value near 1 keeps it.

Input Gate

The input gate decides how much new information to write:

The candidate memory proposes what could be added:

Cell State Update

The cell state update is the most important equation in the chapter:

In words:

new memory = old memory we kept + new memory we decided to write

This additive path is why LSTMs preserve information and gradients better than vanilla RNNs. Information does not have to be completely rewritten at every step.

Output Gate

The output gate decides what part of long-term memory becomes visible as the hidden state:

The LSTM may remember many things in the cell state, but only some of them need to be exposed for the current prediction.

Concrete Memory Example

Suppose the memory contains “Jessie lives in Toronto.” Later the sequence says “Jessie moved to New York.” The forget gate should weaken Toronto, the input gate should write New York, and the output gate should reveal whichever location matters for the current prediction.

What Is A GRU?

A gated recurrent unit, or GRU, is a simpler gated recurrent network. It has no separate cell state. The hidden state acts as both working memory and longer-term memory.

GRUs use fewer parameters and often train faster, while performing similarly to LSTMs on many tasks.

LSTM Versus GRU

Historical Context

LSTMs dominated sequence modeling for years. They powered speech recognition, handwriting recognition, language modeling, and machine translation. Transformers later replaced them for many language tasks, but LSTMs and GRUs remain important for understanding the memory problem that attention was designed to escape.

PyTorch provides LSTM and GRU modules. The model structure is similar to the RNN classifier from the previous chapter, but the recurrent layer has a stronger memory mechanism.

import torch
from torch import nn

class LSTMClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            batch_first=True,
        )
        self.classifier = nn.Linear(hidden_dim, num_classes)

    def forward(self, token_ids):
        x = self.embedding(token_ids)

        outputs, (final_hidden, final_cell) = self.lstm(x)
        last_hidden = final_hidden[-1]

        return self.classifier(last_hidden)

model = LSTMClassifier(
    vocab_size=10000,
    embed_dim=64,
    hidden_dim=128,
    num_classes=2,
)

batch = torch.randint(0, 10000, (32, 40))
logits = model(batch)
print(logits.shape)  # (32, 2)

Switching To A GRU

Replacing the LSTM with a GRU is often a one-line architectural change:

self.gru = nn.GRU(
    input_size=embed_dim,
    hidden_size=hidden_dim,
    batch_first=True,
)

outputs, final_hidden = self.gru(x)
last_hidden = final_hidden[-1]

A useful experiment is to train the RNN, LSTM, and GRU on the same sequence task and compare validation accuracy, parameter count, and training speed.