Backpropagation Through Time › Implementing BPTT in PyTorch
Language AI HandbookAll chapters Backpropagation Through Time › Implementing BPTT in PyTorchOne section · Sections in this chapter7 / 15 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
PyTorch's autograd system handles full BPTT automatically when you call .backward() on a loss computed over a sequence. The computational graph is built during the forward pass and traversed during the backward pass. All you need to do is ensure you are not breaking the graph with unnecessary detach() calls.
Setting Up the RNN
Let's build a simple character-level language model to demonstrate BPTT:
import numpy as np
import torch
import torch.nn as nn
# Set random seed for reproducibility
torch.manual_seed(42)
np.random.seed(42)
# Simple RNN for character-level language modeling
class SimpleRNN(nn.Module):
def __init__(self, input_size, hidden_size, output_size):
super().__init__()
self.hidden_size = hidden_size
self.rnn = nn.RNN(input_size, hidden_size, batch_first=True)
self.fc = nn.Linear(hidden_size, output_size)
def forward(self, x, hidden=None):
# x: (batch, seq_len, input_size)
output, hidden = self.rnn(x, hidden)
# output: (batch, seq_len, hidden_size)
logits = self.fc(output)
# logits: (batch, seq_len, output_size)
return logits, hiddenModel parameters: 7,707 RNN parameters: 5,952 FC parameters: 1,755
The model has a standard RNN layer followed by a linear output layer. The number of parameters depends only on the vocabulary size and hidden dimension, not the sequence length, because the weights are shared across all timesteps. This weight sharing is what gives RNNs their efficiency in parameters, but it also creates the gradient challenges we have been examining.
The nn.RNN module in PyTorch implements the vanilla RNN we have been deriving. Internally, it runs the recurrence at each timestep and accumulates the hidden states. When .backward() is called on the final loss, PyTorch automatically unrolls this computation and computes BPTT gradients through the accumulated graph.
Full BPTT with PyTorch Autograd
Full BPTT is the default behavior in PyTorch. When you run a forward pass through the entire sequence and then call .backward(), PyTorch automatically propagates gradients through all timesteps:
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
# Create synthetic sequence data: random one-hot inputs
def make_batch(batch_size, seq_len, vocab_size):
# Random input token indices
input_idx = torch.randint(0, vocab_size, (batch_size, seq_len))
# One-hot encode
x = torch.zeros(batch_size, seq_len, vocab_size)
x.scatter_(2, input_idx.unsqueeze(2), 1.0)
# Target: next-character prediction (shift by 1)
target_idx = torch.randint(0, vocab_size, (batch_size, seq_len))
return x, target_idx
x, targets = make_batch(batch_size, seq_length, vocab_size)
# Full BPTT: forward through entire sequence, then backward
optimizer.zero_grad()
hidden = None # Start with zero hidden state
logits, hidden = model(x, hidden)
# logits: (batch, seq_len, vocab_size)
# Reshape for cross-entropy: (batch * seq_len, vocab_size)
loss = criterion(logits.reshape(-1, vocab_size), targets.reshape(-1))
loss.backward() # Full BPTT: gradients flow through all seq_length steps
optimizer.step()Loss: 3.2914 Gradient norm (W_hh): 0.051048 Gradient norm (W_ih): 0.051249 Sequence length processed: 20
The gradient norms show how much signal reaches the weight matrices from backpropagating through all seq_length timesteps. For short sequences, these gradients are healthy. For very long sequences (hundreds of steps), they often collapse toward zero, illustrating the vanishing gradient problem.
Implementing Truncated BPTT with detach()
The key to truncated BPTT is the detach() operation. When you detach the hidden state before passing it to the next chunk, you pass the values forward (preserving memory across chunks) without passing the gradient history (stopping gradient flow at chunk boundaries):
def train_truncated_bptt(
model, optimizer, criterion, seq_length, chunk_size, vocab_size, batch_size
):
"""
Demonstrates truncated BPTT.
Processes a long sequence in chunks of chunk_size.
Hidden state is carried across chunks; gradients are truncated.
"""
# Create a long sequence
total_steps = seq_length
x_full, targets_full = make_batch(batch_size, total_steps, vocab_size)
total_loss = 0.0
num_chunks = 0
# Start with zero hidden state
hidden = None
for start in range(0, total_steps, chunk_size):
end = min(start + chunk_size, total_steps)
# Extract chunk
x_chunk = x_full[:, start:end, :]
t_chunk = targets_full[:, start:end]
# CRITICAL: detach hidden state before each chunk
# This passes values forward but stops gradient flow backward
if hidden is not None:
hidden = hidden.detach()
optimizer.zero_grad()
# Forward pass through this chunk only
logits, hidden = model(x_chunk, hidden)
chunk_loss = criterion(
logits.reshape(-1, vocab_size), t_chunk.reshape(-1)
)
chunk_loss.backward() # Backward only through this chunk
optimizer.step()
total_loss += chunk_loss.item()
num_chunks += 1
return total_loss / num_chunks, num_chunks
avg_loss, n_chunks = train_truncated_bptt(
model,
optimizer,
criterion,
seq_length=200,
chunk_size=35,
vocab_size=vocab_size,
batch_size=batch_size,
)Average chunk loss: 3.2992 Number of chunks processed: 6 Chunk size: 35 steps Total sequence: 200 steps Gradient truncation: every 35 steps
The detach() call is the single most important line in truncated BPTT. Without it, gradients would flow back through the entire sequence history, turning the training into full BPTT but with multiple optimizer steps mid-sequence. With it, each chunk is a self-contained backward pass.
One common mistake when implementing truncated BPTT is calling optimizer.zero_grad() before detaching the hidden state. The order matters: you should detach the hidden state from the previous chunk before zeroing gradients, not after. The code above handles this correctly, but it is easy to get wrong when refactoring.
Monitoring Gradient Flow
A useful diagnostic is to track gradient norms over training. Healthy gradients have norms that do not consistently collapse to zero or explode to large values:
def train_with_gradient_monitoring(model, seq_lengths, n_steps=100):
"""
Train model on sequences of different lengths and monitor gradient norms.
Returns gradient norm history for each sequence length.
"""
grad_histories = {}
for seq_len in seq_lengths:
model_copy = SimpleRNN(vocab_size, hidden_size, vocab_size)
opt = torch.optim.Adam(model_copy.parameters(), lr=1e-3)
grad_norms = []
for step in range(n_steps):
x, targets = make_batch(4, seq_len, vocab_size)
opt.zero_grad()
logits, _ = model_copy(x)
loss = criterion(
logits.reshape(-1, vocab_size), targets.reshape(-1)
)
loss.backward()
# Record gradient norm for recurrent weights
grad_norm = model_copy.rnn.weight_hh_l0.grad.norm().item()
grad_norms.append(grad_norm)
opt.step()
grad_histories[seq_len] = grad_norms
return grad_histories
seq_lengths_to_compare = [5, 15, 30, 60]
grad_histories = train_with_gradient_monitoring(
model, seq_lengths_to_compare, n_steps=80
)
The gradient norms reveal a clear pattern: longer sequences produce smaller gradients for the recurrent weight matrix. This is the vanishing gradient problem in action. The gradient signal must travel through more matrix multiplications to reach the earliest timesteps, and each multiplication by (combined with the tanh derivative) tends to shrink the signal. This degradation motivates the architectural solutions covered in the following chapters on LSTMs and GRUs.
Full vs. Truncated BPTT: Side-by-Side Comparison
Let's compare the gradient behavior of full BPTT and truncated BPTT on the same sequence:
def compare_bptt_methods(
seq_length, chunk_size, vocab_size, batch_size, n_trials=5
):
"""
Compare gradient norms from full BPTT vs. truncated BPTT.
"""
full_norms = []
trunc_norms = []
for _ in range(n_trials):
x, targets = make_batch(batch_size, seq_length, vocab_size)
# Full BPTT
model_full = SimpleRNN(vocab_size, hidden_size, vocab_size)
opt_full = torch.optim.Adam(model_full.parameters(), lr=1e-3)
opt_full.zero_grad()
logits, _ = model_full(x)
loss = criterion(logits.reshape(-1, vocab_size), targets.reshape(-1))
loss.backward()
full_norms.append(model_full.rnn.weight_hh_l0.grad.norm().item())
# Truncated BPTT (using only the last chunk for comparison)
model_trunc = SimpleRNN(vocab_size, hidden_size, vocab_size)
opt_trunc = torch.optim.Adam(model_trunc.parameters(), lr=1e-3)
opt_trunc.zero_grad()
# Use only the last chunk_size steps for backward
x_chunk = x[:, -chunk_size:, :]
t_chunk = targets[:, -chunk_size:]
logits_chunk, _ = model_trunc(x_chunk)
loss_chunk = criterion(
logits_chunk.reshape(-1, vocab_size), t_chunk.reshape(-1)
)
loss_chunk.backward()
trunc_norms.append(model_trunc.rnn.weight_hh_l0.grad.norm().item())
return (
np.mean(full_norms),
np.std(full_norms),
np.mean(trunc_norms),
np.std(trunc_norms),
)
seq_len = 60
chunk = 20
full_mean, full_std, trunc_mean, trunc_std = compare_bptt_methods(
seq_len, chunk, vocab_size, batch_size
)Sequence length: 60, Chunk size: 20 Full BPTT gradient norm: 0.040708 +/- 0.012397 Truncated BPTT gradient norm: 0.062081 +/- 0.011493 Ratio (trunc / full): 1.53x
Truncated BPTT tends to produce larger gradient norms than full BPTT on long sequences, precisely because it avoids the extended chain of matrix products that attenuates the signal in full BPTT. This counter-intuitive result is why truncated BPTT can sometimes train faster: the gradients are noisier (they do not account for long-range dependencies), but they are not vanishing.
The trade-off here is between gradient fidelity and gradient magnitude. Full BPTT computes the exact gradient of the loss with respect to , but on long sequences this exact gradient may be so small that learning is effectively stalled. Truncated BPTT computes an approximate gradient that ignores long-range contributions, but the resulting gradient is often large enough to update the weights. For tasks where long-range dependencies matter, this approximation can harm final model quality; for tasks dominated by local context, it often works as well or better.
Key Parameters
The key parameters for BPTT training are:
- chunk_size (truncation window): Controls how many timesteps are unrolled for each backward pass. Typical values are 20-200. Larger values capture longer dependencies but use more memory.
- hidden_size: The dimensionality of the hidden state. Larger hidden states give the model more capacity but increase both forward and backward pass costs as .
- gradient clipping threshold: As covered in the Gradient Clipping chapter, clipping by global norm (typical threshold 1.0-5.0) prevents exploding gradients from destabilizing training.