Embedding layers and tied weights: Visual walkthrough

This page rebuilds the whole story of the parent topic as a sequence of pictures. We start with a token ID — literally just a number — and end with a probability over the next word, using the same table of vectors twice. Every symbol is drawn before it is used.

You should already have met token IDs from 4.2.01-Tokenization-fundamentals and subword pieces from 4.2.04-Subword-tokenization-BPE-WordPiece. Everything else, we build here.

Step 1 — What a token ID actually is

WHAT. We take the word "dog" and, via the tokenizer, get a plain integer $x$ .

WHY. A neural network cannot add or multiply the letters d-o-g. It needs numbers. But we cannot feed the raw ID $x = 1$ directly either — the network would think ID $2$ is "twice" ID $1$ , which is nonsense ("the" is not twice "dog"). So the ID is only a pointer, and Step 2 fixes what it points at.

PICTURE. Below, four words each get a row-number slot. The number is an address, nothing more.

Figure — Embedding layers and tied weights

Step 2 — The lookup table (the embedding matrix)

WHAT. We build a table $E$ with one row per token. Each row is a list of $d$ real numbers — a point in space.

WHY. We want IDs to carry meaning, and meaning that the model can learn and adjust. So instead of a fixed rule, we give each token its own tunable vector. The number $d$ (how long each row is) is how many independent "dials" describe a token.

PICTURE. The table is $V$ rows tall and $d$ columns wide. Retrieving token $x$ means: walk to row $x$ and read it off. No arithmetic — just reading.

Step 3 — Why retrieval equals a matrix multiply (and why we skip it)

WHAT. We show that "reading row $x$ " is secretly the same as multiplying a special one-hot vector by $E$ .

WHY. This matters for two reasons. (1) It lets us treat embeddings with the same calculus we use for every other layer, so gradients flow through them. (2) It reveals the trick: because the one-hot is all zeros except a single $1$ , the multiply is $99.99%$ wasted work — so in practice we just index.

First, define the one-hot vector — a switch board with exactly one switch flipped on.

Before multiplying, we need one small piece of linear-algebra notation: the transpose.

Now multiply it into the table:

PICTURE. The lit-up cell (amber) picks its row; all the grey rows are killed by their zeros.

Step 4 — The two jobs that bracket the whole model

WHAT. The embedding table is used at the start (turn ID into a vector) and there is a second matrix at the end (turn a vector into scores over all tokens). Let us name that second matrix.

WHY. A language model's job is: read some tokens, predict the next one. So it must (a) convert incoming tokens to vectors — the input side — and (b) convert its final internal vector back into a preference over every possible next token — the output side. These are opposite directions of travel.

Definition The pieces of the pipeline

$E_{in}$ — the input embedding table. This is exactly the same object as the $E$ from Steps 2–3; the subscript "in" is only a reminder that here it plays the input-side role. So $E_{in} \in R^{V \times d}$ and $E_{in} [x] = e_{x}$ .
$L$ — the number of transformer layers stacked in the model (a fixed whole number, e.g. $12$ for GPT-2 small). We count layers $1, 2, \dots, L$ ; the subscript on a hidden state says how many layers it has passed through.
$h_{0}$ — the input vector, $= E_{in} [x]$ , the embedded token before any layer (hence subscript $0$ ).
$h_{L}$ — the final hidden state, the vector after all $L$ transformer layers have mixed context (from 4.4.03-Self-attention-mechanism). It summarises "what should come next?".
$W_{out} \in R^{d \times V}$ — the output projection: a matrix that turns the $d$ -long $h_{L}$ into $V$ scores. In the untied model it is a separate second table.
$z$ — the vector of logits, one raw score per vocabulary token.

PICTURE. IDs enter left, a hidden vector emerges right, and $W_{out}$ fans it back out to $V$ scores. Notice we are storing two big $V \times d$ tables.

Step 5 — The tying idea: reuse the table, transposed

WHAT. We delete $W_{out}$ and instead reuse the input table $E_{in}$ (our original $E$ ), flipped on its side using the transpose from Step 3: $W_{out} = E_{in}^{T}$ .

WHY. Look at the two jobs again. Input asks "what vector means this token?" Output asks "which token does this vector mean?" Those are inverse questions about the same meaning-space. If a token's meaning-vector is good for reading it in, it should be good for scoring it out. Forcing one shared table makes the space consistent and, as a bonus, halves the parameters — and because a transpose stores no new numbers (Step 3), the output side now costs zero extra memory: it is the very same $V d$ numbers, merely read column-wise instead of row-wise.

Why does the transpose give scores? Read the result one entry at a time:

PICTURE. The context vector $h_{L}$ (amber) sits among the token vectors. Whichever token vector it aligns with best gets the biggest logit — here "ran".

Step 6 — Watch the numbers: logits then probabilities

WHAT. Use the toy table and one context vector to get real scores, then squash them into probabilities.

WHY. Seeing arithmetic land on the picture from Step 5 turns "dot product = alignment" from a slogan into something you can check by hand.

Toy table ( $V = 4, d = 3$ ), rows = cat, dog, the, ran: $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, h_{L} = [0.5, 0.2, - 0.3]$

Now turn scores into probabilities. We need a rule that (a) makes everything positive, (b) makes them sum to $1$ , (c) keeps the biggest score the most likely. That rule is the softmax (its output is what 5.1.03-Cross-entropy-loss then grades).

PICTURE. Bars for the four logits, and beside them the softmax probabilities: the ordering is preserved, the values are now a valid distribution.

Step 7 — Gradients: the shared table learns twice

WHAT. Because the one table $E = E_{in}$ appears at both ends, learning signals from both ends land on it and add up.

WHY. During training the model measures its error (cross-entropy) and pushes each parameter to reduce it — this push is the gradient (a derivative: "how much does the loss change if I nudge this number?"). With tying, $E$ gets nudged by two messages: "represent the input token better" and "score the correct next token higher". Two supervisions on one table means richer, more consistent learning.

PICTURE. Two arrows of "correction" flow back into a single table and merge.

Step 8 — The edge cases you must not trip on

The one-picture summary

This single diagram is the whole page: an ID becomes a row, layers stir in context, and the same table — transposed — turns the result back into scores; gradients from both ends feed the one table.

Recall Feynman retelling — say it in plain words

A tokenizer hands us a bare number, an address with no meaning. We keep a big notebook, one page per token, each page holding a short list of numbers we are allowed to tune. To "embed" a token we just flip to its page and copy the numbers — that copying is secretly a multiply by a switch-vector that's all zeros except one, but since almost everything is zero we skip the arithmetic and flip pages directly. Weird words we can't spell just get filed under one shared UNK page. The transformer stirs these pages together with their neighbours across all $L$ layers until one final vector says "here's the gist of the next word." To turn that gist into a preference over every possible word, we would normally keep a second big notebook. But the two jobs are mirror images — "what does this token mean?" versus "which token does this mean?" — so we reuse the first notebook, tipped on its side (that tipping is the transpose, and it stores no new paper), comparing the gist against every page with a dot product: the page that points the same way wins. Softmax makes those comparisons into honest probabilities. And because the one notebook is used at both ends, whenever the model is wrong, corrections arrive from both directions and stack up on the same pages — so it learns faster and stays consistent, while storing only half as many numbers. That single reused notebook — the embedding table $E$ , read forwards to embed and transposed to score — is the entire idea of weight tying.

Recall

What does a token ID carry on its own? ::: Nothing but an address (a row number); meaning lives in the embedding vector it points to. Why is embedding lookup not implemented as a matrix multiply? ::: The one-hot multiply is almost all zeros — $O (V d)$ wasted work — so we index the row directly in $O (1)$ . What does the transpose $(\cdot)^{T}$ do, and why is it "free"? ::: It swaps rows and columns (tips the array on its side) without changing or storing any new numbers — so reusing $E^{T}$ costs no extra memory. What exactly is tied in weight tying? ::: The output projection is set to the transpose of the input embedding: $W_{out} = E_{in}^{T}$ . What does a single logit $z_{j}$ measure? ::: The dot product (alignment) of the context vector $h_{L}$ with token $j$ 's embedding $e_{j}$ , plus a bias. What is $L$ in $h_{L}$ ? ::: The number of transformer layers; the subscript counts how many layers the hidden state has passed through ( $0$ = none, $L$ = all). Why do tied embeddings learn from two sources? ::: The shared table appears at input and output, so gradients from both paths add onto $E$ . How are special / unknown tokens handled? ::: PAD/UNK/BOS/EOS each get their own learned row in the same table; out-of-vocabulary words map to the UNK ID. Where does position information enter, if not the token embedding? ::: It is added separately as a positional encoding: $h_{0} = E [x] + PosEnc(pos)$ .