复现SimCSE

### 如何从头开始实现或训练 SimCSE 模型 SimCSE 是一种用于无监督和有监督语义相似度学习的方法，它通过对比学习来增强句子嵌入的质量。以下是使用 PyTorch 和 Hugging Face 的 `transformers` 库实现 SimCSE 的方法。 #### 数据准备 为了训练 SimCSE 模型，需要构建一个包含重复句子的数据集。对于无监督版本，可以通过随机打乱词序等方式生成正样本[^1]。而对于有监督版本，则可以使用标注好的成对相似句作为输入数据[^3]。 ```python from datasets import load_dataset, DatasetDict def create_unsupervised_data(sentences): augmented_sentences = [] for sentence in sentences: # 随机扰动或其他方式生成正样本 pass return {"sentence": sentences + augmented_sentences} raw_datasets = load_dataset("your_dataset_name") unsupervised_datasets = raw_datasets.map(create_unsupervised_data) ``` #### 加载预训练模型与 tokenizer 加载 Hugging Face 提供的预训练语言模型以及对应的 Tokenizer： ```python from transformers import AutoTokenizer, AutoModelForSequenceClassification model_checkpoint = "bert-base-uncased" tokenizer = AutoTokenizer.from_pretrained(model_checkpoint) model = AutoModelForSequenceClassification.from_pretrained( model_checkpoint, num_labels=1 # 对于对比学习任务通常设置为单标签回归形式 ).to('cuda') ``` #### 构建对比损失函数 SimCSE 使用的是 InfoNCE Loss 来优化目标函数。该损失计算每个句子与其自身的余弦相似度与其他负样例之间的差异[^4]。 ```python import torch import torch.nn.functional as F class ContrastiveLoss(torch.nn.Module): def __init__(self, temperature=0.05): super().__init__() self.temperature = temperature def forward(self, embeddings): cos_sim = F.cosine_similarity(embeddings.unsqueeze(1), embeddings.unsqueeze(0), dim=-1) / self.temperature labels = torch.arange(cos_sim.size(0)).long().to(embeddings.device) loss = F.cross_entropy(cos_sim, labels) return loss ``` #### 训练过程配置 采用混合精度训练能够显著降低显存占用并加速收敛速度[^2]。此外还可以考虑分布式训练进一步提升效率。 ```python from transformers import Trainer, TrainingArguments training_args = TrainingArguments( output_dir="./results", evaluation_strategy="epoch", learning_rate=3e-5, per_device_train_batch_size=64, weight_decay=0.01, save_total_limit=2, fp16=True, # 启用半精度浮点数运算 dataloader_drop_last=True, logging_steps=10, max_grad_norm=None, ) trainer = Trainer( model=model, args=training_args, train_dataset=unsupervised_datasets["train"], eval_dataset=unsupervised_datasets["validation"], compute_metrics=lambda p: {"loss": p.losses.mean()}, data_collator=lambda examples: { **tokenizer([ex['sentence'] for ex in examples], padding='max_length', truncation=True), 'labels': None}, optimizers=(None,None)) trainer.train() ``` #### 总结 上述代码展示了如何基于 Hugging Face 工具链完成 SimCSE 的复现流程。需要注意实际应用时可能还需要调整超参数以适应具体场景需求。