Padding, Packing, and Masking for Variable-Length Sequence Batches

Tensor operations over a batch need a rectangular layout, and sequences in a corpus do not have the same length. A batch of $B$ sequences must therefore be forced into a tensor of shape $B \times T \times n$ for some $T$, and the only question is which positions hold real data and how the rest of the computation is told to ignore the others. Getting that bookkeeping wrong does not usually produce an exception; it produces a model that trains on a slightly different objective than intended, which is why the details are worth stating precisely.

The rectangular batch and what padding wastes

The straightforward approach pads every sequence to the length of the longest member of the batch, using a reserved index whose embedding exists but is never a target. If lengths are drawn independently and uniformly from $1$ to $L$, the expected fill ratio of a batch padded to its maximum is around one half in the large-batch limit, so roughly half the arithmetic in a padded batch is spent on positions that contribute nothing. Sequence lengths in real corpora are not uniform, and the waste is worse when a single long sequence sets the bound for a batch of short ones.

Sorting batches by length, or bucketing by length quantile, raises the fill ratio because the maximum of each batch grows more slowly than the maximum of the whole corpus. The cost is a correlation between batch composition and content, which matters if length correlates with a label and the label needs balanced batches, and it interacts with shuffling: the batch order can still be shuffled even when sorting happens within a batch-forming window. An alternative is a token budget: form a batch by adding sequences until the product of batch size and current maximum length would exceed a fixed budget, which bounds the arithmetic per batch directly instead of bounding the count of sequences.

The two layouts are not interchangeable across layer types. A recurrence processes positions strictly in order, so a packed layout is natural there and removes real arithmetic; a layer that mixes positions through a weight matrix over the time axis requires the rectangular layout, because the mixing operation is a dense product across that axis and the mask is the only thing standing between padding and the result. A model that interleaves the two has to convert between layouts at the boundary, and the conversion is where the state and the mask must agree about which positions are real.

A batch that mixes very short and very long sequences is the worst case for padding, and it is also the common case for corpora whose lengths are long-tailed. Padding each batch to its own maximum bounds the waste by the spread of lengths inside that batch instead of the spread across the corpus, which is the whole effect of bucketing: a batch drawn from a narrow band of lengths has a fill ratio near one, while a batch drawn from the full range can spend most of its arithmetic on filler. The same reasoning applies to the memory bound on a batch, which is set by the batch size multiplied by the maximum length in it rather than by the total number of tokens.

What must be masked, and with what

Padding enters the computation at several places, and each needs its own treatment. A mask of shape $B \times T$ with true at real positions is enough for all of them.

The loss is the first place. Summing a per-position cross-entropy over the padded tensor and dividing by $B \cdot T$ understates the loss and, worse, gives padded positions weight one while real positions share their weight with filler. Multiplying the per-position loss by the mask and dividing by the number of real positions gives the same value as training each sequence separately. The same accounting applies to a decoder’s target sequence, which is shifted by one position relative to its input and masked for the same reasons, as described in encoder-decoder models and teacher forcing. Any per-batch normalization, including the denominator of a mean-squared error and the statistics of a batch-normalized layer, needs the same correction, and the resulting differences are small enough to be invisible until the model is evaluated.

The last-step readout is the second. For a task whose label belongs to the whole sequence, the state that matters is the one at step $length_i$, not at step $T$, so the state is selected per sequence by indexing the output tensor at index $length_i - 1$ for each row. Because an LSTM or GRU layer produced no meaningful value at padded steps, the choice does not corrupt anything, but any pooling over positions (mean, max) has to be masked likewise, and a global max over padded positions would pick up whatever values the padded steps produced.

Attention over the padded axis is the third, and it is the one with a numerical trap. Padding positions must receive zero attention weight. Adding a very negative constant to the scores before the softmax achieves that: the exponential of the offset is zero in exact arithmetic and below the smallest normal float in practice, so the resulting weights are exactly zero. Multiplying weights by the mask after the softmax does not achieve the same thing, because the row then no longer sums to one and the attended vector is scaled down by however much padding was in the row. Using an actual negative infinity is also hazardous, since a row whose real positions are all masked produces the undefined $0/0$ in the softmax; a large finite negative number sized to the dtype avoids both problems.

import numpy as np

def masked_softmax(scores, mask, neg=1e9):
    """scores, mask: (B, T). Padded positions get exactly zero weight."""
    s = np.where(mask, scores, -neg)          # mask before the softmax, not after
    s = s - s.max(axis=-1, keepdims=True)     # shift for numerical stability
    e = np.exp(s)
    return e / e.sum(axis=-1, keepdims=True)

def masked_mean(values, mask):
    """values: (B, T, d), mask: (B, T). Returns (B, d)."""
    m = mask[..., None].astype(values.dtype)
    return (values * m).sum(axis=1) / np.maximum(m.sum(axis=1), 1.0)

The division by the count of real positions rather than by $T$ is the part that is easy to omit, and it is the difference between a per-sequence average and a length-weighted one.

Packed sequences

Padding can be removed entirely by flattening a batch into a single ordered run of its real positions. Padding sequences of total length $N_{\text{total}}$ in a batch of size $B$ gives a tensor of shape $1 \times N_{\text{total}} \times n$ for zero-padded input, together with an index per step naming the sequence it belongs to and a list of the original lengths. The recurrence then steps over $N_{\text{total}}$ positions instead of $B \cdot T$, so the arithmetic is exactly proportional to the number of real symbols.

The complication is the state. After each step the state of each active sequence must be assembled into the batch dimension of the step and the state of the sequences that ended must be removed and stored, which is a gather and a scatter per step rather than a single block operation. The gain is proportional to the padding that would otherwise be computed, which for highly variable lengths is close to a factor of two. Frameworks expose the layout as a packed representation whose interface requires the sequences to be ordered by decreasing length, since that ordering is what makes the per-step slicing valid; sorting the batch and restoring the original order afterwards is part of the contract. Because length-sorted batches produce gradient estimates that are correlated with length, many implementations shuffle the batch order and resample the sort each epoch.

The packed layout also constrains where a mask is still required. Steps within one sequence are contiguous, so no mask is needed inside the recurrence, but any operation that mixes across the flattened axis, including attention over a batch, needs the boundaries recovered from the index tensor.

Failure modes that survive correct batching

Three mistakes recur even when the mask is constructed correctly. The first is applying the mask to the loss but not to a regularizer or an auxiliary head, so the padded positions still contribute gradient through a secondary term. The second is detaching the carried state at the wrong place when a sequence is split across batches, which either leaks gradient between unrelated sequences or loses the continuity that the split was meant to preserve. The third is a length that exceeds the model’s positional capacity or the chunk size the recurrence was trained with, which turns inference into an out-of-distribution condition even though the tensor shapes are all legal.

A subtler issue is that padding changes the effective batch statistics of every normalization layer. Statistics computed over padded positions differ from statistics over real ones, and the effect is largest for the shortest sequences in a batch. Masked statistics fix it, and the correction has to be applied consistently in training and evaluation, which is the same requirement that appears for inference-time state handling in what a recurrent cell computes. Sequence lengths that vary by an order of magnitude within one batch make this more than a rounding concern, and bucketing by length is what keeps the correction small.