Demystifying Packed Sequences: A Guide to Efficient RNN Processing in PyTorch

When working with sequences of varying lengths in neural networks, it's common to pad shorter sequences with a special value (e.g., 0) to make them all the same length. This allows them to be batched together and fed into an RNN. However, including padding tokens during RNN processing can distort the results because the RNN shouldn't process these irrelevant values.

PyTorch's Packed Sequence Solution

To address this, PyTorch offers the nn.utils.rnn.pack_padded_sequence function. It takes two arguments:

• Padded sequences: A tensor of shape (seq_len, batch_size, input_size), where seq_len is the maximum sequence length in the batch (including padding), batch_size is the number of sequences, and input_size is the dimensionality of each element in the sequence.
• Sequence lengths: A list or tensor of length batch_size containing the actual lengths of each sequence (excluding padding).

Internal Structure of a Packed Sequence

pack_padded_sequence returns a PackedSequence object, which is essentially a tuple of two elements:

1. Data: A tensor containing the elements of all sequences, interleaved by time step. This means elements from the first sequence at each time step are concatenated, followed by elements from the second sequence, and so on. Padding elements are excluded.
2. Batch sizes: A list or tensor of length max_seq_len (the same seq_len from the input). Each element represents the number of non-padded elements at that specific time step across the batch. This lets the RNN know how many valid elements to process at each step.

How RNNs Process Packed Sequences

When you pass a PackedSequence to an RNN in PyTorch, it performs the following steps:

1. Iterates through time steps: The RNN processes the data element-by-element, but it only considers the elements corresponding to non-padded sequences at each time step. This is guided by the batch sizes information within the PackedSequence.
2. Ignores padding: Padding elements are effectively skipped during the RNN's calculations, ensuring that only the relevant sequence information is used to update the hidden state.
3. Maintains hidden state: The RNN updates its hidden state based on the processed non-padded elements at each time step.

Unpacking the Packed Sequence (Optional)

After the RNN has processed the packed sequence, you might want to recover the individual sequences (without padding) and their hidden states. PyTorch provides the nn.utils.rnn.pad_packed_sequence function for this purpose. It takes the PackedSequence and the original batch sizes (before packing) as input and returns the following:

• Output: A tensor of shape (seq_len, batch_size, output_size), where output_size is the size of the RNN's output (typically the same as input_size). This tensor contains the RNN's output for each time step in the original sequences (excluding padding).
• Batch sizes (optional): If provided, the original batch sizes can be returned, identical to the input to pack_padded_sequence.

Benefits of Packed Sequences

• Efficiency: PyTorch's RNN implementation leverages the information in the packed sequence to efficiently process only the non-padded elements, leading to faster computation.
• Accuracy: By ignoring padding elements during processing, the RNN focuses on the actual content of the sequences, resulting in more accurate predictions.

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

# Sample sequences of different lengths
sequences = [torch.tensor([1, 2, 3]), torch.tensor([4, 5]), torch.tensor([6, 7, 8, 9])]
sequence_lengths = [len(seq) for seq in sequences]

# Pad sequences to the maximum length (optional, but recommended for batching)
padded_sequences = pad_sequence(sequences, batch_first=True, padding_value=0)

# Define an LSTM model
class MyLSTM(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(MyLSTM, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)

    def forward(self, packed_input):
        # Pack the padded sequences (assuming they are already sorted)
        packed_sequences = pack_padded_sequence(padded_sequences, sequence_lengths, batch_first=True)

        # Pass the packed sequence to the LSTM
        output, (hidden, cell) = self.lstm(packed_sequences)

        # Optionally unpack the output if needed
        # unpacked_output, _ = pad_packed_sequence(output, batch_first=True)

        return output, hidden, cell

# Create an instance of the LSTM model
model = MyLSTM(input_size=sequences[0].size(1), hidden_size=128)

# Run the model with the packed sequence
packed_output, hidden, cell = model(packed_sequences)

# Process the output and hidden state as needed
print(packed_output.shape)  # Output shape: (seq_len, batch_size, hidden_size)
print(hidden.shape)        # Hidden state shape: (num_layers, batch_size, hidden_size)

This code:

1. Defines sample sequences with varying lengths.
2. Pads the sequences (optional, but useful for batching).
3. Creates an LSTM model class (MyLSTM).
4. Defines a forward method that packs the padded sequences and passes them to the LSTM.
5. Unpacks the output (commented out as an optional step).
6. Creates an instance of the model and runs it with the packed sequence.
7. Prints the output and hidden state shapes.

• Involves creating a mask tensor of the same shape as the padded sequences. This mask has elements set to 1 for valid sequence positions and 0 for padding positions.
• During the RNN's forward pass, you can multiply the input sequences with the mask element-wise. This effectively sets the padding elements to zero before feeding them to the RNN.
• While masking achieves similar results to packed sequences, it can be less efficient computationally, especially for very long sequences or large batches.

Here's a code snippet illustrating masking:

# Create a mask tensor
mask = (padded_sequences != 0).float()  # Convert to float for multiplication

# Pass the masked input to the RNN
output, (hidden, cell) = model(padded_sequences * mask)

RNN with Variable Length Inputs:

• Some RNN implementations in PyTorch (e.g., nn.GRU) allow specifying variable-length sequences directly without packing.
• However, this approach might not be available for all RNN types (e.g., standard LSTM) and may have limitations in terms of batching efficiency. Refer to PyTorch documentation for supported functionalities.

Choosing the Right Method:

• Packed sequences are generally the most efficient and recommended approach for handling variable-length sequences in PyTorch RNNs due to their memory and computational advantages.
• Consider using masking if packed sequences are not supported by the specific RNN type you're using. However, be aware of potential efficiency drawbacks.
• RNNs with variable length inputs might be an option in specific scenarios, but check for limitations in batching and compatibility with your desired RNN architecture.