Skip to main content

Word and Sentence Embeddings

An embedding table maps each token ID to a dense vector. Similarity in that space is learned from an objective and corpus; it is evidence about that representation, not a universal statement that two words or sentences mean the same thing.

Prerequisites​

Complete Text Data and Tokenization and be comfortable with tensor axes.

Learning goals​

  • trace (batch, tokens) IDs to (batch, tokens, embedding) vectors.
  • pool only non-padding positions into a fixed-size sentence vector.
  • state why embedding similarity depends on model, pooling, language, domain, and version.

Worked example​

The weights below are hand-written, not pretrained. That makes masked mean pooling and cosine similarity auditable without a download.

import torch
from torch import nn
from torch.nn import functional as F

table = nn.Embedding(6, 3, padding_idx=0)
with torch.no_grad():
    table.weight.copy_(torch.tensor([
        [0.0, 0.0, 0.0], [0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0], [1.0, 1.0, 0.0],
        [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0],
    ]))

input_ids = torch.tensor([[2, 3, 0], [2, 4, 5]])
mask = input_ids.ne(0)
token_vectors = table(input_ids)
sentence_vectors = (token_vectors * mask.unsqueeze(-1)).sum(dim=1) / mask.sum(dim=1, keepdim=True)
similarity = F.cosine_similarity(sentence_vectors[0], sentence_vectors[1], dim=0)

print(f"ids={tuple(input_ids.shape)} tokens={tuple(token_vectors.shape)}")
print(f"mask={mask.tolist()} pooled={sentence_vectors.tolist()}")
print(f"cosine={similarity.item():.3f}")

Output:

ids=(2, 3) tokens=(2, 3, 3)
mask=[[True, True, False], [True, True, True]] pooled=[[1.0, 0.5, 0.0], [0.0, 0.3333333432674408, 0.0]]
cosine=0.447

Mean pooling discards order. A contextual encoder instead produces a token vector conditioned on surrounding tokens, but its pooled sentence representation still needs an explicit rule and validation for the intended similarity task.

Common mistakes​

  • Averaging padding vectors or dividing by the padded length.
  • Comparing raw dot products when vector norms vary and cosine was intended.
  • Calling nearest neighbors synonyms without checking polysemy and context.
  • Using a sentence embedding from one model with an index built by another version.
  • Inferring fairness or factual correctness from a plausible geometric neighborhood.

Practice​

Replace the first sentence's final padding ID with token 5. Predict the direction of the cosine change before running it, then explain what mean pooling cannot represent about token order.

Next steps​

Sequence Models before Transformers adds order-dependent state.

Primary sources​

Checked 2026-10-16: PyTorch Embedding, PyTorch cosine similarity, word2vec paper, and Sentence-BERT paper.