突破BERT局限:SimCSE无监督训练中文句子向量实战指南

在自然语言处理领域,语义相似度计算一直是核心任务之一。传统方法通常依赖标注数据进行监督学习,但高质量标注数据的获取成本极高。想象一下,当你接手一个客服对话分析项目,手头只有海量未标注的聊天记录,却需要快速构建一个能够准确理解用户意图的语义匹配系统——这正是SimCSE大显身手的场景。

SimCSE(Simple Contrastive Learning of Sentence Embeddings)通过创新的"Dropout即数据增强"思想,仅需无标注文本就能训练出高质量的句子向量表示。与需要复杂数据增强策略的对比学习方法不同,SimCSE的实现异常简洁,却能达到甚至超越监督学习的效果。本文将带你从零开始,用PyTorch实现一个完整的中文SimCSE训练流程。

1. 环境准备与数据加载

1.1 搭建基础环境

首先确保你的开发环境已安装必要依赖。推荐使用Python 3.8+和PyTorch 1.8+版本:

pip install torch transformers pandas tqdm scikit-learn

对于中文任务,我们选择bert-base-chinese作为基础模型。这个预训练模型在中文文本表示上已有不错的表现,适合作为SimCSE的起点。

1.2 准备无标注数据

SimCSE的美妙之处在于它对数据要求极低——只需要大量自然文本,无需任何标注。假设我们有一个客服对话数据集customer_service.txt,每行包含一条用户query或客服回复:

import pandas as pd

# 加载原始文本数据
with open('customer_service.txt', 'r', encoding='utf-8') as f:
    texts = [line.strip() for line in f if len(line.strip()) > 10]  # 过滤过短文本

print(f"共加载{len(texts)}条文本样例")
print("示例文本:", texts[0])

提示:数据质量对最终效果影响显著。建议预处理时去除明显噪声(如特殊符号、乱码等),但保留文本的原始表达方式,这对学习真实的语义表示很重要。

2. 模型架构设计与实现

2.1 构建SimCSE核心结构

SimCSE的核心思想非常简单:将同一个句子两次输入BERT,由于Dropout的随机性,会得到两个略有不同的向量表示,将这些表示作为正样本对,同时将其他句子的表示作为负样本。

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

class SimCSE(nn.Module):
    def __init__(self, model_name='bert-base-chinese'):
        super(SimCSE, self).__init__()
        self.bert = BertModel.from_pretrained(model_name)
        self.tokenizer = BertTokenizer.from_pretrained(model_name)
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.bert.to(self.device)
        
    def forward(self, input_ids, attention_mask):
        # 通过BERT获取句子表示
        outputs = self.bert(input_ids, attention_mask=attention_mask)
        # 使用[CLS]token的表示作为整个句子的表示
        cls_rep = outputs.last_hidden_state[:, 0]
        return cls_rep

2.2 实现对比损失函数

InfoNCE(Noise Contrastive Estimation)损失是SimCSE训练的关键,它鼓励正样本对的相似度高于负样本对:

def contrastive_loss(embeddings, temperature=0.05):
    # embeddings: [batch_size, hidden_size]
    batch_size = embeddings.size(0)
    # 计算相似度矩阵
    sim_matrix = torch.matmul(embeddings, embeddings.T)  # [batch_size, batch_size]
    # 对角线元素是正样本对的相似度
    pos_sim = torch.diag(sim_matrix) / temperature
    # 构造标签:对角线位置为正样本
    labels = torch.arange(batch_size).to(embeddings.device)
    # 计算交叉熵损失
    loss = nn.CrossEntropyLoss()(sim_matrix/temperature, labels)
    return loss

注意:温度参数temperature对模型性能影响很大。较小的温度会使模型更关注困难的负样本,通常设置在0.01到0.1之间。

3. 训练流程优化技巧

3.1 高效数据加载与处理

为了充分利用GPU资源,我们需要实现一个高效的数据加载器:

from torch.utils.data import Dataset, DataLoader

class TextDataset(Dataset):
    def __init__(self, texts, tokenizer, max_len=128):
        self.texts = texts
        self.tokenizer = tokenizer
        self.max_len = max_len
    
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        text = self.texts[idx]
        inputs = self.tokenizer(
            text, 
            max_length=self.max_len, 
            padding='max_length', 
            truncation=True, 
            return_tensors='pt'
        )
        return {
            'input_ids': inputs['input_ids'].squeeze(),
            'attention_mask': inputs['attention_mask'].squeeze()
        }

# 初始化数据集和数据加载器
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
dataset = TextDataset(texts, tokenizer)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True)

3.2 训练循环实现

下面是一个完整的训练循环,包含梯度累积和模型保存逻辑:

def train(model, dataloader, optimizer, epochs=3, accumulation_steps=4):
    model.train()
    total_steps = len(dataloader) * epochs
    current_step = 0
    
    for epoch in range(epochs):
        for batch_idx, batch in enumerate(dataloader):
            # 将数据移动到设备
            input_ids = batch['input_ids'].to(model.device)
            attention_mask = batch['attention_mask'].to(model.device)
            
            # 前向传播 - 同一个batch输入两次利用Dropout得到不同表示
            embeddings1 = model(input_ids, attention_mask)
            embeddings2 = model(input_ids, attention_mask)
            
            # 合并两次的结果作为最终batch
            embeddings = torch.cat([embeddings1, embeddings2], dim=0)
            
            # 计算损失
            loss = contrastive_loss(embeddings)
            loss = loss / accumulation_steps  # 梯度累积
            
            # 反向传播
            loss.backward()
            
            # 梯度累积到指定步数后更新参数
            if (batch_idx + 1) % accumulation_steps == 0:
                optimizer.step()
                optimizer.zero_grad()
            
            current_step += 1
            if current_step % 100 == 0:
                print(f"Epoch {epoch+1}, Step {current_step}/{total_steps}, Loss: {loss.item()*accumulation_steps:.4f}")
    
    # 保存训练好的模型
    torch.save(model.state_dict(), 'simcse_bert_chinese.bin')

4. 评估与应用实践

4.1 语义相似度评估

训练完成后,我们需要评估模型在语义相似度任务上的表现。一个常用的方法是计算Spearman相关系数:

from sklearn.metrics.pairwise import cosine_similarity
from scipy.stats import spearmanr

def evaluate(model, eval_pairs):
    """评估模型在语义相似度任务上的表现
    eval_pairs: [(text1, text2, human_score), ...]
    """
    model.eval()
    pred_scores = []
    human_scores = []
    
    with torch.no_grad():
        for text1, text2, score in eval_pairs:
            # 编码两个句子
            inputs1 = model.tokenizer(text1, return_tensors='pt', padding=True, truncation=True).to(model.device)
            inputs2 = model.tokenizer(text2, return_tensors='pt', padding=True, truncation=True).to(model.device)
            
            emb1 = model(inputs1['input_ids'], inputs1['attention_mask']).cpu().numpy()
            emb2 = model(inputs2['input_ids'], inputs2['attention_mask']).cpu().numpy()
            
            # 计算余弦相似度
            sim = cosine_similarity(emb1, emb2)[0][0]
            pred_scores.append(sim)
            human_scores.append(score)
    
    # 计算Spearman相关系数
    spearman = spearmanr(human_scores, pred_scores).correlation
    print(f"Spearman相关系数: {spearman:.4f}")
    return spearman

4.2 实际应用示例

训练好的SimCSE模型可以广泛应用于各种语义相关任务。以下是一个简单的语义搜索实现:

class SemanticSearch:
    def __init__(self, model, texts):
        self.model = model
        self.texts = texts
        self.embeddings = self._encode_all(texts)
    
    def _encode_all(self, texts):
        self.model.eval()
        embeddings = []
        with torch.no_grad():
            for text in texts:
                inputs = self.model.tokenizer(text, return_tensors='pt', 
                                           padding=True, truncation=True).to(self.model.device)
                emb = self.model(inputs['input_ids'], inputs['attention_mask'])
                embeddings.append(emb.cpu().numpy())
        return np.concatenate(embeddings, axis=0)
    
    def search(self, query, top_k=5):
        with torch.no_grad():
            inputs = self.model.tokenizer(query, return_tensors='pt', 
                                       padding=True, truncation=True).to(self.model.device)
            query_emb = self.model(inputs['input_ids'], inputs['attention_mask']).cpu().numpy()
        
        # 计算余弦相似度
        similarities = cosine_similarity(query_emb, self.embeddings)[0]
        # 获取最相似的top_k个结果
        top_indices = similarities.argsort()[-top_k:][::-1]
        return [(self.texts[i], similarities[i]) for i in top_indices]

在实际项目中,我发现几个关键点对最终效果影响很大:首先是batch size要尽可能大(受限于GPU内存),因为更多的负样本能提供更好的对比信号;其次是学习率不宜过大,通常设置在1e-5到5e-5之间;最后是数据质量比数量更重要,清洗掉低质量文本能显著提升模型性能。

Logo

码道开发者社区,聚焦华为云码道 CodeArts 代码智能体,沉淀 Agent、Skill、鸿蒙开发实战内容,供开发者查阅资料、交流技术、分享工程实践

更多推荐