Embeddings and positions

The tokenizer turns text into integers. The first layer of the model turns each integer into a vector, and every layer after that works on those vectors. This post is about that first layer, and the second piece of info it has to add, which is where in the sequence each token is sitting.

An embedding is just a row lookup

An embedding table is a learned matrix $E$ with one row per vocabulary entry and $C$ columns, where $C$ is the model width. Token ID $i$ becomes row $E_i$:

For the Day 1 model, $V = 1{,}024$ and $C = 128$, so the table holds 131,072 numbers. A batch of token IDs shaped [B, T] becomes vectors shaped [B, T, C].

Mathematically the lookup is a one-hot vector times the matrix. In practice PyTorch's nn.Embedding just indexes the row directly, which skips multiplying by 1,023 zeros. The gradient only flows back to the rows that got used, so a token that never shows up in training keeps its initial random row.

The rows start random. Training moves them so that tokens used in similar contexts end up with similar vectors, because the layers above are the same function for every token.

Attention doesn't know the order

Attention compares every position's query with every earlier position's key. Nothing in that comparison says which key is one step back and which one is fifty back. The causal mask gives you one fact, that a key isn't in the future. Beyond that, if you shuffled the earlier tokens, each position would get the exact same weighted average.

And for language that's fatal. "The dog bit the man" and "The man bit the dog" would look the same to the last position. So the model needs position added into the vector before attention sees it.

A second table, indexed by position

The simplest fix is a second table, $P \in \mathbb{R}^{T_{\max} \times C}$, with one row per position. The input to the first block is the sum:

This is what GPT-2 does and what octlm's Day 1 baseline does. The Day 1 table has 128 rows, one per position of its 128-token context.

So what happened ?

• The same word at two positions gets the same token row. Only the position row tells them apart in the sum.
• The sum keeps both signals in one vector of the same width. Nothing gets concatenated, so the model width doesn't grow.
• Switching to sinusoidal rows only changes the middle table, the token table isn't touched.

A learned table has two limits though:

1. It stops at its last row. A model with a 128-row table has no vector for position 128. octlm raises an error instead of inventing one. Day 2 wrote this down as the reason the length experiment couldn't use learned positions as its control.
2. Each row only learns from positions training actually reaches. octlm packs training text into full-length blocks, so all 128 rows train. But a model trained on shorter sequences would leave its later rows at their random starting values.

A formula instead of a table

The 2017 Transformer paper used fixed waves instead of a learned table. Column pair $i$ of position $t$ is a sine and a cosine at a frequency that falls with $i$:

The first pair turns once every $2\pi$ positions, the last pair turns about 10,000 times slower. So each position gets a unique pattern, and you can compute it for any length.

So what happened ?

• In the table, the left columns change every row and the right columns barely change. Fast pairs separate neighbors, slow pairs separate positions that are far apart.
• In the similarity map, the band along the diagonal has the same shape all the way down. The dot product of two sinusoidal position vectors is a sum of $\cos(\omega_i (t - s))$ terms, so it only depends on the distance $t - s$.
• Wider tables separate positions more sharply, because more frequencies take part.

That second point is why sinusoids were picked in the first place. A fixed offset is a fixed linear map, so a model could learn "attend 3 back" once and use it anywhere. In practice though, the signal gets added to the token vector and then mixed by the projections, which blurs it. Day 2 replaced both approaches with RoPE, which applies position inside attention itself.

The code

The embedding step handles all three position types. Learned and sinusoidal positions get added here. RoPE adds nothing here and acts later inside attention.

def _embed(self, token_ids: Tensor, offset: int = 0) -> Tensor:
    x = self.token_embedding(token_ids)
    if self.position_embedding is not None:
        positions = torch.arange(offset, offset + token_ids.shape[1], device=token_ids.device)
        return x + self.position_embedding(positions)
    if self.config.position == "sinusoidal":
        table = sinusoidal_positions(
            offset + token_ids.shape[1], self.config.d_model, token_ids.device
        )[offset:]
        return x + table
    return x

The sinusoidal table is built from the sequence length on every forward pass, so it has no maximum length:

def sinusoidal_positions(length: int, width: int, device: torch.device) -> Tensor:
    position = torch.arange(length, device=device, dtype=torch.float32).unsqueeze(1)
    index = torch.arange(0, width, 2, device=device, dtype=torch.float32)
    angle = position / torch.pow(ROPE_BASE, index / width)
    table = torch.zeros(length, width, device=device)
    table[:, 0::2] = angle.sin()
    table[:, 1::2] = angle.cos()
    return table

A test on this site checks the TypeScript table on this page against this function to 1e-5.

What we measured, and what we chose not to build

Problem, we needed a way to look at embeddings without pulling in a plotting dependency.

The benchmark prints the nearest token to token 0 and the nearest position to position 0 by cosine similarity. On the untrained benchmark model those neighbors are random, which is what they should be. So the diagnostic only proves the code path works. The Day 4 model is the first one trained long enough for its neighbors to actually mean something.

The Day 1 reading plan also covered pooling, which turns a sequence of vectors into one vector. Mean pooling averages [B, T, C] over $T$ into [B, C], and learned pooling adds weights to decide the mix. Both are for classifiers, which need one answer per sequence. A language model needs a prediction at every position, so neither one belongs in its forward pass. The note keeps both as written exercises, and octlm has no pooling code.

What we didn't build, and why

• Pooling layers. Explained above, we don't write code nobody uses.
• Embedding plots. An untrained model's embedding space has no structure worth plotting.
• Separate input and output tables. octlm ties the output head to this same table. The baseline block covers why.