RNN Architecture › Practical Engineering Considerations

Building an RNN for a real task involves several engineering decisions beyond the core architecture.

Real text comes in variable lengths. A naive approach is to pad all sequences in a batch to the maximum length, but this wastes computation on padding tokens and can introduce noise into the gradients (if the loss is accidentally computed over padding positions).

PyTorch handles this through packed sequences. The torch.nn.utils.rnn.pack_padded_sequence function rearranges the batch into a compact representation that skips padding tokens during the RNN forward pass. After the RNN, pad_packed_sequence reverses this packing. The key inputs are the sequence tensor (padded to maximum length) and the actual lengths of each sequence.

When using packed sequences, sort the sequences by length in descending order before packing. PyTorch's pack_padded_sequence requires this (unless you set enforce_sorted=False). The length vector tells the packing routine when to transition between sequences.

For the loss computation, use a loss mask to ignore positions corresponding to padding tokens. The simplest way is to set ignore_index in nn.CrossEntropyLoss to the padding token's index, which instructs the loss function to skip those positions.

Gradient clipping is almost universally used when training RNNs. The exploding gradient problem, where the gradient norm grows exponentially across time steps, can cause catastrophic weight updates that destabilize training.

The clip_grad_norm_ function in PyTorch clips the global gradient norm: if the $L^2$ norm of all gradient vectors exceeds max_norm, each gradient is scaled down proportionally so that the total norm equals max_norm. This is different from clipping each gradient individually (which would change the gradient direction). Global norm clipping preserves the relative directions of the gradient components.

Choosing the right max_norm requires some tuning. A common heuristic is to log the gradient norm during training (before clipping) and set max_norm to roughly the 90th to 95th percentile of the unclipped norms. If gradients are clipped on most steps, the learning rate may be too high; if they are never clipped, the threshold may be unnecessary. The Part X chapter on Gradient Clipping covers this in more detail.

The DataLoader in our implementation shuffles the training windows randomly at each epoch. This works well for independent windows (where the sequence context starts fresh at each window boundary), but it is the wrong approach when you want to carry hidden state across windows (as in stateful BPTT).

For stateful processing, you should arrange the data into non-overlapping chunks and ensure that consecutive batches correspond to consecutive positions in the original sequence. This requires a custom Dataset and sampler that does not shuffle across sequence boundaries. The hidden state from the end of each batch is then detached from the computation graph (using .detach()) and passed as the initial hidden state for the next batch. The .detach() is important: it prevents gradients from flowing backward across the chunk boundary, implementing truncated BPTT.

The softmax computation in the output layer can suffer from numerical instability when logits are large. PyTorch's nn.CrossEntropyLoss combines log_softmax and the negative log-likelihood in a numerically stable way (using the log-sum-exp trick internally), so you should always use CrossEntropyLoss rather than manually computing softmax and then taking the log.

Similarly, avoid using float16 (half precision) for vanilla RNNs unless you have tested carefully. The relatively small values produced by $\tanh$ can underflow in float16 for large hidden sizes, causing training instability. LSTM and GRU with careful numerical design handle mixed precision better.