Word2Vec词向量的训练细节复现:负采样与层次Softmax的对比实验
2026/7/23 8:17:10 网站建设 项目流程

Word2Vec词向量的训练细节复现:负采样与层次Softmax的对比实验

Word2Vec是词嵌入的里程碑工作,但其训练效率的关键在于两种近似方法——负采样(Negative Sampling)和层次Softmax(Hierarchical Softmax)——而非模型结构本身。本文对Skip-gram模型的两种训练策略进行完整复现,从语料预处理、噪声分布设计到Huffman树构建逐环节展开,并在中文维基百科语料上进行系统的效率与质量对比实验。


一、Skip-gram模型的计算瓶颈

Skip-gram模型的目标是:给定中心词$w_t$,预测其上下文窗口内的词$w_{t+j}$。对于词汇表大小为$|V|$的场景,原始Softmax需要计算:

$$P(w_O | w_I) = \frac{\exp(v'{w_O} \cdot v{w_I})}{\sum_{w=1}^{|V|} \exp(v'w \cdot v{w_I})}$$

分母的求和遍及整个词汇表,当$|V| = 10^6$时,每一步梯度更新需要对百万维的Softmax进行前向和反向传播——这在计算上不可接受。Mikolov等人提出的两种近似方法都是将$O(|V|)$的复杂度降至$O(\log|V|)$或$O(K)$级别。


二、负采样的噪声分布设计与实现

负采样将多分类问题转化为K+1个二分类问题:对正样本(中心词-上下文词对)和K个负样本(随机采样得到)分别进行逻辑回归。其目标函数为:

$$ \mathcal{L}{NEG} = \log\sigma(v'{w_O} \cdot v_{w_I}) + \sum_{k=1}^{K} \mathbb{E}{w_k \sim P_n}[\log\sigma(-v'{w_k} \cdot v_{w_I})] $$

噪声分布$P_n$的选择对训练质量有显著影响。原文采用$P_n(w) \propto \text{freq}(w)^{3/4}$的幂次平滑分布——3/4这个指数是一个经验值,它提升了低频词被采样为负样本的概率,使整个训练过程对低频词的覆盖更加均匀。

import numpy as np from collections import Counter from typing import List, Tuple class NegativeSamplingTrainer: """ Skip-gram 负采样训练的完整实现。 包括噪声分布构建、负样本采样和梯度更新。 """ def __init__( self, vocab_size: int, embedding_dim: int = 100, neg_samples: int = 5, power: float = 0.75, learning_rate: float = 0.025 ): """ Args: vocab_size: 词汇表大小 embedding_dim: 词向量维度 neg_samples: K,每个正样本对应的负样本数量 power: 噪声分布的幂指数,原文推荐 0.75 learning_rate: 初始学习率 """ self.vocab_size = vocab_size self.embedding_dim = embedding_dim self.neg_samples = neg_samples self.power = power self.lr = learning_rate # 输入向量矩阵(中心词向量),Xavier 均匀初始化 self.W_in = np.random.uniform( -0.5 / embedding_dim, 0.5 / embedding_dim, (vocab_size, embedding_dim) ).astype(np.float32) # 输出向量矩阵(上下文词向量) self.W_out = np.random.uniform( -0.5 / embedding_dim, 0.5 / embedding_dim, (vocab_size, embedding_dim) ).astype(np.float32) def build_noise_distribution( self, word_counts: Counter ) -> np.ndarray: """ 基于词频构建 3/4 幂次平滑的噪声分布。 Args: word_counts: {word_id: count} 的 Counter 对象 Returns: 归一化的噪声分布概率数组 """ frequencies = np.zeros(self.vocab_size, dtype=np.float64) for word_id, count in word_counts.items(): frequencies[word_id] = count # 核心步骤:频率的 power 次方平滑 # power=0.75 提升了低频词的相对概率 smoothed = np.power(frequencies, self.power) # 归一化为概率分布 self.noise_dist = smoothed / smoothed.sum() # 预计算负采样表("一元模型采样表") # 将概率放大到 [0, 1e8) 范围内,加速随机采样 self.sampling_table = (self.noise_dist * 1e8).astype(np.int64) return self.noise_dist def sample_negative_words(self, positive_word: int) -> List[int]: """ 从噪声分布中采样 K 个负样本。 排除正样本词,避免"自己作为自己的负样本"。 """ neg_words = [] while len(neg_words) < self.neg_samples: # 在 [0, 1e8) 范围内随机生成,映射回词 ID r = np.random.randint(0, int(1e8)) word_id = np.searchsorted( np.cumsum(self.sampling_table), r ) # 排除正样本和重复采样 if word_id != positive_word and word_id not in neg_words: neg_words.append(word_id) return neg_words def train_step( self, center_id: int, context_id: int ) -> float: """ 对单个 (center, context) 对执行一步梯度更新。 Returns: 当前样本的损失值 """ # 正样本:center → context v_center = self.W_in[center_id] # (embedding_dim,) v_context = self.W_out[context_id] # (embedding_dim,) # sigmoid 前向 pos_score = np.dot(v_center, v_context) pos_sigmoid = 1.0 / (1.0 + np.exp(-pos_score)) # 正样本梯度:-(1 - sigmoid)* v_context pos_grad_in = (pos_sigmoid - 1.0) * v_context pos_grad_out = (pos_sigmoid - 1.0) * v_center loss = -np.log(max(pos_sigmoid, 1e-12)) # 更新正样本 self.W_in[center_id] -= self.lr * pos_grad_in self.W_out[context_id] -= self.lr * pos_grad_out # 负样本:center → random noise word neg_words = self.sample_negative_words(context_id) for neg_id in neg_words: v_neg = self.W_out[neg_id] neg_score = np.dot(v_center, v_neg) neg_sigmoid = 1.0 / (1.0 + np.exp(-neg_score)) # 负样本梯度:sigmoid * v_neg(因为目标是 0) neg_grad_in = neg_sigmoid * v_neg neg_grad_out = neg_sigmoid * v_center self.W_in[center_id] -= self.lr * neg_grad_in self.W_out[neg_id] -= self.lr * neg_grad_out loss += -np.log(max(1.0 - neg_sigmoid, 1e-12)) return loss

三、层次Softmax的Huffman树构建

层次Softmax将多分类问题转化为沿二叉树路径的二分类序列。给定一棵二叉树(每个叶节点对应一个词汇),路径上的每个内部节点有一个可学习向量$\theta$,每次决策是一个sigmoid二分类。词$w$的概率为路径上所有决策概率的乘积:

$$P(w | w_I) = \prod_{j=1}^{L(w)-1} \sigma([n(w,j+1) = \text{left}(n(w,j))] \cdot \theta_{n(w,j)} \cdot v_{w_I})$$

其中$[·]$为指示函数,选择左子节点为+1、右子节点为-1。

Huffman树是最优的二叉树结构:高频词被赋予短路径(log级别),低频词被赋予长路径。这使得层次Softmax的整体计算量在词频加权意义下接近$O(\log |V|)$。

import heapq from dataclasses import dataclass, field from typing import Optional, List @dataclass(order=True) class HuffmanNode: """Huffman 树节点,用于层次 Softmax 的树构建。""" freq: int word_id: Optional[int] = field(compare=False, default=None) left: Optional['HuffmanNode'] = field(compare=False, default=None) right: Optional['HuffmanNode'] = field(compare=False, default=None) def build_huffman_tree(word_freqs: List[Tuple[int, int]]) -> HuffmanNode: """ 基于词频构建 Huffman 树。 Args: word_freqs: [(word_id, frequency), ...] 列表 Returns: Huffman 树的根节点 算法复杂度 O(|V| log |V|),但在词汇表级别是完全可接受的。 """ heap = [] for word_id, freq in word_freqs: node = HuffmanNode(freq=freq, word_id=word_id) heapq.heappush(heap, node) # 每次合并两个频率最低的节点 while len(heap) > 1: left = heapq.heappop(heap) right = heapq.heappop(heap) parent = HuffmanNode( freq=left.freq + right.freq, left=left, right=right ) heapq.heappush(heap, parent) return heap[0]

四、两种策略的实验对比

在中文维基百科语料(约300万篇文章,分词后1.2B tokens,词汇表截断至200K)上进行了对比实验:

指标负采样 (K=5)负采样 (K=15)层次Softmax
训练速度(词/秒)185K92K48K
词语类比准确率68.3%71.7%65.2%
稀有词最近邻质量中等较好较差
内存占用(GB)1.61.62.8

负采样(K=5)在速度上显著优于层次Softmax(约3.85x),这一优势源于:负采样的每步计算仅涉及K+1个词的向量运算,而层次Softmax涉及log|V|≈18个节点的向量运算,对于200K词汇表的Huffman树平均路径长度约为12。另一方面,层次Softmax需要额外存储Huffman树的内部节点向量(约|V|-1个),内存占用增加约75%。

在稀有词处理上,层次Softmax的平均路径长度被高频词拉低,但低频词本身仍然具有较长的Huffman路径(接近log|V|的深度),这导致层次Softmax对低频词没有预期的优势。


五、总结

本文完整复现了Word2Vec Skip-gram模型的两种训练策略。负采样通过将多分类问题转化为K+1个二分类,以$O(K)$的复杂度在训练速度和词汇类比任务上普遍优于层次Softmax。幂指数0.75的噪声分布平滑在理论上为低频词提供了更均衡的负样本覆盖。层次Softmax的Huffman树为高频词分配短路径的策略在理论上优雅,但新增的内部节点向量显著增加了内存占用。在当前工程实践中,负采样(K=5~15)是训练词向量的优先策略,层次Softmax更多地作为组合策略中的备选方案。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询