# How to Perform Text Classification Using Word Embeddings in PyTorch

> Learn text classification with word embeddings in PyTorch. Map tokens to vectors, aggregate sequences, and train a linear classifier for effective text analysis.

- Repository: [Microsoft/AI-For-Beginners](https://github.com/microsoft/AI-For-Beginners)
- Tags: how-to-guide
- Published: 2026-08-08

---

**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:

```python
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`:

```python
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`:

```python
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

```python
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`:

```python
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:

```python
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:

```python
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:

```python
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.