Weight Tying › Summary

Weight tying exploits the structural symmetry between input embedding and output projection matrices in language models. Both matrices have shape $V \times d$, where $V$ is vocabulary size and $d$ is the hidden dimension. Instead of maintaining separate matrices $\mathbf{E}$ for input and $\mathbf{W}_{\text{out}}$ for output, weight tying sets $\mathbf{W}_{\text{out}} = \mathbf{E}$, so a single shared matrix serves both roles.

The key insights from this chapter:

• Parameter reduction: Weight tying eliminates one full embedding matrix, saving $V \times d$ parameters. For a 50K vocabulary with 768-dimensional embeddings, that is 38 million fewer parameters, often more than 15 percent of a small model's total capacity.

• Theoretical justification: Input and output embeddings both answer the question "what does this token mean?" Tying them forces a consistent semantic representation where the hidden state that produces a token is similar to that token's embedding. This is not a constraint; it is a clarification of what good representations should look like.

• Dot product as similarity: With tied weights, generating a token corresponds to finding which embedding most closely aligns with the current hidden state. The transformer's task is to produce hidden states that are similar to the embeddings of the correct next tokens.

• Gradient dynamics: Tied weights receive gradients from both the input path (through transformer layers) and the output path (directly from cross-entropy loss). This provides more learning signal per token occurrence and acting as implicit multi-task learning.

• Encoder-decoder extension: In sequence-to-sequence models, all three embedding matrices (encoder input, decoder input, decoder output) can share weights when the vocabulary is shared, eliminating two-thirds of embedding parameters. T5 uses this approach.

• Embedding scale: The $\sqrt{d}$ scaling required for input stability should be applied at forward time, not baked into the weight matrix, so that the output projection sees unscaled embeddings.

• When to use: Weight tying is recommended when vocabularies are large relative to model depth and input and output domains are symmetric. Avoid it when vocabularies differ, when fine-tuning models trained without tying, or when empirical evidence in your specific task suggests separate matrices perform better.

Modern language models nearly universally adopt weight tying, making it a fundamental architectural decision rather than an optimization. Understanding why it works deepens your understanding of what these models are learning: a unified representation space where token identity is consistently encoded whether a token is being read or generated.