word2vec的分层softmax的原理

在自然语言处理（NLP）领域，Word2Vec是实现词嵌入的经典模型。然而，当词汇量 $V$ 巨大时（例如数十万或数百万），模型训练阶段的输出层——标准的Softmax函数——会成为严重的性能瓶颈。这是因为Softmax的归一化计算需要遍历整个词汇表，导致时间复杂度高达 $O(V)$。

分层Softmax（Hierarchical Softmax, HS）是解决这一问题的核心优化技术，它将计算复杂度从 $O(V)$ 降低到了 $O(\log V)$，极大地提高了训练速度。

一、 标准Softmax的局限性

标准的Softmax函数计算目标词 $w_o$ 的概率 $P(w_o|w_I)$ 如下：

$$P(w_o|w_I) = \frac{\exp(v_{w_o}^T v_{w_I})}{\sum_{w=1}^{V} \exp(v_w^T v_{w_I})}$$

分母的求和项需要计算所有 $V$ 个词汇的内积，导致每一步梯度更新都依赖于 $V$ 的大小，这在实践中是不可接受的。

二、 分层Softmax的核心原理：霍夫曼树

分层Softmax用一棵二叉树（通常是霍夫曼树，Huffman Tree）来代替传统的扁平Softmax层。这棵树满足以下特点：

1. 叶子节点即词汇： 树中的每个叶子节点代表词汇表中的一个词 $w$。
2. 内部节点即分类器： 树中的每个内部节点都代表一个二元分类器，用于决定向左走还是向右走。
3. 计算路径： 词 $w$ 的概率计算不再是遍历 $V$ 个词，而是计算从根节点到该词叶子节点的路径上所有二元分类的概率乘积。

霍夫曼树的选择依据

使用霍夫曼树作为结构是关键。霍夫曼编码基于词频构建，高频词汇距离根节点更近，因此它们的路径更短，平均路径长度得到优化，从而使得平均计算复杂度趋近于 $O(\log V)$。

三、 概率计算和梯度更新

对于目标词 $w$，假设从根节点到 $w$ 的路径上有 $L(w)$ 个内部节点 $n_1, n_2, \dots, n_{L(w)-1}$。在每个内部节点 $n_j$ 处，我们使用Sigmoid函数进行二元分类判断走向。

每个内部节点 $n_j$ 都有一个与之相关的向量 $v’_{n_j}$（类似于Softmax中的权重矩阵行）。

目标词 $w$ 的概率 $P(w|w_I)$ 被定义为路径上所有二元决策概率的乘积：

$$P(w|w_I) = \prod_{j=1}^{L(w)-1} P(\text{decision}_j | w_I)$$

其中，第 $j$ 个决策 $P(\text{decision}_j | w_I)$ 使用Sigmoid函数表示：

$$P(\text{decision}j | w_I) = \sigma( [2 \cdot \text{target}_j – 1] \cdot v{w_I}^T v’_{n_j} )$$

这里：
* $v_{w_I}$ 是输入词的向量。
* $v’_{n_j}$ 是第 $j$ 个内部节点的权重向量。
* $\text{target}_j$ 是指示目标路径方向的标签（通常左边为 1，右边为 0）。
* $[2 \cdot \text{target}_j – 1]$ 将标签转换为 $-1$ 或 $1$，确保Sigmoid函数计算的是“正确路径”的概率。

四、 实践示例：模拟分层Softmax计算

在实际的Word2Vec实现（如Gensim）中，HS通过预先构建霍夫曼树并存储路径和目标标签来实现。我们通过一个简化示例来演示这个计算过程。

假设我们有一个隐藏层向量 $v_{w_I}$ 和一个词 ‘cat’ 的路径，该路径经过两个内部节点 $n_1$ 和 $n_2$。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
import numpy as np

from scipy.special import expit # expit is the sigmoid function: 1 / (1 + exp(-x))



# 1. 模拟输入词的隐藏层向量 v_I

hidden_vector = np.array([0.8, -0.4, 0.1]) 



# 2. 模拟内部节点的权重向量 (这些向量是训练中要更新的参数)

n1_weight = np.array([0.1, 0.3, -0.2]) # Node 1 vector

n2_weight = np.array([-0.5, 0.2, 0.0]) # Node 2 vector



# 3. 目标路径的决策 (Target path to 'cat'):

# Node 1: Target is 1 (e.g., Left)

# Node 2: Target is 0 (e.g., Right)

path_targets = [1, 0]

path_nodes = [n1_weight, n2_weight]



# 4. 计算 P('cat' | v_I)



probability = 1.0



for j, n_weight in enumerate(path_nodes):

    target_j = path_targets[j]



    # 核心计算：v_I^T * v'_n

    score = np.dot(hidden_vector, n_weight)



    # 将 target 转换为 -1 或 1

    sign = 2 * target_j - 1 



    # 计算正确路径的概率 P(d_j | w_I) = sigma( sign * score )

    P_correct_turn = expit(sign * score)



    probability *= P_correct_turn



    print(f"--- Node {j+1} ---")

    print(f"Dot Score: {score:.4f}")

    print(f"P(Correct Turn | w_I): {P_correct_turn:.4f}")



print(f"\nFinal Probability P('cat' | w_I) = {probability:.6f}")

在训练过程中，梯度更新只发生在路径上的 $L(w)-1$ 个内部节点向量 $v’_{n_j}$ 上，而不是 $V$ 个输出词向量上，这正是 HS 带来巨大速度提升的关键所在。

五、 总结

分层Softmax通过引入基于霍夫曼编码的二叉树结构，巧妙地将大规模的多分类问题转化为一系列高效的二元分类问题。这种结构将计算复杂度从线性 $O(V)$ 优化到对数 $O(\log V)$，使得在拥有数百万词汇表的语料库上训练Word2Vec成为可能，是AI模型基础设施中常用的优化手段之一。