ContraBERT: Enhancing Code Pre-trained Models via Contrastive Learning

全文摘要

本文介绍了一种名为ContraBERT的方法，旨在通过对比学习提高代码预训练模型的鲁棒性。由于现有的大规模预训练模型在代码表示方面表现出色，因此它们被广泛应用于下游任务中，如克隆检测、代码搜索和代码翻译等。然而，这些先进的预训练模型也容易受到简单的攻击，例如变量名重命名等。这种弱点可能会被下游模型继承并放大到前所未有的规模。为了解决这个问题，作者设计了九种简单和复杂的编程语言和自然语言数据增强操作，构建不同的变体，并使用掩码语言建模（MLM）和对比预训练任务对现有预训练模型进行继续训练，以增强其鲁棒性。实验结果表明，ContraBERT可以有效地提高现有预训练模型的鲁棒性，并且在四个流行的下游任务上提供了改进。

方法描述

该论文提出了一种名为ContraBERT的预训练模型，用于程序代码和注释的表示学习。其主要思路是通过程序代码和注释之间的对比来增强模型的鲁棒性和泛化能力。具体来说，他们设计了一系列的程序代码和注释增强操作，包括简单的删除、替换等以及复杂的插入死代码、重排语句等操作。然后将这些增强后的样本与原始样本一起输入到预训练模型中进行训练，以提高模型的表现。

方法改进

相比于之前的CodeBERT和GraphCodeBERT，ContraBERT在预训练过程中引入了对比学习任务，即通过对比相似和不相似的数据来学习更好的特征表示。此外，他们在数据增强上也做了更多的工作，设计了一些更复杂的增强操作，如插入死代码、重排语句等，从而提高了模型的泛化能力和鲁棒性。

解决的问题

ContraBERT的目标是提高程序代码和注释的表示学习效果，从而提高下游任务（如代码搜索、缺陷检测、代码翻译等）的表现。传统的预训练模型只考虑了单个样本的学习，而ContraBERT则通过对样本进行增强和对比学习，使得模型能够更好地理解不同样本之间的关系，并从中学习到更好的特征表示。这种方法已经在多个下游任务上进行了验证，并取得了比传统方法更好的表现。

论文实验

本文主要介绍了使用增强模型鲁棒性的方法来提高预训练模型在下游任务上的性能的实验研究。该研究使用了四种不同的评估指标（MAP@R、Acc、BLEU-4和MRR）来评估预训练模型在四个不同下游任务（克隆检测、缺陷检测、代码翻译和代码搜索）上的表现，并将结果与原始的CodeBERT和GraphCodeBERT进行了比较。

首先，作者通过验证样本对攻击的准确性来探究增强模型鲁棒性的效果。他们选择了一个零射程设置下的克隆检测任务来进行实验，发现所有增强模型都比原模型表现更好，这表明他们的设计是有效的。

其次，作者还使用T-SNE可视化技术来验证预训练模型是否能够更好地学习程序表示空间。结果显示，增强模型比原模型更能够聚类相似的数据并推开不相似的数据，从而更好地学习程序表示空间。

最后，作者还进行了大量的下游任务实验来比较增强模型和原模型的表现。结果表明，增强模型在克隆检测、缺陷检测和代码翻译等任务上表现更好，但在代码搜索任务上表现略微优于原模型。此外，作者还进行了掩码语言建模和对比学习任务的单独实验，发现这两种任务在增强模型中都有其独特的作用，结合起来可以进一步提高模型的性能。

综上所述，本文的研究结果表明，增强模型鲁棒性是一种有效的方法来提高预训练模型在下游任务上的性能。这种方法可以通过增强模型的鲁棒性来更好地学习程序表示空间，并且可以在多个下游任务上取得显著的改进。

论文总结

文章优点

• 论文提出了一种新的预训练模型ContraBERT，通过对比实验验证了其在增强代码语义表示鲁棒性方面的有效性。
• 作者设计了九种简单的或复杂的编程语言和自然语言序列的数据增强操作，并证明每种操作都对提高模型的鲁棒性有帮助。
• 在四个下游任务上的广泛研究显示，这些增强鲁棒性的模型相对于原始模型提供了显著的改进。

方法创新点

• ContraBERT采用了无监督的对比学习框架来增强现有预训练模型的鲁棒性。
• 作者设计了九种简单或复杂的数据增强操作，包括对程序和注释的变换等，以帮助模型更好地学习相似样本并区分不同的样本。
• ContraBERT同时利用了掩码语言建模（MLM）和对比学习预训练任务，提高了模型的表现。

未来展望

• 可以进一步探索更复杂的增强策略，例如在多个操作上构建更复杂的增强变体。
• 对于内部有效性和外部有效性等方面的问题，可以考虑使用更多的数据集和评估指标来进行深入的研究。
• 该方法还可以应用于其他领域，如自然语言处理等。