LSTMs, GRUs, and Gated Memory

A vanilla RNN sends gradients through a long chain of state updates. Those gradients can shrink or grow across the chain, making distant dependencies difficult to learn. LSTMs add a memory path whose retention the network can control.

Hochreiter and Schmidhuber introduced the LSTM in 1997 with a separate memory state and gates controlling writes and reads. Gers, Schmidhuber, and Cummins later added a forget gate.

The equations below describe the later LSTM with a forget gate, which lets the cell reset memory on continual streams. LSTMs became widely used in NLP, speech, and time-series modeling before attention-based models became common.

The cell state highway

An LSTM keeps a separate cell state, $c_t$. Gates control how much of its previous value is retained, how much new information is added, and how much is exposed as the hidden state.

The vanilla RNN computes a new hidden state through a nonlinear transformation at each step. The LSTM cell-state update instead adds a gated candidate to a gated copy of the old cell state.

The scalar comparison below tracks how strongly the initial state affects a later state. Both architectures receive the same input stream; the derivative measures sensitivity to that initial state.

Here are the equations for 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 (how much of the previous cell state to retain):

$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. The cell-state update adds the retained old state to the new candidate, scaled by the input gate.

Why addition helps gradient flow

In a vanilla RNN, the gradient flows backward through a chain of matrix multiplications, one per timestep. Repeated contraction along the Jacobian chain can shrink the gradient exponentially.

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}$, so this is one useful path through the derivative, not the full derivative.

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 the improvement can be substantial on tasks that require retaining information across many steps.

Walking through one timestep

Suppose an LSTM has processed "the cat" and now receives "sat". Its previous hidden state $h_{t-1}$ depends on the preceding words; its current input $x_t$ is the embedding for "sat". The visualization uses chosen values to trace the gates.

In phase 1, the chosen forget-gate values are $[0.9, 0.1, 0.8, 0.7]$. The second coordinate retains only 10% of its previous cell-state value; the others retain 70–90%. These coordinates have no assigned linguistic meaning in this example.

In phase 2, the second coordinate receives an input-gate value of 0.8 and a candidate value of 0.9. Their product adds 0.72 to the retained cell value, replacing much of what the forget gate removed.

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 more of that coordinate reaches the hidden state. Dimension 4 has a low output gate (0.3), keeping most of that cell state internal, which can preserve information for later steps.

A long-range dependency, slot by slot

Across a sentence, a gate can retain a feature until a later prediction needs it.

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.

Karpathy's 2015 post visualized a quote-detection cell in a trained character-level LSTM. Radford et al. (2017) found a sentiment-sensitive neuron in a multiplicative LSTM trained on Amazon reviews.

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$

Here z means the amount to write. Some sources, including PyTorch, use the opposite convention: z weights the old state. PyTorch also applies the reset gate after its recurrent candidate projection; check the library equations when comparing implementations.

In the final 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.

The GRU has no separate cell state or 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 arithmetic, though wall-clock speed depends on the implementation and hardware. Chung et al. (2014) found no consistent winner in their GRU–LSTM comparison. In a separate study, Greff et al. (2017) tested eight LSTM variants on three tasks and found no consistent improvement over the standard LSTM.

Implementation from scratch

This PyTorch cell implements the four affine transforms and gated updates explicitly:

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 implementation uses a 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. Karpathy's character-level examples also showed difficulty with dependencies spanning thousands of characters.

"GRUs are always worse than LSTMs." Chung et al. found that performance depended on the task. A GRU has fewer parameters at the same hidden size, but that alone does not determine model quality or training speed.

"The forget gate decides what to forget." Its value is a retention fraction: 1 keeps the old cell value and 0 erases it. The name describes the operation it controls, rather than the direction of the gate value.

What's next

Gates help preserve information, but an encoder still has to summarize its input for a decoder. The next post, Seq2seq and Bahdanau attention, lets the decoder revisit individual source positions instead of relying on one final state.

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.
• PyTorch. GRU equations and implementation conventions. Update-gate direction and reset-gate placement.
• 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.