使用Torch实现CoSENT实验

这周五最后的时间看完了苏神刚更新的CoSENT，主要是为了解决目前Sentence-BERT等模型训练和预测目标不一致的问题，模型细节可以参考苏神的博客。本着学习的态度，用ark-nlp复现了一下苏神的实验，并使用CHIP-STS测试该模型在医疗场景下的效果。

代码地址：https://github.com/xiangking/PyTorch_CoSENT

实验数据

• ATEC、BQ、LCQMC和PAWSX：https://github.com/bojone/BERT-whitening/tree/main/chn
• CHIP-STS（平安医疗科技疾病问答迁移学习）：https://tianchi.aliyun.com/dataset/dataDetail?dataId=95414

CoSENT在形式上和无监督的simcse还是比较相似的，由于笔者在实现无监督的simcse时是通过句子分别通过模型的形式，所以大部分的结构都与CoSENT，更多的更改在损失这一块：

# bert_embedding其实就是句子输入BERT后生成的向量经过池化或直接使用CLS对应的向量后的表示，
bert_embedding_a = self.get_pooled_embedding(
    input_ids_a,
    token_type_ids_a,
    position_ids_ids_a,
    attention_mask_a
)

bert_embedding_b = self.get_pooled_embedding(
    input_ids_b,
    token_type_ids_b,
    position_ids_b,
    attention_mask_b
)

cosine_sim = torch.sum(bert_embedding_a * bert_embedding_b, dim=1) * 20
cosine_sim = cosine_sim[:, None] - cosine_sim[None, :]
    
labels = label_ids[:, None] < label_ids[None, :]
labels = labels.long()
    
cosine_sim = cosine_sim - (1 - labels) * 1e12

# 上面的代码对照公式便可理解，该部分添加0，主要是为了防止例如[-1e12, -1e12, -1e12]经过logsumexp的问题
cosine_sim = torch.cat((torch.zeros(1).to(cosine_sim.device), cosine_sim.view(-1)), dim=0)
loss = torch.logsumexp(cosine_sim.view(-1), dim=0)

参数设置

代码参数设置如下：

句子截断长度：64（PAWSX数据集截断长度为128）
batch_size：32
epochs：5

效果

使用spearman系数作为测评指标，ATEC、BQ、LCQMC和PAWSX使用test集进行测试实验，CHIP-STS则使用验证集

	ATEC	BQ	LCQMC	PAWSX	CHIP-STS
BERT+CoSENT（ark-nlp）	49.80	72.46	79.00	59.17	76.22
BERT+CoSENT（bert4keras）	49.74	72.38	78.69	60.00
Sentence-BERT（bert4keras）	46.36	70.36	78.72	46.86

PS：上表ark-nlp展示的是5轮里最好的结果，由于没有深入了解bert4keras，所以设置参数可能还是存在差异，因此对比仅供参考

针对CHIP-STS测试数据集，选择阈值为0.6生成结果进行提交，结果如下：

	Precision	Recall	Macro-F1
BERT+CoSENT	82.89	81.528	81.688
BERT句子对分类	84.331	83.799	83.924

致谢

感谢苏神无私的分享