""" ============================================================================= CS349E · 第 2 讲 · 01_word2vec.py —— 把词变成向量：查表、word2vec 的代价账、在合成语料上训练 skip-gram ============================================================================= 【这个脚本要回答的问题】 Transformer 的第一步是"把 token 变成向量"。这件事的思想来自 word2vec（本讲阅读材料之一，Mikolov et al. 2013）。 演示 1：词嵌入就是查表——"独热向量 × 矩阵"和"取矩阵的一行"完全相同，代价却差几万倍。 以及 Llama-3.1-8B 的词嵌入表和 lm_head：一个是查表，一个是真矩阵乘法。 演示 2：word2vec 论文的代价账 Q（每个训练样本要碰多少参数）：NNLM、RNNLM、CBOW、skip-gram； 为什么"去掉隐藏层"能让训练数据从几亿词涨到几十亿词；完整 softmax vs 层次 softmax vs 负采样。 演示 3：在一个合成语料上训练 skip-gram + 负采样（numpy 手写梯度），看最近邻和 "king − man + woman ≈ queen"式的类比是怎样从"共现统计"里长出来的。 演示 4：维度和数据量：维度太小、数据太少，类比准确率都会下降（论文表 2 的现象，在玩具上复现）。 【合成语料是怎么造的（重要：先看懂它，才知道类比为什么能成立）】 真实语料太大，这里造一个"规则已知"的小语料。每个句子围绕一个中心词，周围放几个上下文词和虚词，再打乱顺序： - 性别对（king/queen、man/woman、father/mother …）：上下文 = 这个"角色"的词（palace、crown …） + 这个"性别"的词（he/his … 或 she/her …）。所以 king ≈ "王室" + "男性"，queen ≈ "王室" + "女性"。 - 国家 / 首都（france/paris …）：国家词的上下文 = 该国的专属词（france_lang …）+ "国家"类词； 首都词的上下文 = 同样的专属词 + "城市"类词；另有 "paris capital france" 这类句子。 - 形容词 / 比较级（big/bigger …）：上下文 = 这个形容词的话题词 + "原级"类词（very …）或"比较级"类词（than …）。 skip-gram 只看"谁和谁一起出现"。如果 king 和 queen 共享"王室"上下文、king 和 man 共享"男性"上下文， 向量空间里就会出现一个大致恒定的"男 → 女"方向，类比就成立。真实语料里的规律没这么干净，原理相同。 【怎么读输出】全部是确定性的计算（固定种子），每次运行结果相同。 【运行】python3 01_word2vec.py （约 15 秒） 【依赖】numpy；从 ../lecture01 导入 console、llm_math。 """ import math import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import fmt_bytes, fmt_num, print_header, print_observations, print_table # noqa: E402 from llm_math import LLAMA31_8B # noqa: E402 # --------------------------------------------------------------------------- # 合成语料的"词典" # --------------------------------------------------------------------------- GENDER_PAIRS = [("king", "queen"), ("prince", "princess"), ("man", "woman"), ("boy", "girl"), ("father", "mother"), ("son", "daughter"), ("brother", "sister"), ("uncle", "aunt"), ("husband", "wife"), ("actor", "actress")] ROLE_CONTEXT = [["palace", "crown", "throne", "kingdom"], ["castle", "heir", "royal", "young"], ["adult", "works", "street", "job"], ["toys", "playground", "school", "small"], ["children", "family", "parent", "home"], ["parents", "child", "born", "raised"], ["siblings", "older", "younger", "twin"], ["nephew", "niece", "visits", "relative"], ["married", "wedding", "spouse", "ring"], ["film", "stage", "movie", "role"]] GENDER_CONTEXT = [["he", "his", "him", "mr"], ["she", "her", "hers", "mrs"]] CAPITAL_PAIRS = [("france", "paris"), ("germany", "berlin"), ("italy", "rome"), ("japan", "tokyo"), ("china", "beijing"), ("spain", "madrid"), ("egypt", "cairo"), ("russia", "moscow"), ("canada", "ottawa"), ("greece", "athens"), ("india", "delhi"), ("peru", "lima")] COUNTRY_CONTEXT = [[f"{c}_{s}" for s in ("lang", "food", "river")] for c, _ in CAPITAL_PAIRS] NATION_WORDS = ["nation", "border", "government", "population"] CITY_WORDS = ["city", "streets", "downtown", "mayor"] ADJ_PAIRS = [("big", "bigger"), ("small", "smaller"), ("fast", "faster"), ("slow", "slower"), ("cheap", "cheaper"), ("warm", "warmer"), ("cold", "colder"), ("strong", "stronger"), ("weak", "weaker"), ("tall", "taller")] ADJ_CONTEXT = [[f"{a}_t1", f"{a}_t2"] for a, _ in ADJ_PAIRS] BASE_WORDS = ["very", "quite", "so"] COMPARATIVE_WORDS = ["than", "even", "much"] FILLERS = ["the", "a", "is", "of", "and", "in", "to", "was", "it", "that"] # 类比测试：三组关系。capital 组写成 (首都, 国家)，问题形如 "paris : france = berlin : ?" ANALOGY_SETS = { "性别（语义）": GENDER_PAIRS, "首都 → 国家（语义）": [(city, country) for country, city in CAPITAL_PAIRS], "比较级（句法）": ADJ_PAIRS, } def make_corpus(n_sentences, seed=0): """ 生成合成语料：一个列表，每个元素是一个句子（词的列表）。 【思路】见文件开头。40% 的句子讲性别对，35% 讲国家 / 首都，25% 讲形容词。每个句子 = 中心词 + 4 个 上下文词 + 2–4 个虚词，打乱顺序（skip-gram 只关心窗口里的共现，不关心语法）。 【参数】n_sentences：句子数；seed：随机种子 【返回】list[list[str]] 【Python 小知识】rng.choice(列表, k) 从列表里有放回地抽 k 个；rng.shuffle(lst) 原地打乱。 """ rng = np.random.default_rng(seed) sents = [] for _ in range(n_sentences): r = rng.random() if r < 0.40: i, g = rng.integers(len(GENDER_PAIRS)), rng.integers(2) center = GENDER_PAIRS[i][g] ctx = list(rng.choice(ROLE_CONTEXT[i], 2)) + list(rng.choice(GENDER_CONTEXT[g], 2)) elif r < 0.75: i, kind = rng.integers(len(CAPITAL_PAIRS)), rng.integers(3) country, city = CAPITAL_PAIRS[i] if kind == 0: center, ctx = country, list(rng.choice(COUNTRY_CONTEXT[i], 2)) + list(rng.choice(NATION_WORDS, 2)) elif kind == 1: center, ctx = city, list(rng.choice(COUNTRY_CONTEXT[i], 2)) + list(rng.choice(CITY_WORDS, 2)) else: center, ctx = city, ["capital", country] else: i, g = rng.integers(len(ADJ_PAIRS)), rng.integers(2) center = ADJ_PAIRS[i][g] ctx = list(rng.choice(ADJ_CONTEXT[i], 2)) + list(rng.choice(BASE_WORDS if g == 0 else COMPARATIVE_WORDS, 2)) words = [center] + [str(w) for w in ctx] + [str(w) for w in rng.choice(FILLERS, rng.integers(2, 5))] rng.shuffle(words) sents.append(words) return sents def skipgram_pairs(sents, window): """ 把语料变成 (中心词 id, 上下文词 id) 对：每个词和它前后 window 个词各组成一对。 【返回】(vocab 列表, {词: id}, pairs：[P, 2] 整数数组, counts：每个词出现的次数) 【边界】句首 / 句尾的窗口被截断（不跨句子）。 【Python 小知识】sorted(set) 让词表顺序固定（集合本身的遍历顺序依赖哈希种子，不固定）。 """ vocab = sorted({w for s in sents for w in s}) idx = {w: i for i, w in enumerate(vocab)} counts = np.zeros(len(vocab)) pairs = [] for s in sents: ids = [idx[w] for w in s] for i, c in enumerate(ids): counts[c] += 1 for j in range(max(0, i - window), min(len(ids), i + window + 1)): if j != i: pairs.append((c, ids[j])) return vocab, idx, np.array(pairs, dtype=np.int64), counts def train_sgns(pairs, counts, dim, negatives=5, epochs=3, lr0=0.025, batch=512, seed=0): """ skip-gram + 负采样（SGNS，Mikolov et al. 2013b）的 numpy 实现。 【目标】对每一对 (中心词 c, 真实上下文词 o)，再随机抽 k 个"负样本"词 n_1 … n_k，最大化 log σ(v_c · u_o) + Σ_j log σ(−v_c · u_{n_j}) v 是"输入向量"（最后当作词向量用），u 是"输出向量"。直觉：真实上下文的点积要大，随机词的点积要小。 负样本按 词频^0.75 抽（论文的经验值：比按词频抽更常抽到少见词）。 【为什么不用完整 softmax】完整 softmax 要对词表里每个词算一次点积（V 次），负采样只算 k+1 次（演示 2）。 【梯度（手推）】设 s = σ(v·u)。对正样本，∂(−log s)/∂(v·u) = s − 1；对负样本，∂(−log σ(−v·u))/∂(v·u) = σ(v·u)。 于是 ∂/∂v = (s_pos − 1)·u_o + Σ_j s_neg_j·u_{n_j}，∂/∂u_o = (s_pos − 1)·v，∂/∂u_{n_j} = s_neg_j·v。 【学习率】从 lr0 线性降到接近 0（论文 §4.2 的做法）。 【参数】pairs：[P, 2]；counts：词频；dim：向量维度；negatives：k；epochs：遍数；batch：每次更新多少对 【返回】输入向量矩阵 [V, dim] 【Python 小知识】np.add.at(A, idx, vals) 把 vals 加到 A[idx] 上，idx 里有重复的行也会**全部累加**； 普通的 A[idx] += vals 遇到重复下标只会加一次（后写的覆盖先写的）。 """ rng = np.random.default_rng(seed) V = len(counts) noise = counts ** 0.75 noise /= noise.sum() w_in = (rng.random((V, dim)) - 0.5) / dim # word2vec 原始实现的初始化：小的均匀随机数 w_out = np.zeros((V, dim)) # 输出向量从 0 开始 P = len(pairs) total, done = epochs * P, 0 for _ in range(epochs): perm = rng.permutation(P) for b0 in range(0, P, batch): bi = perm[b0:b0 + batch] c, o = pairs[bi, 0], pairs[bi, 1] lr = lr0 * max(1e-4, 1.0 - done / total) done += len(bi) neg = rng.choice(V, size=(len(bi), negatives), p=noise) v, up, un = w_in[c], w_out[o], w_out[neg] s_pos = 1.0 / (1.0 + np.exp(-np.sum(v * up, axis=1))) s_neg = 1.0 / (1.0 + np.exp(-np.einsum("bd,bkd->bk", v, un))) g_pos = (s_pos - 1.0)[:, None] g_neg = s_neg[:, :, None] grad_v = g_pos * up + np.sum(g_neg * un, axis=1) np.add.at(w_out, o, -lr * g_pos * v) np.add.at(w_out, neg, -lr * g_neg * v[:, None, :]) np.add.at(w_in, c, -lr * grad_v) return w_in def normalize_rows(E): """每一行除以自己的长度（之后点积 = 余弦相似度）。""" return E / np.linalg.norm(E, axis=1, keepdims=True) def analogy_accuracy(pairs_list, idx, E): """ 类比测试（论文 §4 的方法）：对同一组关系里的两对 (a, a*)、(b, b*)，计算 x = v(a*) − v(a) + v(b)， 在全部词里找余弦相似度最大的词（排除 a、a*、b 三个问题词），等于 b* 才算对。 【返回】(准确率, 题目数) 【边界】同义词也算错（和论文一样严格）。 """ En = normalize_rows(E) ok = tot = 0 for a, a2 in pairs_list: for b, b2 in pairs_list: if a == b: continue x = En[idx[a2]] - En[idx[a]] + En[idx[b]] sims = En @ x for w in (a, a2, b): sims[idx[w]] = -np.inf ok += int(np.argmax(sims) == idx[b2]) tot += 1 return ok / tot, tot def nearest(word, idx, vocab, E, k=4): """余弦相似度最大的 k 个词（不含自己）。""" En = normalize_rows(E) sims = En @ En[idx[word]] sims[idx[word]] = -np.inf top = np.argsort(-sims)[:k] return ", ".join(f"{vocab[i]}({sims[i]:.2f})" for i in top) # --------------------------------------------------------------------------- # 演示 1：嵌入 = 查表 # --------------------------------------------------------------------------- def demo1_embedding_is_lookup(): """ 独热向量 e_i（第 i 个位置是 1，其余是 0）乘以嵌入矩阵 E[V, d]，结果就是 E 的第 i 行。 数学上是矩阵乘法，实现上是"按下标取一行"（gather）。 """ print_header("演示 1：词嵌入 = 查表（独热向量 × 矩阵 = 取一行）") rng = np.random.default_rng(0) V, d = 1000, 64 E = rng.standard_normal((V, d)) ids = np.array([3, 141, 592]) one_hot = np.zeros((len(ids), V)) one_hot[np.arange(len(ids)), ids] = 1.0 diff = np.max(np.abs(one_hot @ E - E[ids])) print(f"V = {V}, d = {d}：独热 × E 与 E[ids] 的最大差 = {diff:.1e}") rows = [["独热向量 × E（矩阵乘法）", fmt_num(2 * V * d, 0), fmt_bytes(V * d * 4)], ["E[ids]（查表 / gather）", "0", fmt_bytes(d * 4)]] print_table(["做法（每个 token）", "FLOPs", "要读的字节（FP32）"], rows) m = LLAMA31_8B table = m.vocab * m.d_model * 2 print(f"\nLlama-3.1-8B（BF16）：词表 {m.vocab:,} × d = {m.d_model}") rows = [ ["词嵌入表（输入端）", fmt_bytes(table), "查 B 行", fmt_bytes(64 * m.d_model * 2), "0"], ["lm_head（输出端，不共享）", fmt_bytes(table), "整张矩阵乘", fmt_bytes(table), fmt_num(2 * 64 * m.d_model * m.vocab)], ] print_table(["矩阵", "大小", "decode 时怎么用", "B = 64 时一步读多少", "FLOPs"], rows) print_observations([ "数学上嵌入是矩阵乘法，实现上只是取一行：不做任何运算，只读 d 个数。所以第 1 讲的公式里，" "'参与矩阵乘法的参数'不含词嵌入表。", f"输出端正好相反：要给词表里每个词打分，lm_head 必须整张读、整张乘。Llama-3.1-8B 的 lm_head 有 " f"{fmt_bytes(table)}，占全部权重的 {table / m.weight_bytes() * 100:.1f}%——decode 每一步都要读一遍。", "这也是为什么小模型常常让 lm_head 和词嵌入共享同一个矩阵（tied）：省下的是参数和显存，" "但 decode 时 lm_head 仍然要整张读，带宽一点没省。", ]) # --------------------------------------------------------------------------- # 演示 2：word2vec 论文的代价账 # --------------------------------------------------------------------------- def demo2_complexity(): """ 论文定义：训练代价 O = E × T × Q，E 是遍数，T 是训练词数，Q 是每个样本要碰的参数个数： NNLM ：Q = N·D + N·D·H + H·V （输入 N 个词、投影维 D、隐藏层 H、词表 V） RNNLM ：Q = H·H + H·V CBOW ：Q = N·D + D·log2(V) （去掉隐藏层；输出用层次 softmax） skip-gram：Q = C·(D + D·log2(V)) （C 是最远上下文距离） 用层次 softmax（Huffman 树）时 H·V 换成 H·log2(V)。另一篇论文（2013b）的负采样把 D·log2(V) 换成 D·(k+1)。 """ print_header("演示 2：word2vec 的代价账 Q（每个训练样本要碰多少参数）") N, D, H, C, V, k = 10, 500, 500, 10, 1_000_000, 5 lg = math.log2(V) rows = [ ["NNLM（完整 softmax）", "N·D + N·D·H + H·V", N * D + N * D * H + H * V], ["NNLM（层次 softmax）", "N·D + N·D·H + H·log2V", N * D + N * D * H + H * lg], ["RNNLM（层次 softmax）", "H·H + H·log2V", H * H + H * lg], ["CBOW（层次 softmax）", "N·D + D·log2V", N * D + D * lg], ["skip-gram（层次 softmax）", "C·(D + D·log2V)", C * (D + D * lg)], ["skip-gram（负采样 k=5）", "C·(D + D·(k+1))", C * (D + D * (k + 1))], ] base = rows[0][2] T = 1e9 table = [[r[0], r[1], fmt_num(r[2]), f"{base / r[2]:,.0f}×", fmt_num(2 * r[2] * T)] for r in rows] print(f"取 N = {N}、D = H = {D}、C = {C}、V = {V:,}（log2 V ≈ {lg:.1f}）、T = 10 亿词、1 遍：") print_table(["模型", "Q 的公式", "Q", "比完整 softmax 的 NNLM 省", "总 FLOPs ≈ 2·Q·T"], table) print_observations([ "完整 softmax 的 H·V 一项是 5 亿：给 100 万个词每个都打一次分。层次 softmax 把它降到 H·log2V ≈ 1 万，负采样（每个上下文词）降到 D·(k+1) = 3 千。", "降完之后，NNLM 的瓶颈变成 N·D·H（非线性隐藏层，250 万）。word2vec 的关键一步是**整个去掉隐藏层**：" "CBOW 的 Q 只有约 1.5 万，比 NNLM 小两个数量级——同样的算力能多看两个数量级的数据。", "这是全课第一次出现的主题：**先算账，再设计**。模型变简单了，但能吃下更多数据，最后效果更好" "（论文表 6：skip-gram 在 60 亿词上训练，类比准确率 65.6%，NNLM 只有 50.8%）。", "Transformer 走的是相反的方向：每个 token 的运算很多，但全部是能并行的大矩阵乘法（02 脚本演示 6）。", ]) # --------------------------------------------------------------------------- # 演示 3：在合成语料上训练 skip-gram # --------------------------------------------------------------------------- def demo3_train(): print_header("演示 3：在合成语料上训练 skip-gram + 负采样") sents = make_corpus(20000, seed=0) vocab, idx, pairs, counts = skipgram_pairs(sents, window=3) print(f"语料：{len(sents):,} 句，{int(counts.sum()):,} 个词，词表 {len(vocab)}，训练对 {len(pairs):,}") print(f"每个训练对：完整 softmax 要算 {len(vocab)} 个点积，负采样（k = 5）只算 6 个；" f"词表 100 万时是 100 万 vs 6。") E = train_sgns(pairs, counts, dim=30, seed=0) E_rand = np.random.default_rng(1).standard_normal(E.shape) print("\n最近邻（余弦相似度）：") for w in ("king", "paris", "bigger", "she"): print(f" {w:>7} → {nearest(w, idx, vocab, E)}") rows = [] for name, pl in ANALOGY_SETS.items(): acc, tot = analogy_accuracy(pl, idx, E) acc_r, _ = analogy_accuracy(pl, idx, E_rand) rows.append([name, tot, f"{acc * 100:.1f}%", f"{acc_r * 100:.1f}%"]) print("\n类比测试（a : a* = b : ?，在全部词里找最近的，排除三个问题词）：") print_table(["关系", "题数", "训练后的向量", "随机向量"], rows) En = normalize_rows(E) offs = np.array([En[idx[b]] - En[idx[a]] for a, b in GENDER_PAIRS]) offs_n = normalize_rows(offs) cos = offs_n @ offs_n.T mean_cos = (cos.sum() - len(offs)) / (len(offs) * (len(offs) - 1)) print(f"\n10 个'男 → 女'差向量两两之间的平均余弦相似度：{mean_cos:.2f}（随机方向约为 0）") print_observations([ "训练后的向量把'性别'、'首都'、'比较级'三种关系都学成了**近似恒定的方向**：10 个 '男 → 女' 差向量" "彼此高度对齐，所以 v(queen) ≈ v(king) − v(man) + v(woman)。随机向量的准确率接近 0。", "模型从来没被告诉'性别'是什么，它只看到了共现：king 和 queen 共享王室词，king 和 man 共享男性词。" "这就是**分布式假设**（distributional hypothesis）：一个词的意思由它的上下文决定。", "最近邻却几乎都是'经常和它一起出现的话题词'（king → throne、palace；paris → france_food …），" "看不出性别或'首都'关系：这些关系藏在**方向**里，不在距离里。这正是论文用类比、而不是最近邻来评价词向量的原因" "（§1：词之间有 multiple degrees of similarity）。", "这个语料是刻意造得很干净的，所以准确率接近 100%；论文在真实语料上最好的结果约 65%。", ]) return idx, E # --------------------------------------------------------------------------- # 演示 4：维度和数据量 # --------------------------------------------------------------------------- def demo4_dim_and_data(): print_header("演示 4：维度 × 数据量 → 类比准确率（三组关系的平均）") sizes = (1000, 4000, 20000) dims = (2, 5, 10, 30) table = [] for n in sizes: sents = make_corpus(n, seed=100 + n) vocab, idx, pairs, counts = skipgram_pairs(sents, window=3) row = [f"{n:,} 句（{int(counts.sum()):,} 词）"] for d in dims: E = train_sgns(pairs, counts, dim=d, seed=0) accs = [analogy_accuracy(pl, idx, E)[0] for pl in ANALOGY_SETS.values()] row.append(f"{np.mean(accs) * 100:.1f}%") table.append(row) print_table(["语料"] + [f"d = {d}" for d in dims], table) print_observations([ "d = 2 时无论多少数据都学不好（最高 3.7%）：两维空间里放不下几十个词、好几种互相独立的关系。", "数据太少（1,000 句）时，维度从 2 加到 30 也只有 3.2%；d = 5 时，数据从 4,000 句加到 20,000 句只从 25.5% " "涨到 38.7%，d = 10 时却从 60.9% 涨到 100%。维度和数据要一起加（论文 §4.2 表 2 的结论）。", "【系统联系】维度决定每个 token 的向量有多长，也就决定了后面每个矩阵乘法的大小。Transformer 的 d_model " "（Llama-3.1-8B 是 4096）就是这个维度，它直接进入第 1 讲的每一条公式。", ]) if __name__ == "__main__": demo1_embedding_is_lookup() demo2_complexity() demo3_train() demo4_dim_and_data()