从词表到语义：Softmax 权重共享 (Weight Tying)

在阅读 Transformer 的源码实现（如 Hugging Face 的 BERT/GPT）时，一个常被忽视但至关重要的细节是：Embedding 层的权重矩阵 $\mathbf{E}$ 与输出层（LM Head）的权重矩阵 $\mathbf{W}_{\text{out}}$ 是共享的。

这意味着：

$$\mathbf{W}_{\text{out}} = \mathbf{E}^\top$$

在 PyTorch 实现中，它们通常指向同一个 nn.Parameter 对象。

这一设计不仅仅是为了节省参数量，其背后蕴含着深刻的几何直觉和正则化（Regularization）原理。本文将探讨权重共享在语言建模中的理论基础及其必要性。

权重共享的定义

Transformer 模型的输入端和输出端都涉及词表（Vocabulary）空间的映射：

• 输入端（Embedding）：将离散的 Token ID 映射为连续的词向量。
  • 映射矩阵 $\mathbf{E} \in \mathbb{R}^{V \times d}$，其中 $V$ 是词表大小，$d$ 是隐层维度。
• 输出端（LM Head）：将最后一层的隐状态 $\mathbf{h}$ 映射回词表概率分布（Logits）。
  • 映射矩阵 $\mathbf{W}_{\text{out}} \in \mathbb{R}^{d \times V}$。
  • 下一个词的概率分布：$P(\text{next\_token}) = \text{Softmax}(\mathbf{h} \mathbf{W}_{\text{out}} + \mathbf{b})$。

在不共享权重的情况下，这两个矩阵是独立的参数。权重共享（Weight Tying）强制 $\mathbf{W}_{\text{out}}$ 使用 $\mathbf{E}$ 的转置，从而将两个参数矩阵耦合在一起。

几何一致性：输入与输出空间的对齐

从几何角度来看，$\mathbf{E}$ 和 $\mathbf{W}_{\text{out}}$ 分别定义了词在输入空间和输出空间的表示。

• $\mathbf{E}$ 的行向量 $\mathbf{e}_i$：代表第 $i$ 个词在语义输入空间中的位置。
• $\mathbf{W}_{\text{out}}$ 的列向量 $\mathbf{w}_j$：代表第 $j$ 个词在输出空间中的“方向”。

当模型预测下一个词时，实际上是在计算隐状态 $\mathbf{h}$ 与所有输出向量 $\mathbf{w}_j$ 的内积（相似度）：

$$\text{Logits}_j = \mathbf{h} \cdot \mathbf{w}_j$$

内积越大，意味着 $\mathbf{h}$ 与 $\mathbf{w}_j$ 的方向越一致，预测该词的概率就越大。

若不共享权重：
输入端的“苹果”向量 $\mathbf{e}_{\text{apple}}$ 和输出端的“苹果”向量 $\mathbf{w}_{\text{apple}}$ 可能会学习到完全不同的几何表示。虽然模型理论上可以通过训练让它们趋于一致，但这显著增加了优化的难度和自由度。

若共享权重（$\mathbf{W}_{\text{out}} = \mathbf{E}^\top$）：
这就施加了一个强约束：如果隐状态 $\mathbf{h}$ 指向“苹果”的词向量方向，那么模型预测“苹果”的概率最大。
这保证了输入空间和输出空间的几何一致性（Geometric Consistency）。当模型预测下一个词应该是“水果”类时，生成的隐状态 $\mathbf{h}$ 会自然地靠近“苹果”、“香蕉”等词向量聚类的区域，而这些区域正是输入 Embedding 已经学习好的语义簇。

理论依据与正则化效应

除了几何直觉，权重共享还起到了强大的正则化（Regularization）作用，尤其是在处理长尾分布数据时。

缓解稀疏词训练难题（Rare Word Problem）

自然语言的词频分布遵循齐普夫定律（Zipf's Law），绝大多数词（长尾部分）在训练语料中出现的频率极低。

• 对于输入 Embedding $\mathbf{E}$：只有当词 $i$ 作为输入出现时，对应的行 $\mathbf{e}_i$ 才会获得梯度更新。
• 对于输出权重 $\mathbf{W}_{\text{out}}$：只有当词 $i$ 作为目标标签（Target）出现时，对应的列 $\mathbf{w}_i$ 才会获得显著的梯度更新。

如果不共享权重，低频词的 $\mathbf{e}_i$ 和 $\mathbf{w}_i$ 都面临训练样本不足的问题。
共享权重后，无论该词是出现在 Source 端还是 Target 端，梯度都会更新同一个参数。这相当于将低频词的有效训练样本数加倍，显著提升了低频词表示的质量。

困惑度（Perplexity）的改善

Press & Wolf (2017) 的研究表明，共享权重不仅能大幅减少参数量（对于词表巨大的模型，参数量可减少近 50%），还能显著降低困惑度（PPL）。
这表明，该约束不仅仅是一种工程上的“省参数”技巧，更是一种有效的归纳偏置（Inductive Bias），帮助模型学习到更鲁棒的语言表示。

适用场景与局限性

尽管权重共享是 BERT、GPT、T5、LLaMA 等主流模型的标准配置，但也存在例外情况。

跨语言与多模态模型

在跨语言模型（如 Encoder-Decoder 翻译模型）中，如果 Source 和 Target 语言差异较大，或者词表完全不同，共享权重可能不再适用。同理，在多模态模型中，文本 Embedding 与图像 Projection 层显然无法共享。

优化动力学的潜在冲突

部分研究指出，$\mathbf{E}$ 和 $\mathbf{W}_{\text{out}}$ 在优化动力学上可能存在潜在冲突：

• $\mathbf{E}$ 倾向于在空间中分散以区分不同语义。
• $\mathbf{W}_{\text{out}}$ 倾向于聚集以适应 Softmax 的分类边界。
  因此，XLNet 等部分模型曾尝试解除这种绑定，或引入额外的非线性投影层。然而，在当前的自回归大模型（LLM）范式下，共享权重的收益（参数效率与泛化性）远大于其潜在的动力学冲突，因此依然是绝对的主流选择。

总结

Softmax 权重共享体现了深度学习模型设计中“奥卡姆剃刀”原则的胜利。
通过一个简单的线性约束 $\mathbf{W}_{\text{out}} = \mathbf{E}^\top$，它同时实现了：

1. 参数效率：显著降低显存占用。
2. 几何对齐：统一输入与输出的语义空间。
3. 正则化：改善长尾词汇的表征学习。

参考文献