首页
/ 深入理解Word2Vec-SkipGram模型实现(基于graykode/nlp-tutorial项目)

深入理解Word2Vec-SkipGram模型实现(基于graykode/nlp-tutorial项目)

2025-07-06 01:55:35作者:温玫谨Lighthearted

1. Word2Vec与SkipGram模型简介

Word2Vec是自然语言处理中经典的词向量学习算法,它通过神经网络模型将单词映射到低维连续向量空间。SkipGram是Word2Vec的两种主要架构之一(另一种是CBOW),它的核心思想是通过中心词预测上下文词,从而学习到有意义的词向量表示。

SkipGram模型的特点:

  • 输入是中心词(target word)
  • 输出是上下文词(context words)
  • 采用滑动窗口机制构建训练样本
  • 适合处理低频词

2. 代码实现解析

2.1 数据准备

首先定义了一些简单的句子作为训练数据:

sentences = ["apple banana fruit", "banana orange fruit", "orange banana fruit",
             "dog cat animal", "cat monkey animal", "monkey dog animal"]

这些句子经过处理后构建词汇表,并为每个单词分配唯一索引:

word_sequence = " ".join(sentences).split()
word_list = list(set(word_sequence))
word_dict = {w: i for i, w in enumerate(word_list)}
voc_size = len(word_list)

2.2 SkipGram样本构建

SkipGram模型需要构建(中心词,上下文词)的训练对。这里采用窗口大小为1的方式:

skip_grams = []
for i in range(1, len(word_sequence) - 1):
    target = word_dict[word_sequence[i]]
    context = [word_dict[word_sequence[i - 1]], word_dict[word_sequence[i + 1]]]
    for w in context:
        skip_grams.append([target, w])

2.3 模型架构

实现了一个简单的神经网络模型:

class Word2Vec(nn.Module):
    def __init__(self):
        super(Word2Vec, self).__init__()
        self.W = nn.Linear(voc_size, embedding_size, bias=False)
        self.WT = nn.Linear(embedding_size, voc_size, bias=False)

    def forward(self, X):
        hidden_layer = self.W(X)
        output_layer = self.WT(hidden_layer)
        return output_layer

这个模型包含两个线性层:

  1. 第一个线性层将one-hot编码的输入词映射到低维嵌入空间
  2. 第二个线性层将嵌入向量映射回词汇表大小的空间

2.4 训练过程

训练采用随机小批量梯度下降:

model = Word2Vec()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(5000):
    input_batch, target_batch = random_batch()
    # 前向传播、计算损失、反向传播
    ...

2.5 可视化词向量

训练完成后,将学习到的词向量可视化:

for i, label in enumerate(word_list):
    W, WT = model.parameters()
    x, y = W[0][i].item(), W[1][i].item()
    plt.scatter(x, y)
    plt.annotate(label, xy=(x, y), xytext=(5, 2), 
                 textcoords='offset points', ha='right', va='bottom')
plt.show()

3. 关键实现细节分析

3.1 随机批次采样

random_batch()函数实现了随机采样训练样本:

def random_batch():
    random_index = np.random.choice(range(len(skip_grams)), batch_size, replace=False)
    random_inputs = [np.eye(voc_size)[skip_grams[i][0]] for i in random_index]
    random_labels = [skip_grams[i][1] for i in random_index]
    return random_inputs, random_labels

这里将中心词转换为one-hot编码作为输入,上下文词索引作为标签。

3.2 损失函数选择

使用交叉熵损失函数:

criterion = nn.CrossEntropyLoss()

这是分类任务的典型选择,因为我们要预测的是上下文词在词汇表中的概率分布。

3.3 参数初始化

PyTorch的Linear层默认使用均匀初始化,这对于Word2Vec模型是合适的。

4. 扩展思考

  1. 负采样优化:实际应用中常使用负采样代替softmax,可以显著提高训练效率
  2. 层次softmax:另一种优化方法是使用二叉树结构计算概率
  3. 动态窗口大小:可以尝试使用变化的窗口大小来捕捉不同范围的上下文信息
  4. 子采样高频词:对高频词进行子采样可以改善低频词的学习效果

5. 总结

这个实现展示了Word2Vec-SkipGram模型的核心思想:通过神经网络学习词向量,使得语义相近的词在向量空间中距离相近。虽然实现相对简单,但包含了SkipGram模型的关键要素:

  • 滑动窗口构建训练样本
  • 简单的神经网络架构
  • 基于中心词预测上下文词的目标
  • 词向量的可视化验证

理解这个基础实现有助于深入掌握更复杂的词嵌入方法,如GloVe、FastText等。