Deep Learning · LSTM and GRU

Gated recurrent units: the LSTM's cell state and three gates, the simpler GRU, and an experiment measuring how far gradients reach back through 100 time steps.

The LSTM (Hochreiter and Schmidhuber, 1997) adds a cell state $c_t$ that is updated additively:

$c_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t, \qquad h_t = o_t \odot \tanh(c_t)$

forget gate f: how much of the old memory to keep
input gate i: how much new information to write
output gate o: how much of the cell to expose

Because the cell path is mostly additions, gradients can flow across many steps. The GRU (Cho et al., 2014) merges these into an update gate and a reset gate, which means fewer parameters.

The additive path only helps if the forget gate actually stays near 1. My first experiment, a memory task over 30 to 100 steps, didn’t separate RNN, GRU and LSTM at all within my small training budget. So instead I measured the mechanism directly: how much gradient from the last output reaches each earlier input step, at initialisation.

Then I set the LSTM’s forget-gate bias high, so f ≈ 1 at the start. Jozefowicz et al. (2015) recommend initialising this bias to 1 for exactly this reason. PyTorch stores the gate biases in the order input, forget, cell, output, so the forget gate is the second block of hidden_size entries.

import torch, torch.nn as nn
torch.set_num_threads(4)

def grad_reach(m):
    x = torch.randn(16, 100, 8, requires_grad=True)       # 100 time steps
    out, _ = m(x)
    out[:, -1].sum().backward()
    g = x.grad.norm(dim=(0, 2))
    return "  ".join(f"t={t}:{g[t]:.1e}" for t in [99, 75, 50, 25, 0])

torch.manual_seed(0); print("RNN                ", grad_reach(nn.RNN(8, 64, batch_first=True)))
torch.manual_seed(0); print("GRU                ", grad_reach(nn.GRU(8, 64, batch_first=True)))
torch.manual_seed(0); print("LSTM (default init)", grad_reach(nn.LSTM(8, 64, batch_first=True)))

torch.manual_seed(0)
lstm = nn.LSTM(8, 64, batch_first=True)
with torch.no_grad():
    lstm.bias_ih_l0[64:128].fill_(3.0)    # forget-gate bias -> f = sigmoid(~3) ~ 0.95
    lstm.bias_hh_l0[64:128].fill_(0.0)
print("LSTM (forget b=3)  ", grad_reach(lstm))

Output:

RNN                 t=99:6.2e+00  t=75:4.3e-05  t=50:1.8e-10  t=25:7.2e-16  t=0:2.7e-21
GRU                 t=99:2.7e+00  t=75:7.3e-05  t=50:3.5e-09  t=25:2.1e-13  t=0:1.3e-17
LSTM (default init) t=99:1.4e+00  t=75:3.9e-05  t=50:3.3e-10  t=25:3.3e-15  t=0:4.9e-20
LSTM (forget b=3)   t=99:1.2e+00  t=75:5.9e-01  t=50:8.3e-01  t=25:1.2e+00  t=0:5.1e+00

At default initialisation, all three cells lose the gradient at a similar rate over 100 steps. The LSTM isn’t magic out of the box. With the forget gate held open, the gradient reaching step 0 stays around the same order of magnitude as at the last step. That’s the “constant error carousel” idea from the original paper, made visible. Training is what teaches the gates when to stay open.

That closes the deep learning phase. NLP comes next, where these sequence models get applied to text.

Deep Learning · LSTM and GRU

Sources