BERT与对比学习结合实战：解决小样本场景下的语义相似度计算难题

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

背景痛点

在实际业务场景中，我们经常遇到标注数据稀缺的问题。传统 BERT 模型直接微调时，容易出现两个典型问题：

1. 过拟合问题：当训练样本不足时（比如每类只有几十个样本），BERT 强大的表征能力反而会导致模型记住训练集噪声
2. 表征空间坍缩：所有句子向量都聚集在超球面的狭窄区域，导致相似度计算失效

通过可视化分析可以发现，普通 BERT 微调后句向量的平均余弦相似度往往高达 0.8 以上，严重影响了语义区分度。

技术方案对比

常见的语义匹配方案主要有三种：

• 孪生网络：结构简单但容易陷入平凡解
• Triplet Loss：需要精心设计 triplet 采样策略
• 对比学习：通过正负样本对比学习，特别适合小样本场景

其中对比学习使用的 InfoNCE 损失函数可以表示为：

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

这里 τ 是温度系数，控制分布尖锐程度，实验表明 τ =0.05~0.2 效果最佳。

核心实现

模型加载

我们基于 HuggingFace Transformers 实现：

from transformers import BertModel, BertTokenizer
import torch.nn as nn

class ContrastiveBERT(nn.Module):
    def __init__(self, model_name='bert-base-chinese'):
        super().__init__()
        self.bert = BertModel.from_pretrained(model_name)
        self.temperature = nn.Parameter(torch.tensor(0.07))

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids, attention_mask)
        # 使用[CLS]token 作为句子表示
        embeddings = outputs.last_hidden_state[:,0,:]  
        return F.normalize(embeddings, p=2, dim=1)

关键训练逻辑

实现动态负采样和混合精度训练：

# 初始化
model = ContrastiveBERT().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
scaler = torch.cuda.amp.GradScaler()

# 训练循环
for batch in train_loader:
    with torch.cuda.amp.autocast():
        # 获取批次内所有样本的 embedding
        embeddings = model(batch['input_ids'], batch['attention_mask'])

        # 计算相似度矩阵
        sim_matrix = embeddings @ embeddings.T  # [batch, batch]

        # 生成标签：对角线为正样本
        labels = torch.arange(len(sim_matrix)).cuda()

        # 计算 InfoNCE 损失
        loss = F.cross_entropy(sim_matrix/model.temperature, labels)

    # 混合精度反向传播    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

生产优化建议

1. 显存优化：当 batch_size 较大时，可以分块计算相似度矩阵：

# 分块计算避免 OOM
chunk_size = 512
for i in range(0, len(embeddings), chunk_size):
    chunk = embeddings[i:i+chunk_size]
    sim_chunk = chunk @ embeddings.T  # 只计算部分行

1. 温度系数调参：建议初始设为 0.1，然后在 0.02~0.2 范围内网格搜索

2. 在线服务：使用 ONNX 导出模型，并启用 CUDA Graph 优化：

python -m transformers.onnx --model=checkpoint/ --feature=sequence-classification onnx_model/

实验效果

在 LCQMC 测试集上的结果对比：

延伸思考

1. 如何结合课程学习 (Curriculum Learning) 进一步提升小样本效果？
2. 在跨语言场景下，对比学习应该如何调整负采样策略？
3. 对于超大规模批次(>1 万)，有哪些高效的负采样实现方案？

完整的实现代码已开源在 GitHub 仓库(伪地址)：github.com/example/contrastive-bert