Every transformer spends most of its depth moving vectors around a d_model-dimensional space — but at the very end it has to answer a concrete question: which of the tens of thousands of tokens comes next? The output projection (also called the unembedding, the LM head, or just the final linear layer) is the bridge from that continuous hidden state to a score for every token in the vocabulary. It is one matrix multiply, yet it is often the single largest weight tensor in a small model, it sets the shape of the loss during training, and it is where softmax, temperature, and sampling all attach. This article walks the whole final stage from first principles: the final normalization, the projection and its shapes, what a logit really is (and is not), the softmax that turns logits into probabilities, weight tying with the input embedding, the parameter and compute cost that scales as vocab × d_model, a fully worked numeric example, and a pointer to temperature and sampling downstream.
Where the head sits in the forward pass
By the time a token’s representation reaches the output projection, it has passed through the embedding layer and every transformer block — attention and feed-forward, residual-added and normalized, L times. The result is a sequence of hidden states h ∈ ℝ^(N × d), where N is the sequence length and d is d_model. Each row h_i is a d-dimensional vector that, in a causal language model, has absorbed information from every token up to and including position i.
The head’s job is narrow and mechanical: take each of those vectors and turn it into a vector of vocabulary-sized scores. Nothing about the head ‘understands’ language — all the representation learning already happened in the blocks below. The head is a readout: a single linear map that projects from the model’s internal geometry onto the axis-aligned coordinate system where each dimension is a token. Everything interesting about how the model decides what to say next is compressed into this one projection and the softmax that follows it, which is exactly why it repays a careful look.
The final normalization before the head
In the pre-norm architectures used by essentially all modern LLMs (GPT-2 onward, LLaMA, Phi, Mistral), normalization is applied at the input of each sub-layer rather than after the residual add. A consequence is that the output of the last transformer block is not normalized — it is the raw running sum of the residual stream. Feeding that directly into the projection would be unstable, so pre-norm models insert one final LayerNorm (or RMSNorm in LLaMA-style models) after the last block and before the output projection.
Concretely the computation is h_final = Norm(h_L), then logits = h_final · W_out. This final norm matters more than it looks: it re-centers and re-scales the residual stream so the projection sees inputs in a well-behaved range, which keeps the logits — and therefore the softmax and the gradients — numerically sane. Forgetting it is a classic re-implementation bug: the model still runs, but the logits drift to large magnitudes, the softmax saturates, and training either diverges or crawls. If you are counting parameters, the final norm adds only d (RMSNorm) or 2d (LayerNorm) — negligible next to the projection itself.
The output projection: shapes and the matmul
The projection is a single linear layer with no bias in most implementations. Fixing a convention and holding it for the whole article: let W_out ∈ ℝ^(d × V), where V is the vocabulary size. Then
h_final : [N, d] (final hidden states, post-norm)
W_out : [d, V] (the unembedding / LM head)
logits = h_final · W_out
logits : [N, V]Each output row logits_i has one entry per vocabulary token. Read column-wise, the projection is a stack of V weight vectors, one per token, each living in the same d-dimensional space as the hidden state. The score for token t is just the dot product logits[i, t] = h_final_i · W_out[:, t] — a measure of how aligned the current hidden state is with token t’s learned direction. That is the whole mechanism: the head asks, for every token in the vocabulary, “how much does the model’s current summary point in your direction?” The bigger the alignment, the higher the score. Frameworks that use nn.Linear store the transpose — a [V, d] weight read row-wise — but the arithmetic is identical; only the memory layout differs.
What a logit actually is
The V raw scores that come out of the projection are the logits. It is tempting to call a logit a “log-probability,” and that is almost right but subtly wrong in a way worth pinning down. A logit is an unnormalized score. After softmax, the relationship is
p_i = exp(logit_i) / Σ_j exp(logit_j)
log p_i = logit_i − logsumexp(logits)
logit_i = log p_i + logsumexp(logits)So a logit equals the log-probability only up to a shared additive constant — the logsumexp term, which is the same for every token at that position. This is why softmax is shift-invariant: adding any constant c to all logits leaves the probabilities unchanged, because exp(logit_i + c) factors the exp(c) out of both numerator and denominator. Only the differences between logits carry information. A logit of 12 means nothing on its own; a logit of 12 versus a logit of 4 for a competing token means the first is exp(8) ≈ 2981× more likely. Keep this straight and the rest of the pipeline — temperature, sampling, cross-entropy — falls out cleanly.
Softmax: from logits to a probability distribution
Logits are unbounded real numbers; to sample or to score, you need a proper probability distribution — non-negative entries summing to one. Softmax is the canonical map that does this:
softmax(logits)_i = exp(logit_i) / Σ_j exp(logit_j)
probs = softmax(logits_i) : [V], Σ probs = 1, probs ≥ 0The exponential does two jobs at once: it forces every score positive, and it exaggerates gaps — a token whose logit is a little higher than the rest gets a disproportionately large share of the probability mass. Softmax is monotonic, so the argmax of the logits is always the most probable token; the softmax only decides how peaked the distribution is around it. In practice you never compute softmax naively, because exp of a large logit overflows. The standard trick uses shift-invariance: subtract the max logit first, softmax(x) = softmax(x − max(x)), so the largest exponent is exp(0) = 1 and nothing overflows. This is also why the fused log_softmax and cross-entropy kernels exist — they compute the normalized log-probabilities directly and stably rather than taking a separate log of a softmax that has already lost precision.
Training: cross-entropy over every position
During training the head runs at every position, because a causal language model predicts the next token at each step in parallel. For position i with known target token y_i (the actual next token in the text), the loss is the negative log-probability the model assigned to that target:
probs_i = softmax(logits_i) : [V]
loss_i = −log probs_i[y_i]
loss = (1/N) Σ_i loss_i (mean over positions)This is exactly cross-entropy between the model’s predicted distribution and the one-hot target. Because it is computed at all N positions at once, the projection during training produces the full [N, V] logit tensor — which, for a long sequence and a large vocabulary, is a genuinely large intermediate (more on the cost below). The gradient of cross-entropy with respect to the logits has a famously clean form, ∂loss/∂logit_i = probs_i − onehot(y_i): push the probability of the correct token up, push everything else down, in proportion to how much mass the model currently misplaced. That tidy gradient is a big part of why the softmax-cross-entropy pairing is universal for language modeling.
Inference: only the last position matters
At generation time the picture changes in a way that saves real work. To sample the next token you only need the logits at the last position — the hidden state that has seen the entire prompt. The logits at earlier positions correspond to tokens that are already fixed, so recomputing their distributions buys you nothing. A well-written decode loop therefore projects only the final hidden vector through the head:
train : logits = h_final · W_out over all N positions -> [N, V]
decode : logits = h_last · W_out for 1 position -> [V]The compute difference is exactly a factor of N: 2·N·d·V multiply-adds for the full-sequence head versus 2·d·V for a single decode step. With a 128K vocabulary and a 4096-wide model that is the difference between projecting one vector and projecting thousands, per generated token. It also explains a subtle asymmetry in serving: the head is a small slice of prefill cost (amortized over the whole prompt) but a visible slice of per-token decode cost, because at decode the attention and feed-forward layers only process one new position while the head still touches the entire vocabulary.
Parameter cost: the head is often the biggest tensor
The projection holds d × V parameters, and both factors are large. Modern vocabularies run from about 32K (LLaMA 2, Phi) through 128K (LLaMA 3) to over 250K (some multilingual models), while d ranges from ~768 in the smallest models to 4096 and beyond. Multiply them and the head becomes enormous relative to a single transformer block:
| d_model | Vocab V | Head params (d×V) |
|---|---|---|
| 768 | 50,257 | ~38.6M |
| 2,048 | 32,000 | ~65.5M |
| 3,072 | 32,064 | ~98.5M |
| 4,096 | 128,256 | ~525.3M |
For a small model this one matrix can be several percent of the total parameter count — frequently the single largest weight tensor, larger than any individual attention or feed-forward matrix. The input embedding table E ∈ ℝ^(V × d) has the same shape and the same count, so a naive model carries two vocab-sized tensors. That observation is the entire motivation for weight tying, which the next section covers.
Weight tying: sharing the embedding and the head
The input embedding maps token IDs to vectors (E : [V, d]); the output projection maps vectors back to token scores (W_out : [d, V]). These are inverse-shaped operations over the same vocabulary, and a long line of work (Press & Wolf; the original transformer) showed you can tie them — use one matrix for both:
# untied: two independent tensors
E : [V, d] (input embedding)
W_out : [d, V] (separate output head)
# tied: reuse the embedding, transposed
W_out = Eᵀ logits = h_final · EᵀTying removes the entire d × V block of head parameters — for Phi-3-style dimensions (d = 3072, V = 32064) that is roughly 98M parameters saved, a real fraction of a small model’s budget. Beyond the memory win, tying has an intuitive justification: a token’s input representation and the direction you project onto to predict that token ought to be related, and sharing the matrix enforces that. Most small language models tie by default because the savings are decisive at their scale; some very large models leave them untied, having found a small quality gain worth the extra parameters when the vocab fraction is negligible.
A fully worked shape and parameter example
Make it concrete with a LLaMA-3-8B-shaped configuration: d = 4096, V = 128,256, and a batch processing a sequence of N = 2048 tokens during training.
Head parameters = d × V = 4096 × 128256 = 525,336,576 ≈ 525M
Training logits : [N, V] = [2048, 128256] = 262,668,288 values
at fp32 -> ~1.05 GB just for this one activation
Head FLOPs (train): 2 · N · d · V = 2 · 2048 · 4096 · 128256 ≈ 2.15 TFLOP
Head FLOPs (decode): 2 · d · V = 2 · 4096 · 128256 ≈ 1.05 GFLOP / tokenTwo things jump out. First, at ~525M parameters the head alone is over 6% of an 8B model — and identical in size to the input embedding, so untied you would spend ~1.05B parameters (~13%) on vocabulary tables. Second, the activation [N, V] is a gigabyte-scale intermediate during training even though the weight is “only” 525M — which is why fused cross-entropy kernels that never materialize the full logit tensor in high precision are so valuable. Now a tiny softmax to close the loop: logits [2.0, 1.0, 0.1] exponentiate to [7.39, 2.72, 1.11], sum 11.21, giving probabilities [0.659, 0.242, 0.099] — the top token takes about two-thirds of the mass, and adding the same constant to all three logits would not change that split at all.
Compute cost and where the head bottlenecks
As a matrix multiply the head costs about 2·N·d·V FLOPs (the factor of two counts a multiply and an add per element). Compared with a transformer block, whose cost scales with N·d² for the projections and N·d·d_ff for the feed-forward, the head is linear in d but linear in the typically-much-larger V. Whether it dominates depends on the ratio V versus the per-block work.
Where the head really bites is memory bandwidth, not arithmetic. During autoregressive decode every generated token must read the entire d × V weight from memory and write a V-length logit vector, while the rest of the model only processes a single new position. On a memory-bandwidth-bound decode step — which is the regime CPU inference and small-batch GPU serving live in — streaming half a gigabyte of head weights per token is a measurable share of the time budget. This is a second, performance-oriented reason tying is attractive on constrained hardware: one shared vocab tensor is half the bytes to keep resident and half the traffic through the cache hierarchy, which matters as much as the raw parameter savings on a CPU-hosted small model.
Temperature and sampling: the downstream story
Everything so far produces one probability distribution per position. Choosing a token from it is a separate stage — the sampler — and it hooks directly onto the logits before the softmax. The key knob is temperature T: divide the logits by T before normalizing.
probs = softmax(logits / T)Because softmax cares only about logit differences, dividing by T rescales every gap. A high T > 1 shrinks the gaps and flattens the distribution (more random, more diverse); a low T < 1 stretches them and sharpens it (more confident, more repetitive); and in the limit T → 0 all the mass collapses onto the top logit, which is exactly greedy argmax decoding. On top of temperature sit the truncation samplers — top-k, top-p (nucleus), min-p — that zero out the tail before renormalizing so a rare token never sneaks through. Those are their own topic; the load-bearing point here is that they all operate on the very logits this article produced, which is why a clean mental model of “logit → softmax → probability” is the prerequisite for reasoning about decoding at all.
Common pitfalls and implementation gotchas
The final stage is small but has a dense cluster of bugs. Skipping the final norm feeds an un-normalized residual stream into the head; the model runs but logits blow up and softmax saturates. Getting the tie transpose wrong — using E instead of Eᵀ, or mismatching the framework’s [V, d] storage — produces silently garbage logits with correct shapes, the worst kind of failure. Taking log of softmax separately instead of using log_softmax loses precision and can yield -inf losses; always use the fused cross-entropy.
A few more: applying temperature after softmax instead of before (which does not do what you want — it must divide the logits, not the probabilities); forgetting that softmax is shift-invariant and “fixing” logits by adding a constant, which changes nothing; and materializing the full [N, V] logits in fp32 for a long sequence when a fused kernel would avoid the gigabyte activation. For CPU-hosted small models, the practical checklist is short: tie the embedding to halve the vocab bytes, keep the head weight in the model’s working precision, and remember that at decode the head is a bandwidth cost paid on every single token — it is the last thing the model does, and on constrained hardware it is rarely free.
W_out : [d, V] that turns each final, normalized hidden vector into V raw scores — the logits. A logit is an unnormalized score, equal to a log-probability only up to a shared additive constant, so only the differences between logits matter; softmax (exp(logit_i) / Σ exp(logit_j)) turns them into a real distribution, and it is shift-invariant, which is why the numerically stable version subtracts the max. During training the head runs at every position for cross-entropy; at decode only the last position matters, saving a factor of N. Because the head holds d × V parameters — often the largest tensor in a small model, and a per-token memory-bandwidth cost at decode — weight tying with the input embedding is the standard move on constrained hardware. Temperature and the top-k / top-p samplers all attach to these same logits, so a clean “logit → softmax → probability” picture is the foundation for everything downstream.