一行代码的魔法

在PyTorch中实现一个简单的语言模型时，你可能会看到这样一行看似普通的代码：

self.lm_head.weight = self.embedding.weight

这行代码做的事情非常简单：让输出层的权重指针直接指向嵌入层的权重。但就是这样一行代码，却能够为一个典型的70亿参数大模型节省约2亿个参数——相当于模型总参数量的3%左右。

这个技术叫做权重共享（Weight Tying），有时也称为权重绑定。它几乎被所有现代大语言模型采用：GPT系列、LLaMA系列、BERT等。然而，尽管这个技术如此普遍，很多人可能并没有深入思考过：为什么这样做是合理的？输入嵌入和输出投影明明承担着不同的功能，它们的权重为什么能够共享？

这篇文章将带你深入理解权重共享背后的技术原理，从直觉到数学推导，从历史发展到实践考量。

两个角色的不同使命

要理解权重共享，我们首先需要搞清楚Transformer中两个关键组件的角色。

输入嵌入层：从符号到语义

输入嵌入层（Input Embedding Layer）的工作是将离散的token ID转换成稠密向量。假设词汇表大小为$|V|$，嵌入维度为$d_{model}$，嵌入层的权重矩阵$W_{emb}$的形状就是$(|V|, d_{model})$。

flowchart LR
    A[Token ID: 42] --> B[Embedding Matrix]
    B --> C[Vector: 0.1, -0.3, 0.7, ...]
    
    subgraph B[Embedding Matrix W_emb]
        D[Row 0: word 'the']
        E[Row 1: word 'a']
        F[...]
        G[Row 42: word 'learning']
        H[Row |V|-1: ...]
    end
    
    style G fill:#f9f,stroke:#333,stroke-width:2px

当输入一个token ID时，嵌入层做的事情就是从矩阵中"抽取"对应的行向量。这个过程可以理解为：将人类可理解的符号（单词）翻译成模型可处理的数学语言（向量）。

输出语言模型头：从语义到预测

输出层（Language Model Head）的任务恰好相反。经过多层Transformer处理后，模型产生一个$d_{model}$维的隐藏状态向量，输出层需要将这个向量转换成对词汇表中每个词的预测分数。

输出层的权重矩阵$W_{out}$形状为$(d_{model}, |V|)$。计算过程是：

$$\text{logits} = h \cdot W_{out}$$

其中$h$是最终的隐藏状态向量，$\text{logits}$是每个词的预测分数。

flowchart LR
    A[Hidden State h] --> B[Linear Projection]
    B --> C[Logits for each token]
    
    subgraph C[Logits]
        D['the': 0.2]
        E['a': -0.5]
        F['learning': 1.3]
        G[...]
    end

这个过程可以理解为：模型在问自己，当前的隐藏状态最"接近"哪个词？

直觉：两个互逆的操作

现在，让我们仔细思考这两个操作：

  • 嵌入层：单词 → 语义向量（编码）
  • 输出层：语义向量 → 单词得分（解码）

这不就是一对互逆的操作吗？

更具体地说：

  • 嵌入层的每一行代表一个词的语义向量
  • 输出层的每一列代表"如何判断隐藏状态是否接近这个词"

如果模型学会了"学习"这个词的语义表示，那么：

  • 作为输入时，这个表示告诉我们"学习"意味着什么
  • 作为输出时，这个表示告诉我们什么样的隐藏状态应该预测为"学习"

这就是权重共享的核心直觉：同一个词的语义表示应该是一致的，无论它作为输入还是输出。

数学原理：为什么共享是合理的

几何视角

让我们从几何角度来理解。假设嵌入矩阵$W_{emb}$的第$i$行$w_i$表示第$i$个词的语义向量。

当我们在输出层计算logits时：

$$\text{logits}_i = h \cdot w_i^T = \langle h, w_i \rangle$$

这实际上是在计算隐藏状态$h$与每个词向量的内积。内积越大，说明$h$与该词向量越"接近"。

graph TB
    subgraph "向量空间视角"
        A[隐藏状态 h]
        B[词向量 w_1: 'the']
        C[词向量 w_2: 'cat']
        D[词向量 w_3: 'learning']
        
        A -->|"内积小"| B
        A -->|"内积中"| C
        A -->|"内积大"| D
    end
    
    style D fill:#9f9,stroke:#333

如果输出层使用独立的权重矩阵$W_{out}$（不共享），那么输出层实际上是在学习另一套词向量。这就带来了一个问题：为什么同一个词需要两套不同的向量表示？

梯度更新的视角

让我们看看训练过程中梯度的流动。对于传统的非共享模型：

嵌入层的梯度更新：

$$\frac{\partial L}{\partial W_{emb}[k]} = \text{来自输入端的梯度}$$

只有当第$k$个词出现在输入序列中时，它的嵌入向量才会被更新。

输出层的梯度更新：

$$\frac{\partial L}{\partial W_{out}[k]} = \text{来自输出端的梯度}$$

只有当第$k$个词作为目标词时，它的输出向量才会被更新。

这导致了一个问题：低频词的嵌入向量更新不足。一个词可能作为目标出现很多次，但作为输入出现很少，或者反过来。

共享权重的优势

当使用权重共享时，同一个词向量同时接收来自两端的梯度：

$$\frac{\partial L}{\partial W_{shared}[k]} = \frac{\partial L_{input}}{\partial W[k]} + \frac{\partial L_{output}}{\partial W[k]}$$

这意味着：

  1. 更高效的梯度利用：每个词向量都能从输入和输出两个方向学习
  2. 更一致的语义表示：强制模型为每个词学习统一的表示
  3. 隐式的正则化效果：限制模型的自由度，防止过拟合

参数节省：数字说话

让我们用具体数字来说明权重共享带来的参数节省。

假设一个典型的语言模型配置：

  • 词汇表大小 $|V| = 50,257$（类似GPT-2）
  • 嵌入维度 $d_{model} = 4096$（类似LLaMA-7B）

不使用权重共享

嵌入层参数量：

$$|V| \times d_{model} = 50,257 \times 4096 \approx 206 \text{M}$$

输出层参数量：

$$d_{model} \times |V| = 4096 \times 50,257 \approx 206 \text{M}$$

总计：约4.12亿参数

使用权重共享

只需要一份嵌入矩阵：

总计：约2.06亿参数

节省比例

对于一个70亿参数的模型：

  • 节省参数量：约2.06亿
  • 占总参数比例：$\frac{206}{7000} \approx 2.9\%$
pie title "70亿参数模型的参数分布（示意）"
    "Transformer主体" : 6500
    "Embedding+Output（共享后）" : 206
    "节省的参数" : 206

理论支撑：两篇关键论文

权重共享技术并非凭空产生，它有坚实的理论基础支撑。

Press & Wolf (2016/2017)：经验验证

Ofir Press和Lior Wolf在2016年发表了论文《Using the Output Embedding to Improve Language Models》（发表于ICLR 2017）。这篇论文的核心贡献是：

  1. 实证验证：在多个数据集上验证了权重共享的有效性
  2. 实验结论：共享权重不仅不会损害性能，反而能够提升模型表现
  3. 理论启示：输出层的权重矩阵本身就是一个有效的词嵌入

论文中的关键发现是，在Penn Treebank数据集上，使用权重共享的模型显著优于不共享的模型，同时模型大小减少了一半以上。

Inan et al. (2017)：理论框架

Hakan Inan、Khashayar Khosravi和Richard Socher在2017年发表了论文《Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling》（同样发表于ICLR 2017）。这篇论文提供了更深层的理论解释：

核心创新：增强损失框架

传统的交叉熵损失：

$$L_{CE} = D_{KL}(y^* || \hat{y})$$

其中$y^*$是one-hot目标分布，$\hat{y}$是模型预测分布。

论文提出的增强损失：

$$L_{total} = L_{CE} + \alpha L_{aug}$$

其中增强项$L_{aug}$利用词嵌入空间的度量信息：

$$\tilde{y}_i = \text{softmax}\left(\frac{w_i^T w_{target}}{\tau}\right)$$

这个公式利用了词向量之间的相似性：语义相近的词应该有相似的预测概率。

理论推导的关键结论：

在特定条件下（高温极限），增强损失的最优解自然导出：

$$W_{out} = W_{emb}^T$$

这意味着，权重共享不是人为设计的技巧，而是某种理论框架下的最优解。

flowchart TD
    A[传统交叉熵损失] --> B[增强损失框架]
    B --> C[利用词向量相似性]
    C --> D[理论分析]
    D --> E[最优解: W_out = W_emb^T]
    E --> F[权重共享]
    
    style F fill:#9f9,stroke:#333,stroke-width:2px

各大模型的实践

GPT系列

GPT系列模型（GPT-1、GPT-2、GPT-3、GPT-4）都采用了权重共享。在Hugging Face的transformers库中，GPT-2的配置明确说明：

The GPT2 Model transformer with a language modeling head on top (linear layer with weights tied to the input embeddings).

实现方式非常直接：

class GPT2LMHeadModel(GPT2PreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.transformer = GPT2Model(config)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        
        # Model parallel
        self.model_parallel = False
        self.device_map = None
        self.gradient_checkpointing = False

        # Initialize weights and apply final processing
        self.post_init()
    
    def tie_weights(self):
        # 关键：共享权重
        self.lm_head.weight = self.transformer.wte.weight

BERT系列

BERT作为encoder-only模型，主要用于掩码语言建模（Masked Language Modeling）。尽管任务不同，BERT同样采用了权重共享。

但需要注意，BERT的输出层是预测被掩盖的token，其本质与因果语言模型相同，因此权重共享同样适用。

LLaMA系列

LLaMA系列（包括LLaMA-1、LLaMA-2、LLaMA-3）作为目前最流行的开源大模型之一，同样使用了权重共享。

在Hugging Face的实现中，LLaMA的LlamaForCausalLM类中：

class LlamaForCausalLM(LlamaPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.model = LlamaModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def tie_weights(self):
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight

可以看到，LLaMA通过tie_word_embeddings配置项控制是否共享权重。

T5模型：三层共享

T5（Text-to-Text Transfer Transformer）作为encoder-decoder架构的代表，采用了更激进的三层权重共享：

  1. 编码器的输入嵌入
  2. 解码器的输入嵌入
  3. 解码器的输出投影
graph TB
    subgraph "T5的三层权重共享"
        A[Encoder输入嵌入] 
        B[Decoder输入嵌入]
        C[Decoder输出投影]
        D[共享权重矩阵]
    end
    
    A --> D
    B --> D
    C --> D
    
    style D fill:#9f9,stroke:#333,stroke-width:2px

这种设计基于一个观察：在机器翻译等任务中，源语言和目标语言共享大量子词（特别是使用BPE分词时），因此可以用同一个嵌入空间表示两种语言的词汇。

实现细节与注意事项

PyTorch实现

在PyTorch中实现权重共享非常简单：

import torch
import torch.nn as nn

class SimpleLanguageModel(nn.Module):
    def __init__(self, vocab_size, hidden_size, tie_weights=True):
        super().__init__()
        
        # 输入嵌入层
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        
        # 模型主体（简化为一个线性层）
        self.body = nn.Linear(hidden_size, hidden_size)
        
        # 输出语言模型头
        # 注意：bias=False，共享权重时通常不使用偏置
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)
        
        # 权重共享
        if tie_weights:
            self.lm_head.weight = self.embedding.weight
    
    def forward(self, input_ids):
        # 嵌入
        x = self.embedding(input_ids)
        # 模型处理
        x = self.body(x)
        # 输出logits
        logits = self.lm_head(x)
        return logits

关键注意事项

  1. 偏置项：共享权重时，输出层通常不使用偏置项（bias=False）。这是因为嵌入层没有对应的偏置概念。

  2. 权重缩放：原始Transformer论文中，嵌入权重会乘以$\sqrt{d_{model}}$：

# 嵌入层输出后
x = self.embedding(input_ids) * math.sqrt(self.d_model)

这个缩放的目的是使嵌入向量的尺度与位置编码相当，防止位置编码主导信号。

  1. 初始化：共享权重后，只需要初始化嵌入层即可，输出层会自动使用相同的初始化。

  2. 保存和加载：保存模型时，共享的权重只会保存一份，这进一步减少了存储空间。

Hugging Face Transformers中的处理

在Hugging Face的transformers库中，权重共享通过_tie_or_clone_weights函数实现：

def _tie_or_clone_weights(self, output_embeddings, input_embeddings):
    """Tie or clone module weights depending of whether we are using TorchScript or not"""
    if self.config.torchscript:
        output_embeddings.weight = nn.Parameter(input_embeddings.weight.clone())
    else:
        # 直接共享权重
        output_embeddings.weight = input_embeddings.weight

    if hasattr(output_embeddings, "out_features") and hasattr(input_embeddings, "num_embeddings"):
        output_embeddings.out_features = input_embeddings.num_embeddings

何时不使用权重共享

虽然权重共享被广泛采用，但在某些情况下，不共享权重可能更合适。

1. 表达能力的损失

权重共享会减少模型的自由度。对于非常大的模型和非常复杂的数据分布，独立的两套嵌入可能提供更好的表达能力。

有研究发现，在某些任务上，非共享权重的模型可能表现略好，特别是：

  • 模型已经足够大，参数节省的边际效益递减
  • 训练数据非常丰富，不需要额外的正则化
  • 输入和输出的语义空间确实存在差异

2. 多语言/多模态场景

在某些特殊场景中，输入和输出的语义空间可能确实不同：

  • 跨语言翻译：源语言和目标语言的词汇可能有完全不同的语义结构
  • 多模态模型：输入可能是图像，输出是文本，显然不能共享
  • 特殊token处理：某些模型可能为输入和输出设计不同的特殊token集合

3. 微调时的考量

在进行迁移学习或微调时，有时会考虑是否解除权重共享：

# 微调前解除权重共享
model.lm_head.weight = nn.Parameter(model.embedding.weight.clone())

这样做可以让模型更灵活地适应新任务，但也会增加过拟合的风险。

4. 实验发现：输出嵌入更优

Press & Wolf (2016) 的论文中有一个有趣的发现：如果单独使用输出嵌入（而非输入嵌入）作为词向量表示，效果往往更好。

这说明在语言模型中，输出端的表示可能比输入端学习到更丰富的语义信息。权重共享后，共享的嵌入会受到两端梯度的影响，但实验表明它更偏向于输出表示。

权重共享与模型规模的关系

一个有趣的问题是：随着模型规模的增长，权重共享的重要性如何变化？

小模型：节省比例更大

对于参数量较小的模型，权重共享节省的参数比例更大：

模型规模 词汇表大小 隐藏维度 节省参数 节省比例
125M 50,257 768 38.6M 30.9%
1.5B 50,257 2048 103M 6.9%
7B 32,000 4096 131M 1.9%
70B 32,000 8192 262M 0.4%

大模型：正则化作用更重要

对于超大模型，参数节省的边际效益递减，但权重共享的正则化作用仍然重要：

  1. 防止过拟合：限制参数空间，提高泛化能力
  2. 训练稳定性：统一的语义表示有助于梯度稳定
  3. 工程便利：减少显存占用，加快加载速度

词汇表大小的影响

2024年的一项研究《Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies》指出，更大的模型应该使用更大的词汇表。这意味着：

$$V_{optimal} \propto N^{\alpha}$$

其中$N$是模型参数量，$\alpha$是一个待确定的指数。

如果词汇表随模型规模增长，那么权重共享节省的参数量也会相应增长，其重要性可能不会随模型规模增长而单调递减。

权重共享的历史脉络

早期神经网络语言模型

权重共享的思想可以追溯到早期的神经网络语言模型。Bengio等人在2001年的经典论文《A Neural Probabilistic Language Model》中，就已经探索了输入和输出表示之间的关系。

Word2Vec的启示

Word2Vec（Mikolov et al., 2013）实际上隐含了类似的思想。在skip-gram模型中，每个词有两个向量：输入向量和输出向量。尽管它们不共享，但实验表明它们的性质相似——都能捕捉语义关系。

AWD-LSTM

在AWD-LSTM（Merity et al., 2017）中，权重共享是一个关键技巧。论文的实验表明，使用权重共享可以在不损失性能的情况下将模型大小减少一半以上。

Transformer的确立

Transformer论文《Attention Is All You Need》（Vaswani et al., 2017）在3.4节明确提到：

We share the same weight matrix between the two embedding layers and the pre-softmax linear transformation.

这里"two embedding layers"指的是编码器和解码器的嵌入层，“pre-softmax linear transformation"就是输出投影层。原始Transformer实际上实现了三层共享。

总结

权重共享（Weight Tying）是一个看似简单却蕴含深意的技术。它通过一行代码：

self.lm_head.weight = self.embedding.weight

实现了：

  1. 参数节省：减少约vocab_size × hidden_dim的参数量
  2. 正则化效果：强制输入和输出使用统一的语义表示
  3. 训练效率：梯度从两端同时更新嵌入向量
  4. 理论支撑：在特定理论框架下是最优解

这个技术的普及告诉我们：好的设计往往是简洁的。当我们深入理解了模型的内部运作机制——输入嵌入和输出投影实际上是在执行互逆的操作——权重共享就是一个自然的选择。

当然，权重共享并非万能。在某些特殊场景下，独立的输入和输出嵌入可能提供更好的表达能力。但对于绝大多数语言模型应用，权重共享是一个经过时间检验的优秀设计。

下次当你看到模型配置中tie_word_embeddings=True时，你知道这不仅仅是一个节省内存的技巧——它背后是多年的研究积累和深刻的理论洞察。

参考文献

  1. Press, O., & Wolf, L. (2017). Using the Output Embedding to Improve Language Models. ICLR 2017. arXiv:1608.05859

  2. Inan, H., Khosravi, K., & Socher, R. (2017). Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling. ICLR 2017. arXiv:1611.01462

  3. Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762

  4. Merity, S., Keskar, N. S., & Socher, R. (2017). Regularizing and Optimizing LSTM Language Models. ICLR 2018. arXiv:1708.02182

  5. Radford, A., et al. (2019). Language Models are Unsupervised Multitask Learners. OpenAI.

  6. Touvron, H., et al. (2023). LLaMA: Open and Efficient Foundation Language Models. arXiv:2302.13971

  7. Raffel, C., et al. (2020). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. JMLR, 21(140), 1-67.

  8. Bengio, Y., Ducharme, R., & Vincent, P. (2001). A Neural Probabilistic Language Model. NeurIPS.

  9. Mikolov, T., et al. (2013). Distributed Representations of Words and Phrases and their Compositionality. NeurIPS.

  10. Hugging Face Transformers Documentation: GPT-2, LLaMA, T5.

  11. Tao, J., & Geng, X. (2024). Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies. NeurIPS 2024. arXiv:2407.13623

  12. Stack Overflow: Why are weight matrices shared between embedding layers in “Attention is All You Need”?

  13. Zaremba, W., Sutskever, I., & Vinyals, O. (2014). Recurrent Neural Network Regularization. arXiv:1409.2329