共享参数让嵌入层更偏向输出预测,影响输入表示质量。
Weight Tying Biases Token Embeddings Towards the Output Space
- 通过共享输入输出嵌入矩阵,模型更优化输出预测而非输入表示。
- 训练初期输出梯度主导,导致嵌入空间向输出方向偏移。
- 调整输入梯度可缓解偏差,对小模型训练有重要启示。
权值共享(weight tying)是语言模型设计中的常见做法,但其对学习到的嵌入空间的影响仍不明确。本文发现,共享的嵌入矩阵与输出(解嵌入)矩阵的对齐程度高于未共享模型的输入嵌入,表明共享矩阵主要服务于输出预测而非输入表征。这种解嵌入偏差源于训练早期输出梯度占主导地位。通过调优的透镜分析(tuned lens analysis),我们发现这会削弱早期层对残差流的贡献。在训练中放大输入梯度可减少该偏差,提供了梯度不平衡的因果证据。这一机制性发现说明,权值共享实质上优化了嵌入矩阵用于输出预测,损害了其输入表征功能。这些结果解释了为何权值共享在大模型中可能降低性能,并对嵌入矩阵占比高的小型语言模型训练具有重要影响。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。