BGE对比学习训练实战：从数据准备到模型优化的全流程指南

1次阅读

没有评论

共计 2421 个字符，预计需要花费 7 分钟才能阅读完成。

对比学习在 NLP 领域已经证明能显著提升模型性能，但在实际训练 BGE（Bidirectional Generative Encoder）时，开发者常会遇到几个典型问题：

数据稀疏性 ：正样本对数量有限，导致模型难以学习到足够丰富的表示。
负样本质量差 ：随机采样的负样本往往与正样本差异过大，无法提供有意义的对比信号。
训练不稳定 ：损失波动大、收敛慢，尤其在训练初期容易出现梯度爆炸。

这些问题直接影响模型最终效果和训练效率。接下来，我将分享一套经过实战验证的优化方案。

传统的数据增强方法如简单的 token masking 可能不够充分。我们采用组合策略：

动态 token masking：随机 mask 输入序列中 15%-30% 的 token，比例随训练轮次动态调整。
token shuffling：在 mask 基础上，对未被 mask 的部分 token 进行局部重排（限制在 3 -token 窗口内）。
同义替换 ：对小部分非关键实体词使用同义词库替换，增加语义多样性。

这种组合策略能在保持语义一致性的同时，有效增加数据多样性。

BGE 使用的对比损失函数是 InfoNCE 的变体：

$$\mathcal{L} = -\log\frac{e^{sim(q,k^+)/\tau}}{e^{sim(q,k^+)/\tau} + \sum_{k^-}e^{sim(q,k^-)/\tau}}$$

我们的实现中有几个关键点：

温度参数 τ 的动态调整 ：初始设为 0.1，每 5 个 epoch 根据验证集表现调整（±0.02）。
相似度计算优化 ：采用双向最大余弦相似度而非简单点积：
$$sim(q,k) = \max(cos_sim(q,k), cos_sim(k,q))$$
梯度裁剪 ：对对比损失部分的梯度实施 L2 norm 裁剪（阈值设为 1.0）。

渐进式负采样 ：
前 5 个 epoch：使用 in-batch negatives
5-10 个 epoch：增加 hard negatives（相似度 top50% 的样本）
10 个 epoch 后：引入跨 batch 的 memory bank negatives

学习率 warmup+ 衰减 ：

scheduler = get_cosine_schedule_with_warmup(
    optimizer, 
    num_warmup_steps=1000, 
    num_training_steps=total_steps
)

class BGEDataset(Dataset):
    def __init__(self, texts, tokenizer, aug_prob=0.3):
        self.texts = texts
        self.tokenizer = tokenizer
        self.aug_prob = aug_prob

    def augment(self, text):
        # 组合增强策略
        tokens = text.split()
        if random.random() < self.aug_prob:
            # Token masking
            mask_idx = random.sample(range(len(tokens)), 
                           k=int(len(tokens)*0.2))
            tokens = ["[MASK]" if i in mask_idx else t 
                     for i,t in enumerate(tokens)]

            # 局部 shuffling
            if len(tokens) > 4 and random.random() < 0.5:
                start = random.randint(0, len(tokens)-3)
                tokens[start:start+3] = random.sample(tokens[start:start+3], k=3)
        return " ".join(tokens)

class ContrastiveLoss(nn.Module):
    def __init__(self, temp=0.1):
        super().__init__()
        self.temp = temp

    def forward(self, q, k_pos, k_negs):
        # q: [batch, dim], k_pos: [batch, dim], 
        # k_negs: [batch, neg_num, dim]
        pos_sim = torch.cosine_similarity(q, k_pos, dim=-1)
        neg_sim = torch.cosine_similarity(q.unsqueeze(1), k_negs, dim=-1)

        logits = torch.cat([pos_sim.unsqueeze(-1)/self.temp, 
                           neg_sim/self.temp], dim=1)
        labels = torch.zeros(q.size(0), dtype=torch.long).to(q.device)
        return F.cross_entropy(logits, labels)

在 MSMARCO 数据集上的实验结果：