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. Those IDs label vocabulary entries. Adding token 4821 to token 917 tells us nothing about the relationship between their text pieces; the integers are indices into a table.

An embedding table holds a learned vector for each token ID. The lookup supplies the vectors that attention and feedforward layers will transform in Arc 4.

Like word2vec and GloVe/FastText, a transformer learns these vectors during training. Its input table is part of a deeper network, trained jointly with the layers that follow it.

Lookup as matrix multiply

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$.

Token ID $k$ selects row $k$ of $E$, written E[k] in PyTorch. The same lookup can be expressed as a matrix multiplication.

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

PyTorch's nn.Embedding indexes the weight tensor without constructing a one-hot vector. The equivalent multiplication connects lookup to the linear maps introduced earlier: a basis vector in $V$ dimensions selects a vector in $d_\text{model}$ dimensions.

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.

The parameter share depends on the architecture. At fixed vocabulary size, the embedding table grows linearly with $d_\text{model}$, while the transformer stack grows roughly with depth times $d_\text{model}^2$. Increasing width or depth can therefore reduce the fraction spent on embeddings.

What training does to embeddings

At initialization, embedding vectors are random. Many transformer implementations use a small normal initialization, but the scale is model-specific. The random rows differ numerically; they have not yet learned token-specific roles.

Backpropagation reaches the embedding table through the model. The input lookup contributes gradients to rows used in the batch. If weights are tied to a full-softmax output layer, that output use can give every row a gradient. Momentum and weight decay can also change a row even when it was not looked up. Token frequency matters, but occurrence count is not the same as optimizer update count.

Training can give tokens with overlapping predictive roles similar vectors. For example, "dog" and "cat" share some contexts, providing related training signals. Their final geometry also depends on how the rest of the network uses those inputs.

Word2vec directly trains word and context vectors through a shallow prediction objective. A transformer's input rows receive gradients through its attention and feedforward layers.

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.

Weight tying

Press and Wolf (2017) studied weight tying, setting $U = E^\top$, and reported improved language-model results while reducing parameters. With this constraint, the same vector used to represent "cat" at the input also supplies its output score.

GPT-2 and BERT use weight tying. Storing one matrix instead of two reduces vocabulary parameters from $2 \times V \times d_\text{model}$ to $V \times d_\text{model}$.

Some models, including Llama 2 and Llama 3, use separate input and output matrices. Other families and model sizes retain tying. Check the checkpoint configuration rather than assuming that one convention has replaced the other.

The input table supplies vectors for the first layers; the output matrix converts final hidden states to token scores. Weight tying requires one matrix to serve both functions.

The parameter saving is $Vd$. Its share of the model depends on vocabulary size, width, and the size of the transformer stack. Whether independent matrices are worth that cost is an empirical design choice.

Tying constrains the two operations to use one matrix throughout training. Initializing two separate matrices equally and then training them independently is a different choice.

Geometry: the narrow cone

What geometry does training produce? We need to distinguish the input embedding table, which has one fixed row per token, from contextual hidden states, which change with the surrounding sequence and layer. Results about one do not automatically apply to the other.

An isotropic distribution has no preferred direction. Independent normalized vectors drawn uniformly in high dimensions usually have cosine similarity near zero. A distribution with a shared directional component can instead have positive average cosine, one form of anisotropy.

Studies of contextual representations have found anisotropy and dimensions that dominate cosine similarity. Godey et al. investigate its relationship with transformer mechanisms; Timkey and van Schijndel show why raw cosine can mislead in contextual spaces.

To inspect an input table, load a checkpoint and measure pairwise cosine similarities for selected token rows:

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.get_input_embeddings().weight.detach()  # shape: (vocab_size, d_model)
 
# Pick some diverse tokens
words = ["the", "cat", "dog", "Python", "def", "return",
         ".", ",", "!", "1234", "Obama", "quantum"]
# Preserve every subtoken rather than mislabeling the first piece as a word.
ids = list(dict.fromkeys(
    token_id for word in words
    for token_id in tokenizer.encode(word, add_special_tokens=False)
))
pieces = tokenizer.convert_ids_to_tokens(ids)
vecs = E[ids]
 
# Pairwise cosine similarity
norms = vecs / vecs.norm(dim=1, keepdim=True)
cos_sim = norms @ norms.T
 
print("Pairwise cosine similarities:")
for i, piece in enumerate(pieces):
    print(f"{piece:>10}", *[f"{cos_sim[i,j]:6.2f}" for j in range(len(ids))])
 
# 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}")

Compare the measured similarities to the random-pair baseline from that same table. Report the checkpoint, tokenizer, token selection, and whether special tokens are included. A small group of hand-selected words cannot establish the geometry of the entire vocabulary.

Does anisotropy hurt?

Anisotropy is not automatically a defect for the language model: its layers learn to process those representations. It can complicate comparisons based on raw cosine, especially when a few shared components dominate the score.

For retrieval, evaluate a model trained for sentence or document embeddings rather than assuming that an average of input-token rows is an effective semantic representation. Centering or whitening changes the geometry; whether it helps is a task-specific measurement.

Worked example

We can inspect a model's embedding table by comparing three groups of tokens: punctuation, common English, and code keywords.

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.get_input_embeddings().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):
    ids = []
    for word in words:
        encoded = tokenizer.encode(word, add_special_tokens=False)
        if len(encoded) != 1:
            raise ValueError(f"{word!r} is not a single token: {encoded}")
        ids.append(encoded[0])
    return ids
 
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}")

This code tests whether within-group similarities exceed between-group similarities in the selected checkpoint. It deliberately rejects multi-token words, so replace any rejected item with an actual token rather than silently using its first piece.

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 embedding table is just a lookup." Inference selects a row by index. Training changes the values in those rows, so the table also contains learned information about how tokens should enter the network.

"High cosine similarity proves semantic similarity." Compare scores to a baseline for the same representation. Shared directions, tokenization, and the training objective all affect cosine. A useful similarity measure needs validation on the intended task.

What's next

The tokenizer produces IDs, and the input table supplies their learned vectors. A conversation also needs markers that distinguish messages and roles.

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. Studies shared input/output weights and their effects on language-model quality and parameter count.
• 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. Investigates how self-attention can produce anisotropy, including in untrained models; this is not a universal measurement of input embedding tables.
• 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, evaluated for sentence-representation tasks.