Skip to content

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.

transformerstraining

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.