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. This is mechanistic evidence that weight tying optimizes the embedding matrix for output prediction, compromising its role in input representation. These results help explain why weight tying can harm performance at scale and have implications for training smaller LLMs, where the embedding matrix contributes substantially to total parameter count.

翻译：权重绑定（即共享输入与输出嵌入矩阵的参数）是语言模型设计中的常见做法，但其对所学嵌入空间的影响机理尚不明确。本文表明，相较于参数不绑定的可比较模型，绑定后的嵌入矩阵与输出矩阵（解嵌入矩阵）的排列一致性显著高于其与输入嵌入的对应关系，这表明共享矩阵主要服务于输出预测而非输入表示。这种解嵌入偏差源于训练初期输出梯度的主导作用。通过调谐透镜分析，我们发现该偏差对早期层计算产生负面影响，导致其对残差流的贡献效率下降。在训练过程中放大输入梯度可缓解此偏差，为梯度失衡的因果作用提供了证据。这一机制性证据表明，权重绑定为优化输出预测而调整嵌入矩阵，从而削弱其在输入表示中的作用。这些发现有助于解释为何权重绑定在大规模场景下可能损害模型性能，并对训练中小型语言模型（其中嵌入矩阵参数占总参数量显著比例）具有重要启示。