项目 02：embedding 与语义几何

{{ intro_card( points=[ "用 numpy 演示 one-hot → embedding 的关系（一次矩阵乘法）", "训练一个 toy next-token 模型，观察 embedding 从随机到结构化的过程", "可视化训练前后 cosine similarity 矩阵的变化", "破坏实验：冻结 embedding / 打乱 ID / 维度太小，看模型怎么坏", ], audience="想理解 LLM 内部表示、想自己训练 embedding 的工程师", prerequisites="PyTorch 基础，知道 nn.Linear 和 nn.Embedding", time="35 分钟", ) }}

项目目标

理解三件事：

1. token ID 没有语义——42 和 43 只是两个编号，距离上没有任何含义
2. embedding table 把 ID 映射成可学习的向量，这才是模型内部真正"看见"的东西
3. embedding 不是预先定义好的，是被训练目标塑形出来的——next-token objective 会把"经常一起出现"的 token 拉到向量空间的近邻

做完这个项目，你能在脑子里画出"为什么调整 embedding 维度会影响 vocab embedding 参数量""为什么训练初期 loss 高、embedding 是噪声"。

背景与原理

one-hot 是 embedding 的退化情形

最朴素的"把 token 喂进神经网络"方法是 one-hot：

词表大小 = 5: ['cat', 'dog', 'fish', 'apple', 'pear']
'cat'   → [1, 0, 0, 0, 0]
'dog'   → [0, 1, 0, 0, 0]

任意两个 one-hot 向量的内积都是 0、距离都是 √2，意味着模型从输入层就接收到"所有词彼此完全无关"的假设。

接到 nn.Linear(5, hidden_dim) 后：

[0, 1, 0, 0, 0] @ W   # W shape: (5, hidden_dim)
= W 的第 2 行

这正是 embedding lookup 在数学上等价的操作——把 one-hot 乘以权重矩阵 = 从权重矩阵里取出对应行。nn.Embedding(vocab_size, dim) 只是把这个操作变成 O(1) 的查表，避免在大词表上做 O(V) 的矩阵乘。

embedding 是什么

import torch.nn as nn
emb = nn.Embedding(vocab_size=10000, embedding_dim=128)
# 内部就是一个 (10000, 128) 的可学习权重矩阵
# emb(torch.tensor([3, 7])) → (2, 128)，取第 3 行和第 7 行

embedding table 在训练开始时是随机初始化的——所有 token 在向量空间里都散在各处，cosine similarity 趋近于 0（因为高维随机向量天生接近正交）。

训练过程把这张表"塑形"——梯度通过损失函数倒回来，告诉它"这些 token 之间应该更近，那些应该更远"。

为什么 next-token 训练能产生语义结构

核心机制：模型为了预测下一个 token，必须让"在相似上下文里出现"的 token 在某个内部表示里靠近。否则后续层就没法用统一的逻辑处理它们。

这和语言学里的 分布假设（distributional hypothesis） 一致——"你只要看一个词的伙伴，就知道这个词是什么"。

Word2Vec、GloVe、现代 LLM 的 input embedding 本质都是这个机制的不同形式。

动手实现

任务 A：手算一遍 one-hot 与 embedding 的等价

import numpy as np

vocab_size, dim = 5, 4
W = np.random.randn(vocab_size, dim) * 0.1   # embedding table

# 方式 1：one-hot + matmul
def via_onehot(token_id):
    onehot = np.zeros(vocab_size)
    onehot[token_id] = 1
    return onehot @ W     # shape (dim,)

# 方式 2：直接查表
def via_lookup(token_id):
    return W[token_id]    # shape (dim,)

# 两者完全相等
for i in range(vocab_size):
    assert np.allclose(via_onehot(i), via_lookup(i))

# 但 one-hot 方式要 (1, V) @ (V, D)，V 大时极其浪费
# 查表 O(1)，是工程上唯一可行的方式

任务 B：训练一个 toy next-token 模型，看 embedding 怎么变

用最小语料训练一个 4-token 上下文窗口的字符模型，观察训练前 / 训练后的 embedding 相似度。

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

# 1) 准备小语料 + 字符 tokenizer
text = ("the cat sat on the mat the dog sat on the rug "
        "the bird flew over the cat the cat watched the bird "
        "the dog barked at the cat ") * 50

chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
data = torch.tensor([stoi[c] for c in text], dtype=torch.long)

VOCAB, DIM, CTX = len(chars), 16, 8

# 2) 极简模型：embedding → 平均池化 → 输出层
class ToyLM(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed = nn.Embedding(VOCAB, DIM)
        self.head = nn.Linear(DIM, VOCAB)
    def forward(self, x):
        # x: (B, CTX)
        h = self.embed(x).mean(dim=1)   # (B, DIM)
        return self.head(h)             # (B, VOCAB)

model = ToyLM()

# 3) 保存训练前的 embedding
emb_before = model.embed.weight.detach().clone()

# 4) 训练
def get_batch(batch_size=32):
    idx = torch.randint(0, len(data) - CTX - 1, (batch_size,))
    x = torch.stack([data[i:i+CTX] for i in idx])
    y = torch.stack([data[i+CTX] for i in idx])
    return x, y

opt = torch.optim.Adam(model.parameters(), lr=1e-2)
for step in range(2000):
    x, y = get_batch()
    logits = model(x)
    loss = F.cross_entropy(logits, y)
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 200 == 0:
        print(f"step {step}: loss={loss.item():.3f}")

emb_after = model.embed.weight.detach().clone()

# 5) 对比：哪几对字符 cosine similarity 变化最大
def cos(a, b):
    return (a @ b) / (a.norm() * b.norm() + 1e-9)

interesting = [('t', 'h'), ('c', 'a'), ('d', 'o'), ('o', 'n'), ('a', 'z')]
print(f"\n{'pair':<8} | before  | after   | delta")
for a, b in interesting:
    ia, ib = stoi[a], stoi[b]
    before = cos(emb_before[ia], emb_before[ib]).item()
    after = cos(emb_after[ia], emb_after[ib]).item()
    print(f"{a}+{b}     | {before:+.3f}  | {after:+.3f}  | {after-before:+.3f}")

典型输出：

step 0: loss=2.95
step 200: loss=2.41
step 400: loss=2.18
step 1800: loss=1.85

pair     | before  | after   | delta
t+h      | -0.043  | +0.612  | +0.655
c+a      | +0.015  | +0.421  | +0.406
d+o      | -0.027  | +0.388  | +0.415
o+n      | +0.008  | +0.557  | +0.549
a+z      | -0.011  | -0.082  | -0.071     ← 几乎没变（语料里 'a' 'z' 不挨着）

关键观察：

• 训练前所有 pair 的 cosine 都接近 0（高维随机向量近似正交）
• 训练后，在语料里经常相邻的字符对（th ca do on）相似度大幅上升
• 跟语料里不挨着的对（a+z）几乎没变
• 这就是"embedding 是被目标函数塑形"的最小直接证据

任务 C：可视化整个 embedding 空间

用 PCA 或 t-SNE 把 16 维降到 2 维：

from sklearn.decomposition import PCA
import matplotlib.pyplot as plt

pca = PCA(n_components=2)
coords = pca.fit_transform(emb_after.numpy())

plt.figure(figsize=(8, 8))
for i, ch in enumerate(chars):
    if ch == ' ':
        continue
    plt.scatter(coords[i, 0], coords[i, 1], s=80, alpha=0.6)
    plt.annotate(ch, (coords[i, 0], coords[i, 1]), fontsize=14)
plt.title("Toy LM embedding (after 2000 steps)")
plt.savefig("embedding.png", dpi=120)

在小语料训练的 toy 模型上，元音通常会聚成一簇，常见辅音组合会形成另一簇。这种结构不是被显式指定的，是被 next-token 目标函数"逼出来"的。

观察指标

每次训练都记录这些数字：

破坏实验

实验 1：冻结 embedding 不让训

model.embed.weight.requires_grad = False
# 训练 loss 依然会下降（输出层在学），但比可训 embedding 慢很多
# 验证集 loss 通常显著差

结论：embedding 是模型表达能力的核心来源之一，冻结它等于把语义层封住。

实验 2：随机打乱 token ID 后训练

# 把 data 里所有 ID 做一个固定置换
perm = torch.randperm(VOCAB)
data_shuffled = perm[data]
# 然后用 data_shuffled 训练

结论：模型能力完全不受影响——因为 embedding 是根据数据共现关系自适应学习的，token ID 只是个标签，叫什么都行。这反过来证明"ID 没有语义"。

实验 3：embedding 维度压到极小

DIM = 2     # 从 16 降到 2
# 验证集 loss 显著上升
# PCA 图上不同 token 互相挤在一起，没有可分结构

结论：dim 太小时模型容量不够，无法把 vocab_size 个 token 安排进区分明显的空间。但 dim 也不是越大越好——dim 太大会过拟合 + 浪费参数。生产 LLM 常见 dim = 512 ~ 8192，跟模型规模匹配。

实验 4：训练用 tokenizer A，推理用 tokenizer B

# 训练时 stoi: {'a': 5, 'b': 12, ...}
# 推理换一套 stoi: {'a': 18, 'b': 3, ...}
# embedding 的第 5 行对应训练里的 'a'，但推理时变成了 token 18 的代表
# 输出立刻变成乱码

结论：embedding 行号和 tokenizer ID 是强绑定的合同。换 tokenizer 必须重训或同步更新 embedding 行的对应。

交付物

• toy_lm.py：上面的完整训练脚本
• cos_compare.py：训练前/后高频 vs 低频对的 cosine 对比表
• embedding.png：训练后 embedding 空间的 PCA 2D 图
• 一段 200 字短复盘：为什么"打乱 ID 训练能跑出同样的 loss"是 embedding 机制的最强证据

与本站其他内容连接

• 项目 01：从零实现 tokenizer——上一步：怎么生成那些 ID
• ai-for-python 番外 09：embedding 来龙去脉——更详细的历史脉络
• ai-for-python 番外 11：从 One-Hot 到 LLM——离散到连续的范式转换
• llm-app/Embedding 嵌入向量——应用层（sentence embedding）与本篇（model-internal embedding）的区别

延伸阅读

• Word2Vec 原始论文
• GloVe: Global Vectors for Word Representation
• The Illustrated Word2vec
• Karpathy: makemore（embedding 与 next-char 训练）