ContraBERT: Enhancing Code Pre-trained Models via Contrastive Learning

全文摘要

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

方法描述

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

方法改进

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

解决的问题

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