🎓 A1Learner
🔍
0 of 87 lessons completed

PyTorch Bag of Embeddings

A bag of embeddings is a sentence vector made by averaging the vectors of the sentence's words. A small classifier on that average can already judge the sentiment of a review.

Average three word rows

The three indices stand for great, lovely, and warm. Their mean is one vector for the whole group of words.

torch.manual_seed(0)
embedding = nn.Embedding(18, 4, padding_idx=0)
rows = embedding(torch.tensor([1, 2, 5]))
mean = rows.mean(dim=0)
print([round(value, 4) for value in mean.tolist()])

Order changes nothing

The same three words in a different order average to the same vector. The bag keeps which words appeared and drops the order they came in.

other = embedding(torch.tensor([5, 1, 2])).mean(dim=0)
print(torch.allclose(mean, other))

Padding adds nothing

A mask of 1s and 0s marks the real words. The masked average over a padded row matches the plain average.

padded = torch.tensor([1, 2, 5, 0, 0])
mask = (padded != 0).unsqueeze(-1).float()
masked = (embedding(padded) * mask).sum(dim=0) / mask.sum()
print(torch.allclose(mean, masked))

Example

import random

import torch
from torch import nn

positive = ["great", "lovely", "bright", "fun", "warm"]
negative = ["awful", "dull", "cold", "boring", "grim"]
filler = ["the", "a", "movie", "story", "scene", "really", "very"]

vocab = ["<pad>"] + positive + negative + filler
word_to_index = {word: index for index, word in enumerate(vocab)}


def make_review(rng):
    label = rng.randrange(2)
    side = positive if label == 1 else negative
    words = [rng.choice(side) for _ in range(rng.randint(2, 3))]
    words += [rng.choice(filler) for _ in range(rng.randint(1, 3))]
    rng.shuffle(words)
    return [word_to_index[word] for word in words], label


rng = random.Random(42)
train_reviews = [make_review(rng) for _ in range(240)]
test_reviews = [make_review(rng) for _ in range(120)]

max_len = 6


def to_batch(reviews):
    indices = torch.zeros(len(reviews), max_len, dtype=torch.long)
    for i, (words, _) in enumerate(reviews):
        indices[i, :len(words)] = torch.tensor(words)
    labels = torch.tensor([label for _, label in reviews])
    return indices, labels


train_idx, train_y = to_batch(train_reviews)
test_idx, test_y = to_batch(test_reviews)

torch.manual_seed(0)
embedding = nn.Embedding(len(vocab), 16, padding_idx=0)
head = nn.Linear(16, 2)
optimizer = torch.optim.Adam(
    list(embedding.parameters()) + list(head.parameters()), lr=0.05)
loss_fn = nn.CrossEntropyLoss()


def sentence_vectors(indices):
    vectors = embedding(indices)
    mask = (indices != 0).unsqueeze(-1).float()
    return (vectors * mask).sum(dim=1) / mask.sum(dim=1)


for epoch in range(30):
    optimizer.zero_grad()
    loss = loss_fn(head(sentence_vectors(train_idx)), train_y)
    loss.backward()
    optimizer.step()

with torch.no_grad():
    predictions = head(sentence_vectors(test_idx)).argmax(dim=1)
accuracy = float((predictions == test_y).float().mean())
print("Training reviews:", len(train_reviews))
print("Test reviews:", len(test_reviews))
print(f"Test accuracy: {accuracy:.3f}")

Key points

  • Bag of embeddings: the mean of a sentence's word vectors
  • Word order is lost in the average, so two orders of the same words give the same vector
  • Use padding_idx=0 so the pad row stays at zero and adds nothing
  • A mask counts only the real words in the average
  • The head is one Linear layer from the sentence vector to the class scores

In this example, a seeded generator builds 240 training reviews and 120 test reviews. Each review mixes two or three sentiment words from one side with one to three filler words, then shuffles them. Short reviews are filled with the pad index up to length 6. The embedding table has width 16 and a padding index of 0. A mask averages only the real words of each review. One Linear head scores the average, and Adam trains the table and the head together for 30 epochs. The test accuracy prints as 1.000, because the sentiment words never mix sides in this data. A later lesson, Reading Text Classifier Errors, shows reviews where word order matters and this kind of model fails.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does a bag of embeddings throw away?