The Output Side: Unembedding and Logits
Lesson 4 of 4 in Embeddings: Tokens Become Vectors.
After the final block, each position holds one last Hidden state — a vector of d_model (model dimension) numbers. To predict the next token, the model has to convert that vector into an opinion about every entry in the Vocabulary. That’s the job of the Unembedding — also called the output projection or the LM head.
Mechanically it’s one matrix multiply: the final vector times a d_model × |V| matrix, producing one raw score per vocabulary token. Those scores are the Logits. A Softmax then squashes them into probabilities that sum to 1, and the Sampling procedure picks the actual next token. Notice the symmetry: the embedding took you from token space into vector space; the unembedding takes you back out.
A vertical stack read bottom to top: token ids from the tokenizer; the embedding matrix lookup producing one vector of d_model numbers per token; N transformer blocks that reshape each vector in context; the final hidden state; the unembedding projection producing one score per vocabulary entry; and logits passed through softmax into next-token probabilities.
Look at the two bookend matrices side by side. The embedding matrix is |V| × d_model; the unembedding is d_model × |V| — the same shape, transposed. That invites an elegant trick called weight tying: reuse the input embedding matrix (transposed) as the output projection instead of learning a second one. The original transformer did exactly this (Vaswani et al., 2017, following Press & Wolf, 2016), and the intuition is satisfying — one shared geometry serves for both reading tokens in and writing them out, and the model saves an entire matrix’s worth of parameters.
Whether to tie is a per-architecture design choice, and model families genuinely differ — untied matrices give the output side freedom to specialize, at the cost of more parameters. What generalizes is the shape of the trade, not any one family’s answer.
One practical note: when a model is generating, only the final position’s logits matter — that’s where the next token comes from. What the sampler does with those logits (temperature, top-p and friends) is the Inference & Serving domain’s territory.
From vector to distribution
Let h be the final hidden state at the last position (a row vector of width d_model) and W_U the unembedding matrix of shape d_model × |V|. The logits are z = h · W_U (some architectures add a bias term). Then softmax turns scores into probabilities: p_i = exp(z_i) / Σ_j exp(z_j). Every token gets a nonzero probability — softmax never assigns exactly zero — which is one reason models can always surprise you.
Weight tying sets W_U = W_Eᵀ. Under tying, the logit for token i is just the dot product h · e_i between the hidden state and that token’s embedding row — the model literally asks “which token’s vector does my current state point toward?” It also saves |V| × d_model parameters, a meaningful fraction in smaller models (Press & Wolf, 2016, arXiv:1608.05859).
Interactive checkpoint quiz (2 questions) — open this page in a browser to take it.