RNN Vanishing Gradients: BLEU Drop 42→29 on Long Inputs

Language translation, real-time speech recognition, stock price forecasting, music generation — every one of these tasks shares a property that standard feedforward networks fundamentally cannot handle: the output depends not just on the current input, but on a sequence of past inputs. When Google Translate converts a sentence from German to English, word order shifts dramatically between languages, so the model must carry meaning across dozens of tokens simultaneously. That is a sequence problem, and it is everywhere in production ML.

The feedforward network processes each input in isolation. Feed it the word 'bank' with no context and it cannot tell you whether the answer is a financial institution or a river bank. Recurrent Neural Networks solve this by threading a hidden state through time — each timestep reads the current input and the previous hidden state together, creating a rolling summary of everything seen so far. The problem is that 'rolling summary' degrades fast. After thirty timesteps, the gradient signal needed to teach the network about something that happened at timestep one has been multiplied by a weight matrix thirty times over, and it either vanishes to zero or explodes to infinity. Long Short-Term Memory networks, introduced by Hochreiter and Schmidhuber in 1997, are the engineering answer to that mathematical catastrophe.

By the end of this article you'll understand exactly why vanilla RNNs fail on long sequences, how LSTM gates control information flow at the mathematical level, how to implement and train both in PyTorch with production-quality code, and the real mistakes that silently destroy model performance in live systems. You'll also walk away with the precise vocabulary to answer LSTM questions in a senior ML engineering interview.

The Recurrent Cell: How RNNs Process Sequences

A Recurrent Neural Network processes a sequence of inputs by maintaining a hidden state vector that is updated at each timestep. At time t, the hidden state h_t is computed as h_t = tanh(W_ih x_t + b_ih + W_hh h_{t-1} + b_hh). The same weight matrices W_ih and W_hh are reused at every step — that's the recurrence. This parameter sharing is both the power and the curse: it lets the model generalise across varying sequence lengths, but it also means the gradient involves repeated multiplication by the same matrix, causing it to grow or shrink exponentially with sequence length.

Production engineers often forget that the hidden state dimension must be large enough to capture the information bottleneck. If your sequence contains 200 words of financial news, a hidden size of 32 will force severe compression — you'll lose sentiment cues from the third sentence. Rule of thumb: hidden size >= vocabulary size * 0.05 for language tasks, but measure empirically via ablation.

In PyTorch, a single-layer RNN is trivial to instantiate, but the default initialisation for the recurrent weight matrix uses uniform distribution in [-1/sqrt(hidden), 1/sqrt(hidden)]. That range is too narrow to preserve gradient magnitude beyond 10 steps. Always override with orthogonal initialisation for the recurrent kernel.

import torch
import torch.nn as nn

class SimpleRNN(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.rnn = nn.RNN(
            input_size=embed_dim,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            nonlinearity='tanh'
        )
        # Override default uniform initialisation with orthogonal for gradient flow
        for name, param in self.rnn.named_parameters():
            if 'weight_hh' in name:
                nn.init.orthogonal_(param, gain=1.0)
            elif 'bias' in name:
                nn.init.zeros_(param)
        self.fc = nn.Linear(hidden_size, vocab_size)

    def forward(self, x, hidden=None):
        # x shape: (batch, seq_len)
        emb = self.embedding(x)  # (batch, seq_len, embed_dim)
        out, hidden = self.rnn(emb, hidden)  # out: (batch, seq_len, hidden)
        logits = self.fc(out)    # (batch, seq_len, vocab_size)
        return logits, hidden

The Vanishing & Exploding Gradient Problem — Why It's the Core Issue

Backpropagation through time (BPTT) unrolls the recurrent computation into a deep feedforward network with T layers, each sharing the same weight matrix W_hh. The gradient with respect to the loss at timestep T, when propagated to timestep 1, involves the matrix product W_hh^(T-1). The singular values of W_hh determine whether this product shrinks to zero (vanishing) or grows to infinity (exploding).

In practice, tanh activation compresses outputs to (-1,1), and if any singular value of W_hh is less than 1, the gradient decays exponentially. Exploding happens when singular values exceed 1 — common if the RNN weights grow during training. The fix for exploding is gradient clipping: cap the gradient norm at some threshold. The fix for vanishing is architectural: change the recurrence itself to allow gradients to flow unchanged.

LSTMs solve vanishing by introducing a cell state that is only modified by additive gates, not multiplied by a weight matrix at each step. The cell state's contribution to the gradient is the forget gate value, which is a learned sigmoid between 0 and 1 — but it's additive, not multiplicative by a matrix. That breaks the repeated multiplication chain.

Production mistake: engineers clip gradients after computing the loss but forget to clip after the backward pass. The right order: loss.backward(), then clip, then optimizer.step(). Torch's autograd does not clip automatically.

import torch

model = SimpleRNN(vocab_size=10000, embed_dim=256, hidden_size=256)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

# Training loop with gradient clipping
for batch in dataloader:
    x, y = batch  # x: (batch, seq_len), y: (batch, seq_len)
    logits, _ = model(x)
    loss = loss_fn(logits.view(-1, logits.size(-1)), y.view(-1))
    
    optimizer.zero_grad()
    loss.backward()
    
    # MUST clip after backward, before step
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
    
    optimizer.step()
    
    # Check gradient norm for debugging
    total_norm = 0
    for p in model.parameters():
        if p.grad is not None:
            total_norm += p.grad.norm().item() ** 2
    total_norm = total_norm ** 0.5
    if total_norm > 100:
        print(f"Warning: high gradient norm {total_norm:.2f}")

LSTM Gate Mechanics: Forget, Input, Output, and Cell State

An LSTM cell has four neural network layers that control the flow of information. The three gates — forget, input, output — each produce values between 0 and 1 via sigmoid activation. The cell state update is:

f_t = sigmoid(W_f [h_{t-1}, x_t] + b_f) i_t = sigmoid(W_i [h_{t-1}, x_t] + b_i) o_t = sigmoid(W_o [h_{t-1}, x_t] + b_o) C_tilde = tanh(W_C [h_{t-1}, x_t] + b_C) C_t = f_t C_{t-1} + i_t C_tilde h_t = o_t * tanh(C_t)

The forget gate determines what fraction of the previous cell state to keep. The input gate decides how much of the new candidate to add. The output gate controls what part of the cell state is exposed to the next layer.

A common practice is to initialise the forget gate bias to 1.0 (or a large positive number) so that the network starts in a state of 'remember almost everything'. This prevents catastrophic forgetting early in training. In production, if you see the model losing long-range dependencies, check the forget gate biases: if they've drifted below 0, reset them to 1 and freeze for the first 10 epochs.

import torch
import torch.nn as nn
import torch.nn.functional as F

class LSTMCell(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size
        
        # Combined weight matrices for gates and candidate
        # Input-to-hidden: 4 * hidden_size x input_size
        # Hidden-to-hidden: 4 * hidden_size x hidden_size
        self.W_ih = nn.Linear(input_size, 4 * hidden_size, bias=True)
        self.W_hh = nn.Linear(hidden_size, 4 * hidden_size, bias=False)
        
        # Initialise forget gate bias to 1.0
        nn.init.unsqueeze_(self.W_ih.bias.data[hidden_size:2*hidden_size], 1.0)
        
    def forward(self, x, state):
        h_prev, c_prev = state
        # Combine inputs
        gates = self.W_ih(x) + self.W_hh(h_prev)  # (batch, 4*hidden)
        # Split into forget, input, candidate, output
        f_gate, i_gate, c_tilde, o_gate = gates.chunk(4, dim=-1)
        
        f = torch.sigmoid(f_gate)
        i = torch.sigmoid(i_gate)
        c_hat = torch.tanh(c_tilde)
        o = torch.sigmoid(o_gate)
        
        c = f * c_prev + i * c_hat
        h = o * torch.tanh(c)
        return h, (h, c)

Visual Gate Logic Diagrams: Forget, Input, Output Flows

The three gates of an LSTM can be understood visually as decision points in the information flow. The diagram below shows the full LSTM cell with the forget gate (red), input gate (green), output gate (blue), and the cell state highway. The forget gate decides what to keep from the previous cell state, the input gate decides what new information to add, and the output gate decides what part of the cell state to output as the hidden state. Each gate is a sigmoid layer that outputs values between 0 (block) and 1 (pass through), followed by element-wise multiplication with the respective signal.

Forget Gate Flow: The previous hidden state and current input are concatenated and passed through a sigmoid layer. The output is multiplied element-wise with the previous cell state. If the gate outputs 0, the information is forgotten; if 1, it is fully retained.

Input Gate Flow: The same concatenated input passes through a sigmoid (input gate) and a tanh (candidate cell state). The tanh outputs values between -1 and 1. The input gate output is multiplied by the candidate, and the result is added to the cell state.

Output Gate Flow: The new cell state is passed through tanh and then multiplied by the output gate's sigmoid output. This produces the new hidden state, which also goes to the next timestep and the output layer.

Sequence Architecture Comparison: RNN, LSTM, GRU, Transformer

Each architecture for sequence modelling has a different trade-off between memory, speed, and gradient stability. The table below compares the four major architectures used in production systems today. Note that the Transformer replaces recurrence with self-attention, allowing parallel computation over all timesteps, but at quadratic cost in sequence length.

The Transformer's key advantage is that it does not share weights across timesteps, so gradients never vanish due to recursion. However, the quadratic self-attention makes it impractical for very long sequences without approximations (e.g., Longformer, Performer). For sequences under 512 tokens, the Transformer is now the default choice in NLP. For time series or streaming scenarios where recurrence is natural, LSTM and GRU remain competitive.

Production Rule of Thumb: If your sequence length < 100 and you need real-time inference, use GRU. If accuracy is paramount and you have GPU memory for attention, use a Transformer. If you are already using LSTM and it works, there is no urgent need to switch — but monitor per-length performance.

Peephole Connections and LSTM Variants

The original LSTM introduced by Hochreiter & Schmidhuber had peephole connections: the gates received the cell state directly, not just the hidden state. The forget gate becomes f_t = sigmoid(W_f * [C_{t-1}, h_{t-1}, x_t]). This allows the network to 'look into' the cell state when deciding to forget. In practice, peephole connections add parameters and often give marginal gains. They are rarely used in modern implementations because the standard LSTM already performs well for most tasks.

Other variants include the GRU (Gated Recurrent Unit) which merges the forget and input gates into a single update gate and removes the separate cell state. GRU has fewer parameters and trains faster but can struggle on very long sequences. The bidirectional LSTM (BiLSTM) runs two LSTMs in opposite directions and concatenates their outputs — essential for tasks where context from both past and future matters, like named entity recognition.

Production tip: If you need speed and your sequence length is moderate (<100), use GRU instead of LSTM. At batch size 64 with sequence length 50, GRU is ~20% faster in training. If you need maximum accuracy and have long sequences, use LSTM with peephole connections and gradient clipping.

A common mistake is stacking too many LSTM layers. Three layers are often enough; beyond that, the gradient decays through the vertical dimension, not just the temporal. Use residual connections between layers to help.

import torch.nn as nn

# Standard LSTM
lstm = nn.LSTM(input_size=256, hidden_size=256, num_layers=2, batch_first=True)

# Bidirectional LSTM
bilstm = nn.LSTM(input_size=256, hidden_size=256, num_layers=2, 
                 bidirectional=True, batch_first=True)
# Output hidden size will be 512 (256*2) for each direction

# GRU (approx same params as LSTM with hidden=256)
gru = nn.GRU(input_size=256, hidden_size=256, num_layers=2, batch_first=True)

# For peephole, you need a custom cell or use a library like `torch-rnn`
# PyTorch's built-in LSTM does not support peephole connections out of the box

Scheduled Sampling: Bridging the Train/Inference Gap

Teacher forcing during training feeds the ground-truth token as input to the next timestep. At inference, the model must use its own predictions. This mismatch, called exposure bias, causes the model to never learn to correct its own errors. Scheduled sampling gradually reduces the probability of using ground-truth tokens during training, forcing the model to learn from its own outputs.

A simple schedule: start with a high probability (e.g., 1.0) of teacher forcing, and decay it linearly to 0.0 over training. For example, decay from 1.0 to 0.2 over 100k steps, then hold at 0.2. The code below implements scheduled sampling for a sequence-to-sequence model using a step-based schedule.

Production note: Scheduled sampling can destabilise training if the schedule is too aggressive. Start with a small decay per step (e.g., 1e-5) and monitor the loss. If the loss spikes, slow down the decay or use a curriculum where short sequences are teacher-forced longer.

import torch
import torch.nn.functional as F

def scheduled_sampling(model, input_tokens, target_tokens, step, total_steps, teacher_forcing_start=1.0, teacher_forcing_end=0.2):
    """
    Apply scheduled sampling during training.
    
    Args:
        model: nn.Module that takes (input_tokens, prev_token) and returns logits.
        input_tokens: (batch, seq_len) ground-truth input sequence (e.g., source).
        target_tokens: (batch, seq_len) ground-truth target sequence.
        step: current training step.
        total_steps: total scheduled sampling steps.
    Returns:
        loss: scalar loss for this batch.
    """
    batch_size, seq_len = target_tokens.shape
    
    # Compute teacher forcing probability (linear decay)
    ratio = min(1.0, step / total_steps)  # 0 at start, 1 at end
    teacher_prob = teacher_forcing_start + (teacher_forcing_end - teacher_forcing_start) * ratio
    
    # Start with first token (target_tokens[:, 0] is start token)
    prev_token = target_tokens[:, 0]
    total_loss = 0.0
    
    for t in range(1, seq_len):
        logits = model(input_tokens, prev_token)  # (batch, vocab_size)
        loss = F.cross_entropy(logits, target_tokens[:, t], reduction='sum')
        total_loss += loss
        
        # Decide whether to use teacher forcing or model prediction
        if torch.rand(1).item() < teacher_prob:
            prev_token = target_tokens[:, t]
        else:
            prev_token = torch.argmax(logits, dim=-1)  # greedy decode
    
    return total_loss / (seq_len - 1)  # average over steps

# Usage in training loop:
# for step, (src, tgt) in enumerate(dataloader):
#     loss = scheduled_sampling(model, src, tgt, step, total_steps=100000)
#     loss.backward()
#     optimizer.step()

Training and Production Pitfalls: What Actually Breaks Your Model

The most common production failure for sequence models is training/inference mismatch. During training, you often use teacher forcing: the true previous token is fed as input. At inference, you use the model's own predicted token. This discrepancy compounds errors quickly, especially in RNNs where a single wrong token at step 5 can derail the remaining 95 steps.

Another silent killer: padding and masking. If you pad sequences to equal length, the RNN will process padding tokens and produce meaningless outputs. You must apply a mask to the loss so that padding positions contribute zero gradient. PyTorch's nn.utils.rnn.pack_padded_sequence handles this, but many engineers forget to also mask the loss.

Memory management: LSTMs maintain a state for every sequence in the batch. If you have long sequences and large batches, you can run out of GPU memory. Use gradient checkpointing to trade compute for memory. For very long sequences, truncate backpropagation: after each segment, detach the hidden state graph to avoid storing the full computation graph.

import torch
import torch.nn.functional as F

def masked_cross_entropy(logits, targets, lengths):
    """Compute cross-entropy loss ignoring padding positions."""
    # logits: (batch, max_seq_len, vocab_size)
    # targets: (batch, max_seq_len)
    # lengths: (batch,) integer sequence lengths
    
    max_len = logits.size(1)
    # Create mask of shape (batch, max_len) — 1 for valid positions, 0 for padding
    mask = torch.arange(max_len, device=logits.device).unsqueeze(0) < lengths.unsqueeze(1)
    mask = mask.float()
    
    # Compute loss per token
    loss = F.cross_entropy(logits.permute(0, 2, 1), targets, reduction='none')
    # Apply mask
    loss = loss * mask
    # Average over valid positions only
    return loss.sum() / lengths.sum()

# Usage in training loop:
# logits, _ = model(x_padded, lengths)  # x_padded is already packed
# loss = masked_cross_entropy(logits, y_padded, lengths)

Implementing LSTM in Keras/TensorFlow

While PyTorch dominates research, TensorFlow's Keras API is still widely used in production for its ease of deployment via TF Serving and TF Lite. The Keras LSTM layer handles sequence masking, state management, and batching internally. The key differences from PyTorch: Keras uses a Masking layer or mask_zero=True in Embedding, and the LSTM layer accepts return_sequences=True/False.

Below is a minimal LSTM model for language modelling using TensorFlow/Keras. The model is identical in architecture to the PyTorch version but uses Keras' functional API. Note that Keras layers automatically determine input shapes from the data, but we specify input_shape for clarity. Also, Keras' Embedding layer can mask padding tokens if mask_zero=True is set, which propagates through the LSTM and can be used to ignore padding in the loss (we apply a manual mask in the custom training loop as a best practice).

import tensorflow as tf
from tensorflow.keras.layers import Embedding, LSTM, Dense, Masking
from tensorflow.keras.models import Sequential

vocab_size = 10000
embed_dim = 256
hidden_size = 256
max_seq_len = 50

model = Sequential([
    Embedding(input_dim=vocab_size, output_dim=embed_dim, mask_zero=True),
    LSTM(units=hidden_size, return_sequences=True, dropout=0.2),
    Dense(vocab_size, activation='softmax')
])

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

# Custom training loop with masking (alternative to built-in masking)
# Assuming padded dataset: (batch, seq_len) int tensors, (batch,) lengths
optimizer = tf.keras.optimizers.Adam()
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False)

def train_step(x_batch, y_batch, lengths):
    with tf.GradientTape() as tape:
        logits = model(x_batch, training=True)  # (batch, seq_len, vocab)
        # Build mask
        mask = tf.sequence_mask(lengths, maxlen=max_seq_len, dtype=tf.float32)  # (batch, seq_len)
        loss = loss_fn(y_batch, logits, sample_weight=mask)
    grads = tape.gradient(loss, model.trainable_variables)
    # Gradient clipping in TF
    grads, _ = tf.clip_by_global_norm(grads, clip_norm=5.0)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

# Example usage:
# for batch in dataset:
#     loss = train_step(batch['source'], batch['target'], batch['length'])

# For inference, use model.predict() or export with tf.saved_model