6.5 LSTM and GRU · LearnAI

The LSTM (1997) kept recurrent networks alive for two decades and powered the first neural translation systems, speech recognizers and Google's 2016 translation engine. Its gates are the clearest example of a design that fixes a gradient problem by changing the architecture rather than the optimizer — the same philosophy behind residual connections (4.6 Initialization and normalization) — and the gating idea itself returns in modern state-space models (7.12 Beyond vanilla attention: state space models, long context and hybrids).

You'll be able to

Write and explain every LSTM gate equation and the GRU equations
Show why the additive cell-state path lets gradients survive many steps
Count LSTM parameters and compare LSTM and GRU trade-offs
Make sensible architecture choices (stacking, bidirectionality, dropout) and know where RNNs still fit in 2026

Key ideas

An LSTM keeps a memory cell cₜ updated additively: cₜ = fₜ⊙cₜ₋₁ + iₜ⊙c̃ₜ; the hidden state is hₜ = oₜ⊙tanh(cₜ).
Forget, input and output gates are sigmoids of [hₜ₋₁, xₜ]; they are learned per dimension and per step.
∂cₜ/∂cₜ₋₁ ≈ diag(fₜ): with fₜ ≈ 1 the Jacobian is near the identity, so gradients survive many steps (constant error carousel).

Show 4 more key ideas

LSTM parameters: 4(n(n+d)+n); GRU: 3(n(n+d)+n). GRU merges forget/input into one update gate and drops the output gate.
Initialize the forget-gate bias positive; stack 2–4 layers; use bidirectional RNNs for encoders but never for generation.
LSTM and GRU perform about equally; the forget gate and output nonlinearity matter most.
RNNs lost NLP to transformers (parallelism, O(1) path length) but recurrence returned in SSMs (Mamba) and small-footprint settings.

Suggested path through this lesson

The problem to solve

In a vanilla RNN the state is overwritten at every step, $\mathbf{h}_t=\tanh(\mathbf{W}_{hh}\mathbf{h}_{t-1}+\dots)$ , and the gradient between distant steps is a product of Jacobians that shrinks or blows up (6.4 Recurrent neural networks and BPTT). Hochreiter and Schmidhuber's insight (1997) was to give the network a separate memory cell $\mathbf{c}_t$ that is updated additively, with learned gates deciding what to erase, what to write and what to reveal.

The LSTM, gate by gate

All gates read the previous hidden state and the current input, $[\mathbf{h}_{t-1},\mathbf{x}_t]$ (concatenated), through their own weights, and squash to $(0,1)$ with a sigmoid so they act as soft switches:

\begin{aligned} \mathbf{f}_t &= \sigma(\mathbf{W}_f[\mathbf{h}_{t-1},\mathbf{x}_t]+\mathbf{b}_f) &&\text{forget gate: what to erase from memory}\\ \mathbf{i}_t &= \sigma(\mathbf{W}_i[\mathbf{h}_{t-1},\mathbf{x}_t]+\mathbf{b}_i) &&\text{input gate: what to write}\\ \tilde{\mathbf{c}}_t &= \tanh(\mathbf{W}_c[\mathbf{h}_{t-1},\mathbf{x}_t]+\mathbf{b}_c) &&\text{candidate content}\\ \mathbf{c}_t &= \mathbf{f}_t\odot\mathbf{c}_{t-1}+\mathbf{i}_t\odot\tilde{\mathbf{c}}_t &&\text{new memory: keep some old, add some new}\\ \mathbf{o}_t &= \sigma(\mathbf{W}_o[\mathbf{h}_{t-1},\mathbf{x}_t]+\mathbf{b}_o) &&\text{output gate: what to reveal}\\ \mathbf{h}_t &= \mathbf{o}_t\odot\tanh(\mathbf{c}_t) &&\text{hidden state / output} \end{aligned}

( $\odot$ is element-wise multiplication.) Read it as a tiny computer: the cell is RAM, the forget gate clears bits, the input gate writes bits, the output gate decides which bits the rest of the network gets to see this step. A concrete story: to predict the verb in "The keys to the old wooden cabinet are…", one cell dimension can store "subject is plural" when it sees keys (input gate open), hold it through five unrelated words (forget gate ≈ 1, input gate ≈ 0), and release it at the verb (output gate open).

Parameter count. Each of the four blocks has a weight matrix of shape $n\times(n+d)$ and a bias of size $n$ , so an LSTM has $4\big(n(n+d)+n\big)$ parameters. With hidden size $n=512$ and input size $d=256$ : $4\times(512\times768+512)=1{,}574{,}912\approx1.57$ M. Implementations fuse the four matrices into one $4n\times(n+d)$ matrix multiply per step.

Why gradients survive: the additive highway

Differentiate the cell update with respect to the previous cell:

\frac{\partial\mathbf{c}_t}{\partial\mathbf{c}_{t-1}}=\mathrm{diag}(\mathbf{f}_t)+\underbrace{\big(\text{terms through }\mathbf{h}_{t-1}=\mathbf{o}_{t-1}\odot\tanh(\mathbf{c}_{t-1})\big)}_{\text{small when gates are saturated}} .

When the forget gate is near 1 and the other gates are not changing much, the Jacobian is close to the identity, so the product over many steps stays near the identity instead of shrinking like $(\lambda\gamma)^{t-k}$ (6.4 Recurrent neural networks and BPTT). Hochreiter called this the constant error carousel: error signals circulate through the cell unchanged until a gate decides otherwise. The gates are learned, so the network itself chooses, per dimension and per step, how long to remember. Practical corollary: initialize the forget-gate bias to a positive value (≈1) so memory is kept by default early in training (Jozefowicz et al., 2015).

Intuition. Vanilla RNN: $\mathbf{h}_t = F(\mathbf{h}_{t-1})$ — multiplicative, every step reshapes the memory. LSTM: $\mathbf{c}_t = \mathbf{c}_{t-1} + \Delta_t$ (when $\mathbf{f}_t\approx1$ ) — additive, like a residual connection through time (4.6 Initialization and normalization).

This is not a complete cure. LSTMs reliably learn dependencies over hundreds of steps, not tens of thousands; the hidden state is still a fixed-size bottleneck; and computation remains sequential, one step after another.

The GRU: fewer gates, similar power

The gated recurrent unit (Cho et al., 2014) merges the cell and hidden state and uses two gates:

\begin{aligned} \mathbf{z}_t &= \sigma(\mathbf{W}_z[\mathbf{h}_{t-1},\mathbf{x}_t]) &&\text{update gate: how much to overwrite}\\ \mathbf{r}_t &= \sigma(\mathbf{W}_r[\mathbf{h}_{t-1},\mathbf{x}_t]) &&\text{reset gate: how much past to use for the candidate}\\ \tilde{\mathbf{h}}_t &= \tanh(\mathbf{W}[\mathbf{r}_t\odot\mathbf{h}_{t-1},\mathbf{x}_t]) &&\text{candidate}\\ \mathbf{h}_t &= (1-\mathbf{z}_t)\odot\mathbf{h}_{t-1}+\mathbf{z}_t\odot\tilde{\mathbf{h}}_t &&\text{convex blend of old and new} \end{aligned}

The update gate plays both the forget and input roles (they sum to one), there is no output gate, and the parameter count drops to $3\big(n(n+d)+n\big)$ — about 1.18M for the example above. Large comparisons ("LSTM: A Search Space Odyssey", 2015) found that the standard LSTM and the GRU perform about equally across tasks, that the forget gate and the output nonlinearity are the most important components, and that no variant wins consistently. In practice: GRU when data or compute is scarce, LSTM when in doubt.

Architecture choices

Stacking 2–4 layers adds depth per time step; add residual connections between layers if deeper.
Bidirectional RNNs run one network forward and one backward and concatenate states — ideal for encoders and taggers that see the whole input, impossible for generation (you cannot condition on the future).
Dropout on inputs, outputs and between layers; for the recurrent connection use variational dropout (same mask at every step) or none.
Layer normalization inside the recurrence stabilizes training of deep stacks (4.6 Initialization and normalization).
Use fused cuDNN kernels (nn.LSTM) rather than Python loops over steps; pack padded batches.

RNNs in 2026

Transformers replaced LSTMs for almost all of NLP after 2018 because attention trains in parallel and connects any two positions in one step (7.1 Attention Is All You Need: the big picture). But recurrence did not die: it came back as linear recurrences whose training can be parallelized — state-space models such as Mamba and architectures like RWKV (7.12 Beyond vanilla attention: state space models, long context and hybrids) — and even as a revived "xLSTM" with exponential gating. Gated RNNs still win in small-footprint settings (on-device keyword spotting, sensor time series), inside some reinforcement-learning agents (4.11 Reinforcement learning essentials), and as a teaching tool: understanding why the LSTM works is understanding why residual streams and gates are everywhere.

Where this shows up later

Encoder–decoder LSTMs for translation — 6.6 Seq2seq, encoder–decoder and beam search; attention on top of them — 6.7 Attention is born.
Gating as "learned forgetting" reappears in Mamba's selective state spaces — 7.12 Beyond vanilla attention: state space models, long context and hybrids; additive update paths are residual connections — 4.6 Initialization and normalization, 7.4 Inside the block: residuals, LayerNorm, FFN and embeddings.
Contextual word embeddings from biLSTM language models (ELMo) — 6.3 Word embeddings, 7.5 Encoder-only: BERT and friends (and sentence embeddings).

Common pitfalls (4)

Using a bidirectional RNN in a language model or decoder — it would see the token it is supposed to predict.
Applying ordinary dropout to the recurrent connection with a new mask each step; it destroys the memory.
Expecting an LSTM to carry information across thousands of steps; hundreds is the realistic range.
Writing a Python for-loop over time steps in production code instead of using the fused cuDNN implementation.

Recap: all 7 key points

An LSTM keeps a memory cell cₜ updated additively: cₜ = fₜ⊙cₜ₋₁ + iₜ⊙c̃ₜ; the hidden state is hₜ = oₜ⊙tanh(cₜ).
Forget, input and output gates are sigmoids of [hₜ₋₁, xₜ]; they are learned per dimension and per step.
∂cₜ/∂cₜ₋₁ ≈ diag(fₜ): with fₜ ≈ 1 the Jacobian is near the identity, so gradients survive many steps (constant error carousel).
LSTM parameters: 4(n(n+d)+n); GRU: 3(n(n+d)+n). GRU merges forget/input into one update gate and drops the output gate.
Initialize the forget-gate bias positive; stack 2–4 layers; use bidirectional RNNs for encoders but never for generation.
LSTM and GRU perform about equally; the forget gate and output nonlinearity matter most.
RNNs lost NLP to transformers (parallelism, O(1) path length) but recurrence returned in SSMs (Mamba) and small-footprint settings.

Watch

Long Short-Term Memory (LSTM), Clearly Explained

StatQuest · 21 min · open on YouTube ↗

Walks a number through every gate; pause and predict each value before he reveals it.

3 videos

Read

Articles

Understanding LSTM Networks ↗Christopher Olah · 2015The canonical diagrams; every gate explained in one picture each. Read it twice.

Papers

Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation ↗Cho et al. · 2014Introduces the GRU (Section 2.3) and the encoder–decoder framing used in the next lesson.LSTM: A Search Space Odyssey ↗Greff, Srivastava, Koutník, Steunebrink & Schmidhuber · 2015Large-scale ablation of LSTM variants; the forget gate and output activation are what matter.xLSTM: Extended Long Short-Term Memory ↗Beck et al. · 2024The 2024 attempt to scale LSTMs with exponential gating and matrix memory; read the introduction for the modern context.

Books

Dive into Deep Learning — Chapter 10: Modern Recurrent Neural Networks ↗Zhang, Lipton, Li & Smola · 2023LSTM and GRU implemented from scratch and with the framework API; deep and bidirectional RNNs.Deep Learning — Chapter 10.10: The Long Short-Term Memory and Other Gated RNNs ↗Goodfellow, Bengio & Courville · 2016Formal treatment with the self-loop view of the cell state.

Exercises

Trace one LSTM step by hand
Take a 1-dimensional LSTM with all weights 1 and biases 0, previous cell $c_{t-1}=0.5$ , previous hidden $h_{t-1}=0.2$ and input $x_t=1$ . Compute $f_t, i_t, \tilde c_t, c_t, o_t, h_t$ . Expected: $f_t=i_t=o_t=\sigma(1.2)\approx0.769$ , $\tilde c_t=\tanh(1.2)\approx0.834$ , $c_t\approx0.5\cdot0.769+0.769\cdot0.834\approx1.026$ , $h_t\approx0.769\cdot\tanh(1.026)\approx0.593$ .
Swap the cell and watch the gradient
Repeat the gradient-decay plot from 6.4 Recurrent neural networks and BPTT (‖∂L_T/∂h_k‖ vs T−k) for a vanilla RNN, an LSTM and a GRU of the same hidden size on the same text. Expected: the RNN decays exponentially; the gated cells decay far more slowly, and the LSTM with forget-bias 1 slowest of all.
Parameter accounting
Count the parameters of nn.LSTM(256, 512) and nn.GRU(256, 512) by formula, then verify with sum(p.numel() for p in m.parameters()). Expected: 1,574,912 and 1,181,184 (PyTorch keeps two bias vectors per gate, so you'll see slightly more — explain the difference).

Quiz

Question 1 of 5

Which LSTM gate decides how much of the previous memory cell to keep?

Flashcards

Card 1 of 7

These cards are now in your spaced-repetition queue. Review all due cards →

Scored 80% or more and reviewed the cards? You're done here.