Vanishing Gradients › Code Implementation: Observing Vanishing Gradients Empirically
Language AI HandbookAll chapters Vanishing Gradients › Code Implementation: Observing Vanishing Gradients EmpiricallyOne section · Sections in this chapter11 / 18 ↓
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
To see vanishing gradients in a real training scenario, we can train a small RNN on a synthetic long-range dependency task and examine both the gradient norms at each timestep and the model's performance.
Synthetic Long-Range Dependency Task
We create a task where the correct prediction at the final timestep depends on a token that appeared many steps earlier. The network must remember information across an extended blank sequence. This is the classic copying task used in many papers to benchmark recurrent architectures.
import torch
torch.use_deterministic_algorithms(False)
def generate_memory_task(n_samples, delay, vocab_size=8, seed=42):
"""
Generate sequences where the last token must match the first token.
Input: [signal, zeros*delay, query]
Target: predict signal value at the query position.
"""
torch.manual_seed(seed)
signals = torch.randint(0, vocab_size, (n_samples,))
X = torch.zeros(n_samples, delay + 2, dtype=torch.long)
X[:, 0] = signals
X[:, -1] = vocab_size # special query token
y = signals
return X, y
delay_short = 5
delay_long = 40
vocab_size = 8
n_train = 800
n_test = 200
X_short, y_short = generate_memory_task(n_train, delay_short, vocab_size)
X_long, y_long = generate_memory_task(n_train, delay_long, vocab_size)
X_short_test, y_short_test = generate_memory_task(
n_test, delay_short, vocab_size, seed=99
)
X_long_test, y_long_test = generate_memory_task(
n_test, delay_long, vocab_size, seed=99
)The task has a deliberate structure. The first token in each sequence is the signal: one of eight possible symbols. The next delay tokens are all zeros, acting as blank filler. The final token is a special query marker. At the query position, the network must predict the symbol that appeared at the very beginning. For a short delay, this is feasible. For a long delay, the gradient cannot reach the signal position.
import torch.nn as nn
class SimpleRNN(nn.Module):
def __init__(self, vocab_size, embed_dim, hidden_size, n_classes):
super().__init__()
self.embed = nn.Embedding(vocab_size + 1, embed_dim)
self.rnn = nn.RNN(embed_dim, hidden_size, batch_first=True)
self.classifier = nn.Linear(hidden_size, n_classes)
def forward(self, x):
emb = self.embed(x)
out, _ = self.rnn(emb)
return self.classifier(out[:, -1, :])
embed_dim = 16
hidden_size = 32
n_classes = vocab_size
model_short = SimpleRNN(vocab_size, embed_dim, hidden_size, n_classes)
model_long = SimpleRNN(vocab_size, embed_dim, hidden_size, n_classes)import torch
import torch.optim as optim
def train_model(model, X_train, y_train, X_test, y_test, n_epochs=60, lr=0.01):
optimizer = optim.Adam(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()
train_losses = []
test_accs = []
for epoch in range(n_epochs):
model.train()
logits = model(X_train)
loss = criterion(logits, y_train)
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_losses.append(loss.item())
if (epoch + 1) % 5 == 0:
model.eval()
with torch.no_grad():
test_logits = model(X_test)
preds = test_logits.argmax(dim=1)
acc = (preds == y_test).float().mean().item()
test_accs.append(acc)
return train_losses, test_accs
train_loss_short, test_acc_short = train_model(
model_short, X_short, y_short, X_short_test, y_short_test
)
train_loss_long, test_acc_long = train_model(
model_long, X_long, y_long, X_long_test, y_long_test
)Task: predict which token appeared 5 steps ago Final test accuracy: 1.000 Baseline (random): 0.125 Task: predict which token appeared 40 steps ago Final test accuracy: 0.110 Baseline (random): 0.125
The training curves show a clean contrast. The short-delay model reduces its loss consistently and achieves high accuracy. The long-delay model struggles: its loss barely decreases and its accuracy hovers near the random baseline. The only difference between the two tasks is where the relevant information sits in the sequence. The vanilla RNN architecture provides no mechanism to bridge the gap.
This result is not a function of insufficient training time or a poorly chosen learning rate. We could train for 600 epochs instead of 60, and the long-delay model would still fail. The gradient reaching the signal position is so small that no amount of gradient descent can accumulate useful weight updates.
Measuring Per-Timestep Gradient Norms During Training
We can also look directly at the gradient norms assigned to each timestep by extracting gradients from the computation graph after a backward pass.
def measure_gradient_norms_training(model, X, y):
"""
Run one forward-backward pass and collect gradient norms
for the hidden state at each timestep.
"""
model.train()
model.zero_grad()
embed_out = model.embed(X)
h = torch.zeros(1, X.size(0), model.rnn.hidden_size)
hiddens = []
for t in range(X.size(1)):
xt = embed_out[:, t : t + 1, :]
_, h = model.rnn(xt, h)
h.retain_grad()
hiddens.append(h)
logits = model.classifier(hiddens[-1].squeeze(0))
loss = nn.CrossEntropyLoss()(logits, y)
loss.backward()
norms = []
for h_t in hiddens:
if h_t.grad is not None:
norms.append(h_t.grad.norm().item())
else:
norms.append(0.0)
return norms
model_for_grad_short = SimpleRNN(vocab_size, embed_dim, hidden_size, n_classes)
model_for_grad_long = SimpleRNN(vocab_size, embed_dim, hidden_size, n_classes)
grad_norms_short = measure_gradient_norms_training(
model_for_grad_short, X_short[:64], y_short[:64]
)
grad_norms_long = measure_gradient_norms_training(
model_for_grad_long, X_long[:64], y_long[:64]
)
These bar plots provide the direct empirical evidence. In the short-delay case, gradient norms are spread across the sequence and remain measurable at timestep 0 where the signal appeared. In the long-delay case, essentially all gradient signal is concentrated in the last few steps. The gradient at timestep 0, where the signal lives, is negligible. The network has no gradient to guide it toward learning the dependency.
Looking at these plots, you can almost feel the network's blindness. The bar chart for the long-delay task shows a tall column at the last timestep, a few smaller bars at the preceding steps, and then essentially nothing for the first 40 timesteps. The signal is there in the data. The network processes it during the forward pass. But the backward pass cannot reach it. The 40 steps of blank tokens between the signal and the query are not just noise: they are a 40-step gradient barrier.