On Contrastive Learning of Semantic Similarity for Code to Code Search

Anthony Saieva, Saikat Chakraborty, Gail Kaiser,

Columbia University, Microsoft Research, Columbia University

引用

Saieva A, Chakraborty S, Kaiser G. On Contrastive Learning of Semantic Similarity forCode to Code Search[J]. arXiv preprint arXiv:2305.03843, 2023.

论文：https://arxiv.org/abs/2305.03843

仓库：https://github.com/reinforest-team/REINFOREST

摘要

本文介绍了一种新颖的代码到代码搜索技术，通过结合静态和动态特征，并在训练过程中利用相似和不相似的示例，提高了大型语言模型（LLMs）的性能。该方法在训练过程中编码动态运行时信息，无需在推断时执行搜索，同时也是第一种在正负参考样本上训练的代码搜索技术。研究表明，这种方法在跨语言代码搜索方面表现出色，优于现有工具高达44.7%。重要的是，作者强调了模型微调的重要性，并提供了一个名为COSCO的开源实现工具，以确保研究的可重现性和可扩展性。

1 引言

随着自BERT论文发布以来强大的大型语言模型（LLM's）的崛起，近几年来，将机器学习应用于源代码分析任务变得越来越流行。同时，随着大量公共代码存储库的出现，代码到代码搜索的潜力在常见软件维护和开发任务中变得更加重要，如代码迁移、转译、代码修复、错误检测、教育和重构。

本文提出了一系列新颖的技术，用于增强LLMs进行代码到代码搜索和相关任务。我们介绍了一种新的代码到代码搜索技术，可以在训练过程中对源代码和动态运行时信息进行编码，无需在推断时执行搜索语料库或搜索查询中的任何内容。我们展示了这一方法的基本原理。在训练过程中，模型从训练语料库中学习静态和动态相似性，最大化不相似代码之间的距离，最小化相似代码之间的距离。我们生成了一个称为语义相似性分数（SSS）的单一动态相似性特征，在训练之前执行尽可能多的训练语料库，并比较具有相同输入类型的函数的输出。这是唯一执行代码的时刻。然后，将该模型用于创建搜索语料库的预先计算嵌入，最后在推断时，模型将搜索查询嵌入，从而简单比较预先计算的嵌入以得出搜索结果。传统的动态分析技术要求在训练期间执行代码，以及在搜索过程中执行搜索语料库和搜索查询。这会导致在实践中运行代码时出现的所有与环境相关的问题。相比之下，我们的方法只增加了在受控环境中执行训练集代码的开销，一旦模型分发，就无需再次执行代码。值得注意的是，我们的方法仍然适用于在训练集的部分或全部不可执行的情况。

为了验证我们方法的有效性，我们考虑了跨语言搜索问题，其中查询和语料库来自不同语言，就像COSAL一样。我们认为这将最好地展示我们技术的有效性，因为我们最关心行为相似性，跨语法不同的语言进行搜索意味着模型不能依赖源代码相似性。

我们进行了一系列消融研究，展示了LLMs在各种情况下推断动态行为的能力，即使推断仅基于静态信息。这些情况包括多个查询和语料库语言、各种基础LLM模型，包括开源和专有的模型，是否有动态运行时信息，以及在训练中正负参考样本的数量变化。我们的评估显示，我们的技术在Atcoder数据集上比最先进的跨语言搜索技术表现出高达44.7%的显著改进，该数据集包含361个问题对应的Java和Python解决方案，来自83个编程竞赛，并且在各种LLM架构中，包括CodeBERT和OpenAI的Codex模型中，保持着接近或更好的性能，即使LLMs无法进行微调。值得注意的是，我们发现一个较小的经过微调的LLM明显优于更大的专有LLMs，突显了开源模型对研究社区的价值和重要性。此外，我们发现只要在训练过程中至少有一个正参考和一个负参考可用，我们的技术与基准相比表现出色，允许在实践中可能遇到的更稀疏的数据集，并展示了考虑相似和不相似参考样本的重要性。

我们的贡献如下：

首次提出了一种代码到代码搜索技术，在训练期间编码动态运行时信息，无需在推断过程中执行搜索语料库或搜索查询中的任何代码。首次提出了一种在训练期间考虑相似和不相似参考样本的代码到代码搜索技术。进行了一系列消融研究，展示了LLMs在跨语言搜索中推断动态行为的能力，即使推断仅基于静态信息。评估结果显示我们的技术在许多模型架构中仍然有效。消融研究表明，即使在训练中只有一个正参考和一个负参考样本，也能实现有效的性能。评估结果显示，精心调整的模型优于更大的未经微调的LLM模型，即使是最大的LLMs。我们开源了一个名为COSCO的工具和训练流程的实现，代表了跨语言代码搜索的新技术水平。

2 技术介绍

我们将COSCO描述为两个不同的步骤 - (a) 训练和 (b) 在线查询嵌入和搜索。在代码层面，COSCO包括两个不同的编码器：查询代码编码器（Eq）和文档代码编码器（Ed）。每个编码器在给定输入代码x的情况下生成x的固定大小的向量表示。

图 1：训练和推理过程概述

A. COSCO的训练

我们训练COSCO的目标是优化两个不同的正交目标。首先，我们希望COSCO能够转换由查询代码编码器（）和文档代码编码器（）生成的嵌入，以最小化相关代码之间的距离，并最大化不相关代码之间的距离。我们的第二个目标是在训练期间为COSCO提供动态代码相似性信息，以便在推断（搜索）期间我们不必计算查询代码与数据库中所有其他代码之间的语义（I/O）相似性。图2展示了COSCO训练过程的概述。

图 2：COSCO训练流程

对比训练。为了实现第一个目标，我们从对比学习中汲取灵感。我们假设存在一个数据集D，其中包含一组元组。是用源语言s（例如Java）编写的一段代码。和是用目标语言（例如Python）编写的代码集。集合是与解决相同问题的解决方案。相反，是与解决不同问题的解决方案。我们将代码样本称为正样本，代码样本称为负样本。直观地说，我们希望将中每个元素的嵌入带到的嵌入附近，并希望将的元素的嵌入远离的嵌入。

首先，我们使用查询代码编码器对进行编码，使用文档代码编码器对和集合进行编码。形式上，，，,其中是查询代码的嵌入，和分别是对应于正样本和负样本的嵌入集合。然后，我们计算查询代码与每个正样本和负样本之间的余弦相似度。我们将两个向量a和b的余弦相似度计算为：

我们的下一个目标是同时最大化与中所有元素之间的相似性，并最小化与中所有元素之间的相似性。我们通过最小化以下损失函数来实现这样的目标：

这里，和是正样本和负样本的目标标签。和的值应取决于系统中使用的相似性函数。对于无界相似性函数，和应分别设置为1和0。在COSCO的情况下，我们设置= 1和= 0，因为这些值是我们使用的相似性函数的最大值和最小值（方程（1））。

语义相似性分数。作为我们的第二个训练目标是为模型提供有关语义（I/O）相似性的知识，我们假设一个函数，计算和之间的语义相似性。我们通过修改方程（2）中的损失函数，为模型赋予这样的语义相似性：

这里，代表语义训练的相对重要性。在理想情况下，我们可以计算所有和对之间的语义相似性，我们认为将设置为1.0会产生最佳性能。然而，假设对于所有这样的和对，存在是不切实际的。因此，我们通过将设置为0.2来在实践中找到最佳验证性能。无论如何， 在COSCO的实现中仍然是一个用户定义的超参数。

图 3：相似性评分过程

用于训练的语义相似性分数计算。COSCO的架构和过程在训练期间使用基于运行时信息的语义相似性分数计算。我们使用SLACC的相似性评分方法，即匹配输出的数量除以输入的数量。图3概述了我们生成相似性分数的两个过程步骤。首先，如图3a所示，我们从代码语料库中生成输入语料库。对于任何一段代码，我们静态分析系统调用并推导出该特定代码的输入结构。在提取语料库中所有输入结构之后，我们为每个结构基于原始类型生成大量随机值的输入，并创建一个输入语料库。然后，如3b中所示，对于任意两个样本，我们提取输入结构，使用与其输入结构匹配的输入运行它们，然后执行前面描述的相似性计算。与SLACC一样，如果两段代码具有不同的输入结构，它们将获得相似性分数为零。

B. 使用COSCO 进行代码搜索

一旦训练过程结束，我们就有了两个经过训练的编码器，和，用于分别编码查询和文档。训练结束后，我们立即使用对搜索数据库中的所有代码进行编码。由于编码后的表示是实值向量，索引工具如 FAISS 可以将它们存储在索引数据库中，以便进行高效的搜索。对于每个查询代码，我们使用查询编码器生成查询嵌入。然后，我们根据方程（1）中的相似性分数在嵌入数据库中进行搜索，并返回具有最高相似性分数的前 n 个候选项。

3 实验评估

3.1 实验设置

研究问题。在本文中，我们研究以下研究问题：

RQ1：COSCO 的性能如何与其他跨语言代码搜索技术的性能相比？

RQ2：COSCO 的方法论和性能是否可以推广到不同的模型？

RQ3：在训练过程中包含语义相似性分数是否会提高代码搜索的效果？

RQ4： 改变可用于训练的正样本和负样本比较样本数量如何影响 COSCO 的性能？

评估数据集。为了评估 COSCO，我们使用了 Atcoder 数据集，该数据集包含18644个Java解决方案和22317个Python解决方案，对应于361个编程竞赛问题。为了验证我们的概念，我们通过将数据集划分到不同问题上创建了训练、验证和测试集，即训练集、验证集和测试集不共享相同问题的解决方案。由于这些问题彼此独立，这种划分方式可以防止跨集合的数据泄漏。在论文中展示的所有实验中，我们基于相同的数据集划分进行评估。表 I 显示了数据集的详细统计信息。所有实验均在具有216 GB RAM 和 2.44 GHz AMD 64核处理器的 A100 NVIDIA 上执行。

评估指标。我们使用了三种不同的评估指标。第一种是 N 精度（PR@N），其中 N 的取值范围是从一到五（包括边界值）。N 精度是指在前 N 个搜索结果中真正正例的平均数量，使用测试集中的每个样本作为查询，针对测试集中的其余部分进行搜索。这些指标是代码搜索中的基准标准 [27]，但我们认为它们最重要，因为它们反映了用户在查询后查看搜索结果列表时的体验。我们还研究了一种称为平均排名差距（ARG）的指标，它定义为负例的平均排名减去正例的平均排名。与 PR@N 指标不同，ARG 可以让我们了解语料库中所有搜索结果的整体情况。通过 PR@N 指标和 ARG 指标，我们可以了解用户的体验以及 COSCO 在整个数据集上的性能。最后，我们根据正例结果的平均第一个位置（AFP）来评估工具的性能（数值越低越好）。ARG 提供了整个结果集的背景信息，而 AFP 则提供了结果集中排名最高部分的背景信息。

RQ1总体表现

首先，我们考察了模型的总体表现，即模型在查找跨语言克隆方面的表现如何。使用前面定义的评估指标，我们针对既需要训练又不需要训练的技术进行了评估。为了进行对照比较，我们实现了一个使用Python的BM25库进行 TF-IDF 搜索的基准方法，该方法将源代码视为自然语言并执行基于标记的搜索。我们还实现了一种基于AST的搜索，改编自 SLACC，其中我们将查询和语料库中文档的AST表示，将它们从特定于语言的 AST 转换为具有通用节点类型的通用 AST（在可能的情况下）。然后，我们使用标准的树距离算法来确定两个样本的相似性。在我们的非训练比较中的最后一个比较中，我们重新实现了 COSAL 的标记子集技术，其中函数被分成较小的片段，然后这些较小的片段被聚类以确定函数之间的相似性。

为了与其他需要训练的代码搜索技术进行比较，我们选择了三种在各种代码相关任务中表现出色的最新 LLM 模型。CodeBERT是由微软亚洲研究院开发的用于代码表示和生成的预训练模型。它可以针对各种下游任务进行微调，并在几个基准数据集上取得了最先进的性能。CodeBERT 跨编程语言传递知识的能力使其在代码生成任务中特别有用，并在软件工程和编程语言研究中有许多潜在应用。GraphCodeBERT是CodeBERT的扩展，它将控制流图（CFGs）整合到其预训练过程中。通过使用 CFGs，GraphCodeBERT 可以更好地捕捉代码标记之间的关系，并为代码分析和生成任务生成更准确的表示。它在几个代码分析基准测试中取得了最先进的性能，展示了其在建模复杂代码结构方面的有效性。UnixCoder是由 Facebook AI 开发的用于自然语言到 shell 命令转换的预训练模型。它在大量命令行使用示例语料库上进行训练，可以在给定自然语言输入的情况下生成 shell 命令。UnixCoder 在几个基准测试中取得了最先进的性能，展示了其理解和将复杂自然语言命令转换为可执行 shell 命令的能力。

图 4：分类器样式搜索

我们使用三个模型中的每一个实现了搜索算法，如图4所示。这些模型在训练集上作为分类器进行训练，然后对测试集中的每个文档，分类器确定文档与给定查询匹配的概率，并根据最高概率对结果进行排序。因此，在推断时，这个过程需要对整个测试语料库进行分类，并产生巨大的计算开销。相比之下，COSCO 可以使用预先计算的语料库嵌入，如图 1b 所示，只需要从查询创建一个嵌入向量，而不需要考虑其余语料库。要获取搜索结果，只需选择与查询向量最接近的预先计算的嵌入。我们的最佳版本的 COSCO 的性能实验大约花费了 14 小时，而其他经过训练的模型代码搜索比较则需要多天才能完成。为了解决这个问题，我们从测试集中随机抽取了 200 个查询，并报告了这 200 个样本的结果。通过这种分层抽样，实验所花费的时间大致与使用 COSCO 进行的实验所花费的时间相同。

图 5：COSCO整体表现

图 5详细说明了我们针对 PR@N 指标和 ARG 与 COSCO 最佳表现版本针对所有比较技术的实验结果。为了视觉清晰，将 PR@N 指标归一化为百分比比例，ARG 指标放大了 100 倍，AFP 指标被省略，但原始值包含在补充材料中。图 5a 显示了针对 Python 代码库的 Java 查询的结果，而图 5b 显示了针对 Java 代码库的 Python 查询的结果。在所有评估指标上，COSCO 的表现均优于所有比较技术，与次佳方法相比，性能提高了高达 40%。令人惊讶的是，我们发现不需要训练的技术在大多数情况下表现优于经过训练的技术，除了 UnixCoder 和从 Java 到 Python 的搜索。尽管每种方法的相对性能取决于评估标准，但我们并不奇怪地发现，当前最先进的 COSAL 通常是次佳表现的。由于 COSCO 在所有评估指标上比现有的跨语言代码搜索技术表现出色很多，我们得出对 RQ1 的答案如下：

结果 1：COSCO 的表现明显优于现有技术，在 Java 到 Python 搜索中比最接近基线的方法提高了 40.6%，在 Python 到 Java 搜索中比最接近基线的方法提高了 11.5%。我们的表现优于 COSAL 的最新技术高达 44.7%。

RQ2 模型通用性

为了进行性能比较，我们使用与上一小节描述的相同的指标：PR@N、ARG和AFP。我们选择了最先进的可用LLM模型进行比较。我们进行了两个实验。首先，我们比较了各种模型的嵌入在不对COSCO的性能进行任何修改的情况下，在相同任务上训练在这些模型上的嵌入的结果。任何性能提升都应归因于COSCO的程序，而不是基本LLM。其次，我们将在每个评估指标上比较COSCO在使用各种基础模型嵌入进行训练时的性能与下一个最佳代码搜索技术的性能。要了解我们的技术的通用性的重要性，我们应该相对于所有模型架构的最新技术进行检查。在所有基础模型上实现最先进的性能将进一步表明我们的技术独立于基础LLM的价值。

图 6：COSCO训练影响(PR@1)

图6比较了训练和未训练模型之间的性能。显然，COSCO的训练总是改善Java到Python和Python到Java查询的代码搜索性能。改进范围从5.2%到1044%。不出所料，我们的训练对我们在训练过程中能够微调的模型产生了最大影响。然而，更有趣的是，我们发现，尽管没有经过训练的CodeBERT模型在训练过程中表现最差，但在Java和Python跨语言查询中，它的微调版本优于所有其他模型。这表明，在训练过程中微调嵌入的影响比整体LLM的大小更大，并展示了微调的重要性。

图 7：COSCO在不同模型上的表现

图7显示了使用COSCO技术相对于COSAL（SOTA）的各种模型的总体性能。图7a详细说明了Java到Python查询的结果，并显示COSCO在所有指标上实现了最先进的性能。有趣的是，对于Python到Java查询，情况并非如此，因为在某些情况下COSAL表现优于COSCO，但从未超过所有版本的COSCO。由于我们不了解OpenAI模型的内部情况，尚不清楚为什么其中一些模型的性能比最先进技术更好，而其他模型则不是。然而，最重要的是，我们发现，在所有跨语言查询的每个指标上，经过微调的CodeBERT优于所有其他模型。这巩固了经过微调的CodeBERT作为COSCO的最高性能版本，并展示了开源模型的强大。这也表明，如果专有模型被开源，我们可能会看到更大的改进。

这些实验的结果表明：1）COSCO的训练技术总是改善基础LLM模型，2）COSCO的技术在大多数情况下表现优于最先进技术，并且在能够在训练过程中微调时总是如此。因此，我们得出结论：

结果2：COSCO的技术可以泛化到多个LLM模型，将Java到Python搜索的性能提高了至少26.2%，将Python到Java搜索的性能提高了5.17%，并且在训练过程中微调基础LLM时，Java到Python搜索的性能提高了7.69倍，Python到Java搜索的性能提高了9.96倍。

RQ3 SSS的影响

我们比较了在训练过程中包含SSS和不包含SSS时COSCO训练版本的性能。为了进行这些实验，我们简单地在训练过程中保留了SSS。任何性能改进都必须来自模型从SSS中学习运行时行为。我们还测试了CodeBERT以及OpenAI模型，以查看学习运行时行为是否实际上是现代LLM的固有属性还是特定于单个模型。

图 8：SSS对COSCO的影响

图8显示了SSS对COSCO在各种基础模型和跨语言查询中性能的影响。我们发现，包含SSS，无论查询语言或用于嵌入的模型如何，都会始终改善性能。这意味着这些模型实际上在学习动态行为，尽管在推断过程中从未运行代码。此外，这表明这种能力是现代LLM的固有特性，而不是单个模型的属性。我们还看到它对OpenAI的专有模型（davinci）中的改进最大，这意味着这种技术继续改进性能，无论模型在训练过程中是否进行微调。因此，我们得出结论：

结果3：包含SSS提高了COSCO在所有嵌入中的代码搜索性能。将SSS与COSCO的性能最佳版本（CodeBERT）一起使用，为Java到Python查询贡献了7%的改进，为Python到Java查询贡献了4.8%的改进。

RQ4 不同的训练样本

我们通过改变训练过程中可用的正面和负面样本的最大数量来研究RQ4。如果模型只是从正面样本中学习，那么当删除负面样本时性能不应下降，反之亦然。此外，如果更多的训练数据可以提高性能，那么在训练过程中可用的参考样本数量增加时，性能也应该增加。我们对一个、三个和五个正面和负面参考样本进行了性能实验。然后我们用零个正面参考样本和五个负面参考样本以及反之进行了实验。我们在训练过程中对进行了微调的CodeBERT模型以及RQ2中使用的所有OpenAI Codex模型进行了这些实验。

图9展示了这些结果，包括CodeBERT和OpenAI的Ada模型在PR@1指标上进行的Java到Python搜索和Python到Java搜索的评估。x轴和y轴显示了训练过程中可用的正面和负面样本的数量，圆圈的大小代表了模型在x轴和y轴描述的情况下的PR@1性能。圆圈越大表示性能越好。

两个有趣的发现在两个可以进行微调和不能进行微调的模型中是一致的。第一，当在训练过程中省略正面或负面样本时，性能会急剧下降。这表明COSCO在训练过程中的双重最大化和最小化技术实际上是其性能的原因。如果只有正面样本对性能负责，那么删除负面样本不会产生影响。这显然不是事实。第二，我们发现当只有一个正面和一个负面参考样本可用时，性能变化不大。这意味着增加训练集的大小对结果不会产生有意义的影响。但我们确实看到了轻微的增加，表明它有一定影响。这可能可以解释为在训练过程中源代码标签是二元的，这意味着额外的正面和负面样本并没有明确提供额外的上下文。如果使用一个更明确量化样本之间距离的不同训练指标，那么可用样本的数量可能会产生更大影响。然而，这一发现表明，即使在训练过程中可用的数据集稀疏的情况下，COSCO的性能仍然有效，即使只有一个正面和一个负面参考样本可用。

结果4：COSCO的性能取决于正面和负面样本。使用正面和负面样本的组合可将性能提高至Python到Java搜索为10.2倍和17.8倍，Java到Python搜索为15.5倍和12.2倍，分别比仅使用正面和负面参考样本。在训练过程中添加超过一个正面和负面样本后，性能改善减弱，分别为Java到Python和Python到Java搜索分别提高了5.4%和0.6%。

转述人：宗兆威