10.2 RNN, LSTM, and BPTT: State Can Carry Context, But Also Leak

The logs in the Prediction Chamber are no longer isolated lines, each second of a task must be evaluated in relation to the prior nine. In the Model Workshop, the same update rule must be reused across time, while maintaining a finite-dimensional state.

The strength and risk of RNNs both stem from the same source: the hidden state. It can summarize the past, but it can also carry forward outdated information (like the identity of a previous user, future intentions, or padding artifacts) into contexts where they shouldn’t belong.

Learning Objectives

• Write the vanilla RNN recurrence and corresponding tensor shapes;
• Understand BPTT and gradient vanishing/explosion from an unrolled computation graph;
• Derive the LSTM cell dynamics, hidden state, and gate computations;
• Properly handle padding, sequence length, bidirectional processing, and streaming state;
• Distinguish between sequence classification, token prediction, and text generation.

1. RNN Shares Parameters Across Time

For input $x_t$ and the previous state $h_{t-1}$:

$$ h_t=\phi(W_{xh}x_t+W_{hh}h_{t-1}+b_h), $$

$$ o_t=W_{ho}h_t+b_o. $$

The same weights $W_{xh}$ and $W_{hh}$ are reused across all time steps $t$, enabling handling of variable-length sequences and decoupling the number of parameters from sequence length. This design assumes that the update rule is time-invariant; when strong seasonal or phase-dependent patterns exist, explicit time features or alternative architectures are still required.

h_t represents a compressed state and does not guarantee preservation of all historical information. The retained content depends on hidden size, training objective, and gradient flow.

2. Shape Contract

Common at batch_first=True:

input:  [B, T, input_size]
output: [B, T, directions * hidden_size]
h_n:    [layers * directions, B, hidden_size]

batch_first Doesn't change the layout of the hidden state. In multi-layer/bidirectional setups, directly accessing h_n[-1] might only retrieve the last layer in a specific direction; you must reshape according to [layers, directions, B, H] and explicitly select or concatenate.

Embedding inputs are typically long token IDs [B,T], outputs [B,T,E]. Padding ID must match the vocabulary/embedding configuration.

3. BPTT is Backprop on an Unrolled Graph

After unrolling the recurrence $T$ steps, the loss with respect to an early hidden state contains a Jacobian product:

$$ \frac{\partial L}{\partial h_t} =\sum_{k\ge t} \frac{\partial L_k}{\partial h_k} \prod_{j=t+1}^{k} \frac{\partial h_j}{\partial h_{j-1}}. $$

This product repeatedly multiplies the weight matrix $W_{hh}$, activation derivatives, and gate outputs. When typical singular values are less than 1, the gradient contribution decays; when greater than 1, it explodes. Gradient clipping helps prevent explosive updates but cannot recover long-term signal that has already vanished.

The longer the sequence length, the higher the activation memory and the backward computation cost. Truncated BPTT detaches the hidden state every few steps, performing backpropagation only within a finite window. This reduces computational cost and explicitly cuts off credit assignment across windows.

4. LSTM Provides a Controlled Addition Path to the Cell State

A common LSTM formulation:

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

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

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

$$ g_t=\tanh(W_g[x_t,h_{t-1}]+b_g), $$

$$ c_t=f_t\odot c_{t-1}+i_t\odot g_t, $$

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

The forget/input/output gates respectively regulate the retention of previous cell state, the candidate write-in, and the exposure of the output. The addition operation in the cell state allows certain gradients to flow more persistently through time, but gates can still saturate, meaning LSTM does not guarantee infinite memory or immunity to gradient vanishing problems.

GRUs merge portions of state and gates, reducing parameter count. The choice between them should be based on validation performance, latency, and data volume, not on a fixed notion that "LSTM is more advanced than RNN."

5. Padding Is Not Empty Information

Variable-length sequences are commonly padded to match the maximum batch length. If an RNN is run directly without special handling:

• Padding steps continue to update the hidden state;
• output[:, -1] might correspond to padding tokens, not the final valid token;
• If token loss is not masked, padding tokens are treated as training targets;
• Batch statistics or metrics may be skewed due to inclusion of padding.

The effective output can be gathered using lengths, or pack_padded_sequence can be used to skip over padding entirely. In PyTorch, when passing a lengths tensor to pack, it is typically moved to the CPU.

import torch
from torch import nn
from torch.nn.utils.rnn import pack_padded_sequence

class LSTMClassifier(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, classes, pad_id=0):
        super().__init__()
        self.embedding = nn.Embedding(
            vocab_size,
            embedding_dim,
            padding_idx=pad_id,
        )
        self.lstm = nn.LSTM(
            embedding_dim,
            hidden_dim,
            batch_first=True,
        )
        self.head = nn.Linear(hidden_dim, classes)

    def forward(self, token_id, length):
        embedded = self.embedding(token_id)
        packed = pack_padded_sequence(
            embedded,
            length.cpu(),
            batch_first=True,
            enforce_sorted=False,
        )
        _, (h_n, _) = self.lstm(packed)
        last_layer_hidden = h_n[-1]
        return self.head(last_layer_hidden)

If the model is bidirectional or multilayered, the forward pass must be reprocessed for h_n directions, direct reuse of the last row is not valid.

6. Bidirectional RNN Sees Into the Future

A bidirectional encoder processes the entire sequence both left-to-right and right-to-left, making it well-suited for offline tagging or classification tasks. However, it is unsuitable for causal online prediction or autoregressive generation, because the backward pass relies on future tokens.

Enabling bidirectional training during feature extraction inadvertently introduces temporal leakage. Before deploying, always verify whether the full sequence is available at the time of inference.

7. Many-to-one, Many-to-many, and Causal LM

• Sequence classification: the entire sequence outputs a single label;
• Token labeling: each valid position outputs a label;
• Forecasting: predicts future values based on past inputs;
• Autoregressive language models: predicts the probability distribution of the next token;
• Seq2seq: decoder state or outputs are conditioned on encoder states or outputs.

Loss placement differs. Supervision is applied only at the final step, resulting in longer information paths earlier in the sequence; step-wise supervision provides denser gradients but requires alignment with task labels.

Teacher forcing feeds the true previous token during decoder training, while during inference it feeds the model’s own output, creating an exposure mismatch. Techniques like scheduled sampling alter the objective, but no automatic solution exists.

8. The Boundary of Stateful Streaming

Stateful streaming RNNs can retain state across chunks, reducing redundant computations. However, it is essential to define when state should be reset:

• New users or new sessions;
• Device restarts or long idle periods;
• Changes in the streaming order within a batch;
• Model version updates;
• Backfill or out-of-order events.

When training truncated chunks, the state must be explicitly reset detach() to prevent the computation graph from growing indefinitely. State must be zeroed out or properly indexed between independent sequences. The state cache itself is a stateful service component and requires TTL (time-to-live), consistency, privacy, and fault recovery mechanisms.

9. RNN's Sequential Cost and Alternatives

Each $h_t$ depends on $h_{t-1}$, making the temporal dimension difficult to fully parallelize; however, batch and layer matrix operations can still be parallelized. For long sequence training, throughput may not match that of attention or temporal convolution models.

Alternative approaches:

• 1D or dilated causal convolution: fully parallelizable with a fixed receptive field;
• Transformers: enable global content-based interactions, though attention memory may incur secondary costs;
• State-space models: offer alternative long-sequence computation and state design paradigms;
• Feature aggregation combined with tabular models: may provide greater stability for small-data tasks.

RNNs still hold value in low-latency streaming scenarios, with limited state requirements, and for shorter sequence lengths.

10. Training and Evaluation

Records:

• Length distribution and length-bucket metrics;
• Padding fraction/pack throughput;
• Hidden/cell norm, gradient norm, and clip rate;
• State reset correctness;
• Causal cutoff and label maturity;
• Teacher-forced loss and free-running generation/forecast error.

Random row splitting may inadvertently place adjacent time windows from the same session into train and test sets. Splitting by entity, time, or horizon helps avoid highly overlapping window leakage.

Common Misconceptions

• RNNs remember the entire history: The hidden state is a finite compression of past inputs; training does not guarantee that target information is preserved.
• LSTMs solve gradient vanishing: They improve gradient flow, but do not provide a guarantee of infinite memory.
• output[:, -1] is the final effective state: This is often not true when padding is present.
• Bidirectional models are just stronger: In causal sequential tasks, they look into the future, which violates the causality assumption.
• Preserving state always improves continuity: Incorrect reuse of state across entities can lead to severe information leakage.

Exercise

1. Expand the 4-step RNN computation graph and label the shared parameters.
2. Verify the output shape of batch_first=True against h_n.
3. Compare the effects of padding on output[:,-1], length-based gather, and packed LSTM.
4. Implement truncated BPTT and observe the graph and memory differences before and after detach.
5. Define state keys, reset logic, and TTL rules for a streaming multi-user service.

Summary

RNNs push historical information into a hidden state through shared recurrence, and BPTT computes gradients along the unrolled time diagram. LSTM gates and cells improve long-range dependencies, but whether the model truly respects deployment causality depends on factors like padding, bidirectionality, truncation, and state lifecycle.

The next lesson removes the bottleneck of single recurrence: attention allows each query to directly read information from a set of keys and values, while introducing new challenges such as masking, quadratic matrix complexity, and interpretability boundaries.