Language-model head
The last layer of a language model. It turns the final vector at each position into one score per token in the vocabulary.
What it means
The language model head, or LM head, is the last layer of a language model. Everything before it, the body, turns the text into one vector for each position. The head turns each of those vectors into one score, called a logit, for every token in the vocabulary. A softmax then turns those scores into the probabilities of what comes next.
It's usually a single linear layer: multiply the vector by a table with one row per token.
An example with NumPy
import numpy as np rng = np.random.default_rng(0) vocab = ["the", "mug", "is", "light"] hidden = rng.normal(size=8) # the model's final vector at one position head = rng.normal(size=(len(vocab), 8)) # the head's weights: one row per token logits = head @ hidden # one raw score per token probs = np.exp(logits) / np.exp(logits).sum() # softmax for token, score, p in zip(vocab, logits, probs): print(f"{token:<6} score {score:>6.2f} probability {p:.2f}") print("next token:", vocab[int(np.argmax(probs))])
the score -1.47 probability 0.01 mug score 0.37 probability 0.08 is score -1.66 probability 0.01 light score 2.82 probability 0.90 next token: light
The weights here are random, so "light" means nothing yet. Training is what makes the right tokens score highly. In a real model the shapes go from (batch, length, width) to (batch, length, vocabulary size).
Worth knowing
The head can be large: a width of 1,000 and a vocabulary of 50,000 means 50 million weights. So many models reuse the token embedding table for it, which is called weight tying. In Hugging Face Transformers, AutoModel loads the body alone, and AutoModelForCausalLM adds this head. To classify text instead, you keep the body and swap the head for a smaller one.
Try it now
Set it to 50,257 / 1,024 / 768 / 12 and you get GPT-2 small's roughly 124 million. Width costs the most, because it's squared.
Learn it properly
The GPT architecture · A stack of decoder blocks that predicts the next token. About 35 minutes.
Other terms from the same lesson
- GPT: Generative pre-trained transformer: a stack of transformer blocks trained to predict the next token.
- weight tying: Reusing the input embedding table as the output layer, so the model needs fewer numbers to learn.
Never trained an AI?
Teach a small one to write, right in your browser. It takes about 20 seconds and needs no account.