Practical build 4

Lab: train and evaluate embeddings

Introduction

A token ID is only a label until training gives its lookup row useful geometry. This build turns 800 generated token pairs into eight-feature vectors and tests whether related tokens become retrievable neighbors.

Learning goal: Train a contrastive embedding table and evaluate whether its geometry recovers the generated relationship rule.
Before you start: Tensors, embeddings as table lookup, binary classification loss, gradients, and data splits.

Lesson plan

Generate independent pair sets and define inputs, targets, baseline, and success criteria.
Build the embedding scorer and trace IDs, vectors, pair logits, and loss shapes.
Train, evaluate classification and retrieval, then state what the toy evidence cannot prove.

Prerequisites

Complete Lessons 1 to 4 and read Lesson 7 first. You should know that a tensor has a shape, gradients store information for parameter updates, and an optimizer changes parameters. This build also uses a class. Lesson 3 introduced the bridge: a class is a reusable package of data and functions. In a PyTorch model, __init__ creates and stores trainable layers. The forward method says how one input becomes one output. Calling model(x) runs forward; you do not call it directly.

Plain-language start: labels are not meanings

Imagine twenty labelled drawers. Drawer 0 contains the vector for cat. Drawer 1 contains the vector for dog. The number on a drawer only locates it. The fact that 1 is close to 0 does not make dog similar to cat. An embedding lookup starts by selecting a drawer. Learning happens later, when the loss changes the numbers stored inside the selected drawers.

This build creates five artificial topics: animals, fruit, music, weather, and travel. A positive pair contains two different words from one topic, such as cat and puppy. A negative pair crosses topics, such as cat and flute. The model must score positive pairs above zero and negative pairs below zero. This is a small contrastive task: it learns by comparing examples that should be close with examples that should be apart.

Important limit

The topics are labels chosen by this program. If the learned vector for cat is near dog, the model recovered our generated co-occurrence rule. It did not discover the full natural-language meaning of either word. Real word embeddings need a large, varied corpus and careful evaluation.

Input, target, output, and success

Set an expectation before training: a balanced positive/negative task has a 50% always-one-class baseline. The trained model must beat that result on held-out pairs. Training accuracy alone cannot show that the vectors learned a rule that transfers.

One example has a centre word ID, a context word ID, and a binary target. For example, (cat, dog, 1) is positive and (cat, piano, 0) is negative. With a batch of B = 800 training pairs, centre IDs have shape (800,), context IDs have shape (800,), and labels have shape (800,). Their dtypes are torch.long, torch.long, and torch.float32.

The vocabulary has V = 20 entries and each vector has D = 8 features. nn.Embedding(V, D) therefore owns a parameter matrix of shape (20, 8). Looking up a batch of centre IDs returns (B, 8). Looking up contexts returns another (B, 8). A row-wise dot product reduces the feature axis and returns one logit per pair, shape (B,).

Success has two parts. First, binary accuracy on 400 held-out pairs should beat the 50% majority baseline. Second, the three nearest vectors for each word should usually come from its generated topic. We report same-topic precision@3: correct retrieved neighbours divided by all retrieved neighbours.

Generate separate training and validation pairs

The complete program uses separate random seeds. Seed 10 creates training pairs. Seed 20 creates held-out pairs that are not used for optimizer updates. Both sets follow the same known rule. They are balanced: half positive and half negative. A majority-class baseline therefore reaches exactly 0.5 accuracy.

train = make_contrastive_pairs(800, seed=10)
validation = make_contrastive_pairs(400, seed=20)

train_centres, train_contexts, train_labels = train
assert train_centres.shape == (800,)
assert train_centres.dtype == torch.long
assert train_labels.shape == (800,)
assert train_labels.sum().item() == 400

The generator makes the example repeatable. Repeating the same seed gives identical tensors. Changing the seed gives a different sample. Determinism is useful for debugging because a code change is not confused with a data change.

The score and loss

Let e_i be the vector for the centre ID and e_j the vector for the context ID. We first divide each vector by its length. The resulting cosine similarity is the dot product of the normalized vectors:

cosine(i, j) = (e_i / ||e_i||) dot (e_j / ||e_j||)
logit(i, j) = 5 * cosine(i, j)

A cosine value is between -1 and 1. Multiplying by five gives binary cross-entropy a wider logit range. Binary cross-entropy with logits applies a stable sigmoid calculation internally. For target y, it penalizes a low score when y = 1 and a high score when y = 0. “With logits” matters: do not apply sigmoid yourself before this loss.

The model class

class ContrastiveEmbedding(nn.Module):
    def __init__(self, vocabulary_size, embedding_size=8):
        super().__init__()
        self.embedding = nn.Embedding(vocabulary_size, embedding_size)

    def forward(self, centre_ids, context_ids):
        centre = F.normalize(self.embedding(centre_ids), dim=-1)
        context = F.normalize(self.embedding(context_ids), dim=-1)
        return 5.0 * (centre * context).sum(dim=-1)

super().__init__() initializes the nn.Module part of the object. Assigning the embedding to self.embedding registers its matrix as a parameter. As a result, model.parameters() finds it and Adam can update it. dim=-1 means the final axis, which is the eight embedding features. Element-wise multiplication keeps shape (B, 8); summing that last axis gives (B,).

This model deliberately shares one embedding table for centres and contexts. That makes retrieval direct: every word lives in one vector space. Some skip-gram systems use separate input and output tables. That can work well, but then you must decide which table, or which combination, to use for retrieval.

Training, gradients, and evaluation

model = ContrastiveEmbedding(vocabulary_size=20, embedding_size=8)
optimizer = torch.optim.Adam(model.parameters(), lr=0.04)

for step in range(250):
    optimizer.zero_grad()
    logits = model(train_centres, train_contexts)  # (800,)
    loss = F.binary_cross_entropy_with_logits(logits, train_labels)
    loss.backward()
    optimizer.step()

zero_grad clears gradients from the previous step. backward calculates how each used embedding value affected the loss. step changes the table. Looking up a vector is not itself semantic learning; the loss and optimizer produce the learned arrangement.

Evaluation calls model.eval() and uses torch.no_grad(). This model has no dropout, so eval does not change its arithmetic, but using the evaluation pattern now prevents errors when later models do contain training-only behaviour. Pair predictions use logit >= 0. Retrieval normalizes the full (20, 8) matrix, multiplies it by its transpose to obtain a (20, 20) similarity matrix, hides each word’s self-score, and selects the top three remaining scores.

Run the complete program

Download build_embeddings.py. It contains data generation, the model, training, held-out evaluation, retrieval, assertions, CLI arguments, and no download or hidden state.

python examples/pytorch/build_embeddings.py
python examples/pytorch/build_embeddings.py --steps 100
python examples/pytorch/build_embeddings.py --smoke-test

The smoke test uses at most three optimizer steps. It checks that the program connects and runs, but it skips learning thresholds. It is not evidence that the model learned. On the tested CPU runtime, the normal 250-step command produced:

vocabulary size: 20
lookup shape for 4 IDs: (4, 8)
majority baseline accuracy: 0.500
held-out pair loss: 0.1438
held-out pair accuracy: 0.988
same-topic precision@3: 1.000
nearest to cat: puppy, kitten, dog
nearest to piano: drum, violin, flute

Small floating-point differences are possible on another PyTorch version. The important comparisons are that held-out accuracy clearly beats 0.5 and retrieval mostly returns words from the same generated topic.

Failure drills

Drill 1: use floating-point IDs. Change a centre tensor to .float(). nn.Embedding raises a dtype error because table rows need integer indices. Fix it by preserving torch.long.

Drill 2: remove negative pairs. Train only target-one examples. The model can move every vector toward the same direction and still satisfy the loss. Nearest neighbours then carry little topic information. Contrastive learning needs meaningful contrasts.

Drill 3: include the query itself in retrieval. Remove fill_diagonal_(-inf). Every word retrieves itself with cosine similarity one, which inflates the metric without testing generalization. A retrieval metric must exclude the query.

Drill 4: normalize along dim=0. This normalizes across examples instead of features. Scores then depend on which other examples happen to be in the batch. The correct feature axis is -1.

Solved exercises

1. What shapes result from 32 pairs and embedding size 12?

Each ID tensor is (32,). Each lookup is (32, 12). Element-wise products remain (32, 12). Summing feature axis 1 returns logits of shape (32,). The embedding parameter has shape (20, 12).

2. Why can a token ID not be sent directly into a linear layer as a number?

A raw ID gives a false numeric order. A linear formula would treat ID 10 as larger than ID 2. IDs are categories, not measurements. An embedding uses an ID only to select a trainable row, so no numeric distance between IDs is assumed.

3. Add one word to each topic. What must change?

Add the strings to TOPIC_WORDS. The derived vocabulary and maps grow automatically. The model must be created with the new vocabulary length. Precision@3 can remain unchanged because three neighbours are still requested, though precision@4 would now test the complete four same-topic alternatives.

4. Why is held-out pair accuracy not enough?

A classifier might separate pair types while still producing poor nearest-neighbour ordering for a chosen downstream use. The retrieval check tests the geometry directly. Both metrics remain limited because they use the same artificial topic rule.

What this build does and does not prove

Nearby learned vectors are evidence about this generated context rule. They are not proof that the tokens share every human meaning. Change the pair-generation rule and the geometry may change, because the training evidence changed.

You built a complete learned representation pipeline: generated data, baseline, model, loss, optimization, independently resampled evaluation, and retrieval. The vocabulary is small, so evaluation repeats many pairs that occurred during training. A separate random seed does not make those pairs unseen. This checks recovery of the chosen topic-group rule, not generalization to new words or relationships. It does not establish broad language understanding, fairness, robustness to rare words, or transfer to a real corpus. Those claims need real data and stronger evaluations.

← Previous build: linear regression Next build: LSTM →