每日算法 — 使用java实现Word2Vec：Skip-Gram模型与负采样优化

引言

Word2Vec是2013年Google提出的经典词嵌入模型，能够将词语映射到低维连续向量空间，使得语义相近的词在向量空间中距离相近。本文用Java实现Word2Vec的核心训练算法——Skip-Gram模型结合负采样（Negative Sampling），帮助读者理解词向量从无到有的训练过程。

核心原理

Skip-Gram模型

Skip-Gram的目标是根据中心词预测其上下文词。给定中心词 w_t，模型预测窗口内的上下文词 w_{t-c}, …, w_{t+c}（不包括自身）。

负采样优化

原始Softmax分母需要遍历整个词表，计算代价极高。负采样通过将多分类问题转化为二分类问题来加速训练：

• 正样本：中心词与真实上下文词的组合 (w_t, w_c)
• 负样本：中心词与从噪声分布中随机采样的词的组合 (w_t, w_n)

目标函数变为最大化正样本的sigmoid得分，同时最小化负样本的sigmoid得分。

目标函数

J = log σ(v’{w_O} · v{w_I}) + Σ_{i=1}^{k} E_{w_i ~ P_n(w)} [log σ(-v’{w_i} · v{w_I})]

其中 v_{w_I} 是中心词向量，v’_{w_O} 是上下文词向量，k 是负采样数量。

Java实现

项目结构

src/
├── Word2Vec.java         // 核心训练类
├── HuffmanTree.java      // 霍夫曼树（用于层次Softmax，本文使用负采样）
├── Word2VecDemo.java     // 演示主类
└── utils/
    └── DataUtils.java    // 数据预处理工具

核心训练类

import java.util.*;

/**
 * Word2Vec Skip-Gram + 负采样实现
 */
public class Word2Vec {
    private final int vectorSize;       // 词向量维度
    private final int windowSize;       // 上下文窗口大小
    private final int negativeSamples;  // 负采样数量
    private final double learningRate;  // 学习率
    private final int minCount;         // 最小词频阈值

    // 词表
    private Map<String, Integer> wordToId;
    private Map<Integer, String> idToWord;
    private int vocabSize;

    // 词向量: syn0是中心词向量, syn1Neg是负采样上下文向量
    private double[][] syn0;
    private double[][] syn1Neg;

    // 词频与负采样噪声分布
    private double[] noiseDistribution;
    private int[] wordFreq;
    private int totalWords;

    private final Random rand;

    public Word2Vec(int vectorSize, int windowSize, int negativeSamples, 
                    double learningRate, int minCount) {
        this.vectorSize = vectorSize;
        this.windowSize = windowSize;
        this.negativeSamples = negativeSamples;
        this.learningRate = learningRate;
        this.minCount = minCount;
        this.rand = new Random(42);
    }

    /**
     * 构建词表
     */
    public void buildVocab(List<String> sentences) {
        Map<String, Integer> freqMap = new HashMap<>();
        for (String sentence : sentences) {
            String[] words = sentence.trim().toLowerCase().split("\\s+");
            for (String word : words) {
                if (word.isEmpty()) continue;
                freqMap.merge(word, 1, Integer::sum);
            }
        }

        wordToId = new HashMap<>();
        idToWord = new HashMap<>();
        List<Integer> freqs = new ArrayList<>();

        for (Map.Entry<String, Integer> entry : freqMap.entrySet()) {
            if (entry.getValue() >= minCount) {
                int id = wordToId.size();
                wordToId.put(entry.getKey(), id);
                idToWord.put(id, entry.getKey());
                freqs.add(entry.getValue());
            }
        }

        vocabSize = wordToId.size();
        wordFreq = new int[vocabSize];
        for (int i = 0; i < vocabSize; i++) {
            wordFreq[i] = freqs.get(i);
            totalWords += wordFreq[i];
        }

        // 构建负采样的噪声分布 (使用词频的0.75次幂)
        noiseDistribution = new double[vocabSize];
        double sum = 0;
        for (int i = 0; i < vocabSize; i++) {
            noiseDistribution[i] = Math.pow(wordFreq[i], 0.75);
            sum += noiseDistribution[i];
        }
        for (int i = 0; i < vocabSize; i++) {
            noiseDistribution[i] /= sum;
        }

        System.out.println("词表大小: " + vocabSize);
    }

    /**
     * 初始化词向量
     */
    public void initVectors() {
        syn0 = new double[vocabSize][vectorSize];
        syn1Neg = new double[vocabSize][vectorSize];

        for (int i = 0; i < vocabSize; i++) {
            for (int j = 0; j < vectorSize; j++) {
                syn0[i][j] = (rand.nextDouble() - 0.5) / vectorSize;
            }
        }
    }

    /**
     * 根据噪声分布采样负样本词ID
     */
    private int sampleNegative() {
        double r = rand.nextDouble();
        double cumulative = 0;
        for (int i = 0; i < vocabSize; i++) {
            cumulative += noiseDistribution[i];
            if (r <= cumulative) {
                return i;
            }
        }
        return vocabSize - 1;
    }

    /**
     * Sigmoid函数
     */
    private double sigmoid(double x) {
        if (x > 6) return 1.0;
        if (x < -6) return 0.0;
        return 1.0 / (1.0 + Math.exp(-x));
    }

    /**
     * 训练一个中心词-上下文词对
     */
    private void trainPair(int centerId, int contextId) {
        double[] centerVec = syn0[centerId];
        double[] neu1e = new double[vectorSize]; // 误差累积

        // 正样本训练
        double dot = 0;
        for (int j = 0; j < vectorSize; j++) {
            dot += centerVec[j] * syn1Neg[contextId][j];
        }
        double g = (1 - sigmoid(dot)) * learningRate;
        for (int j = 0; j < vectorSize; j++) {
            neu1e[j] += g * syn1Neg[contextId][j];
            syn1Neg[contextId][j] += g * centerVec[j];
        }

        // 负样本训练
        for (int n = 0; n < negativeSamples; n++) {
            int negId = sampleNegative();
            if (negId == contextId) continue; // 跳过正样本

            dot = 0;
            for (int j = 0; j < vectorSize; j++) {
                dot += centerVec[j] * syn1Neg[negId][j];
            }
            g = (0 - sigmoid(dot)) * learningRate;
            for (int j = 0; j < vectorSize; j++) {
                neu1e[j] += g * syn1Neg[negId][j];
                syn1Neg[negId][j] += g * centerVec[j];
            }
        }

        // 更新中心词向量
        for (int j = 0; j < vectorSize; j++) {
            centerVec[j] += neu1e[j];
        }
    }

    /**
     * 训练模型
     */
    public void train(List<String> sentences, int epochs) {
        for (int epoch = 0; epoch < epochs; epoch++) {
            int processed = 0;
            for (String sentence : sentences) {
                String[] words = sentence.trim().toLowerCase().split("\\s+");
                List<Integer> ids = new ArrayList<>();
                for (String word : words) {
                    Integer id = wordToId.get(word);
                    if (id != null) ids.add(id);
                }

                // Skip-Gram: 每个词作为中心词，预测窗口内的上下文
                for (int i = 0; i < ids.size(); i++) {
                    int center = ids.get(i);
                    int start = Math.max(0, i - windowSize);
                    int end = Math.min(ids.size(), i + windowSize + 1);
                    for (int j = start; j < end; j++) {
                        if (i == j) continue;
                        trainPair(center, ids.get(j));
                    }
                }
                processed++;
            }
            System.out.println("Epoch " + (epoch + 1) + "/" + epochs + " 完成");
        }
    }

    /**
     * 计算两个词的余弦相似度
     */
    public double similarity(String word1, String word2) {
        Integer id1 = wordToId.get(word1.toLowerCase());
        Integer id2 = wordToId.get(word2.toLowerCase());
        if (id1 == null || id2 == null) return -1;

        double dot = 0, norm1 = 0, norm2 = 0;
        for (int i = 0; i < vectorSize; i++) {
            dot += syn0[id1][i] * syn0[id2][i];
            norm1 += syn0[id1][i] * syn0[id1][i];
            norm2 += syn0[id2][i] * syn0[id2][i];
        }
        return dot / (Math.sqrt(norm1) * Math.sqrt(norm2));
    }

    /**
     * 获取词的向量
     */
    public double[] getVector(String word) {
        Integer id = wordToId.get(word.toLowerCase());
        return id != null ? syn0[id].clone() : null;
    }
}

演示主类

import java.util.*;

public class Word2VecDemo {
    public static void main(String[] args) {
        // 构造简单语料（实际应用应使用大规模文本）
        List<String> sentences = Arrays.asList(
            "king is a strong man",
            "queen is a wise woman",
            "boy is a young man",
            "girl is a young woman",
            "prince is a young king",
            "princess is a young queen",
            "man is strong",
            "woman is wise",
            "prince is a boy will be king",
            "princess is a girl will be queen"
        );

        // 训练Word2Vec模型
        Word2Vec model = new Word2Vec(10, 2, 5, 0.025, 1);
        model.buildVocab(sentences);
        model.initVectors();
        model.train(sentences, 1000);

        // 测试语义相似度
        System.out.println("\n=== 语义相似度测试 ===");
        System.out.printf("king vs queen: %.4f%n", model.similarity("king", "queen"));
        System.out.printf("king vs prince: %.4f%n", model.similarity("king", "prince"));
        System.out.printf("king vs princess: %.4f%n", model.similarity("king", "princess"));
        System.out.printf("boy vs girl: %.4f%n", model.similarity("boy", "girl"));
        System.out.printf("man vs woman: %.4f%n", model.similarity("man", "woman"));
        System.out.printf("king vs man: %.4f%n", model.similarity("king", "man"));
        System.out.printf("queen vs woman: %.4f%n", model.similarity("queen", "woman"));
    }
}

复杂度分析

负采样将Softmax的 O(V) 复杂度降为 O(k)，其中 V 是词表大小，k 通常取5-20，极大地加速了训练。

总结

本文用Java实现了Word2Vec的Skip-Gram模型与负采样优化，完整覆盖了词表构建、向量初始化、噪声采样、前向传播与反向更新。读者可基于此扩展实现CBOW模型、层次Softmax，或将其应用到文本分类、推荐系统等下游任务中。