论文笔记:CodeRetriever: Large-scale Contrastive Pre-training for Code Search

833 阅读6分钟

导语

1 简介

代码搜索(Code search)旨在检索给定自然语言查询该功能的相关代码,现有的Code预训练的方法多关注于token级别的MLM任务,这导致了以下两种不足:

  1. 各向异性的表示问题。token级别的自训练方法导致高频token的嵌入聚类并占据表示空间,极大地限制了预训练模型中长尾低频token的表达能力;
  2. 跨语言表示问题。由于混合编程语言的代码很难出现在相同的上下文中,因此对预训练的模型来说,学习具有相同功能但使用不同编程语言的代码的统一语义表示是一个挑战。

本文提出CodeRetriever模型,专注于学习函数级别的代码表示,特别是用于代码搜索场景。该模型由一个文本编码器和一个代码编码器组成,将文本/代码编码为单独的向量。代码与文本(或代码与代码)之间的语义相关性通过向量之间的相似度来衡量。

在CodeRetriever的训练中,通过最小化两种类型的对比损失来优化码/文本编码器。

  • 单模态对比损失,鼓励模型将功能相似的代码在表示空间中推得更近。
  • 双模态对比损失,有助于建模代码和文本之间的相关性。

本文采用常用的CodeSearchNet语料库来训练CodeRetreiver。CodeSearchNet主要包括成对数据集(与文档配对的函数)和非成对数据集(只有一个函数)。配对数据集可以直接用于双模态对比学习。对于CodeRetriver中的单模态对比学习,采用无监督语义引导方法构建正码解码对。图1(a)显示了一个code-code对示例。此外,生成的code-code对可以使用不同的编程语言,这可以减轻跨语言表示问题。为了进一步利用未配对数据和配对数据中的大规模代码,作者提取了代码和内联注释对,以增强CodeRetriever中的双模态对比学习。图1(b)显示了一个示例:内嵌注释也可以反映代码的语义和内部逻辑。

通过对比这些单模态和双模态对,CodeRetriever可以更好地学习函数级代码语义表示,缓解各向异性表示问题,同时 显式地建模不同编程语言的代码的相关性,并将统一的自然语言作为一个支点,以减轻跨语言表示问题。本文在11个代码搜索数据集上评估了CodeRetriever,这些数据集涵盖了6种编程语言、真实场景和不同粒度的代码(函数级、代码段级和语句级),结果表明CodeRetriever实现了最新的性能。

2 预备知识:Code Search

Code Search旨在通过一段自然语言描述,找到与其最相符合的Code片段的任务。

3 方法

CodeRetriever采用Siamese Code/Text编码器架构,将代码/文本表示为密集向量。EcodeE_{code}EtextE_{text}分别表示代码编码和文本编码器。code-code对(c, c+)和text-code对(t, c+)的语义相似度计算如下:

3.1 单模态

单模态的损失函数为:

其中 τ\tau设为1,为温度参数。

3.2 双模态

双模态损失函数与单模态相仿,

3.3 整体训练目标

整体训练的loss即两部分损失相加,其中,双模态部分不仅有code-document的,也有code-comment的损失函数,所以整体目标函数为3部分:

4 正样本对的建立

使用对比学习的重要一点是如何选取正负样本对,对于双模态部分,这个比较好选择,直接选取Code和它对应部分的文档说明或者注释即可。而对于单模态,由于没有这种标注信息,则需要使用一些预处理方法来进行设置。主要步骤如下:

  1. 通过匹配函数名称和文档内容来收集有噪声的Code-code对。主要方法如下:分别使用函数名和文档内容作为语料训练两个SimCSE模型NameMatcher、DocMatcher(这是一个检索模型,给出一个函数名,NameMatcher将返回Top K个语料中最相关的函数名)。设定阈值为0.75,超过这个阈值的就认为是正样本对。
  2. 使用CrossModel进行去噪。1中得到的正样本对是具有很多噪声的,因此作者又使用1中通过Doc得到的样本对作为训练集来训练一个二分类模型,这个模型将一对代码进行输入,返回他们的语义相似程度。作者使用这个模型对1中所有的正样本对进行去噪过滤。

5 实验

采用GraphCodeBERT作为Backbone,基准数据集采用CodeSearch、Adv、CoSQA、CoNaLa、SO-DS和StaQC。评估指标采用MRR。在微调阶段,对比三种不同的方式:

  • In-Batch Negative。即选取同一个Batch中的其他样例作为负例;
  • Hard Negative。即选取更加“困难”的样例作为负例;
  • AR2。采用对抗学习的方法选择的“困难”样例作为负例。

实验结果如下图所示:

作者还对比了Low-resource下的实验结果:

并展示了CodeRetriever模型成功解决了文章开始提出的各向异性(图5)和跨语言表示(图4)问题。

作者还进行了消融实验(表3)发现,使用不去噪的码解码对进行单模态对比学习会带来轻微的性能下降,而使用去噪,它可以获得显著的性能提升。这证明了去噪步骤的有效性,并表明单模态对比学习依赖于正对构建的质量。从使用doc-code和comment-code进行双模态对比学习的结果来看,该模型的性能得到了进一步的提高,说明双模态对比学习可以利用文档或注释中的关键语义信息来帮助更好地理解代码。

6 相关工作

7 总结

本文引入了将单模态和双模态对比学习相结合的CodeRetriever作为代码搜索的预训练任务。对于单模态对比学习,作者提出了一种语义引导方法来构建正的代码对。对于双模态对比学习,本文利用文档和内嵌注释构建正文本代码对。在几个公开基准上的大量实验结果表明,所提出的CodeRetriever带来了显著的改进,在所有基准上都达到了最新的水平。进一步的分析结果表明,CodeRetriever在低资源和跨语言代码搜索任务中也具有强大的功能,并证明了单模态和双模态对比学习的有效性。