The last matrix multiply is often the largest single one in the model, and on a 2B model with a 256k vocabulary it costs more than several transformer blocks.
Why you'd careThe thing you have already noticed
You swap a 7B model for a 2B one expecting roughly triple the decode speed and get maybe 1.6x. Or you turn on echo with logprobs to score a long prompt and the server runs out of memory on a request that prefilled fine a second earlier. Both are the LM head. It is one matrix multiply of shape [d_model, vocab_size], and vocab_size has grown much faster than d_model — 32,000 in Llama 2, 128,256 in Llama 3, roughly 256k in Gemma. On a 2B-class model with a 256k vocabulary the head touches close to a quarter of the parameters and reads a quarter of the weight bytes on every single decode step. Shrinking the body does not shrink it.
In and outWhat goes in, what comes out
| In | The final hidden state after the last transformer block and the final RMSNorm: a bf16 tensor of shape [batch, 1, d_model] at decode, or [batch, seq_len, d_model] during prefill. Plus the unembedding matrix of shape [d_model, vocab_size], possibly sharded across tensor-parallel ranks by vocabulary. |
|---|---|
| Process | One dense GEMM, logits = h W_U. Usually accumulated or cast to fp32 before the softmax. Under vocabulary-parallel tensor parallelism each rank computes its own vocabulary slice and the result is all-gathered or reduced. During prefill most engines run the head only on the last position of each sequence. |
| Out | An unnormalised logits tensor of shape [batch, vocab_size], typically fp32, with values in roughly the −20 to +25 range. No softmax applied yet. If logprobs were requested, a top-n slice plus the log-softmax normaliser is retained for the response. |
The residual-stream direction is preserved exactly: logits are a lossless linear read-out of the hidden state. What is lost is everything below the retained top-n once logprobs are truncated, and the hidden state itself, which is freed immediately. You cannot go back and recompute the head at higher precision, and positions skipped during prefill can never be scored afterwards without a second forward pass.
ConceptThe idea underneath
Think of the residual stream as a single vector in d_model dimensions that every layer has been adding to. The unembedding matrix has one column per vocabulary token, and each column is a direction in that same space. The head computes logit_i = h · W_U[:, i] — a dot product. So the question it asks is literally “how far does the current residual point along the direction that means cat?” Nothing is learned or decided here; it is a similarity read-out of work already done.
Press and Wolf showed the input embedding and this output matrix can be the same weights (weight tying), which is why small models often pay for the vocabulary only once. Gemma-2-2B ties them: the shared matrix is 2304 × 256,128 ≈ 590M parameters against a 2.6B total. Tying saves storage, not work — decode still reads all 590M and does the full matmul on every step.
There is a real expressiveness limit here too. A single linear head means the matrix of log-probabilities it can produce has rank at most d_model. Yang et al. called this the softmax bottleneck: with a 2048-dimensional residual and a 256k vocabulary, the model cannot represent arbitrary next-token distributions, only a low-rank family of them.
Cost: the head is 2 × d_model × vocab_size FLOPs per token. A transformer block is about 24 × d_model². At d_model = 2304 and a 256k vocabulary, the head is worth roughly ten blocks — and Gemma-2-2B has 26 of them. That is the whole reason this stage has a page.
At a glanceSee it
The final hidden vector is multiplied by the unembedding matrix to score every token in the vocabulary.
The knobsHyperparameters and nuance
- tie_word_embeddings(HuggingFace config) —
truereuses the input embedding as the head, savingd_model × vocab_sizeparameters. Small models tie; most frontier models untie. Untying costs memory and adds a second matrix to load; tying couples input and output representations, which measurably hurts very large models. - vocab_size(tokenizer choice, fixed at training) — a bigger vocabulary means fewer tokens per document, so cheaper attention and shorter sequences, but a linearly bigger head. Too small and sequences balloon; too large and the head dominates decode while rare tokens stay undertrained.
- logprobs / top_logprobs(OpenAI-compatible), prompt_logprobs (vLLM) — how many alternatives are kept, typically 0–20.
prompt_logprobsforces a head pass over every prompt position, so peak memory scales asseq_len × vocab_size. This is the single most common way to OOM a healthy server. - logits dtype / fp32 upcast(engine-level, e.g. vLLM computes the softmax in fp32) — a bf16 head has roughly 1e−2 relative error, which is larger than the gap between the top two candidates at flat positions. Cheaper and faster; changes which token wins on near-ties.
- --tensor-parallel-size(vLLM) — shards the vocabulary across ranks and inserts a collective into every decode step. Higher degrees cut per-rank head memory but add fixed latency per token and change the floating-point reduction order, which moves logits in the last bits.
EffectHow this stage moves the answer
Most of the time this stage is invisible, because the top token wins by a wide margin. It turns the answer at exactly the positions where it does not. Run the head in bf16 instead of fp32 and near-ties flip: the model picks a different synonym, a different digit, a different closing token, and from there the whole continuation diverges. The place this is worst is logprob-based classification — the pattern where you ask for a single Yes/No token and read its probability rather than its identity. That probability is the head's output, and a bf16 head makes it noisy in the third decimal, which quietly breaks any threshold you calibrated. Vocabulary size turns answers in a different way: it decides how numbers and identifiers are split, so two models with identical architectures but different tokenizers make different arithmetic mistakes.
EvalsWhat it does to your measurements
Perplexity is computed directly from this tensor, so anything that perturbs it moves perplexity by construction. More subtly, a large share of multiple-choice benchmarks — MMLU as scored by the EleutherAI lm-evaluation-harness, for instance — do not generate at all; they compare the logprobs of the candidate continuations. Those scores are a direct measurement of this stage, so head dtype, vocabulary-parallel reduction order and logprob truncation all move them. The silent invalidation: change --tensor-parallel-size from 2 to 4 between runs and the vocabulary shards, and therefore the reduction order, change. Logits shift in the last bits, near-ties flip, and you get a few tenths of a point of difference with nothing in the run config to explain it. Also check whether your harness scores by logprob or by generation — the two give different numbers for the same model and are routinely compared as if they did not.
Failure modesWhen it goes wrong
- Out-of-memory on a request that prefilled fine
prompt_logprobsorechomaterialised an fp32[seq_len, vocab_size]tensor instead of a single row. - A 2B model is barely faster than a 7B one at batch 1the head is a fixed cost that does not shrink with the body, and it dominates when the body is small and the vocabulary is large.
- The same prompt returns different top-2 logprobs on two deploymentsone computes the head in bf16, the other upcasts to fp32 before the softmax.
- Rare garbage tokens appear at high temperaturethe vocabulary was padded to a multiple of 64 or 128 for kernel alignment and the padding rows were never masked to −inf.
- Benchmark scores drift after a parallelism changevocabulary sharding changed the all-reduce order; the model weights are identical.
PapersWhere this comes from
- Using the Output Embedding to Improve Language ModelsPress & Wolf, 2016 (arXiv:1608.05859; EACL 2017). Established weight tying between the input embedding and the output projection, and showed it improves quality as well as saving parameters. This is why the head and the embedding table are frequently the same tensor.
- Breaking the Softmax Bottleneck: A High-Rank RNN Language ModelYang et al., 2017 (arXiv:1711.03953). Showed that a single softmax over a linear projection caps the rank of the representable log-probability matrix at
d_model. It is the reason the head is an expressiveness constraint and not only a cost centre. - Efficient softmax approximation for GPUsGrave et al., 2016 (arXiv:1609.04309; ICML 2017). Introduced adaptive softmax, and established the point that still holds: the cost of the output layer grows linearly with vocabulary size, which is what makes the layer worth restructuring at all. Cite it for the vocabulary-scaling argument only. It is an RNN language-model paper that predates the Transformer and contains no attention layer, so it cannot be used to support a head-versus-attention comparison — that framing belongs to later serving measurements, not to this paper.