LSTM Networks

"Everything that arrives, I weigh: should this stay, or has it served its purpose? I am not cruel, only deliberate. The old context that no longer earns its place, I let slip toward zero, gently, one multiply at a time. What I keep, I keep almost perfectly, because the road I guard runs nearly straight and the gradient that walks it does not stumble."

An LSTM Forget Gate Deciding What to Let Go

In Section 10.1 we built the Elman recurrent network and watched it learn short-range structure well and long-range structure not at all. The diagnosis came one chapter earlier: Section 9.3 showed that backpropagation through time multiplies one hidden-to-hidden Jacobian per step, and Section 9.4 showed that this repeated product makes gradients vanish (Jacobian norms below one) or explode (norms above one), so credit for an event hundreds of steps in the past never reaches the weights that should learn from it. Clipping tames explosion, but vanishing is the harder disease, and no amount of careful initialization cures it in a plain RNN: the squashing $\tanh$ guarantees that the local derivative $1 - h^2$ is at most one and usually well below it. The LSTM, introduced by Hochreiter and Schmidhuber in 1997 and refined with the forget gate by Gers, Schmidhuber, and Cummins in 2000, is the architectural answer. We use the unified notation of Appendix A throughout: $\mathbf{x}_t$ the input, $\mathbf{h}_t$ the hidden state, and now $\mathbf{c}_t$ the cell state.

Why does adding one more vector to the recurrence change everything? Because the problem was never that the RNN lacked capacity; it was that the path connecting a distant cause to a present effect ran through a multiplicative bottleneck that destroyed gradient magnitude. The LSTM does not make that path wider; it builds a second path that is almost straight. The four competencies this section installs are these: to explain how a separate, additively updated cell state preserves gradient magnitude where the hidden-state recurrence destroys it; to write the forget, input, and output gates and the cell and hidden updates from memory; to argue why the additive update is the constant error carousel that Section 9.4 said we needed; and to implement and train an LSTM cell that solves a long-lag task on which the vanilla RNN of Section 10.1 flatlines.

1. The Long-Term-Dependency Cure: A Separate Cell State Beginner

The single idea behind the LSTM is to split memory into two vectors that play different roles. The hidden state $\mathbf{h}_t$ stays as before: it is the network's working output at each step, the thing the readout reads and the next step consumes. Alongside it the LSTM carries a cell state $\mathbf{c}_t$, an internal memory that is not directly exposed and whose job is to transport information across many steps with as little distortion as possible. The hidden state is what the network says now; the cell state is what the network remembers.

What makes the cell state special is the way it is updated. Recall from Section 9.4 that the vanilla RNN updates its only memory by $\mathbf{h}_t = \tanh(\mathbf{W}\mathbf{h}_{t-1} + \mathbf{U}\mathbf{x}_t)$, a fully multiplicative rewrite: the old state is matrix-multiplied and then squashed, so every backward step multiplies the gradient by $\operatorname{diag}(1 - \mathbf{h}_t^2)\,\mathbf{W}$, a product that shrinks geometrically. The cell state instead is updated by a near-additive rule of the form $\mathbf{c}_t = \mathbf{f}_t \odot \mathbf{c}_{t-1} + \mathbf{i}_t \odot \tilde{\mathbf{c}}_t$, where $\mathbf{f}_t$ and $\mathbf{i}_t$ are gates in $[0,1]$ and $\odot$ is the elementwise product. The crucial structural fact is that the previous cell state $\mathbf{c}_{t-1}$ is carried forward by an elementwise multiply by the forget gate and then added to the new content, with no matrix multiplication and no squashing nonlinearity between $\mathbf{c}_{t-1}$ and $\mathbf{c}_t$. When the forget gate is near one, $\mathbf{c}_t \approx \mathbf{c}_{t-1} + (\text{new stuff})$, and the cell state behaves like a conveyor belt that carries memory forward essentially untouched. This near-linear path is the "carousel" of the constant error carousel, and it is the whole reason the LSTM can remember.

The contrast with Section 10.1 could not be sharper. There, the only way to remember something for a hundred steps was to keep matrix-multiplying and squashing it a hundred times, and the signal decayed at every multiply. Here, remembering for a hundred steps requires only that the forget gate stay near one along those steps, which the network can learn to do, and the cell content then flows forward almost unchanged. The hidden state still does the squashed, expressive, short-range work; the cell state is the long-range highway laid alongside it. This division of labor between a working memory and a protected long-term memory is the conceptual heart of the architecture, and everything in the next subsection is machinery for controlling the two flows.

2. The Three Gates and the Cell Update Intermediate

An LSTM cell at step $t$ takes the previous hidden state $\mathbf{h}_{t-1}$, the previous cell state $\mathbf{c}_{t-1}$, and the current input $\mathbf{x}_t$, and produces the new $\mathbf{h}_t$ and $\mathbf{c}_t$. It does so through three gates and one candidate, each a small learned function of the concatenated $[\mathbf{h}_{t-1}, \mathbf{x}_t]$. A gate is just a vector in $[0,1]$ produced by a sigmoid $\sigma$, used as an elementwise multiplier: a coordinate near one lets information through, a coordinate near zero blocks it. Writing each gate with its own weight matrix and bias, the forget, input, and output gates are

The candidate cell content, the new information the cell might write, is produced by a $\tanh$ rather than a sigmoid, because it is a value to store rather than a gate to multiply, so it ranges over $[-1, 1]$:

The cell state is then updated by the near-additive rule that is the whole point of the architecture: forget part of the old cell, add a gated part of the candidate.

Finally the hidden state is the output gate applied to a squashed view of the updated cell, so the cell state is exposed to the rest of the network only through a controllable, bounded window:

These six equations are the complete LSTM cell. Read them as a small story repeated every step: the forget gate $\mathbf{f}_t$ looks at the current situation and decides which coordinates of long-term memory are now obsolete and should decay toward zero; the input gate $\mathbf{i}_t$ decides which coordinates of the freshly computed candidate $\tilde{\mathbf{c}}_t$ are worth writing; the cell update commits both decisions with a multiply-and-add; and the output gate $\mathbf{o}_t$ decides how much of the resulting memory to reveal as this step's hidden state. Each gate is a full learned layer with its own parameters, so an LSTM cell has four times the parameters of an Elman cell of the same width (three gates plus one candidate), which is the price of the control. Table 10.2.1 names each gate's role and activation.

One implementation detail matters for both speed and the code in subsection five. The four pre-activations share the same inputs $[\mathbf{h}_{t-1}, \mathbf{x}_t]$, so in practice the four weight matrices are stacked into one big matrix and the four gates are computed in a single matrix multiply, then sliced apart. This is exactly what nn.LSTM does internally, and it is why a library LSTM is far faster than a naive Python loop over four separate multiplies. We write the stacked form in Code 10.2.1 so the from-scratch cell matches the library's arithmetic.

3. Why the Additive Update Keeps Gradients Alive Advanced

We can now make precise the claim that the cell state preserves gradients where the hidden state destroys them. The object that decided everything in Section 9.3 was the hidden-to-hidden Jacobian, whose repeated product across steps either vanished or exploded. For the LSTM the analogous object is the cell-to-cell Jacobian $\partial \mathbf{c}_t / \partial \mathbf{c}_{t-1}$, and its structure is qualitatively different. Differentiating the cell update $\mathbf{c}_t = \mathbf{f}_t \odot \mathbf{c}_{t-1} + \mathbf{i}_t \odot \tilde{\mathbf{c}}_t$ with respect to $\mathbf{c}_{t-1}$, and noting that $\mathbf{f}_t$, $\mathbf{i}_t$, and $\tilde{\mathbf{c}}_t$ depend on $\mathbf{c}_{t-1}$ only indirectly through $\mathbf{h}_{t-1}$ (not at all in the standard, non-peephole cell), the dominant, direct term is simply

Compare this with the vanilla RNN's $\operatorname{diag}(1 - \mathbf{h}_t^2)\,\mathbf{W}$. The RNN Jacobian is a saturating activation derivative (at most one, usually less) times a fixed weight matrix whose spectrum the network cannot freely choose; the product of many such matrices is forced toward zero or infinity. The LSTM Jacobian, by contrast, is just the forget gate, a vector the network chooses at every step and can drive close to one. Backpropagating the cell-state gradient from step $t$ to step $k$ multiplies the chain of these Jacobians,

which is a product of numbers in $[0,1]$ with no weight matrix and no activation derivative in sight. If the forget gates along that span are all near one, the product stays near one and the gradient arrives at step $k$ with its magnitude essentially intact across an arbitrary number of steps. This is the constant error carousel: an almost-identity recurrence on the cell state along which error can circulate without decay. It is the precise structural cure for the vanishing gradient that Section 9.4 diagnosed.

It would overstate the case to call the cell path perfectly linear, and the honest qualifier matters. The forget, input, and output gates all depend on $\mathbf{h}_{t-1}$, and $\mathbf{h}_{t-1}$ depends on $\mathbf{c}_{t-1}$ through $\mathbf{h}_{t-1} = \mathbf{o}_{t-1} \odot \tanh(\mathbf{c}_{t-1})$, so there is an indirect, gated, squashed coupling from $\mathbf{c}_{t-1}$ back into the gates of step $t$. That coupling contributes additional, smaller terms to the exact Jacobian. The point is not that these terms are zero but that they are a correction on top of a dominant $\operatorname{diag}(\mathbf{f}_t)$ term, whereas in the vanilla RNN the squashed, matrix-multiplied term is the only term. The LSTM did not abolish the multiplicative recurrence; it added an additive one alongside it and let the additive one carry the long-range load. That is enough: a single near-identity highway through the graph is all gradient descent needs to assign credit across long lags.

4. Variants: Peephole and Coupled Gates Intermediate

The cell of subsection two is the standard modern LSTM, but two named variants appear often enough in code and papers to recognize. The first is the peephole LSTM of Gers and Schmidhuber (2000), which lets the gates see the cell state directly, not only through the hidden state. Each gate gains an extra term that reads $\mathbf{c}_{t-1}$ (or $\mathbf{c}_t$ for the output gate) through a diagonal weight, for example $\mathbf{f}_t = \sigma(\mathbf{W}_f \mathbf{x}_t + \mathbf{U}_f \mathbf{h}_{t-1} + \mathbf{V}_f \odot \mathbf{c}_{t-1} + \mathbf{b}_f)$ with $\mathbf{V}_f$ a learned diagonal. The peephole gives the gates precise access to the cell's magnitude, which helps on tasks that require counting or precise timing, at the cost of a few more parameters and a slightly more tangled Jacobian.

The second is the coupled input-forget gate (CIFG), which ties the input gate to the forget gate by setting $\mathbf{i}_t = \mathbf{1} - \mathbf{f}_t$, so the cell only writes new content in proportion to how much old content it forgets: $\mathbf{c}_t = \mathbf{f}_t \odot \mathbf{c}_{t-1} + (\mathbf{1} - \mathbf{f}_t) \odot \tilde{\mathbf{c}}_t$. This is a convex blend of old and new, it removes one gate's worth of parameters, and it is exactly the simplification that the Gated Recurrent Unit of Section 10.3 takes further. A widely cited empirical study (Greff and colleagues, 2017) ran thousands of LSTM variants and found that none reliably beat the standard cell, that the forget gate and the output activation are the components that matter most, and that the peephole connections and the coupling are usually safe to drop. The practical lesson: use the standard cell of subsection two unless a specific task argues otherwise, and reserve the variants for when you have a measured reason.

5. Worked Example: An LSTM That Solves a Long-Lag Task Advanced

We now prove the central claim with running code. The task is the canonical long-lag memory test: at the first step the network sees a signal bit (a $+1$ or $-1$), then a long run of $L$ zero-information noise steps, and at the final step it must output the signal bit it saw at the start. Solving it requires carrying one bit of memory across $L$ steps with no reinforcement in between, exactly the long-range dependency the vanilla RNN of Section 10.1 cannot learn once $L$ grows past a few tens of steps. Code 10.2.1 implements an LSTM cell entirely from scratch, with every gate equation of subsection two written explicitly, and trains it on this task.

import torch

torch.manual_seed(0)
L = 60                 # lag: signal at step 0, recalled at step L; vanilla RNN fails here
d_in, d_h = 1, 16      # scalar input, hidden/cell width

def make_batch(n):
    """Signal bit at step 0, then L noise steps; target is the bit. Shape (n, L+1, 1)."""
    x = torch.zeros(n, L + 1, d_in)
    bit = (torch.randint(0, 2, (n,)) * 2 - 1).float()   # +1 or -1
    x[:, 0, 0] = bit                                    # signal lives only at step 0
    x[:, 1:, 0] = 0.05 * torch.randn(n, L)              # pure noise afterward
    return x, (bit > 0).long()                          # label: was the bit +1?

class ScratchLSTM(torch.nn.Module):
    """One LSTM cell, all gates explicit; weights stacked [f,i,c-tilde,o] (note: nn.LSTM uses [i,f,g,o]; see Code 10.2.2)."""
    def __init__(self, d_in, d_h):
        super().__init__()
        self.d_h = d_h
        self.Wx = torch.nn.Linear(d_in, 4 * d_h)        # input  -> 4 stacked pre-activations
        self.Wh = torch.nn.Linear(d_h,  4 * d_h)        # hidden -> 4 stacked pre-activations
        self.readout = torch.nn.Linear(d_h, 2)
        # Forget-gate bias to +1 so the carousel starts OPEN (subsection 3 key insight).
        with torch.no_grad():
            self.Wx.bias[:d_h].fill_(1.0); self.Wh.bias[:d_h].fill_(1.0)

    def forward(self, x):
        n = x.size(0)
        h = x.new_zeros(n, self.d_h); c = x.new_zeros(n, self.d_h)
        for t in range(x.size(1)):
            z = self.Wx(x[:, t]) + self.Wh(h)           # one matmul, then slice the 4 gates
            f, i, g, o = z.chunk(4, dim=1)
            f = torch.sigmoid(f); i = torch.sigmoid(i)  # forget, input gates in [0,1]
            g = torch.tanh(g);    o = torch.sigmoid(o)  # candidate in [-1,1], output gate
            c = f * c + i * g                           # ADDITIVE cell update: the carousel
            h = o * torch.tanh(c)                       # exposed hidden state
        return self.readout(h)                          # classify from the last hidden state

model = ScratchLSTM(d_in, d_h)
opt = torch.optim.Adam(model.parameters(), lr=5e-3)
lossfn = torch.nn.CrossEntropyLoss()
for step in range(400):
    x, y = make_batch(128)
    opt.zero_grad(); loss = lossfn(model(x), y); loss.backward(); opt.step()

xte, yte = make_batch(2000)
acc = (model(xte).argmax(1) == yte).float().mean().item()
print("scratch LSTM: final loss = %.4f, test accuracy = %.3f" % (loss.item(), acc))

scratch LSTM: final loss = 0.0091, test accuracy = 0.999

The result is the whole argument in one number: an LSTM cell, differing from the Elman cell of Section 10.1 only by the gated additive cell state, learns a sixty-step dependency that the Elman cell cannot. Now the library pair. Code 10.2.2 replaces the hand-written cell with torch.nn.LSTM, which implements exactly the same six equations (with the same stacked-weight arithmetic) in optimized, batched C++.

import torch

torch.manual_seed(0)

class LibLSTM(torch.nn.Module):
    def __init__(self, d_in, d_h):
        super().__init__()
        self.lstm = torch.nn.LSTM(d_in, d_h, batch_first=True)   # the whole cell + time loop
        self.readout = torch.nn.Linear(d_h, 2)
        # Same forget-gate bias trick; nn.LSTM stacks gates as [i, f, g, o].
        with torch.no_grad():
            self.lstm.bias_ih_l0[d_h:2*d_h].fill_(1.0)
            self.lstm.bias_hh_l0[d_h:2*d_h].fill_(1.0)

    def forward(self, x):
        out, (hN, cN) = self.lstm(x)        # out: all hidden states; hN: last hidden state
        return self.readout(hN[-1])         # classify from the final hidden state

model = LibLSTM(d_in, d_h)                  # d_in, d_h, L, make_batch reused from Code 10.2.1
opt = torch.optim.Adam(model.parameters(), lr=5e-3)
lossfn = torch.nn.CrossEntropyLoss()
for step in range(400):
    x, y = make_batch(128)
    opt.zero_grad(); loss = lossfn(model(x), y); loss.backward(); opt.step()

xte, yte = make_batch(2000)
acc = (model(xte).argmax(1) == yte).float().mean().item()
print("nn.LSTM: final loss = %.4f, test accuracy = %.3f" % (loss.item(), acc))

nn.LSTM: final loss = 0.0067, test accuracy = 1.000

Read the two code blocks together: Code 10.2.1 made every gate equation of subsection two executable and showed the additive cell state solving a task the vanilla RNN cannot; Code 10.2.2 reproduced the result with the production module, demonstrating that the library is the same six equations with the bookkeeping hidden. The from-scratch version is for understanding; the library version is for use.

import torch.nn as nn

# A 2-layer LSTM sequence classifier: the whole architecture of this section in 4 lines.
lstm = nn.LSTM(input_size=1, hidden_size=64, num_layers=2,
               batch_first=True, dropout=0.2)   # stacking, dropout, fusion all internal
head = nn.Linear(64, 2)
out, (hN, cN) = lstm(x)            # x: (batch, time, 1); runs full BPTT-ready recurrence
logits = head(hN[-1])              # classify from the top layer's last hidden state