Embedding — Module 11.1

How This Lesson Fits the Module & Volume

After the tokenizer emits IDs, the first neural layer of a language model maps each ID to a dense vector—the token embedding. Vol. 09 taught static Word2Vec/GloVe and nn.Embedding; Vol. 10 added positional embeddings. Here embeddings are the entry point to the causal Transformer that will produce contextual hidden states and logits.

The next lecture, Embedding Space, studies geometry; this one focuses on the lookup mechanism and wiring.

Learning Objectives

By the end of this lesson, students should be able to:

Describe token embedding as a learned |V|×d lookup table indexed by token IDs.
Combine token + positional (and optional segment) embeddings as model input.
Implement embedding lookup in PyTorch and inspect shapes through a tiny LM stem.
Explain weight tying between input embeddings and the LM head.
Contrast static pretrained vectors with end-to-end LM-trained embeddings.
Connect embedding dimension d_model to the rest of the Transformer width.

Definition

An embedding (token embedding) is a dense vector representation of a vocabulary ID, usually stored as a row of a matrix E ∈ R^|V|×d and retrieved by index. In Transformers, the input to layer 0 is typically token embedding + positional encoding/embedding (plus scaling/dropout).

From ID to Vector

Token ID

Integer from tokenizer.

Lookup

Row of E: e = E[id].

+ Position

Add PE / RoPE prep.

Stack

Enter Transformer blocks.

Piece	Shape	Role
Token embedding E	(\|V\|, d)	Meaning prior per type
Positional embedding	(T_max, d) or RoPE	Order / distance cues
Input to block 0	(B, T, d)	Sum (or concat schemes)
LM head W	(d, \|V\|) or tied E^T	IDs ← hidden states

Static (Vol. 09)

Word2Vec / GloVe / FastText.
Frozen or lightly tuned.
One vector per word type.

LM Input Embeddings

Trained with next-token loss.
Subword rows, not only words.
Still non-contextual until blocks run.

Contextual States

Post-attention hidden vectors.
Same ID → different states.
What people mean by “contextual embeddings.”

Code: Embedding Stem + Tied Head

import torch
from torch import nn
import math

class EmbeddingStem(nn.Module):
    def __init__(self, vocab_size, d_model, max_len=512, tie_weights=True):
        super().__init__()
        self.tok = nn.Embedding(vocab_size, d_model)
        self.pos = nn.Embedding(max_len, d_model)
        self.drop = nn.Dropout(0.1)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        if tie_weights:
            self.lm_head.weight = self.tok.weight  # tie

    def forward(self, idx):
        B, T = idx.shape
        positions = torch.arange(T, device=idx.device)
        x = self.tok(idx) * math.sqrt(self.tok.embedding_dim)
        x = self.drop(x + self.pos(positions))
        # ... Transformer blocks would go here ...
        return self.lm_head(x)

stem = EmbeddingStem(vocab_size=1000, d_model=128)
ids = torch.randint(0, 1000, (2, 16))
logits = stem(ids)
print(logits.shape)  # (2, 16, 1000)
print("tied:", stem.lm_head.weight.data_ptr() == stem.tok.weight.data_ptr())

Why Embeddings

Dense, trainable, GPU-friendly.
Share statistical strength across contexts.
Far smaller than one-hot inputs.

Caveats

|V|×d can dominate small models.
Input embedding alone is not contextual.
Extending vocab needs careful init.

Common Misconception

“The embedding layer outputs contextual meaning like BERT’s final states.” The embedding table only provides a type-level starting vector. Context enters through attention and feed-forward stacks. Saying “the embedding of bank in this sentence” usually refers to a hidden state—not the raw lookup row.

Knowledge Check

Short Answer: What is the shape of a token embedding matrix? Answer: (|V|, d) — vocab size by embedding dimension.
True/False: Embedding lookup is equivalent to multiplying by a one-hot vector. Answer: True (without materializing the one-hot).
Multiple Choice: Input to the first Transformer block is typically: (a) raw IDs, (b) token (+ positional) vectors, (c) softmax probs. Answer: (b).
Short Answer: What is weight tying? Answer: Sharing the token embedding matrix with the LM-head projection weights.
True/False: Two identical token IDs always keep identical vectors after layer 12. Answer: False—context changes hidden states.
Multiple Choice: Positional embeddings exist because: (a) softmax needs them, (b) attention alone is permutation-tolerant without position cues, (c) vocab size depends on T. Answer: (b).
Short Answer: Name one difference from Word2Vec vectors. Answer: LM embeddings are trained with next-token loss on subword IDs (and used inside a deep stack)—any clear contrast.
Short Answer: Why scale embeddings by sqrt(d) in some implementations? Answer: Stabilizes magnitudes when adding positional encodings / following Transformer conventions.
Multiple Choice: Extending the tokenizer vocab requires: (a) nothing, (b) resizing embedding (and usually LM head) rows, (c) dropping LayerNorm. Answer: (b).
True/False: d_model is the width shared by embeddings and attention blocks. Answer: True (in standard designs).

Key Takeaways

Token embeddings map IDs → dense d-dimensional vectors via a learned table.
Positions are fused early; then the Transformer creates contextual states.
Tying embeddings to the LM head is a common parameter-saving trick.
Do not confuse lookup embeddings with post-stack contextual representations.
Next: Embedding Space—geometry, similarity, and structure.

Trainer’s Guide

Hands-on idea: Print model.get_input_embeddings().weight.shape on GPT-2 and estimate parameter count vs total model size.

Discussion prompt: If embeddings are tied, what happens to the LM head when you fine-tune only the last layers?

Recap: Embeddings lift token IDs into the vector space the Transformer operates on. Continue with Embedding Space.