The Embedding Table and Its Geometry

By the end of the vocabulary-design post we had a tokenizer that turns text into a sequence of integer IDs. We spent a lot of energy getting those IDs right: choosing vocabulary sizes, worrying about fertility, understanding merge order. Notice something about what we actually produced.

Token 4821 means nothing. It is an arbitrary integer. You cannot add token 4821 to token 917 and get anything meaningful. You cannot compute a dot product between two integers. The token ID is a pointer into a table, and nothing more.

So how does the model start computing with these things?

The answer is an embedding table: a matrix of learned vectors, one row per token in the vocabulary. This is the very first thing that happens in a transformer. Before attention, before layer norms, before any of the machinery we will build up in Arc 4, every token ID gets swapped out for a dense vector. That vector is what the rest of the model actually sees.

This post is about that table: what it is mathematically, what training does to it, how it relates to the output layer, and what the resulting geometry looks like. We covered the ancestors of this idea in Post 2.3 (word2vec) and Post 2.4 (GloVe/FastText). Transformer embeddings live a different life.

Lookup as matrix multiply

Here is the entire mechanism. We have a vocabulary of $V$ tokens and a model dimension $d_\text{model}$. The embedding table $E$ is a matrix of shape $V \times d_\text{model}$. Each row $i$ is the learned embedding vector for token $i$.

When token ID $k$ comes in, we grab row $k$ from $E$. In PyTorch this is literally E[k], an indexing operation. But there is a cleaner way to think about it that matters later.

Define a one-hot vector $\mathbf{e}_k$ of length $V$: all zeros except a 1 in position $k$. Then:

Multiplying a one-hot vector by the embedding matrix selects a single row. That is all an embedding lookup is, really: a matrix multiply where one of the inputs is extremely sparse (exactly one nonzero entry).

Try it yourself:

In practice, nobody actually constructs the one-hot vector and does a full matmul. nn.Embedding in PyTorch is implemented as a direct index into the weight tensor, which is much faster. The mathematical equivalence matters because it tells you the embedding lookup is a linear map from a discrete, sparse, $V$-dimensional space into a dense, continuous, $d_\text{model}$-dimensional space. If you've been following the vectors-and-spaces post, that should feel significant. The token enters the model as a basis vector in token-space and exits as a point in representation-space.

How big is this thing?

The parameter count is straightforward: $V \times d_\text{model}$. Real numbers:

For Llama 3 8B, the embedding table alone accounts for over 500 million parameters, roughly 7% of the total model. When people expand vocabularies (as we discussed in Post 3.4), this is the cost they are paying. Every new token adds a full $d_\text{model}$-dimensional vector of learnable weights.

But that 7% is not a universal number. The same kind of table can be 30% of a small model or well under 1% of a huge one. The transformer stack scales roughly with depth times $d_\text{model}^2$ while the embedding scales linearly with $d_\text{model}$ (the vocabulary stays roughly fixed once you've picked a tokenizer). So as models grow, the embedding's slice of the budget shrinks fast.

What training does to embeddings

At initialization, embedding vectors are random. Most frameworks use a normal distribution with a small standard deviation (something like $\sigma = 0.02$). At this point, every token's vector is basically noise. "the" and "cat" and "}" all look the same.

Then training happens. The model is predicting next tokens, and backpropagation pushes gradients through every layer, all the way back to the embedding table. When token $k$ appears in a training batch, only row $k$ receives a gradient update. Tokens that never show up in a batch get no update at all that step. Rare tokens get trained much less than common ones, and their embeddings stay closer to random initialization. I'll come back to this when we talk about geometry.

Over millions of training steps, tokens that serve similar predictive roles end up with similar embedding vectors. "dog" and "cat" both predict articles, verbs of motion, possessive constructions. The gradients that shape them are drawn from overlapping contexts, so they drift toward overlapping regions of the embedding space.

If you've been through the word2vec posts, this probably sounds familiar. Word2vec also learned vectors by predicting context words, and we saw the famous arithmetic: $\text{king} - \text{man} + \text{woman} \approx \text{queen}$. Transformer embeddings inherit some of that structure, but there's a critical difference. In word2vec the embedding is the model, and the entire training objective directly shapes the vectors. In a transformer, the embedding table is just the first layer, shaped by gradients that have flowed backward through dozens of attention and FFN layers. The embedding of "cat" in GPT-2 reflects not just what contexts "cat" appears in, but how the entire 12-layer network has learned to process those contexts.

Transformer embeddings carry less clean linear structure than word2vec embeddings as a result. Vector arithmetic still sort of works for common words, but it gets muddier, because the embedding doesn't need to encode everything about the token. It just needs to provide a useful starting point for the layers above.

Embedding vs. unembedding

The embedding table maps token IDs to vectors at the start of the model. At the end, we need the reverse: map the final hidden state back to a probability distribution over the vocabulary. This is the unembedding layer (also called lm_head).

The unembedding matrix $U$ has shape $d_\text{model} \times V$. Given a final hidden state $\mathbf{h}$ (a vector of dimension $d_\text{model}$), the logits are:

Each column of $U$ is a "target vector" for one token. The dot product $\mathbf{h} \cdot U[:,k]$ measures how much the hidden state points toward token $k$. We softmax these logits to get probabilities and sample or argmax to pick the next token.

Here is the interesting question: should $E$ and $U$ share weights?

Weight tying

Press and Wolf showed in 2017 that you can set $U = E^\top$ and get models that train just as well with significantly fewer parameters. This is called weight tying or shared embeddings. Intuitively, the same vector that represents "cat" in the input space is also the target vector for "cat" in the output space. If the model's hidden state points toward the "cat" embedding, the model predicts "cat".

GPT-2 used weight tying. BERT used weight tying. It was standard practice for years. The parameter savings are real: instead of $2 \times V \times d_\text{model}$, you pay $V \times d_\text{model}$ once.

But most modern models have untied the two matrices. Llama, Llama 2, Llama 3, Mistral, and Qwen all use separate embedding and unembedding weights. Why?

The input and output spaces serve different functions. The embedding table maps discrete token IDs into a space where the first attention layer can start finding useful patterns. The unembedding layer maps the final hidden state into a space where token probabilities are well calibrated. Those are different jobs, and forcing them to share a single matrix means both have to compromise.

There's also a practical argument. The parameter savings matter less as models get larger. Llama 2 7B has 6.7 billion parameters; untying the embedding adds 131 million, less than 2% of the total. The representational flexibility is worth it.

There's also a middle ground: pseudo-inverse tying, where the unembedding matrix is initialized as the pseudo-inverse of the embedding matrix but allowed to diverge during training. This gives you the warm-start benefit of shared structure without the constraint of permanent coupling.

Geometry: the narrow cone

So we have billions of gradient updates shaping the embedding table. What does the resulting geometry actually look like? This is where things get weird.

If you take a trained embedding table and compute the cosine similarity between random pairs of tokens, you'd expect the average to be near zero. In a well-distributed high-dimensional space, random vectors are approximately orthogonal. That's the basic geometric intuition from the norms-and-similarity post: in high dimensions, most directions are far apart.

That's not what happens. In practice, the average cosine similarity between random token pairs in a trained transformer is surprisingly high, often around 0.5 or even 0.6. Embeddings cluster in a narrow cone pointing away from the origin, and almost all of them land in roughly the same angular neighborhood.

This phenomenon is called anisotropy, and it has been observed consistently across models and architectures. The representation space is not being used efficiently. Instead of spreading out to fill the available dimensions, embeddings collapse into a low-dimensional subspace.

Why does this happen? A few contributing factors:

Frequency imbalance. Common tokens ("the", "a", "is") appear in nearly every training batch and accumulate massive gradient updates. They develop large-norm embedding vectors. Rare tokens stay near the origin with small norms. The common tokens dominate the space, pulling the mean embedding away from zero, and every token's vector ends up with a substantial component in the direction of that mean. Timkey and van Schijndel found that much of this effect concentrates in a handful of "rogue dimensions" with outsized variance. Zero out those dimensions and the embedding space becomes much more isotropic.

The softmax bottleneck. The unembedding layer maps hidden states to logits via dot products. If the model uses weight tying, the embedding vectors must simultaneously serve as input representations and output targets. The output softmax pushes token vectors apart only along dimensions that help discriminate between likely next tokens. Dimensions that don't contribute to prediction are wasted, and the effective dimensionality drops.

Self-attention amplifies it. Godey et al. showed in 2024 that the self-attention mechanism itself is inherently anisotropic. Even with isotropic inputs, self-attention tends to produce outputs that cluster in a narrow cone. The effect compounds across layers, but it starts at the embedding level.

You can see this yourself. Load a real model's embedding table, grab vectors for a diverse set of tokens, and look at the pairwise cosine similarities.

import torch
from transformers import AutoTokenizer, AutoModel
 
model_name = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, torch_dtype=torch.float32)
 
# Grab the embedding table
E = model.embed_tokens.weight.detach()  # shape: (vocab_size, d_model)
 
# Pick some diverse tokens
words = ["the", "cat", "dog", "Python", "def", "return",
         ".", ",", "!", "1234", "Obama", "quantum"]
ids = [tokenizer.encode(w, add_special_tokens=False)[0] for w in words]
vecs = E[ids]  # shape: (12, d_model)
 
# Pairwise cosine similarity
norms = vecs / vecs.norm(dim=1, keepdim=True)
cos_sim = norms @ norms.T
 
print("Pairwise cosine similarities:")
for i, w in enumerate(words):
    print(f"{w:>10}", *[f"{cos_sim[i,j]:6.2f}" for j in range(len(words))])
 
# Global anisotropy: average cosine of 1000 random pairs
rand_idx = torch.randint(0, E.shape[0], (1000, 2))
v1, v2 = E[rand_idx[:, 0]], E[rand_idx[:, 1]]
avg_cos = torch.nn.functional.cosine_similarity(v1, v2).mean()
 
print(f"\nAverage cosine similarity (1000 random pairs): {avg_cos:.4f}")
print(f"Embedding shape: {E.shape}")
print(f"Mean norm: {E.norm(dim=1).mean():.4f}")
print(f"Std of norms: {E.norm(dim=1).std():.4f}")

When you run this you'll typically see that "cat" and "dog" have a cosine similarity around 0.7 or 0.8, which makes sense. But "cat" and "," might also be at 0.4 or 0.5, which is weirdly high for two tokens with nothing in common. That's anisotropy. The global average for random pairs will likely be in the 0.3 to 0.6 range, depending on the model. In a truly isotropic space with 2048 dimensions, you'd expect something close to 0.

Does anisotropy hurt?

It depends on what you're doing. For the transformer's own internal processing, anisotropy isn't necessarily harmful. The model has learned to work within this geometry. The attention mechanism and FFN layers have adapted to the cone-shaped input space and route information effectively despite the geometric inefficiency.

When you try to use embeddings directly for tasks like retrieval or semantic search, anisotropy is a real problem. If all vectors point in roughly the same direction, cosine similarity loses its discriminative power. This is one reason dedicated embedding models (covered later in the embeddings-and-vector-search post) apply extra training specifically to spread the vectors out.

Worked example

Let me walk through a more complete worked example. We will load a model's embedding table, compute similarities across three groups of tokens (punctuation, common English, code keywords), and look at the structure.

import torch
from transformers import AutoTokenizer, AutoModel
 
model_name = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, torch_dtype=torch.float32)
E = model.embed_tokens.weight.detach()
 
# Three groups of tokens
punctuation = [".", ",", "!", "?", ";", ":"]
common_en   = ["the", "and", "is", "was", "they", "from"]
code_tokens = ["def", "return", "import", "class", "self", "None"]
 
def get_ids(words):
    return [tokenizer.encode(w, add_special_tokens=False)[0] for w in words]
 
groups = {
    "punct":   get_ids(punctuation),
    "english": get_ids(common_en),
    "code":    get_ids(code_tokens),
}
 
# Average within-group and between-group cosine similarity
def avg_cos(ids_a, ids_b, emb):
    va = emb[ids_a]
    vb = emb[ids_b]
    na = va / va.norm(dim=1, keepdim=True)
    nb = vb / vb.norm(dim=1, keepdim=True)
    sims = na @ nb.T
    if ids_a == ids_b:
        mask = ~torch.eye(len(ids_a), dtype=torch.bool)
        return sims[mask].mean().item()
    return sims.mean().item()
 
print("Within-group avg cosine similarity:")
for name, ids in groups.items():
    print(f"  {name}: {avg_cos(ids, ids, E):.4f}")
 
print("\nBetween-group avg cosine similarity:")
names = list(groups.keys())
for i in range(len(names)):
    for j in range(i + 1, len(names)):
        sim = avg_cos(groups[names[i]], groups[names[j]], E)
        print(f"  {names[i]} vs {names[j]}: {sim:.4f}")
 
# Token norms by group
print("\nAverage embedding norm by group:")
for name, ids in groups.items():
    norms = E[ids].norm(dim=1)
    print(f"  {name}: mean={norms.mean():.3f}, std={norms.std():.3f}")

What you typically find: within-group similarity is higher than between-group similarity, which is good. Punctuation tokens cluster together. Code keywords cluster together. Common English words cluster together. The structure is real. But the between-group similarities are also non-trivially positive (often 0.3 to 0.5), which is the anisotropy showing through.

You will also notice that embedding norms vary a lot. Common tokens like "the" tend to have larger norms than rare code tokens. Frequent tokens receive more gradient updates and their vectors grow. Some researchers have argued that embedding norm encodes a kind of frequency signal that the model uses downstream.

Misconceptions

"Embeddings are basically word2vec." They share the same core idea (tokens as learned dense vectors) but they're shaped by fundamentally different training signals. Word2vec embeddings are the final product of a shallow model. Transformer embeddings are the input layer of a deep model, shaped by gradients flowing backward through dozens of layers of attention and feedforward computation. The resulting geometry is different: less clean linear structure, more task-adapted representations.

"The embedding table is just a lookup." Technically true at inference time, but it misses what makes it interesting. The embedding table is a learned linear map from a $V$-dimensional discrete space to a $d_\text{model}$-dimensional continuous space. Calling it "just a lookup" is like calling a neural network "just matrix multiplications." Correct, but it leaves out all the structure.

"High cosine similarity means the tokens are semantically similar." In a well-spread embedding space this would be true. In an anisotropic one, every random pair has a cosine around 0.5 just from the cone geometry. The relative cosine between pairs still carries information, but the absolute number doesn't, so you have to compare to a baseline. This is why papers that use transformer embeddings directly for similarity tasks usually subtract out the mean or apply whitening first.

What's next

We now have the full picture of what happens before the first transformer layer. Text goes in, the tokenizer chops it into subword tokens (Posts 3.1 through 3.4), each token gets an integer ID, and that ID selects a learned vector from the embedding table. That vector lives in $\mathbb{R}^{d_\text{model}}$, and it is what the rest of the model actually sees.

The next post covers Special Tokens, Chat Templates, and Input Formatting: the BOS, EOS, and padding markers, the chat templates, and the formatting conventions that turn a multi-turn conversation into a flat sequence of token IDs.

Additional reading (and watching)

• Press, O. & Wolf, L. (2017). Using the Output Embedding to Improve Language Models. EACL 2017. The paper that introduced weight tying and made it the default for years.
• Ethayarajh, K. (2019). How Contextual are Contextualized Word Representations?. EMNLP 2019. Early characterization of anisotropy in transformer representations.
• Godey, N., Clergerie, É., & Sagot, B. (2024). Anisotropy Is Inherent to Self-Attention in Transformers. EACL 2024. Showed that self-attention is inherently anisotropic even with isotropic inputs, explaining why the cone geometry starts at the embedding level.
• Timkey, W. & van Schijndel, M. (2021). All Bark and No Bite: Rogue Dimensions in Transformer Language Models. EMNLP 2021. Showed that a handful of "rogue dimensions" with outsized variance drive much of the observed anisotropy, and that zeroing them out recovers near-isotropic geometry.
• Su, J. et al. (2022). Whitening Sentence Representations for Better Semantics and Faster Retrieval. Proposed a simple whitening transformation to counteract anisotropy in embedding spaces, widely adopted in production embedding pipelines.