2 min readZenith
Weight Tying in Language Models
Sharing the input embedding and output projection — why it works, when it stops working, and how it interacts with logit scale.
Weight tying shares one matrix between the input embedding and the output projection of a language model. It's one line of code with surprisingly deep consequences.
self.lm_head.weight = self.tok_embeddings.weight
Why it works
Both matrices map between the same two spaces — vocabulary and hidden state — just in opposite directions. Tying them says: the representation used to recognize a token should be the representation used to predict it. For a 50k vocab and 4096-dim model, that's ~200M parameters saved, and in small models the regularization effect actually improves perplexity.
When it stops working
Large models increasingly untie, for two well-understood reasons:
- Scale mismatch. Input embeddings want small norms (they enter a residual stream);
output projections want norms tuned for logit magnitudes. One matrix can't serve both
masters, which is why tied models often multiply embeddings by
sqrt(d_model)on the way in — the input and output roles need different effective scales. - Capacity. At multi-billion-parameter scale, the saved parameters are a rounding error, while the flexibility of separate matrices is measurable.
Zenith's decoder ties by default — at 10.7M parameters and a small vocabulary, tying is squarely in the regime where it helps.
Implementation pitfalls
- Tie by assigning the parameter, not copying data — otherwise the weights drift apart silently after the first optimizer step.
- Weight decay hits the shared matrix from two gradient paths; account for it when comparing tied vs. untied runs.
- Checkpoint loaders must preserve the tie — a save/load round-trip that silently duplicates the matrix is an easy bug to ship.
The takeaway
Weight tying is a bet that recognition and prediction are the same geometry. The bet pays off when parameters are scarce and loses gracefully when they aren't — which is why you see it in small open models and not in frontier ones.