LSTM Architecture

Why Vanilla RNNs Forget

A vanilla RNN has exactly one knob for memory: the hidden state $h_t$. Every timestep it is completely overwritten by $h_t = \tanh(W_{hh} \, h_{t-1} + W_{xh} \, x_t + b)$. The word “completely” is the problem. To remember something from ten steps ago, that information must survive ten consecutive multiplications by $W_{hh}$ and ten applications of the tanh nonlinearity.

Mathematically, the gradient that flows back from step $T$ to step $t$ picks up a factor $\prod_{k=t+1}^{T} W_{hh}^{\top} \, \text{diag}(\tanh'(\cdot))$. If $\|W_{hh}\| < 1$ this product shrinks to zero — the vanishing gradient. If it is bigger than one, the product explodes. Either way, learning long-range dependencies is brittle.

The core problem: a vanilla RNN has no way to selectively preserve information. Every step is a full rewrite of memory through a squashing nonlinearity. There is no highway for gradients.

The Long Short-Term Memory cell, introduced by Hochreiter and Schmidhuber in 1997, is the classic fix. It adds a second state variable — the cell state $C_t$ — whose update is additive rather than multiplicative, and three learned gates that decide how that memory is read, written, and exposed. Once you understand those four pieces, the entire architecture falls into place.

The Cell State: A Memory Highway

Picture a conveyor belt running along the top of each LSTM cell. On that belt rides the cell state $C_t$, a vector the same size as the hidden state. The entire mechanism of the LSTM is about carefully editing what is on that belt:

• The forget gate $f_t$ scales $C_{t-1}$ element-wise, erasing parts of memory the cell no longer needs.
• The input gate $i_t$ and the candidate $\tilde{C}_t$ together decide what new content to write onto the belt.
• The output gate $o_t$ decides what part of the belt's current contents to expose through the hidden state $h_t$.

The whole update is captured in one short, additive line: $C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t$. That plus sign is the entire reason LSTMs work. When we backpropagate through it, the derivative $\partial C_t / \partial C_{t-1} = f_t$ has no squashing nonlinearity wrapped around it — so if the network learns to keep $f_t \approx 1$, gradients flow unchanged across hundreds of steps.

The Four Sub-Networks Inside a Cell

Every LSTM cell contains four tiny feed-forward layers, and each of them takes exactly the same input: the concatenation $z_t = [h_{t-1}, x_t]$. Three of them produce sigmoid-gated vectors in $(0, 1)$; one produces a tanh-bounded candidate in $(-1, 1)$.

The choice of activation is not arbitrary. Gates are switches, and sigmoid's range of $(0, 1)$ is exactly the vocabulary of switches — “let 90% through” is a meaningful thing to say. The candidate is content, and content needs a sign — the negative half of tanh lets a memory slot move in the opposite direction (e.g. a word like not subtracting from an accumulated sentiment).

The Forget Gate

$f_t = \sigma(W_f \, [h_{t-1}, x_t] + b_f)$. For each dimension of the cell state, $f_t$ returns a number between 0 and 1 that will multiply $C_{t-1}$:

• $f_t[k] \approx 1$ — channel $k$ of memory passes through essentially unchanged.
• $f_t[k] \approx 0$ — channel $k$ is erased.
• Intermediate values fade memory gradually.

A classic intuition: suppose one channel of $C$ encodes the subject's gender (for subject-verb agreement). When a new sentence begins, the forget gate on that channel should go to zero — “flush the gender, a new subject is coming.” The network learns to do this from data.

Input Gate and Candidate Memory

Writing to memory is split into how much and what — the input gate and the candidate.

• Input gate $i_t = \sigma(W_i \, [h_{t-1}, x_t] + b_i)$. A sigmoid-valued vector. Close to 1 means “yes, write heavily into this channel this step”; close to 0 means “skip this channel.”
• Candidate $\tilde{C}_t = \tanh(W_C \, [h_{t-1}, x_t] + b_C)$. A tanh-bounded vector of proposed new content. Unlike the gates, this one is signed.

The two combine as $i_t \odot \tilde{C}_t$. Think of it as (how loud) × (what to say). Splitting these two concerns is crucial: a single tanh-of-linear cannot express “I have strong content but I don't want to write it hard right now.” With the gate, the network can.

Updating the Cell State

Everything so far has been preparation for one line — the line that defines LSTM:

That's it. Two element-wise products and an element-wise sum. The first term erases; the second term writes. There is no activation wrapped around the result. That is what makes $C_t$ the highway.

Take the derivative along the highway: $\partial C_t / \partial C_{t-1} = f_t$. No tanh derivative (which can be as small as 0.0), no repeated weight matrix (which magnifies or shrinks across steps). Just a direct element-wise multiplier the network can learn to keep near 1 when it needs to remember — and the memory of the sequence survives.

Constant Error Carousel. Hochreiter and Schmidhuber's original name for this mechanism. When $f_t = 1$ and $i_t = 0$, the cell state literally just circulates: $C_t = C_{t-1}$, and gradients pass through with factor 1. Carousels go around forever; so do these gradients.

The Output Gate and Hidden State

The cell state $C_t$ can hold values much larger than 1 if the network wants (it is bounded only by what the gates write over time). But downstream layers expect a bounded, stable signal. So the LSTM exposes a second, tamer view of memory:

Two things happen here. First, $\tanh(C_t)$ squashes memory into $(-1, 1)$ — a “readable” form. Second, $o_t$ masks it: channels the network chooses not to reveal this step are zeroed out in $h_t$. The cell state itself is untouched.

Because $h_t$, not $C_t$, is what the next timestep's gates see, the output gate is also the LSTM's way of saying “this part of memory is not relevant to the next prediction” while still keeping it on the belt for later.

Interactive: LSTM Architecture

Click any gate in the diagram below to see its equation, description, and the intuition behind it. Press Animate Flow to watch information move through the cell in the canonical order: forget → write → update → expose.

Interactive: The Gates at Work

Numbers bring intuition. Drag the sliders to set the previous hidden state $h_{t-1}$, the current input $x_t$, and the previous cell state $C_{t-1}$. Watch how each gate responds and how the two final outputs — $C_t$ and $h_t$ — are assembled from them. Try the perfect remembering configuration in the hint below the visualization: push the forget gate toward 1 and the input gate toward 0. You will see $C_t$ essentially equal $C_{t-1}$ — memory preserved.

A Full Forward Pass with Real Numbers

Before we write code, let's run the cell by hand for one step. Take hidden size $H = 2$ and input size $I = 2$. Our first input is the token “I” with embedding $x_1 = [0.5, -0.2]$, and both states start at zero.

1. Concatenate: $z_1 = [h_0, x_1] = [0, 0, 0.5, -0.2]$.
2. Forget gate: $W_f z_1 + b_f = [0.28, 0.04]$, so $f_1 = \sigma([0.28, 0.04]) = [0.5695, 0.5100]$.
3. Input gate: $i_1 = \sigma([-0.23, 0.21]) = [0.4428, 0.5523]$.
4. Candidate: $\tilde{C}_1 = \tanh([0.11, -0.05]) = [0.1096, -0.0500]$.
5. Output gate: $o_1 = \sigma([-0.16, 0.12]) = [0.4601, 0.5300]$.
6. Cell state: $C_1 = f_1 \odot 0 + i_1 \odot \tilde{C}_1 = [0.0485, -0.0276]$.
7. Hidden state: $h_1 = o_1 \odot \tanh(C_1) = [0.0223, -0.0146]$.

These are not made-up numbers; they are what you will see printed when you run the NumPy code below. The point is that once the formulas are written out, running the cell is just algebra — five affine layers and a few element-wise products.

LSTM From Scratch in Python

Here is the complete cell, implemented with nothing but NumPy. Click any line to see the exact tensor values at that point — including the actual matrix entries of each gate's weight matrix, the pre-activations, and the gate outputs for all three timesteps of our toy sequence “I love it”.

    h0     h1     x0     x1
c0  0.2  -0.3   0.4    0.1
c1  0.1   0.5  -0.2    0.3

    h0     h1     x0     x1
c0  0.3   0.2  -0.1    0.4
c1 -0.2   0.1   0.5    0.2

    h0     h1     x0     x1
c0  0.1  -0.4   0.3    0.2
c1  0.4   0.2  -0.1    0.5

    h0     h1     x0     x1
c0  0.5   0.1  -0.2    0.3
c1  0.2  -0.3   0.4   -0.1

1import numpy as np
2
3def sigmoid(x):
4    return 1.0 / (1.0 + np.exp(-x))
5
6# ---- One LSTM cell: processes a single timestep ----
7def lstm_cell(x_t, h_prev, C_prev, params):
8    # Concatenate previous hidden state and current input: shape (hidden + input,)
9    z = np.concatenate([h_prev, x_t])
10
11    # Four gates — each is a linear layer on z, then an activation
12    f_t     = sigmoid(params["W_f"] @ z + params["b_f"])   # forget gate
13    i_t     = sigmoid(params["W_i"] @ z + params["b_i"])   # input gate
14    C_tilde = np.tanh(params["W_C"] @ z + params["b_C"])   # candidate memory
15    o_t     = sigmoid(params["W_o"] @ z + params["b_o"])   # output gate
16
17    # Update the long-term memory (the cell state)
18    C_t = f_t * C_prev + i_t * C_tilde
19
20    # Produce the new hidden state (the exposed output)
21    h_t = o_t * np.tanh(C_t)
22
23    return h_t, C_t
24
25# ---- Run over a sequence ----
26params = {
27    "W_f": np.array([[ 0.2, -0.3,  0.4,  0.1],
28                     [ 0.1,  0.5, -0.2,  0.3]]),
29    "b_f": np.array([ 0.1,  0.2]),
30    "W_i": np.array([[ 0.3,  0.2, -0.1,  0.4],
31                     [-0.2,  0.1,  0.5,  0.2]]),
32    "b_i": np.array([-0.1,  0.0]),
33    "W_C": np.array([[ 0.1, -0.4,  0.3,  0.2],
34                     [ 0.4,  0.2, -0.1,  0.5]]),
35    "b_C": np.array([ 0.0,  0.1]),
36    "W_o": np.array([[ 0.5,  0.1, -0.2,  0.3],
37                     [ 0.2, -0.3,  0.4, -0.1]]),
38    "b_o": np.array([ 0.0, -0.1]),
39}
40
41sequence = [np.array([0.5, -0.2]),   # "I"
42            np.array([0.8,  0.3]),   # "love"
43            np.array([0.1,  0.9])]   # "it"
44
45h = np.zeros(2)
46C = np.zeros(2)
47for t, x_t in enumerate(sequence, 1):
48    h, C = lstm_cell(x_t, h, C, params)
49    print(f"step {t}: h={h.round(4)}, C={C.round(4)}")

Run the code. You should see:

Notice how the magnitudes of $C_t$ are accumulating across steps — memory building up — while the magnitudes of $h_t$ stay small, bounded by $\tanh$ and masked by $o_t$. That is the highway-vs-readout distinction playing out in actual numbers.

The Same LSTM in PyTorch

Now watch how the same mechanism becomes two lines in PyTorch. The NumPy version makes the math concrete; the PyTorch version is what you will actually use. We show both ways: nn.LSTMCell (one step at a time, mirroring our NumPy loop) and nn.LSTM (the whole sequence in a single call).

1import torch
2import torch.nn as nn
3
4torch.manual_seed(42)
5
6# ---- Option A: one cell, manual loop over timesteps ----
7cell = nn.LSTMCell(input_size=2, hidden_size=2)
8
9h = torch.zeros(1, 2)           # (batch=1, hidden=2)
10C = torch.zeros(1, 2)
11seq = [torch.tensor([[0.5, -0.2]]),
12       torch.tensor([[0.8,  0.3]]),
13       torch.tensor([[0.1,  0.9]])]
14
15for t, x_t in enumerate(seq, 1):
16    h, C = cell(x_t, (h, C))
17    print(f"step {t}: h={h.detach().numpy().round(4)}, "
18          f"C={C.detach().numpy().round(4)}")
19
20# ---- Option B: full nn.LSTM — one call processes the whole sequence ----
21lstm = nn.LSTM(input_size=2, hidden_size=2, batch_first=True)
22batch = torch.tensor([[[0.5, -0.2], [0.8, 0.3], [0.1, 0.9]]])   # (1, 3, 2)
23
24output, (h_n, c_n) = lstm(batch)
25print("output shape:", tuple(output.shape))     # (1, 3, 2)
26print("final h_n   :", h_n.detach().numpy().round(4))
27print("final c_n   :", c_n.detach().numpy().round(4))

The output values differ from the NumPy version because PyTorch uses its own random initialization for the weight matrices. The mechanism is identical, line for line — the same four gates, the same additive cell update, the same gated hidden state.

Why LSTMs Actually Learn Long Memory

We have talked about the gradient highway abstractly; the visualization below makes it quantitative. It compares the magnitude of the gradient $\partial L / \partial h_1$ (for a vanilla RNN) against $\partial L / \partial C_1$ (for an LSTM) as a function of sequence length. Drag the sliders and watch the RNN gradients collapse while the LSTM gradients survive.

Read off the math from the visualization. In a vanilla RNN the gradient at timestep 1 is $\prod_{k} W_{hh}^{\top} \tanh'(\cdot)$— a product of many potentially-small terms. In the LSTM the dominant term is $\prod_{k} f_k$, and the network is free to learn $f_k \approx 1$ whenever it needs memory to survive.

Interactive: BPTT Gradient Unroller

Unroll the cell in 3D and watch a gradient particle travel backward along the memory highway. Drag $T$, flip between LSTM and vanilla RNN, and compare the final gradient magnitude displayed above the canvas. The LSTM's particle stays bright while the RNN's fades — that is the exponential shrink from $\prod_k |W_{hh}| \cdot \tanh'(\cdot)$ playing out in pixels.

BPTT by Hand vs Autograd

The highway argument says $\partial C_t / \partial C_{t-1} = \text{diag}(f_t)$ — but that is only the dominant backward path. The full Jacobian also includes smaller contributions that flow through $h_{t-1}$ into the next step's gates. We check both against PyTorch's autograd.

tensor([[0.5000, -0.2000]])

tensor([[0.0, 0.0]])

tensor([[0.0, 0.0]])

1import torch
2
3# Re-use the §17.1 hand-picked weights, cast to torch.
4# We only need the forget-gate outputs to make the highway argument, so rather
5# than reconstruct the full LSTM, we start from the pre-computed f_t values.
6#
7# From the §17.1 trace:
8f_1 = torch.tensor([0.5695, 0.5100])    # forget gate at t=1
9f_2 = torch.tensor([0.6127, 0.5312])    # forget gate at t=2
10f_3 = torch.tensor([0.5577, 0.6186])    # forget gate at t=3
11
12# ---- Analytical highway term ----
13# The LSTM cell-state update is additive: C_t = f_t * C_{t-1} + i_t * g_t.
14# Taking the Jacobian while holding the other paths fixed:
15#     ∂C_t / ∂C_{t-1} = diag(f_t)
16# So the backward product from C_3 all the way back to C_1 is simply
17#     ∂C_3 / ∂C_1 (diag) = f_3 * f_2
18# — two element-wise multiplications, no tanh, no W_hh, no repeated squash.
19dC3_dC1_highway = f_3 * f_2
20print(f"highway term diag(∂C_3/∂C_1) = {dC3_dC1_highway.numpy().round(4)}")
21
22# ---- Autograd verification ----
23# Run the forward pass again, this time with C_1 marked as requires_grad.
24# Then sum C_3 and call .backward() — C_1.grad gives the true gradient,
25# which also includes the (smaller) paths through h_1, h_2 into the gates.
26from torch import nn
27torch.manual_seed(42)
28cell = nn.LSTMCell(input_size=2, hidden_size=2)
29
30x1 = torch.tensor([[0.5, -0.2]])
31x2 = torch.tensor([[0.8,  0.3]])
32x3 = torch.tensor([[0.1,  0.9]])
33
34h0 = torch.zeros(1, 2)
35C0 = torch.zeros(1, 2)
36h1, C1 = cell(x1, (h0, C0))
37C1 = C1.detach().clone().requires_grad_(True)   # seed autograd here
38h2, C2 = cell(x2, (h1, C1))
39h3, C3 = cell(x3, (h2, C2))
40
41# Sum C_3 so .backward() populates C1.grad.
42C3.sum().backward()
43print(f"autograd  ∂ΣC_3/∂C_1 = {C1.grad.detach().numpy().round(4)}")
44print(f"highway   approx     = {dC3_dC1_highway.numpy().round(4)}")
45print("Difference is the h-path contributions autograd captures exactly.")

The autograd result is a few percent larger than the highway approximation — those extra percent are the h-path contributions your hand-derivation explicitly ignored. The key point stands: the dominant term is the diagonal $f_3 \odot f_2$, and it is the one the network can keep near 1 by learning appropriate gate weights.

The Forget-Bias Trick, Ablated

The forget-bias-to-+1 heuristic (Jozefowicz, Zaremba & Sutskever, 2015) initializes $b_f$ so that at the start of training the forget gate is not halving memory every step. The following zero-input experiment isolates the effect: we seed one unit of memory and watch it decay over 20 steps with $b_f = 0$ vs $b_f = 1$.

[[0.0, 0.0]] — (hidden, hidden+input) = (1, 2)

1import numpy as np
2
3def sigmoid(x):
4    return 1.0 / (1.0 + np.exp(-x))
5
6# A minimal LSTM, zero inputs, with a CHOICE of forget-bias.
7# We write a non-zero cell state at t=0 and watch it decay.
8def decay_experiment(b_f_value, steps=20, C0=1.0):
9    # All weights zero → gates depend only on biases.
10    W = np.zeros((1, 2))               # shape (hidden=1, hidden+input=2)
11    b_i, b_g, b_o = 0.0, 0.0, 0.0      # neutral
12
13    h = np.array([0.0])
14    C = np.array([C0])                 # seed one unit of memory at t=0
15    trace = [C[0]]
16
17    for _ in range(steps):
18        z = np.concatenate([h, np.array([0.0])])   # zero input
19        f = sigmoid(W @ z + b_f_value)   # the only gate that matters here
20        i = sigmoid(W @ z + b_i)
21        g = np.tanh(W @ z + b_g)
22        o = sigmoid(W @ z + b_o)
23        C = f * C + i * g                # i*g = 0 because g = tanh(0) = 0
24        h = o * np.tanh(C)
25        trace.append(float(C[0]))
26    return trace
27
28no_bias    = decay_experiment(b_f_value=0.0)
29with_bias  = decay_experiment(b_f_value=1.0)
30
31print(f"b_f = 0:  C_0={no_bias[0]:.4f}   C_20={no_bias[-1]:.3e}")
32print(f"b_f = 1:  C_0={with_bias[0]:.4f}   C_20={with_bias[-1]:.3e}")
33print(f"ratio (with / without): {with_bias[-1] / no_bias[-1]:.0f}x larger")

The Copy Task: An Experimental Proof

Hochreiter and Schmidhuber introduced the copy task in the original 1997 paper as a stress test for long-range memory. The network sees $N$ content tokens, then $L$ distractor steps, then a cue telling it to reproduce the original tokens. A vanilla RNN cannot solve this when $L$ exceeds ~10; an LSTM solves it at $L = 100$ and beyond (Arjovsky, Shah & Bengio, 2016; Bai, Kolter & Koltun, 2018).

LSTM Variants — Peephole, CIFG, ConvLSTM

The LSTM you have just learned is the 1997 / 2000 canonical cell. Over the next two decades researchers proposed dozens of mutations. Greff et al. (2017) — “LSTM: A Search Space Odyssey” — ran 5200 training runs sweeping eight of the most plausible variants on three sequence-modelling tasks and reported a striking result: none of them consistently beat the vanilla cell by a meaningful margin. The ones worth knowing about are below, not because you should usually reach for them, but because you will see them in the wild and the names should mean something.

Check Your Understanding

Four short questions to make sure the pieces are connected. The answers lead with a one-sentence justification, not just “A”.

Summary

• Two states, not one. $C_t$ is long-term memory (no activation on its update), $h_t$ is the gated, bounded output.
• Four sub-networks, one input. All of $f_t, i_t, \tilde{C}_t, o_t$ are affine functions of the same vector $[h_{t-1}, x_t]$.
• Gates use sigmoid; content uses tanh. Switches want $(0, 1)$; signed content wants $(-1, 1)$.
• The whole point is one equation: $C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t$. Additive updates make the gradient highway possible.
• Hidden state reads memory, doesn't own it: $h_t = o_t \odot \tanh(C_t)$.

With the LSTM cell in hand, the next section unpacks the GRU — a later simplification that merges the forget and input gates into a single update gate and drops the cell state entirely. You will see how many LSTM ideas survive in a smaller package, and where the two architectures diverge.

References

• Hochreiter, S. & Schmidhuber, J. (1997). Long Short-Term Memory. Neural Computation 9(8), 1735–1780. [Introduces the LSTM cell and the Constant Error Carousel.]
• Gers, F. A., Schmidhuber, J. & Cummins, F. (2000). Learning to Forget: Continual Prediction with LSTM. Neural Computation 12(10), 2451–2471. [Adds the forget gate to the original 1997 architecture.]
• Jozefowicz, R., Zaremba, W. & Sutskever, I. (2015). An Empirical Exploration of Recurrent Network Architectures. ICML 2015. [Shows empirically that initializing the forget-gate bias to +1 consistently improves training.]
• Pascanu, R., Mikolov, T. & Bengio, Y. (2013). On the difficulty of training recurrent neural networks. ICML 2013 / arXiv:1211.5063. [The canonical analysis of the vanishing/exploding gradient problem in RNNs.]
• Greff, K., Srivastava, R. K., Koutník, J., Steunebrink, B. R. & Schmidhuber, J. (2017). LSTM: A Search Space Odyssey. IEEE TNNLS 28(10), 2222–2232 / arXiv:1503.04069. [Large-scale ablation of LSTM variants — the vanilla cell holds up.]
• Gers, F. A., Schraudolph, N. N. & Schmidhuber, J. (2002). Learning Precise Timing with LSTM Recurrent Networks. Journal of Machine Learning Research 3, 115–143. [Introduces peephole connections.]
• Shi, X., Chen, Z., Wang, H., Yeung, D.-Y., Wong, W.-K. & Woo, W. (2015). Convolutional LSTM Network: A Machine Learning Approach for Precipitation Nowcasting. NeurIPS 2015 / arXiv:1506.04214. [ConvLSTM.]
• Ba, J. L., Kiros, J. R. & Hinton, G. E. (2016). Layer Normalization. arXiv:1607.06450. [Includes the Layer-Norm LSTM variant.]
• Schuster, M. & Paliwal, K. K. (1997). Bidirectional Recurrent Neural Networks. IEEE Transactions on Signal Processing 45(11), 2673–2681. [Introduces the bidirectional recurrent architecture.]
• Graves, A. & Schmidhuber, J. (2005). Framewise phoneme classification with bidirectional LSTM and other neural network architectures. Neural Networks 18(5-6), 602–610. [Bidirectional LSTM applied to speech.]
• Arjovsky, M., Shah, A. & Bengio, Y. (2016). Unitary Evolution Recurrent Neural Networks. ICML 2016 / arXiv:1511.06464. [Benchmarks the copy task at long distractor lengths — LSTM solves L=100+, vanilla RNN fails at L>10.]
• Bai, S., Kolter, J. Z. & Koltun, V. (2018). An Empirical Evaluation of Generic Convolutional and Recurrent Networks for Sequence Modeling. arXiv:1803.01271. [Additional copy-task benchmarks cited by the Copy-Task demo.]