LSTM Architecture › LSTMs for Long Sequences

Michael BrenndoerferMay 11, 20256 min read
Language AI HandbookAll chapters
Part 11 of 67Recurrent Neural Networks
Chapter 77LSTM Architecture › LSTMs for Long Sequences
Reading section 6 of 14 · LSTMs for Long Sequences · 6 minSection 6 of 14 · 6 min

Book contents

Section 6 of 14 · 6 min read

LSTM Architecture › LSTMs for Long SequencesOne section · Sections in this chapter6 / 14 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The advantage of LSTMs becomes measurable when we compare gradient propagation to vanilla RNNs. Recall from the Vanishing Gradients chapter that vanilla RNN gradients decay exponentially: each timestep multiplies the gradient by a factor typically less than 1. For typical random initialization, this factor is around 0.7 to 0.8, making gradients negligible after 50 or so steps.

LSTMs change this substantially. The additive cell state update allows gradients to flow backward through the addition operation without multiplication by a small value. The forget gate introduces some scaling, but a well-trained forget gate learns values close to 1 for important information, creating a near-constant gradient highway.

The mathematical story is worth tracing explicitly. In a vanilla RNN, the gradient of the loss with respect to the hidden state at step tt, given the loss depends on step t+kt+k, involves a product of kk Jacobian matrices, each associated with the tanh nonlinearity and the weight matrix:

∂L∂ht=(∏j=tt+k−1∂hj+1∂hj)∂L∂ht+k\frac{\partial L}{\partial h_t} = \left(\prod_{j=t}^{t+k-1} \frac{\partial h_{j+1}}{\partial h_j}\right) \frac{\partial L}{\partial h_{t+k}}

Each factor in this product contains the derivative of tanh, which saturates near 0 for large inputs. In the LSTM, the analogous quantity for the cell state path involves only the forget gates:

∂L∂Ct=(∏j=tt+k−1fj+1)∂L∂Ct+k\frac{\partial L}{\partial C_t} = \left(\prod_{j=t}^{t+k-1} f_{j+1}\right) \frac{\partial L}{\partial C_{t+k}}

A product of forget gate values is much gentler than a product of tanh Jacobians. If the network learns forget gates close to 1, this product stays near 1 for many steps, preserving the gradient signal over arbitrarily long distances.

In[12]:
Code
def simulate_gradient_magnitude(seq_length, architecture="vanilla"):
    """
    Estimate gradient magnitude reaching the first timestep.

    Vanilla RNN: exponential decay per step due to tanh saturation
    and matrix multiplication. Typical factor around 0.75.

    LSTM: much gentler decay because cell state gradients bypass
    repeated matrix multiplications. Forget gate factor near 0.99.
    """
    if architecture == "vanilla":
        decay_factor = 0.75
    else:
        decay_factor = 0.992
    return decay_factor**seq_length


seq_lengths = np.arange(1, 201)
vanilla_gradients = [
    simulate_gradient_magnitude(l, "vanilla") for l in seq_lengths
]
lstm_gradients = [simulate_gradient_magnitude(l, "lstm") for l in seq_lengths]

threshold = 1e-6
vanilla_cutoff = next(
    (l for l, g in zip(seq_lengths, vanilla_gradients) if g < threshold), None
)
lstm_cutoff = next(
    (l for l, g in zip(seq_lengths, lstm_gradients) if g < threshold), None
)
Out[13]:
Console
Gradient falls below 1e-06:
  Vanilla RNN: at sequence length ~49
  LSTM: remains above threshold for all tested lengths
Out[14]:
Visualization
Semi-log plot comparing gradient decay in vanilla RNNs versus LSTMs for sequence lengths from 1 to 200.
Gradient magnitude at the first timestep as a function of sequence length, on a logarithmic scale. Vanilla RNNs fall below the practical training threshold (dashed line) around 50 timesteps, making long-range learning impossible. LSTM gradients decay far more slowly and remain trainable across sequences an order of magnitude longer.

The difference is substantial. Vanilla RNN gradients become negligible after around 50 timesteps. LSTM gradients remain orders of magnitude larger and stay above the practical training threshold for sequences well over 100 steps. This gradient stability is why LSTMs enabled practical sequence modeling at scale: tasks like machine translation, speech recognition, and language modeling all involve dependencies that span many dozens of tokens.

Stacking LSTMs

For complex tasks, it is common to stack multiple LSTM layers on top of each other. The output hidden states of the first LSTM layer become the input to the second layer, and so on. Each layer operates on the sequence of hidden states produced by the layer below, allowing deeper layers to model longer-range or more abstract patterns.

A two-layer LSTM processes a sequence as follows. The first layer sees the raw input features x1,x2,…,xTx_1, x_2, \ldots, x_T and produces hidden states h1(1),h2(1),…,hT(1)h_1^{(1)}, h_2^{(1)}, \ldots, h_T^{(1)}. These serve as the input to the second layer, which produces h1(2),h2(2),…,hT(2)h_1^{(2)}, h_2^{(2)}, \ldots, h_T^{(2)}. The second layer's outputs can be used for prediction.

Stacking layers tends to improve performance up to a point. Two or three layers are common in many NLP applications. Beyond three layers, training stability can become an issue: the gradients that must flow backward through both time and depth face a compounding attenuation. Dropout between layers helps regularize stacked LSTMs and is almost always used in practice when stacking more than one layer. In PyTorch, the num_layers argument handles stacking automatically, and the dropout argument applies dropout between all layers except the last.

One subtlety with stacked LSTMs: each layer has its own cell state and hidden state. When you initialize the LSTM with hidden = (h_0, c_0), both tensors have shape (num_layers, batch, hidden_size) to accommodate one state per layer. This is easy to overlook when moving from single-layer to multi-layer configurations.

Bidirectional LSTMs

A standard LSTM processes sequences left to right. At each timestep, it has seen all previous tokens but none of the future ones. For tasks like named entity recognition, sentiment analysis, or document classification, the model has access to the entire input at inference time. In those cases, each token's representation can benefit from both past and future context.

A bidirectional LSTM (Bi-LSTM) addresses this by running two LSTMs over the same sequence: one in the forward direction (left to right) and one in the backward direction (right to left). At each timestep, the forward LSTM has seen all tokens up to position tt, and the backward LSTM has seen all tokens from position tt onward. The hidden states from both directions are concatenated to produce a representation that encodes bidirectional context.

For position tt in a bidirectional LSTM with hidden size dd, the combined representation has dimension 2d2d:

htbi=[ht→;ht←]h_t^{\text{bi}} = [h_t^{\rightarrow}; h_t^{\leftarrow}]

where ht→h_t^{\rightarrow} is the forward hidden state and ht←h_t^{\leftarrow} is the backward hidden state.

The gains from bidirectionality can be substantial for understanding tasks. Consider labeling the word "bank" in a sentence. In "He walked to the river bank to fish," the word "bank" is a geographical feature. In "He visited the bank to deposit a check," it is a financial institution. A unidirectional LSTM would need to predict the label for "bank" before seeing the disambiguating words later in the sentence. A bidirectional LSTM incorporates both the context before and after "bank" simultaneously, making the correct label much easier to determine.

The key limitation of bidirectional LSTMs is that they cannot be used for autoregressive generation. When generating text one token at a time, there is no future context to pass to the backward layer. Bi-LSTMs are therefore primarily used in encoder roles: understanding and representing existing text, not producing new text. In PyTorch, setting bidirectional=True in nn.LSTM enables bidirectional processing, and the output has shape (batch, seq_len, 2 * hidden_size).