Lesson 9

LSTM and GRU gates: controlling memory

Why add gates?

A plain RNN repeatedly transforms one hidden state. During backpropagation through many positions, gradient signals are repeatedly multiplied. They can shrink toward zero or grow without control. A small gradient makes it hard to learn that an early token matters much later. LSTM and GRU networks use learned gates to create more controlled paths through time.

A gate is a vector of numbers between zero and one, produced by a sigmoid function. Values near zero block information; values near one pass it. The gate is learned from the current input and previous state. It is not a human-written if statement.

LSTM and GRU vocabulary

For an LSTM, forward returns output, (h_n, c_n). The output has a hidden vector at every time step. h_n is the final hidden state and c_n is the final cell state. A GRU returns output, h_n because it does not expose a separate cell state.

Worked example: LSTM shapes

Interpret one gate first: if a forget gate is 0.8 and the old cell value is 5, that path carries 0.8 × 5 = 4 forward before new information is added. A value near 1 keeps most of the old value. A value near 0 removes most of it.

We first embed integer tokens. With batch_first=True, input vectors have shape (B,T,C). The LSTM output is (B,T,H). Both final states are (layers,B,H). This example has one layer, two sequences, four positions, and hidden size six.

import torch
from torch import nn

torch.manual_seed(9)
vocabulary_size = 5
embedding_size = 3
hidden_size = 6

embedding = nn.Embedding(vocabulary_size, embedding_size)
lstm = nn.LSTM(
    input_size=embedding_size,
    hidden_size=hidden_size,
    batch_first=True,
)
output_layer = nn.Linear(hidden_size, vocabulary_size)

token_ids = torch.tensor([
    [0, 1, 2, 3],
    [3, 2, 1, 0],
], dtype=torch.long)

vectors = embedding(token_ids)
states, (final_hidden, final_cell) = lstm(vectors)
logits = output_layer(states)

print(tuple(vectors.shape))
print(tuple(states.shape))
print(tuple(final_hidden.shape), tuple(final_cell.shape))
print(tuple(logits.shape))

# Expected:
# (2, 4, 3)
# (2, 4, 6)
# (1, 2, 6) (1, 2, 6)
# (2, 4, 5)

The cell and hidden states have the same shape but different roles and values. Do not replace one with the other when carrying state into the next chunk. To continue, call lstm(next_vectors, (final_hidden, final_cell)). When independent sequences begin, start with no supplied state so PyTorch uses zeros.

Compare a GRU

A gate value is a multiplier, not a yes-or-no switch. The common wrong answer “0.8 means keep the memory” misses that the kept amount is 4 in this example. Always multiply the gate by the value it controls.

A GRU combines some LSTM decisions into update and reset gates. It usually has fewer parameters and can train faster. It is not always worse or better; validation results on the actual task should decide.

gru = nn.GRU(embedding_size, hidden_size, batch_first=True)
gru_states, gru_hidden = gru(vectors)
print(tuple(gru_states.shape))  # Expected: (2, 4, 6)
print(tuple(gru_hidden.shape))  # Expected: (1, 2, 6)

lstm_parameters = sum(p.numel() for p in lstm.parameters())
gru_parameters = sum(p.numel() for p in gru.parameters())
print(gru_parameters < lstm_parameters)  # Expected: True

Both layers can produce next-token logits by sending every output state through a linear layer. Training uses shifted token targets and cross-entropy exactly as in the RNN lesson. Gates improve gradient flow but do not guarantee unlimited memory. Very long sequences remain difficult and sequential processing limits parallel speed.

Padding and real batches

Real sequences often have different lengths. Padding adds a special token so they fit a rectangular batch. If padded positions contribute to loss, the model learns meaningless padding targets. Use a padding mask or ignore_index in the loss. PyTorch also provides packed sequences, but begin with padding and careful masking because the shapes are easier to inspect.

Try it

A two-layer bidirectional LSTM uses batch size three, sequence length seven, and hidden size five. What are the shapes of its output, final hidden state, and final cell state?

layer = nn.LSTM(4, 5, num_layers=2,
                bidirectional=True, batch_first=True)
x = torch.zeros(3, 7, 4)
output, (hidden, cell) = layer(x)
print(tuple(output.shape))  # Expected: (3, 7, 10)
print(tuple(hidden.shape))  # Expected: (4, 3, 5)
print(tuple(cell.shape))    # Expected: (4, 3, 5)

Recap

LSTM and GRU are recurrent networks with learned gates. An LSTM carries both hidden and cell state; a GRU carries one state and normally uses fewer parameters. Batch-first affects input and output but not the ordering of final-state axes. Gating helps preserve useful information and gradients, though it does not remove every long-sequence limitation.

References: PyTorch nn.LSTM and PyTorch nn.GRU.