LSTMs, GRUs, and Gated Memory

So the vanilla RNN has a problem. The hidden state gets multiplied by the same weight matrix at every timestep, so gradients either shrink exponentially (vanishing) or blow up (exploding) as you backpropagate through time. In practice, an RNN can barely learn dependencies longer than about 10-20 timesteps. If the subject of a sentence appeared 30 words ago and the verb needs to agree with it, the vanilla RNN has already forgotten.

The question Hochreiter and Schmidhuber asked in 1997 was: what if the network could decide what to remember and what to forget? What if, instead of blindly squashing the entire hidden state through a tanh every step, we gave the cell explicit gates that control the flow of information?

That idea became the Long Short-Term Memory cell. The 1997 original introduced the constant-error cell-state path with input and output gates. The forget gate (the part most modern presentations treat as the heart of the model) was added a few years later by Gers, Schmidhuber, and Cummins (2000) so the cell could learn to reset itself on continual streams. The equations in this post are that modern LSTM, not exactly the 1997 cell. With those refinements in place, LSTMs became the dominant neural sequence model in much of NLP, speech, and time-series modeling before attention-heavy and Transformer models displaced them.

The cell state highway

The core insight of an LSTM is the cell state, $c_t$. Think of it as a highway running above the main computation. Information flows along this highway mostly unchanged, with gates acting as on-ramps and off-ramps that merge new data in or erase old data out.

In a vanilla RNN, the hidden state is completely overwritten every step. In an LSTM, the cell state is additively updated. Old information plus new information.

The easiest way I've found to feel the difference is to plant a signal at step 0 and ask, "how much of it is left a few steps later?" This is a scalar, toy picture of sensitivity, not a claim that real high-dimensional hidden states preserve a clean separable chunk of the original input. With that caveat: both architectures see the same stream of inputs. The RNN squashes the old state through a tanh each step. The LSTM scales the old cell-state piece by a number close to 1 and adds the new stuff on top.

Now let's look at the four equations that define one LSTM timestep. Each gate takes the concatenation of the previous hidden state $h_{t-1}$ and the current input $x_t$, then applies its own learned weight matrix and bias.

Forget gate (what to erase from cell state):

$f_t = \sigma(W_f \cdot [h_{t-1}, x_t] + b_f)$

Input gate (how much of the new candidate to write):

$i_t = \sigma(W_i \cdot [h_{t-1}, x_t] + b_i)$

Candidate values (what new information looks like):

$\tilde{c}_t = \tanh(W_c \cdot [h_{t-1}, x_t] + b_c)$

Output gate (what part of cell state to reveal as hidden state):

$o_t = \sigma(W_o \cdot [h_{t-1}, x_t] + b_o)$

Then the cell state update and hidden state output:

$c_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t$

$h_t = o_t \odot \tanh(c_t)$

The $\sigma$ is the sigmoid function (outputs between 0 and 1, so each gate value acts like a dial). The $\odot$ is element-wise multiplication. And that cell state equation is the key: it's an addition. The old cell state, scaled by the forget gate, plus new information, scaled by the input gate.

Why addition fixes vanishing gradients

In a vanilla RNN, the gradient flows backward through a chain of matrix multiplications, one per timestep. If the spectral radius of the weight matrix is less than 1, those multiplied gradients shrink exponentially toward zero.

In an LSTM, the gradient can flow backward along the cell state highway via that additive update. Along the direct cell-state path, treating the gates at that timestep as fixed, $\partial c_t / \partial c_{t-1} = f_t$. The full Jacobian has more terms because $f_t$, $i_t$, and $\tilde{c}_t$ all depend on $h_{t-1}$ and therefore indirectly on $c_{t-1}$, but the dominant story is the diagonal scaling by the forget gate.

As long as the forget gate stays close to 1, the direct cell-state gradient decays slowly instead of being repeatedly transformed by the full recurrent matrix and the tanh derivative. That decay is still a product of forget-gate values, so a forget gate of 0.97 over 40 steps gives roughly $0.97^{40} \approx 0.30$ along that path. LSTMs mitigate vanishing gradients rather than eliminating them, but mitigating them is enough to learn dependencies orders of magnitude longer than a vanilla RNN can.

Everything else (the input gate, the output gate, the candidate) is machinery for controlling what gets onto and off of that highway.

Walking through one timestep

Let's make this concrete. Suppose our LSTM has been processing the sentence "the cat" and is now seeing the word "sat". The previous hidden state $h_{t-1}$ encodes what the network learned from "the cat", and the current input $x_t$ is the embedding for "sat".

Step through the visualization below to watch each gate fire in sequence:

Look at what happened in phase 1. The forget gate output was $[0.9, 0.1, 0.8, 0.7]$. That second dimension got a value of 0.1, meaning the cell almost completely erased whatever was stored there. Maybe that dimension was tracking something about the previous word's part of speech that is no longer relevant now that we've seen a verb. The other dimensions stayed mostly intact (0.7-0.9), preserving the memory of "the cat".

In phase 2, the input gate decides how much new information to write. The second dimension gets a high input gate value (0.8) paired with a strong positive candidate (0.9), so a lot of new information floods in right where the old information was just erased. The network cleared a slot and immediately filled it with verb-related features.

By phase 5, the output gate filters the updated cell state to produce the new hidden state. Dimension 3 has a high output gate (0.9), so the network is confident about exposing that information. Dimension 4 has a low output gate (0.3), keeping most of that cell state internal, perhaps saving it for a later timestep.

A long-range dependency, slot by slot

One timestep is fine for showing mechanics. The interesting question is what happens across many timesteps, when one piece of information needs to survive a whole sentence before it becomes useful.

Consider subject-verb agreement with a relative clause in the middle: "the keys that the man on the bench forgot are on the table". The verb "are" has to agree with "keys" (plural), not with "man" or "bench" which are closer but irrelevant. A trained LSTM can learn mechanisms like this. In some language models, researchers have found interpretable units that track number or syntax (Lakretz et al. found long-distance number information largely managed by two "number units" in a studied LSTM, and Linzen et al. showed LSTMs capture non-trivial grammatical structure even when sequential and structural cues conflict), though real models often distribute the same information across several dimensions.

The viz below is hand-built, not trained, but it shows the mechanism. Watch the c[0] slot when "keys" arrives. The forget gate for that slot drops near zero (erase whatever was there), the input gate writes a large positive value, and from that point on the forget gate holds near 0.97 through every other token. By the time "are" arrives, the plural signal is still intact.

Real trained LSTMs do develop interpretable neurons like this. Karpathy's 2015 post visualized a "quote-detector" neuron and a "line-length" neuron inside a character-level LSTM trained on Linux kernel source. Radford et al. (2017) later found an "unsupervised sentiment neuron" inside a large multiplicative LSTM trained only on Amazon reviews. The network discovers, on its own, which features are worth carrying across time, and dedicates specific cell-state dimensions to them.

The GRU: a streamlined alternative

In 2014, Cho et al. proposed the Gated Recurrent Unit as a simpler alternative. The GRU merges the cell state and hidden state into a single vector and uses only two gates instead of three.

Reset gate (how much of the previous hidden state to expose when computing the candidate):

$r_t = \sigma(W_r \cdot [h_{t-1}, x_t] + b_r)$

Update gate (the interpolation factor between old and new):

$z_t = \sigma(W_z \cdot [h_{t-1}, x_t] + b_z)$

Candidate hidden state:

$\tilde{h}_t = \tanh(W_h \cdot [r_t \odot h_{t-1}, x_t] + b_h)$

Final hidden state:

$h_t = (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t$

Notice that last equation. The update gate $z_t$ interpolates between keeping the old hidden state and adopting the new candidate. When $z_t$ is close to 0, the hidden state passes through unchanged (like the LSTM's forget gate being close to 1). When $z_t$ is close to 1, the old state is replaced with new information.

So the GRU accomplishes something similar to the LSTM but with fewer parameters. There's no separate cell state, no output gate. The update gate $z_t$ is the main keep-vs-write mechanism: it decides how much old hidden state to preserve versus how much candidate state to adopt, the same job the LSTM splits between forget and input gates. The reset gate $r_t$ does something different: it controls how much of the previous hidden state is visible while computing the candidate, gating what counts as "context" for the new write.

Fewer parameters reduce parameter memory and usually reduce compute, though actual wall-clock speed depends on the implementation, fused kernels, and the hardware. And empirically? Chung et al. (2014) ran a careful comparison and found that neither architecture consistently outperforms the other. GRUs tend to do better on smaller datasets, LSTMs tend to have more capacity for complex tasks, but honestly the differences are often within noise. Greff et al. (2017) did an extensive hyperparameter search and reached a similar conclusion: the vanilla LSTM is remarkably robust, and most proposed modifications (including the GRU simplification) don't reliably improve on it.

Implementation from scratch

Here's an LSTM cell in PyTorch, not using nn.LSTM but building the gates manually. This makes the four equations completely explicit.

import torch
import torch.nn as nn
 
class LSTMCell(nn.Module):
    def __init__(self, input_size: int, hidden_size: int):
        super().__init__()
        self.hidden_size = hidden_size
        # One big linear layer for all four gates at once.
        # Input: [h_{t-1}, x_t] of size (hidden_size + input_size)
        # Output: 4 * hidden_size (forget, input, candidate, output)
        self.gates = nn.Linear(input_size + hidden_size, 4 * hidden_size)
 
    def forward(self, x_t, h_prev, c_prev):
        # x_t:    (batch, input_size)
        # h_prev: (batch, hidden_size)
        # c_prev: (batch, hidden_size)
 
        combined = torch.cat([h_prev, x_t], dim=1)
        gate_values = self.gates(combined)  # (batch, 4 * hidden_size)
 
        # Split into four chunks along the last dimension
        f, i, c_tilde, o = gate_values.chunk(4, dim=1)
 
        f = torch.sigmoid(f)                # forget gate
        i = torch.sigmoid(i)                # input gate
        c_tilde = torch.tanh(c_tilde)       # candidate
        o = torch.sigmoid(o)                # output gate
 
        # Cell state update: the additive highway
        c_t = f * c_prev + i * c_tilde
 
        # Hidden state output
        h_t = o * torch.tanh(c_t)
 
        return h_t, c_t
 
class LSTMLanguageModel(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_size, num_layers=1):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.cells = nn.ModuleList([
            LSTMCell(embed_dim if layer == 0 else hidden_size, hidden_size)
            for layer in range(num_layers)
        ])
        self.output_proj = nn.Linear(hidden_size, vocab_size)
        self.hidden_size = hidden_size
        self.num_layers = num_layers
 
    def forward(self, token_ids):
        # token_ids: (batch, seq_len)
        batch_size, seq_len = token_ids.shape
        embeds = self.embedding(token_ids)  # (batch, seq_len, embed_dim)
 
        # Initialize hidden and cell states to zeros
        h = [torch.zeros(batch_size, self.hidden_size, device=token_ids.device)
             for _ in range(self.num_layers)]
        c = [torch.zeros(batch_size, self.hidden_size, device=token_ids.device)
             for _ in range(self.num_layers)]
 
        outputs = []
        for t in range(seq_len):
            x = embeds[:, t, :]
            for layer_idx, cell in enumerate(self.cells):
                h[layer_idx], c[layer_idx] = cell(x, h[layer_idx], c[layer_idx])
                x = h[layer_idx]  # feed this layer's output to the next
            outputs.append(x)
 
        out = torch.stack(outputs, dim=1)       # (batch, seq_len, hidden_size)
        logits = self.output_proj(out)           # (batch, seq_len, vocab_size)
        return logits

The key trick is the single nn.Linear(..., 4 * hidden_size). All four gate computations run as one matrix multiply, and then we split the output with .chunk(4). Efficient LSTM implementations usually fuse the four gate affine transforms into one larger matrix multiply, because one big matmul is much faster on a GPU than four small ones.

Misconceptions

"LSTMs completely solve vanishing gradients." They mitigate the problem significantly, but they don't eliminate it. The gradient still flows through the forget gate at each timestep, and if the forget gate learns values consistently below 1.0, the gradient will still decay over very long sequences. In practice, LSTMs work well up to maybe a few hundred timesteps. Beyond that, performance degrades. Karpathy's famous blog post showed LSTMs doing remarkable things with character-level text, but even he noted they struggled with dependencies spanning thousands of characters.

"GRUs are always worse than LSTMs." The comparison is actually task-dependent, as Chung et al. showed. GRUs have fewer parameters, which means they're less prone to overfitting on small datasets and faster to train. For many practical tasks, especially when you're not memory-constrained and the sequences aren't extremely long, GRUs are a perfectly reasonable choice.

"The forget gate decides what to forget." A naming issue. The forget gate actually decides what to keep. A forget gate value of 1.0 means "keep everything", and 0.0 means "erase everything". Hochreiter and Schmidhuber named it the forget gate because it controls the forgetting process, but when you're reading the equations, think of it as a "remember gate". Values close to 1 preserve information. (Side note on naming history: the gate Hochreiter and Schmidhuber introduced in 1997 was the input gate; the forget gate is the Gers, Schmidhuber, and Cummins addition from 2000.)

What's next

The gradient highway solved one problem, but not the hardware problem. LSTMs and GRUs still process tokens in order, so training cannot parallelize across sequence positions the way a Transformer can. Vaswani et al. framed this directly: recurrence's sequential nature precludes parallelization within training examples, which becomes painful for long sequences and large batches. Transformers won largely because attention gave every token a short path to every other token and let training fan out across all positions in parallel. That's the next arc.

Additional reading (and watching)

• Hochreiter, S. & Schmidhuber, J. (1997). Long Short-Term Memory. Neural Computation, 9(8), 1735-1780. The original LSTM, with input and output gates and the constant-error cell-state path. The forget gate isn't here yet.
• Gers, F.A., Schmidhuber, J., & Cummins, F. (2000). Learning to Forget: Continual Prediction with LSTM. Neural Computation, 12(10), 2451-2471. Added the forget gate so an LSTM cell can learn to reset itself on continual streams. This is the equation set most modern presentations call "the LSTM."
• Jozefowicz, R., Zaremba, W., & Sutskever, I. (2015). An Empirical Exploration of Recurrent Network Architectures. ICML 2015. Found that initializing the LSTM forget-gate bias to 1 closed much of the gap between LSTM and GRU.
• Cho, K., et al. (2014). Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation. arXiv:1406.1078.
• Chung, J., et al. (2014). Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling. arXiv:1412.3555.
• Greff, K., et al. (2017). LSTM: A Search Space Odyssey. IEEE Transactions on Neural Networks and Learning Systems.
• Karpathy, A. (2015). The Unreasonable Effectiveness of Recurrent Neural Networks. Blog post. Visualized interpretable neurons inside a character-level LSTM.
• Radford, A., Jozefowicz, R., & Sutskever, I. (2017). Learning to Generate Reviews and Discovering Sentiment. arXiv:1704.01444. A single LSTM unit in a large language model trained on Amazon reviews turned out to track sentiment.
• Linzen, T., Dupoux, E., & Goldberg, Y. (2016). Assessing the Ability of LSTMs to Learn Syntax-Sensitive Dependencies. TACL. Showed LSTMs capture non-trivial grammatical structure but still err when sequential and structural cues conflict.
• Lakretz, Y., et al. (2019). The Emergence of Number and Syntax Units in LSTM Language Models. NAACL 2019. Identified two specific units in a trained LSTM that carry long-distance number information for subject-verb agreement.
• OpenAI. (2019). Dota 2 with Large Scale Deep Reinforcement Learning. arXiv:1912.06680. Used a 4096-unit LSTM as the core memory of the Dota 2 agent.
• Gu, A. & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752. Selective SSMs that bring back the RNN-style recurrence while remaining parallelizable in training.