19ContraBERT：顶会ICSE23 数据增强+对比学习+代码预训练模型，提升NLP模型性能与鲁棒性：处理程序变异（变量重命名）【网安AIGC专题11.15】

写在最前面

随着大规模代码的崛起，无监督学习成为了提高代码预训练模型性能的有效手段。这些预训练模型在广泛的下游任务中表现出色，如自然语言处理和程序语言处理。例如，像CodeBERT和GraphCodeBERT这样的模型在预训练阶段通过大规模代码数据学到通用的表示，并在下游任务上进行微调，取得了优于传统监督学习方法的成绩。然而，这些模型在面对代码变异等挑战时，鲁棒性仍然有待提高。

该论文关注的问题是：现有的代码预训练模型是否能够学到鲁棒的代码表示？研究表明，这些模型在处理程序变异，如变量重命名等情况时表现不佳。为了提高模型的鲁棒性，作者提出了一种新的方法，称为ContraBERT。

本文为邹德清教授的《网络安全专题》课堂笔记系列的文章，本次专题主题为大模型。

李仕昕同学分享了ContraBERT: Enhancing Code Pre-trained Models via Contrastive Learning
《ContraBERT：通过对比学习增强代码预训练模型》
关于提升代码预训练模型鲁棒性的文章，发表在ICSE2023，作者单位是南洋理工大学Nanyang Technological University, Singapore
分享时的PPT简洁大方，重点突出

论文：https://dl.acm.org/doi/10.5555/3618408.3620098
ICSE（International Conference on Software Engineering）是软件工程领域的顶级国际会议之一。ICSE会议每年都吸引来自世界各地的顶尖研究人员和学者，以展示和讨论最新的软件工程研究成果。

对nlp领域其他方向研究的启发

《ContraBERT: Enhancing Code Pre-trained Models via Contrastive Learning》侧重于增强代码预训练模型的鲁棒性，提供了一些有益的启发可以应用到自然语言处理（NLP）领域的其他方向研究中：

采用对比学习、数据增强和多任务学习等方法以提高模型性能和鲁棒性的示例

1. 对比学习的应用：这篇论文通过对比学习方法增强了代码预训练模型的性能。
   可以尝试将对比学习引入到各种任务中，以提高模型的性能和鲁棒性。

2. 数据增强的重要性：数据增强是提高模型性能和鲁棒性的关键因素之一。
   通过设计合适的数据增强方法，可以生成语义等价的样本，以提高模型的泛化能力。
   这种方法可以应用于各种NLP任务，如文本分类、命名实体识别和机器翻译等。

3. 多任务学习：ContraBERT采用了多任务学习的策略，将掩码语言模型（MLM）和对比学习结合在一起。
   这种方法可以用来训练更通用的NLP模型，以适应多种下游任务。可以尝试将多任务学习引入到NLP任务中，以提高模型的效率和性能。

4. 预训练任务的设计：这篇论文的成功经验表明，通过设计具有挑战性的预训练任务，可以帮助模型更好地捕捉语言的特征。
   可以探索不同的预训练任务，以进一步改进NLP模型的性能。

介绍

大代码时代的到来使得目前有很多工作开始利用无监督学习，来通过大量的代码数据学习到通用的预训练模型，然后再针对不同的下游任务对预训练模型进行微调。这些预训练模型对不同的下游任务有比较强的泛化性，并且得到了广泛的使用。
比如预训练模型CodeBert和GraphCodeBert都是在数据集codesearchnet上进行预训练的模型，然后在下游任务上进行微调，达到了优于监督学习的方法。
（CodeBert在自然语言NL和程序语言PL对上进行学习，(结构：多层双向transformer),GraphCodeBert利用代码的语义结构来学习代码表征，利用了代码的数据流来进行编码。）

文章针对这两个预训练模型做了一个鲁棒性的初步研究，结论是发现这些预训练模型对于变量重名名这种程序变异都不够鲁棒。
具体实验是使用了克隆检测的数据集，将程序中的多个变量随机重命名，然后计算程序在变量重命名前后预训练模型输出的embedded vectors的余弦相似度，从这个图中可以看出，在重命名的变量数从1增加到8时，红线和蓝线代表的两个预训练模型对于程序变异前后的输出相似度由1降到了0.4，说明他们对这种对抗样本是十分敏感的。

为了提高预训练模型鲁棒性，文章提出的方法ContraBert
如上图所示，上面两条线就是文章提出的方法应用在codebert和grapgcodebert上的结果，方法的鲁棒性是优于原始模型的。

Contrabert

ContraBERT采用了一种基于对比学习的无监督学习框架，旨在增强现有代码预训练模型的鲁棒性。该方法结合了两个预训练任务：掩码语言模型（MLM）和对比学习。

• MLM任务旨在通过将一部分句子中的标记进行掩码，然后尝试从剩余部分还原这些标记，从而学习更好的标记表示。
• 对比学习通过最小化相似样本之间的距离，同时最大化不同样本之间的距离来学习表示，以提高模型的鲁棒性。
  对比学习来提高模型鲁棒性背后的直觉就是：程序变异前后作为相似样本进行拉进，使得模型对程序变异这种对抗攻击更加鲁棒。

方法Method

ContraBERT

• MLM is utilized to help the model learn better token representations
• Contrastive is utilized to help the model group the similar vector representations to enhance model robustness.

ContraBERT的训练过程包括两个关键任务：MLM和对比学习contrastive learning。

1. 首先，原始输入样本对以及由数据增强算法生成的增强样本对被送入模型。
2. 然后，模型执行MLM任务，其中一组标记被随机掩码，模型需要预测被掩码的标记。
3. 接下来，模型执行对比学习任务，通过最小化相似样本之间的距离，同时最大化不同样本之间的距离，以增强模型的鲁棒性。
4. 最后，MLM和对比学习两个任务的损失被加权相加以进行训练。

对比学习是需要拉进相似样本的，所以使用对比学习的时候需要用数据增强在原始样本的基础上构造一系列的相似样本。

可以从这个流程图看到：

1. 原始的输入C，W这种函数和注释样本对在进行训练前，会使用文章定义的一系列程序语言和自然语言的数据增强算法f和g来构造c和w的增强样本，
2. 然后使用原始样本对和增强样本对这个四元组作为输入，
3. 之后使用两个预训练任务MLM和contrastive对给定的模型进行训练。

所以整个方法具体分为两个部分，

1. 首先，是如何设计的这个程序语言和自然语言的数据增强操作f和g，
2. 第二步，是如何设计Contrabert模型的训练。

数据增强和训练细节

数据增强在ContraBERT中发挥关键作用，以提高模型鲁棒性。
对于程序语言（PL）和自然语言（NL）的数据增强，作者提出了一系列操作，包括函数重命名、变量重命名、插入未使用的语句、随机交换函数中的语句和随机删除函数中的语句。
这些增强方法旨在构造语义等价的变种样本，以供对比学习使用。

首先介绍数据增强操作f和g，数据增强主要是为了在原始样本的基础上构造语义等价的变体来提供给对比学习这个任务
给定一个函数和注释的样本对
首先针对程序语言PL的数据增强，文章提出了五种增强操作，
第一个是函数重命名，将函数名称进行替换（有额外词汇集，从预训练数据集中提取），
第二个是变量重命名，将每个函数里的变量进行随机重命名，
第三个是在函数中插入未使用的语句，构造一个没有在函数中出现的变量的赋值语句插入函数中
第四个是随机交换函数中没有相互依赖关系的语句（遍历 AST 并分析数据依赖性以进行提取。）
第五个是随机抽取函数里的一行语句进行删除，这个主要是作为正则化来避免训练的过拟合（？）
这些增强方法都可以构造出语意等价的相似样本，被视为原始样本的正样本。

针对自然语言的数据增强主要是

• 一个是反向翻译，就是将源语句翻译成另一种语言，然后将翻译后的语言再转换为原语言，比如文章就是讲英语转换成德语在翻译回英语。
• 除此之外还有一些简单的如随机删除注释里的一个单词，随机交换注释里的两个单词，或者是复制一个单词插入注释里面。

文章会使用这些增强方法讲PL-NL对进行增强形成一个增强集合，然后在每个训练step中随机选择一对出来进行训练。

（和之前的工作比增加了自然语言的数据增强）

模型的训练需要两个编码器M和M’，它们具有相同的模型结构和初始参数。在训练过程中，原始输入样本对和增强样本对被用于执行MLM和对比学习任务。MLM任务通过随机掩码标记来训练模型，而对比学习任务则使用对比学习损失函数（infonce loss）来最小化相似样本之间的距离，同时最大化不同样本之间的距离。

将样本进行数据增强后，就是正式进行模型训练了。

具体设计细节如图，contrabert需要两个encoder
M和M’，他们的模型结构一样，并且初始参数相同，文章具体会使用codebert或者是graphcodebert

• 每一个样本对c，w通过数据增强得到四元组(c,w,c’,w’)后，构建两个输入序列，一个是原始样本对序列x和增强样本对序列x’,其中用CLS标识符来标记序列开始，SEP来标记链接自然语言和程序语言两种序列的符号。

• 构造了输入之后,整个训练过程有两个训练任务，首先是MLM，在图中的左半部分，训练时将原始输入序列X中一组随机位置的token mask掉然后训练模型能够学到被mask掉的标记，这个训练任务为了让模型学到更好的特征表示，MLM是被广泛使用的训练任务，我就不过多介绍了。

• 第二个训练任务就是对比学习，对比学习的损失函数使用的事infonce，可以看到这个函数表示，q指的是图中的query，是原始样本序列X的encoder的[cls]位置的向量表示，k+指的是增强样本序列X’encode的向量表示，被视为原始样本的正样本，i=1到n的ki值得是其他样本的增强样本，被视为原始样本的负样本，通过infoloss这个公式可以将分子上原始样本和正样本拉进，分母上原始样本和其他负样本拉远，使模型能够将原始样本所有的变异样本都学到相似的语义特征。
  该文章参考了对比学习的一个有代表性的方法moco，设置了一个队列来储存之前batch的负样本来动态更新进行学习。

• 最后MLM和对比学习两个loss进行加权相加来进行训练。在训练过程中M使用梯度下降进行参数更新，M’使用动量来更新weight

模型训练好后，就可以根据不同的下游任务进行微调了，文章关注了三类下游任务
第一类是检索任务，包含克隆检测和代码搜索（根据自然语言query所需程序），第二类是分类任务，如缺陷检测，第三类是生成任务，比如代码到代码的翻译。
可以根据不同任务添加特定模块进行微调

实验Experiment

作者从四个方面来评估了他们的方法，

RQ1 ：Robustness Enhancement鲁棒性提升

Performance of different augmentation operators in enhancing the robustness of the pre-trained model

在第一个实验中，作者评估了不同增强操作对模型鲁棒性的影响。实验结果表明，通过使用设计的PL-NL数据增强方法，模型的鲁棒性得到了显著提升。将这些增强方法组合使用可以获得最佳效果。

第一个方面是提出的不同的数据增强算法对提高模型鲁棒性的表现。
实验通过检测contrabert对克隆检测的对抗攻击的鲁棒性来评估（因为克隆检测也是从其他干扰项中识别出语义等效的样本）
如表所示，contrabert在训练模型时使用了提出的所有增强方法，所以这个表每一行分别去掉一个提出增强方法来观察不同增强方法带来的鲁棒性提升，具体做法是将模型能够正确预测的样本（num）进行随机次数的变异来观察模型的预测准确性。
可以看到设计的 PL-NL 数据增强每一个方法都对模型鲁棒性有帮助（纵向对比，每一个和最后一行对比），当这些增强合并时，预训练模型的鲁棒性能获得最佳效果。

存在问题：没有给对比实验，就是提供直接用codebert和graphcodebert在对抗攻击下的鲁棒性进行参考。

RQ2: Visualization for Code Embeddings代码嵌入可视化

第二个实验通过可视化ContraBERT的向量空间来观察模型是否学到了更好的向量表示。结果显示，ContraBERT学习到的向量空间更加紧凑，能够更好地区分不同问题的程序表示，从而提高了模型的性能。

第二个实验室可视化contrabert的向量空间来观察此方法是否学到了更好的向量表示
从图中我们可以看出，contrabert比单纯的codebert和graphcodebert学习到的向量空间更加紧凑，代表同一个问题的程序表示更想死，不同问题之间也隔的更开。所以ContraBERT 能够对语义等效的样本进行分组并排除不相似的样本，从而学习更好的向量表示。

（也是使用了克隆检测，可以观察样本表示的相似度，该数据集由 104 个编程问题组成，其中每个问题都有 500 个语义等效的程序，但具有不同的实现。随机选择 5 个不同的问题和 100 个样本）

RQ3: Performance of ContraBERT on Downstream Tasks下游任务性能

第三个实验评估了ContraBERT在四个下游任务上的性能，包括克隆检测、缺陷检测、代码翻译和代码搜索。实验结果表明，ContraBERT在这些任务上取得了显著的性能提升，特别是在克隆检测和缺陷检测任务中表现出色。

表2是在克隆检测和缺陷检测上的结果，表3是代码翻译，表4是代码搜索

可以观察到ContraBERT 全面提升了原始 CodeBERT 和 GraphCodeBERT 在四个下游任务上的性能，我们将这些改进归因于模型鲁棒性的增强，在这些任务上具有更好的性能。
同时也可以看到这些结果都展示了模型分别使用两个预训练任务MLM和contrastive对下游任务的影响，可以看出两个任务结合能获得最好的效果。
（代码搜索提升不高是任务本身难度，程序和自然语言中的映射是很难提升的）
（MAP@R：评估在给定查询的情况下检索集合中R个最相似样本的结果。 MAP@R 用于克隆检测，其中 R 设置为 499 进行评估。）
（BLEU-4:生成的序列与真实值之间的文本相似性。）
（MRR：平均倒数排名，评估检索系统性能，第一个结果匹配，分数为1，第二个匹配分数为0.5，第n个匹配分数为1/n）

RQ4: Ablation Study for Pre-training Tasks预训练任务的消融研究

作者还进行了预训练任务的消融研究，验证了MLM和对比学习任务的重要性，结果表明两者的结合能够获得最佳效果。

讨论

尽管ContraBERT在提高代码预训练模型的鲁棒性方面取得了显著的成果，但仍存在一些潜在限制。例如，作者采用了与计算机视觉任务中的对比学习框架相似的参数设置，这可能不是最适用于代码模型的参数。此外，作者在评估模型鲁棒性时只使用了克隆检测任务，未提供与传统模型的对比实验。

然而，这项工作为提高代码预训练模型的性能和鲁棒性提供了有价值的思路。未来的研究可以进一步探索更复杂的数据增强方法，以进一步提高模型的鲁棒性，并改进对代码搜索任务的性能提升方法。

通过对比学习和数据增强，ContraBERT为提高代码预训练模型的性能和鲁棒性提供了一种新的方法，为代码领域的研究和应用提供了有力支持。

这篇文章主要就是提倡在训练代码模型的时候除了考虑模型在下游任务上的表现还要考虑模型的鲁棒性。
这篇文章在实现上还是存在一些limitation
这篇论文参考的对比学习框架MOCO，MOCO是为计算机视觉中的图像分类任务设计的，而这篇文章的参数照搬了moco的原始设置，可能这些参数对于代码模型来说不是最佳的。
第二点是在实验时评估模型鲁棒性的时候只使用克隆检测任务来说明问题（POJ-104），并且没有设置baseline对比实验。

关于工作的思考，可以发现在对下游任务的提升中，contrabert对代码搜索的下游任务的提升相对于其他下游任务的性能提升来说是比较微小。
文章关于这个现象的解释是，对于代码搜索任务，由于模型需要学习查询与其相应代码之间的语义映射，是比较难的，然而设计的增强算法都仅仅修改代码或查询语言本身，无法捕获它们的相关性，这导致提升受到限制。
所以对于代码搜索任务，进一步提高性能的一个可能的解决方案是为增强变体建立 PL 和 NL 之间的标记关系，然而，这需要大量工作来分析程序和自然语言注释之间的关系。

除此之外，文章证明了提出的数据增强能够提高预训练模型的鲁棒性，后续还可以探究更多复杂的数据增强，使得模型鲁棒性有更大的提升的空间

---

希望这篇博客对您理解《ContraBERT: Enhancing Code Pre-trained Models via Contrastive Learning》论文提供了有益的信息。如果您对该论文或相关领域有更多疑问，欢迎随时交流。