Long Short-Term Memory: Beating the Vanishing Gradient

In the recurrent network post, a vanilla RNN learned the cycle a → b → c → a by carrying a hidden state hth_t forward in time. It worked, but it ended on a wall. To send information from step 11 to step 2020, the gradient has to travel back through twenty multiplications by the same recurrent matrix WhhW_{hh}. Multiply a number by itself twenty times and it either explodes or collapses to nothing. In practice it collapses: the gradient vanishes, and the network cannot learn anything that depends on the distant past. Short-range patterns it can manage. Long-range memory is hopeless.

This post is the fix. We build a Long Short-Term Memory cell from scratch in NumPy and point it at a task the vanilla RNN fails: remember the very first character and recall it after a long run of filler. Same NumPy, same one-hot vocabulary, same training loop. Only the architecture changes.

The idea: a conveyor belt for memory

The vanishing gradient comes from one place: the hidden state is rewritten at every step. ht=tanh⁡(xtWxh+ht−1Whh+bh)h_t = \tanh(x_t W_{xh} + h_{t-1} W_{hh} + b_h) mashes the past and the present together through a matrix and a squashing tanh⁡\tanh, over and over. Nothing survives that mash for long.

The LSTM adds a second piece of memory that is not rewritten this way: a cell state CtC_t. Picture it as a conveyor belt running straight across the top of the cell (the line across the top of the hero diagram). Information is placed on the belt, rides along mostly untouched, and is read off later. The only edits allowed are gentle ones: scale what is already there, and add something new.

That scale-and-add is the whole idea. The update is

Ct=ft⊙Ct−1+it⊙C~tC_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t

where ⊙\odot is the elementwise product. Compared to the RNN’s tanh⁡(…Whh… )\tanh(\dots W_{hh} \dots), this update is almost additive. When the belt should just remember, ft≈1f_t \approx 1 and it≈0i_t \approx 0, and it becomes Ct≈Ct−1C_t \approx C_{t-1}: the value is copied, not transformed. Copying a value has gradient 11, and multiplying a gradient by 11 twenty times leaves it exactly where it started. The belt doesn’t shrink the signal, so the signal doesn’t vanish. This falls straight out of the backprop math later, and it is the whole reason the post exists.

A gate is just a neuron

What decides how much to scale and how much to add? Gates. A gate is nothing new: it is the same small sigmoid neuron we have been building since the first perceptron. It takes the current input and the previous hidden state, computes a weighted sum, and squashes it with σ into a number between 0 and 1. That number is a soft valve. 0 means block this completely, 1 means let it all through, and everything in between is a partial open.

An LSTM cell has three such gates plus one candidate:

  • The forget gate ftf_t decides how much of the old cell state to keep.
  • The input gate iti_t decides how much of the new candidate to write.
  • The candidate C~t\tilde{C}_t is the new content on offer (a tanh⁡\tanh, so it can be positive or negative).
  • The output gate oto_t decides how much of the cell state to expose as the hidden state.

Every gate reads the same thing: the concatenation of the current input xtx_t and the previous hidden state ht−1h_{t-1}. Let zt=[ xt ; ht−1 ]z_t = [\,x_t \,;\, h_{t-1}\,] be that stacked vector.

The gate equations

ft=σ(ztWf+bf)forget gateit=σ(ztWi+bi)input gateC~t=tanh⁡(ztWc+bc)candidate cellot=σ(ztWo+bo)output gateCt=ft⊙Ct−1+it⊙C~tcell state (the belt)ht=ot⊙tanh⁡(Ct)hidden state (the output)\begin{aligned} f_t &= \sigma(z_t W_f + b_f) & &\text{forget gate} \\[0.4em] i_t &= \sigma(z_t W_i + b_i) & &\text{input gate} \\[0.4em] \tilde{C}_t &= \tanh(z_t W_c + b_c) & &\text{candidate cell} \\[0.4em] o_t &= \sigma(z_t W_o + b_o) & &\text{output gate} \\[0.4em] C_t &= f_t \odot C_{t-1} + i_t \odot \tilde{C}_t & &\text{cell state (the belt)} \\[0.4em] h_t &= o_t \odot \tanh(C_t) & &\text{hidden state (the output)} \end{aligned}

Read top to bottom, that is the entire cell. Four little sigmoid/tanh neurons feed two combining steps: the belt update CtC_t and the readout hth_t. The hidden state hth_t is now a filtered view of the cell state, not the memory itself. The memory lives on the belt. hth_t is only what we choose to reveal of it.

The data: recall across a gap

We need a task that is impossible without long-term memory, so the two posts are directly comparable. Here it is:

The sequence starts with a cue, either a or b. Then a long run of filler dots. Then a query ?. At the ?, output the original cue.

With a gap of 12, one example looks like this:

a . . . . . . . . . . . . ?   ->   output 'a'
b . . . . . . . . . . . . ?   ->   output 'b'

The cue appears once, at the very first step, and is needed twelve steps later. Everything in between is noise that has to be ignored. A vanilla RNN cannot bridge that gap: its gradient to step 0 has vanished by the time it reaches the ?. An LSTM can park the cue on the belt and roll it to the end.

import numpy as np
import numpy.typing as npt

vocab: list[str] = ["a", "b", ".", "?"]
stoi: dict[str, int] = {c: i for i, c in enumerate(vocab)}
vocab_size: int = len(vocab)

GAP: int = 12               # filler characters between cue and query
seq_len: int = GAP + 2      # cue + fillers + query

rng = np.random.default_rng(1)


def make_example(cue: str) -> tuple[list[int], int]:
    chars = [cue] + ["."] * GAP + ["?"]
    xs = [stoi[c] for c in chars]
    target = stoi[cue]      # recalled at the final '?' step
    return xs, target


def one_hot(idx: int) -> npt.NDArray[np.float64]:
    v = np.zeros((1, vocab_size))
    v[0, idx] = 1.0
    return v

The vocabulary is four one-hot characters, the same small clean setup as the RNN post, only with a task that reaches further back in time.

The network

Each gate maps the stacked vector zt=[xt ; ht−1]z_t = [x_t \,;\, h_{t-1}] to the hidden size, so every gate is one matrix multiply. Stacking input and previous hidden into one vector lets us write z @ W instead of two separate products.

hidden_size: int = 16
concat_size: int = vocab_size + hidden_size


def init(rows: int, cols: int) -> npt.NDArray[np.float64]:
    return rng.standard_normal((rows, cols)) * 0.1


Wf = init(concat_size, hidden_size)   # forget gate
Wi = init(concat_size, hidden_size)   # input gate
Wo = init(concat_size, hidden_size)   # output gate
Wc = init(concat_size, hidden_size)   # candidate cell
bf = np.ones((1, hidden_size))        # start the forget gate OPEN
bi = np.zeros((1, hidden_size))
bo = np.zeros((1, hidden_size))
bc = np.zeros((1, hidden_size))

Wy = init(hidden_size, vocab_size)    # hidden -> output logits
by = np.zeros((1, vocab_size))

One detail earns its place: bf starts at ones, not zeros. With a positive forget bias, ft=σ(…)f_t = \sigma(\ldots) starts near 11, so the belt remembers by default and only learns to forget where forgetting helps. Start it at zero and the freshly initialised cell erases half of itself every step, which is the very leak we are trying to avoid. That one bias is the difference between a belt that holds its cargo and one that drops it immediately.

def sigmoid(x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    return 1.0 / (1.0 + np.exp(-x))


def softmax(x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    e = np.exp(x - np.max(x))
    return e / np.sum(e)

Forward pass

The forward pass walks the sequence one character at a time, updating the belt C and the hidden state h. We cache every intermediate value because backprop through time needs all of them.

def forward(xs: list[int]) -> dict:
    h = np.zeros((1, hidden_size))
    C = np.zeros((1, hidden_size))
    cache = {
        "z": {}, "f": {}, "i": {}, "o": {}, "g": {},
        "C": {-1: C}, "h": {-1: h}, "x": {},
    }
    for t, idx in enumerate(xs):
        x = one_hot(idx)
        z = np.concatenate([x, h], axis=1)   # [x_t ; h_{t-1}]

        f = sigmoid(z @ Wf + bf)
        i = sigmoid(z @ Wi + bi)
        o = sigmoid(z @ Wo + bo)
        g = np.tanh(z @ Wc + bc)             # candidate C̃

        C = f * C + i * g                    # the belt update
        h = o * np.tanh(C)                   # the filtered readout

        cache["z"][t], cache["x"][t] = z, x
        cache["f"][t], cache["i"][t] = f, i
        cache["o"][t], cache["g"][t] = o, g
        cache["C"][t], cache["h"][t] = C, h
    return cache

g is the candidate C~t\tilde{C}_t (named g in code to keep lines short). The two lines that matter are the belt update C = f * C + i * g and the readout h = o * np.tanh(C), the equations from the table transcribed verbatim.

Backpropagation through time

Only the final ? step is supervised: that is where we ask for the cue back, apply softmax and cross-entropy, and send the gradient backward.

The forward pass built the cell state left to right; BPTT unwinds it right to left. At each step two gradients arrive: dh, flowing into the hidden state, and dC, flowing along the belt from the future. Watch the belt term, because that is where the vanishing gradient gets cured.

For the cell update Ct=ft⊙Ct−1+it⊙C~tC_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t, differentiate with respect to Ct−1C_{t-1}:

∂Ct∂Ct−1=ft\frac{\partial C_t}{\partial C_{t-1}} = f_t

The gradient handed back to the previous cell state is the incoming dC scaled by ftf_t. Elementwise, no matrix, no tanh⁡\tanh derivative. When the forget gate is open (ft≈1f_t \approx 1), the gradient passes through untouched. Chain twelve steps together and the gradient to step 0 is multiplied by f1f2⋯f12≈1f_1 f_2 \cdots f_{12} \approx 1, not by twelve copies of WhhW_{hh} squeezed through tanh⁡′\tanh'. The RNN vanished because that product shrank geometrically. The LSTM survives because along the belt the product is a string of ones. That is the wall the previous post ended on, and this one line is how the belt gets around it.

Everything else is the ordinary chain rule through the gates. The hidden state ht=ot⊙tanh⁡(Ct)h_t = o_t \odot \tanh(C_t) splits dh into a piece for the output gate and a piece that rejoins the belt:

dot=dh⊙tanh⁡(Ct)dC+=dh⊙ot⊙(1−tanh⁡2(Ct))\begin{aligned} d o_t &= dh \odot \tanh(C_t) \\[0.3em] dC &\mathrel{+}= dh \odot o_t \odot \bigl(1 - \tanh^2(C_t)\bigr) \end{aligned}

Then the belt update distributes dC to the forget gate, the input gate, and the candidate:

dft=dC⊙Ct−1,dit=dC⊙C~t,dC~t=dC⊙itdf_t = dC \odot C_{t-1}, \qquad di_t = dC \odot \tilde{C}_t, \qquad d\tilde{C}_t = dC \odot i_t

Each of those passes back through its neuron’s own nonlinearity (σ(1−σ)\sigma(1-\sigma) for the gates, 1−tanh⁡21 - \tanh^2 for the candidate), lands on the weights, and propagates into z=[x ; ht−1]z = [x \,;\, h_{t-1}], whose hidden slice becomes next iteration’s dh. And dC_next = dC * f_t carries the belt gradient one step further back, undiminished.

def loss_and_grads(
    xs: list[int], target: int
) -> tuple[float, dict, npt.NDArray[np.float64]]:
    cache = forward(xs)
    T = len(xs)
    last = T - 1

    logits = cache["h"][last] @ Wy + by     # supervise only the final step
    probs = softmax(logits)
    loss = float(-np.log(probs[0, target] + 1e-12))

    dlogits = probs.copy()
    dlogits[0, target] -= 1.0               # d(cross-entropy)/d(logits)

    grads = {name: np.zeros_like(p) for name, p in {
        "Wf": Wf, "Wi": Wi, "Wo": Wo, "Wc": Wc,
        "bf": bf, "bi": bi, "bo": bo, "bc": bc,
        "Wy": Wy, "by": by,
    }.items()}

    grads["Wy"] += cache["h"][last].T @ dlogits
    grads["by"] += dlogits

    dh_next = dlogits @ Wy.T                 # gradient into h[last]
    dC_next = np.zeros((1, hidden_size))     # gradient along the belt

    for t in reversed(range(T)):
        f, i, o, g = cache["f"][t], cache["i"][t], cache["o"][t], cache["g"][t]
        C, C_prev = cache["C"][t], cache["C"][t - 1]
        z = cache["z"][t]

        tanhC = np.tanh(C)
        do = dh_next * tanhC
        dC = dC_next + dh_next * o * (1 - tanhC**2)   # belt + readout paths

        df = dC * C_prev
        di = dC * g
        dg = dC * i

        df_raw = df * f * (1 - f)            # back through each nonlinearity
        di_raw = di * i * (1 - i)
        do_raw = do * o * (1 - o)
        dg_raw = dg * (1 - g**2)

        grads["Wf"] += z.T @ df_raw
        grads["Wi"] += z.T @ di_raw
        grads["Wo"] += z.T @ do_raw
        grads["Wc"] += z.T @ dg_raw
        grads["bf"] += df_raw
        grads["bi"] += di_raw
        grads["bo"] += do_raw
        grads["bc"] += dg_raw

        dz = (df_raw @ Wf.T + di_raw @ Wi.T
              + do_raw @ Wo.T + dg_raw @ Wc.T)
        dh_next = dz[:, vocab_size:]         # hidden slice of [x ; h_prev]
        dC_next = dC * f                      # <-- the belt, scaled only by f_t

    return loss, grads, probs

That last line, dC_next = dC * f, is the vanishing-gradient cure written out in one expression.

Trying it out

The training loop is plain SGD with gradient clipping (to tame the occasional large step), averaging the gradient over both cues each step so the network sees a balanced signal.

def clip(grads: dict, limit: float = 5.0) -> None:
    for g in grads.values():
        np.clip(g, -limit, limit, out=g)


learning_rate = 0.1
steps = 4000

params = {
    "Wf": Wf, "Wi": Wi, "Wo": Wo, "Wc": Wc,
    "bf": bf, "bi": bi, "bo": bo, "bc": bc,
    "Wy": Wy, "by": by,
}


def apply(grads: dict) -> None:
    for name, p in params.items():
        p -= learning_rate * grads[name]


for step in range(1, steps + 1):
    batch = {name: np.zeros_like(p) for name, p in params.items()}
    total = 0.0
    for cue in ["a", "b"]:
        xs, target = make_example(cue)
        loss, grads, _ = loss_and_grads(xs, target)
        for name in batch:
            batch[name] += grads[name] / 2.0
        total += loss / 2.0
    clip(batch)
    apply(batch)
    if step % 500 == 0 or step == 1:
        print(f"step {step:5d}  loss {total:.4f}")

Running it prints a story in the loss curve:

step     1  loss 1.3851
step   500  loss 0.6960
step  1000  loss 0.6943
step  1500  loss 0.6936
step  2000  loss 0.6923
step  2500  loss 0.6921
step  3000  loss 0.0112
step  3500  loss 0.0023
step  4000  loss 0.0013

For the first two thousand steps the loss sits right at 0.693, which is ln(2), the loss of a coin flip. The network is stumped. It hasn’t worked out how to get the cue across the gap, so it guesses. Then, somewhere around step 2500, something clicks: the cell learns to store the cue at step 0, hold it on the belt through the twelve fillers, and read it back at the ?. The loss falls off a cliff to near zero.

Then we check whether it actually recalls the cue:

def predict(xs: list[int]) -> int:
    cache = forward(xs)
    logits = cache["h"][len(xs) - 1] @ Wy + by
    return int(np.argmax(softmax(logits)))


correct = 0
trials = 200
for _ in range(trials):
    cue = rng.choice(["a", "b"])
    xs, target = make_example(cue)
    if predict(xs) == target:
        correct += 1

print(f"\nrecall accuracy across gap={GAP}: {correct}/{trials} = {correct/trials:.1%}")

for cue in ["a", "b"]:
    xs, target = make_example(cue)
    pred = predict(xs)
    print(f"cue '{cue}' + {GAP} fillers + '?'  ->  '{vocab[pred]}'  (want '{vocab[target]}')")
recall accuracy across gap=12: 200/200 = 100.0%
cue 'a' + 12 fillers + '?'  ->  'a'  (want 'a')
cue 'b' + 12 fillers + '?'  ->  'b'  (want 'b')

Perfect recall across a twelve-step gap.

For contrast, I ran a vanilla RNN on this identical task: the same recurrent cell from the previous post, same 16 hidden units, same 4000 steps, trained to read off the cue at the final ?. Its loss never budges from ln(2):

vanilla RNN recall accuracy gap=12: 98/200 = 49.0%

Chance. The gradient to step 0 vanishes before the RNN can learn to carry the cue, so it never does better than a coin flip. Same task, same effort. The only difference is the belt.

What each piece bought us

As we did with the three triangles, it is worth pausing on what each part earned its keep:

  • The cell state is the whole reason this works. By making memory an additive belt instead of a rewritten hidden state, it turns the recurrent gradient from a shrinking product of matrices into a string of near-ones. That is the direct answer to the vanishing-gradient wall.
  • The forget gate decides what to keep. Starting it open (bias one) means the belt holds its cargo by default and only learns to drop things deliberately.
  • The input gate and candidate decide what to write. Together they are how the cue gets onto the belt at step 0, and how the fillers get kept off it.
  • The output gate decides what to reveal. Memory and readout are separated, so the cell can hold the cue quietly through the gap while hth_t stays uncommitted, then expose it exactly at the ?.

Four small sigmoid-and-tanh neurons, each of them the same “gate is just a neuron” we have built since the first perceptron, arranged around one additive belt. That is all an LSTM is.

Where does this lead? Two directions, both posts for another day. You can stack these cells, feeding one LSTM’s hidden states into another, to learn features of features across time the way the MLP stacked layers across space. Or you can question the belt itself. An LSTM still funnels the entire past through one fixed-width cell state, a bottleneck that strains on very long sequences. The idea that replaced it, letting each step look directly at every other step with no belt at all, is attention. But we have earned the belt first.

Complete code

Here is the whole thing in one runnable file.

import numpy as np
import numpy.typing as npt

vocab: list[str] = ["a", "b", ".", "?"]
stoi: dict[str, int] = {c: i for i, c in enumerate(vocab)}
vocab_size: int = len(vocab)

GAP: int = 12
seq_len: int = GAP + 2

rng = np.random.default_rng(1)


def make_example(cue: str) -> tuple[list[int], int]:
    chars = [cue] + ["."] * GAP + ["?"]
    xs = [stoi[c] for c in chars]
    return xs, stoi[cue]


def one_hot(idx: int) -> npt.NDArray[np.float64]:
    v = np.zeros((1, vocab_size))
    v[0, idx] = 1.0
    return v


hidden_size: int = 16
concat_size: int = vocab_size + hidden_size


def init(rows: int, cols: int) -> npt.NDArray[np.float64]:
    return rng.standard_normal((rows, cols)) * 0.1


Wf = init(concat_size, hidden_size)
Wi = init(concat_size, hidden_size)
Wo = init(concat_size, hidden_size)
Wc = init(concat_size, hidden_size)
bf = np.ones((1, hidden_size))
bi = np.zeros((1, hidden_size))
bo = np.zeros((1, hidden_size))
bc = np.zeros((1, hidden_size))

Wy = init(hidden_size, vocab_size)
by = np.zeros((1, vocab_size))


def sigmoid(x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    return 1.0 / (1.0 + np.exp(-x))


def softmax(x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    e = np.exp(x - np.max(x))
    return e / np.sum(e)


def forward(xs: list[int]) -> dict:
    h = np.zeros((1, hidden_size))
    C = np.zeros((1, hidden_size))
    cache = {
        "z": {}, "f": {}, "i": {}, "o": {}, "g": {},
        "C": {-1: C}, "h": {-1: h}, "x": {},
    }
    for t, idx in enumerate(xs):
        x = one_hot(idx)
        z = np.concatenate([x, h], axis=1)

        f = sigmoid(z @ Wf + bf)
        i = sigmoid(z @ Wi + bi)
        o = sigmoid(z @ Wo + bo)
        g = np.tanh(z @ Wc + bc)

        C = f * C + i * g
        h = o * np.tanh(C)

        cache["z"][t], cache["x"][t] = z, x
        cache["f"][t], cache["i"][t] = f, i
        cache["o"][t], cache["g"][t] = o, g
        cache["C"][t], cache["h"][t] = C, h
    return cache


def loss_and_grads(
    xs: list[int], target: int
) -> tuple[float, dict, npt.NDArray[np.float64]]:
    cache = forward(xs)
    T = len(xs)
    last = T - 1

    logits = cache["h"][last] @ Wy + by
    probs = softmax(logits)
    loss = float(-np.log(probs[0, target] + 1e-12))

    dlogits = probs.copy()
    dlogits[0, target] -= 1.0

    grads = {name: np.zeros_like(p) for name, p in {
        "Wf": Wf, "Wi": Wi, "Wo": Wo, "Wc": Wc,
        "bf": bf, "bi": bi, "bo": bo, "bc": bc,
        "Wy": Wy, "by": by,
    }.items()}

    grads["Wy"] += cache["h"][last].T @ dlogits
    grads["by"] += dlogits

    dh_next = dlogits @ Wy.T
    dC_next = np.zeros((1, hidden_size))

    for t in reversed(range(T)):
        f, i, o, g = cache["f"][t], cache["i"][t], cache["o"][t], cache["g"][t]
        C, C_prev = cache["C"][t], cache["C"][t - 1]
        z = cache["z"][t]

        tanhC = np.tanh(C)
        do = dh_next * tanhC
        dC = dC_next + dh_next * o * (1 - tanhC**2)

        df = dC * C_prev
        di = dC * g
        dg = dC * i

        df_raw = df * f * (1 - f)
        di_raw = di * i * (1 - i)
        do_raw = do * o * (1 - o)
        dg_raw = dg * (1 - g**2)

        grads["Wf"] += z.T @ df_raw
        grads["Wi"] += z.T @ di_raw
        grads["Wo"] += z.T @ do_raw
        grads["Wc"] += z.T @ dg_raw
        grads["bf"] += df_raw
        grads["bi"] += di_raw
        grads["bo"] += do_raw
        grads["bc"] += dg_raw

        dz = (df_raw @ Wf.T + di_raw @ Wi.T
              + do_raw @ Wo.T + dg_raw @ Wc.T)
        dh_next = dz[:, vocab_size:]
        dC_next = dC * f

    return loss, grads, probs


def clip(grads: dict, limit: float = 5.0) -> None:
    for g in grads.values():
        np.clip(g, -limit, limit, out=g)


learning_rate = 0.1
steps = 4000

params = {
    "Wf": Wf, "Wi": Wi, "Wo": Wo, "Wc": Wc,
    "bf": bf, "bi": bi, "bo": bo, "bc": bc,
    "Wy": Wy, "by": by,
}


def apply(grads: dict) -> None:
    for name, p in params.items():
        p -= learning_rate * grads[name]


for step in range(1, steps + 1):
    batch = {name: np.zeros_like(p) for name, p in params.items()}
    total = 0.0
    for cue in ["a", "b"]:
        xs, target = make_example(cue)
        loss, grads, _ = loss_and_grads(xs, target)
        for name in batch:
            batch[name] += grads[name] / 2.0
        total += loss / 2.0
    clip(batch)
    apply(batch)
    if step % 500 == 0 or step == 1:
        print(f"step {step:5d}  loss {total:.4f}")


def predict(xs: list[int]) -> int:
    cache = forward(xs)
    logits = cache["h"][len(xs) - 1] @ Wy + by
    return int(np.argmax(softmax(logits)))


correct = 0
trials = 200
for _ in range(trials):
    cue = rng.choice(["a", "b"])
    xs, target = make_example(cue)
    if predict(xs) == target:
        correct += 1

print(f"\nrecall accuracy across gap={GAP}: {correct}/{trials} = {correct/trials:.1%}")

for cue in ["a", "b"]:
    xs, target = make_example(cue)
    pred = predict(xs)
    print(f"cue '{cue}' + {GAP} fillers + '?'  ->  '{vocab[pred]}'  (want '{vocab[target]}')")