Embedding lookup and learned vectors

Introduction

Use categorical token IDs to select trainable vectors without inventing numeric order between IDs.

Learning goal

Track embedding shapes and explain how lookup gradients update selected rows.

Before you start

Token IDs, trainable parameters, backpropagation, and vector similarity.

Lesson plan

  1. Index rows by token ID
  2. Track appended feature dimensions
  3. Inspect sparse row gradients

The problem

Token IDs are arbitrary category labels. Feeding ID 20 as a measured number would falsely imply it is twice ID 10. An embedding table instead assigns every vocabulary item a trainable feature vector.

Lookup itself does not create meaning. The training objective changes selected rows so their geometry becomes useful for predicting contexts, labels, or another target.

Work through a small example

Let a vocabulary have four tokens and embedding width two. The parameter table has shape (V=4,D=2). Looking up IDs [2,0,2] selects rows 2, 0, and 2, producing shape (T=3,D=2).

Integer IDs select rows from an embedding tableThree scalar ID circles two, zero, and two point to rows in a four-by-two pastel table. The selected rows form a three-by-two output, with row two repeated. 202 row 0row 1row 2row 31, 00, 12, 3−1, 1 [2, 3][1, 0][2, 3]

The general rule

For IDs shaped (B,T), nn.Embedding(V,D) returns (B,T,D). It is mathematically equivalent to multiplying one-hot vectors by a (V,D) matrix, but direct indexing avoids materializing huge sparse one-hot tensors.

During backpropagation, only rows used in the batch receive lookup gradients. Repeated IDs accumulate gradient contributions. With ordinary dense optimizer state, memory may still scale with all V*D parameters.

Implement and inspect

import torch
from torch import nn

table = nn.Embedding(4, 2)
with torch.no_grad():
    table.weight.copy_(torch.tensor([[1., 0.], [0., 1.], [2., 3.], [-1., 1.]]))

ids = torch.tensor([2, 0, 2])
vectors = table(ids)
loss = vectors.sum()
loss.backward()

print(vectors.detach().tolist())
print(table.weight.grad.tolist())

Expected vectors are [[2,3],[1,0],[2,3]]. Gradient rows are [[1,1],[0,0],[2,2],[0,0]]: row 2 appears twice, row 0 once, and unused rows receive zero.

Engineering checks

Go deeper

Embedding geometry is identified only relative to the objective. Rotating every vector and compensating later weights can preserve model behavior, so individual coordinates usually lack fixed human meaning.

Rare tokens receive fewer direct updates. Frequency-aware sampling, subword sharing, and regularization can help, but each changes the learned relation. Static embeddings assign one vector per token; contextual models create a position-specific vector using surrounding tokens.

Practice and recap

Question: IDs have shape (32,10), vocabulary size 5,000, and embedding width 64. Name the parameter and output shapes.

Worked answer

The table parameter has shape (5000,64). Lookup appends the feature axis, so output has shape (32,10,64). The IDs remain categorical addresses; the returned values are floating-point learned features.

An embedding is an indexed parameter table. The objective and data—not the ID values—determine what relations its vectors encode.