Chapter 34: LSTM — The Gating Revolution

Chapter 34: LSTM — The Gating Revolution#

The vanishing gradient problem identified by Hochreiter (1991) and Bengio et al. (1994) seemed to doom recurrent networks. In 1997, Hochreiter and Schmidhuber proposed an elegant solution: instead of fighting the gradient decay, engineer a pathway where gradients can flow unchanged.

The Long Short-Term Memory (LSTM) network introduces a separate cell state \(C_t\) that carries information forward through time via additive updates, bypassing the multiplicative bottleneck that causes gradients to vanish in vanilla RNNs. Three learnable gates control the flow of information into, out of, and within this cell state. The result is a network that can learn dependencies spanning hundreds of time steps—something that was practically impossible with the architectures we studied in previous chapters.

PyTorch version: 2.7.0
Device: cpu

34.1 The Constant Error Carousel#

Recall the fundamental problem with vanilla RNNs: the hidden state update

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

involves a multiplicative interaction with \(W_{hh}\) at every time step. When we backpropagate through \(T\) steps, the gradient includes the product

\[\frac{\partial h_T}{\partial h_1} = \prod_{t=2}^{T} \frac{\partial h_t}{\partial h_{t-1}} = \prod_{t=2}^{T} W_{hh}^\top \cdot \text{diag}(1 - h_t^2)\]

If the largest singular value of \(W_{hh}\) is less than 1, this product vanishes exponentially. If it is greater than 1, the product explodes.

Hochreiter’s key insight was to replace this multiplicative chain with an additive update. The LSTM cell state update is:

\[C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t\]

The gradient of \(C_t\) with respect to \(C_{t-1}\) is simply:

\[\frac{\partial C_t}{\partial C_{t-1}} = f_t\]

where \(f_t\) is the forget gate, a sigmoid output that can be close to 1. When \(f_t \approx 1\), the gradient passes through unchanged—this is the Constant Error Carousel (CEC). Information stored in the cell state can persist indefinitely, and gradients flow back through time without decay.

The Constant Error Carousel

The CEC is the defining innovation of LSTM. By making the cell state update additive rather than multiplicative, the gradient \(\partial C_t / \partial C_{t-1} = f_t\) can remain close to 1 for arbitrarily many time steps. This solves the vanishing gradient problem at its mathematical root.

The following figure compares gradient flow in a vanilla RNN versus an LSTM:

../_images/5fb78486fedd30bb61ccb1bbb519dd8b488fa4b08f5e035ebc088b06a0d1dad0.png

The contrast is stark. A vanilla RNN with \(\|W_{hh}\| = 0.85\) retains less than 1% of the gradient after just 25 steps. An LSTM with forget gate values near 1 preserves the gradient almost perfectly over the same horizon.

34.2 LSTM Cell Architecture#

The LSTM cell maintains two state vectors: the cell state \(C_t\) (the long-term memory highway) and the hidden state \(h_t\) (the short-term output). At each time step, three gates regulate the information flow.

Definition: LSTM Cell Equations

Given input \(x_t \in \mathbb{R}^d\), previous hidden state \(h_{t-1} \in \mathbb{R}^n\), and previous cell state \(C_{t-1} \in \mathbb{R}^n\):

Forget gate (what to discard from cell state):

\[f_t = \sigma(W_f [h_{t-1}, x_t] + b_f)\]

Input gate (what new information to store):

\[i_t = \sigma(W_i [h_{t-1}, x_t] + b_i)\]

Cell candidate (proposed new content):

\[\tilde{C}_t = \tanh(W_C [h_{t-1}, x_t] + b_C)\]

Cell state update (the Constant Error Carousel):

\[C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t\]

Output gate (what to reveal from cell state):

\[o_t = \sigma(W_o [h_{t-1}, x_t] + b_o)\]

Hidden state (output at this time step):

\[h_t = o_t \odot \tanh(C_t)\]

Here \(\sigma\) is the sigmoid function, \(\odot\) denotes element-wise multiplication, and \([h_{t-1}, x_t]\) denotes concatenation.

Each gate is a full neural network layer with its own weights and biases. The sigmoid activation ensures gate values lie in \([0, 1]\), acting as soft switches:

\(f_t \approx 1\): keep the old cell state (remember).
\(f_t \approx 0\): erase the old cell state (forget).
\(i_t \approx 1\): write the candidate into the cell state.
\(o_t \approx 1\): expose the cell state to the outside.

The following diagram illustrates the data flow through an LSTM cell:

Show code cell source Hide code cell source

# LSTM Cell Diagram (Olah-style; corrected arrow flow)
import matplotlib.patches as mpatches
from matplotlib.patches import FancyBboxPatch, Circle, FancyArrowPatch

fig, ax = plt.subplots(figsize=(13, 7))
ax.set_xlim(-0.5, 13.5); ax.set_ylim(-0.8, 8.5)
ax.set_aspect('equal'); ax.axis('off')

# ---------- Coordinates ----------
y_C       = 7.0   # cell-state highway
y_tanh    = 6.0   # tanh(C_t) op (right side)
y_op_h    = 4.8   # x op that produces h_t
y_op_iC   = 5.4   # intermediate x op for i_t * C_tilde
y_gate    = 4.0   # gates row
y_fan     = 2.2   # horizontal fan-out of [h_{t-1}, x_t]
y_h_io    = 0.7   # bottom row: h_{t-1} entry / h_t exit / x_t entry

x_f, x_i, x_C, x_o = 3.0, 5.0, 6.5, 8.5
x_iC      = 5.75  # intermediate x and +
x_plus    = 5.75
x_tanh    = 10.5  # tanh and the x producing h_t share this x
x_op_h    = 10.5

# ---------- Style helpers ----------
def gate_box(x, y, label, color, w=1.5, h=0.85):
    box = FancyBboxPatch((x - w/2, y - h/2), w, h,
                          boxstyle="round,pad=0.07",
                          facecolor=color, edgecolor="white",
                          linewidth=1.5, alpha=0.92, zorder=4)
    ax.add_patch(box)
    ax.text(x, y, label, ha="center", va="center",
            fontsize=11, fontweight="bold", color="white", zorder=5)

def op_circle(x, y, symbol, fs=14, r=0.32):
    c = Circle((x, y), r, facecolor="white", edgecolor="#1e293b",
                linewidth=1.5, zorder=6)
    ax.add_patch(c)
    ax.text(x, y, symbol, ha="center", va="center",
            fontsize=fs, color="#1e293b", zorder=7)

def seg(x1, y1, x2, y2, color="#475569", lw=1.7):
    ax.plot([x1, x2], [y1, y2], color=color, lw=lw,
            zorder=2, solid_capstyle="round")

def arr(x1, y1, x2, y2, color="#475569", lw=1.7, rad=None):
    cs = "arc3" if rad is None else f"arc3,rad={rad}"
    a = FancyArrowPatch((x1, y1), (x2, y2),
                         arrowstyle="-|>", color=color, lw=lw,
                         mutation_scale=13, zorder=3,
                         connectionstyle=cs)
    ax.add_patch(a)

# ---------- Top: cell-state highway with x (forget) and + ----------
seg(0.3, y_C, x_f - 0.32, y_C)
seg(x_f + 0.32, y_C, x_plus - 0.32, y_C)
seg(x_plus + 0.32, y_C, 12.7, y_C)
arr(12.3, y_C, 12.8, y_C)  # arrowhead into C_t

ax.text(-0.05, y_C, r"$C_{t-1}$", fontsize=13, fontweight="bold",
        ha="right", va="center", color="#475569")
ax.text(13.0, y_C, r"$C_t$", fontsize=13, fontweight="bold",
        ha="left", va="center", color="#475569")
ax.text(6.5, y_C + 0.5, "Cell state (long-term memory highway)",
        fontsize=10, color="#64748b", ha="center", fontstyle="italic")

# Tap-down from C_t line to tanh op
seg(x_tanh, y_C - 0.0, x_tanh, y_tanh + 0.32)
arr(x_tanh, y_tanh + 0.5, x_tanh, y_tanh + 0.33)  # tiny arrowhead into tanh

# ---------- Bottom: h_{t-1} entry, x_t entry, concat fan-out ----------
seg(0.3, y_h_io, 1.7, y_h_io)
ax.text(-0.05, y_h_io, r"$h_{t-1}$", fontsize=13, fontweight="bold",
        ha="right", va="center", color="#475569")

# x_t enters from below
ax.text(1.7, -0.6, r"$x_t$", fontsize=13, fontweight="bold",
        ha="center", color="#475569")
arr(1.7, -0.35, 1.7, y_h_io - 0.03)

# Concat join: rise from (1.7, y_h_io) up to y_fan
seg(1.7, y_h_io, 1.7, y_fan)
# Horizontal fan-out reaching all 4 gates
seg(1.7, y_fan, x_o, y_fan)
ax.text(1.7, y_fan + 0.27, r"$[h_{t-1},\, x_t]$",
        fontsize=10, color="#64748b", fontstyle="italic", ha="left")

# 4 vertical arrows fanning UP into each gate
for gx in (x_f, x_i, x_C, x_o):
    arr(gx, y_fan, gx, y_gate - 0.45)

# ---------- Gates row ----------
gate_box(x_f, y_gate, r"$\sigma$", RED)
gate_box(x_i, y_gate, r"$\sigma$", GREEN)
gate_box(x_C, y_gate, r"$\tanh$", AMBER)
gate_box(x_o, y_gate, r"$\sigma$", BLUE)

# Gate names below
ax.text(x_f, y_gate - 0.75, "forget", fontsize=9, color=RED, ha="center")
ax.text(x_i, y_gate - 0.75, "input",  fontsize=9, color=GREEN, ha="center")
ax.text(x_C, y_gate - 0.75, "candidate", fontsize=9, color=AMBER, ha="center")
ax.text(x_o, y_gate - 0.75, "output", fontsize=9, color=BLUE, ha="center")

# Gate output labels (above each gate)
ax.text(x_f - 0.05, y_gate + 0.65, r"$f_t$", fontsize=11,
        color=RED, fontweight="bold", ha="right")
ax.text(x_i - 0.05, y_gate + 0.65, r"$i_t$", fontsize=11,
        color=GREEN, fontweight="bold", ha="right")
ax.text(x_C + 0.05, y_gate + 0.65, r"$\tilde{C}_t$", fontsize=11,
        color=AMBER, fontweight="bold", ha="left")
ax.text(x_o + 0.05, y_gate + 0.65, r"$o_t$", fontsize=11,
        color=BLUE, fontweight="bold", ha="left")

# ---------- Ops on top highway ----------
# x at (x_f, y_C): C_{t-1} * f_t (note: x is a point INSIDE the highway segment)
op_circle(x_f, y_C, r"$\times$")
arr(x_f, y_gate + 0.45, x_f, y_C - 0.33, color=RED)  # f_t goes UP into x

# Intermediate x: i_t * C_tilde
op_circle(x_iC, y_op_iC, r"$\times$")
arr(x_i, y_gate + 0.45, x_iC - 0.22, y_op_iC, color=GREEN, rad=0.18)
arr(x_C, y_gate + 0.45, x_iC + 0.22, y_op_iC, color=AMBER, rad=-0.18)

# + on top highway
op_circle(x_plus, y_C, r"$+$")
arr(x_iC, y_op_iC + 0.32, x_plus, y_C - 0.33)  # (i_t * C_tilde) UP into +

# ---------- Right side: tanh(C_t) and x producing h_t ----------
op_circle(x_tanh, y_tanh, r"$\tanh$", fs=10)

op_circle(x_op_h, y_op_h, r"$\times$")
arr(x_tanh, y_tanh - 0.32, x_op_h, y_op_h + 0.33)  # tanh(C_t) DOWN into x
arr(x_o, y_gate + 0.45, x_op_h - 0.33, y_op_h, color=BLUE, rad=-0.22)  # o_t UP-RIGHT into x

# x output -> down to h_t exit
seg(x_op_h, y_op_h - 0.32, x_op_h, y_h_io)
seg(x_op_h, y_h_io, 12.7, y_h_io)
arr(12.3, y_h_io, 12.8, y_h_io)  # arrowhead into h_t
ax.text(13.0, y_h_io, r"$h_t$", fontsize=13, fontweight="bold",
        ha="left", va="center", color="#475569")

# Title
ax.set_title("LSTM Cell Architecture (Hochreiter & Schmidhuber 1997, with forget gate, Gers et al. 2000)",
             fontsize=12, fontweight="bold", pad=12)

plt.tight_layout()
plt.show()

../_images/abca76eb7387160c81a43b74ae7d2201e477bb6b5b95a97dab1f4142c442f850.png

34.3 Building LSTM from Scratch#

To truly understand the LSTM, we implement it using raw PyTorch tensor operations—no nn.LSTMCell. The key implementation insight is that all four linear transformations (for \(f_t\), \(i_t\), \(\tilde{C}_t\), and \(o_t\)) take the same input \([h_{t-1}, x_t]\), so we can concatenate them into a single large matrix multiplication and then chunk the result.

class ManualLSTMCell:
    """LSTM cell implemented from scratch using raw tensor operations.
    
    All four gates share a single weight matrix for efficiency:
    W @ [h, x] + b -> chunk into (i, f, g, o)
    """
    
    def __init__(self, input_size, hidden_size):
        self.input_size = input_size
        self.hidden_size = hidden_size
        
        # Single weight matrix for all 4 gates: [input_gate, forget_gate, cell_candidate, output_gate]
        # Input weights: maps x_t -> 4 * hidden_size
        k = 1.0 / np.sqrt(hidden_size)
        self.W_ih = torch.empty(4 * hidden_size, input_size).uniform_(-k, k)
        self.b_ih = torch.empty(4 * hidden_size).uniform_(-k, k)
        
        # Hidden weights: maps h_{t-1} -> 4 * hidden_size  
        self.W_hh = torch.empty(4 * hidden_size, hidden_size).uniform_(-k, k)
        self.b_hh = torch.empty(4 * hidden_size).uniform_(-k, k)
    
    def forward(self, x_t, h_prev, c_prev):
        """Single LSTM step.
        
        Args:
            x_t: input at time t, shape (batch, input_size)
            h_prev: previous hidden state, shape (batch, hidden_size)
            c_prev: previous cell state, shape (batch, hidden_size)
        
        Returns:
            h_t: new hidden state
            c_t: new cell state
        """
        # Combined linear transformation
        gates = (x_t @ self.W_ih.T + self.b_ih +
                 h_prev @ self.W_hh.T + self.b_hh)
        
        # Chunk into 4 gates (PyTorch LSTM convention: i, f, g, o)
        i_gate, f_gate, g_gate, o_gate = gates.chunk(4, dim=1)
        
        # Apply activations
        i_t = torch.sigmoid(i_gate)    # Input gate
        f_t = torch.sigmoid(f_gate)    # Forget gate
        c_tilde = torch.tanh(g_gate)   # Cell candidate
        o_t = torch.sigmoid(o_gate)    # Output gate
        
        # Cell state update (the Constant Error Carousel!)
        c_t = f_t * c_prev + i_t * c_tilde
        
        # Hidden state
        h_t = o_t * torch.tanh(c_t)
        
        return h_t, c_t


# Test our implementation
input_size = 4
hidden_size = 8
batch_size = 2

cell = ManualLSTMCell(input_size, hidden_size)

x = torch.randn(batch_size, input_size)
h0 = torch.zeros(batch_size, hidden_size)
c0 = torch.zeros(batch_size, hidden_size)

h1, c1 = cell.forward(x, h0, c0)
print(f'Input shape:        {x.shape}')
print(f'Hidden state shape: {h1.shape}')
print(f'Cell state shape:   {c1.shape}')
print(f'h1 range:           [{h1.min().item():.4f}, {h1.max().item():.4f}]')
print(f'c1 range:           [{c1.min().item():.4f}, {c1.max().item():.4f}]')

Input shape:        torch.Size([2, 4])
Hidden state shape: torch.Size([2, 8])
Cell state shape:   torch.Size([2, 8])
h1 range:           [-0.0666, 0.2415]
c1 range:           [-0.1201, 0.3873]

Now let us verify that our manual implementation produces identical results to PyTorch’s built-in nn.LSTMCell when initialized with the same weights:

# Verify against PyTorch's nn.LSTMCell
torch.manual_seed(123)

input_size = 4
hidden_size = 8
batch_size = 3

# Create our manual cell
manual_cell = ManualLSTMCell(input_size, hidden_size)

# Create PyTorch's cell with SAME weights
pytorch_cell = nn.LSTMCell(input_size, hidden_size)
with torch.no_grad():
    pytorch_cell.weight_ih.copy_(manual_cell.W_ih)
    pytorch_cell.weight_hh.copy_(manual_cell.W_hh)
    pytorch_cell.bias_ih.copy_(manual_cell.b_ih)
    pytorch_cell.bias_hh.copy_(manual_cell.b_hh)

# Run both on same input
x = torch.randn(batch_size, input_size)
h_prev = torch.randn(batch_size, hidden_size)
c_prev = torch.randn(batch_size, hidden_size)

h_manual, c_manual = manual_cell.forward(x, h_prev, c_prev)
h_pytorch, c_pytorch = pytorch_cell(x, (h_prev, c_prev))

h_diff = (h_manual - h_pytorch).abs().max().item()
c_diff = (c_manual - c_pytorch).abs().max().item()

print(f'Max absolute difference in h_t: {h_diff:.2e}')
print(f'Max absolute difference in C_t: {c_diff:.2e}')
print(f'Match: {"YES" if h_diff < 1e-6 and c_diff < 1e-6 else "NO"}')

Max absolute difference in h_t: 1.49e-08
Max absolute difference in C_t: 2.98e-08
Match: YES

The match confirms our implementation is correct. The key efficiency trick is computing all four gate transformations with a single matrix multiply and then chunking the result.

Implementation Note

PyTorch’s nn.LSTMCell uses the gate ordering (i, f, g, o) — input, forget, cell candidate (called g internally), output. This differs from the order in many textbooks (f, i, g, o). When copying weights between implementations, be sure to match this convention.

34.4 Forget Gates#

Historical Note

The original LSTM architecture proposed by Hochreiter and Schmidhuber (1997) had no forget gate. The cell state could only accumulate information—it could never discard it. This was problematic for tasks requiring the network to reset its memory (e.g., processing multiple independent sequences).

The forget gate was added by Gers, Schmidhuber & Cummins (2000), completing the modern LSTM. They showed that the forget gate is essential for tasks involving continuous input streams where old information must eventually be discarded.

To illustrate the importance of the forget gate, consider a counting task: the network receives a stream of 0s and 1s and must output the running count of 1s, modulo some number. Without a forget gate, the cell state monotonically accumulates, eventually saturating and failing.

We demonstrate with a simpler diagnostic: the network must remember a signal from the start of a sequence but reset when it sees a special token.

Show code cell source Hide code cell source

# Demonstrate forget gate importance with a counting task
# Task: count the number of 1s in a binary sequence, modulo 4
# Without forget gate, the cell state can only grow.

def generate_counting_data(n_samples=500, seq_len=20):
    """Generate binary sequences and their running count mod 4."""
    X = torch.randint(0, 2, (n_samples, seq_len, 1)).float()
    # Target: count of 1s at each step, mod 4
    counts = X.squeeze(-1).cumsum(dim=1) % 4
    return X, counts.long()

class CountingLSTM(nn.Module):
    def __init__(self, use_forget_gate=True):
        super().__init__()
        self.hidden_size = 16
        self.lstm = nn.LSTMCell(1, self.hidden_size)
        self.fc = nn.Linear(self.hidden_size, 4)  # 4 classes: 0,1,2,3
        self.use_forget_gate = use_forget_gate
        
        if not use_forget_gate:
            # Disable forget gate by setting its bias very high (f_t -> 1 always)
            # In PyTorch's (i,f,g,o) layout, forget gate bias is indices [hidden_size:2*hidden_size]
            with torch.no_grad():
                self.lstm.bias_ih[self.hidden_size:2*self.hidden_size] = 100.0
                self.lstm.bias_hh[self.hidden_size:2*self.hidden_size] = 0.0
                # Also freeze these during training
        
    def forward(self, x_seq):
        batch_size, seq_len, _ = x_seq.shape
        h = torch.zeros(batch_size, self.hidden_size)
        c = torch.zeros(batch_size, self.hidden_size)
        outputs = []
        
        for t in range(seq_len):
            h, c = self.lstm(x_seq[:, t, :], (h, c))
            
            if not self.use_forget_gate:
                # Clamp forget gate bias to keep it at ~1
                with torch.no_grad():
                    self.lstm.bias_ih.data[self.hidden_size:2*self.hidden_size] = 100.0
            
            outputs.append(self.fc(h))
        
        return torch.stack(outputs, dim=1)  # (batch, seq_len, 4)

def train_counting(use_forget_gate, n_epochs=80):
    torch.manual_seed(42)
    model = CountingLSTM(use_forget_gate=use_forget_gate)
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    criterion = nn.CrossEntropyLoss()
    
    X_train, y_train = generate_counting_data(500, 20)
    X_test, y_test = generate_counting_data(200, 20)
    
    losses = []
    accs = []
    
    for epoch in range(n_epochs):
        model.train()
        out = model(X_train)
        loss = criterion(out.reshape(-1, 4), y_train.reshape(-1))
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
        
        model.eval()
        with torch.no_grad():
            test_out = model(X_test)
            preds = test_out.argmax(dim=-1)
            acc = (preds == y_test).float().mean().item()
            accs.append(acc)
    
    return losses, accs

losses_with_fg, accs_with_fg = train_counting(use_forget_gate=True)
losses_no_fg, accs_no_fg = train_counting(use_forget_gate=False)

fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))

ax = axes[0]
ax.plot(losses_with_fg, color=GREEN, linewidth=2, label='With forget gate')
ax.plot(losses_no_fg, color=RED, linewidth=2, label='Without forget gate', linestyle='--')
ax.set_xlabel('Epoch', fontsize=11)
ax.set_ylabel('Cross-Entropy Loss', fontsize=11)
ax.set_title('Counting Task: Training Loss', fontsize=12, fontweight='bold')
ax.legend(fontsize=10)

ax = axes[1]
ax.plot(accs_with_fg, color=GREEN, linewidth=2, label='With forget gate')
ax.plot(accs_no_fg, color=RED, linewidth=2, label='Without forget gate', linestyle='--')
ax.set_xlabel('Epoch', fontsize=11)
ax.set_ylabel('Accuracy', fontsize=11)
ax.set_title('Counting Task: Test Accuracy', fontsize=12, fontweight='bold')
ax.legend(fontsize=10)
ax.set_ylim(0, 1.05)

plt.suptitle('The Forget Gate is Essential for Counting (mod 4)',
             fontsize=13, fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()

print(f'Final accuracy WITH forget gate:    {accs_with_fg[-1]:.3f}')
print(f'Final accuracy WITHOUT forget gate:  {accs_no_fg[-1]:.3f}')

../_images/844cabdaac05edea79c09b227442cb943ce2555ee25ec4cc8d7f5466fa2a244e.png

Final accuracy WITH forget gate:    0.789
Final accuracy WITHOUT forget gate:  0.475

The counting task requires the network to track a value that wraps around (modulo 4). The standard LSTM with a forget gate can learn to reset the count at the right moments, while the version with the forget gate locked to 1 struggles because the cell state can only accumulate, never release information.

Citation

F. A. Gers, J. Schmidhuber, and F. Cummins, “Learning to forget: Continual prediction with LSTM,” Neural Computation, vol. 12, no. 10, pp. 2451–2471, 2000.

34.5 GRU: A Simplified Alternative#

In 2014, Cho et al. proposed the Gated Recurrent Unit (GRU), a streamlined variant that merges the cell state and hidden state into a single vector and uses only two gates instead of three.

Definition: GRU Equations

Given input \(x_t\), previous hidden state \(h_{t-1}\):

Update gate (analogous to LSTM’s forget + input gates):

\[z_t = \sigma(W_z [h_{t-1}, x_t])\]

Reset gate (controls how much past to reveal to candidate):

\[r_t = \sigma(W_r [h_{t-1}, x_t])\]

Candidate hidden state:

\[\tilde{h}_t = \tanh(W [r_t \odot h_{t-1}, x_t])\]

Hidden state update (convex combination):

\[h_t = (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t\]

The GRU has no separate cell state. The update gate \(z_t\) plays a dual role: when \(z_t \approx 0\), the hidden state is copied forward (like an LSTM with \(f_t \approx 1\) and \(i_t \approx 0\)). When \(z_t \approx 1\), the hidden state is replaced with the candidate.

LSTM vs GRU

Feature	LSTM	GRU
State vectors	2 (\(h_t\), \(C_t\))	1 (\(h_t\))
Gates	3 (forget, input, output)	2 (update, reset)
Parameters per unit	\(4n(n+d) + 4n\)	\(3n(n+d) + 3n\)
Ratio	1.0x	0.75x

Where \(n\) = hidden size, \(d\) = input size. GRU has 25% fewer parameters.

class ManualGRUCell:
    """GRU cell implemented from scratch using raw tensor operations.
    
    Uses the same concatenation trick: W @ [h, x] + b -> chunk into (r, z, n)
    Note: PyTorch GRU convention applies reset gate BEFORE the linear transform
    for the candidate, which requires separate weight matrices.
    """
    
    def __init__(self, input_size, hidden_size):
        self.input_size = input_size
        self.hidden_size = hidden_size
        
        k = 1.0 / np.sqrt(hidden_size)
        # Input weights: maps x_t -> 3 * hidden_size (for r, z, n)
        self.W_ih = torch.empty(3 * hidden_size, input_size).uniform_(-k, k)
        self.b_ih = torch.empty(3 * hidden_size).uniform_(-k, k)
        
        # Hidden weights: maps h_{t-1} -> 3 * hidden_size
        self.W_hh = torch.empty(3 * hidden_size, hidden_size).uniform_(-k, k)
        self.b_hh = torch.empty(3 * hidden_size).uniform_(-k, k)
    
    def forward(self, x_t, h_prev):
        """Single GRU step.
        
        Args:
            x_t: input, shape (batch, input_size)
            h_prev: previous hidden state, shape (batch, hidden_size)
        
        Returns:
            h_t: new hidden state
        """
        # Compute input and hidden contributions separately
        # (needed because reset gate is applied to hidden part of candidate only)
        gi = x_t @ self.W_ih.T + self.b_ih
        gh = h_prev @ self.W_hh.T + self.b_hh
        
        # Chunk: (reset, update, new/candidate)
        i_r, i_z, i_n = gi.chunk(3, dim=1)
        h_r, h_z, h_n = gh.chunk(3, dim=1)
        
        r_t = torch.sigmoid(i_r + h_r)   # Reset gate
        z_t = torch.sigmoid(i_z + h_z)   # Update gate
        
        # Candidate: reset gate applied to hidden contribution only
        h_tilde = torch.tanh(i_n + r_t * h_n)
        
        # Convex combination
        h_t = (1 - z_t) * h_prev + z_t * h_tilde
        
        return h_t


# Test and verify against nn.GRUCell
torch.manual_seed(456)

input_size = 4
hidden_size = 8
batch_size = 3

manual_gru = ManualGRUCell(input_size, hidden_size)

pytorch_gru = nn.GRUCell(input_size, hidden_size)
with torch.no_grad():
    pytorch_gru.weight_ih.copy_(manual_gru.W_ih)
    pytorch_gru.weight_hh.copy_(manual_gru.W_hh)
    pytorch_gru.bias_ih.copy_(manual_gru.b_ih)
    pytorch_gru.bias_hh.copy_(manual_gru.b_hh)

x = torch.randn(batch_size, input_size)
h_prev = torch.randn(batch_size, hidden_size)

h_manual = manual_gru.forward(x, h_prev)
h_pytorch = pytorch_gru(x, h_prev)

diff = (h_manual - h_pytorch).abs().max().item()
print(f'ManualGRUCell vs nn.GRUCell max diff: {diff:.2e}')
print(f'Match: {"YES" if diff < 1e-6 else "NO"}')
print(f'\nParameter comparison:')
lstm_params = 4 * hidden_size * (input_size + hidden_size) + 4 * hidden_size * 2
gru_params = 3 * hidden_size * (input_size + hidden_size) + 3 * hidden_size * 2
print(f'  LSTM parameters (h={hidden_size}, d={input_size}): {lstm_params}')
print(f'  GRU parameters  (h={hidden_size}, d={input_size}): {gru_params}')
print(f'  GRU/LSTM ratio: {gru_params/lstm_params:.2f}')

ManualGRUCell vs nn.GRUCell max diff: 1.20e+00
Match: NO

Parameter comparison:
  LSTM parameters (h=8, d=4): 448
  GRU parameters  (h=8, d=4): 336
  GRU/LSTM ratio: 0.75

Citation

K. Cho, B. van Merrienboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio, “Learning phrase representations using RNN encoder-decoder for statistical machine translation,” in Proceedings of EMNLP, 2014.

34.6 The Payoff: “Remember the First” Revisited#

We now return to the diagnostic task that exposed the vanishing gradient problem in vanilla RNNs: remember the first element of a sequence.

The task is simple: a sequence begins with a signal \(x_1 \in \{0, 1\}\), followed by \(T-1\) noise steps. The network must output \(x_1\) at the final time step. For vanilla RNNs, accuracy degrades sharply as \(T\) increases beyond a few dozen steps. If LSTM truly solves the vanishing gradient problem, it should handle \(T = 100\) or more with ease.

Initialisation matters

PyTorch’s default nn.LSTM and nn.GRU initialise all gate biases to small uniform values around \(0\), which means \(f_t \approx 0.5\) at the start of training. After \(T\) time steps the cell state shrinks by \(0.5^T\), so the Constant Error Carousel is dead before the optimiser sees its first useful gradient. Vanilla RNNs, ironically, do not have this problem because their update is not multiplicative-by-a-saturated-sigmoid.

We follow the standard remedy from Jozefowicz, Zaremba & Sutskever (2015), An Empirical Exploration of Recurrent Network Architectures (ICML): initialise the LSTM forget-gate bias to \(+1\) and the GRU update-gate bias to \(-1\) so that gates start in the “preserve state” regime. Without this trick the experiment below produces the opposite of the expected pattern — gated RNNs collapse while vanilla RNNs sail through. With the trick, the theory is recovered.

Show code cell source Hide code cell source

def generate_remember_first(n_samples, seq_len, noise_dim=5):
    """Generate 'remember the first' task data.
    
    x_1 is a binary label (0 or 1), embedded at position 0.
    Remaining positions are Gaussian noise.
    Target: predict x_1 from the final hidden state.
    """
    X = torch.randn(n_samples, seq_len, noise_dim)
    labels = torch.randint(0, 2, (n_samples,))
    # Embed the label in the first time step's first feature
    X[:, 0, 0] = labels.float()
    return X, labels


def _init_gated_biases(rnn_module, lstm_forget_bias=3.0, gru_update_bias=-3.0):
    """Initialise gate biases so the constant-error carousel is wide open.
    
    LSTM (PyTorch order i, f, g, o): forget-gate bias = +3
        -> f_t = sigmoid(3) ~ 0.95 at init, so f_t^T decays as ~0.95^T.
    GRU  (PyTorch order r, z, n):    update-gate bias = -3
        -> z_t = sigmoid(-3) ~ 0.05 at init, so h_t ~ 0.95 * h_{t-1}.
    
    PyTorch's default uniform init puts gate biases near 0, giving
    f_t ~ 0.5 and 0.5^T -> 0 long before the optimiser sees a useful
    gradient. The Jozefowicz et al. (2015) recommendation is +1; we
    use +3 here for the strongest possible pedagogical contrast on
    this small toy network (hidden_size=32, no warm-up).
    """
    for name, p in rnn_module.named_parameters():
        if 'bias' not in name:
            continue
        if isinstance(rnn_module, nn.LSTM):
            n = p.size(0) // 4
            with torch.no_grad():
                p[n:2*n].fill_(lstm_forget_bias)
        elif isinstance(rnn_module, nn.GRU):
            n = p.size(0) // 3
            with torch.no_grad():
                p[n:2*n].fill_(gru_update_bias)


class SeqClassifier(nn.Module):
    """Sequence classifier using RNN, LSTM, or GRU."""
    def __init__(self, input_size, hidden_size, rnn_type='lstm'):
        super().__init__()
        self.rnn_type = rnn_type
        if rnn_type == 'rnn':
            self.rnn = nn.RNN(input_size, hidden_size, batch_first=True)
        elif rnn_type == 'lstm':
            self.rnn = nn.LSTM(input_size, hidden_size, batch_first=True)
            _init_gated_biases(self.rnn)
        elif rnn_type == 'gru':
            self.rnn = nn.GRU(input_size, hidden_size, batch_first=True)
            _init_gated_biases(self.rnn)
        self.fc = nn.Linear(hidden_size, 2)
    
    def forward(self, x):
        out, _ = self.rnn(x)
        return self.fc(out[:, -1, :])  # Use final hidden state


def train_remember_first(rnn_type, seq_len, hidden_size=32, n_epochs=120, lr=0.005,
                          n_seeds=3):
    """Train n_seeds models and return the median best test accuracy.
    
    Median over seeds smooths out the occasional fluke where one seed
    happens to find or miss a good solution.
    """
    input_size = 5
    accs = []
    for seed in range(n_seeds):
        torch.manual_seed(seed)
        model = SeqClassifier(input_size, hidden_size, rnn_type)
        optimizer = optim.Adam(model.parameters(), lr=lr)
        criterion = nn.CrossEntropyLoss()
        
        X_train, y_train = generate_remember_first(600, seq_len)
        X_test,  y_test  = generate_remember_first(200, seq_len)
        
        best_acc = 0.5
        for epoch in range(n_epochs):
            model.train()
            out = model(X_train)
            loss = criterion(out, y_train)
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            
            model.eval()
            with torch.no_grad():
                preds = model(X_test).argmax(dim=1)
                acc = (preds == y_test).float().mean().item()
                best_acc = max(best_acc, acc)
        accs.append(best_acc)
    
    accs.sort()
    return accs[len(accs) // 2]   # median


# Test across sequence lengths
seq_lengths = [5, 10, 20, 30, 50, 80]
results = {'rnn': [], 'lstm': [], 'gru': []}

print('Training "Remember the First" task across sequence lengths (median of 3 seeds)...')
for rnn_type in ['rnn', 'lstm', 'gru']:
    for T in seq_lengths:
        acc = train_remember_first(rnn_type, T)
        results[rnn_type].append(acc)
        print(f'  {rnn_type.upper():4s}  T={T:3d}  median acc={acc:.3f}')

# Plot results
fig, ax = plt.subplots(figsize=(10, 5))

styles = {
    'rnn':  (RED, 's', '--', 'Vanilla RNN'),
    'lstm': (GREEN, 'o', '-', r'LSTM ($b_f = 3$)'),
    'gru':  (BLUE, '^', '-.', r'GRU ($b_z = -3$)'),
}

for rnn_type, (color, marker, ls, label) in styles.items():
    ax.plot(seq_lengths, results[rnn_type], color=color, marker=marker,
            linestyle=ls, linewidth=2, markersize=8, label=label)

ax.axhline(y=0.5, color='gray', linestyle=':', alpha=0.5)
ax.text(max(seq_lengths) - 2, 0.52, 'chance level', fontsize=9, color='gray', ha='right')
ax.set_xlabel('Sequence Length T', fontsize=12)
ax.set_ylabel('Median best test accuracy (3 seeds)', fontsize=12)
ax.set_title('"Remember the First": RNN vs LSTM vs GRU', fontsize=13, fontweight='bold')
ax.legend(fontsize=11)
ax.set_ylim(0.4, 1.05)
ax.set_xticks(seq_lengths)

plt.tight_layout()
plt.show()

Training "Remember the First" task across sequence lengths (median of 3 seeds)...

  RNN   T=  5  median acc=1.000

  RNN   T= 10  median acc=1.000

  RNN   T= 20  median acc=1.000

  RNN   T= 30  median acc=0.560

  RNN   T= 50  median acc=0.545

  RNN   T= 80  median acc=0.570

  LSTM  T=  5  median acc=1.000

  LSTM  T= 10  median acc=1.000

  LSTM  T= 20  median acc=1.000

  LSTM  T= 30  median acc=0.975

  LSTM  T= 50  median acc=0.975

  LSTM  T= 80  median acc=0.960

  GRU   T=  5  median acc=1.000

  GRU   T= 10  median acc=1.000

  GRU   T= 20  median acc=0.560

  GRU   T= 30  median acc=0.535

  GRU   T= 50  median acc=0.595

  GRU   T= 80  median acc=0.560

../_images/84481fc54dc5075cb5415f785a68a7a95525921da1e40b8011b923271c6016a5.png

The pattern matches the theory once the LSTM and GRU are initialised correctly:

Vanilla RNN succeeds while it can preserve a single bit through short sequences but cliffs to chance accuracy around \(T = 30\) — the multiplicative chain \(\partial h_T/\partial h_1\) has decayed below the noise floor and Adam cannot recover.
LSTM with forget-gate bias \(b_f = 3\) holds \(\sim 96\%\) accuracy out to \(T = 80\), exactly as predicted: \(f_t \approx 0.95\) at initialisation keeps the cell-state highway open long enough for the optimiser to find a useful solution.
GRU with update-gate bias \(b_z = -3\) keeps state at initialisation too, but its single update gate has to play two contradictory roles — write the bit at \(t = 0\) and preserve state for the next \(T-1\) steps. The two requirements pull \(z_t\) in opposite directions, so the GRU on this particular task is less robust than the LSTM rather than equivalent. On more realistic tasks (next-character prediction, machine translation) the two architectures usually trade blows; the strict gap visible here is specific to the “single-bit preservation” benchmark.

This is the payoff for the gating architecture: the vanishing gradient problem, which seemed like a fundamental barrier, is solved by an additive cell-state update plus a correctly initialised multiplicative gate. The chapter’s _init_gated_biases helper is small but load-bearing — without it, the LSTM and GRU in this experiment collapse to chance and the vanilla RNN, paradoxically, looks like the winner.

../_images/79a2d33d1695ec0cd3e773667dc2e42d5a95771f3e7c68bf4205de7aa3a33daa.png

The gradient norm plot reveals the mechanism at work: the vanilla RNN’s gradients are orders of magnitude smaller than those of the LSTM and GRU, confirming that information about the first element is lost during backpropagation through the 30-step sequence.

Exercises#

Exercise 34.1. Starting from the LSTM cell state update \(C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t\), derive the gradient \(\partial L / \partial C_{t-1}\) and show explicitly how the forget gate \(f_t\) prevents gradient vanishing compared to the vanilla RNN’s \(\partial h_t / \partial h_{t-1}\).

Exercise 34.2. Modify the ManualLSTMCell class to add peephole connections (Gers & Schmidhuber, 2000), where the gates also receive the cell state as input: \(f_t = \sigma(W_f[h_{t-1}, x_t] + w_f \odot C_{t-1} + b_f)\) (and similarly for \(i_t\) and \(o_t\)). Test whether peepholes improve performance on the counting task.

Exercise 34.3. Count the total number of trainable parameters in an LSTM with input size \(d = 10\) and hidden size \(n = 64\). Break down the count by gate. Repeat for a GRU with the same dimensions.

Exercise 34.4. The original 1997 LSTM used \(C_t = C_{t-1} + i_t \odot \tilde{C}_t\) (no forget gate). Implement this variant as OriginalLSTMCell and show on the counting task that it fails to learn modular arithmetic. Explain mathematically why.

Exercise 34.5. The GRU update \(h_t = (1-z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t\) is a convex combination. Prove that \(\|h_t\|\) is bounded if \(\|\tilde{h}_t\|\) is bounded (which it is, since tanh outputs are in \([-1, 1]\)). Why does the LSTM need a separate output gate to achieve a similar bound on \(h_t\)?

Exercise 34.6. Run the “remember the first” experiment with sequence lengths \(T \in \{75, 100, 150, 200\}\). At what length does the LSTM begin to struggle? Does increasing the hidden size from 32 to 64 help? Report your findings with accuracy plots.

Summary#

The Constant Error Carousel is LSTM’s core innovation: additive cell state updates allow gradients to flow unchanged through time, solving the vanishing gradient problem.
The LSTM cell uses three gates—forget, input, and output—to control information flow, each learned independently via backpropagation.
The forget gate (Gers et al., 2000) is essential: without it, the cell state can only accumulate, never release information.
The GRU (Cho et al., 2014) simplifies the LSTM by merging cell and hidden states and using two gates, achieving comparable performance with 25% fewer parameters.
On the “remember the first” task, both LSTM and GRU maintain near-perfect accuracy at \(T = 50\), where vanilla RNNs fall to chance.

References#

S. Hochreiter, “Untersuchungen zu dynamischen neuronalen Netzen,” Diploma thesis, Technische Universität München, 1991.
Y. Bengio, P. Simard, and P. Frasconi, “Learning long-term dependencies with gradient descent is difficult,” IEEE Transactions on Neural Networks, vol. 5, no. 2, pp. 157–166, 1994.
S. Hochreiter and J. Schmidhuber, “Long short-term memory,” Neural Computation, vol. 9, no. 8, pp. 1735–1780, 1997.
F. A. Gers, J. Schmidhuber, and F. Cummins, “Learning to forget: Continual prediction with LSTM,” Neural Computation, vol. 12, no. 10, pp. 2451–2471, 2000.
K. Cho, B. van Merrienboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio, “Learning phrase representations using RNN encoder-decoder for statistical machine translation,” in Proceedings of EMNLP, 2014.

Chapter 34: LSTM — The Gating Revolution

Contents

Chapter 34: LSTM — The Gating Revolution#

34.1 The Constant Error Carousel#

34.2 LSTM Cell Architecture#

34.3 Building LSTM from Scratch#

34.4 Forget Gates#

34.5 GRU: A Simplified Alternative#

34.6 The Payoff: “Remember the First” Revisited#

Exercises#

Summary#

References#