""" ============================================================================= 第 2 讲 · 代码 2：训练真正的 word2vec —— 负采样、高频词下采样、动态窗口、CBOW ============================================================================= 对应讲义：lectures/lecture02_word_vectors_advanced.md 第 4、5 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 代码 01 推导并验证了负采样的损失和梯度。这个文件把它用起来，并加上原始 word2vec 里的另外几个"工程技巧"，看看每一个都起什么作用： 1) 高频词下采样（subsampling）：以一定概率丢掉 the、a 这类高频词， 既省时间，又让"有信息量"的上下文占更大比重。 2) 动态窗口（dynamic window）：每个位置的窗口半径在 1 ~ m 之间随机取， 效果相当于离中心词越近的上下文权重越大。 3) CBOW（continuous bag of words）：反过来，用上下文词的平均向量预测中心词。 4) 计数与预测的统一（讲义第 5 节）：Levy & Goldberg（2014）证明 SGNS 的最优解满足 u_o · v_c ≈ PMI(c, o) - log K 我们在玩具语料上直接检验这个关系。 ----------------------------------------------------------------------------- 【整体流程】 ----------------------------------------------------------------------------- 每一轮（epoch）： 对每个句子做下采样 → 按动态窗口抽 (中心词, 上下文词) 对 → 打乱 → 每批 B 个样本，每个样本配 K 个负例 → sgns_batch 求梯度 → 只更新涉及的那些行 学习率从 lr 线性衰减到接近 0（和原始 word2vec 一样）。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 02_sgns_training.py 需要 numpy。运行时间约 40 秒~1 分钟（会把 8 种设置各训练 3 次做对比，包括第 1 讲的朴素 softmax 版本）。 """ import importlib import time import numpy as np from shared import (NegativeSampler, build_vocab, evaluate_vectors, generate_corpus, nearest_neighbors, noise_distribution, print_header, sentences_to_ids, sigmoid, skipgram_pairs, word_categories) ns = importlib.import_module("01_negative_sampling") # ============================================================================= # 第 1 部分：高频词下采样与动态窗口 # ============================================================================= def keep_probabilities(counts, t=1e-2): """ 下采样时每个词被**保留**的概率：P(keep w) = min(1, sqrt(t / f(w)))，f(w) 是 w 的词频（占比）。 【思路】（讲义第 4.2 节） word2vec 论文的公式是"丢弃概率 = 1 - sqrt(t / f(w))"。 - 词频 f(w) <= t 的词：sqrt(t/f) >= 1，永远保留； - 词频越高，保留概率越低。例如 f = 100·t 时只保留 10%。 为什么开平方？如果直接按 t/f 保留，所有高频词在下采样后的出现次数都会变成一样多， 把"高频"这个信息完全抹掉了；开平方只是"压一压"，保留了高低的顺序。 【参数】 counts : np.ndarray (V,)，每个词的出现次数 t : float，阈值。论文对几十亿词的语料推荐 1e-5 左右； 我们的玩具语料只有 171 个词，每个词的占比都很大，所以用 1e-2。 【返回】 np.ndarray (V,)，每个词的保留概率，在 (0, 1] 之间。 【边界】 - t 太大：几乎不丢任何词，下采样失效；t 太小：内容词也被大量丢弃，语料变得太稀。 - 原始 C 代码用的公式略有不同（(sqrt(f/t) + 1) · t/f），思路一样。 """ f = counts / counts.sum() return np.minimum(1.0, np.sqrt(t / f)) def make_skipgram_pairs(id_sentences, window, rng, dynamic_window=True, keep_prob=None): """ 为一轮训练抽取 (中心词, 上下文词) 对：先下采样，再按（动态）窗口取上下文。 【参数】 id_sentences : list[list[int]] window : int，最大窗口半径 m rng : np.random.Generator dynamic_window : bool，True 时每个中心词的窗口半径在 1..m 中均匀随机取 keep_prob : None 或 np.ndarray (V,)。None 表示不下采样 【返回】 np.ndarray[int]，形状 (N, 2)，每行是 (center, outside) 【边界】 - 下采样是在**抽窗口之前**做的：被丢掉的词直接从句子里消失， 所以它两边的词会"靠得更近"，窗口实际上覆盖了更远的词。这是原始 word2vec 的做法。 - 每一轮都重新随机下采样、重新随机窗口，所以每一轮看到的样本略有不同。 【NLP 联系】 动态窗口的效果：距离为 1 的上下文总会被取到，距离为 m 的只有 1/m 的概率被取到。 这和第 3 节 GloVe 用 1/距离 给共现加权是同一个想法：离得近的词关系更紧密。 """ pairs = [] for sent in id_sentences: if keep_prob is not None: sent = [w for w in sent if rng.random() < keep_prob[w]] for t, center in enumerate(sent): b = int(rng.integers(1, window + 1)) if dynamic_window else window for j in range(max(0, t - b), min(len(sent), t + b + 1)): if j != t: pairs.append((center, sent[j])) return np.array(pairs, dtype=int).reshape(-1, 2) # ============================================================================= # 第 2 部分：训练 skip-gram（负采样） # ============================================================================= def train_sgns(id_sentences, counts, dim=10, window=2, K=5, epochs=20, lr=0.025, batch_size=128, subsample_t=1e-2, dynamic_window=True, noise_power=0.75, seed=0, verbose=True): """ 训练 skip-gram 负采样（SGNS）词向量。 【思路】 初始化：中心词向量 V 用 ±0.5/d 的小随机数，上下文向量 U 全部为 0（原始 word2vec 的做法）。 为什么 U 可以是 0，而第 1 讲说"不能全设成 0"？因为 V 是随机的， 第一步 u 的梯度 = (系数) × v_c，不同的词会收到不同的梯度，对称性已经被 V 打破了。 每一轮：抽样本对 → 打乱 → 每批采负例 → 求梯度 → 用 np.add.at 只更新涉及的行。 【参数】 id_sentences : list[list[int]]，编号句子 counts : np.ndarray (V,)，每个词的出现次数（用于噪声分布和下采样） dim : int，词向量维度 d window : int，最大窗口半径 K : int，每个正例配几个负例。论文建议小语料 5~20、大语料 2~5 epochs : int，训练轮数 lr : float，初始学习率（每个样本的梯度不取平均，所以沿用 word2vec 的 0.025） batch_size : int，每批样本数 subsample_t : float 或 None，下采样阈值；None 表示不下采样 dynamic_window : bool，是否使用动态窗口 noise_power : float，噪声分布的指数（0.75） seed : int，随机种子 verbose : bool，是否打印每轮的平均损失 【返回】 (center_vectors, outside_vectors, history)：两个 (V, d) 矩阵和每轮的平均损失列表 【边界】 - 打印的"平均损失"是 SGNS 的损失（1 个正例 + K 个负例的二分类损失之和）， 和第 1 讲朴素 softmax 的损失（交叉熵）**不能直接比较**。 - 学习率太大时，np.add.at 在一批里对高频词累加很多次梯度，可能"一步走太远"导致不稳定； 可以减小 lr 或 batch_size。 【NLP 联系】 真实的 word2vec 用多线程、每个线程一个样本一个样本地更新共享的参数（"Hogwild" 式异步 SGD）， 不加锁也能正常训练，因为每次只改很少几行，不同线程很少同时改同一行。 """ V = len(counts) rng = np.random.default_rng(seed) C = (rng.random((V, dim)) - 0.5) / dim U = np.zeros((V, dim)) sampler = NegativeSampler(noise_distribution(counts, noise_power), seed=seed + 1) keep = keep_probabilities(counts, subsample_t) if subsample_t else None # 先抽一次样本对，估计总步数（用于学习率线性衰减） first_pairs = make_skipgram_pairs(id_sentences, window, rng, dynamic_window, keep) total_steps = epochs * (len(first_pairs) // batch_size + 1) step, history = 0, [] for epoch in range(epochs): pairs = first_pairs if epoch == 0 else make_skipgram_pairs(id_sentences, window, rng, dynamic_window, keep) pairs = pairs[rng.permutation(len(pairs))] loss_sum = 0.0 for start in range(0, len(pairs), batch_size): batch = pairs[start:start + batch_size] alpha = lr * max(1e-4, 1.0 - step / total_steps) negatives = sampler.sample((len(batch), K)) loss, g_vc, rows, g_u = ns.sgns_batch(batch[:, 0], batch[:, 1], negatives, C, U) np.add.at(C, batch[:, 0], -alpha * g_vc) np.add.at(U, rows.ravel(), -alpha * g_u.reshape(-1, dim)) loss_sum += loss step += 1 history.append(loss_sum / len(pairs)) if verbose: print(f" 第 {epoch + 1:>2d} 轮：{len(pairs):>6d} 个样本对，平均损失 {history[-1]:.4f}") return C, U, history # ============================================================================= # 第 3 部分：CBOW（负采样） # ============================================================================= def make_cbow_examples(id_sentences, window, rng, dynamic_window=True, keep_prob=None): """ 为 CBOW 抽样本：每个位置一个样本 = (中心词, 它的全部上下文词)。 【参数】同 make_skipgram_pairs。 【返回】 (centers, contexts) centers : np.ndarray[int] (N,) contexts : np.ndarray[int] (N, 2·window)，不足的位置用 -1 填充 【边界】 - 没有任何上下文词的位置（句子只剩 1 个词）会被跳过。 """ centers, contexts = [], [] for sent in id_sentences: if keep_prob is not None: sent = [w for w in sent if rng.random() < keep_prob[w]] for t, center in enumerate(sent): b = int(rng.integers(1, window + 1)) if dynamic_window else window ctx = [sent[j] for j in range(max(0, t - b), min(len(sent), t + b + 1)) if j != t] if ctx: centers.append(center) contexts.append(ctx + [-1] * (2 * window - len(ctx))) return np.array(centers, dtype=int), np.array(contexts, dtype=int).reshape(-1, 2 * window) def train_cbow(id_sentences, counts, dim=10, window=2, K=5, epochs=20, lr=0.025, batch_size=128, subsample_t=1e-2, dynamic_window=True, noise_power=0.75, seed=0, verbose=True): """ 训练 CBOW（负采样）词向量：用上下文词向量的**平均** h 去预测中心词。 【思路】（讲义第 4.1 节） h = (1/|ctx|) Σ_{j ∈ ctx} v_j （上下文向量的平均） J = -log σ(u_c · h) - Σ_k log σ(-u_{w_k} · h) （和 SGNS 一样的二分类损失，只是 v_c 换成了 h） ∂J/∂h 的算法和 SGNS 里的 ∂J/∂v_c 完全相同； 因为 h 是平均值，每个上下文词向量得到 ∂J/∂h ÷ |ctx| 的梯度。 【参数】同 train_sgns。 【返回】(context_vectors, center_vectors, history) 注意 CBOW 里角色互换了：输入矩阵（第一个返回值）存的是"作为上下文时"的向量， 输出矩阵（第二个）存的是"作为被预测的中心词时"的向量。最终同样用两者之和。 【边界】 - CBOW 每个位置只产生 1 个样本（skip-gram 产生 2m 个），所以训练更快； 但它把上下文"平均"了，对罕见词的学习一般不如 skip-gram 细致（论文中的经验结论）。 """ V = len(counts) rng = np.random.default_rng(seed) Cin = (rng.random((V, dim)) - 0.5) / dim Uout = np.zeros((V, dim)) sampler = NegativeSampler(noise_distribution(counts, noise_power), seed=seed + 1) keep = keep_probabilities(counts, subsample_t) if subsample_t else None first = make_cbow_examples(id_sentences, window, rng, dynamic_window, keep) total_steps = epochs * (len(first[0]) // batch_size + 1) step, history = 0, [] for epoch in range(epochs): centers, ctx = first if epoch == 0 else make_cbow_examples(id_sentences, window, rng, dynamic_window, keep) order = rng.permutation(len(centers)) centers, ctx = centers[order], ctx[order] loss_sum = 0.0 for start in range(0, len(centers), batch_size): c, cx = centers[start:start + batch_size], ctx[start:start + batch_size] alpha = lr * max(1e-4, 1.0 - step / total_steps) mask = cx >= 0 # (B, 2m)，哪些位置是真的上下文词 n_ctx = mask.sum(axis=1) # (B,) h = (Cin[np.where(mask, cx, 0)] * mask[:, :, None]).sum(axis=1) / n_ctx[:, None] # (B, d) # 和 SGNS 一样的计算，只是"中心词向量"换成了 h，"上下文词"换成了中心词 c negatives = sampler.sample((len(c), K)) rows = np.concatenate([c[:, None], negatives], axis=1) u = Uout[rows] signs = np.ones(rows.shape) signs[:, 1:] = -1.0 probs = sigmoid(signs * np.einsum("bkd,bd->bk", u, h)) loss_sum += -np.sum(np.log(np.maximum(probs, 1e-12))) g = -(1.0 - probs) * signs grad_h = np.einsum("bk,bkd->bd", g, u) # (B, d) grad_u = g[:, :, None] * h[:, None, :] # (B, K+1, d) np.add.at(Uout, rows.ravel(), -alpha * grad_u.reshape(-1, dim)) per_ctx = np.repeat((grad_h / n_ctx[:, None])[:, None, :], cx.shape[1], axis=1) # (B, 2m, d) np.add.at(Cin, cx[mask], -alpha * per_ctx[mask]) step += 1 history.append(loss_sum / len(centers)) if verbose: print(f" 第 {epoch + 1:>2d} 轮：{len(centers):>6d} 个样本，平均损失 {history[-1]:.4f}") return Cin, Uout, history # ============================================================================= # 第 4 部分：检验 SGNS ≈ 分解"平移的 PMI 矩阵" # ============================================================================= def shifted_pmi_check(C, U, id_sentences, window, K, min_count=20): """ 检验 Levy & Goldberg 的结论：SGNS 的最优解满足 u_o · v_c ≈ PMI(c, o) - log K。 【思路】（讲义第 5.2 节有推导） 从语料里数出所有 (c, o) 对的次数，算出 PMI(c, o) = log [ #(c,o) · N / (#c · #o) ]， 其中 #c、#o 是它们作为中心词 / 上下文词的次数，N 是总对数。 （严格地说，负采样的噪声分布是词频的 0.75 次方，PMI 里的 #o 也应该换成相应的分布； 这里用最简单的版本，所以只期望看到"强相关"，而不是完全相等。） 对出现次数 >= min_count 的词对，比较点积 u_o · v_c 和 PMI - log K。 【参数】 C, U : 训练好的中心词 / 上下文向量 id_sentences : 训练用的编号句子 window, K : 训练时的窗口与负例数 min_count : 只看至少共现这么多次的词对（次数太少的 PMI 估计很不可靠） 【返回】 (pearson_r, mean_dot, mean_target, n_pairs) 【边界】 - 理论结论假设向量维度足够大（能精确表示任意矩阵）。d 很小时模型只能做低秩近似， 相关性会变弱。演示 4 比较了 d = 10 和 d = 50。 - 结论是在"固定窗口、不下采样"的设定下推出的，所以检验时用这样训练出的模型。 """ pairs = skipgram_pairs(id_sentences, window) V = C.shape[0] M = np.zeros((V, V)) np.add.at(M, (pairs[:, 0], pairs[:, 1]) if isinstance(pairs, np.ndarray) else tuple(np.array(pairs).T), 1) N = M.sum() pmi = np.log(np.maximum(M, 1e-12) * N / np.maximum(M.sum(1, keepdims=True) * M.sum(0, keepdims=True), 1e-12)) sel = M >= min_count dots = (C @ U.T)[sel] target = pmi[sel] - np.log(K) r = np.corrcoef(dots, target)[0, 1] return float(r), float(dots.mean()), float(target.mean()), int(sel.sum()) # ============================================================================= # 演示 # ============================================================================= def main(): sents = generate_corpus() word2id, id2word, cnt = build_vocab(sents) counts = np.array([cnt[w] for w in id2word], dtype=float) ids = sentences_to_ids(sents, word2id) content = list(word_categories()) # ---------------- 演示 1：下采样 ---------------- print_header("演示 1：高频词下采样（t = 0.01）") keep = keep_probabilities(counts, 1e-2) for w in ["the", "a", "is", "i", "queen", "cat", "paris", "spanish"]: i = word2id[w] print(f" {w:<8s}词频 {counts[i] / counts.sum():6.2%} 保留概率 {keep[i]:.2f}") rng = np.random.default_rng(0) kept = sum(rng.random() < keep[w] for s in ids for w in s) print(f"\n 下采样前 {int(counts.sum())} 个词，一次下采样后约剩 {kept} 个（{kept / counts.sum():.0%}）") print(" → 被丢掉的几乎全是虚词；内容词（cat、queen、paris……）全部保留。") # ---------------- 演示 2：训练 SGNS ---------------- print_header("演示 2：训练 skip-gram + 负采样（d = 10，K = 5，下采样 + 动态窗口）") t0 = time.perf_counter() C, U, hist = train_sgns(ids, counts) print(f"\n 训练用时 {time.perf_counter() - t0:.1f} 秒（取决于你的电脑）") vec = C + U print("\n 最近邻（只在内容词中找）：") for q in ["cat", "apple", "car", "king", "girl", "paris", "french"]: nb = nearest_neighbors(q, vec, word2id, id2word, k=4, restrict_to=content) print(f" {q:<7s}→ " + ", ".join(f"{w}({s:.2f})" for w, s in nb)) # ---------------- 演示 3：对比 ---------------- print_header("演示 3：各种变体的对比（同样的语料，d = 10，每种跑 3 个随机种子）") l1_train = importlib.import_module("05_word2vec_training") pairs_fixed = skipgram_pairs(ids, 2) configs = [ ("第 1 讲：朴素 softmax（15 轮）", lambda s: l1_train.train_word2vec(pairs_fixed, len(id2word), dim=10, seed=s, verbose=False)), ("SGNS（下采样 + 动态窗口）", lambda s: train_sgns(ids, counts, seed=s, verbose=False)), ("SGNS（不下采样）", lambda s: train_sgns(ids, counts, subsample_t=None, seed=s, verbose=False)), ("SGNS（固定窗口）", lambda s: train_sgns(ids, counts, dynamic_window=False, seed=s, verbose=False)), ("SGNS（不下采样 + 固定窗口）", lambda s: train_sgns(ids, counts, subsample_t=None, dynamic_window=False, seed=s, verbose=False)), ("SGNS（噪声分布 power = 1）", lambda s: train_sgns(ids, counts, noise_power=1.0, seed=s, verbose=False)), ("SGNS（噪声分布 power = 0）", lambda s: train_sgns(ids, counts, noise_power=0.0, seed=s, verbose=False)), ("CBOW（负采样）", lambda s: train_cbow(ids, counts, seed=s, verbose=False)), ] print(f"\n {'方法':<30s}{'用时/次':>8s}{'区分度':>9s}{'类比（3 次平均）':>16s}{' 最低 ~ 最高':>14s}") for name, fn in configs: rs, t0 = [], time.perf_counter() for seed in [0, 1, 2]: Ck, Uk, _ = fn(seed) rs.append(evaluate_vectors(name, Ck + Uk, word2id, id2word, verbose=False)) secs = (time.perf_counter() - t0) / 3 accs = [r["analogy"] for r in rs] print(f" {name:<30s}{secs:>7.1f}s{np.mean([r['gap'] for r in rs]):>9.3f}" f"{np.mean(accs):>14.1%}{min(accs):>10.1%} ~ {max(accs):.1%}") print(""" 说明： - 用时取决于你的电脑。在只有 171 个词的玩具语料上，朴素 softmax 并不慢； 负采样的优势要在词表很大时才能体现（见代码 01 演示 6）。 - 注意"最低 ~ 最高"那一列：只换一个随机种子，类比准确率就能相差 10 个百分点左右。 96 道题太少，单次运行的差异很大一部分是**随机波动**。比较方法时一定要多跑几次（讲义第 7.4 节）。""") # ---------------- 演示 4：平移的 PMI ---------------- print_header("演示 4：SGNS 在分解'平移的 PMI 矩阵'吗？（固定窗口、不下采样）") K = 5 for d in [10, 50]: Cd, Ud, _ = train_sgns(ids, counts, dim=d, K=K, epochs=20, subsample_t=None, dynamic_window=False, noise_power=1.0, verbose=False) r, md, mt, n = shifted_pmi_check(Cd, Ud, ids, window=2, K=K) print(f" d = {d:>2d}：{n} 个高频词对，u·v 与 PMI - log K 的相关系数 r = {r:.3f}；" f"平均值 {md:+.2f} vs {mt:+.2f}") print(""" 说明：为了贴近理论推导的设定，这里用固定窗口、不下采样、噪声分布 power = 1。 相关系数越接近 1，说明学到的点积越接近"PMI - log K"。维度越大，模型越能精确地表示这个矩阵。""") if __name__ == "__main__": main()