The language model head, often written as lm_head, is the last major numerical transformation before sampling. By the time execution reaches this point, the model has already tokenized the prompt, looked up embeddings, passed through many attention and feed-forward blocks, and applied a final RMSNorm. What remains is simple to state: take the final hidden state of the current token position and compare it against every token in the vocabulary. The result is one score per possible next token.
That is not a metaphor. It really is just a giant projection. If the hidden width is d_model and the vocabulary size is V, then the output weight matrix has shape [V, d_model] or an equivalent transposed layout depending on the framework. Every row corresponds to one vocabulary item. The model computes a dot product between the current hidden state and each row. Those dot products are the logits.
This is why the output head matters conceptually even though it is mechanically plain. Earlier layers spend all their effort building the hidden state into something useful. The lm_head is where that usefulness is cashed out into an explicit choice space. If the hidden state after the prompt "What is the capital of France?" is well-formed, the row corresponding to " Paris" should produce one of the highest scores.
Key idea: the lm_head does not contain a tiny symbolic dictionary that says France → Paris. It contains one learned row per token, and scoring is just a huge bank of similarity checks between the current hidden state and those rows.
Why this stage exists
The transformer core works in hidden-state space, not vocabulary space. Attention and FFN layers manipulate vectors whose dimensions are optimized for the model's internal representation, not for human-readable token identities. Eventually the system has to return to the vocabulary. The lm_head is the bridge from internal feature space back to discrete next-token choices.
Without this projection, the model could build a sophisticated latent representation of the prompt but would have no mechanism for turning that representation into candidate outputs. The lm_head is therefore not optional plumbing. It is the decision surface. Sampling cannot happen until the model has produced explicit scores for actual tokens.
That is also why the final RMSNorm usually sits immediately before this stage. The output projection compares one hidden vector against tens of thousands of learned rows. If the hidden vector's scale drifts unpredictably, the logits can become noisier or harder to interpret downstream. Normalizing right before the final score computation improves stability for this very wide comparison.
The mathematics, immediately translated
logits = h · W_vocabᵀ
Here h is the final hidden state for the current token position, with shape [d_model]. W_vocab is the output matrix with shape [V, d_model]. The multiplication produces a vector logits of shape [V]. Each output element is one dot product between h and one vocabulary row.
In English: compare the current hidden state to every possible next token identity. If a row aligns strongly with h, that token gets a high score. If it aligns poorly, it gets a low score. The logits are not probabilities yet. They are just raw scores. Sampling will later convert them into a choice distribution.
Because the calculation is a dot product, every row of W_vocab acts like a learned direction in model space. The question being asked for each token is: how similar is the current hidden state to this token's output direction? The best-scoring rows are the candidates the model believes fit the context most strongly.
Tensor shapes
| Tensor | Meaning | Typical shape | Example |
|---|---|---|---|
| h | Final hidden state for the last position | [d_model] | [4096] |
| W_vocab | Vocabulary projection matrix | [V, d_model] | [32000, 4096] or [128000, 4096] |
| row_i | Output row for token i | [d_model] | One row for " Paris", one for ".", one for " banana", and so on |
| logits | One score per vocabulary token | [V] | [32000] or [128000] |
If the vocabulary is 32K and the hidden width is 4096, the output matrix contains over 131 million parameters. That is one reason the lm_head deserves its own chapter. Even though the operation is conceptually one GEMV, it touches a lot of memory. On bigger vocabularies the cost becomes even more obvious.
A running example: why “ Paris” should win
Imagine the model has processed the full prompt "What is the capital of France?" and produced a final hidden state h for the last token position. That vector is not literally the string Paris. It is a compact numerical summary of the whole prompt: question form, geography context, expectation that an answer should follow, and many subtler internal features created by prior layers.
Now the lm_head computes one dot product per vocabulary token. The row for " Paris" may line up strongly with those answer features. The row for " London" may line up somewhat, because it is also a capital city. The row for " banana" should line up poorly. The row for "." might receive a nonzero score, but ideally not the top one. In a good model, the ordering of those scores reflects contextual plausibility.
Nothing mystical happens in the final step. The model does not suddenly reason symbolically at the output head. It simply exposes, through dot products, what the hidden state already makes likely. If the hidden state is excellent, the lm_head reveals that excellence. If the hidden state is weak or confused, the lm_head faithfully reveals that weakness too.
lm_head weight types
The output head uses the same basic idea regardless of storage format: compare h to vocabulary rows and emit logits. But implementations differ in how those rows are stored, whether they are tied to embeddings, and whether precision is reduced. The storage choice affects memory, bandwidth, and occasionally output quality.
The vocabulary projection is the final decision layer. It does not care whether the rows come from a standalone output matrix, a shared embedding table, or a quantized format. In all cases the conceptual job is the same: one learned row per token, one score per row.
Full-precision lm_head weights are the straightforward dense matrix form. FP32 is accurate but large; FP16 and BF16 reduce memory and bandwidth. Some systems keep the output head at a slightly higher precision than internal FFN weights because final-token ranking quality can be sensitive to aggressive compression.
Quantized lm_heads store vocabulary rows in INT8, INT4, or similar formats with scales and possibly zero-points. This can reduce memory traffic significantly, especially for large vocabularies. The trade-off is that logit ordering, particularly near the top of the distribution, may become more fragile if quantization is too aggressive.
Tied weights mean W_vocab = Eᵀ, where E is the token embedding table. In plain English, the matrix used to look tokens up on the way in is reused, transposed, on the way out. This saves memory and sometimes improves consistency between input and output token geometry.
Untied weights use a separate learned output matrix. This costs more memory but gives the model freedom to learn a different geometry for deciding outputs than for representing inputs. Some architectures prefer that extra expressiveness even though tied weights are elegant and economical.
Tied versus untied is especially worth understanding because it changes parameter count dramatically for large vocabularies. If the embedding table is already enormous, reusing it as the output matrix can save a meaningful amount of memory. But tying is not always chosen. Some model families want separate input and output roles, and that extra freedom can matter enough to justify the cost.
Why weight tying is elegant but not universal
Weight tying appeals to engineers because it removes duplication. The token geometry learned on the way into the model becomes the token geometry used to score outputs on the way out. That can improve memory efficiency and make the whole architecture feel cleaner. If a token is represented in one place, why not reuse that representation rather than store a second giant table that often wants similar structure anyway?
But the output job is not identical to the input job. Input embeddings help the model represent tokens as ingredients in a context. Output rows help the model discriminate among candidate next tokens under highly conditioned hidden states. Those goals overlap, but they are not perfectly the same. Untied output heads buy freedom at the cost of memory. Tied heads buy efficiency at the cost of some expressiveness.
That trade-off becomes more important as vocabulary size grows. With 100K or 200K tokens, the output matrix is too large to treat casually. A design that saves one full copy of it may unlock a deployment target that would otherwise be awkward. So tying is not merely aesthetic; it is often a systems decision.
Why the lm_head can be expensive
The lm_head is only one matrix multiply per decoding step, but it is a wide one. The hidden state is small compared with the vocabulary axis. You have to touch one row per token in the vocabulary, and each row has d_model elements. If V = 128K and d_model = 4096, that is a lot of weight data to stream just to choose the next token.
During single-token decode, this is typically a matrix-vector operation rather than a large matrix-matrix operation. That matters because matrix-vector work tends to be less compute-dense and more bandwidth-sensitive. There is not enough arithmetic reuse per loaded weight to hide all memory costs. The core of the problem is simple: each output row is used once for this token, then discarded.
This is the same reason the output head is a natural place to think about quantization, weight tying, and chunked execution. The logic is easy. The data movement is not. When people say the lm_head is \"just another matrix multiply,\" that is true, but it should not mislead you into thinking it is free.
Streaming and chunked implementations
One practical implementation strategy is to process the vocabulary matrix in chunks. Instead of loading all rows conceptually at once, you stream through a block of rows, compute their logits, update your running top-k or output buffer, then move to the next block. This does not reduce the exact amount of arithmetic for full logits, but it improves memory locality and can make large vocabularies easier to handle on constrained hardware.
for (int base = 0; base < V; base += BLOCK_ROWS) {
int rows = min(BLOCK_ROWS, V - base);
load_rows(W_vocab + base, rows, d_model);
for (int r = 0; r < rows; ++r) {
logits[base + r] = dot(h, W_vocab[base + r]);
}
}Chunking is especially useful when you only need the best candidates rather than every logit stored permanently. You can compute logits block by block and maintain a running heap of top values. The full scores still exist mathematically, but operationally you do not have to materialize a giant dense vector in the most naive way.
On accelerators, the same idea becomes tiled GEMV or small GEMM execution. The tuning problem is to choose a tile size that respects memory bandwidth, cache or shared-memory capacity, and vectorized dot-product efficiency. Too small a tile wastes launch and bookkeeping overhead. Too large a tile may spill or underutilize the memory hierarchy.
Batched serving changes the shape of the problem
When many requests are decoded together, the lm_head stops looking like one vector against a huge matrix and starts looking more like a small matrix against a huge matrix. That shift matters. Batched dot products reuse output rows across multiple hidden states, which improves arithmetic density and often makes accelerators happier. The output head may still be bandwidth-sensitive, but the cost profile becomes less pathological than single-request decode.
That is why production servers often try to batch tokens from different users into the same decode step. The model logic is unchanged, but the hardware sees a better-shaped workload. Inference serving is therefore partly a scheduling problem: how do you align arrival times so enough hidden states hit the lm_head together without making users wait too long?
This is also where prefill and decode diverge sharply. During prefill you may score many positions or many prompts in a throughput-friendly shape. During interactive decode you often care about the latency of the next token for one user. The exact same lm_head math therefore lives two different lives depending on serving mode.
Can you avoid scoring all V tokens?
For exact next-token logits, usually no: the plain definition requires one score per vocabulary token. If you want the true argmax or a true sampled distribution over the whole vocabulary, you eventually need the scores or equivalent bounds for all rows. There are approximate tricks, but they are approximate or architecture-specific.
Some systems maintain only top-k candidates while scanning the vocabulary, which saves memory but not the fundamental need to inspect all rows. Other ideas use approximate nearest-neighbor search, hierarchical vocabularies, candidate pruning from a smaller draft model, or restricted candidate sets in specialized applications. These can work, but they change the inference contract. The default transformer head is full-vocabulary scoring.
This is one reason speculative decoding is interesting. A smaller draft model can propose candidate continuations, and the larger model verifies them. That does not eliminate the output head forever, but it can amortize how often the expensive full model must perform the whole decode cycle. The lm_head remains central; the system merely changes when and how often it is exercised.
Speculative decoding connection
Speculative decoding does not replace the lm_head math with magic. It rearranges the workflow. A cheaper model predicts several likely next tokens. The larger model then checks whether those guesses are acceptable. If they are, multiple tokens are accepted with fewer full-model passes. If they are not, the larger model falls back to ordinary scoring and sampling. In both cases the final authority is still the large model's probability structure, which ultimately flows through its vocabulary projection.
The benefit is systems-level amortization. If the large model can verify several tokens in one go or avoid some redundant single-token passes, the cost of repeatedly running the output head and all preceding blocks is reduced per accepted token. That is a serving optimization, not a change in what the lm_head fundamentally is.
Logits are not probabilities, and that separation is useful
The lm_head should ideally stop at logits. Leaving them unnormalized preserves flexibility. The next stage may apply temperature, top-k filtering, top-p filtering, repetition penalties, or domain-specific masking before any probability distribution is finalized. If the output head tried to do too much, it would entangle raw scoring with policy decisions that belong to decoding.
This separation is good engineering hygiene. The output head answers the question what scores does the model assign? Sampling answers the question how should the system turn those scores into a choice? Keeping them distinct makes debugging easier, lets you compare raw model quality separately from decoding policy, and avoids baking sampling assumptions into the projection kernel itself.
Implementation pseudocode
void lm_head(float* logits,
const float* h,
const Matrix W_vocab, // [V, d_model]
int V,
int d_model) {
for (int token = 0; token < V; ++token) {
float acc = 0.0f;
for (int j = 0; j < d_model; ++j) {
acc += h[j] * W_vocab[token][j];
}
logits[token] = acc;
}
}Again, the pseudocode is intentionally plain. A production kernel will tile, vectorize, quantize, maybe fuse bias handling, maybe accumulate in higher precision, and probably avoid the most literal storage pattern. But the logical operation remains one dot product per vocabulary row.
Hardware behavior: CPU, GPU, FPGA
On CPUs, the lm_head is usually bandwidth-sensitive during decode. The hidden vector is small and hot; the vocabulary matrix is huge and cold. Good implementations rely on vectorized dot products, careful blocking, and often quantized rows to reduce memory traffic. If the vocabulary is very large, memory locality dominates almost every other concern.
On GPUs, the output head benefits from highly optimized GEMV or GEMM-like kernels, especially during batched prefill or batched decode. But single-token, single-request decode can still leave efficiency on the table because there is only one live hidden vector. Chunking, batching across requests, and low-bit weight formats matter a lot here.
On FPGAs or custom accelerators, the lm_head is a streaming dot-product farm. The design is regular, but the weight volume is substantial. If the device cannot feed rows quickly enough, compute lanes idle. That makes the output head a good example of an operator whose mathematical simplicity does not guarantee cheap deployment.
Memory bandwidth analysis
Suppose the vocabulary has 32,000 rows and d_model = 4096. In FP16, each row is 8192 bytes. The full matrix is roughly 256 MB. You do not necessarily move every byte from DRAM for every token if caches and tiling help, but the order of magnitude makes the serving challenge obvious. This is not a tiny epilogue. It is a major memory object consulted every decoding step.
That is why quantization can help disproportionately. If INT8 or INT4 weights cut the bytes per row, you reduce the pressure on the memory system. Whether the speedup materializes depends on kernel maturity and dequantization overhead, but the direction is clear: the output head rewards bandwidth reduction.
Database and systems analogies
The lm_head resembles a giant indexed candidate scoring pass. Imagine a query engine that has already built a compact representation of what the answer should look like and now scans a dictionary table of possible outputs, computing a similarity score for each row. The heavy reasoning happened upstream. This stage is the final ranking over explicit candidates.
Another analogy is a search engine reranker with a fixed candidate set equal to the entire vocabulary. The hidden state is the query embedding. The vocabulary rows are item embeddings. The dot products are relevance scores. The highest scores become the next-token candidates. The transformer literature uses its own words for this, but the geometric pattern is familiar to anyone who has built vector search systems.
What can be cached
The weight matrix itself is persistent model state and benefits from normal device residency, packing, and ordinary hardware caches. But the logits are ephemeral. Each new token position has a new hidden state, so you recompute the scores. There is no long-lived analogue to the KV cache here.
If weights are tied, the embedding table and output matrix share storage conceptually, which can simplify residency planning. But that still does not create a semantic output cache. The only reusable state is the parameter data and any packed low-level representation of it.
Operation and bound classification
| Question | Answer for lm_head |
|---|---|
| Primary operation class | Projection from hidden-state space into vocabulary-score space |
| Secondary interpretation | Massive bank of dot products or similarity checks |
| Dominant bound | Often memory bandwidth and latency during decode, especially for large vocabularies |
| Produces probabilities? | No, it produces logits; softmax and sampling come afterward |
| Mixes tokens together? | No, it uses only the current position's hidden state |
Common implementation mistakes
One mistake is forgetting that the matrix orientation may be stored transposed relative to the math notation. Another is applying softmax too early and then trying to do top-k or sampling in a numerically awkward way. The output head should emit logits first. Probability normalization belongs to the sampling stage.
Another common mistake is quantizing the output head too aggressively without checking top-token ranking quality. A tiny average error in logits can still change the winning token when the top candidates are close. If your model suddenly confuses " Paris" and " Lyon" in borderline contexts, the output head precision may be part of the problem.
A more subtle systems mistake is optimizing only the core dot product and ignoring batching strategy. In production, the output head's speed depends on whether you can batch requests, how weights are packed, and whether streaming top-k avoids unnecessary memory writes. The kernel is important, but the serving context around the kernel is often equally important.
If you were implementing this yourself
Start with a reference dense implementation that produces exact logits for a tiny model and compare against a known framework. Verify not just the top token but several nearby logits. Then add softmax and sampling in a separate stage so the responsibilities stay clean. The lm_head should be a pure score generator.
Once the reference works, decide whether your deployment cares more about memory or precision. If memory dominates, consider quantization or tied weights. If output ranking quality is fragile, keep the head at a safer precision even if the inner layers are more aggressively compressed. The output head is one of the last places where small errors can visibly change user-facing token choices.
Most importantly, remember the mental model: the lm_head is not a magical language oracle. It is a final projection from [d_model] to [V]. Every optimization, every bug, and every hardware trade-off becomes easier to reason about when you keep that brutally simple fact in view.
That simplicity is a gift: if the model predicts badly here, the hidden state, the rows, or the serving mechanics are wrong, and you can investigate those causes directly.