Vol. V · The Transformer Plate 7 of 10output head, unembedding, weight tying, logits, embedding tableAL

The Output Head.

The same table can read tokens in and score them out. Whether it should is a trade the course measured.

After the last block and the final norm, the model has one vector of width d per position. The output head, or unembedding, multiplies it by a d × V matrix to get a logit for every token. That matrix has exactly the shape of the embedding table transposed, and since 2016 many models have simply reused it: weight tying. A token’s row then serves twice — as the vector that represents it when read, and as the direction the final vector must point in for it to be predicted.

Tying saves V × d parameters, which for a small model with a large vocabulary is a great deal — for micro it is more than half the model — and it gives rarely seen tokens’ output rows some training through their input role. Untying gives the model separate freedom in the two roles. Large models increasingly untie, because the extra parameters are a small share for them. The workshop compares two recorded runs of micro’s shape that differ only in this choice.

Below: two recorded runs of micro’s shape for 6 000 steps on the same data, one with a tied output head and one untied. Curves show training (light) and held-out (dark) loss; the readings give parameter counts and the final difference.

What you will learn

Describe the output head as a d × V matrix producing logits.
Explain weight tying and what it saves.
Compare tied and untied training runs.
Say when tying is worth it.

Workshop

Recorded runs of micro’s shape for 6 000 steps: tied output head (blue) and untied (red); light lines are training loss, dark lines held-out loss.

One table or two

Tied, micro’s shape has 1.85 million parameters; untied, 2.9 million — 57% more, all in the new output matrix. After 6 000 identical steps the untied model’s held-out loss is 4.148 nats against 4.208 for the tied one: an improvement of 0.060, bought with half again as many parameters. Whether that is worth it depends on what else those parameters could buy; spent on more layers or width instead, the same budget would probably do as well or better, which is why small models tie.

In large models the arithmetic flips. For Llama 3 8B an untied head of 128 256 × 4 096 adds about 0.5 billion parameters, about 7% of the total and very little compute per token compared with the blocks, and the separate matrix lets input and output representations specialise; Llama 3 is untied, as are most large open models of 2024–2025, while some smaller ones still tie. A tied head has one more quirk visible in micro: the logit for a token is a dot product with that token’s own embedding, so the model must learn to suppress the easy prediction “the next token is the same as this one” — one reason repeated-token loops are a characteristic failure of small tied models.

Weight tying for language models was proposed independently in 2016 by Ofir Press and Lior Wolf and by Hakan Inan, Khashayar Khosravi and Richard Socher, who showed it reduced perplexity as well as parameters for the LSTM models of the time.

Terms

Output head (unembedding): The d × V matrix that turns the final vector into a logit for every token.
Weight tying: Using the embedding table, transposed, as the output head.
Logit: The dot product of the final vector with a token’s output row.
Untied head: A separate output matrix with its own V × d parameters.

The output head is a d × V matrix; tying it to the embedding saves V d parameters, which matters most for small models. Untied, micro learns slightly better for 57% more parameters.