How to Perform Text Classification Using Word Embeddings in PyTorch

Text classification using word embeddings in PyTorch requires mapping token indices to dense vectors with nn.Embedding, aggregating variable-length sequences via mean pooling or EmbeddingBag, and training a linear classifier on the resulting fixed-size representation.

The EmbeddingsPyTorch.ipynb notebook in the microsoft/AI-For-Beginners repository demonstrates a complete pipeline for building text classifiers that leverage both learned and pre-trained word embeddings. This implementation handles variable-length sequences efficiently while supporting initialization from external vector sources like Word2Vec and GloVe.

The Three-Stage Architecture

According to the source code in lessons/5-NLP/14-Embeddings/EmbeddingsPyTorch.ipynb, the classifier architecture consists of three distinct stages:

1. Embedding Layer – torch.nn.Embedding maps each token index to a dense vector of fixed dimensionality.
2. Sequence Aggregation – Collapses a variable-length sequence of embeddings into a single fixed-size vector representation.
3. Linear Classifier – torch.nn.Linear projects the aggregated vector into the target class space.

This architecture appears in two variants within the notebook: one using explicit padding with mean pooling (lines 77-86), and an optimized version using EmbeddingBag (lines 95-104).

Method 1: Mean Pooling with Padding

The first approach pads all sequences in a batch to the length of the longest sample, then computes the mean across the sequence dimension.

The Padding Collate Function

The padify function (lines 106-119) handles batch construction by padding shorter sequences with zeros:

def padify(batch):
    # batch = [(label, text), ...]

    seqs = [encode(sample[1]) for sample in batch]
    max_len = max(len(s) for s in seqs)
    padded = torch.stack([
        torch.nn.functional.pad(torch.tensor(s),
                                (0, max_len - len(s)),
                                mode='constant', value=0)
        for s in seqs
    ])
    labels = torch.LongTensor([sample[0] - 1 for sample in batch])
    return labels, padded

Model Implementation

The corresponding model uses nn.Embedding followed by torch.mean:

class EmbedClassifier(nn.Module):
    def __init__(self, vocab_sz, embed_dim, n_classes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_sz, embed_dim)
        self.fc = nn.Linear(embed_dim, n_classes)

    def forward(self, x):
        # x: (batch, seq_len) after padding

        emb = self.embedding(x)               # (batch, seq_len, embed_dim)

        pooled = torch.mean(emb, dim=1)       # (batch, embed_dim)

        return self.fc(pooled)

Method 2: Efficient Aggregation with EmbeddingBag

The EmbeddingBag layer (lines 99-104) eliminates the need for explicit padding by internally aggregating embeddings from a flattened tensor using offset indices. This reduces memory overhead and simplifies the forward pass.

The Offset Collate Function

The offsetify function (lines 119-131) constructs the flat tensor and offset vector required by EmbeddingBag:

def offsetify(batch):
    # batch = [(label, text), ...]

    seqs = [torch.tensor(encode(sample[1])) for sample in batch]
    offsets = torch.cat([
        torch.tensor([0]), 
        torch.cumsum(torch.tensor([len(s) for s in seqs]), dim=0)[:-1]
    ])
    labels = torch.LongTensor([sample[0] - 1 for sample in batch])
    text = torch.cat(seqs)
    return labels, text, offsets

EmbeddingBag Model Architecture

class EmbedBagClassifier(nn.Module):
    def __init__(self, vocab_sz, embed_dim, n_classes):
        super().__init__()
        self.embedding = nn.EmbeddingBag(vocab_sz, embed_dim, mode='mean')
        self.fc = nn.Linear(embed_dim, n_classes)

    def forward(self, text, offsets):
        # text: concatenated token ids, offsets: start idx of each sample

        emb = self.embedding(text, offsets)    # (batch, embed_dim)

        return self.fc(emb)

Leveraging Pre-Trained Embeddings

Beyond training embeddings from scratch, the notebook demonstrates loading external vector matrices to bootstrap semantic knowledge.

Loading Word2Vec Vectors

Lines 29-31 show how to load Word2Vec vectors using gensim:

import gensim.downloader as api
w2v = api.load('word2vec-google-news-300')

Lines 64-70 then populate the model's embedding matrix with these vectors, falling back to random initialization for out-of-vocabulary words.

Loading GloVe via TorchText

Lines 142-144 utilize torchtext.vocab.GloVe for direct access to pre-trained vectors:

import torchtext
glove = torchtext.vocab.GloVe(name='6B', dim=50)

Initializing the Embedding Weight

Lines 222-224 assign the pre-trained vectors directly to the embedding layer:

model = EmbedClassifier(len(glove.stoi), glove.dim, len(classes))
model.embedding.weight.data = glove.vectors

This assignment ensures the model benefits from semantic relationships encoded in the external embeddings before fine-tuning.

Training Loop Implementation

The train_epoch_emb function (lines 73-94) provides a generic training loop that accommodates both padding-based and EmbeddingBag variants. The implementation handles the different batch signatures—either (labels, padded) or (labels, text, offsets)—while computing cross-entropy loss and tracking accuracy:

def train_one_epoch(model, loader, optimizer, loss_fn, device):
    model.train()
    total_loss, correct, total = 0, 0, 0
    for batch in loader:
        optimizer.zero_grad()
        if len(batch) == 3:  # EmbeddingBag case

            labels, text, offsets = [t.to(device) for t in batch]
            preds = model(text, offsets)
        else:                # Padding case

            labels, padded = [t.to(device) for t in batch]
            preds = model(padded)
        loss = loss_fn(preds, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        _, predicted = torch.max(preds, 1)
        correct += (predicted == labels).sum().item()
        total += labels.size(0)
    return total_loss / total, correct / total

Summary

• Architecture: Text classification using word embeddings requires an embedding layer, sequence aggregation (mean pooling or EmbeddingBag), and a linear classifier.
• Sequence Handling: Use padify with torch.mean for explicit padding, or offsetify with nn.EmbeddingBag for memory-efficient aggregation without padding.
• Pre-trained Vectors: Initialize embedding.weight.data with Word2Vec (via gensim) or GloVe (via torchtext) vectors to leverage external semantic knowledge.
• Source Location: Complete implementations reside in lessons/5-NLP/14-Embeddings/EmbeddingsPyTorch.ipynb within the microsoft/AI-For-Beginners repository.

Frequently Asked Questions

What is the difference between nn.Embedding and EmbeddingBag?

nn.Embedding retrieves embeddings for each token index and returns a 3D tensor (batch, seq_len, embed_dim), requiring manual aggregation. nn.EmbeddingBag combines the lookup and aggregation (sum, mean, or max) into a single operation, accepting a 1D tensor of token IDs and offset indices to delimit sequences, producing a 2D tensor (batch, embed_dim) directly.

When should I use pre-trained embeddings versus training from scratch?

Use pre-trained embeddings (GloVe, Word2Vec) when your dataset is small or when you need to capture general semantic relationships (e.g., king - man + woman ≈ queen) without extensive training. Train from scratch when your domain uses specialized vocabulary or when you have sufficient data to learn task-specific representations that outperform general-purpose vectors.

How does the offsetify collate function work with EmbeddingBag?

The offsetify function flattens all token sequences in a batch into a single 1D tensor, then computes cumulative lengths to generate offset indices marking where each sequence begins. EmbeddingBag uses these offsets to segment the flat tensor internally, aggregating embeddings per sequence without constructing a padded 2D matrix, as shown in lines 119-131 of the notebook.

Why use mean pooling instead of sum or max pooling?

Mean pooling (the default in the notebook's EmbeddingBag with mode='mean') produces a representation proportional to the semantic average of the document, invariant to sequence length. Sum pooling accumulates magnitude with longer texts, while max pooling extracts the most salient feature per dimension but loses frequency information. Mean pooling provides stable gradients and consistent magnitude across variable-length inputs for classification tasks.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}