The Embedding Layer

The tokenizer hands the model a list of integers: [464, 3797, 3332] for "the cat sat". Here is the problem. Those integers are labels, not quantities. Token 3797 is not "more" than token 464, and token 3798 has nothing to do with "cat" just because it is next door. But a neural network can only add and multiply. Feed it the raw ids and it will happily conclude that 3797 is about eight times as much as 464, which is nonsense.

So we need a way to turn "which token is it?" into numbers that can be added and multiplied meaningfully. That is the job of the embedding layer, the first learned component in the pipeline. It is also the simplest, which makes it the right place to build intuitions we will need for everything after.

The problem: ids are names, not numbers

Let's be precise about what goes wrong. Suppose we feed the id directly as a one-dimensional input $x = 3797$. Any weight $w$ the model applies gives $w \cdot 3797$, and for the neighbouring token 3798 it gives $w \cdot 3798$, almost the same value. The model is forced to treat tokens with nearby ids as nearly identical, but ids were assigned by the order in which BPE happened to learn merges. The numbering carries no information, so any computation that uses it as a number is hallucinating structure.

What we actually want is a representation with two properties. First, it must not impose any accidental relationship between tokens: each token should start out as free to be similar or dissimilar to any other. Second, it should be able to learn relationships from data, so that "cat" and "dog" end up close and "cat" and "parliament" end up far. The first property is easy; the second is where the magic happens.

One-hot encoding: correct but wasteful

The textbook way to feed a categorical label into a network is a one-hot vector: a vector of length $V$ (the vocabulary size) that is all zeros except for a single 1 at the position of the token's id. With a five-token vocabulary [the, cat, sat, on, mat], "cat" is $[0, 1, 0, 0, 0]$ and "mat" is $[0, 0, 0, 0, 1]$.

This solves the first problem beautifully. Every one-hot vector is the same length, they are all at the same distance from each other, and their dot products are all zero: no token is accidentally "close" to any other. The representation is perfectly neutral.

It fails the second problem just as completely. Because every pair is orthogonal, one-hot vectors can never express that "cat" and "dog" are related. Similarity is exactly zero for every pair, forever. And they are enormous: for GPT-2, each token is a 50,257-dimensional vector with 50,256 zeros in it. Storing a 1,024-token context as one-hots would take fifty million numbers, almost all of them zero. The first thing any network would do with such an input is multiply it by a weight matrix to shrink it, so let's look at what that multiplication does.

Dense embeddings: a lookup table that is also a matrix multiply

Here is the key idea, and it is small. Multiply the one-hot vector $e_i$ (a 1 at position $i$) by a matrix $E$ of shape $V \times d$. What comes out?

The result is exactly the $i$-th row of $E$: a dense vector of $d$ numbers. A one-hot vector times a matrix selects a row. All the zeros contribute nothing, and the single 1 copies its row out. So instead of ever materialising the one-hot vector, we can simply index into $E$ and take row $i$. That is what nn.Embedding does: it is literally a $V \times d$ matrix, and the forward pass is E[ids].

Click a token below and watch the three views line up: the one-hot vector, the row it selects in $E$, and the resulting dense vector. The numbers in $E$ are made up, but the mechanism is the real one.

In PyTorch the whole layer is a few lines. The second function does the same thing the slow way, to prove the identity:

import torch, torch.nn as nn, torch.nn.functional as F

class Embedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(vocab_size, d_model) * 0.02)   # V x d

    def forward(self, ids):            # ids: (batch, seq) of integers
        return self.weight[ids]        # (batch, seq, d): pure row lookup

def embed_via_onehot(weight, ids):
    onehot = F.one_hot(ids, num_classes=weight.shape[0]).to(weight.dtype)   # (batch, seq, V)
    return onehot @ weight             # (batch, seq, d): same numbers, V times slower

emb = Embedding(50257, 768)
ids = torch.tensor([[464, 3797, 3332]])
assert torch.allclose(emb(ids), embed_via_onehot(emb.weight, ids))
print(emb(ids).shape)

What the numbers mean

Now the question everyone asks: what is dimension 17 of the embedding for "cat"? The honest answer is: nothing in particular. Nobody designed the coordinates. The 768 numbers in each row start as small random values and are adjusted by gradient descent, thousands of times, in whatever direction makes next-token prediction slightly less wrong. The coordinates are whatever they need to be.

And yet structure appears. Think about the pressure on two tokens like "cat" and "dog". They appear in similar contexts ("the ___ sat on", "I fed the ___"), so the model must produce similar next-token predictions after each. The cheapest way to do that is to give them similar vectors, so the rest of the network can treat them alike. Similarity of use becomes similarity of geometry, without anyone asking for it. This is the distributional hypothesis (Firth, 1957: "you shall know a word by the company it keeps") turned into an optimisation outcome.

Measuring similarity: cosine

To say "cat is close to dog" we need a number. The natural candidate is the dot product, but it grows with the lengths of the vectors, and lengths in an embedding table vary a lot (frequent tokens often get longer vectors simply because they receive more gradient). To compare directions only, divide out the lengths. The result is the cosine of the angle between the two vectors:

It ranges from $1$ (same direction) through $0$ (perpendicular, unrelated) to $-1$ (opposite). The lengths cancel, so a token that is twice as "loud" but points the same way scores the same.

Vector arithmetic and the famous analogy

The result that made embeddings a household name in NLP: in word2vec, $\text{vec}(\text{king}) - \text{vec}(\text{man}) + \text{vec}(\text{woman})$ lands close to $\text{vec}(\text{queen})$. The offset between "man" and "king" (roughly, "royalty") is the same as the offset between "woman" and "queen". Relations become directions.

Three honest caveats. First, the standard evaluation finds the nearest word to the result excluding the three input words; without that exclusion the nearest word is very often just "king" itself (Nissim et al., 2020). Second, the result is approximate; "queen" is usually near, not exactly there, and the analogy works far less reliably for most relations than the headline example suggests. Third, LLM input embeddings are trained for a different job than word2vec and their analogy structure is weaker; the cleanest "concept directions" in a modern model tend to live in the residual stream deeper in the network, not in row 0. The picture is real, but it is a picture.

The map below is a hand-placed two-dimensional toy: about thirty words in clusters. Because the space is only 2-d, cosine similarities are cruder than in a real 768-d table, but the mechanics of measuring and of vector arithmetic are exactly what you would run on real embeddings.

Choosing the embedding dimension

How wide should each row be? In a transformer the embedding width is the width of the residual stream, $d_{model}$, because the embedding is the initial residual stream. So $d$ is set by the overall model design, not by the embedding layer alone. Wider means every token can carry more information at once and the attention and MLP layers have more room, at a cost that grows as $d^2$ in every block.

The pattern: the embedding table is a large fraction of a small model and a small fraction of a large one, because it grows with $V \times d$ while the rest of the model grows with $N \times d^2$. Use the calculator to see the split, including the cost of an untied output layer and the memory at half precision.

Weight tying with the output layer

Look again at Figure 1 of the introduction. The pipeline starts with a $V \times d$ matrix (embedding: id → vector) and ends with a $d \times V$ matrix (unembedding: vector → one score per id). Same shape, transposed. Both are maps between "which token" and "a point in $d$-dimensional space". Do we really need two of them?

Often, no. Weight tying (Press & Wolf, 2016; Inan et al., 2016) uses one matrix for both jobs: the logits are $h E^\top$, where $h$ is the final hidden vector. The score for token $i$ is then the dot product between $h$ and token $i$'s embedding row. Read that way, the output layer says: "predict tokens whose embeddings point in the direction the network ended up pointing". That is a sensible thing for the network to want.

Why tie

Three reasons. It removes $V \times d$ parameters, which for GPT-2 small would otherwise be another 31% of the model. It acts as a regulariser: every token's row now gets gradient both when the token appears as input and when it is the target, so rare tokens learn faster. And it tends to improve perplexity in small and medium models, which is why GPT-2 and many models of that generation tied.

Why untie

The input and output jobs are not quite the same. The input embedding wants tokens that are used similarly to be close. The output embedding wants tokens that are predicted in the same contexts to be close. These overlap heavily but not perfectly (think of "a" and "an": interchangeable as inputs, never interchangeable as predictions). In large models the parameter saving is negligible (a fraction of a percent), so the extra flexibility wins, and most large models today, including the Llama 3 family, use separate input and output matrices. Small models, where the table is a big share of the parameters, still commonly tie.

Embeddings are context-free (and how that gets fixed)

Here is something people find surprising. The embedding row for "bank" is the same vector whether the sentence is "the river bank" or "the bank approved the loan". The lookup does not know the context; it cannot, because it only receives the id. So at layer 0, every occurrence of a token carries an identical vector (plus its positional information, next chapter).

The transformer's job is to fix this. In the first block, attention lets the vector at "bank" read from "river" or "loan" and mix in what it finds; the MLP then transforms the mixed vector. By the last block, the two occurrences of "bank" have drifted to very different points, one near "shore" and "water" concepts, the other near "money" and "loan". The embedding gives every token a starting point; the blocks move it to where it belongs in this sentence. You can see this in Figure 3, and the attention chapter shows the mechanism.

Initialisation: why the numbers start small

Before training, $E$ has to contain something. GPT-2 initialises the token embedding from a normal distribution with mean 0 and standard deviation 0.02 (and the positional table with 0.01). Why so small, and why not zero?

Not zero, because then every token would start identical and the first layer's gradient would be the same for every row: symmetric, and stuck. Random breaks the symmetry. Small, because of what happens downstream. The embedding is the residual stream's starting value, and the attention scores in the first block are dot products of projections of these vectors. The dot product of two random $d$-dimensional vectors with per-coordinate standard deviation $\sigma$ has standard deviation about $\sigma^2 \sqrt{d}$.

The 0.02 is not sacred. Some models scale by $1/\sqrt{d}$, and models that use pre-normalisation are more forgiving. Some multiply the embedding output by $\sqrt{d}$ (the original transformer did this, with tied weights, so that the tied output logits and the input scale both come out reasonable). What matters is the principle: the initial residual stream should have a scale the rest of the network's initialisation expects. code/lumen/embeddings.py exposes the standard deviation as an argument so you can experiment.

Looking at embeddings: PCA, t-SNE, and a warning

A 768-dimensional table cannot be plotted. To look at it, people project it down to two dimensions. PCA finds the two directions of greatest variance and drops the rest; it is linear, deterministic, and preserves large-scale structure but usually squashes clusters into a blob because two directions cannot hold much of the variance. t-SNE (van der Maaten & Hinton, 2008) and its relative UMAP instead try to keep each point's nearest neighbours near it in the picture, which produces pretty, well-separated clusters.

The warning: those pictures are seductive and easy to over-read. t-SNE cluster sizes mean nothing, distances between clusters mean nothing, and different random seeds or perplexity settings give different maps from the same data (Wattenberg et al., 2016). A two-dimensional shadow of a 768-dimensional object discards almost everything. Use the plots to get a feel, then verify any claim with numbers: cosine similarities, nearest-neighbour lists, or a probe trained on the real vectors.

Practice

The companion file is code/lumen/embeddings.py. It has a from-scratch Embedding module, a cosine_neighbours(E, token, k) helper, and a flag for tied or untied output projections that the GPT-2 build in code/lumen/gpt2.py reads.

Further reading

• Mikolov et al. (2013). Efficient Estimation of Word Representations in Vector Space. word2vec and the analogy result.
• Pennington, Socher & Manning (2014). GloVe: Global Vectors for Word Representation. Embeddings from co-occurrence statistics; a clear account of why geometry emerges.
• Press & Wolf (2016). Using the Output Embedding to Improve Language Models. The case for weight tying.
• Nissim, van Noord & van der Goot (2020). Fair is Better than Sensational: Man is to Doctor as Woman is to Doctor. Why analogy results must exclude the inputs, and how easily they mislead.
• van der Maaten & Hinton (2008). Visualizing Data using t-SNE. The method, and Wattenberg, Viégas & Johnson (2016), How to Use t-SNE Effectively, on how not to read it.
• Bengio et al. (2003). A Neural Probabilistic Language Model. The origin of learned embeddings inside a language model.