Embeddings and Sequence Representation

Embeddings and Sequence Representation is the bridge between symbolic sequence data and the tensors that PyTorch models can optimize. Words, product IDs, diagnoses, clicks, and characters are discrete labels; neural networks need dense numeric vectors with gradients. By the end of this lesson, you should be able to build an embedding layer, batch variable-length sequences, keep padding from contaminating learning, and choose a representation for pooling, recurrent models, or transformer-style attention.

Purpose and Outcome

A sequence model starts with an ordered list of items. The model does not understand the string learning or the category ID 42; it only receives integer indices. An embedding table maps each index to a trainable row vector. During training, backpropagation updates only the rows that were used in the batch, so tokens that appear in similar contexts can move toward useful locations in vector space. In this course, embeddings are the first step toward recurrent networks, attention, and transformers because they define the numeric representation that all later sequence layers consume.

The practical outcome is a tensor with shape usually written as [batch, time, embedding_dim]. batch counts examples, time counts positions in each sequence, and embedding_dim is the width of each learned vector. That tensor can be pooled into one vector per example, processed position by position by an RNN or GRU, or combined with positional information and passed to attention.

How Embeddings Work Internally

nn.Embedding is conceptually a lookup table, not a matrix multiplication over one-hot vectors. Its parameter is a weight matrix with shape [num_embeddings, embedding_dim]. When the input tensor contains integer IDs, PyTorch gathers the corresponding rows and appends the embedding dimension to the input shape. If the input IDs have shape [2, 4], the output has shape [2, 4, embedding_dim].

This lookup behavior matters for performance and learning. A one-hot representation for a vocabulary of 100,000 tokens would allocate a huge sparse vector for every token. An embedding stores only the dense table and gathers rows by integer index. In the backward pass, gradients are accumulated for rows that appeared in the batch. Common rows may receive frequent updates; rare rows may remain poorly trained unless the dataset, initialization, or pretraining strategy handles them.

Most sequence systems reserve special IDs. A padding token lets examples of different lengths share one rectangular batch. An unknown token handles items missing from the vocabulary. Start and end tokens can mark generation boundaries. The padding row is special because it should not behave like a real word or event. In PyTorch, padding_idx initializes that row to zeros and prevents it from receiving gradient updates through the embedding layer.

API Anatomy

Piece Meaning
num_embeddings Number of rows in the table, usually the vocabulary size including special tokens.
embedding_dim Width of each learned vector and the feature size passed to the next sequence layer.
padding_idx Index that represents padding and is kept fixed by the embedding layer.
input dtype Token IDs must be integer tensors, normally torch.long.
output shape The input shape plus one final dimension containing the selected vectors.

The usual workflow is tokenize or otherwise index the raw sequence, build a vocabulary with reserved IDs, convert each example to integer IDs, pad a batch to a common length, run the embedding layer, then pass the result to pooling, a recurrent layer, convolution, or attention. The model owns the embedding weights unless you explicitly load pretrained vectors and decide whether to freeze or fine-tune them.

Example 1: Token IDs to Vectors

The first example shows the core lookup. Two short sequences are padded to length three. The output shape confirms that the embedding layer preserves the batch and time dimensions while adding a vector dimension. The padding vector is zero because index 0 was declared as padding_idx. The final line prints True because no backward pass has happened yet.

import torch
import torch.nn as nn

torch.manual_seed(7)
vocab = {"<pad>": 0, "deep": 1, "learning": 2, "rocks": 3}
ids = torch.tensor([[1, 2, 3], [1, 0, 0]])
embedding = nn.Embedding(num_embeddings=len(vocab), embedding_dim=4, padding_idx=0)

vectors = embedding(ids)
print(vectors.shape)
print(vectors[1, 1])
print(embedding.weight.grad is None)

Expected output: torch.Size([2, 3, 4]), then a zero padding vector with grad_fn=<SelectBackward0>, then True.

Example 2: Masked Mean Pooling

A common baseline for text classification or sequence tagging features is to average token embeddings into one vector per example. The important detail is the mask. Averaging over padded positions silently changes the representation of shorter examples, especially when the padding vector is not exactly zero or when later layers add bias. This example multiplies by a boolean mask and divides by the true sequence length.

import torch
import torch.nn as nn

torch.manual_seed(3)
ids = torch.tensor([[1, 2, 3, 0], [4, 5, 0, 0]])
lengths = torch.tensor([3, 2])
embedding = nn.Embedding(6, 5, padding_idx=0)
vectors = embedding(ids)
mask = (ids != 0).unsqueeze(-1)
mean_vectors = (vectors * mask).sum(dim=1) / lengths.unsqueeze(-1)

print(mean_vectors.shape)
print(torch.isfinite(mean_vectors).all().item())

Expected output: torch.Size([2, 5]) and True.

The output is deterministic in shape and finiteness. The actual vector values depend on the random initialization. This representation discards order, so deep learning rocks and rocks learning deep can become very similar if they contain the same tokens. That trade-off is acceptable for some bag-of-words style baselines and inadequate for tasks where order changes meaning.

Example 3: Packed Sequences for a GRU

Recurrent layers read sequences step by step. If a batch is padded, the recurrent layer will otherwise process padding positions as if they were real time steps. pack_padded_sequence tells the GRU how many valid positions each example has. After processing, pad_packed_sequence restores a rectangular tensor for downstream code.

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

torch.manual_seed(11)
ids = torch.tensor([[1, 2, 3, 0], [4, 5, 0, 0]])
lengths = torch.tensor([3, 2])
embedding = nn.Embedding(6, 8, padding_idx=0)
gru = nn.GRU(input_size=8, hidden_size=6, batch_first=True)

embedded = embedding(ids)
packed = pack_padded_sequence(embedded, lengths.cpu(), batch_first=True, enforce_sorted=False)
packed_output, hidden = gru(packed)
output, restored_lengths = pad_packed_sequence(packed_output, batch_first=True)

print(output.shape)
print(hidden.shape)
print(restored_lengths.tolist())

Expected output: torch.Size([2, 3, 6]), torch.Size([1, 2, 6]), and [3, 2].

The restored time dimension is three because the longest true sequence has length three, even though the original padded input had width four. The hidden state has shape [num_layers, batch, hidden_size] for this one-layer GRU. Packing is often worthwhile for recurrent models because it avoids wasted computation and prevents hidden states from being shaped by padding.

Example 4: Token Plus Position for Attention

Self-attention compares positions directly and does not contain recurrence by itself. If the model sees only token embeddings, it has no built-in notion that one token came before another. A common representation adds a learned or fixed positional vector to each token vector. A padding mask then tells attention layers which keys are padding.

import torch
import torch.nn as nn

torch.manual_seed(19)
ids = torch.tensor([[1, 2, 0, 0], [3, 4, 5, 0]])
positions = torch.arange(ids.size(1)).unsqueeze(0)
token_embedding = nn.Embedding(6, 8, padding_idx=0)
position_embedding = nn.Embedding(4, 8)
transformer_input = token_embedding(ids) + position_embedding(positions)
key_padding_mask = ids == 0

print(transformer_input.shape)
print(key_padding_mask.tolist())

Expected output: torch.Size([2, 4, 8]) and [[False, False, True, True], [False, False, False, True]].

The token and position embeddings must have the same width because they are added elementwise. The mask has shape [batch, time]; True marks positions that should be ignored as padding. This representation keeps order available to attention while preserving the same [batch, time, features] convention used throughout PyTorch sequence models.

Design Choices and Trade-offs

The first choice is vocabulary construction. Word-level vocabularies are readable but produce many rare tokens. Character-level vocabularies are small but create longer sequences. Subword tokenization often balances those extremes, although it adds preprocessing complexity. For non-language sequences, the same idea applies: decide whether an item ID, a bucketed numeric value, or a composed set of categorical features is the right token.

The embedding dimension controls capacity and cost. A tiny dimension may force unrelated tokens to share too little representational space. A very large dimension increases memory, bandwidth, and overfitting risk. A useful starting point is to treat the dimension as a hyperparameter tied to vocabulary size, dataset size, and downstream model width, then validate it against a baseline rather than assuming larger is better.

Pooling, recurrence, and attention represent sequence information differently. Mean pooling is fast and stable but loses order. Recurrent layers preserve order with a compact hidden state, but long sequences can be slower because steps are sequential. Attention can model direct interactions between distant positions, but memory grows with the number of pairwise position comparisons. Your representation should match the task: sentiment over a sentence may tolerate pooling; next-token prediction and translation usually need order-sensitive representations.

Failure Modes and Troubleshooting

  • Symptom: IndexError: index out of range in self. Cause: a token ID is negative or at least num_embeddings. Diagnose: print ids.min(), ids.max(), and the embedding table size for a failing batch. Correct: fix vocabulary serialization, map missing tokens to <unk>, and keep training and inference vocabularies identical.
  • Symptom: the model learns sequence length instead of content. Cause: padding positions are included in pooling, loss, or attention. Diagnose: compare predictions for the same sequence padded to different lengths. Correct: use padding_idx, masks, packed sequences, and loss functions configured to ignore padding labels.
  • Symptom: validation quality is unstable for rare tokens. Cause: rare embedding rows receive few updates. Diagnose: count token frequencies and inspect examples containing high-loss rare tokens. Correct: increase data coverage, use subword units, share features, load pretrained vectors, or group rare items into an unknown bucket.
  • Symptom: a recurrent model runs slowly on batches with mixed lengths. Cause: much of the batch is padding. Diagnose: log true lengths and padded lengths. Correct: bucket examples by length, pack padded sequences, or cap maximum sequence length after measuring task impact.

Security, Performance, and Reliability

Embedding tables can become large enough to dominate model memory. The table size is num_embeddings * embedding_dim parameters, so a million IDs with 256-dimensional vectors means 256 million learned values before any recurrent or attention layer is added. Large tables also make checkpoints bigger and can slow loading. Sparse gradients may help some optimizers, but optimizer support differs, so verify the chosen optimizer with the embedding configuration.

Reliability depends on stable preprocessing. If the vocabulary order changes between training and inference, the model will still run but token IDs will point at the wrong rows. Store the vocabulary or tokenizer artifact with the model checkpoint and validate special-token IDs during startup. For user-provided text or event streams, bound maximum sequence length and handle unknown items deliberately so malformed or unexpected inputs do not create crashes or excessive memory use.

Hands-on Lab

Prerequisites: Python, PyTorch, and a shell where you can run short scripts. No dataset download is required.

  1. Create a small vocabulary with <pad> at index 0 and <unk> at index 1.
  2. Convert three short token sequences into ID lists, mapping any missing token to 1.
  3. Pad the batch to the longest sequence and build a length tensor.
  4. Run nn.Embedding with padding_idx=0.
  5. Compute a masked mean vector for each example and print its shape.
  6. Change one sequence by appending extra padding and verify that its pooled vector does not change.

Verification: the embedded batch should have shape [3, max_length, embedding_dim], the pooled tensor should have shape [3, embedding_dim], and the same real tokens with more trailing padding should produce the same pooled vector within floating-point tolerance.

Cleanup: remove the scratch script or notebook. If you saved a checkpoint for experimentation, delete it unless you also saved the exact vocabulary artifact that created its token IDs.

Assessment Exercises

  1. A classifier using mean-pooled embeddings improves on training data but fails when sentences are negated. Explain why the representation may be the problem and propose an order-sensitive replacement.
  2. You load a checkpoint and validation accuracy collapses without a shape error. What vocabulary or tokenizer checks would you run first?
  3. Given a batch with lengths [40, 7, 6, 5], explain one batching strategy that could reduce wasted recurrent computation.
  4. Design a test proving that padding does not affect pooled representations or attention outputs for an otherwise identical example.
  5. For a catalog with millions of item IDs and many rare items, compare a large item embedding table with a representation built from shared categorical features.

Summary

Embeddings turn integer IDs into trainable dense vectors, but the surrounding sequence representation determines whether the model sees padding, order, and length correctly. In PyTorch, the critical habits are keeping token IDs within the table, reserving stable special tokens, masking or packing padded positions, and choosing pooling, recurrence, or attention based on the task. These choices set the foundation for the sequence and transformer models that follow in this course.