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 — GPT-2, Llama, Mistral — use weight tying: the input embedding matrix and the output head share the same parameters. Mathematically, W_lm = E.T. This saves vocab × d_model parameters (524M for Llama 3 8B) and creates a satisfying symmetry: tokens with similar input embeddings get similar output logit profiles. Some larger models (GPT-3, GPT-4) don't tie weights — 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

The Output Head: From Hidden States to Logits

Weight tying

Code

External links

Exercise

Progress

Comments 0