Weight Tying Biases Token Embeddings Towards the Output Space
Field: NLP Authors: Antonio Lopardo, Avyukth Harish, Catherine Arnett arXiv: 2503.23753
Abstract
Weight tying, i.e. sharing parameters between input and output embedding matrices, is common practice in language model design, yet its impact on the learned embedding space remains poorly understood. In this paper, we show that tied embedding matrices align more closely with output (unembedding) matrices than with input embeddings of comparable untied models, indicating that the shared matrix is shaped primarily for output prediction rather than input representation. This unembedding bias arises because output gradients dominate early in training. Using tuned lens analysis, we show this negatively affects early-layer computations, which contribute less effectively to the residual stream. Scaling input gradients during training reduces this bias, providing causal evidence for the role of gradient imbalance. These results provide mechanistic evidence that weight tying optimizes the embedding matrix for output prediction at the expense of its role in input representation. They help explain why weight tying may harm performance at scale and have implications for training smaller LLMs, where the embedding matrix accounts for a large fraction of total parameters.
Key Takeaways
- Tied embeddings lean towards the output space: tied embedding matrices align more closely with unembedding matrices than with input embeddings from comparable untied models.
- Gradient imbalance is the cause: output gradients dominate early in training, shaping the shared matrix primarily for prediction.
- Downstream effect: tuned lens analysis shows the bias weakens early-layer contributions to the residual stream.
- Causal evidence: scaling input gradients during training reduces the bias, confirming the gradient imbalance mechanism.
- Implications: results help explain why weight tying may hurt performance at scale and matter for small LLMs, where embeddings are a large parameter share.