Embedding layers and tied weights: Exercises — AI-ML deep dive

Definition Names we use on this page

$E$ = the embedding matrix below, shape $V \times d$ . Its row $i$ , written $e_{i}$ or $E [i]$ , is the vector for token $i$ .
$E_{in}$ = the input embedding matrix — the copy of $E$ used to look up a token when it enters the model.
$W_{out}$ = the output projection matrix, shape $d \times V$ — the matrix used to turn a context vector into one score per token when the model predicts. Weight tying means $W_{out} = E_{in}^{T}$ , i.e. both are built from the one matrix $E$ .
$o_{x}$ = the one-hot vector for token $x$ : a length- $V$ list that is $0$ everywhere except a single $1$ in slot $x$ (e.g. for $V = 4$ , token 2 gives $o_{2} = [0, 0, 1, 0]$ ). It is the "pick token $x$ " selector.
$h_{L}$ = the context vector, length $d$ , produced after the transformer layers.
$L$ = the loss: a single number measuring how wrong a prediction is (smaller = better). We use cross-entropy loss $L = - ln P_{correct}$ .

We reuse the tiny world from the parent note so nothing new is assumed:

$E = 0.2 0.3 - 0.9 0.6 - 0.5 - 0.4 0.1 0.3 0.8 0.7 0.2 - 0.5 cat, ID = 0 dog, ID = 1 the, ID = 2 ran, ID = 3$

Here $V = 4$ rows (one per token) and $d = 3$ columns (one per feature). Everything below is built only from indexing a row and dot products (multiply matching entries, add them up) — the two operations the parent note earned.

Figure 1 below shows $E$ as a coloured lookup table: the green box marks row 2 ("the"), the row we retrieve in Problem 1.2, so you can see that "embedding = reading one row".

Figure — Embedding layers and tied weights

L1 — Recognition

Problem 1.1

The matrix $E$ above has shape $V \times d$ . State $V$ , state $d$ , and say in one sentence what a single row of $E$ physically represents.

Recall Solution

$V = 4$ (the number of tokens in the vocabulary — count the rows). $d = 3$ (the embedding dimension — count the columns). A single row $E [i]$ is the learned vector that token $i$ is always mapped to; e.g. row 0 is the meaning-vector for "cat", and the green-boxed row in Figure 1 is row 2, "the".

Problem 1.2

Without any multiplication, write down $Embed (2)$ , the embedding of the token "the".

Recall Solution

Embedding is retrieval, not multiplication — we just read row 2 (the green-boxed row in Figure 1): $Embed (2) = E [2] = [- 0.9, 0.1, 0.2] .$ No arithmetic happens. That is the whole point of a lookup table.

Problem 1.3

True or false: with weight tying we set $W_{out} = E_{in}$ (the same matrix, unchanged). Here $E_{in}$ and $W_{out}$ are the input and output matrices defined in the "Names we use" box above. If false, correct it.

Recall Solution

False. We set $W_{out} = E_{in}^{T}$ — the transpose. The input matrix $E_{in}$ is $V \times d$ (rows = tokens); to multiply a $d$ -length context vector and get $V$ logits we need a $d \times V$ matrix, which is exactly $E_{in}^{T}$ .

L2 — Application

Problem 2.1

Given context vector $h_{L} = [1, 0, 0]$ , compute the tied logit for token "cat" (ID 0), i.e. $z_{0} = h_{L} \cdot e_{0}$ (ignore bias).

Recall Solution

Dot product = multiply matching entries, sum: $z_{0} = (1) (0.2) + (0) (- 0.5) + (0) (0.8) = 0.2.$ Because $h_{L}$ points purely along the first axis, the logit just reads off the first coordinate of "cat"'s embedding.

Problem 2.2

With $h_{L} = [0.5, 0.2, - 0.3]$ (same as parent Example 2), compute all four logits $z_{0}, z_{1}, z_{2}, z_{3}$ under weight tying (no bias). You should reproduce the parent's numbers.

Recall Solution

Each $z_{j} = h_{L} \cdot e_{j}$ (dot with row $j$ ):

$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$

Logits $= [- 0.24, - 0.14, - 0.49, 0.51]$ . "ran" scores highest because its row is most aligned with $h_{L}$ .

Problem 2.3

Turn those logits into probabilities with softmax: $P_{j} = \frac{e ^{z_{j}}}{\sum _{k} e ^{z_{k}}}$ . Report each to 2 decimals.

Recall Solution

Softmax turns raw scores into a probability distribution (positive, sums to 1) by exponentiating then normalising — the exponential is chosen because it is always positive and monotonic, so a bigger logit always means a bigger probability. $e^{- 0.24} = 0.787, e^{- 0.14} = 0.869, e^{- 0.49} = 0.613, e^{0.51} = 1.665$ . Sum $= 3.934$ . $P \approx [0.20, 0.22, 0.16, 0.42] .$ "ran" wins with $\approx 0.42$ . (Small rounding differs slightly from the parent's $0.40$ ; the exact softmax values were confirmed by the symbolic check described in the "How to use this page" box.)

L3 — Analysis

Problem 3.1

The tied logit is $z_{j} = h_{L} \cdot e_{j}$ . Explain geometrically why a token whose embedding points in the same direction as $h_{L}$ gets a high logit, using the identity $a \cdot b = ∥ a ∥ ∥ b ∥ cos θ$ .

Recall Solution

The dot product equals (length of $h_{L}$ ) $\times$ (length of $e_{j}$ ) $\times cos θ$ , where $θ$ is the angle between the two arrows. Holding the lengths fixed, $cos θ$ is largest ( $= 1$ ) when $θ = 0$ — arrows point the same way — and smallest ( $= - 1$ ) when they point opposite ( $θ = 18 0^{\circ}$ ). So alignment = high score. The output layer literally asks: "which token embedding points most like my context vector?" Figure 2 below draws the context arrow (gray) and each token arrow in colour so you can see which one it points most like.

Problem 3.2

Take $h_{L} = [0.3, - 0.4, 0.7]$ — this is exactly the embedding of "dog" (row 1). Compute all logits. Which token wins, and what does that reveal about tied weights?

Recall Solution

$z_{0} = 0.3 (0.2) + (- 0.4) (- 0.5) + 0.7 (0.8) = 0.06 + 0.20 + 0.56 = 0.82$
$z_{1} = 0.3 (0.3) + (- 0.4) (- 0.4) + 0.7 (0.7) = 0.09 + 0.16 + 0.49 = 0.74$
$z_{2} = 0.3 (- 0.9) + (- 0.4) (0.1) + 0.7 (0.2) = - 0.27 - 0.04 + 0.14 = - 0.17$
$z_{3} = 0.3 (0.6) + (- 0.4) (0.3) + 0.7 (- 0.5) = 0.18 - 0.12 - 0.35 = - 0.29$

Winner: "cat" (ID 0) with $z_{0} = 0.82$ , narrowly beating "dog" itself ( $0.74$ ). The self-score $z_{1} = ∥ e_{1} ∥^{2} = 0.74$ is just the squared length of "dog". "cat" wins here because its row is nearly parallel to "dog"'s and slightly longer, so alignment × length pushes it above. This shows tied weights make the geometry of embeddings directly control next-token scores — similar tokens (cat/dog rows are close) score similarly.

Problem 3.3

Recall from the "Names we use" box that $L$ is the loss (how wrong the prediction is). With tying, the gradient into $E$ is a sum: $\frac{\partial L}{\partial E} = \frac{\partial L}{\partial E _{in}} + (\frac{\partial L}{\partial W _{out}})^{T}$ . In plain words, why does a single row of $E$ receive two contributions per step?

Recall Solution

The same matrix $E$ is used twice in the forward pass: once at the input (row selected when the token appears) and once at the output (every row dotted with $h_{L}$ to make logits). By the chain rule, a parameter used in two places gets one gradient term per use, and they add. So an embedding learns simultaneously "what this token means when I read it" and "when I should predict it" — richer supervision from one weight.

L4 — Synthesis

Problem 4.1

GPT-2 small has $V = 50257$ , $d = 768$ . Compute (a) parameters in the untied input + output layers, and (b) parameters in the tied version (shared matrix + one output bias of length $V$ ). (c) State the saving.

Recall Solution

One matrix of shape $V \times d$ has $V d$ entries. (a) Untied: input $V d = 50257 \times 768 = 38, 597, 376$ , output another $38, 597, 376$ . Total $= 77, 194, 752$ . (b) Tied: one shared $V d = 38, 597, 376$ plus bias $V = 50257$ . Total $= 38, 647, 633$ . (c) Saving $= 77, 194, 752 - 38, 647, 633 = 38, 547, 119$ parameters (~50% of the embedding/projection block). This is why tying matters most for large vocabularies — see 4.2.04-Subword-tokenization-BPE-WordPiece for where $V$ comes from.

Problem 4.2

A model uses token embedding plus positional encoding: $h_{0} = E [x] + Pos (pos)$ . Using our tiny $E$ , suppose $Pos (0) = [0.1, 0.0, - 0.1]$ . Compute $h_{0}$ for "cat" at position 0. Does the token part depend on position?

Recall Solution

$h_{0} = E [0] + Pos (0) = [0.2, - 0.5, 0.8] + [0.1, 0.0, - 0.1] = [0.3, - 0.5, 0.7] .$ The token embedding $E [0]$ is position-independent — "cat" always contributes $[0.2, - 0.5, 0.8]$ . Only the added 4.3.02-Positional-encoding term changes with position. Separation of concerns: what the token is vs. where it sits.

Problem 4.3

Reuse Problem 2.2's logits $[- 0.24, - 0.14, - 0.49, 0.51]$ . The true next token is "ran" (ID 3). Compute the cross-entropy loss $L = - ln P_{3}$ (with $L$ the loss from the "Names we use" box) using the softmax probabilities.

Recall Solution

From Problem 2.3, $P_{3} \approx 0.4232$ (exact softmax). Cross-entropy for the correct class is just the negative log of its probability: $L = - ln (0.4232) \approx 0.860.$ Lower is better; a perfect prediction ( $P_{3} = 1$ ) would give $L = 0$ . This links tied embeddings straight to 5.1.03-Cross-entropy-loss.

L5 — Mastery

Problem 5.1

Prove that the untied model with matrices $E_{in}$ and $W_{out}$ can represent every function the tied model can, but not vice-versa — and state the one-line reason tying is still preferred in practice.

Recall Solution

Untied ⊇ tied: the tied model is the special case $W_{out} = E_{in}^{T}$ . Any tied model is obtained from the untied model by choosing those specific weights, so the untied family contains every tied model — it is strictly more expressive (it also allows $W_{out} \neq = E_{in}^{T}$ ). Why tie anyway: more expressive ≠ better. Tying (i) roughly halves parameters, cutting memory and overfitting, and (ii) gives each embedding two gradient signals (Problem 3.3), which empirically improves generalisation. Capacity you don't need is capacity that overfits. (Relevant when adapting models cheaply — see 6.2.04-Parameter-efficient-fine-tuning.)

Problem 5.2

Build a $2 \times 2$ toy embedding matrix $E_{2} = [a c b d]$ and the one-hot vector $o_{1} = [0, 1]$ (the selector for token 1, as defined in the "Names we use" box). Show by explicit multiplication that $o_{1}^{T} E_{2}$ equals row 1 of $E_{2}$ , confirming "indexing = one-hot matmul". Then verify numerically with $a = 0.7, b = - 0.2, c = 0.4, d = 0.9$ .

Recall Solution

Multiply the row vector $o_{1}^{T} = [0, 1]$ into $E_{2}$ , entry by entry down each column: $o_{1}^{T} E_{2} = [0, 1] [a c b d] = [0 \cdot a + 1 \cdot c, 0 \cdot b + 1 \cdot d] = [c, d] = E_{2} [1] .$ The $0$ kills row 0 entirely; the single $1$ copies row 1 straight through. Numeric check with $a = 0.7, b = - 0.2, c = 0.4, d = 0.9$ : $o_{1}^{T} E_{2} = [0 (0.7) + 1 (0.4), 0 (- 0.2) + 1 (0.9)] = [0.4, 0.9],$ which is exactly row 1 $= [c, d] = [0.4, 0.9]$ . This is why we never build the one-hot in code — it computes a row we could just read.

Problem 5.3

Under tying, the self-logit of a token equals the squared length of its embedding: $z_{j j} = e_{j} \cdot e_{j} = ∥ e_{j} ∥^{2}$ when $h_{L} = e_{j}$ . Verify this for "ran" (row 3) and explain what it means for a model asked to "predict the token it just saw".

Recall Solution

$e_{3} = [0.6, 0.3, - 0.5]$ , so $∥ e_{3} ∥^{2} = 0. 6^{2} + 0. 3^{2} + (- 0.5)^{2} = 0.36 + 0.09 + 0.25 = 0.70.$ Set $h_{L} = e_{3}$ : the logit for "ran" is $0.70$ , a token's alignment with itself ( $θ = 0$ , $cos θ = 1$ ), the maximum possible cosine. So a context that "looks like" a token strongly prefers predicting that token — the geometric backbone of tied language models, and exactly the mechanism 4.4.03-Self-attention-mechanism exploits when copying context.

Recall Explain it back (self-test)

Embedding = retrieval ::: pick row $x$ of $E$ ; no multiply needed. Tied output logit for token $j$ ::: $z_{j} = h_{L} \cdot e_{j}$ — alignment of context with token $j$ 's embedding. Why transpose, not equal ::: input is $V \times d$ (select a row); output needs $d \times V$ to turn a length- $d$ context into $V$ logits, which is $E^{T}$ . Two gradients per row ::: $E$ is used at input and output, so chain-rule terms add. Tied param count ::: $V d + V$ (shared matrix + output bias).