Machine Learning · Transformers

Language Modeling Head

Reference entry · last updated September 19, 2026

Language modeling head is the linear map from the top Transformer block hidden vector to \(V\) logits, one per vocabulary entry, then a softmax that yields next-token probabilities.[1] It sits after the stack. It is not a layer inside a block.

1. First Principles: Vocabulary-Sized Output, Not a Lookup

Let \(h \in \mathbb{R}^{d_{\text{model}}}\) be the hidden vector at one position after the last block. The head applies a learned linear map to a vector of \(V\) unnormalized scores (logits), then a softmax:

where \(W \in \mathbb{R}^{d_{\text{model}} \times V}\) (or the transpose of an embedding matrix of shape \(V \times d_{\text{model}}\)). Each coordinate of \(p\) is a probability for one vocabulary token. Vaswani et al. describe this as a learned linear transformation and softmax that convert decoder output to predicted next-token probabilities.[1]

The input embedding is a lookup. A token ID selects one row of an embedding table. The language modeling head is the inverse direction: a dense multiply that scores every vocabulary row against \(h\). Cost scales with \(V \times d_{\text{model}}\), not with a constant-time index. That contrast is already stated on the Transformer architecture page.

2. Weight Tying

Vaswani et al. share the same weight matrix between the two embedding layers and the pre-softmax linear transformation, following Press and Wolf.[1][2] Tied weights mean the vector that represents a token on the way in is the same vector used to score that token on the way out, up to transpose.

Tying removes a separate \(d_{\text{model}} \times V\) parameter block. Untied heads keep an independent output matrix. The 2017 Transformer uses tying. Later models may choose either arrangement. This page does not claim tying is required for every decoder.

3. Decoding Is Not the Head

The head produces a probability distribution over the vocabulary. Decoding is the rule that turns that distribution into the next token: greedy arg max, temperature sampling, or truncated sampling such as top-k. Those rules are not part of the linear map. They are described under decoding strategies on the architecture page.

Generation remains autoregressive: one token is chosen, appended, and the stack plus head run again. The head does not store past keys or values. That cache lives in the attention sub-layers of each block.

4. Limits

• The head does not mix information across positions. Mixing already happened inside the stack.
• Softmax over \(V\) scores every vocabulary entry. Truncation (top-k, nucleus) is a decoding choice applied after those scores exist.
• Weight tying is a parameter-sharing choice in Vaswani et al. and Press and Wolf. It is not a proof that every language model ties input and output embeddings.[1][2]