Deep Learning · RNN

Text, audio and time series are sequences. An RNN reads one step at a time and keeps a hidden state:

$h_t = \tanh(W_{xh}x_t + W_{hh}h_{t-1} + b)$

The same weights are used at every step. Unrolled over time, it’s a very deep network with shared weights, trained by backpropagation through time.

The weakness is that the gradient for step 1 passes through many $W_{hh}$ multiplications and tanh derivatives, so it tends to vanish (or occasionally explode). I measured how much the output at the last step depends on each input step, using autograd.

import torch, torch.nn as nn

torch.manual_seed(0)
rnn = nn.RNN(input_size=4, hidden_size=16, batch_first=True)
x = torch.randn(2, 50, 4, requires_grad=True)       # batch 2, 50 time steps
out, h_n = rnn(x)
print("output", tuple(out.shape), "| final hidden", tuple(h_n.shape))

out[:, -1].sum().backward()
g = x.grad.norm(dim=(0, 2))
for t in [49, 45, 40, 30, 10, 0]:
    print(f"gradient reaching step {t:>2}: {g[t]:.2e}")

# manual cell equals nn.RNN
h = torch.zeros(1, 16)
for t in range(50):
    h = torch.tanh(x[0, t] @ rnn.weight_ih_l0.T + rnn.bias_ih_l0 + h @ rnn.weight_hh_l0.T + rnn.bias_hh_l0)
print("manual loop matches nn.RNN:", torch.allclose(h, h_n[0, 0:1], atol=1e-5))

Output:

output (2, 50, 16) | final hidden (1, 2, 16)
gradient reaching step 49: 1.05e+00
gradient reaching step 45: 7.57e-02
gradient reaching step 40: 4.58e-03
gradient reaching step 30: 1.42e-05
gradient reaching step 10: 5.46e-11
gradient reaching step  0: 3.23e-13
manual loop matches nn.RNN: True