RNNs and LSTMs

Problem

Text, audio, sensor readings and clicks arrive as sequences, and their lengths vary. A feed-forward network from neurons and layers takes a fixed-size input, so it would need a separate weight for every position and could never see a longer sequence than it was built for. We want one function that consumes the sequence an element at a time, remembers what matters about the past in a fixed-size vector, and is trained with the same backpropagation as everything else.

The plain recurrent network does exactly that, but its memory is short: the gradient that tells step 1 how it affected the loss at step 30 is a product of 29 per-step factors, and it usually vanishes or explodes. The LSTM's gated, additive cell state is the classic fix. Before the Transformer, encoder-decoder LSTMs with attention were the state of the art in machine translation, and attention itself was invented to help them.

Intuition

Unroll the loop and an RNN is a very deep feed-forward network with one layer per time step, all layers sharing the same weights. Information from step 1 reaches step T only by passing through T − 1 copies of that layer, each multiplying it by the recurrent weight and squashing it through tanh. That is a game of telephone: the message fades unless every relay repeats it at full volume.

An LSTM adds a second, separate line, the cell state c, that does not pass through a weight matrix or a squashing function from step to step. Each step only scales it by a forget gate between 0 and 1 and adds new content through an input gate. If the forget gate stays near 1, the cell carries information, and gradient, across many steps nearly unchanged. Learning when to keep, write and read the cell is the whole trick.

One cell, reused at every step. Backprop through time runs the dashed path, and each step multiplies by the same Jacobian.

Mechanics

The recurrence

PyTorch's nn.RNN computes, for each step t,

h_t = tanh(x_t W_ih^⊤ + b_ih + h_t−1 W_hh^⊤ + b_hh)

with nonlinearity="relu" as the alternative to tanh. The two bias vectors always appear as a sum, so one would do; the source keeps the second "for CuDNN compatibility". The output at each step is h_t itself, and a model reads off a prediction with a linear layer on top: from every step for tagging, or from the last one for classifying the whole sequence. Stacking (num_layers) feeds layer k's outputs as layer k+1's inputs. A bidirectional RNN runs a second copy right to left and concatenates the two states, so each output sees the whole sequence.

Backprop through time and vanishing gradients

Training unrolls the loop and applies the chain rule along the unrolled graph; this is backpropagation through time (BPTT). The gradient from step T back to step k is a product of Jacobians:

∂h_T / ∂h_k = Π_t=k+1..T diag(1 − h_t²) · W_hh

Each factor has the same W_hh. If its largest singular value times the tanh slope is below 1, the product shrinks geometrically with the distance T − k; if it is above 1 it can grow geometrically. Since tanh′ ≤ 1, a large W_hh does not reliably help: it pushes the units into saturation, where the slope is near 0. Bengio, Simard and Frasconi (1994) showed this tension is fundamental for learning long-range dependencies by gradient descent. Exploding gradients are tamed by gradient clipping (Pascanu et al., 2013), rescaling the gradient when its norm passes a threshold. Vanishing gradients need a change of architecture. In practice long sequences are also cut into chunks, carrying the hidden state across chunks but stopping the gradient at each boundary (truncated BPTT).

LSTM gates, derived

To keep a gradient alive, we want a state whose step-to-step Jacobian is close to the identity. The simplest such update is additive, c_t = c_t−1 + (new content), so ∂c_t/∂c_t−1 = 1: Hochreiter and Schmidhuber's "constant error carousel". A pure sum never forgets and lets everything in, so each part gets a learned, data-dependent gate, a sigmoid between 0 and 1. The forget gate was added by Gers, Schmidhuber and Cummins (2000). PyTorch's LSTM:

i_t = σ(W_iix_t + b_ii + W_hih_t−1 + b_hi)   input gate
f_t = σ(W_ifx_t + b_if + W_hfh_t−1 + b_hf)   forget gate
g_t = tanh(W_igx_t + b_ig + W_hgh_t−1 + b_hg)   candidate
o_t = σ(W_iox_t + b_io + W_hoh_t−1 + b_ho)   output gate
c_t = f_t ⊙ c_t−1 + i_t ⊙ g_t     h_t = o_t ⊙ tanh(c_t)

The direct path from c_t−1 to c_t has Jacobian diag(f_t): no weight matrix, no squashing. Over many steps it multiplies by Π f_t, which stays near 1 for as long as the network keeps its forget gates open. The gates also depend on h_t−1, which depends on c_t−1, so there are indirect paths too; the worked example measures the full derivative. The four gates' weights are stored stacked: weight_ih_l0 is (W_ii|W_if|W_ig|W_io), shape (4·hidden_size, input_size), so one matrix multiply computes all four. PyTorch initializes every LSTM weight and bias from U(−1/√hidden, 1/√hidden), forget bias included. Jozefowicz, Zaremba and Sutskever (2015) recommend setting the forget-gate bias to 1 or more; in PyTorch you do it by hand. With proj_size > 0, h_t is multiplied by a learned W_hr to a smaller size (Sak et al., 2014), while c_t keeps hidden_size.

The green line is the cell state: one multiply by f and one add per step.

GRU

The gated recurrent unit (Cho et al., 2014) merges the cell and hidden state and uses two gates, a reset gate r and an update gate z:

r_t = σ(W_irx_t + b_ir + W_hrh_t−1 + b_hr)    z_t = σ(W_izx_t + b_iz + W_hzh_t−1 + b_hz)
n_t = tanh(W_inx_t + b_in + r_t ⊙ (W_hnh_t−1 + b_hn))    h_t = (1 − z_t) ⊙ n_t + z_t ⊙ h_t−1

The last line is an interpolation, so ∂h_t/∂h_t−1 contains diag(z_t) on its direct path: the update gate plays the forget gate's role. A GRU has three gate blocks instead of four, so 3/4 of an LSTM's parameters at the same size, and in comparisons the two usually perform similarly. The PyTorch docs note that its n_t applies r_t after multiplying by W_hn, unlike the original paper's W_hn(r_t ⊙ h_t−1), "on purpose for efficiency": all three hidden-side products then come from one matrix multiply.

Sequence to sequence, and attention as the bridge

For translation, an encoder RNN reads the source sentence and a decoder RNN, started from the encoder's final state, generates the target one token at a time, feeding each output back in as the next input (Sutskever, Vinyals and Le, 2014). The whole source must squeeze through one fixed-size vector, and quality dropped on long sentences. Bahdanau, Cho and Bengio (2014) let the decoder, at every step, score all encoder states against its current state, turn the scores into weights with a softmax, and read their weighted average, the context vector. That is attention. The Transformer kept the attention and dropped the recurrence, which removed the sequential dependency that stops an RNN from processing all positions of a sequence in parallel during training.

Padding and packed sequences

A batch holds sequences of different lengths. Padding them to the longest wastes compute on the padding and, worse, runs the recurrence through padding, so the "last" hidden state of a short sequence is polluted. A PackedSequence stores only real elements, time step by time step: at each step it holds the elements of every sequence still running, with sequences sorted from longest to shortest. batch_sizes[t] is how many sequences are still running at step t, so the kernel just processes a shrinking batch. pack_padded_sequence builds it from a padded tensor plus lengths, pad_packed_sequence undoes it, and with enforce_sorted=False the sort is done for you and recorded in sorted_indices/unsorted_indices so the outputs come back in your order. The RNN then returns h_n at each sequence's true last step.

Worked example

Only a forget gate held near 1 keeps both the content and the gradient alive across thirty steps.

Implementation

import torch
import torch.nn as nn
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence, pad_packed_sequence

class Tagger(nn.Module):
    def __init__(self, vocab, dim=64, hidden=128, n_tags=10):
        super().__init__()
        self.embed = nn.Embedding(vocab, dim, padding_idx=0)
        self.lstm = nn.LSTM(dim, hidden, num_layers=2, batch_first=True,
                            dropout=0.1, bidirectional=True)
        self.out = nn.Linear(2 * hidden, n_tags)
        for name, p in self.lstm.named_parameters():
            if name.startswith("bias_ih"):     # gate order i, f, g, o
                p.data[hidden:2 * hidden] = 1.0
            elif name.startswith("bias_hh"):
                p.data[hidden:2 * hidden] = 0.0

    def forward(self, tokens, lengths):      # tokens (N, L) padded with 0
        x = self.embed(tokens)
        packed = pack_padded_sequence(x, lengths.cpu(), batch_first=True,
                                      enforce_sorted=False)
        out, (h_n, c_n) = self.lstm(packed)  # h_n: (2 layers * 2 directions, N, hidden)
        out, _ = pad_packed_sequence(out, batch_first=True)
        return self.out(out)                 # (N, L, n_tags)

seqs = [torch.tensor([5, 2, 9]), torch.tensor([7, 1]), torch.tensor([3])]
tokens = pad_sequence(seqs, batch_first=True)          # (3, 3), zeros as padding
lengths = torch.tensor([len(s) for s in seqs])
model = Tagger(vocab=10)
loss = model(tokens, lengths).sum()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # tame exploding gradients

Pass lengths on the CPU: batch_sizes is always a CPU int64 tensor. batch_first=True applies to input and output but not to h_n and c_n, which stay (D·num_layers, N, H). For a bidirectional LSTM, out[:, -1] is not the same as h_n: the reverse direction's final state is at time step 0. Dropout here applies between the two layers only; PyTorch warns if you set it with num_layers=1.

Tradeoffs

• Sequential by nature. Step t needs step t − 1, so training cannot be parallelized across time the way a Transformer can. On modern accelerators that, more than accuracy, is why Transformers replaced RNNs for large models.
• Constant memory at inference. An RNN's state is a fixed-size vector, whatever the context length, while a Transformer's KV cache grows with every token. That is the appeal that state space models and other modern linear recurrences revisit.
• LSTM versus GRU. GRUs have three gate blocks to the LSTM's four and are a little cheaper; in most comparisons neither wins consistently. Start with an LSTM, set the forget bias, and try a GRU if speed matters.
• Memory is still finite. Gates slow the decay but do not remove it, and everything must pass through a fixed-size state. Attention over the encoder states sidesteps the bottleneck.
• Packing versus padding. Packing is correct and saves compute on uneven lengths, but costs a sort and some bookkeeping. With similar lengths (bucketed batches), padding plus a mask is simpler.

Connections

References

• Elman. Finding Structure in Time. Cognitive Science, 1990.
• Werbos. Backpropagation Through Time: What It Does and How to Do It. Proceedings of the IEEE, 1990.
• Bengio, Simard and Frasconi. Learning Long-Term Dependencies with Gradient Descent is Difficult. IEEE Transactions on Neural Networks, 1994.
• Hochreiter and Schmidhuber. Long Short-Term Memory. Neural Computation, 1997.
• Gers, Schmidhuber and Cummins. Learning to Forget: Continual Prediction with LSTM. Neural Computation, 2000.
• Pascanu, Mikolov and Bengio. On the difficulty of training Recurrent Neural Networks. ICML 2013. arXiv:1211.5063
• Cho et al. Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation. EMNLP 2014. arXiv:1406.1078
• Sutskever, Vinyals and Le. Sequence to Sequence Learning with Neural Networks. NeurIPS 2014. arXiv:1409.3215
• Bahdanau, Cho and Bengio. Neural Machine Translation by Jointly Learning to Align and Translate. 2014. arXiv:1409.0473
• Sak, Senior and Beaufays. Long Short-Term Memory Based Recurrent Neural Network Architectures for Large Vocabulary Speech Recognition. 2014. arXiv:1402.1128
• Jozefowicz, Zaremba and Sutskever. An Empirical Exploration of Recurrent Network Architectures. ICML 2015.
• PyTorch v2.15.0: torch/nn/modules/rnn.py, torch/nn/utils/rnn.py.