How Embeddings Are Learned

When a model is initialised, every token in the vocabulary gets assigned a random vector — just noise. A randomly initialised vector for cat might be [0.03, -0.11, 0.07, ...], while dog gets [-0.08, 0.14, -0.02, ...]. These two vectors are not especially close. They have no meaning yet.

By the end of training on a large corpus, the same vectors become neighbours in a high-dimensional space. No one hand-crafted these numbers. The model adjusted them, one gradient step at a time, to reduce prediction error. The geometry is a byproduct of the training objective.

Because the right geometry depends on the task, the language, and the domain. There is no universal "correct" vector for bank — the right representation depends on whether the corpus is financial news, river ecology papers, or general English conversation. A pre-fixed table could not capture these specific patterns.

Learning the embedding table means the model discovers the geometry that best serves the training objective. Changing the task or data changes what gets learned.

"A word is known by the company it keeps" — J.R. Firth, 1957. This is the distributional hypothesis, and it is the conceptual foundation of learned embeddings.

If cat and dog appear repeatedly in similar sentence contexts — "The [word] sat on the mat", "The [word] slept on the sofa", "She gave the [word] some food" — then during training, the gradients for both tokens will be nudged in similar directions. Over many steps, their vectors converge.

Contexts containing "cat":
  "The cat sat on the mat."
  "The cat slept on the sofa."
  "She gave the cat some food."

Contexts containing "dog":
  "The dog sat on the mat."
  "The dog slept on the sofa."
  "She gave the dog some food."

→ Similar contexts → similar gradient direction → similar vectors

The embedding table is a matrix $E \in \mathbb{R}^{V \times d}$ where $V$ is vocabulary size and $d$ is embedding dimensionality. Each row $E[t]$ is the vector for token $t$. This matrix is a learnable parameter — it is updated by gradient descent exactly like the weights of any other layer.

Token ID vs vector: which is which?

Think of the embedding table as a database table. To keep it small, each vector below has just 4 numbers:

Illustrative values only. Real embedding values are learned during model training.

Token ID 1  → identifies the row for "cat"
Row 1       = [ 0.80, −0.30, 0.40, 0.15 ]   ← the vector for "cat"

A token ID does not contain vector numbers. ID 1 is not the vector [0.80, −0.30, 0.40, 0.15]; it simply tells the model which vector to retrieve.

Where do the vector numbers come from?

1. Initially, the values are typically set by a numerical initialization scheme.
2. During training, the model adjusts them through optimization and backpropagation (the update rule above).
3. Over time, the values become useful for the model's language-processing task.

In a real LLM each vector may contain hundreds or thousands of numbers. Token ID = which row? Vector = the learned numbers stored in that row. What “4 numbers” means is explained in Vector Dimensions.

Each training iteration follows the same loop. The embedding table is updated as a side-effect of training the model on its objective (next-token prediction, masked language modelling, or another task).

1. Sample a batch of text from the corpus.
2. Tokenize → convert to token IDs.
3. Look up embedding rows for those IDs: $x_i = E[t_i]$ for each token $t_i$ in the batch.
4. Pass embeddings through the model layers (attention, feed-forward, etc.).
5. Compute the training loss $\mathcal{L}$ (e.g. cross-entropy against the target tokens).
6. Backpropagate: compute $\nabla_{E[t_i]} \mathcal{L}$ for each token $t_i$ seen in the batch.
7. Optimizer updates those embedding rows: $E[t_i] \leftarrow E[t_i] - \eta \cdot \nabla_{E[t_i]} \mathcal{L}$.
8. Repeat for millions of batches.

Consider a batch containing tokens $\{t_1, t_2, \ldots, t_B\}$ (with possible repetitions). The full gradient matrix $\nabla_E \mathcal{L}$ has the same shape as $E$ — but most rows are exactly zero.

Consider a minimal vocabulary of two tokens: cat (ID 0) and dog (ID 1). The embedding table starts as:

Training repeatedly sees the sentences "the cat sat on the mat" and "the dog sat on the mat". For the target word sat, both cat and dog are the preceding context. The gradient updates for both rows point in directions that better predict sat from the context.

Each loop through the training data performs one gradient update. Over millions of iterations, the embedding table is gradually shaped by the training signal.

Each loop iteration slightly adjusts the vectors for tokens seen in that batch. After millions of iterations the geometry reflects distributional patterns in the corpus.

The code below implements one step of embedding learning from scratch — no PyTorch, just NumPy arithmetic. It shows the gradient computation, sparse update, and how two tokens in similar contexts converge.

import numpy as np

# ── Toy setup ─────────────────────────────────────────────────────────────
# Vocabulary: {cat: 0, dog: 1, sat: 2}
# Task: given context token, predict target token (simplified skip-gram)
VOCAB_SIZE = 3    # cat, dog, sat
EMBED_DIM   = 2   # 2D for easy visualisation
LR          = 0.1

# Initialise embedding table randomly
rng = np.random.default_rng(42)
E = rng.uniform(-0.5, 0.5, (VOCAB_SIZE, EMBED_DIM))

# ── Simplified forward pass ────────────────────────────────────────────────
# Predict target from context via dot product + softmax
def forward(E, context_id, target_id):
    """Return loss and probability of target given context."""
    context_vec = E[context_id]            # shape: (d,)
    logits = E @ context_vec               # shape: (V,)  — dot with all rows
    logits -= logits.max()                 # numerical stability
    probs = np.exp(logits) / np.exp(logits).sum()
    loss = -np.log(probs[target_id] + 1e-9)
    return loss, probs

# ── Training loop ──────────────────────────────────────────────────────────
# Training data: (context, target) pairs
# "cat sat" → (cat→sat) and "dog sat" → (dog→sat)
data = [(0, 2), (1, 2)] * 500    # 1000 pairs, same pattern for both

for step, (ctx, tgt) in enumerate(data):
    loss, probs = forward(E, ctx, tgt)

    # ── Gradient computation ──────────────────────────────────────────────
    # d_loss / d_logits = probs - one_hot(target)
    d_logits = probs.copy()
    d_logits[tgt] -= 1.0                  # gradient of cross-entropy + softmax

    # d_loss / d_E[context] = E.T @ d_logits  (from E @ context_vec = logits)
    d_E_ctx = E.T @ d_logits              # shape: (d,)

    # ── Sparse update: only context row changes ───────────────────────────
    E[ctx] -= LR * d_E_ctx

    if step % 200 == 0:
        sim = np.dot(E[0], E[1]) / (np.linalg.norm(E[0]) * np.linalg.norm(E[1]))
        print(f"step {step:4d}  loss={loss:.3f}  cos(cat,dog)={sim:.3f}")

print("\nFinal cosine similarity between cat and dog:", end=" ")
sim = np.dot(E[0], E[1]) / (np.linalg.norm(E[0]) * np.linalg.norm(E[1]))
print(f"{sim:.4f}  (should be close to 1.0 — same contexts)")

The geometry of learned embeddings depends heavily on the training objective, data, and whether the table is frozen or fine-tuned downstream.

Training shapes the geometry, but it also encodes the flaws of the training data and process. Understanding these limitations is essential for responsible use.

• Bias amplification: Embeddings trained on biased corpora reflect those biases. Occupational stereotypes, racial associations, and gender biases found in text are baked into the vectors.
• Infrequent tokens: Rare tokens receive far fewer gradient updates. Their rows may not converge — their geometry can be nearly random relative to frequent tokens.
• Catastrophic forgetting: Fine-tuning on a new domain can shift embedding rows away from their pretrained geometry, degrading performance on the original domain.
• Non-interpretable dimensions: There is no clean correspondence between individual dimensions and human-readable concepts. The geometry emerges from the training signal, not from human design.
• Static embeddings are context-free: A single trained row for "bank" cannot capture both "river bank" and "financial bank" — the same vector is used regardless of context.

Common misconceptions to tackle:

• "Embeddings are programmed with word meanings" — No: they start random and are adjusted by gradient descent to reduce prediction error. Meaning is never explicitly encoded.
• "Training makes all similar words neighbours" — No: "similarity" in embedding space means distributional similarity (similar contexts), not conceptual or semantic similarity in every sense. Antonyms can be neighbours.
• "A token's embedding is stable across all models" — No: the same word has different embeddings in BERT vs GPT-2 vs LLaMA because the training objectives, data, and architectures differ.
• "You can look at an embedding vector and read its meaning" — No: individual dimensions are not interpretable. The meaning lives in the geometric relationships between vectors.
• "Random init means training always produces the same geometry" — No: training is stochastic. Different random seeds, batch orderings, and hyperparameters produce different final geometries, though they may encode similar information.

Summary