Lesson 8

Recurrent neural networks and backpropagation through time

Why recurrence helps

A fixed-window MLP sees only a chosen number of previous tokens. A recurrent neural network, or RNN, processes an entire sequence in order. At each position it combines the current token features with a hidden state passed from the previous position. This hidden state is a compact, changing summary of what the network has read.

For the sequence h i space, the RNN reads h and creates a state. It reads i together with that state and creates a new state. The same weights are reused at every position. Reusing weights allows the same model to process different sequence lengths, although a batch still needs compatible lengths or padding.

Terms and shapes

With batch_first=True, embedded input has shape (B, T, C): batch, time, embedding features. RNN output has shape (B, T, H). For one layer, the final hidden state has shape (1, B, H). The leading axis counts recurrent layers, not time positions.

Worked example: predict at every position

Run two steps on paper: imagine the update h_t = x_t + 0.5h_{t-1}, with h_0=0 and inputs 2 then 4. The states are 2 and 5. The second state contains the new input 4 plus one unit carried from the past. A real RNN learns the matrices that perform this update.

This module embeds token IDs, runs an RNN, and converts every hidden state to vocabulary logits. The class only organizes layers; the main idea is the sequence of tensor shapes. Targets for next-token training would also have shape (B, T).

import torch
from torch import nn

class CharacterRNN(nn.Module):
    def __init__(self, vocabulary_size, embedding_size, hidden_size):
        super().__init__()
        self.embedding = nn.Embedding(vocabulary_size, embedding_size)
        self.rnn = nn.RNN(
            input_size=embedding_size,
            hidden_size=hidden_size,
            batch_first=True,
        )
        self.output = nn.Linear(hidden_size, vocabulary_size)

    def forward(self, token_ids, hidden=None):
        vectors = self.embedding(token_ids)
        states, final_hidden = self.rnn(vectors, hidden)
        logits = self.output(states)
        return logits, final_hidden

torch.manual_seed(8)
model = CharacterRNN(vocabulary_size=3, embedding_size=4, hidden_size=6)
x = torch.tensor([[0, 1, 2], [1, 0, 1]], dtype=torch.long)
logits, hidden = model(x)

print(tuple(logits.shape))
print(tuple(hidden.shape))

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

The output contains a prediction for each of two sequences and each of three positions. Each prediction has three vocabulary scores. The final hidden tensor keeps one six-feature summary per batch item. If we pass that hidden tensor into a later chunk, processing can continue from the earlier context.

Loss over a sequence

cross_entropy expects class scores with shape (N, V) and targets with shape (N,). We combine batch and time into N = B × T. We do not combine the vocabulary axis.

targets = torch.tensor([[1, 2, 0], [0, 1, 2]], dtype=torch.long)
loss = nn.functional.cross_entropy(
    logits.reshape(-1, 3),
    targets.reshape(-1),
)
loss.backward()
print(loss.ndim)  # Expected: 0

Backpropagation follows the reused RNN computation through time, so a later loss can affect how earlier states were formed. This is called backpropagation through time. On very long sequences, gradients can become extremely small or large. LSTM and GRU layers add gates that help control this flow.

State across chunks

Carrying state across chunks is a data decision. Carry it when chunk two truly continues chunk one. Reset it when the next sequence is unrelated. A common wrong choice carries one document's state into another and creates information that the task should not have.

Passing hidden state between chunks can preserve context, but training code usually calls hidden = hidden.detach() at chunk boundaries. Detaching keeps the numeric state while disconnecting its old computation history. Without it, the graph may grow across the entire dataset and consume increasing memory. Reset hidden state when one independent document ends and another begins.

Try it

Use two recurrent layers, batch size four, sequence length five, and hidden size seven. What are the output and final-hidden shapes when batch_first=True?

rnn = nn.RNN(3, 7, num_layers=2, batch_first=True)
x = torch.zeros(4, 5, 3)
output, hidden = rnn(x)
print(tuple(output.shape))  # Expected: (4, 5, 7)
print(tuple(hidden.shape))  # Expected: (2, 4, 7)

Recap

An RNN reuses one transition at every token and carries a hidden state forward. With batch-first layout, shapes move from (B,T,C) to (B,T,H), while final hidden state is (layers,B,H). Flatten batch and time only when calculating token-level cross-entropy. Recurrent state adds memory, but plain RNNs often struggle with long dependencies.

Reference: PyTorch nn.RNN documentation.