Embedding layers and tied weights

Overview

Embedding layers convert discrete token IDs into continuous vector representations, while weight tying shares parameters between the input embedding and output projection layers, reducing model size and improving generalization.

[!intuition] The Core Idea

Think of embedings as a lookup table where each token ID is a row number, and you retrieve that row's vector. The token "cat" (ID=42) always maps to the same learned vector $\mathbf{e}_{42} \in \mathbb{R}^{d}$.

Why tie weights? The input embedding asks "what does this token mean?" and the output projection asks "which token should come next?" These are inverse operations—so using the same weight matrix (transposed) forces the model to learn a consistent representation space.

Analogy: If learning English→French and French→English translation, using the same dictionary for both directions ensures consistency.

[!definition] Embedding Layer

An embedding layer is a learned matrix $\mathbf{E} \in \mathbb{R}^{V \times d}$ where:

• $V$ = vocabulary size
• $d$ = embedding dimension
• Each row $\mathbf{E}[i]$ is the embedding vector for token $i$

Operation: For input token ID $x \in \{0, 1, \ldots, V-1\}$, the embedding is: $\text{Embed}(x) = \mathbf{E}[x] = \mathbf{e}_x \in \mathbb{R}^d$

This is not a matrix multiplication—it's an indexing operation (retrieval).

[!formula] Deriving the Embedding Operation

Goal: Convert token ID $x$ to vector $\mathbf{e}_x$.

Step 1: One-Hot Encoding (Conceptual)

Represent token $x$ as one-hot vector $\mathbf{o}_x \in \{0,1\}^V$ where: $\mathbf{o}_x[i] = \begin{cases} 1 & \text{if } i = x \\ 0 & \text{otherwise} \end{cases}$

Step 2: Matrix Multiplication Interpretation

If we compute $\mathbf{o}_x^T \mathbf{E}$: $\mathbf{o}_x^T \mathbf{E} = [0, \ldots, 0, 1, 0, \ldots, 0] \begin{bmatrix} — \mathbf{e}_0 — \\vdots \\ — \mathbf{e}_x — \\ \vdots \\ — \mathbf{e}_{V-1} — \end{bmatrix} = \mathbf{e}_x$

Why this step? This shows embedings are mathematically equivalent to matrix multiplication with one-hot vectors, but we never actually create the one-hot (memory inefficient).

Step 3: Efficient Implementation

In practice:

# Conceptual: e_x = one_hot(x) @ E  # O(V*d) - slow!
e_x = E[x]  # O(1) indexing - fast!

Why this matters: For $V=5000$, $d=512$, one-hot multiplication would waste $5000 \times 512 = 2.56\text{M}$ multiplications (most against zeros). Direct indexing is instant.

[!formula] Weight Tying: Full Derivation

The Standard Architecture (No Tying)

Forward pass:

1. Input embedding: $\mathbf{h}_0 = \mathbf{E}_{\text{in}}[x]$ where $\mathbf{E}_{\text{in}} \in \mathbb{R}^{V \times d}$
2. Transformer layers: $\mathbf{h}_L = \text{Transformer}(\mathbf{h}_0)$
3. Output projection: $\mathbf{z} = \mathbf{h}_L \mathbf{W}_{\text{out}} + \mathbf{b}$ where $\mathbf{W}_{\text{out}} \in \mathbb{R}^{d \times V}$
4. Softmax: $P(x_{t+1} | x_{\leq t}) = \text{softmax}(\mathbf{z})$

Total parameters: $V \cdot d$ (input) $+ V \cdot d$ (output) $= 2Vd$

Weight Tying Constraint

Set $\mathbf{W}_{\text{out}} = \mathbf{E}_{\text{in}}^T$ (transpose). Now: $\mathbf{z} = \mathbf{h}_L \mathbf{E}_{\text{in}}^T + \mathbf{b}$

Component-wise: The logit for token $j$ is: $z_j = \mathbf{h}_L \cdot \mathbf{e}_j + b_j$

Interpretation: Score for token $j$ = similarity between context representation $\mathbf{h}_L$ and token embedding $\mathbf{e}_j$.

Why this step? The dot product $\mathbf{h}_L \cdot \mathbf{e}_j$ measures how "aligned" the context is with token $j$'s meaning—high similarity → high probability.

Parameter reduction: Now only $Vd + V$ parameters (embedings + bias), saving $Vd - V \approx Vd$ parameters.

Gradient Flow with Tying

During backpropagation: $\frac{\partial \mathcal{L}}{\partial \mathbf{E}} = \frac{\partial \mathcal{L}}{\partial \mathbf{E}_{\text{in}}} + \left(\frac{\partial \mathcal{L}}{\partial \mathbf{W}_{\text{out}}}\right)^T$

Why this step? Since $\mathbf{E}$ is used in two places (input and output), gradients from both paths accumulate. This provides richer supervision—embedings learn from both "what tokens mean in context" and "what tokens are likely next."

[!example] Example 1: Small Vocabulary Embedding

Setup: $V=4$ tokens ["cat", "dog", "the", "ran"], $d=3$ dimensions.

Embedding matrix: $\mathbf{E} = \begin{bmatrix} 0.2 & -0.5 & 0.8 \\ 0.3 & -0.4 & 0.7 \\ -0.9 & 0.1 & 0.2 \\ 0.6 & 0.3 & -0.5 \end{bmatrix} \begin{matrix} \text{(cat, ID=0)} \\ \text{(dog, ID=1)} \\ \text{(the, ID=2)} \\ \text{(ran, ID=3)} \end{matrix}$

Query: Embed "dog" (ID=1).

Solution: $\text{Embed}(1) = \mathbf{E}[1] = [0.3, -0.4, 0.7]$

Why this step? We just index row 1. No computation needed.

[!example] Example 2: Output Projection with Weight Tying

Context: After transformer, we have $\mathbf{h}_L = [0.5, 0.2, -0.3]$.

Compute logits (using tied weights, $\mathbf{W}_{\text{out}} = \mathbf{E}^T$): $\mathbf{z} = \mathbf{h}_L \mathbf{E}^T = [0.5, 0.2, -0.3] \begin{bmatrix} 0.2 & 0.3 & -0.9 & 0.6 \\ -0.5 & -0.4 & 0.1 & 0.3 \\ 0.8 & 0.7 & 0.2 & -0.5 \end{bmatrix}$

Calculate each logit:

• $z_0 = 0.5(0.2) + 0.2(-0.5) + (-0.3)(0.8) = 0.1 - 0.1 - 0.24 = -0.24$
• $z_1 = 0.5(0.3) + 0.2(-0.4) + (-0.3)(0.7) = 0.15 - 0.08 - 0.21 = -0.14$
• $z_2 = 0.5(-0.9) + 0.2(0.1) + (-0.3)(0.2) = -0.45 + 0.02 - 0.06 = -0.49$
• $z_3 = 0.5(0.6) + 0.2(0.3) + (-0.3)(-0.5) = 0.3 + 0.06 + 0.15 = 0.51$

Probabilities (after softmax): $P(\text{next token}) = \text{softmax}([-0.24, -0.14, -0.49, 0.51]) \approx [0.19, 0.21, 0.15, 0.40]$

Interpretation: "ran" (ID=3) has highest probability because its embedding is most aligned with $\mathbf{h}_L$.

Why this step? Weight tying makes the model ask: "Which token embedding is closest to the context vector?" This is geometrically intuitive.

[!example] Example 3: Parameter Count Comparison

Scenario: GPT-2 small with $V=50257$, $d=768$.

Without weight tying:

• Input embedings: $50257 \times 768 = 38,597,376$
• Output projection: $768 \times 50257 = 38,597,376$
• Total: $77,194,752$ parameters

With weight tying:

• Shared embeddings: $50257 \times 768 = 38,597,376$
• Output bias: $50257$
• Total: $38,647,633$ parameters

Savings: $77.2M - 38.6M = 38.6M$ parameters (~50% reduction in embedding/projection).

Why this step? For large vocabularies, embedings dominate parameter count. Tying cuts this nearly in half without hurting (often improving!) performance.

[!mistake] Common Mistake 1: "Embedings are just random vectors"

The wrong intuition: Embeddings are arbitrary—any random initialization works.

Why it feels right: Initially, embedings are random. And even random embedings give some signal.

The fix: Embedings are learned representations. Through gradient descent:

• Similar tokens (cat/dog) get similar vectors
• Embedings capture semantic and syntactic relationships
• The model learns a geometry where "king - man + woman ≈ queen"

Evidence: Pre-trained embedings (Word2Vec, GloVe) transfer across tasks because they capture real linguistic structure.

[!mistake] Common Mistake 2: "Weight tying forces input=output exactly"

The wrong intuition: With weight tying, $\mathbf{E}_{\text{in}} = \mathbf{W}_{\text{out}}$, so input and output are identical.

Why it feels right: We use the same matrix, so they must be the same.

The fix: We tie $\mathbf{W}_{\text{out}} = \mathbf{E}_{\text{in}}^T$ (transpose!). The operations are:

• Input: $\mathbf{h}_0 = \mathbf{E}[x]$ (select row $x$)
• Output: $\mathbf{z} = \mathbf{h}_L \mathbf{E}^T$ (compute all dot products)

These are inverse operations geometrically:

• Input: "What vector represents this token?"
• Output: "Which token vectors are closest to this context?"

Why this works: The output layer is asking "which embedding is most similar to my hidden state?" If embedings are well-learned, this is exactly what we want.

[!mistake] Common Mistake 3: "Embeddings are per-position"

The wrong intuition: Token "cat" at position 5 gets a different embedding than "cat" at position 10.

Why it feels right: Transformers use positional encodings, so position matters.

The fix: The token embedding is position-independent: "cat" always gets $\mathbf{e}_{\text{cat}}$. Position is added separately: $\mathbf{h}_0 = \mathbf{E}[x] + \text{PositionalEncoding}(\text{pos})$

Why this step? Separating token and position information is cleaner—the embedding captures what the token is, positional encoding captures where it is.

[!recall]- Explain to a 12-Year-Old

Imagine you have a magic dictionary. Every word has a secret code (a list of numbers). When you look up "dog," you always get the same code, like [5, -2, 8].

Now, you're writing a story. You've written "The dog ran," and you want to predict the next word. Your brain (the transformer) thinks about the story and produces its own code, like [6, -1, 7].

To pick the next word, you compare your brain's code to every word's code in the dictionary. Whichever word's code is closest wins! Maybe "fast" has code [6, -1.5, 7.2]—super close—so you predict "fast."

Weight tying means using the same dictionary for looking up words at the start AND comparing at the end. Why? Because if "dog" and "pupy" have similar codes in the dictionary, you want your brain to treat them similarly when predicting too. It's like using one map for going to a place and coming back—more efficient and consistent!

[!mnemonic] Remember Weight Tying

"Shared Shoes": Input and output wear the same shoes (weight matrix), but one walks forward (embed tokens), the other walks backward (predict tokens). Same shoes, opposite directions. Saves money (parameters), still gets you there!

Key Formulas

Embedding Lookup

$\mathbf{e}_x = \mathbf{E}[x] \quad \text{where } \mathbf{E} \in \mathbb{R}^{V \times d}$

Output Logits (Weight Tying)

$z_j = \mathbf{h}_L \cdot \mathbf{e}_j + b_j = \sum_{k=1}^{d} h_{L,k} \cdot e_{j,k} + b_j$

Probability Distribution

$P(x_{t+1} = j \mid x_{\leq t}) = \frac{\exp(z_j)}{\sum_{i=1}^{V} \exp(z_i)}$

Connections

• 4.2.01-Tokenization-fundamentals — Tokens are converted to IDs, then embedded
• 4.2.04-Subword-tokenization-BPE-WordPiece — Vocabulary size $V$ determines embedding matrix size
• 4.3.02-Positional-encoding — Added to token embeddings to inject position info
• 4.4.03-Self-attention-mechanism — Operates on embedded token vectors
• 5.1.03-Cross-entropy-loss — Loss function that backprops through embeddings
• 6.2.04-Parameter-efficient-fine-tuning — Adapter methods often freeze embeddings

Practical Notes

When to use weight tying:

• ✅ Standard language models (GPT, BERT)
• ✅ Small to medium models where parameter count matters
• ✅ When vocabulary is large ($V > 10k$)

When NOT to use:

• ❌ Encoder-decoder models with different source/target vocabularies
• ❌ Models where input/output dimensions differ ($d_{\text{in}} \neq d_{\text{out}}$)
• ❌ Some multimodal models (vision+language) where output is not token prediction

Implementation tip: Most frameworks (PyTorch, TensorFlow) make weight tying trivial:

self.output_proj.weight = self.embedding.weight  # Share reference

#flashcards/ai-ml

Concept Map

Hinglish (regional understanding)

Hinglish (regional understanding)