PyTorch实现Skip-gram模型：从零训练词向量的完整指南

1. 理解这个项目到底在做什么

如果你翻过任何一本自然语言处理相关的教材或课程大纲，看到类似"6.4 PyTorch实现Skip-gram模型"这种编号，大概率是某个章节的核心实验课。标题里的6.4通常是章节号，而真正的主角是"PyTorch实现"和"Skip-gram"这两个词。这个项目解决的是一个非常朴素但极其关键的问题：如何让计算机理解一个词的意思。

很多初学者会问，现在大模型都那么火了，直接调现成的词向量工具不就行了吗？为什么还要自己从零实现一个Skip-gram？答案是：调包很容易，但理解内部机制是另一回事。Skip-gram是word2vec词向量方法中最经典的两个模型之一，它的主角不是分类器，不是生成模型，而是训练过程中"不小心"产出的副产品——词向量。这个项目会带你亲手搭建一个完整的训练流程：从语料准备、词表构建、训练样本生成，到定义PyTorch模型、负采样、训练循环，最后输出一组真正表示语义的词向量。整套流程做完，你对嵌入层（Embedding）怎么工作、训练样本怎么组织、词向量为什么有语义，都会有一个非常具体的认识。

适合谁来学？我建议三类人认真做一遍：刚学完PyTorch基础、想吃透Embedding机制的人；正在看NLP教材、对word2vec章节有疑问的人；以及需要做一个"能跑、能展示、能讲原理"的课程设计或面试项目的人。它不像调个大模型那样需要海量资源，一台普通笔记本就够了，代码量控制在200行以内，但信息密度非常高，属于那种"麻雀虽小五脏俱全"的经典教学项目。

2. 拆解Skip-gram的核心机制

2.1 从"预测上下文"到"得到词向量"的逆向思维

Skip-gram这个英文名直译是"跳元模型"，意思是跳过一些词汇、只取目标词周围的词来建模。给定一段文本"我今天喝了一杯咖啡"，如果中心词是"喝"，设定窗口大小为2，那么上下文就是"我"、"今天"、"了"、"一杯"（左右各两个词）。模型要做的事情，是看到"喝"这个输入，预测这些上下文词会出现。听起来像一个文本分类或者预测任务，对吧？但整个项目的关键点在于：我们根本不在乎预测得准不准，我们在乎的是训练结束后留下来的那张参数表。

打个比方。你想认识一位新同学，不需要记住他每天上过什么课、吃过什么饭（对应分类的准确率），你只需要在一次次见面中，慢慢形成对他"性格和外貌"的概括印象（对应词向量）。Skip-gram就是通过让"喝"反复去"预测"与它共现的词，把这种共现规律压缩成一个固定维度的向量。语义相近的词经常出现在相似的上下文里，所以它们的向量就会在空间上靠得近。比如"咖啡"和"茶"都常出现在"喝"、"一杯"、"热"附近，训练后二者的向量距离就小。

这里要特别提醒一个初学者常见的误解：训练样本不是"中心词-整段上下文文本"这种序列，而是"中心词-单个上下文词"的配对。也就是说"喝"有四个上下文词，就生成四对训练样本（喝, 我）、（喝, 今天）、（喝, 了）、（喝, 一杯）。上下文词之间是独立处理的，顺序不敏感。这个细节直接影响你后面怎么写数据加载器。

2.2 负采样：把偌大的多分类问题变成巧妙的二分类

从直觉上说，最直接的建模方式是：输入中心词的向量，输出一个词表大小的概率分布，用交叉熵训练。但在真实语料中，词表轻易就能达到几万甚至几十万，每次前向传播都要对整个词表做softmax打分，计算量大到普通机器根本跑不动。举个例子，词表大小为10万时，每个训练样本要算10万次的向量点积再加softmax，一次训练迭代动辄几千万样本，这个开销是灾难级的。

word2vec提出的解决方案叫负采样（Negative Sampling）。它的思路很直接：不预测"这个词是所有词中的哪一个"，而是判断"给定中心词和候选词，它们是否真的构成上下文关系"。这是一个二分类问题。正样本就是真实出现在窗口里的上下文词，负样本则是在词表中随机挑几个词冒充上下文。模型只需要对这两个词的向量做点积，再接一个sigmoid输出概率，计算量从O(V)降到O(k)，k通常是个位数，省了几个数量级。

负样本怎么选很讲究。如果完全均匀随机挑，那些高频词（如"的"、"是"）会被过度采样，导致模型总是觉得它们很"万能"，最后词向量质量变差。word2vec的做法是给每个词一个采样概率，对原始词频做0.75次方修正：

P(w) = f(w)^0.75 / Σ f(w')^0.75

这里的f(w)是词w在语料中出现的频率（归一化后）。为什么要加0.75这个指数？因为原始频率分布太偏斜了：少数高频词占了大量概率，绝大多数词概率趋近于零。取0.75次方有一个效果，就是把概率曲线往平缓方向拉一拉，高频词的概率相对下降一点，低频词的概率相对抬升一点，让负样本分布更均衡。这个公式不需要你手算，实现时一般是预先给每个词按这个概率生成一个采样表，训练时直接从表里随机取。

2.3 窗口与共现：训练数据是怎么从文本里长出来的

训练样本的形式确定了，接下来要解决的是数据怎么来。整个流程可以分成三步走。第一步，把原始语料切分成词序列。中文要做分词，英文按空格切就行。第二步，统计词频，过滤掉那些出现次数太少（比如少于5次）的罕见词，给剩下的词建一个"词-索引"映射表。第三步，遍历语料里的每个词作为中心词，取它前后window_size范围内的所有词作为上下文，生成训练对。

窗口大小是一个需要根据语料规模来定的超参数。窗口小（比如2-3），学到的词向量偏重于近邻的句法关系，比如"我"和"你"会在位置上接近；窗口大（比如5-10），能够捕捉更远的主题相关性，比如"咖啡"能和"早晨"稍微靠近一点。经验值一般在5到10之间，如果语料不大，建议往小设，不然后面负采样会遇到很多正样本重叠的问题，训练时间也会拉长。

3. PyTorch实现的关键环节

3.1 一句话语料到批量训练样本的全过程

假设我们手上有一小段中文语料，是若干条短文本拼起来的。第一步代码要做三件事：分词、计数、映射。

import jieba from collections import Counter import random import torch import torch.nn as nn import torch.optim as optim # 模拟用的小语料，实际项目中可以用几十万行文本替换 corpus = [ "我今天喝了一杯咖啡", "他喜欢喝热茶", "咖啡和茶都是饮品", "我在咖啡店工作", "她买了一杯奶茶", "茶和咖啡提神醒脑", ] # 1. 分词 sentences = [list(jieba.cut(s)) for s in corpus] # 2. 统计词频并过滤低频词 word_count = Counter() for sent in sentences: word_count.update(sent) min_count = 1 # 语料小时设为1，真实项目通常设5 vocab = {word for word, cnt in word_count.items() if cnt >= min_count} word_to_idx = {word: i for i, word in enumerate(vocab)} idx_to_word = {i: word for word, i in word_to_idx.items()} vocab_size = len(vocab)

这里有一个很容易忽略但很关键的点：词表（vocab）一旦构建好，就要固定下来变成常数，训练过程中不能变。如果你边训练边往词表里加词，Embedding层参数矩阵大小对不上，代码直接报错。所以任何对词频的过滤、删词操作，都应该在训练开始前完成。

接下来生成训练样本。你需要遍历每个句子，对每个位置的中心词取窗口内的上下文词，生成（中心词索引，上下文词索引）的组合对。一个细节是处理句子边界：中心词靠句子开头或结尾时，窗口会越界，实际做法是取有效的部分作为上下文，而不是截断或者补零。补零会让模型学到"空字符"这种不该存在的语义，属于负优化。

def generate_training_pairs(sentences, word_to_idx, window_size=2): pairs = [] for sent in sentences: idx_list = [word_to_idx[w] for w in sent if w in word_to_idx] for i, center in enumerate(idx_list): left = max(0, i - window_size) right = min(len(idx_list), i + window_size + 1) for j in range(left, right): if i == j: continue # 自己不能预测自己 context = idx_list[j] pairs.append((center, context)) return pairs training_pairs = generate_training_pairs(sentences, word_to_idx) # 每个元素是 (center_idx, context_idx)

窗口大小在这里设置为2。我经常看到初学者把"窗口"理解成上下文总数，实际上它是中心词两侧各自的距离。window_size=2意味着左边取2个，右边取2个，加上中心词本身总共覆盖长度为5的局部区域。你也可以用动态窗口：每次训练时随机取1到window_size之间的值作为实际窗口大小，这样能让不同距离的词都有机会被采样，一定程度上缓解"距离越远，相关性越弱"的问题。

3.2 Embedding层的本质就是查表

这是整个项目最值得玩味的地方。PyTorch里的nn.Embedding(num_embeddings, embedding_dim)本质上只是一个可训练的参数矩阵，形状为（词表大小，嵌入维度），每次把词的索引输入进去时，它做的事情就是按索引取对应行。

Skip-gram标准实现需要两个独立的Embedding层：一个负责中心词，一个负责上下文词。为什么要两个？因为训练时中心词和上下文词的角色是不对称的。如果只用一个嵌入矩阵，那么"咖啡"作为中心词和作为上下文词的向量是同一份参数，模型在梯度更新时就会在两种角色之间拉扯，最后学不到清晰的语义边界。用两个矩阵就相当于让每个词有两个"工牌"：一张在当"主语"的时候用，一张在当"语境"的时候用。等训练结束，中心词那张工牌就是我们最终要的词向量。

class SkipGramModel(nn.Module): def __init__(self, vocab_size, embedding_dim): super().__init__() self.center_embeddings = nn.Embedding(vocab_size, embedding_dim) self.context_embeddings = nn.Embedding(vocab_size, embedding_dim) def forward(self, center_idx, context_idx, neg_idx): # center_idx: (batch_size,) # context_idx: (batch_size,) # neg_idx: (batch_size, num_neg_samples) center_vec = self.center_embeddings(center_idx) # (batch, embedding_dim) context_vec = self.context_embeddings(context_idx) # (batch, embedding_dim) pos_scores = torch.sum(center_vec * context_vec, dim=1) # (batch,) # 负样本打分 neg_vec = self.context_embeddings(neg_idx) # (batch, num_neg, embedding_dim) center_expand = center_vec.unsqueeze(1) # (batch, 1, embedding_dim) neg_scores = torch.sum(center_expand * neg_vec, dim=2) # (batch, num_neg) return pos_scores, neg_scores

打分方式用的是点积，也就是向量对应元素相乘再求和。两个语义相近的词，向量方向一致，点积值就大。训练的目标就是让正样本（真实上下文）的点积分值尽可能大，负样本（随机挑来的词）的点积分值尽可能小。

损失函数用二分类交叉熵（BCE）对sigmoid后的分数计算。如果你的代码里用了nn.CrossEntropyLoss，多半是把问题又绕回了多分类，那样就没有负采样的意义了。正向分数要看成"正样本的logit"，取sigmoid后希望靠近1；负向分数希望靠近0。最终损失把正、负两部分的损失加起来取平均。

def loss_fn(pos_scores, neg_scores): pos_loss = torch.mean(torch.nn.functional.binary_cross_entropy_with_logits( pos_scores, torch.ones_like(pos_scores))) neg_loss = torch.mean(torch.nn.functional.binary_cross_entropy_with_logits( neg_scores, torch.zeros_like(neg_scores))) return pos_loss + neg_loss

3.3 训练循环：一个batch需要准备三份数据

训练之前，需要为每一对正样本生成负样本索引。最常用的高效方式是预构建采样表：按之前提到的0.75次方概率，把每个词的索引复制若干份到一个大列表里，列表长度（比如几百万）就是采样池。训练的时候，每个batch随机从采样池里抽num_neg个索引当成负样本。这个设计非常实用，因为随机直接用randint在词表里采样的话，会导致高频词被大量选成负样本，低频词几乎没有机会。

def build_sampling_table(word_count, word_to_idx, sample_pool_size=1000000): # 计算修正后的频率 total_count = sum(word_count.values()) freq_table = [] probs = [] for word, idx in word_to_idx.items(): f = word_count[word] / total_count p = f ** 0.75 probs.append(p) sum_p = sum(probs) probs = [p / sum_p for p in probs] pool = [] for word, idx in word_to_idx.items(): cnt = int(probs[idx] * sample_pool_size) pool.extend([idx] * cnt) return pool sampling_table = build_sampling_table(word_count, word_to_idx) def sample_negatives(batch_size, num_neg, sampling_table): neg_idx = [random.choice(sampling_table) for _ in range(batch_size * num_neg)] return torch.LongTensor(neg_idx).view(batch_size, num_neg)

训练循环本身很常规：把训练对随机打乱，分批次取出中心词索引和上下文词索引，同时为每个batch生成对应的负样本索引，送入模型，计算损失，反向传播，更新参数。

batch_size = 32 num_neg = 5 epochs = 30 embedding_dim = 50 model = SkipGramModel(vocab_size, embedding_dim) optimizer = optim.Adam(model.parameters(), lr=0.01) pairs_tensor = torch.LongTensor(training_pairs) num_batches = len(pairs_tensor) // batch_size for epoch in range(epochs): shuffled = pairs_tensor[torch.randperm(len(pairs_tensor))] total_loss = 0.0 for i in range(num_batches): batch = shuffled[i * batch_size: (i + 1) * batch_size] center = batch[:, 0] context = batch[:, 1] neg = sample_negatives(batch_size, num_neg, sampling_table) pos_scores, neg_scores = model(center, context, neg) loss = loss_fn(pos_scores, neg_scores) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() print(f"Epoch {epoch+1}: loss = {total_loss / num_batches:.4f}")

这里有一个训练细节值得说：训练时不需要对中心词做任何"标记"表示它出现在什么位置，因为索引本身就是答案。Embedding层通过查表拿到向量，然后和上下文向量做点积。这整个过程中，真正被优化的核心是参数矩阵里的向量数值，这也是为什么叫"端到端学习"。

3.4 训练完后，怎么把词向量拿出去用

训练结束，模型里有两个Embedding层的权重：self.center_embeddings.weight和self.context_embeddings.weight。取词向量时最常见的选择是只用中心词Embedding矩阵。原因是中心词是每次前向传播的主角，梯度更新更多、信息更充分；上下文矩阵虽然也学到了语义，但角色上更"配角"，单独拿出来用效果通常稍差。

讲究一点的实践是两者加起来再平均一下：

final_embeddings = (model.center_embeddings.weight.data + model.context_embeddings.weight.data) / 2

这个方法在小语料上的效果往往比只取中心词矩阵更稳，因为它相当于把两种角色中学到的信息做了互补。但要注意：如果你的最终目标是把词向量用于下游任务（比如分类、聚类），必须事先固定一种取法，不要在对比不同实验时换来换去，不然你以为是在对比模型效果，其实对比的是取法差异。

拿到词向量矩阵后，最简单的使用方式是算余弦相似度。查某个词最接近的邻居：

def get_top_k(word, k=5): idx = word_to_idx[word] vec = final_embeddings[idx] # 归一化后算点积就是余弦相似度 norm = vec / (torch.norm(vec) + 1e-8) scores = torch.matmul(final_embeddings / (torch.norm(final_embeddings, dim=1, keepdim=True) + 1e-8), norm) top_idx = torch.topk(scores, k=k + 1).indices return [idx_to_word[i.item()] for i in top_idx if i != idx]

这个操作在语义一致的小语料上，你会看到"咖啡"的最相近词大概率是"茶"、"饮品"这类；大语料上还会出现"国王-男性+女性=女王"这种经典的类比推理现象。

4. 训练细节：让模型真正收敛

4.1 学习率与训练轮数怎么定才靠谱

Skip-gram模型不大，参数规模主要取决于词表大小和嵌入维度，普通优化器就能训练得很好。我实测下来，Adam优化器配0.01到0.05的初始学习率是比较稳的区间。学习率太高，正负样本的loss会震荡，词向量的模会迅速膨胀；太低，几百个epoch都学不出明显结构。如果是教学用的小语料，30轮左右就能看到明显的语义聚类；如果语料有几十万行，5-10轮就已经足够了，不要盲目按小语料的经验去跑大量epoch，语言模型的词频分布很偏斜，跑太多轮反而会让高频词过拟合、低频词学不到。

这里有个观察loss的实用技巧：loss值本身没有绝对意义，但相邻两个epoch的loss差值很有意义。如果你发现loss下降得非常快但后来几乎不动，说明模型已经收敛，继续训练只是在微调；如果你发现loss在某个值上反复横条，说明学习率可能太大了，需要回调到0.005甚至0.001。另外，打印loss时用每个epoch的平均值，不要看每个batch的具体值，否则会被单个batch的正负样本随机性带偏。

4.2 负样本数量到底选几个

负采样数量是Skip-gram最核心的超参数之一。num_neg太小，模型没有足够的"反例"来区分，词向量容易退化，表现为几乎所有词在空间里挤成一片；num_neg太大，训练速度降低，而且高频词被过度惩罚，低频词的语义反而被稀释。经验法则：语料规模越大，负样本可以越多。小语料（几万到几十万词）用5个负样本就够；中等语料可以到10到15个；超大语料（数十亿词级别）用20个左右，但一般也不会超过20。

我见过有人为了加速训练把num_neg设为1，在小语料上可以跑，但词向量质量会明显缩水，最直接的表现是相似度计算结果特别"平"，谁都和谁都差不多。负采样看起来只是给模型制造了难例，实际上决定了词向量空间的结构好坏。对照实验里，把num_neg从5调到20，词的近邻几乎从"靠谱"变成"乱序"，所以在类似项目里我宁多用一点负样本图个心里踏实。

4.3 频繁词下采样：高频词不能原样进训练

如果按照前面说的方法直接生成训练对，很快会发现问题：像"的"、"了"、"是"这类功能词出现的频率太高，占比极大，训练时绝大多数样本都在被它们带偏。模型会把大量参数用于区分高频词的细微差异，而真正有语义的名词和动词反而得不到足够的训练机会。

word2vec的解决方案是：在生成训练样本之前，对每个词按它的频率进行随机丢弃。出现频率越高的词，被丢弃的概率越大。具体公式是这样的：

保留概率 = 1 - sqrt(t / f(w))

其中f(w)是词w在语料中的出现频率（用0到1之间的小数），t是一个超参数，默认经验值在1e-3到1e-5之间。直觉理解：如果某词的出现频率比阈值低很多，保留概率接近1，基本不丢；频率特别高的词，保留概率会掉到很低，甚至被丢弃。你可以把t调低一点，让丢弃更激进，训练语料变小但质量更高。这个技巧对训练速度也有明显帮助，因为高频词产生的样本数量极其庞大，丢弃后整体的训练步数会减少很多。

实现上不要在训练循环内部动态丢弃，效率太低。建议在构建训练对之前，先对分词结果逐词做一次过滤：

t = 1e-3 word_freq = {word: cnt / total_count for word, cnt in word_count.items()} keep_prob = {word: (1 - math.sqrt(t / word_freq[word])) if word_freq[word] > t else 1.0 for word in word_count} filtered_sentences = [] for sent in sentences: filtered = [w for w in sent if w in keep_prob and random.random() < keep_prob[w]] filtered_sentences.append(filtered)

有了下采样这一步，高频词的占比不再畸形，低频词在训练样本中的相对比例也会自然提升，词向量的整体质量会有一个肉眼可见的提升。

5. 常见问题与排查实录

5.1 训练完，词向量全是一样的，或者是纯噪声

这个现象几乎每个初跑Skip-gram的人都会遇到一次。先排查数据层面：语料是否太小？如果只有几百个词，模型根本学不出稳定的共现结构，建议至少准备几千到几万条不重复的句子做演示。再看词表处理：词频过滤阈值是否设得过高，比如把出现次数少于10次的词全删了，剩下的都是高频功能词，它们之间本来就缺乏语义区别，向量自然趋同。最后检查损失函数：如果loss在epoch之间完全没有变化，很可能是Embedding层没有被优化器正确维护，一个经典错误是用了model.parameters()却没把Embedding参数传进去，或者手写梯度更新时掩盖了梯度。

另一个常见原因是用错了损失函数。有人觉得负采样最后的概率是一个（batch, 1+num_neg）的分数矩阵，就想当然地用了nn.BCELoss，但BCELoss要求输入是概率值（0到1之间），而模型输出的是logit。这时应该用binary_cross_entropy_with_logits，或者手动在score上做sigmoid后再查BCELoss。这个问题在loss不为nan的时候不容易被察觉，但会发现词向量质量非常差。

5.2 相似度结果反直觉，怎么排查

如果训练顺利完成，"咖啡"的最近邻却是"了"，那多半是采样分布出了问题。高频词被采成负样本的次数太多了，模型学会了"看到高频词就能知道它不是上下文"，高频词自身的向量变得异常紧凑，挤在空间的一个小角落里。排查步骤是：打印一下训练样本中每个词出现的次数，看看是不是高频词真的被过度训练了；检查负采样表是不是忘了用0.75次方修正；确认下采样公式的阈值t设得是否过小，导致高频词几乎没有被丢弃。还有一个不起眼的地方：窗口大小如果太小（比如1），模型学到的更多是相邻共现关系，相似结果会集中在句法近邻，语义类比效果差。

如果想快速验证你的实现有没有大问题，推荐一个小技巧：找一个词频适中的词（不能太大也不能太小），比如语料里的"咖啡"或者"茶"，打印它的top-5最近邻。如果结果总体上在语义上说得通（哪怕有一两个奇怪的），就说明整个链路是通的，后面再针对语料规模和频率处理做优化。

5.3 内存占用和训练速度远低于预期

Skip-gram的训练样本数量可能是语料规模的数十倍，初学者经常一次性把全部训练对生成出来，放在一个Python列表里，然后一次性转成Tensor。这在语料很小时没问题，但语料一大，内存立刻被撑爆。比较稳妥的做法是写一个数据生成器（generator）或者自定义Dataset，逐句生成训练对，边训练边产出，不把所有对都常驻内存。另一个常见瓶颈在负采样：如果每次训练随机从Python列表里逐个choice，几百个epoch会非常慢。应该提前把采样表建成一个numpy数组或Tensor，然后用torch.multinomial或随机索引批量采样。

5.4 训练震荡不收敛，对照速查

排查时我建议先检查最小的一环：训练样本的数据形状。打印出一个batch的center, context, neg的维度，确认center是(batch,)，context是(batch,)，neg是(batch, num_neg)。三个都很小，但很多loss函数里的形状广播错误就是从这里开始的。比如center_expand在unsqueeze之后是(batch, 1, dim)，而neg_vec是(batch, num_neg, dim)，两者在最后两维做了正确的点积，输出才能按维度求和成(batch, num_neg)。这些形状逻辑一错，要么报错，要么看似能跑但损失算的是错误的东西。

6. 最后补充一点我自己的实操体会

把"6.4 PyTorch实现Skip-gram模型"这个项目完整做一遍，最大的收获不是学会了一个模型，而是理解了词向量的来源。嵌入层不是神秘的黑盒，它是一张可以被训练的查表；词的语义相似度不是玄学，是共现统计在向量空间里的投影。我个人的习惯是，做一个新语料时，先用很小的窗口和很小的词表把整个流程跑通，确认代码无误后再放大规模去调优，而不是一上来就直接上大语料折腾半天。遇到不收敛，先把学习率降到0.005，把负样本数固定为5，把window_size固定为2，把其他超参都限制到最保守的区间，修通之后再逐个放开。这样做的好处是，每次出问题都只面对一个变量，排查效率高很多。这个小技巧，不只在Skip-gram这个项目上适用，几乎所有自监督词向量实验都能用上。