
简介基于PyTorch实现的词向量模型聚焦自然语言处理中的词汇表示学习覆盖Word2Vec的CBOW与Skip-gram两种经典结构适合机器学习初学者、NLP开发者作为入门实践或教学参考。词向量技术将词语映射为高维连续向量使语义相近的词在空间中距离更近可缓解传统One-Hot表示的维度灾难与语义鸿沟问题。压缩包采用rar格式大小仅10KB共6个文件其中3个Python脚本分别负责数据处理、模型定义和训练主流程2个pyc为编译缓存1个TXT收录1803个单词的小型英文语料结构紧凑便于解压后快速验证训练流程。目前已有3105人学习。借助PyTorch自动梯度机制读者可跟代码走通分词、词汇表构建、One-Hot输入转换、交叉熵损失计算、反向传播更新以及权重保存等环节并通过调整向量维度、窗口大小、迭代次数等超参数观察词向量变化训练所得词向量还可迁移到文本分类、情感分析等下游任务中作为输入特征或加载预训练权重完成词语相似度检索为后续NLP项目提供可直接复用的表示层。 前段时间做了一个垂直领域的搜索排序优化刚开始图省事直接用公开的预训练中文词向量初始化Embedding层上线后效果不但没提升部分query的相关性还变差了。排查了半天发现问题的根源在于那套词向量是在新闻语料上训练的而我们的语料里大量是产品和行业术语词表覆盖不到模型等于从零开始猜。后来老老实实在PyTorch里自己实现训练了一套领域词向量效果才起来。这件事给我的体会是词向量这个任务看似基础但真正动手做一遍能让你把Embedding、负采样、subsampling这些概念吃透而这些恰恰是后面理解BERT这类预训练模型的地基。这篇文章我就把从数据预处理、模型搭建到训练评测的完整过程整理出来适合两类人一是刚入门NLP、想亲手实现一个词向量模型的开发者二是需要在自己领域数据上训练词向量的工程师。我用的是PyTorch 2.x配合CUDA 12.x的环境其实纯CPU也能跑只是慢一些不影响理解整个流程。1. 词向量到底在解决什么问题one-hot是死路1.1 one-hot的三大硬伤在说词向量之前得先把为什么需要词向量这个根本问题弄清楚。最早的文本表示方式叫one-hot就是把每个词表示成一个长度等于词表大小的向量当前词对应的位置是1其余全是0。比如词表里有苹果、香蕉、手机三个词苹果就表示为[1, 0, 0]。这个表示有三个致命的问题。第一个是维度爆炸。真实语料的词表动辄几十万量级如果每个词都用几十万维的向量表示内存开销完全不可接受。更麻烦的是模型训练时输入矩阵会极度稀疏大部分位置是0纯粹在做无用计算梯度更新也只落在极少数的几行上。第二个是语义鸿沟。one-hot向量之间是完全正交的所以苹果和香蕉的相似度苹果和手机的相似度统统为0。但人的心智里明显不是这样苹果和香蕉好歹都是水果语义上是相近的。one-hot把这种关系全部抹平了下游模型不得不自己去重新学习词与词之间的关系任务难度直线上升。第三个是表达能力受限。one-hot本质上是一个查表操作它不携带任何语义信息。就算维度再高模型也没办法从这种表示里推断出苹果和香蕉在语义上有关联所有语义关系都必须靠下游模型额外拟合这在数据量有限的时候几乎不可行。分布式表示Distributed Representation就是在这个背景下成为主流的。它的核心目标是把每个词表示成一个低维稠密向量比如100维或300维向量的每一维不直接对应某个词的含义而是某种通过训练自动学出来的隐性特征。模型可以通过向量之间的距离比如余弦相似度判断词与词的相近程度这才是词向量能发挥作用的关键。1.2 分布式表示背后的假设上下文决定语义分布式表示之所以行得通背后的语言学假设是一句话一个词的意义由它经常出现的上下文决定。这句话最经典的出处是J.R. Firth的观点——You shall know a word by the company it keeps。翻译过来就是判断一个词的语义看它身边经常出现哪些词就够了。如果两个词经常出现在相似的上下文里它们的语义往往就是相近的。苹果和香蕉都经常出现在吃、水果、甜这些词旁边所以它们的向量就应该靠得近。基于这个假设事情就变得很具体了训练一个模型让它在预测上下文词或目标词的过程中把词与词之间的共现统计信息压进向量里。word2vec是这个思路的经典实现提供了两种训练路线CBOWcontinuous bag-of-words和Skip-gram。CBOW用上下文词去预测中心词Skip-gram反过来用中心词去预测上下文词。我自己的经验是CBOW训练速度快在足够大的语料上效果稳定适合训练大规模词表Skip-gram对低频词的表示更友好如果语料规模不大或者专业术语多它的表现通常更好。两者的差异我用一张表来整理实际选型时可以对着看对比维度CBOWSkip-gram输入上下文词目标词输出目标词上下文词训练速度更快适合大规模语料较慢但对低频词更友好低频词表现一般更好常见用途通用词向量、超大语料专业领域、小规模语料2. 数据预处理词表细节决定模型上限2.1 语料清洗与分词训练词向量的目标是把语料中的共现信息提炼出来所以语料质量直接决定向量质量。清洗上我一般会做这几件事去掉HTML标签和不可见字符统一全半角按标点符号切分句子。如果语料是从网页抓取的还要过滤掉重复的框架性内容比如导航栏、页脚这种反复出现但不含语义信息的文本它们会干扰词频统计把高频词的排名抬得虚高。分词这一步要按场景选择。中文语料我常用的方案是jieba分词配合自定义词典把领域术语先加进去。比如做医疗相关项目就把心脑血管、阿司匹林这类词提前收进自定义词典避免被错误切分。如果后面要做的是一个需要处理复杂命名实体的场景可以考虑BPE或sentencepiece这类子词切分方法不过对标准词向量训练来说jieba基本够用。分词之后把所有句子按空格重新拼接成类似英文的格式方便后续按行读取和处理。2.2 词频截断与词表构建实际语料的词频分布是典型的长尾分布大量词只出现一两次。这些低频词由于上下文信息不足学出来的向量往往是随机的还会显著拖慢训练速度。我通常设置min_count5词频小于5的词一律不进入词表在预处理阶段直接丢弃。词表大小一般控制在5万到20万之间如果你的业务里有很多低频专业术语可以把min_count降到3但不建议再低了。词表构建的同时需要维护word2idx和idx2word两个字典分别负责词到索引和索引到词的映射。另外设置一个unk索引处理训练时未出现过的词。这里有一个容易踩的坑如果训练数据里大量词都被映射成了unk模型学到的其实是未知词这个符号本身正常词的向量质量会被严重拖累。所以词表构建完最好先扫一遍训练数据统计unk占比超过2%就要考虑放宽min_count或者补充词典。2.3 高频下采样别让的和了霸屏还有一个容易被忽视的预处理步骤叫subsampling。语料里的的、了、是这类高频虚词出现频率过高会导致大量训练样本都围绕它们产生稀释了低频词的信号。word2vec提出了一种按概率丢弃词的方式对于词频为f的词以概率P(drop) 1 - sqrt(t / f)丢弃它其中t是人为设定的阈值通常取1e-3或1e-5。词频超过t的词都会有概率被丢弃词频越高丢得越狠。我在实际项目里发现这个操作能让的、了这类词减少50%以上的出现次数低频词的训练权重自然就上来了。这里要特别提醒subsampling是在生成训练样本时做的而不是直接删掉词表里的高频词否则下游应用时词表就缺了OOV问题会更严重。实操中我会在Dataset的__getitem__里对中心词做一个概率判断命中丢弃规则就跳过这条样本这样既干净又省内存。2.4 负采样表3/4次幂的妙处负采样negative sampling是word2vec里另一个关键trick。简单说训练时除了让模型把真实出现的word对判为正例还要随机抽一些没出现过的词作为反例。但负样本不是均匀随机抽的而是按照词频的某个指数分布来采样P(word) f(word)^0.75 / sum(f(其他词)^0.75)这个3/4次幂有什么讲究直观来说它对词频做了一个压缩高频词被抽到的概率下降了一些低频词被抽到的概率上升了一些让负样本分布比原来的词频分布更平滑。这样一个看起来有点玄学的指数在实际训练中表现确实比均匀采样好很多原因在于它避免了采样结果完全被的、了这类高频词主导。在PyTorch里实现负采样表可以先用词频统计构建一个概率数组然后借助torch.multinomial按概率批量取样效率很高。也可以直接用现成的torch.utils.data.WeightedRandomSampler但我更推荐手动构建采样表因为后续要在模型内部批量采样时会用到自己控制起来更顺手。我写一个标准的预处理与样本生成代码结构复现时可以直接参考import math from collections import Counter from torch.utils.data import Dataset class Word2VecDataset(Dataset): def __init__(self, sentences, vocab, window_size5): self.vocab vocab self.window_size window_size # 把所有token铺平方便按位置取中心词 self.tokens [] for s in sentences: for w in s: if w in vocab.word2idx: self.tokens.append(vocab.word2idx[w]) def __len__(self): return len(self.tokens) def __getitem__(self, i): center self.tokens[i] context [] # 滑动窗口取中心词左右各window_size个词 for j in range(i - self.window_size, i self.window_size 1): if j i or j 0 or j len(self.tokens): continue context.append(self.tokens[j]) return center, context这段代码是按token位置来取样本的每次返回一个中心词和它窗口内的上下文词。这样省去了把所有滑动窗口样本预先生成到内存里的步骤语料特别大的时候非常省内存。后续在__getitem__里加subsampling过滤也很顺手。3. 模型搭建PyTorch里的Embedding和输出层3.1 Embedding层到底做了什么很多初学者会把PyTorch的nn.Embedding当成一个普通神经层但理解它的本质会对训练逻辑更清楚它就是一个可训练的查找表。输入是词在词表中的索引输出是这个索引对应的稠密向量。词向量训练的核心就是通过反向传播不断调整查找表里的向量值让语义相近的词在向量空间中距离靠近。有一个常见的误区需要澄清nn.Embedding和nn.Linear在数学上其实是等价的区别在于Embedding的输入是整数索引Linear的输入通常是one-hot向量。但Embedding的实现更省内存、查询速度更快因为它根本不会去构造巨大的one-hot矩阵而是直接通过索引把对应行取出来。所以在超大词表场景下Embedding是唯一现实的选择。3.2 CBOW前向传播与双Embedding的设计对于CBOW输入是窗口内上下文词对应的索引我们分别从Embedding层取出它们的向量取平均得到上下文表示再经过一个输出层映射到词表大小的得分向量上。word2vec原版实现里有一个trick建议新手注意上下文词和中心词分别使用两套Embedding矩阵训练结束后把两套向量加起来作为最终词向量。这样做的动机很直接——同一个词既可能作为上下文词出现也可能作为中心词出现如果共用一套向量这两类梯度更新会互相拉扯分成两套矩阵训练过程会更稳定。PyTorch实现时可以把两套Embedding放进同一个Module里代码结构如下import torch import torch.nn as nn class CBOWModel(nn.Module): def __init__(self, vocab_size, embedding_dim): super().__init__() # 输入侧上下文词的 Embedding self.in_embedding nn.Embedding(vocab_size, embedding_dim) # 输出侧中心词目标词的 Embedding self.out_embedding nn.Embedding(vocab_size, embedding_dim) # 初始化参数均匀分布范围控制在[-1, 1]之间 nn.init.uniform_(self.in_embedding.weight, -1.0, 1.0) nn.init.uniform_(self.out_embedding.weight, -1.0, 1.0) def forward(self, context_ids, center_ids, neg_ids): # context_ids: [batch, window*2]取上下文向量并求平均 ctx_vec self.in_embedding(context_ids).mean(dim1) # 正样本得分向量内积 positive torch.einsum(be,be-b, ctx_vec, self.out_embedding(center_ids)) # 负样本得分[batch, num_neg] neg_vec self.out_embedding(neg_ids).transpose(1, 2) negative torch.bmm(ctx_vec.unsqueeze(1), neg_vec).squeeze(1) return positive, negative这段代码里我用了in_embedding和out_embedding分别承担上下文侧和中心词侧的向量表示。forward返回的positive和negative是为了配合负采样损失函数。注意初始化原版word2vec对向量初始化是有讲究的范围通常在[-1, 1]之间太大会导致训练初期loss爆炸太小又学不动。3.3 损失函数为什么不直接算Softmax很多人直觉上会觉得CBOW最终输出的应该是词表大小的概率分布用交叉熵训练不就行了问题在于词表动辄几十万每次输出都要算一遍对所有词的Softmax计算量完全不可接受训练一轮的时间会拖到天荒地老。word2vec因此用了两种近似方案层级Softmax和负采样。负采样的思想是把一个多分类问题转化成若干个二分类问题给定上下文判断这个词真的出现在上下文附近还是这是个随机抽出来的假词。正样本我们给标签1随机抽的负样本给标签0然后用二分类logistic损失来训。这样每步训练只需要计算一个正样本和几个负样本的得分计算量一下就降下来了。对应的损失计算代码如下def negative_sampling_loss(positive, negative): # positive: [batch]负样本得分: [batch, num_neg] pos_loss -torch.log(torch.sigmoid(positive) 1e-8) neg_loss -torch.log(torch.sigmoid(-negative) 1e-8).sum(dim1) return (pos_loss neg_loss).mean()这段代码其实还原了word2vec里负采样的原始损失正样本得分通过sigmoid后取log我们希望它接近1负样本得分取负后同样走sigmoid我们希望它接近0。1e-8是为了防止log(0)导致数值不稳定。实测中用它训练收敛速度和稳定性都很好。4. 训练细节参数组合和Loss不下降的排查4.1 给一份可复用的超参数组合词向量模型本身不深真正影响效果的往往是几个不起眼的超参数。我给出自己几轮实验下来比较稳的默认组合新手可以直接抄作业参数推荐值说明embedding_dim100~300语料小用100语料大用300再增大收益很小window_size5偏语义任务用小窗口偏句法任务可适当加大negative5~15语料越大可以适当多取负样本subsample_t1e-3~1e-4高频词多选1e-3稀疏语料选1e-4min_count3~5自定义词表按实际词频分布调整epoch5~15语料越大epoch越少防止过拟合batch_size256~1024显存够就开大吞吐量更高这里特别想说明的是维度选择。很多人觉得维度越高越好但词向量训练的样本量有限维度过高会引入大量冗余参数低频词尤其容易过拟合到随机构造上。实际经验是百万级token的小语料100维就够用了几亿token的大规模语料300维是一个普遍能获得稳定收益的上限。4.2 学习率与优化器的选择RNN时代流传下来的习惯是词向量用SGD配合不断衰减的学习率效果稳定。但PyTorch生态里很多人图省事直接用Adam。我的实测结论是Adam在词向量训练上会让损失快速下降到一个看起来不错的水平但向量质量未必比SGD更好。原因是Adam会对每个参数自适应调整学习率在词频差异大的场景下容易在低频词上步子迈得太大造成向量震荡。如果你用Adam建议学习率设置在1e-3量级同时配合torch.optim.lr_scheduler的余弦衰减或指数衰减如果你用SGD初始学习率可以设到0.025训练过程中按epoch逐步衰减这更接近原版word2vec的设置。这一步对向量质量的影响非常明显值得多试几组。4.3 Loss不下降先怀疑数据而不是模型我踩过最典型的坑是模型搭建完全没有问题但loss就是一直震荡不降。排查到最后发现是数据处理时窗口逻辑写错了窗口内包含了中心词本身导致模型预测的目标就是输入本身训练信号是恒定的自然学不到语义关系。遇到loss异常建议按这个顺序排查第一打印一个batch的输入和目标确认窗口内是否包含了中心词自己第二检查词表里unk的比例如果训练数据里大量词都被映射成了unk模型学到的其实是未知词这个符号第三确认subsampling没有把必要的领域词丢光这个可以通过统计训练轮次中词表低频词的数量变化来判断。按这个顺序查一遍多数问题都能定位。4.4 梯度裁剪与训练加速词向量的训练样本是大量小梯度更新偶尔会出现巨大的梯度尖刺尤其是遇到异常长的上下文或极端词频时。加一个max_grad_norm2的梯度裁剪能让训练过程稳定很多。训练速度方面通用做法是用负采样避免大Softmax再用多线程加载数据。实测中DataLoader的num_workers设到4到8对训练吞吐量的提升非常明显。如果你语料达到GB级别训练完一轮可能要几十分钟到几个小时这时数据管道优化比一味加大batch size更有效。5. 评测与落地余弦相似度、类比推理和可视化5.1 余弦相似度找近义词训练完成后的第一个自检手段是用余弦相似度找最近邻词。比如输入苹果看模型返回的top10是不是香蕉、梨、水果这类词。基本代码很简单def most_similar(word, model, vocab, topk10): weight model.in_embedding.weight.detach() center weight[vocab.word2idx[word]] sims torch.cosine_similarity(weight, center.unsqueeze(0), dim1) top_idx sims.argsort(descendingTrue)[:topk] return [(vocab.idx2word[i.item()], sims[i].item()) for i in top_idx if i.item() ! vocab.word2idx[word]]这里用in_embedding.weight做检索是因为它代表词的语义表示。如果你的最终词向量是两套Embedding相加那在评测阶段也完全可以用叠加后的向量做检索效果往往更好因为信息更完整。5.2 word analogy词向量的关系推理测试除了相似词word analogy任务也是经典评测方式比如国王王后 男人女人这类关系推理。做法是计算vec(王后) - vec(国王) vec(男人)然后从词表里找与结果向量余弦相似度最高的词期望答案是女人。这个评测考查的不只是词本身的相似性还有词与词之间语义关系的保持能力。实测下来CBOW模型在句法类比上有不错的表现比如run:ran swim:swam这类形态变化但在语义类比上相对弱一些比如苹果水果 汽车交通工具。这也是分布式表示的固有局限它捕捉到的更多是统计共现而非真正的世界知识。5.3 可视化验证与下游任务对照把高维向量降到二维常用t-SNE。t-SNE在向量质量检查里特别好使你可以把几组语义相关的词分别取出来降维后观察是否聚在一起。注意t-SNE里perplexity对图形影响很大不要一上来就盲调一般选30到50迭代次数给足再结合余弦相似度数字综合判断避免被可视化的随机性误导。词向量本身更像是半成品最终价值体现在能被下游模型用得多好。我自己做分类、NER、搜索排序这类任务时会做一个严格对照实验一组用随机初始化的Embedding一组用训练好的领域词向量初始化后续模型结构完全一致。如果领域词向量那一组的指标有稳定提升说明词向量确实把领域语义信息编码进去了。根据我的经验这种提升在小规模标注数据上尤为明显因为模型不需要从零开始学词的语义关系。6. 踩过的坑和一些经验总结6.1 什么时候值得自己训练词向量现在NLP生态里BERT、RoBERTa这类预训练模型已经把通用的词和上下文语义关系学得很好了很多场景直接加载中文预训练模型就能达到不错的效果。那自己训练词向量还有没有意义我的判断是当你的语料有很强的领域特性比如法律、医学、工业界的专业文本公开语料训练的通用词向量覆盖不到这些专业词汇即便覆盖了语义理解也是泛化的这时候自己训练一套领域词向量对下游任务往往是实打实的提升。从学习角度来说亲手实现一遍词向量能把Embedding机制、负采样、subsampling这些概念彻底吃透。这些概念不是孤立的它们构成了理解现代预训练模型的基础。6.2 OOV问题与向量后处理词向量模型只能表示训练时词表里见过的词遇到新词就是OOV只能给一个统一的unk向量。缓解方案有几种一是训练时保留更多低频词扩大词表覆盖二是用字符级或子词级别的表示类似fastText的做法用子词信息拼出OOV词的向量三是在下游任务里遇到OOV时用词形相似或上下文向量来近似。最通用的做法还是在预处理阶段就把词表做得大一些配合扎实的清洗流程比事后补救更靠谱。训练出来的向量建议做一个归一化把每个词向量除以它的L2范数让所有向量落在单位球面上。归一化之后余弦相似度就等价于点积数值上更稳定下游任务里做内积计算也不用担心尺度漂移。这个动作很简单但对后续使用的稳定性帮助很明显我在项目里基本都是默认执行。6.3 从静态词向量到预训练模型的连接理解词向量的训练过程再看现代预训练模型会通透很多。BERT这类模型本质上也是在学词的表示区别在于它学的是上下文相关的词表示——同一个词在不同句子里有不同向量。而word2vec这类模型学的是静态词向量一个词只有一个固定向量。两者不是替代关系而是演进关系。很多预训练模型的下游任务微调核心思想依然和词向量有千丝万缕的联系。所以把这套基础方法学好再看transformer时代的模型你会更容易理解它到底在解决词向量的什么问题而不是把每个新模型当成一个全新概念去背。这也是我为什么一直建议刚入门NLP的朋友哪怕论文里已经很少直接讲word2vec了也要自己动手训练一遍。回顾我自己这次从需求到落地的完整过程最大的收获不是最后效果提升了多少而是当你亲自动手把数据、模型、训练、评测每个环节全部打通之后再看模型输出的数字心里是踏实的。词向量的实现代码量不大但它像一道很好的引子把NLP里最核心的几个抽象逐个击破。如果这篇文章能帮你省下一些查资料、试错的时间那就值了。本文还有配套的精品资源点击获取