Worked examples — Embedding layers and tied weights

This is the do-it-by-hand companion to the parent note. Every symbol here is used exactly the way the parent defined it: $\mathbf{E}$ is the embedding matrix (one row per token), $V$ is how many tokens exist, $d$ is how long each token's vector is, and ${\mathbf{h}}_L$ is the "context vector" the transformer produces just before predicting the next token. If any of those feel shaky, re-read the parent first — here we only compute.

The scenario matrix

Every problem this topic can throw at you falls into one of these cells. The examples below are each tagged with the cell they cover, and together they hit all of them.

[!definition] The one word we must earn: dot product

Before any output-side example, we need the dot product. Take two lists of numbers of the same length, say ${\mathbf{h}}_L=[0.5,0.2,-0.3]$ and ${\mathbf{e}}_j=[0.6,0.3,-0.5]$.

The dot product ${\mathbf{h}}_L\cdot {\mathbf{e}}_j$ means: multiply matching slots, then add: ${\mathbf{h}}_L\cdot {\mathbf{e}}_j=(0.5)(0.6)+(0.2)(0.3)+(-0.3)(-0.5).$

That is the whole engine of a tied output layer: logit for token $j$ = dot product of the context with token $j$'s drawer vector.

[!example] Example 1 — Cell A: plain lookup

Statement. With $V=4$ tokens ["cat","dog","the","ran"], $d=3$, and $\mathbf{E}=\left[\begin{array}{ccc}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{array}\right],$ retrieve the embedding of "the".

Forecast: guess — do we compute anything, or just read a row?

1. Find the ID. "the" is the 3rd word, and IDs start at $0$, so ID $=2$. Why this step? The parent stressed IDs are zero-indexed — off-by-one here silently returns the wrong word.
2. Open drawer 2. Take row index $2$: $\mathbf{E}[2]=[-0.9,0.1,0.2]$. Why this step? Embedding is retrieval, not multiplication — there is nothing to calculate.

Answer: $\text{Embed}(2)=[-0.9,0.1,0.2]$.

Verify: The result must be one of the four rows exactly (no arithmetic can invent new numbers). Row 2 matches. ✓

[!example] Example 2 — Cell B: lookup equals one-hot × matrix

Statement. Show that retrieving "dog" (ID $=1$) equals ${\mathbf{o}}_1^{\mathsf{T}}\mathbf{E}$, where ${\mathbf{o}}_1=[0,1,0,0]$.

Forecast: guess — will the one-hot pick out exactly row 1, or blend rows?

1. Write the multiply slot by slot. Column $k$ of the answer is ${\sum }_i{\mathbf{o}}_1[i]\mathbf{E}[i,k]$. Why this step? This is what "row-vector times matrix" means — sum down each column, weighted by the one-hot.
2. Kill every zero. ${\mathbf{o}}_1$ is $0$ everywhere except slot $1$, so every term dies except $\mathbf{E}[1,k]$. Why this step? This is why we never build the one-hot in practice — we'd multiply thousands of numbers by zero. Indexing skips straight to the survivor.
3. Read off the survivor. Answer $=[0.3,-0.4,0.7]$.

Verify: Equals $\mathbf{E}[1]$ from direct lookup — identical to Example 1's method on row 1. ✓

[!example] Example 3 — Cell C: tied output logits

Statement. Context vector ${\mathbf{h}}_L=[0.5,0.2,-0.3]$, tied weights ${\mathbf{W}}_{\text{out}}={\mathbf{E}}^{\mathsf{T}}$, bias $=0$. Find all four logits and the softmax.

Forecast: which token wins — cat, dog, the, or ran?

1. One dot product per token. Logit $z_j={\mathbf{h}}_L\cdot {\mathbf{e}}_j$. Why this step? Tying turns "predict next token" into "score each drawer by alignment with context" — Section: dot product.
2. Compute:
   • $z_0=(0.5)(0.2)+(0.2)(-0.5)+(-0.3)(0.8)=-0.24$
   • $z_1=(0.5)(0.3)+(0.2)(-0.4)+(-0.3)(0.7)=-0.14$
   • $z_2=(0.5)(-0.9)+(0.2)(0.1)+(-0.3)(0.2)=-0.49$
   • $z_3=(0.5)(0.6)+(0.2)(0.3)+(-0.3)(-0.5)=0.51$
3. Softmax (exponentiate, then normalise so they sum to 1): $\approx [0.19,0.21,0.15,0.40]$. Why this step? Softmax turns "alignment scores" into probabilities that add to 1 — see 5.1.03-Cross-entropy-loss for how these are then scored.

Answer: ran (ID 3) wins at $\approx 0.40$.

Verify: Probabilities sum to $1$; the biggest logit ($0.51$) maps to the biggest probability. ✓

[!example] Example 4 — Cell D: sign behaviour of alignment

Statement. Same $\mathbf{E}$. Compare two contexts against the (${\mathbf{e}}_2=[-0.9,0.1,0.2]$): ${\mathbf{h}}_A=[-1,0,0]$ and ${\mathbf{h}}_B=[+1,0,0]$.

Forecast: guess the sign of each score before computing.

1. Score A. ${\mathbf{h}}_A\cdot {\mathbf{e}}_2=(-1)(-0.9)=+0.9$. Why this step? ${\mathbf{h}}_A$ points against the first axis, and so does ${\mathbf{e}}_2$ (its first slot is $-0.9$). Two negatives → positive → aligned.
2. Score B. ${\mathbf{h}}_B\cdot {\mathbf{e}}_2=(+1)(-0.9)=-0.9$. Why this step? ${\mathbf{h}}_B$ points with the axis, opposite to ${\mathbf{e}}_2$ → anti-aligned → negative logit → low probability.

Answer: $+0.9$ vs $-0.9$ — flipping the context sign flips who the model wants.

Verify: The two scores are exact negatives, as expected when only the sign of $\mathbf{h}$ flips. ✓ See the two-arrow figure again for the geometry.

[!example] Example 5 — Cell E: degenerate inputs

Statement. (a) What are all logits when ${\mathbf{h}}_L=[0,0,0]$? (b) If cat and dog had identical rows, what can the model do about their logits?

Forecast: guess whether a zero context prefers any token.

1. (a) Zero context. Every $z_j=0\cdot {\mathbf{e}}_j=0$. Softmax of all-equal $=[0.25,0.25,0.25,0.25]$. Why this step? A dot product with the zero vector is always $0$ — no drawer aligns with "no direction," so the model is maximally unsure. This is the degenerate case.
2. (b) Identical rows. If ${\mathbf{e}}_{\text{cat}}={\mathbf{e}}_{\text{dog}}$, then for any context $z_{\text{cat}}=z_{\text{dog}}$ always. Why this step? Tied weights use the same drawer for scoring, so duplicate rows are permanently indistinguishable outputs — a warning that embeddings must actually differ (Mistake 1 in the parent).

Answer: (a) uniform $0.25$ each; (b) cat and dog get equal probability forever.

Verify: Uniform probabilities sum to $1$ ($4\times 0.25=1$). ✓ Equal logits ⇒ equal softmax entries by definition. ✓

[!example] Example 6 — Cell F: limiting / saturating behaviour

Statement. Keep $z_1=z_2=z_3=0$ but let $z_0=t$ grow without bound. What does softmax do as $t\to \infty$? As $t\to -\infty$?

Forecast: does the winner's probability approach $1$ or some smaller ceiling?

1. Write softmax. $P_0=\frac{e^t}{e^t+3}$ (three others each contribute $e^0=1$). Why this step? Only relative logit gaps matter; the three zeros contribute $3$ total.
2. Push $t\to \infty$. $e^t$ dwarfs $3$, so $P_0\to 1$, others $\to 0$. This is saturation. Why this step? Once one logit runs away, softmax becomes nearly one-hot — the model is essentially certain.
3. Push $t\to -\infty$. $e^t\to 0$, so $P_0\to 0$ and the other three share it: each $\to 1/3$. Why this step? A token with a hugely negative logit is effectively removed from the race.

Answer: $P_0:0\to \frac{1}{4}\to 1$ as $t:-\infty \to 0\to +\infty$.

Verify: At $t=0$, $P_0=\frac{1}{1+3}=0.25$ (matches Example 5a). ✓ Limits computed symbolically. ✓

[!example] Example 7 — Cell G: real-world memory budget

Statement. A model has $V=50257$, $d=768$ (GPT-2 small). How many parameters do the input+output embeddings use without tying, and how many with tying (plus a length-$V$ output bias)? What fraction is saved?

Forecast: guess — roughly half, or more?

1. Untied. Input $=Vd$, output $=dV$, total $=2Vd=2(50257)(768)=77,194,752$. Why this step? Two separate matrices of the same shape.
2. Tied. Share one $Vd$ matrix, add a bias of length $V$: $Vd+V=38,597,376+50,257=38,647,633$. Why this step? One matrix reused (transposed) for both directions — the parent's core claim.
3. Saved fraction. $\frac{77,194,752-38,647,633}{77,194,752}\approx 0.4993$.

Answer: $\approx 49.9\%$ of the embedding/projection block vanishes. Relevant to 6.2.04-Parameter-efficient-fine-tuning.

Verify: Numbers computed exactly in VERIFY; saving $\approx$ half, minus the tiny bias term. ✓

[!example] Example 8 — Cell H: exam twist (two gradient paths)

Statement. Suppose the loss gradient landing on $\mathbf{E}$ as an input table is $g_{\text{in}}=[0.10,-0.20,0.30]$ for dog's row, and the gradient landing on the same row through the output projection (already transposed back) is $g_{\text{out}}=[0.05,0.05,-0.10]$. With weight tying, what single update does row dog receive?

Forecast: does one path override the other, or do they combine?

1. Recall the tying rule. $\frac{\partial \mathcal{L}}{\partial \mathbf{E}}=\frac{\partial \mathcal{L}}{\partial {\mathbf{E}}_{\text{in}}}+(\frac{\partial \mathcal{L}}{\partial {\mathbf{W}}_{\text{out}}}{)}^{\mathsf{T}}$. Why this step? One matrix used in two places ⇒ gradients add (chain rule over a shared variable) — see 4.4.03-Self-attention-mechanism for where ${\mathbf{h}}_L$ came from.
2. Add slot by slot. $[0.10+0.05,-0.20+0.05,0.30-0.10]=[0.15,-0.15,0.20]$. Why this step? This "richer supervision" is exactly why tied models often generalise better, not just smaller.

Answer: row dog gets combined gradient $[0.15,-0.15,0.20]$.

Verify: Element-wise sum; sanity: middle slot cancels partly ($-0.20+0.05=-0.15$). ✓

[!recall]- Self-check