自制Word2Vec模型梯度下降算法无法收敛问题排查

为练习、学习与实验，我基于Jurafsky & Martin的公式与算法从零实现了Word2Vec模型，但遇到以下问题：当训练语料超过20个句子时，损失值无法降到1以下，除非使用极低学习率并训练数百轮epoch，耗时长达6-7小时，远慢于FastText等成熟工具。希望排查代码中的潜在bug，或获取优化建议以提升训练效率与收敛效果。

class Word2Vec:
    def __init__(self, context_window=2, num_dimensions=5, num_negative_samples=2, normalizer_func=None):
        self.num_dimensions = num_dimensions
        self.context_window = context_window
        self.num_negative_samples = num_negative_samples * context_window
        self.normalize = normalizer_func
        self.W = None
        self.V = None

    def train(self, corpus, learning_rate=0.1, decay_rate=0.1, num_epochs=100):
        """
        (1) Build the vocabulary from the corpus.
        (2) Initialize the weight matrices.
        (3) Iterate over the number of epochs.
        (4) Update the weights for each training example (target word, list of context words, list of negative samples).
        (5) Add the two trained matrices.
        """
        self.losses = list()  # List to store the value of the loss function at each epoch (for informative purposes).
        self.count = 0
        self.last_print = ''
        self.num_epochs = num_epochs

        self.build_vocabulary(corpus)
        self.initialize_weights()

        for epoch in range(self.num_epochs):
            self.total_loss = 0.0  # For informative purposes.
            learning_rate = self.get_learning_rate(learning_rate=learning_rate)
            for target, context, negative_samples in self.generate_training_data(corpus):
                self.update_weights(target, context, negative_samples, learning_rate)
            self.losses.append(self.total_loss)
            print(f"Loss at epoch {epoch}: {self.total_loss}")

        self.W += self.V

    def build_vocabulary(self, corpus):
        """
        Create a list of unique words, where the index represents the word's idx.
        """
        self.vocabulary = set()
        self.word_counts = {}
        for sentence in corpus:
            for word in sentence.split():
                self.vocabulary.add(word)
                self.word_counts[word] = self.word_counts.get(word, 0) + 1
        self.vocabulary = list(self.vocabulary)
        # 新增：词到索引的映射，提升查找效率
        self.word_to_idx = {word: idx for idx, word in enumerate(self.vocabulary)}

    def initialize_weights(self):
        """
        Create two weight matrices, W (target words) and V (context words), with dimensions vocab_size * ndims.
        The value of each dimension is randomly chosen within the range [-0.5/ndims, 0.5/ndims].
        """
        vocab_size = len(self.vocabulary)
        # 修复：按照注释要求，除以num_dimensions
        self.W = np.random.uniform(low=-0.5/self.num_dimensions, high=0.5/self.num_dimensions, size=(vocab_size, self.num_dimensions))
        self.V = np.random.uniform(low=-0.5/self.num_dimensions, high=0.5/self.num_dimensions, size=(vocab_size, self.num_dimensions))

    def generate_training_data(self, corpus):
        """
        Convert each sentence in the corpus into an ordered list of words and yield the target word, context words, and negative samples.
        """
        # 预计算负采样概率分布
        if not hasattr(self, 'negative_sampling_probs'):
            self._compute_negative_sampling_probs()

        for sentence in corpus:
            self.print_progress(len(corpus) * self.num_epochs)
            sentence_list = sentence.strip().split()
            for target_idx, target in enumerate(sentence_list):
                context = self.get_context(sentence_list, target_idx)
                negative_samples = self.get_negative_samples(context)
                yield target, context, negative_samples

    def get_context(self, sentence, target_idx):
        """
        Given a list of words (sentence) and the target word's index within that list, return a list of context words.
        """
        context = []
        for i in range(target_idx - self.context_window, target_idx + self.context_window + 1):
            if i != target_idx and i >= 0 and i < len(sentence):
                context.append(sentence[i])
        return context

    def _compute_negative_sampling_probs(self):
        # 基于词频的负采样分布：词频^0.75 / 总词频^0.75
        total = sum(count**0.75 for count in self.word_counts.values())
        self.negative_sampling_probs = [self.word_counts[word]**0.75 / total for word in self.vocabulary]

    def get_negative_samples(self, context):
        """
        Select random words from the vocabulary that are not in the context words.
        """
        context_set = set(context)
        # 批量采样，过滤掉上下文词，直到数量足够
        negative_samples = []
        while len(negative_samples) < self.num_negative_samples:
            # 用预计算的概率分布采样
            samples = np.random.choice(self.vocabulary, size=self.num_negative_samples - len(negative_samples), p=self.negative_sampling_probs)
            negative_samples.extend([s for s in samples if s not in context_set])
        return negative_samples[:self.num_negative_samples]

    def update_weights(self, target, context, negative_samples, learning_rate):
        """
        Iterate over the context words and negative samples.
        (1) Compute the error.
        (2) Add it to a total loss attribute (for later printing).
        (3) Update the weight matrices.
        """
        target_idx = self.word_to_idx[target]
        target_vec = self.W[target_idx].copy()  # 保存原始target向量，避免更新W后影响V的计算

        # 处理正样本（上下文词）
        for context_word in context:
            ctx_idx = self.word_to_idx[context_word]
            ctx_vec = self.V[ctx_idx]
            dot = np.dot(target_vec, ctx_vec)
            sig = 1 / (1 + np.exp(-dot))
            # 正确的正样本损失
            self.total_loss += -np.log(sig)
            error = sig - 1  # 因为label=1，error=pred - label
            # 更新权重：用原始target_vec计算
            self.W[target_idx] -= learning_rate * error * ctx_vec
            self.V[ctx_idx] -= learning_rate * error * target_vec

        # 处理负样本
        for neg_word in negative_samples:
            neg_idx = self.word_to_idx[neg_word]
            neg_vec = self.V[neg_idx]
            dot = np.dot(target_vec, neg_vec)
            sig = 1 / (1 + np.exp(-dot))
            # 正确的负样本损失
            self.total_loss += -np.log(1 - sig)
            error = sig - 0  # label=0，error=pred - label
            # 更新权重：用原始target_vec计算
            self.W[target_idx] -= learning_rate * error * neg_vec
            self.V[neg_idx] -= learning_rate * error * target_vec

    def compute_error(self, target, context_word, label):
        """
        已废弃，整合到update_weights中，避免重复计算
        """
        pass

    def update_weights_for_word(self, target, word, error, learning_rate):
        """
        已废弃，整合到update_weights中
        """
        pass

    def print_progress(self, len_corpus):
        self.count += 1
        trained = round((self.count / len_corpus) * 100)
        if trained != self.last_print and trained % 10 == 0:  # 每10%刷新一次，减少IO操作
            clear_output(wait=True)
            print(f'{trained}% processed', flush=True)
            self.last_print = trained

    def get_word_vector(self, word):
        if word in self.word_to_idx:
            return self.W[self.word_to_idx[word]].reshape(1, -1)
        else:
            print(f"{word} is not in the vocabulary.")

    def compute_most_similar_words(self, word, top_k=5):
        if self.normalize:
            word = self.normalize(word)
        word_vector = self.get_word_vector(word)

        if word_vector is None:
            return []

        # 计算余弦相似度（更合理的相似度指标）
        norm = np.linalg.norm(self.W, axis=1, keepdims=True)
        normalized_W = self.W / norm
        normalized_word_vec = word_vector / np.linalg.norm(word_vector)
        similarities = np.dot(normalized_W, normalized_word_vec.T)

        top_indices = np.argsort(similarities, axis=0)[::-1][:top_k+1]  # 多取一个，跳过自身
        similar_words = [self.vocabulary[idx] for idx in top_indices.flatten() if idx != self.word_to_idx[word]]
        return similar_words[:top_k]

    def get_learning_rate(self, learning_rate, min_value=0.0001):
        # 替换为随步数衰减的策略，更稳定
        step = self.count
        decay_rate = 0.001
        new_lr = learning_rate / (1 + decay_rate * step)
        return max(new_lr, min_value)

• 损失计算错误
  原代码中损失用abs(np.log(error + 1.01))，完全不符合负采样的损失公式。正确的损失应该是：正样本（上下文词）用-log(sigmoid(target_vec · context_vec))，负样本用-log(1 - sigmoid(target_vec · negative_vec))。修复后损失能正确反映模型预测效果，引导权重向最优方向更新。

• 权重更新顺序错误
  原代码先更新W[target_idx]，再用更新后的W去更新V[word_idx]，导致梯度计算基于已修改的权重，偏离正确方向。修复方法是先保存原始的target_vec，用该向量同时计算W和V的更新量。

• 权重初始化错误
  原代码注释说明权重应在[-0.5/ndims, 0.5/ndims]范围内，但实际初始化时未除以num_dimensions，导致初始权重范围过大，训练时梯度波动剧烈，难以收敛。

• 优化词查找效率
  原代码用self.vocabulary.index(word)查找索引，时间复杂度O(n)。新增word_to_idx字典，将查找时间降为O(1)，大幅减少循环中的耗时。

• 改进负采样策略
  原代码采用均匀随机采样，Word2Vec标准实现中使用词频^0.75的加权分布，能降低高频词的采样概率，提升低频词的采样机会，既符合语言规律，又能加快收敛。

• 优化学习率衰减策略
  原代码仅在损失上升时减半学习率，策略过于简单。改为随训练步数线性衰减（lr = initial_lr / (1 + decay_rate * step)），学习率随训练进程平稳下降，避免后期震荡。

• 减少IO操作
  原代码print_progress每次进度变化都刷新输出，频繁IO会拖慢训练速度。改为每10%进度刷新一次，减少不必要的输出操作。

• 向量化与批量计算
  原代码逐词处理上下文和负样本，可进一步优化为批量计算（比如将上下文词的向量一次性取出，用numpy向量化操作计算误差和权重更新），能大幅提升训练速度。

• 向量相似度计算优化
  原代码用点积计算相似度，改为余弦相似度，消除向量长度对相似度的影响，结果更符合Word2Vec的语义相似性预期。

内容的提问来源于stack exchange，提问作者invalid syntax