RNN/Sequence Concepts (LSTM/GRU, Masking, Embeddings)

This lesson on RNN/Sequence Concepts (LSTM/GRU, Masking, Embeddings) is hands-on and example-driven. You will master the fundamental building blocks of deep sequence modeling, including token embeddings, recurrent memory gating in LSTMs and GRUs, and sequence padding. You will learn how to implement masking strategies to process variable-length batches without corrupting recurrent states or loss calculations.

What You'll Be Able To Do

• Map discrete token indices into dense continuous vector spaces using PyTorch nn.Embedding layers.
• Differentiate the internal gating equations and state-tracking mechanics between standard RNNs, LSTMs, and GRUs.
• Implement sequence padding and attention/loss masking to handle variable-length inputs in mini-batches.
• Pack and unpack variable-length sequences using PyTorch recurrent utility functions to skip computation over pad tokens.
• Apply reduction masks to cross-entropy loss calculations to prevent padding tokens from skewing gradient updates.

Detailed Concept Walkthrough

1. Token Embeddings and Continuous Representations

An embedding layer acts as a trainable lookup table that maps high-cardinality discrete categorical indices into low-dimensional continuous vector spaces. This projects sparse, one-hot token representations into semantically rich geometric vectors where semantic similarity correlates with geometric proximity.

• Mechanism: The embedding layer stores a weight matrix of shape (num_embeddings, embedding_dim) and retrieves rows by integer index rather than performing expensive matrix-vector multiplications against one-hot vectors.
• Under the Hood: During backpropagation, gradients only flow into the specific row vectors accessed during the forward pass, utilizing sparse gradient updates if configured to conserve memory bandwidth.
• Best Practice: Always reserve and document dedicated index IDs for special tokens such as <pad> (typically index 0) and <unk> (unknown tokens) to ensure deterministic batch alignment.

import torch
import torch.nn as nn

# Define vocab size of 1000 and 64-dimensional embedding vectors
vocab_size, embed_dim = 1000, 64
embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embed_dim, padding_idx=0)

# Batch of 2 sequences, each of length 4 token IDs
token_ids = torch.tensor([[12, 45, 9, 0], [78, 3, 0, 0]], dtype=torch.long)
embedded = embedding(token_ids)  # Output shape: [2, 4, 64]

Key Takeaway: Embeddings convert discrete categorical symbols into dense, differentiable vectors optimized via standard backpropagation.

2. Recurrent Units: LSTM and GRU Gating Mechanisms

LSTMs and GRUs solve the vanishing gradient problem of standard RNNs by introducing additive linear state updates controlled by multiplicative non-linear gates. This allows gradients to backpropagate across hundreds of time steps without exponentially decaying to zero.

• Mechanism: An LSTM maintains an explicit cell state (c_t) for long-term memory alongside a hidden state (h_t), regulated by forget, input, and output gates that modulate information flow via sigmoid and tanh activations.
• Under the Hood: The GRU simplifies the LSTM architecture by merging cell and hidden states into a single hidden vector h_t, controlling updates using only two gates: reset and update gates.
• Execution Flow: For every timestep t, the recurrent cell computes gate activations from the current input x_t and previous state h_{t-1}, computing an additive update to retain stable gradient highways.

import torch
import torch.nn as nn

# LSTM: takes input_size=64, hidden_size=128
lstm = nn.LSTM(input_size=64, hidden_size=128, batch_first=True)

# Dummy embedded sequence: [batch_size=2, seq_len=4, embed_dim=64]
x = torch.randn(2, 4, 64)
out, (h_n, c_n) = lstm(x)
# out shape: [2, 4, 128], h_n shape: [1, 2, 128], c_n shape: [1, 2, 128]

Key Takeaway: Gating mechanisms create additive linear gradient highways that enable models to capture long-range temporal dependencies.

3. Sequence Padding and PackPaddedSequence

Because neural network batch computation requires rectangular tensors, uneven sequences must be padded with placeholder values to match the longest sequence. Specialized utility wrappers bypass padded steps during recurrent computation to maintain runtime efficiency and state validity.

• Mechanism: Sequences within a batch are padded with zeros to match max_len, and their unpadded valid lengths are tracked in a 1D tensor.
• Under the Hood: nn.utils.rnn.pack_padded_sequence compresses the rectangular padded tensor into a flat continuous data buffer with batch-size step pointers, skipping computation on trailing pad tokens.
• Best Practice: Always specify batch_first=True consistently across both your RNN modules and packing/unpacking calls to avoid shape mismatches.

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

# Embedded sequences of variable lengths (lengths: 4 and 2)
x = torch.randn(2, 4, 64)
lengths = torch.tensor([4, 2], dtype=torch.int64)

# Pack padded batch
packed_x = pack_padded_sequence(x, lengths.cpu(), batch_first=True, enforce_sorted=False)
lstm = nn.LSTM(64, 128, batch_first=True)
packed_out, (h_n, c_n) = lstm(packed_x)

# Unpack back to regular padded tensor
out, output_lengths = pad_packed_sequence(packed_out, batch_first=True) # [2, 4, 128]

Key Takeaway: Packing sequences prevents recurrent units from needlessly computing state transitions over meaningless padding tokens.

4. Masking in Attention and Loss Functions

Masking applies a boolean or binary filter over sequence positions to explicitly neutralize the effect of padding tokens in subsequent computations such as attention weights or cross-entropy loss. It guarantees that dummy padding values do not bias gradient updates or model outputs.

• Mechanism: In loss computation, padding token positions are zeroed out or assigned ignore_index so that non-informative tokens contribute zero error signal.
• Under the Hood: In attention mechanisms, padding positions are set to -inf before softmax normalization, ensuring their resulting attention probabilities evaluate to strictly zero.
• Syntax Rule: PyTorch loss functions like nn.CrossEntropyLoss accept an ignore_index argument (commonly ignore_index=0), automating loss masking across padded targets.

import torch
import torch.nn as nn

# Setup loss ignoring padding index 0
criterion = nn.CrossEntropyLoss(ignore_index=0)

# Logits: [batch_size=2, num_classes=10, seq_len=3]
logits = torch.randn(2, 10, 3, requires_grad=True)
# Targets with pad token 0 at sequence tails
targets = torch.tensor([[5, 8, 0], [3, 0, 0]], dtype=torch.long)

loss = criterion(logits, targets)
loss.backward()

Key Takeaway: Masking guarantees that synthetic padding tokens contribute zero weight to loss calculations and contextual attention distributions.

Topics Covered in RNN/Sequence Concepts (LSTM/GRU, Masking, Embeddings)

1. Sequence Modeling Overview (0:00 - 1:30) — Introduction to sequence structures and challenges with variable-length temporal data.
2. Token Embeddings (1:30 - 3:15) — Explanation of vector representations and PyTorch nn.Embedding lookup tables.
3. LSTM and GRU Gating (3:15 - 5:30) — Analysis of recurrent cell equations, memory highways, and gating mechanisms.
4. Padding and Packing Utilities (5:30 - 7:15) — Walkthrough of padding batches and using pack_padded_sequence to skip invalid tokens.
5. Masking and Loss Computation (7:15 - 9:00) — Techniques for masking attention and ignoring pad indices in objective functions.

DL Basics Cheat Sheet

• nn.Embedding(num_embeddings, embedding_dim, padding_idx=0) — Initializes lookup table for token IDs with zeroed padding vector

  emb = nn.Embedding(5000, 128, padding_idx=0)
  

• nn.LSTM(input_size, hidden_size, batch_first=True) — Creates long short-term memory recurrent layer accepting [B, T, D]

  lstm = nn.LSTM(128, 256, batch_first=True)
  

• nn.GRU(input_size, hidden_size, batch_first=True) — Creates gated recurrent unit layer accepting [B, T, D]

  gru = nn.GRU(128, 256, batch_first=True)
  

• pack_padded_sequence(input, lengths, batch_first=True, enforce_sorted=False) — Packs padded variable-length tensor into contiguous sequence structure

  packed = pack_padded_sequence(x, lens, batch_first=True, enforce_sorted=False)
  

• pad_packed_sequence(packed_input, batch_first=True) — Unpacks packed sequence object back to rectangular padded tensor

  padded_out, out_lens = pad_packed_sequence(packed_out, batch_first=True)
  

• nn.CrossEntropyLoss(ignore_index=0) — Computes cross entropy while ignoring specified pad target index

  criterion = nn.CrossEntropyLoss(ignore_index=0)
  

Comparison Table

Common Pitfalls

• Mistake: Allowing padding tokens to contribute to loss calculations. Avoid: Set ignore_index to your padding ID in CrossEntropyLoss.
• Mistake: Passing padded sequence zeros through an unmasked RNN. Avoid: Use pack_padded_sequence or slice valid sequence lengths.
• Mistake: Forgetting batch_first=True when tensors are shaped [Batch, Time, Dim]. Avoid: Explicitly pass batch_first=True to RNN modules and packing functions.
• Mistake: Initializing embedding padding index with random trainable weights. Avoid: Supply padding_idx to nn.Embedding to freeze pad vectors at zero.

FAQs

• Why use an Embedding layer instead of a Linear layer with one-hot inputs? An embedding layer performs direct indexing O(1) in memory rather than allocating huge sparse matrices and running O(V * D) matrix multiplications.
• What is the key functional difference between LSTM and GRU? LSTMs maintain separate cell and hidden states using three gates, whereas GRUs combine cell and hidden states using only two gates, reducing computation.
• How does enforce_sorted=False affect pack_padded_sequence? It allows passing batches where sequence lengths are not pre-sorted in descending order, at the expense of a minor internal sort-and-restore overhead.