Unembedding

Summary: The final step of a transformer that takes the refined last-token vector out of the residual stream and projects it into a score (logits) for every possible next token, ready for softmax to produce a probability distribution.

The operation

After all attention and MLP blocks have run, the last vector in the sequence, ${\vec{E}}_n^{(L)}$, is multiplied by an unembedding matrix $W_U$:

$\vec{\ell }=W_U{\vec{E}}_n^{(L)}$

$W_U$ has one row per vocabulary token and one column per embedding dimension (unlike $W_E$, which has one column per token — it’s the transposed shape):

$W_U\in {\mathbb{R}}^{|V|\times d_{\text{embed}}}$

For GPT-3: $50,257\text{ rows }\times 12,288\text{ columns }\approx 617\text{M}$ parameters, symmetric in size to $W_E$.

Each entry of the output is the dot product of one row of $W_U$ with ${\vec{E}}_n^{(L)}$ — i.e. how aligned the final vector is with each token’s “unembedding direction”. These raw alignment scores are the logits, which then get passed through softmax to produce the next-token probability distribution.

Image of unembedding matrix $W_U$ and softmax operation in GPT-3

$W_U\in {\mathbb{R}}^{|V|\times d_{\text{embed}}}$

• maps a single residual stream vector ${\vec{E}}_n^{(L)}\in {\mathbb{R}}^{12,288}$ (i.e. of length $d_{\text{embed}}$)
• to a logits vector $\vec{\ell }\in {\mathbb{R}}^{50,257}$ (i.e. one score per vocabulary token)
• via a $(50,257\times 12,288)$ matrix multiply.

💡 Why only the last vector?

The transformer computes a full vector for every position in the context, but at inference time we only care about the prediction for position $n$ (the next token after the input). So we only unembed position $n$.

During training things are different: every position simultaneously predicts its next token. So training runs the unembedding step at every position in the sequence, producing $n$ loss signals per forward pass rather than one. This is what makes training efficient — a length-$n$ sequence yields $n$ training examples for the cost of one forward pass through the stack.

At inference, the other positions’ output vectors are computed but discarded. They did their job earlier, when they contributed context to the positions after them through attention.

Tied vs untied weights

In some implementations, $W_U$ is tied to $W_E^{\top }$ — they share parameters. This:

• Halves the parameter count spent on embeddings/unembeddings.
• Encodes the prior that “the direction for embedding a token” and “the direction for predicting it” should be the same.

Other implementations (GPT-3 included, per 3b1b’s tally) keep them untied — two independent matrices of the same shape, ~1.2B total parameters between them.

Intuition via dot products

Recall from above:

Each entry of the output is the dot product of one row of $W_U$ with ${\vec{E}}_n^{(L)}$ — i.e. how aligned the final vector is with each token’s “unembedding direction”.

Because unembedding is literally “dot product against each row”, you can read $W_U$ as a bank of learned “if you see this direction in the residual stream, it’s likely to be followed by this token” detectors. The final-layer vector is a weighted cocktail of meanings; the row for “Snape” gives a high logit when that cocktail points toward “villainous Hogwarts potions master who is the hero’s least favourite professor”.

This also means any concept the model wants to make predictable must be a direction the unembedding can read — a nontrivial constraint on how the model uses its embedding space.

See also

• word-embedding — the inverse operation at the start of the stack
• logits — the raw output of unembedding
• softmax — applied to the logits to produce a distribution
• transformer-architecture — where unembedding sits

Sources

• src-3b1b-llms-ch2-transformers