Token and Positional Embeddings

A token embedding is a matrix with one row per vocabulary id and one column per model dimension. The model's first contact with a token id is a row lookup in it, and the returned vector is what the rest of the model treats as the meaning of the id. Backprop updates the rows used in the forward pass. Over training, the geometry of those rows comes to encode similarity as directions.

Token ids carry no order, so the model needs a second signal that tells position one from position seventeen. There are two dominant choices. A learned positional embedding is a second lookup table with one row per position. A fixed sinusoidal positional embedding is a math formula with no parameters. The choice has consequences. A learned table is a parameter, bounded by the maximum context length the model was trained on. A sinusoidal formula extends to any position in theory, but this lesson's SinusoidalPositionalEmbedding precomputes a fixed table at max_context_length and its forward raises past that bound. Both modules therefore enforce a maximum context length here. The model may still struggle past its training length even when the table is large enough to index.

This lesson builds both and composes each with the token embedding into a single input for the next lesson's attention block.

The embedding stage takes a batch of token ids of shape (B, T) and returns a tensor of shape (B, T, D), where D is the model dimension. Every batch element has the same context length T, and every position gets a vector of the same dimension D.

The composition is a sum, not a concatenation. Summing keeps D constant through the network and lets the model decide, feature by feature, whether token meaning or position dominates at each layer.

The token embedding is a parameter tensor of shape (V, D), where V is the vocabulary size. PyTorch exposes it as nn.Embedding(V, D). At init the entries are drawn from a small Gaussian with mean zero and, traditionally for transformer-scale models, a standard deviation around 0.02. The exact init matters less than keeping it consistent across runs.

The forward pass is a single indexing operation: PyTorch maps (B, T) int64 ids to (B, T, D) floats by gathering rows. The backward pass accumulates gradients only into the rows touched in the forward pass, so a row that never appeared in the batch receives zero gradient on that step.

A subtle detail: the token embedding and the output projection at the end of the model often share weights (weight tying). When they do, every backward pass touches every row of the embedding through the output side. This lesson builds only the input side, but in a full model the same matrix could also serve as the output projection.

A learned positional embedding is a second nn.Embedding of shape (max_context_length, D), keyed by position id 0, 1, 2, ..., T-1. The forward pass returns the (T, D) position vectors for the first T positions, and EmbeddingComposer broadcasts them across the batch dimension.

The cost of the learned table is that it cannot be queried at position T if the model was only trained up to position T-1. The row does not exist. Production decoder-only models that use this scheme bake the maximum context length into the architecture and refuse longer inputs.

A sinusoidal positional embedding is a function from position to vector, with no parameters. Position p and feature i produce

angle = p / (10000 ** (2 * (i // 2) / D))
emb[p, 2k]     = sin(angle)
emb[p, 2k + 1] = cos(angle)

Every position gets a unique vector. The wavelength varies geometrically across feature dimensions, so the lower dimensions oscillate fast and encode fine position, while the higher dimensions oscillate slowly and encode coarse position.

Using sin and cos together has a useful consequence: the vector at position p + k is a linear function of the vector at position p. That gives the attention layer an easy path to learning relative-position offsets, and the model needs no separate parameter to express "look five tokens back."

The table is computed once at construction and indexed at forward time.

The input pipeline does three things in order: read the token ids, look up the token vectors, add the positional vectors. The sum is the output.

The sum step replicates the (T, D) positional tensor along the batch dimension. PyTorch handles that automatically because the positional tensor has shape (1, T, D) after unsqueeze.

I run both variants on the same inputs and print two diagnostics among the demo's checks.

The first is parameter count. The learned variant adds max_context_length * D parameters on top of the token embedding. The sinusoidal variant adds zero.

The second is cosine similarity between embeddings at neighbouring positions. The sinusoidal variant decays smoothly and predictably because the function is continuous. The learned variant at initialization shows near-random similarity because its rows are drawn independently. After training it typically develops a similar smooth structure, but it has to discover that structure from data.

I don't build rotary positional encoding (RoPE) or AliBi here. Those are the modern choices in production transformers. They follow the same shape contract as the embeddings above (a position-dependent transformation applied to vectors of shape (B, T, D)), but they act at the attention-projection step instead of at the input. The next lesson builds the attention block, and one optional extension is to fold rotary into the query-key projections there.

I also don't train the embedding. Training needs a loss, which needs a model output, which needs attention and an LM head. That is the next lesson and the one after.

main.py defines three modules and a composer. TokenEmbedding wraps nn.Embedding(V, D). LearnedPositionalEmbedding wraps nn.Embedding(L, D). SinusoidalPositionalEmbedding precomputes the table and exposes it as a buffer. EmbeddingComposer ties a token embedding and a positional embedding together. The demo at the bottom prints the shapes, the parameter counts, and the neighbour-position similarity diagnostic. The tests in code/tests/test_embeddings.py pin shape, broadcast behaviour, parameter count, and the sinusoidal formula.

Run the demo. Then change the model dimension D from 64 to 32 and watch how the sinusoidal wavelength bands change.