深入解析BERT对比学习源码：从理论到实践的关键实现细节

1次阅读

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

对比学习（Contrastive Learning）近年来在 NLP 领域取得了显著成果，其核心思想是通过拉近相似样本、推开不相似样本来学习有意义的表示。BERT 作为强大的预训练语言模型，结合对比学习能够进一步提升语义表示的质量。这种结合特别适用于语义相似度计算、文本匹配等任务，可以有效解决传统方法中语义鸿沟和效率低下的问题。

对比学习的核心在于损失函数的设计。在 BERT 中，常用的对比损失函数是 InfoNCE（Noise Contrastive Estimation），其数学形式如下：

def contrastive_loss(anchor, positive, negatives, temperature=0.1):
    """
    anchor: 锚点样本的向量表示 [batch_size, hidden_size]
    positive: 正样本的向量表示 [batch_size, hidden_size]
    negatives: 负样本的向量表示 [batch_size, num_negatives, hidden_size]
    temperature: 温度系数，控制分布的尖锐程度
    """
    # 计算锚点与正样本的相似度
    pos_sim = torch.sum(anchor * positive, dim=-1) / temperature

    # 计算锚点与所有负样本的相似度
    neg_sim = torch.einsum('bd,bnd->bn', anchor, negatives) / temperature

    # 合并所有相似度
    logits = torch.cat([pos_sim.unsqueeze(-1), neg_sim], dim=-1)

    # 计算交叉熵损失
    labels = torch.zeros(logits.shape[0], dtype=torch.long).to(anchor.device)
    loss = F.cross_entropy(logits, labels)

    return loss

这个实现有几个关键点：

温度系数（temperature）控制着相似度分布的尖锐程度，较小的值会使模型更关注困难的负样本。
使用 einsum 进行高效的矩阵运算，计算锚点与所有负样本的相似度。
将正样本作为第一个类别，负样本作为后续类别，使用标准的交叉熵损失。

负样本的质量直接影响对比学习的效果。在 BERT 对比学习中，常见的负样本采样策略有：

同一 batch 内的其他样本作为负样本（in-batch negatives）
从外部存储库中检索的困难负样本（hard negatives）
随机采样的简单负样本（easy negatives）

下面是一个混合采样的实现示例：

class NegativeSampler:
    def __init__(self, memory_size=65536, hard_ratio=0.5):
        self.memory_bank = torch.randn(memory_size, 768)  # 假设隐藏维度是 768
        self.hard_ratio = hard_ratio

    def sample(self, anchor, batch_negatives, num_negatives=16):
        batch_size = anchor.size(0)

        # 从 batch 内采样部分负样本
        num_in_batch = int(num_negatives * (1 - self.hard_ratio))
        in_batch_negs = batch_negatives[torch.randperm(batch_size)[:num_in_batch]]

        # 从 memory bank 采样困难负样本
        with torch.no_grad():
            similarities = torch.einsum('bd,nd->bn', anchor, self.memory_bank)
            _, hard_indices = similarities.topk(num_negatives - num_in_batch, dim=1, largest=False)
            hard_negs = self.memory_bank[hard_indices]

        # 合并负样本
        negatives = torch.cat([in_batch_negs.unsqueeze(1), hard_negs], dim=1)

        return negatives

在对比学习中，向量归一化（L2 normalization）是至关重要的步骤，原因有三：

防止向量范数影响相似度计算
使相似度得分落在 [-1,1] 区间
提高训练稳定性

实现非常简单：

normalized_embeddings = F.normalize(raw_embeddings, p=2, dim=-1)

下面是一个完整的 BERT 对比学习训练步骤示例：

# 1. 准备输入
input_ids = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")

# 2. 获取 BERT 编码（使用两个不同的 dropout mask 实现数据增强）with torch.no_grad():
    anchor_emb = model(input_ids=input_ids).last_hidden_state[:, 0]  # [CLS] token

# 3. 正向传播（使用不同的 dropout mask）positive_emb = model(input_ids=input_ids).last_hidden_state[:, 0]

# 4. 归一化
anchor_emb = F.normalize(anchor_emb, p=2, dim=-1)
positive_emb = F.normalize(positive_emb, p=2, dim=-1)

# 5. 采样负样本
negatives = sampler.sample(anchor_emb, batch_embeddings)
negatives = F.normalize(negatives, p=2, dim=-1)

# 6. 计算对比损失
loss = contrastive_loss(anchor_emb, positive_emb, negatives)

# 7. 反向传播
loss.backward()
optimizer.step()

batch size 对对比学习效果有显著影响，主要体现在：