The embedding layer — from token IDs to vectors

Where the numbers come from

Every layer discussed so far operates on vectors. Text is not vectors. Something has to make the conversion, and it happens exactly once, at the bottom.

The chain is short. Text is split into tokens. Each token is looked up in a vocabulary and becomes an integer. Each integer indexes a row of a big learned table, and that row is the vector — a 512-number description of the token, learned during training rather than designed.

The lookup deserves a moment, because it is often described as a matrix multiplication and that makes it sound more mysterious than it is. Multiplying a one-hot vector by the embedding matrix selects exactly one row. Every implementation does the selection directly. The parameters are real and substantial — 37,000 tokens × 512 dimensions is about 19 million numbers, a serious fraction of the base model — but the operation is a table read.

Two consequences are easy to miss and matter downstream.

The width is fixed for the entire network. Every layer takes 512 in and gives 512 back. This is not an aesthetic choice — it is what makes the stack a stack. Because the shape never changes, you can add a seventh layer, or a ninety-seventh, without redesigning anything. It is also what makes residual connections possible, since you cannot add x to sublayer(x) unless they are the same shape.

Only the bottom layer embeds. Encoder 2 does not receive tokens; it receives encoder 1's output. From layer two upward the input is just vectors that came from below, which is why "what does a token's vector mean at layer 9?" has no clean answer — it stopped being about that token alone at layer 1.

One thing the embedding cannot supply is order. The table is indexed by token identity, so the same word gets the same vector wherever it appears, and attention is blind to position. That gap is filled by adding a positional signal to these vectors before the first layer — covered in the attention course, and assumed from here on.

The first operation in any transformer converts discrete token IDs into continuous vectors. A tokenizer produces integer IDs — "The cat sat" might become [464, 3797, 3332] using GPT-2's BPE vocabulary. These integers are meaningless to the model's linear algebra. The embedding layer maps each integer to a dense vector that the transformer's attention and feed-forward layers can operate on.

The embedding matrix

The embedding layer is a lookup table implemented as a matrix E with shape (vocab_size, d_model). Each row of E is the embedding vector for one token in the vocabulary. Given a token ID i, the embedding is simply E[i] — a row selection, which is mathematically equivalent to multiplying a one-hot vector by E but implemented as an index operation for efficiency.

The matrix E is a learned parameter — its values are initialized randomly and updated during training via backpropagation, just like every other weight in the model. There is no hand-crafted structure; the model discovers what each token's vector should look like by learning to predict text.

Model dimensions and embedding sizes:

• GPT-2 Small: vocab 50,257 × d_model 768 = 38.6M embedding parameters
• GPT-2 XL: vocab 50,257 × d_model 1,600 = 80.4M embedding parameters
• BERT-base: vocab 30,522 × d_model 768 = 23.4M embedding parameters
• Llama 3 8B: vocab 128,256 × d_model 4,096 = 525.3M embedding parameters (6.5% of 8B total)
• Llama 3 70B: vocab 128,256 × d_model 8,192 = 1.05B embedding parameters (1.5% of 70B total)

For small models, the embedding matrix is a significant fraction of total parameters. GPT-2 Small's embeddings are 38.6M out of 124M total — 31%. For large models, the fraction shrinks because the transformer blocks (attention + FFN, which scale with d_model^2) grow much faster than the embedding matrix (which scales with vocab_size × d_model). Llama 3 70B's 1.05B embedding parameters are a rounding error against the 70B total.

The scaling factor

Before adding positional encodings, the original transformer (Vaswani et al. 2017) multiplies embedding vectors by sqrt(d_model). With $d_{\text{model}} = 512$, this multiplies every embedding by ~22.6. The reason: embedding vectors are initialized with small values (standard deviation ≈ 1/sqrt(d_model)), while sinusoidal positional encodings have values in [-1, 1]. Without the scaling factor, the positional signal would dominate the token identity signal. Multiplying by sqrt(d_model) brings the embedding magnitudes to a comparable scale. GPT-2 and most modern models use learned positional embeddings initialized to the same scale as token embeddings, so the scaling factor is unnecessary and not used.

What embeddings learn

After training, the embedding space organizes tokens by semantic and syntactic relationships. Tokens with similar meanings or similar grammatical roles end up near each other in the d_model-dimensional space. This structure is not designed — it emerges because placing similar tokens near each other makes next-token prediction easier.

from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch
import torch.nn.functional as F

model = GPT2LMHeadModel.from_pretrained("gpt2")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

embed_weights = model.transformer.wte.weight  # (50257, 768)
print(f"Embedding matrix shape: {embed_weights.shape}")
print(f"Embedding parameters: {embed_weights.numel():,}")  # 38,597,376

def get_embedding(word):
    ids = tokenizer.encode(word, add_special_tokens=False)
    if len(ids) != 1:
        print(f"Warning: '{word}' tokenizes to {len(ids)} tokens: {ids}")
    return embed_weights[ids[0]]

words = ["dog", "puppy", "cat", "king", "queen", "algorithm", "computer"]
embeddings = {w: get_embedding(w) for w in words}

print("\nCosine similarities:")
for i, w1 in enumerate(words):
    for w2 in words[i+1:]:
        sim = F.cosine_similarity(
            embeddings[w1].unsqueeze(0),
            embeddings[w2].unsqueeze(0)
        ).item()
        print(f"  cos({w1}, {w2}) = {sim:.4f}")

Typical output from GPT-2's trained embeddings:

  cos(dog, puppy) = 0.8734
  cos(dog, cat) = 0.8291
  cos(dog, king) = 0.4012
  cos(dog, algorithm) = 0.1847
  cos(king, queen) = 0.7956
  cos(algorithm, computer) = 0.6538

Semantically related words cluster together. dog/puppy/cat form a tight animal cluster. king/queen are close. algorithm and computer are closer to each other than to any animal. These relationships emerge purely from the language modeling objective — the model learns that dog and puppy appear in similar contexts and assigns them similar vectors.

Weight tying

Weight tying (Press & Wolf 2017) shares the embedding matrix E with the output projection matrix — the linear layer (the "language model head") that converts the transformer's final hidden states back into vocabulary-sized logits for next-token prediction.

Without weight tying, the model has two separate matrices:

• Input embedding: E_in with shape (vocab_size, d_model) — maps token IDs to vectors
• Output projection: W_out with shape (d_model, vocab_size) — maps hidden states to logits

With weight tying, $W_{\text{out}} = E_{\text{in}}^{T}$. The output logit for token i is the dot product between the final hidden state and token i's embedding vector. A token's embedding defines both what it looks like as input and what the model looks for when predicting it as output.

The parameter savings are substantial: vocab_size × d_model fewer parameters. For Llama 3 8B that's 525M parameters — the equivalent of removing an entire small model from the parameter count.

model = GPT2LMHeadModel.from_pretrained("gpt2")

embed_ptr = model.transformer.wte.weight.data_ptr()
head_ptr = model.lm_head.weight.data_ptr()
print(f"Embedding weight pointer:  {embed_ptr}")
print(f"LM head weight pointer:    {head_ptr}")
print(f"Same underlying tensor:    {embed_ptr == head_ptr}")  # True

# They are literally the same tensor in memory — not copies
print(f"Embedding shape: {model.transformer.wte.weight.shape}")  # (50257, 768)
print(f"LM head shape:   {model.lm_head.weight.shape}")          # (50257, 768)

GPT-2, GPT-3, T5, ALBERT, and many modern models use weight tying. Some architectures (notably Llama 2) keep the matrices separate, which adds parameters but allows the input and output representations to specialize independently. Whether to tie weights is an empirical choice — Press & Wolf (2017) showed that tying provides a regularization benefit and consistently matches or outperforms untied weights at the same model size.

The intuition behind weight tying: a good input representation for a token should also be a good target representation. If the embedding for "Paris" encodes the concept of the French capital, then when the model predicts the next token should be "Paris," the hidden state should point in the same direction as that embedding vector. Weight tying enforces this symmetry by construction, rather than hoping the model learns separate but aligned representations.

Subword embeddings and out-of-vocabulary handling

Modern transformers use subword tokenizers (BPE, WordPiece, or Unigram) so the embedding matrix never needs an [UNK] token. Every possible input is decomposed into known subword tokens. The word "transformers" might be represented as two tokens — ["transform", "ers"] — each with its own embedding. The model's attention and FFN layers then compose these subword representations into a contextual whole-word meaning.

This means the embedding matrix only needs entries for the subword vocabulary, not for every possible word. GPT-2's 50,257-token vocabulary covers all possible byte sequences (via byte-level BPE), so no input can ever be out-of-vocabulary. The embedding layer is guaranteed to produce a vector for any input the tokenizer generates.

The downside: rare words that are split into many subwords get weaker initial representations — they start as a bag of subword vectors that the model must compose through its layers. Common words that correspond to a single token get a dedicated, fully learned embedding from the start. This is one reason larger vocabularies (Llama 3's 128K tokens vs. GPT-2's 50K) tend to improve performance on rare and multilingual text — more words get single-token embeddings.

Embedding dimension and model capacity

The embedding dimension d_model is the single most important hyperparameter in a transformer — it determines the width of every representation throughout the network. Every hidden state, every attention output, and every residual stream value is a vector of d_model floats.

• GPT-2 Small (124M params): $d_{\text{model}} = 768$
• GPT-2 XL (1.5B params): $d_{\text{model}} = 1, 600$
• Llama 3 8B: $d_{\text{model}} = 4, 096$
• Llama 3 70B: $d_{\text{model}} = 8, 192$
• GPT-4 (rumored): $d_{\text{model}} = 12, 288$ or higher

Doubling d_model roughly quadruples the parameter count of the attention and FFN layers (they scale as d_model^2), while only doubling the embedding parameters (they scale as vocab_size × d_model). This is why the embedding matrix becomes a smaller fraction of total parameters as models grow — the quadratic layers dominate.

The embedding dimension also determines how much information each token representation can carry. A 768-dim vector can encode approximately $768 \times 32 = 24, 576$ bits of information at float32 precision (though the effective information content is much lower due to correlations between dimensions). Larger d_model means the model can represent finer-grained distinctions between tokens and maintain more information through the residual stream as layers process and transform the representations.

Inspecting embedding geometry

Embedding vectors don't just cluster by topic — they encode structured relationships. The classic word2vec finding (Mikolov et al. 2013) that king - man + woman ≈ queen (as a vector arithmetic operation) holds approximately in transformer embeddings too, though the relationship is weaker because transformer embeddings are only the starting point — the real relational reasoning happens in attention layers.

def analogy(a, b, c, tokenizer, embed_weights, top_k=5):
    """Compute a - b + c and find nearest tokens."""
    e_a = embed_weights[tokenizer.encode(a, add_special_tokens=False)[0]]
    e_b = embed_weights[tokenizer.encode(b, add_special_tokens=False)[0]]
    e_c = embed_weights[tokenizer.encode(c, add_special_tokens=False)[0]]

    target = e_a - e_b + e_c
    sims = F.cosine_similarity(target.unsqueeze(0), embed_weights)
    top_ids = sims.topk(top_k + 3).indices.tolist()

    skip = set(tokenizer.encode(w, add_special_tokens=False)[0] for w in [a, b, c])
    results = [tokenizer.decode([i]).strip() for i in top_ids if i not in skip]
    return results[:top_k]

print(analogy("king", "man", "woman", tokenizer, embed_weights))
# Typical: ['queen', 'princess', 'monarch', 'Queen', 'empress']

These relationships are a byproduct of the training objective. Because king and queen appear in similar linguistic contexts (with a gender offset that mirrors man/woman), gradient descent pushes their embeddings into a geometric arrangement that encodes that structure. The embedding layer converts discrete symbols into a continuous space where linear operations correspond to semantic relationships — and that continuity is what makes the rest of the transformer's computation possible.

The table described here is one of three distinct things the word embedding names, and it is the one inside a model. The other two — a contextual hidden state, and the single pooled vector a search system actually stores — belong to a different subject, and our Embeddings, End to End track is that subject end to end, from TF-IDF through to how today's embedding models are trained.