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.