Train your own embedding model
Tokens to vectors·20 min

Mean pooling

0 / 1
# grade this stage (from your course directory)
npx theinference test --stage mean-pooling --watch

A sentence is a variable number of token vectors, but an embedding is one fixed vector. Mean pooling bridges the gap: average the token rows over the real (non-padding) positions. Sentence-BERT found mean pooling beats both the [CLS] token and max pooling for sentence similarity, and it's still the default in sentence-transformers.

Why the mask matters

masked mean — never average the padding
rows = [[1, 1],
        [3, 3],
        [5, 5]]        mask = [1, 1, 0]   ← position 2 is padding

pooled = ([1,1] + [3,3]) / 2 = [2, 2]    ← divide by 2, not 3

What to build

Your task

Implement mean pooling

Implement runMeanPool(rows, mask): sum the rows where mask[i] === 1, then divide by the number of non-pad positions (not by the sequence length). Return the pooled [32] vector.

tests/adapters.ts
export function runMeanPool(rows: number[][], mask: number[]): number[] {
  // sum rows where mask[i] === 1, divide by the count of those positions
}
How it's graded

Adapter mode: the pooled vector is compared to the reference with allclose, atol 1e-6. The fixture includes padded positions, so the mask handling is exercised.

The classic bug

Dividing by T (the padded sequence length) instead of by the count of non-pad tokens. With padding present, that silently shrinks short sentences toward zero and quietly fails the atol 1e-6 check.

Tokenize + embedding lookupWhy a projection head?