简介:面向自然语言处理入门者与Java开发者的Word2Vec文本向量化实现工程,采用Eclipse项目结构搭建,可直接导入IDE编译运行,用于把语料训练成低维词向量,并支撑相似度计算、文本分类、情感分析等下游任务。压缩包共26个文件、约38KB,以10个java源码与12个class字节码为主体:源码对应CBOW与Skip-gram两套训练逻辑,以及分词、去停用词等预处理环节;class文件则省去自行编译的麻烦。另有.project、.classpath与prefs配置,保证工程在Eclipse中顺利导入。已有1835人浏览学习。读者可从中理清窗口大小、迭代次数、最小词频等参数对训练效果的影响,理解负采样与Hierarchical Softmax如何加速收敛,并参考模型的保存与加载方式,将训练结果接入自己的Java应用。整体体量轻巧,适合作为词向量原理到代码落地的对照样本。
1. 一个词向量训练器为什么值得用 Java 重写一遍
线上服务是清一色的 Java,做搜索联想和召回排序时又绕不开词向量。很多团队的做法是 Python 侧训练、导成文件、Java 侧读进来算相似度,链路上多一个运行环境、多一份模型同步的活。这个纯 Java 版本的 word2vec 把训练和推理塞进同一个 JVM:读语料、建词表、建 Huffman 树、跑 CBOW 或 Skip-gram、落盘向量,一个 main 方法全干完,打包成 jar 就能跑。适合两类人——一类是想在不引入 Python 依赖的前提下拿到自己领域词向量的后端开发,另一类是想把 expTable、负采样、层次 Softmax 这些实现细节看明白的人。下面按我实际拆代码的顺序展开。
2. 项目骨架与运行链路:从 main 到词向量落盘的五个环节
这个实现的代码量不大,但五脏俱全。一开始我以为它只是个教学 demo,翻完源码后发现该有的都有:词频统计、降采样、动态窗口、Huffman 树、负采样表、参数自适应衰减、多线程训练、模型存取,原版 C 实现里的关键分支基本都覆盖全了。所以别被“Java 版”这三个字误导成玩具,它是可以拿去做小规模领域词表的。
2.1 目录结构:八个包,各管一段
我把它 clone 下来第一件事是看包结构,因为它决定了你要改哪一块时不用满仓库找。
src/main/java/ ├── cli/ # 命令行入口,解析 -train -output -cbow 这些参数 ├── corpus/ # 语料读取、分词器接口、句子切分 ├── vocab/ # 词频统计、Vocab 对象、Huffman 树构建 ├── sample/ # 降采样(subsampling)与负采样表 ├── model/ # syn0/syn1/syn1neg 矩阵、初始化 ├── train/ # CBOW 与 Skip-gram 的训练主循环 ├── io/ # 模型落盘与加载、向量导出为文本 └── util/ # expTable、随机数、相似度计算、TopK拆开看职责其实很清晰:corpus只负责把一行行文本变成int[]的句子,词在词表里的下标由vocab分配;model只管三块矩阵的内存和初始化;train是整个项目的核心,所有梯度更新都在这里;io决定你训练完的东西能不能被别的服务读走。我一般改参数会直接去train包里翻,改输出格式去io,改分词去corpus,这三个包基本覆盖了 90% 的二次开发需求。
有一个设计取舍值得说:它没有把词表做成全局单例,而是作为Trainer的成员传递。看上去啰嗦,但好处是你可以在同一个进程里先后训练两个不同语料的模型而不互相污染,做 A/B 对比时很省事。
2.2 从 mvn package 到第一次出向量
依赖很少,基本是 JDK 8+ 加一个编译插件,没有再拖日志框架和集合库。构建和跑第一条命令是这样:
# 打包,跳过测试省点时间 mvn -q clean package -DskipTests # 最小可跑命令:CBOW + 层次 Softmax,100 维,跑 5 轮 java -jar target/word2vec.jar \ -train corpus.txt \ -output vec.bin \ -cbow 1 \ -hs 1 \ -size 100 \ -window 5 \ -negative 0 \ -sample 1e-4 \ -minCount 5 \ -alpha 0.025 \ -epochs 5 \ -threads 4几个命令参数的含义得说清楚,不然容易照着抄但不知道在调什么。-cbow 1选 CBOW,写 0 就是 Skip-gram;-hs 1打开层次 Softmax,-negative就要给 0,两者是互斥的,同时开会浪费一半算力;-minCount是词频截断阈值,低于它的词直接不进词表,中文语料上这个值给 5 比较稳;-threads是训练线程数,注意它跟后面的模型合并逻辑有关,不是随便调大就快。
如果只是想先看看效果,用不着一次跑全量语料。我习惯先切 20 万行出来跑 1 轮,确认 loss 曲线(或者说进度输出)在往下走、输出的向量不是 NaN,再上全量。这一步能省掉大量“跑了四小时发现参数错了”的血泪时间。
2.3 参数表:真正影响收敛的那几个
原版 word2vec 的参数有十几个,但真正会显著改变结果的其实就那么几个。我按踩坑频率排了个表:
| 参数 | 常见取值 | 调错了会怎样 |
|---|---|---|
size | 100 ~ 300 | 太小语义区分度不够,太大在小语料上直接过拟合 |
window | 5(中文)/ 5~10(英文) | 太大把无关词拉进上下文,语义变糊 |
negative | 5 ~ 15 | 太少梯度噪声大,太多训练时间线性上涨 |
sample | 1e-3 ~ 1e-5 | 不设或太大,高频虚词霸占训练预算 |
alpha | 0.025(CBOW)/ 0.05(SG) | 太大直接发散成 NaN,太小几轮下来没动静 |
minCount | 5 | 给 1 会让词表膨胀到几十万,内存和速度双崩 |
epochs | 5 ~ 15 | 太少欠拟合,太多在领域语料上会记住噪声 |
sample这个参数最容易被忽略。它控制的是降采样阈值,语料里“的、了、是”这类词出现频次极高,携带的信息量却很低,不降采样的话模型大部分梯度都花在这些词上,最后相似度一查,返回的全是虚词。这个值越小,丢弃越激进,中文语料上我一般从1e-4开始试。
3. 核心数据结构:词表、Huffman 树与负采样表怎么建
训练开始前的准备工作,决定了后面所有计算的地基。这一章的三块内容——词表、Huffman 树、负采样表——都是“建一次用到底”的结构,一旦建错,后面训练多久都是白费。我见过不少自己改崩的情况,问题八成出在这三步。
3.1 词频统计与 Vocab 对象
词表的基本单元就是一个类,字段不多但每个都有明确用途:
public final class Vocab { public String word; // 词本身 public long count; // 语料中的出现次数 public int[] point; // Huffman 路径上经过的内部节点编号 public byte[] code; // 每一层走左(0)还是走右(1) public int index; // 在 syn0 矩阵里对应的行号 }统计过程是标准的 Map 累加,但排序规则有个坑必须点出来:
Map<String, Long> counter = new HashMap<>(1 << 20); try (BufferedReader br = Files.newBufferedReader(path, StandardCharsets.UTF_8)) { String line; while ((line = br.readLine()) != null) { for (String w : tokenizer.split(line)) { if (w.isEmpty()) continue; counter.merge(w, 1L, Long::sum); // 词频累加 } } } List<Map.Entry<String, Long>> entries = new ArrayList<>(counter.entrySet()); // 必须按频次降序,频次相同再按字典序,保证多次训练可复现 entries.sort((a, b) -> { int cmp = Long.compare(b.getValue(), a.getValue()); return cmp != 0 ? cmp : a.getKey().compareTo(b.getKey()); });为什么强调降序?因为下面建 Huffman 树用的是双队列法,它假设叶子队列是从频次最低的一端开始取的。如果词表是乱序的,树还是能建出来,但 Huffman 的“高频词编码短”这个性质就没了,训练速度会莫名其妙地慢一截,而且没有任何报错。这个坑特别阴,代码跑得通、日志也正常,就是慢。
另外index的分配必须紧跟排序之后,从 0 开始连续递增。后面syn0是按index * size直接偏移取行的,一旦出现空洞,取出来的就是另一行向量,症状是训练不收敛但看不出哪里错。
3.2 用双队列法在 O(n) 内建 Huffman 树
建树这一步是整个项目里最“算法”的部分。朴素做法是每次从堆里取两个最小值,复杂度 O(n log n),词表几十万的时候开销不小。原版用的是双队列法:叶子节点按频次降序排好,从尾部取就是升序;新建的内部节点天然是递增的,两个队列各自有序,取最小值只需要比较两个队头,整体 O(n)。
private static final int MAX_CODE_LENGTH = 40; private void buildHuffman() { int n = words.size(); long[] count = new long[2 * n + 1]; int[] parent = new int[2 * n + 1]; int[] binary = new int[2 * n + 1]; Arrays.fill(parent, -1); for (int i = 0; i < n; i++) count[i] = words.get(i).count; // 内部节点队列初始为空,用 MAX_VALUE 占位,保证首次比较一定选中叶子 for (int i = n; i <= 2 * n; i++) count[i] = Long.MAX_VALUE; int pos1 = n - 1; // 叶子队列:从频次最小的尾部开始往前 int pos2 = n; // 内部节点队列:从最新生成的那个开始 for (int a = 0; a < n - 1; a++) { // 取两个最小值,谁小从谁那儿取 int min1 = (pos1 >= 0 && count[pos1] < count[pos2]) ? pos1-- : pos2++; int min2 = (pos1 >= 0 && count[pos1] < count[pos2]) ? pos1-- : pos2++; count[n + a] = count[min1] + count[min2]; parent[min1] = n + a; // 合并后的新节点成为父亲 parent[min2] = n + a; binary[min2] = 1; // 右孩子记 1,左孩子默认 0 } // 从叶子往上走,把路径反转成「根 → 叶」的顺序 for (int a = 0; a < n; a++) { int b = a, depth = 0; int[] pt = new int[MAX_CODE_LENGTH]; byte[] cd = new byte[MAX_CODE_LENGTH]; while (parent[b] != -1 && depth < MAX_CODE_LENGTH) { pt[depth] = parent[b]; cd[depth] = (byte) binary[b]; b = parent[b]; depth++; } Vocab v = words.get(a); v.point = new int[depth]; v.code = new byte[depth]; for (int i = 0; i < depth; i++) { v.point[i] = pt[depth - 1 - i]; v.code[i] = cd[depth - 1 - i]; } } }几个参数值得单独说明。count数组长度开到2n+1,是因为 n 个叶子加上 n-1 个内部节点最多用到下标2n-2,多留两个位置当哨兵,避免边界判断写错。MAX_CODE_LENGTH = 40是个硬上限,理论上极端不均衡的分布能让树深超过这个值,代码里截断而不是报错,所以词表分布特别畸形时要注意——不过正常语料基本到不了。
binary[min2] = 1这一行的意思是:合并时把较大的那个(min2)放到右子树。左 0 右 1 是约定,反过来说只要编码和训练时的判断一致就行,但如果你的代码里编码和更新逻辑方向搞反了,模型照样能收敛,只是收敛得慢,属于典型的“能跑但不对”。
3.3 负采样表:0.75 次方与 1e8 桶
走负采样分支时不需要建 Huffman 树,但要建一张采样表。核心思想是把每个词的频次取 0.75 次方后归一化,这个指数是为了压低高频词被采中的概率,同时抬高低频词,经验值,原版就是这么干的。
private int[] buildNegativeTable() { double total = 0; for (Vocab v : words) total += Math.pow(v.count, 0.75); // 1e8 个 int 就是 400MB,内存紧就降到 1e7,精度损失可以接受 int[] table = new int[NEGATIVE_TABLE_SIZE]; int i = 0; double cumulative = Math.pow(words.get(0).count, 0.75) / total; for (int a = 0; a < NEGATIVE_TABLE_SIZE; a++) { table[a] = i; if (a / (double) NEGATIVE_TABLE_SIZE > cumulative && i < words.size() - 1) { cumulative += Math.pow(words.get(++i).count, 0.75) / total; } } return table; }这段看着简单,其实是个“用空间换时间”的查表法:把归一化后的概率分布展开成一张定长数组,采样时只要随机取一个下标就行,复杂度 O(1)。NEGATIVE_TABLE_SIZE越大越精确,但 1e8 的 int 数组直接吃 400MB,容器里跑很容易触发 OOM。我一般给 1e7,实测对最终向量的质量影响在可接受范围内。
还有一个细节:采样采到下标 0 的时要重采或者顺移一位,因为 0 号位通常留给换行符或者<unk>这类占位词,让它当负样本没意义。
4. 训练主循环:CBOW 与 Skip-gram 的梯度更新逐行拆
到这一章就是真正干活的地方了。整个训练循环的骨架其实很短——遍历句子、取上下文、算一次前向、反向更新——但每一步都有几个必须对齐的细节,错一个就是“能跑但效果差”的玄学问题。
4.1 上下文窗口、动态缩窗与子采样
原版有个不太起眼但很有效的技巧:窗口大小不是固定的,每个词随机取1 ~ window之间的一个值。这样近距离的词总是能进上下文,远距离的词只有部分机会进来,相当于给近距离词更高权重。实现上就是一次取模:
int window = (int) (this.window + nextRandom() % this.window); // 实际窗口在 [window, 2*window-1]注意这里是this.window + random % this.window,也就是窗口在window到2*window-1之间浮动,而不是缩小到1 ~ window。不同实现在这里的处理不一样,跟原版对齐就行,别自己发明。
子采样更关键。语料里高频词占比极高,不处理的话模型会把大量梯度花在学“的”和“了”上。判定逻辑是一个基于词频的概率公式:
/** 返回 true 表示这个词本轮被丢弃 */ private boolean shouldDiscard(Vocab v, long trainWordCount) { if (sample <= 0) return false; // sample=0 表示关闭降采样 double ran = (Math.sqrt(v.count / (sample * trainWordCount)) + 1) * (sample * trainWordCount) / v.count; return ran < (nextRandom & 0xFFFF) / 65536.0; }ran算出来的是“保留概率”,低于随机数就丢掉。代入一组数感受一下:假设语料 1 亿词,sample = 1e-4,某词出现 100 万次,那么sample * trainWordCount = 10000,ran = (sqrt(10000/1000000) + 1) * 1000000/10000…这组数算下来保留概率很低,符合预期——出现 1% 的词已经算高频了,该压。而对于只出现几十次的词,ran会大于 1,永远不会被丢。
被丢掉的词在句子里标成-1,后面取上下文时跳过。这个约定很重要,如果忘了判< 0,就会拿-1 * size去访问数组,直接抛数组越界。
学习率的衰减也在这个环节里。它不是每个词都降,而是每处理完一万个词调一次:
float progress = (float) (wordCountActual / (double) (trainWordCount * epochs + 1)); alpha = startingAlpha * (1 - progress); if (alpha < startingAlpha * 0.0001f) alpha = startingAlpha * 0.0001f;初始alpha线性降到万分之一后就不再变。这个下限是必要的,否则最后几轮几乎没有更新,等于白跑。
4.2 层次 Softmax 的一次前向与反向
CBOW 的思路是:把上下文所有词的向量取平均,用这个平均值去预测中心词。预测用的是 Huffman 树上的路径——从根走到目标词,每经过一个内部节点就做一次二分类。
private void trainCbow(int[] sentence, int len, float alpha) { for (int senPos = 0; senPos < len; senPos++) { int word = sentence[senPos]; if (word < 0) continue; // 被降采样丢掉 // 1. 上下文向量求平均,得到隐藏层 neu1 Arrays.fill(neu1, 0f); int cw = 0; int b = (int) (nextRandom() % window); for (int a = b; a < window * 2 + 1 - b; a++) { if (a == window) continue; // 跳过中心词自己 int c = sentence[senPos - window + a]; if (c < 0) continue; for (int e = 0; e < dim; e++) neu1[e] += syn0[c * dim + e]; cw++; } if (cw == 0) continue; for (int e = 0; e < dim; e++) neu1[e] /= cw; // 2. 沿 Huffman 路径逐节点做二分类并累积梯度 Arrays.fill(neu1e, 0f); Vocab target = words.get(word); for (int d = 0; d < target.code.length; d++) { int node = target.point[d]; float f = sigmoid(dot(neu1, syn1, node * dim)); float g = (1 - target.code[d] - f) * alpha; // 关键:残差 × 学习率 for (int e = 0; e < dim; e++) { neu1e[e] += g * syn1[node * dim + e]; // 梯度回传给隐藏层 syn1[node * dim + e] += g * neu1[e]; // 更新输出层节点向量 } } // 3. 把累积梯度原样分发给每个上下文词 for (int a = b; a < window * 2 + 1 - b; a++) { if (a == window) continue; int c = sentence[senPos - window + a]; if (c < 0) continue; for (int e = 0; e < dim; e++) syn0[c * dim + e] += neu1e[e]; } } }g = (1 - code - f) * alpha这一行是整段的核心。code是 0 或 1,f是 sigmoid 输出,也就是模型认为“该走右分支”的概率。当code = 1时,g = -f * alpha,梯度方向是降低走右的概率还是提高?——因为f是走右的概率,目标是让它趋于 1,所以f越小g的绝对值越大,更新力度越大,方向是往“更像右”推。当code = 0时,g = (1 - f) * alpha,目标和实际正好互补。这一行写反了模型照样能跑,只是收敛慢一倍,不容易发现。
sigmoid用查表实现,不用每次算 exp:
private static final int EXP_TABLE_SIZE = 1000; private static final float MAX_EXP = 6f; static { expTable = new float[EXP_TABLE_SIZE]; for (int i = 0; i < EXP_TABLE_SIZE; i++) { expTable[i] = (float) (Math.exp((i / (double) EXP_TABLE_SIZE * 2 - 1) * MAX_EXP)); expTable[i] = expTable[i] / (expTable[i] + 1); // 归一化成 sigmoid 值 } } static float sigmoid(float x) { if (x <= -MAX_EXP) return 0f; if (x >= MAX_EXP) return 1f; return expTable[(int) ((x + MAX_EXP) * (EXP_TABLE_SIZE / MAX_EXP / 2))]; }EXP_TABLE_SIZE / MAX_EXP / 2是整数除法,等于 83,这是原版就有的写法,不要改成浮点。查表把 exp 运算换成了数组访问,在几十亿次更新的量级上差别很明显。
4.3 负采样分支与 Skip-gram 的写法差异
负采样不走树,直接对每个目标词做一次正样本二分类,再随机采negative个负样本各做一次。正负样本的区别只在 label:正样本 label 是 1,负样本是 0。
for (int d = 0; d < negative + 1; d++) { int target; byte label; if (d == 0) { target = word; // 第一个永远是正样本:中心词自己 label = 1; } else { nextRandom = nextRandom * 25214903917L + 11; // 线性同余,和原版一致 target = negTable[(int) ((nextRandom >>> 16) % negTable.length)]; if (target == 0) target = (int) (nextRandom % (words.size() - 1)) + 1; label = 0; } float f = sigmoid(dot(neu1, syn1neg, target * dim)); float g = (label - f) * alpha; for (int e = 0; e < dim; e++) { neu1e[e] += g * syn1neg[target * dim + e]; syn1neg[target * dim + e] += g * neu1[e]; } }nextRandom * 25214903917L + 11这个常数是 Java 里java.util.Random用的同一个 LCG 乘子,用它是为了和原版 C 实现的随机序列尽量接近,方便对比结果。>>> 16是无符号右移,取高位,因为低位随机性差。
Skip-gram 和 CBOW 的差别在隐藏层的来源。CBOW 是上下文取平均去预测中心词,Skip-gram 反过来,用中心词的向量去预测每一个上下文词。代码上表现为:外层循环不变,内层对每个上下文词各跑一次完整的预测和更新,neu1直接就是syn0[word * dim]这一行,不做平均。
// Skip-gram:中心词向量直接当隐藏层 System.arraycopy(syn0, word * dim, neu1, 0, dim); for (int a = b; a < window * 2 + 1 - b; a++) { if (a == window) continue; int c = sentence[senPos - window + a]; if (c < 0) continue; Arrays.fill(neu1e, 0f); // ... 这里对 c 跑一次 HS 或负采样更新,累积到 neu1e for (int e = 0; e < dim; e++) syn0[word * dim + e] += neu1e[e]; // 只回传给中心词 }Skip-gram 在小语料上对低频词的表现通常比 CBOW 好,代价是训练时间大概翻几倍。领域语料只有几百万词的时候,我一般直接上 Skip-gram。
5. 避坑与排查:向量不像话时先查这五处
代码能跑通和模型有用是两件事。下面这几条都是我在实际调试里撞过的,症状、原因、处理方式一一对应,遇到了按顺序排查基本能定位。
第一条:训练日志正常,但向量里全是 NaN。
现象是训练过程不报错,跑完一读模型,所有分量都是NaN。原因通常是学习率给大了,或者syn0初始化用了random()直接当初始值。正确的初始化是(random - 0.5) / dim,把初值压在一个很小的范围里,否则第一次梯度更新就能把数值推到溢出。另外检查一下alpha,CBOW 给 0.05 以上基本必炸。
第二条:查相似度,返回的全是“的、了、是”。
这是降采样没开或者minCount给太低。语料里高频虚词占比太大,模型把预算都花在它们身上了。把sample设成1e-4、minCount设成 5 再重训,通常立竿见影。顺带检查一下分词器,中文没分词直接按字切,也会出现类似症状——每个字都是高频。
第三条:训练速度比预期慢好几倍,CPU 却跑不满。
大概率是 Huffman 树建歪了,树退化成近似链表,每个词的平均路径长度接近词表大小。根因是词表没按频次降序排,双队列法从错误的端点取叶子。还有个可能是MAX_CODE_LENGTH截断生效了,去point数组长度里看一眼,如果大量词的长度等于 40,就是树结构出问题了。
第四条:同样的参数跑两次,结果的相似度排序差很多。
多线程训练的固有现象。多个线程各自更新syn0,没有锁,写冲突导致结果不可复现。原版就是这么设计的,用一点点精度换速度。要可复现就-threads 1,要速度就接受差异。另外注意nextRandom如果是共享的,还会引入跨线程的随机序列竞争,改成每个线程独立种子能缓解。
第五条:跑大语料直接 OOM,堆内存加了几次都不够。
先算一下账:syn0是vocabSize * dim * 4字节,词表 50 万、维度 200 就是 400MB;syn1或syn1neg同量级;负采样表 1e8 个 int 又是 400MB。三个加起来轻松过 1G。处理方式很直接:负采样表降到 1e7,minCount提到 10 砍词表,维度从 300 降到 100。别一上来就-Xmx8g,先把该省的省掉。
6. 二次开发:把训练好的向量接进自己的 Java 服务
训练只是第一步,真正要落地的是让业务代码能方便地查词向量。这块我踩过的坑比训练环节还多,因为训练错了你能看出来,接进服务里错了往往要等线上效果波动才发现。
6.1 落盘格式与模型加载
能用二进制就别用文本。文本格式看着方便排查,但 50 万词 × 200 维大概 400MB,加载时解析浮点字符串能卡住启动线程好几秒。落盘用DataOutputStream,写词表大小、维度、然后每个词的长度前缀字符串加浮点数组:
public static void save(Path path, List<String> words, float[] vectors, int dim) throws IOException { try (DataOutputStream out = new DataOutputStream( new BufferedOutputStream(Files.newOutputStream(path)))) { out.writeInt(words.size()); out.writeInt(dim); for (int i = 0; i < words.size(); i++) { out.writeUTF(words.get(i)); // 自带长度前缀 for (int d = 0; d < dim; d++) { out.writeFloat(vectors[i * dim + d]); } } } }加载端对应着读,读到float[]之后立刻做一件事:把每一行归一化成单位向量。这样后面算余弦相似度就退化成一次点积,省掉每次除模长的开销。
// 预归一化,之后相似度 = 点积 for (int i = 0; i < vocabSize; i++) { double norm = 0; for (int d = 0; d < dim; d++) norm += vectors[i * dim + d] * vectors[i * dim + d]; norm = Math.sqrt(norm); if (norm == 0) continue; // 全零向量跳过,防御性写法 for (int d = 0; d < dim; d++) vectors[i * dim + d] /= (float) norm; }这个预归一化是必须的。训练出来的向量模长差异很大,低频词的模长往往明显偏小,不归一化直接算点积,结果会被模长大的高频词主导。
6.2 线程安全与在线相似度计算
模型加载完之后是只读的,所以线程安全很好保证:所有字段final,不提供任何修改方法,在 Spring 里用@PostConstruct加载成单例就行。真正要小心的是查询时的临时数组——别把float[]缓存成成员变量。
暴力遍历在词表 10 万以内完全够用,一次 TopK 查询大概几毫秒。写起来也简单:
public List<Hit> mostSimilar(String word, int topK) { Integer idx = index.get(word); if (idx == null) return Collections.emptyList(); // OOV 直接返回空,别抛异常 int base = idx * dim; PriorityQueue<Hit> heap = new PriorityQueue<>(topK, Comparator.comparingDouble(h -> h.score)); for (int i = 0; i < words.size(); i++) { if (i == idx) continue; float dot = 0; int off = i * dim; for (int d = 0; d < dim; d++) dot += vectors[base + d] * vectors[off + d]; if (heap.size() < topK) heap.offer(new Hit(words.get(i), dot)); else if (dot > heap.peek().score) { heap.poll(); heap.offer(new Hit(words.get(i), dot)); } } List<Hit> res = new ArrayList<>(heap); res.sort((a, b) -> Float.compare(b.score, a.score)); return res; }用固定大小的最小堆而不是全排序,10 万词表下能省掉一次 O(n log n)。如果词表涨到百万级,就该考虑上近似索引了,但在那之前没必要提前优化。
6.3 领域词表:继续训练还是重训
这是二次开发里最容易被问到的决策。我的结论是:换领域就重训,别在原模型上继续训练。
原因在 Huffman 树。层次 Softmax 的每个内部节点是绑在特定词表上的,你在通用语料上训出来的syn1和新的领域词表对不上号;就算走负采样,syn0的行顺序也变了。硬要继续训练,只能冻结词表、只更新syn0,但那样新词根本进不来,遇到领域专有名词全是 OOV。
真要做增量,正确姿势是:合并两份语料的词频统计,重建词表、重建 Huffman 树,然后用旧模型的向量去初始化新模型的syn0(同名词直接拷贝,新词随机初始化),再跑少量轮次做微调。这个流程代码量不大,但要写对词表映射,我一般是单独写一个WarmStart工具类,不塞进主训练流程里。
验证方式也别偷懒。准备一个小的人工标注集,二十来组“这俩词该相似”和“这俩词不该相似”,每次改完参数跑一遍,看排序有没有退化。这套冒烟测试我一开始嫌麻烦没做,后来有一次调window把一个已经上线的召回效果调崩了,排查了两个晚上才定位到是窗口参数——从那以后我每次动训练参数前,都强制先跑一遍这个二十组的小集合,跑通再上全量。希望这些能帮你少走点弯路。
本文还有配套的精品资源,点击获取