BERT对比学习实战指南：从原理到高效实现

1次阅读

没有评论

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

对比学习（Contrastive Learning）是自监督学习的一种范式，其核心思想是通过构造正负样本对，让模型学习到数据的内在表示。在 NLP 领域，对比学习主要应用于：

语义相似度计算
文本表示学习
句子嵌入生成
跨模态对齐

与传统的有监督学习相比，对比学习的优势在于：

不需要大量标注数据
能学习到更通用的文本表示
对数据噪声更鲁棒

选择 BERT 作为对比学习的基模型有以下几个显著优势：

强大的上下文理解能力：BERT 的双向 Transformer 结构可以捕获丰富的上下文信息
预训练知识迁移：BERT 在预训练阶段已经学习到大量语言知识
灵活的架构适配：可以通过不同的池化方式获取句向量表示

from transformers import BertTokenizer
import torch

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

def prepare_batch(text_pairs, max_length=128):
    """
    准备对比学习的输入批次
    :param text_pairs: 正样本对列表 [(text1, text2), ...]
    :param max_length: 最大序列长度
    :return: 字典形式的模型输入
    """
    texts = [t for pair in text_pairs for t in pair]  # 展开所有文本
    inputs = tokenizer(
        texts, 
        max_length=max_length, 
        padding='max_length', 
        truncation=True, 
        return_tensors='pt'
    )
    return inputs

from transformers import BertModel
import torch.nn as nn

class BertForContrastiveLearning(nn.Module):
    def __init__(self, model_name='bert-base-uncased'):
        super().__init__()
        self.bert = BertModel.from_pretrained(model_name)
        self.projection = nn.Sequential(nn.Linear(768, 768),
            nn.ReLU(),
            nn.Linear(768, 256)  # 降维到更适合对比学习的空间
        )

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        # 使用 [CLS]token 作为句表示
        pooled_output = outputs.last_hidden_state[:, 0, :]
        # 投影到对比学习空间
        projected = self.projection(pooled_output)
        return projected

import torch.nn.functional as F

def contrastive_loss(embeddings, temperature=0.05):
    """
    NT-Xent 损失函数（归一化温度缩放交叉熵损失）:param embeddings: 投影后的嵌入向量 [2N, dim]
    :param temperature: 温度参数
    :return: 对比损失值
    """
    batch_size = embeddings.size(0)
    # 计算余弦相似度
    sim_matrix = F.cosine_similarity(embeddings.unsqueeze(1), embeddings.unsqueeze(0), dim=2)
    # 创建标签：对角线上的元素是正样本对
    labels = torch.arange(batch_size, device=embeddings.device)
    labels = (labels - labels % 2 * 2) + 1  # 0↔1, 2↔3, ...
    # 计算交叉熵损失
    loss = F.cross_entropy(sim_matrix/temperature, labels)
    return loss

# 训练循环示例
def train_one_epoch(model, dataloader, optimizer, device):
    model.train()
    total_loss = 0

    for batch in dataloader:
        inputs = prepare_batch(batch)
        inputs = {k:v.to(device) for k,v in inputs.items()}

        optimizer.zero_grad()
        embeddings = model(inputs['input_ids'], inputs['attention_mask'])
        loss = contrastive_loss(embeddings)

        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    return total_loss / len(dataloader)