BERT对比学习源码解析：从理论到高效实现

1次阅读

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

在自然语言处理任务中，语义相似度计算是一个基础但关键的问题。传统的 BERT 微调方法虽然有效，但在实际应用中往往面临两个主要挑战：一是需要大量标注数据才能达到理想效果，二是模型收敛速度较慢，训练成本高。对比学习作为一种自监督学习方法，通过构建正负样本对，能够更高效地学习语义表示，特别适合数据稀缺的场景。

对比学习的核心思想是拉近相似样本的表示距离，推远不相似样本的表示距离。其数学基础是 InfoNCE 损失函数，公式如下：

$$
L = -\log\frac{\exp(sim(q,k^+)/\tau)}{\sum_{i=1}^N \exp(sim(q,k_i)/\tau)}
$$

其中，$q$ 是查询样本，$k^+$ 是正样本，$k_i$ 包含正样本和负样本，$\tau$ 是温度系数，$sim$ 是相似度函数（通常为余弦相似度）。

在 HuggingFace 的 BertForSequenceClassification 基础上，我们需要进行以下关键改造：

继承 BertPreTrainedModel 创建新的对比学习模型类
重写 forward 方法实现对比学习逻辑
添加自定义的对比损失函数

以下是关键代码片段（带注释）：

class BertForContrastiveLearning(BertPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.bert = BertModel(config)
        self.temperature = config.temperature  # 温度系数
        self.init_weights()

    def forward(self, input_ids, attention_mask, labels=None):
        # 获取句子表示
        outputs = self.bert(input_ids, attention_mask=attention_mask)
        # 使用 [CLS]token 作为句子表示
        embeddings = outputs.last_hidden_state[:, 0, :]

        if labels is not None:
            # 计算对比损失
            loss = self.contrastive_loss(embeddings, labels)
            return loss
        return embeddings

高效的负采样策略对对比学习至关重要。我们采用 batch 内负采样方法，显著减少内存消耗：

在同一 batch 内自动构造负样本
使用掩码避免将正样本误认为负样本
实现内存高效的相似度矩阵计算

def contrastive_loss(self, embeddings, labels):
    # 归一化处理
    embeddings = F.normalize(embeddings, p=2, dim=1)

    # 计算相似度矩阵
    sim_matrix = torch.matmul(embeddings, embeddings.T) / self.temperature

    # 构建正样本掩码
    pos_mask = labels.unsqueeze(0) == labels.unsqueeze(1)
    diag_mask = ~torch.eye(labels.size(0), dtype=torch.bool).to(labels.device)
    pos_mask = pos_mask & diag_mask

    # 计算对比损失
    exp_sim = torch.exp(sim_matrix)
    pos_sim = torch.sum(exp_sim * pos_mask, dim=1)
    neg_sim = torch.sum(exp_sim * (~pos_mask), dim=1)
    loss = -torch.log(pos_sim / (pos_sim + neg_sim)).mean()

    return loss

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    loss = model(input_ids, attention_mask, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

accumulation_steps = 4

for i, (batch, labels) in enumerate(train_loader):
    loss = model(batch, labels)
    loss = loss / accumulation_steps
    loss.backward()

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()