Teaching a network to remember (and forget) on purpose

1 The gating idea

First, a quick recap of the problem

Recall from Session 12 that a plain RNN (Recurrent Neural Network) processes a sequence one step at a time, carrying a hidden state h_t from each step to the next. At every step it does roughly:

# plain RNN update (from Session 12)
h_t = tanh(W_h · h_{t-1} + W_x · x_t + b)

The trouble is in that little h_{t-1} term. To learn from a clue 100 steps ago, the training signal (the gradient) has to travel back through 100 of those multiplications. Each step multiplies by roughly the same weight matrix and a tanh derivative that's less than 1. Multiply 100 numbers that are each < 1 together and you get something microscopic — the vanishing gradient problem. The opposite can also happen (numbers > 1 multiplied many times blow up) — the exploding gradient problem.

The insight: stop forcing it, let it learn

A plain RNN completely rewrites its memory at every single step — it has no choice. The hidden state is squashed through a tanh and replaced. That constant rewriting is exactly what destroys the gradient.

What exactly is a gate?

A gate is a vector of numbers between 0 and 1, produced by a tiny neural layer with a sigmoid activation (σ, which always outputs 0–1). You then multiply it element-wise (the ⊙ symbol) against some information:

• A gate value near 0 = "block this" (multiply by ~0 → information disappears).
• A gate value near 1 = "let it all through" (multiply by ~1 → information passes untouched).
• Values in between = "let some fraction through" — gates are soft, not on/off switches, which is what makes them trainable by gradient descent.

c_new = f ⊙ c + i ⊙ g
      = 0.9 × 5.0 + 0.2 × (−1.0)
      = 4.5 − 0.2
      = 4.3

2 The LSTM cell anatomy

The LSTM (Long Short-Term Memory) cell was introduced by Hochreiter & Schmidhuber in 1997. Its key innovation is keeping two things flowing through time instead of one:

• The cell state c_t — the long-term memory "conveyor belt." It's edited only by gentle additions and removals.
• The hidden state h_t — the short-term, "what I output right now" state (same role as in a plain RNN).

The three gates at a glance

The data flow, as a diagram

The equations, one by one

At each timestep the LSTM first glues together the previous hidden state h_{t-1} and the current input x_t, then runs five small computations. Don't be scared — they're all the same shape: a weighted sum followed by an activation. Let's read each line like a sentence.

# σ = sigmoid (outputs 0..1, a "how much" knob)
# ⊙ = element-wise multiply

f_t = σ( W_f · [h_{t-1}, x_t] + b_f )      # 1. forget gate
i_t = σ( W_i · [h_{t-1}, x_t] + b_i )      # 2. input gate
g_t = tanh( W_g · [h_{t-1}, x_t] + b_g )   # 3. candidate memory
c_t = f_t ⊙ c_{t-1} + i_t ⊙ g_t           # 4. update the cell state
o_t = σ( W_o · [h_{t-1}, x_t] + b_o )      # 5. output gate
h_t = o_t ⊙ tanh( c_t )                    # 6. new hidden state

Line by line, in plain English:

• 1 — Forget gate f_t: looks at the new input and old hidden state and outputs a number 0–1 per memory slot. 0 means "erase this slot," 1 means "keep it."
• 2 — Input gate i_t: decides how much of the brand-new candidate we'll actually write.
• 3 — Candidate memory g_t: the new information we could add, squashed to −1..+1 by tanh so it can push memory up or down.
• 4 — The update (the heart of it): c_t = f_t ⊙ c_{t-1} + i_t ⊙ g_t. Read it as: "keep the fraction f_t of the old belt, then add the fraction i_t of the new content." This is the famous additive update — remember it for Topic 3.
• 5 — Output gate o_t: decides which parts of the (now updated) memory to expose to the outside world this step.
• 6 — Hidden state h_t: the cell state passed through tanh and filtered by the output gate. This is what leaves the cell and feeds the next layer / the prediction.

3 Why LSTMs fix gradient flow

Recall what killed the plain RNN

In a plain RNN the hidden state is transformed every step: h_t = tanh(W · h_{t-1} + ...). When we backpropagate, the gradient of h_t with respect to h_{t-1} includes the weight matrix W and a tanh' factor (always < 1). Chain those across many steps and you're multiplying many small numbers — the product shrinks toward zero. Gradient vanishes.

The additive path is the whole trick

Now look at the LSTM's cell-state update again, and ask: how does c_t depend on c_{t-1}?

c_t = f_t ⊙ c_{t-1} + i_t ⊙ g_t

# the part of the gradient that travels straight back in time:
∂c_t / ∂c_{t-1} = f_t        # just the forget gate!

This is the magic. The connection from one cell state to the next is essentially multiplication by the forget gate — and nothing else. There's no weight matrix and no tanh squashing on this path. So when the gradient flows back through time along the cell state, it gets multiplied by f_t at each step.

gradient after 50 steps ≈ 0.8^50 ≈ 0.000014   # vanished

gradient after 50 steps ≈ 0.99^50 ≈ 0.61       # very much alive

4 GRUs — the leaner gated cell

The GRU (Gated Recurrent Unit), introduced by Cho et al. in 2014, keeps the gating idea but simplifies it. The two big changes versus an LSTM:

• It has no separate cell state — the hidden state h_t carries both long- and short-term memory.
• It uses two gates instead of three: a reset gate and an update gate.

The GRU equations

# σ = sigmoid, ⊙ = element-wise multiply

r_t = σ( W_r · [h_{t-1}, x_t] + b_r )           # 1. reset gate
z_t = σ( W_z · [h_{t-1}, x_t] + b_z )           # 2. update gate
h~_t = tanh( W_h · [r_t ⊙ h_{t-1}, x_t] + b_h ) # 3. candidate state
h_t = (1 − z_t) ⊙ h_{t-1} + z_t ⊙ h~_t          # 4. blend old & new

Reading it as a story:

• Reset gate r_t: when computing the new candidate memory, how much of the past should we let influence it? Near 0 = "ignore the past, this is a fresh start" (great at sentence/clause boundaries).
• Update gate z_t: the star of the GRU. It does the job of the LSTM's forget and input gates at once. Line 4 says: keep fraction (1 − z_t) of the old state and take fraction z_t of the new candidate. They're tied together — what you don't keep, you replace.
• Candidate h~_t: the proposed new memory, computed using a reset-filtered view of the past.

LSTM vs GRU — side by side

5 Bidirectional & stacked RNNs

Bidirectional RNNs — reading both directions

A normal RNN reads left-to-right, so its understanding of word t only uses words before it. But meaning often depends on what comes after. A bidirectional RNN (BiRNN) runs two independent RNNs: one forward (start→end) and one backward (end→start). At each position it concatenates both hidden states, so every output knows about the entire sequence on both sides.

Stacked (deep) RNNs — layers on layers

A stacked RNN (or deep RNN) feeds the sequence of hidden states from one RNN layer as the input sequence to the next RNN layer. Lower layers capture local patterns (letters→words), higher layers capture abstract ones (phrases→meaning) — the same "hierarchy of features" idea you saw with CNNs in earlier sessions.

Real PyTorch: it's basically free

The wonderful part: PyTorch's nn.LSTM gives you stacking and bidirectionality with two arguments. You don't hand-write any of the gate equations from Topic 2 — they're built in.

import torch
import torch.nn as nn

# A 2-layer BIDIRECTIONAL LSTM
lstm = nn.LSTM(
    input_size=100,      # size of each input vector (e.g. word embedding)
    hidden_size=128,     # size of the hidden state per direction
    num_layers=2,        # STACKED: 2 LSTM layers
    batch_first=True,    # input shape is (batch, seq_len, features)
    bidirectional=True,  # BIDIRECTIONAL: read both ways
    dropout=0.2,        # dropout between stacked layers
)

# Fake batch: 32 sequences, each 10 steps long, 100 features per step
x = torch.randn(32, 10, 100)

# output: per-step hidden states (both directions concatenated)
# h_n, c_n: final hidden and cell states for each layer/direction
output, (h_n, c_n) = lstm(x)

print(output.shape)  # (32, 10, 256)  -> 128 * 2 directions
print(h_n.shape)     # (4, 32, 128)   -> 2 layers * 2 directions

★ Putting it all together

This session was really about one elegant idea applied five ways. Here's the story that ties it all together:

Quick self-check

📚 References & Further Reading

Class material

• Your course handout for this session: "SST Deep Learning handout (Session 13)" — the primary reference for this class. (course page)

Papers, docs & deep dives

• Hochreiter & Schmidhuber, "Long Short-Term Memory" (1997) — the original paper that introduced the LSTM and the "constant error carousel" argument behind Topic 3.
• Christopher Olah, "Understanding LSTM Networks" — the legendary, diagram-rich blog post; the best companion to our cell-anatomy section.
• Cho et al., "Learning Phrase Representations using RNN Encoder–Decoder" (2014) — the paper that introduced the GRU (Topic 4).
• PyTorch nn.LSTM documentation — official reference for the arguments and tensor shapes used in Topic 5.
• Chung et al., "Empirical Evaluation of Gated Recurrent Neural Networks" (2014) — head-to-head comparison of LSTM vs GRU, backing the "benchmark both" advice.