深入理解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
这个模型包含两个线性层:
- 第一个线性层将one-hot编码的输入词映射到低维嵌入空间
- 第二个线性层将嵌入向量映射回词汇表大小的空间
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. 扩展思考
- 负采样优化:实际应用中常使用负采样代替softmax,可以显著提高训练效率
- 层次softmax:另一种优化方法是使用二叉树结构计算概率
- 动态窗口大小:可以尝试使用变化的窗口大小来捕捉不同范围的上下文信息
- 子采样高频词:对高频词进行子采样可以改善低频词的学习效果
5. 总结
这个实现展示了Word2Vec-SkipGram模型的核心思想:通过神经网络学习词向量,使得语义相近的词在向量空间中距离相近。虽然实现相对简单,但包含了SkipGram模型的关键要素:
- 滑动窗口构建训练样本
- 简单的神经网络架构
- 基于中心词预测上下文词的目标
- 词向量的可视化验证
理解这个基础实现有助于深入掌握更复杂的词嵌入方法,如GloVe、FastText等。