Packed LSTM sequences and the last valid state

Last updated: 6 Oct 20265 min read

tutorial

Packing makes an LSTM stop at each record’s real length; the final state must represent the last event, not an allocated padding slot.

Make length a first-class input

A batch of event histories has shape batch by allocated steps by features, but each history also needs a measured count of real steps. Padding is storage, not an event. Reject zero-length histories or handle them with a separate fallback; a packed recurrent call expects a positive length. Verify that every length is at most the allocated width and corresponds to the feature rows produced by the same tokenizer or event encoder. The tensor contract should include both the data shape and the length vector.

Pack before recurrence

A packed sequence lets the recurrent cell skip padding timesteps. When inputs are not sorted by length, request internal sorting and let the API restore output order. If lengths are provided as a tensor, put them on CPU for the packing call. The final hidden state from a one-layer, forward-only LSTM then belongs to the last real step of each original sequence. Reading output at the last allocated index instead would mix in padding behavior or return a zero after unpacking. Padding-aware loss is a separate requirement when every step receives a target.

Keep directions and layers explicit

For a one-layer, forward-only LSTM, the last hidden slice and an output gathered at length minus one agree within floating-point tolerance. That shortcut changes for a bidirectional model: its backward final state summarizes a reverse traversal, and output at the last forward position does not contain the backward final summary. A stacked model adds another layer axis. Name the chosen direction and layer before selecting a state; do not treat the first dimension of the hidden tensor as the batch dimension.

Separate batch padding from online state

Packing is useful when independent histories are processed together. A live stream may instead carry hidden and cell state from one request to its next event. That state must be keyed by stream identity and reset when a stream ends, changes owner or times out. A packed batch does not solve state ownership across requests. State reset rules cover that boundary and the gradient history retained during training.

Probe a padding counterfactual

Create two batches with identical valid prefixes but different values in every padded slot. Their final states and predictions should match when lengths are correct. Change one length deliberately and confirm the test fails. Check a one-event history, an unsorted batch and a last batch smaller than the usual batch size. Keep the probe in model tests because a preprocessing change can silently replace the length vector while leaving tensors numerically well formed.

Implementation

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

torch.manual_seed(47)
service_histories = torch.tensor([
    [[0.2, 0.4, 0.1], [0.5, 0.2, 0.3], [0.1, 0.7, 0.2], [0.3, 0.2, 0.8]],
    [[0.9, 0.1, 0.2], [0.4, 0.6, 0.1], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
    [[0.3, 0.5, 0.2], [0.2, 0.1, 0.8], [0.6, 0.3, 0.4], [0.0, 0.0, 0.0]]])
history_lengths = torch.tensor([4, 2, 3], dtype=torch.long)
assert bool(((history_lengths > 0) & (history_lengths <= 4)).all())

event_lstm = nn.LSTM(input_size=3, hidden_size=7, batch_first=True)
packed_histories = pack_padded_sequence(service_histories, history_lengths.cpu(),
                                        batch_first=True, enforce_sorted=False)
packed_outputs, (last_hidden, last_cell) = event_lstm(packed_histories)
event_outputs, restored_lengths = pad_packed_sequence(
    packed_outputs, batch_first=True, total_length=4)
last_valid = event_outputs[torch.arange(3), history_lengths - 1]
torch.testing.assert_close(last_valid, last_hidden[-1])
assert torch.equal(restored_lengths, history_lengths)
assert last_cell.shape == (1, 3, 7)

Performance and operating cost

For sequence length T and hidden width H, a recurrent layer performs work across T dependent steps; four LSTM gates increase per-step matrix work relative to a plain recurrent cell. Packing avoids cell work on padded positions, although sorting and packing add overhead and may limit some kernel choices. Activations stored for backward scale with valid steps and hidden width. Measure valid events per second, not allocated slots per second.

Common Mistakes

Do not use the final allocated padding slot as the last event.
Do not treat a bidirectional final state as though it were a one-direction output at length minus one.
Do not accept zero or incorrect lengths just because the padded tensor has the right shape.