ML//transformer//LM head

The LM head is the final layer of a language model, which converts the last vector of the residual stream, a point in the latent space, into logits over the vocabulary. Its input is called the **final hidden state** (or last hidden state): the residual stream vector after all \(N\) layers, before this projection.

The LM head is the final layer of a language model, which converts the last vector of the residual stream, a point in the latent space, into logits over the vocabulary. Its input is called the final hidden state (or last hidden state): the residual stream vector after all NNN layers, before this projection.

1Residual stream768 dims2LM headWEW_EWE​ transposed3Logits50K vocab4Softmax picks the token

WUW_UWU​ is the embedding matrix transposed (weight tying): the same parameters that encode tokens also decode the output.

Each row represents a token of the vocabulary, and the dot product between the final hidden state and each row measures how aligned the representation is with each possible next token.

WUW_UWU​ is as wide as the last embedding dimension, so the tokenizer and embedding size determine the bottleneck.