Proves that weight tying—a standard LLM efficiency trick—biases embeddings toward output prediction and actively harms early-layer input representations.
March 30, 2026
Original Paper
Weight Tying Biases Token Embeddings Towards the Output Space
arXiv · 2603.26663
The Takeaway
Mechanistic analysis reveals that output gradients dominate the shared matrix, providing a clear explanation for why weight tying fails at larger scales and suggesting that scaling input gradients can fix this imbalance.
From the 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 arise