The Output Head: From Hidden States to Logits

The final hidden state — the residual stream after the last block — has shape (seq_len, d_model). To produce predictions, you must turn that into logits over the vocabulary: shape (seq_len, vocab_size). This is the output head.

Mechanically, it's a linear projection: logits = hidden @ W_lm.T, where W_lm has shape (vocab_size, d_model). Softmax over the last dim gives probabilities. During inference you typically only need the last position's logits (next-token prediction); during training you compute logits at all positions in parallel.

Weight tying

Many models use weight tying: the input embedding matrix and the output head share the same parameters. Mathematically, W_lm = E.T. GPT-2 ties, and so do most of the small variants of the modern families. Tying saves vocab × d_model parameters — for a Llama-3-8B-shaped model that would be 524M — and creates a satisfying symmetry: tokens with similar input embeddings get similar output logit profiles. The flagship open-weight models mostly do not tie, though: Llama 3 8B and Mistral 7B both keep a separate lm_head, which is exactly why Llama 3 8B lands at 8.03B rather than 7.50B. Larger models (GPT-3, GPT-4) don't tie either — at huge scale the parameter savings become a smaller fraction of the total, and untying gives a tiny quality bump.

Code

Tied output head·python

class TiedLM(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
    def forward(self, ids):
        x = self.embed(ids)             # (B, L, d_model)
        # ... transformer body ...
        x = final_norm(x)
        # Reuse embedding weights as the output projection
        logits = x @ self.embed.weight.T     # (B, L, vocab_size)
        return logits

Weight tying

Code

External links

Exercise

Progress

Comments 0