TransformCode: A Contrastive Learning Framework for Code Embedding via Subtree Transformation

最新推荐文章于 2024-08-30 12:54:54 发布

图学习小组

最新推荐文章于 2024-08-30 12:54:54 发布

阅读量903

点赞数 9

文章标签：深度学习代码嵌入自然语言处理软件工程

本文链接：https://blog.csdn.net/qq_41200212/article/details/139964151

版权

TransformCode: A Contrastive Learning Framework for Code Embedding via Subtree Transformation

基本信息

博客贡献人

JokerLin

作者

Zixiang Xian，Rubing Huang，Dave Towey，Chunrong Fang，Zhenyu Chen

摘要

基于现有的利用迁移学习的预训练模型（PTM），本文提出了TransformCode框架，这是一种以对比学习方式学习代码嵌入的新框架。该框架是与编码器和语言无关的，同时还提出了一种新的数据扩充技术——抽象语法树（AST）转换，该技术将语法和语义转换应用于原始代码片段，以生成更多样、更稳健的样本用于对比学习。在本文中，证明了它在几个与代码相关的任务上相对于最先进的方法（如SourcererCC、Code2vec和InferCode）的有效性和优越性。

问题定义

人工智能（AI）通过提高软件开发效率，彻底改变了软件工程（SE）。利用迁移学习的预训练模型（PTM）的出现大大提高了SE的人工智能，经过广泛的代码数据集培训，PTM对编程语言有着深刻的理解，使它们能够针对各式各样的SE任务进行微调。虽然现有的PTM能够执行各种任务，但由于其大量的参数，它们可能需要巨大的GPU资源和高质量的数据集进行微调。它们通常在大规模数据集上进行预训练，然后在特定任务的数据集上微调，相反在许多情况下，最好使用无监督或自监督学习，直接在数据集上训练较小的模型，而不需要太多其他资源。

针对上述问题，本文提出了从任意数据集学习代码嵌入的新框架——TransformCode，该框架可用于各种SE任务，通过将对比学习与动量编码器相结合，从未标记的数据中学习强大而稳健的表示。框架使用队列和动量更新编码器来构建一个动态一致的键字典，这使能够对大量不同的负样本集进行对比学习。本文提出的框架通过结合转换代码的附加特征来改进对比学习，这些特征提高了学习表示的质量和效率。与其他框架相比，本文提出的框架比现有的学习代码表示的方法有如下几个优势：

具有高度的灵活性和适应性。使该框架能够处理需要代码表示的广泛下游任务（例如检测和分类代码克隆，以及以无监督的方式对代码片段进行聚类）
可跨不同编程语言进行扩展，非常高效。框架使用对比学习和数据扩充技术从少量数据中学习代码表示，同时增加了一种自注意机制，可以动态调整代码图中不同节点和边的重要性
可作为用于代码表示学习的通用工具，可以处理无监督和有监督的学习场景
可以根据可用的计算资源调整编码器参数的数量，使其适用于各种部署场景

方法

代码规范化

代码规范化也有助于减少框架编码器的字典大小，并使本文的框架更加关注代码结构。在本文中提出了一种两步代码规范化方法，该方法由以下操作组成：

删除代码中的所有注释（行注释和块注释）：它们不影响代码的执行或功能。
将代码中的所有变量重命名为以var开头、以数字（var1、var2、var3…）结尾的标准格式。这确保代码嵌入不会因不同的命名约定或样式而产生偏差。

图1中显示了按照上述规则进行规范化前后代码的比较。

在这里插入图片描述

（a）归一化前（b）归一化后

图1. 代码规范化示例

AST转换与提取

AST是捕获代码的语法和语义信息的树结构。本文通过应用转换对AST执行数据扩充，例如交换、插入或删除节点，以创建新的AST，称为锚点AST。这些锚点在功能上与原始AST相似，但在语法上不同。

生成锚AST涉及在不改变其语义的情况下修改代码语法。这种转换可以用于创建用于对比学习的多样化且现实的代码样本。在本文中提出了一组通用代码转换方法，可应用于大多数编程语言，同时保留其语义。具体转换方法如下：

切换变量声明语句的顺序
交换二进制运算符的两个操作数，保持语义不变
将算术运算转换为不同的形式
将while语句转换为for语句
将死代码添加到随机选择的语句中，但不会更改代码语义
将try-catch语句添加到随机选择的语句中，但不会更改代码语义
交换随机选择的语句

为了识别原始和锚定AST的关键执行路径，TransformCode框架使用了一组以深度优先的方式遍历AST的规则，并修剪对代码的语法或语义没有贡献的节点。这个过程产生了一个简化和精简的AST，它捕获了代码的基本逻辑，按照遍历顺序将其连接形成一系列令牌。

示例

如图所示，给出基于Java实现的冒泡排序算法的实现代码

图2. 冒泡排序算法的Java实现
提取并简化

从提供的冒泡排序算法代码中提取AST标记，表达式和语句可以通过删除不必要的括号等来简化，简化结果如图3所示

图3. 提取后简化AST标记
归一化处理

为了使模型更加健壮和可推广，在提取AST标记之前应用归一化步骤，按照前面提出的代码规范化的操作步骤对简化后的AST标记再次处理，最终结果如图4所示，处理后的令牌不需要任何特殊的令牌来指示序列的开始或结束，并且可以直接输入到模型中。

图4. 归一化后的AST标记

TransformCode框架

如图5所示，TransformCode框架包括三个主要步骤:

step1：预处理

对代码片段进行预处理，以获得归一化和转换后的代码，该代码用作训练数据。
step2：训练代码标记器

对代码标记器进行训练，并使用经过训练的标记器嵌入所有代码片段。
step3：训练模型

TransformCode框架既有查询编码器又有动量编码器——编码器具有相同的架构，但参数不同；动量编码器由查询编码器的移动平均值更新，这有助于保持表示的一致性和多样性，同样也有助于避免对比损失的崩溃。

图5. 代码嵌入的无监督学习框架

InfoNCE损失

首先为所有查询建立一个负样本队列。队列被维护为先进先出（FIFO）缓冲区，其中当前小批量的编码表示（即训练样本）进入队列，而最旧的表示则移出队列。每个训练样本都被视为一个查询，以从每个输入批次（训练样本）和代码片段队列（负样本队列）中找到最佳匹配键。通过对当前代码样本进行代码转换来获得最佳匹配键，这形成了正对（在对比学习中也称为锚对），最佳匹配键也称为正样本（锚点样本）。本文还通过将查询样本与其他键配对来形成负对：负样本是小批量和队列中的所有其他代码样本。对此在本文中使用InfoNCE损失来最大化正对之间的一致性并且最小化负对之间的一致性。InfoNCE损失定义如下：

$\mathcal{L}_{q,k^+,\lbrace k^-\rbrace} = -\log \frac{\exp(q*k^+/\tau)}{\exp(q*k^+/\tau) + \sum_{k^-}\exp(q*k^-/\tau)} \tag{1}$

其中 $q$ 是查询代码表示（训练样本）； $k^+$ 是正密钥样本（锚样本）的代码表示； $k^−$ 是负密钥样本（不同样本）的表示； $\tau$ 是控制softmax函数柔软度的温度参数。

InfoNCE损失旨在最大化正对相对于负对的可能性，这相当于最大化 $q$ 和 $k^+$ 之间的互信息的下界，因此可以使用InfoNCE损失来最大化代码正对之间的一致性，并最小化代码负对之间的相似性。

相对位置编码

TransformCode框架采用transformer编码器作为的AST编码器，但用相对位置编码取代正弦位置编码。编码器层由多头自注意子层和前馈子层组成，然后是层归一化和残差连接。多头自注意子层允许模型关注代码片段的不同部分，并学习标记之间的语义关系。前馈子层由两个线性变换组成，中间有ReLU函数用来增强模型的表达能力和非线性。相对位置编码定义如下：

$e_{ij}=\frac{x_iW^Q(x_jW^K)^T+x_iW^Q(a_{ij}^K)^T}{\sqrt{d_z}} \tag{2}$

其中 $a_{ij}$ 是输入元素 $x_i$ 和 $x_j$ 之间的边缘。

位置编码使模型能够区分不同位置的相同单词，并产生位置感知表示。这也有助于模型更好地捕捉代码的相对位置结构。

多层感知机

为了进一步捕获代码的结构信息并提高对比学习性能，TransformCode框架在编码器之后引入了MLP头。MLP头将编码器的输出投影到较低维空间 $h_i$ ，其中正对的表示不变，负对的表示更有鉴别性。MLP头公式化如下：

$h_0=E(x) \tag{3}$

$h_1 = \sigma(W_ih_{i-1}+b_i), i=1,...,M-1 \tag{4}$

$W_Mh_{M-1}+b_M \tag{5}$

其中 $x$ 是输入； $E$ 是编码器； $W_i$ 和 $b_i$ 是第 $i$ 个线性层的权重矩阵和偏置矢量； $h_i$ 是第 $i$ 个线性层的输出； $z$ 是最终的代码表示。

MLP头可以被视为增强代码表示的质量和鲁棒性的非线性变换。此外，MLP头还与相对位置编码相互作用，因为相对位置矢量在被馈送到代码片段编码器之前被添加到每个令牌嵌入，因此相对位置编码可以向MLP头提供有用的结构信息。

预训练模型局限性

开源预训练模型（PTM）和参数有效微调（PEFT）方法的最新进展使为下游任务微调PTM成为可能。然而，这些方法在应用于代码嵌入时仍然面临一些限制。例如，大多数像CodeBERT这样的PTM使用特殊的标记，如<s>和</s>，来标记句子的开始和结束，这导致模型更加关注这些标记，而不是实际的代码内容。

对此本文提出一个新方法来解决这个问题，该方法不需要任何预先训练的模型或特殊令牌，而是使用提取的AST路径从头开始训练新模型。如图6中利用Java代码展示了这个方法，图的上部显示了原始代码，下部显示了提取了执行路径的规范化代码。

在这里插入图片描述

图6. Java代码示例

同时在注意力权重图7（图中的不同颜色表示不同的注意力头，较深的阴影表示较高的注意力权重）中可以看出本文提出的方法更均匀地分配注意力，并选择性地关注代码的关键部分，能够更有效地捕捉代码的语义结构和逻辑。

在这里插入图片描述

图7. 提出的方法与CodeBERT的比较

实验

评估指标

代码克隆检测是一项分类任务，涉及确定两个代码片段是否相同。为了评估代码克隆检测的性能，本文使用以下常用于分类任务的指标（ $TP$ 表示真阳性； $TN$ 表示真阴性； $FP$ 表示假阳性； $FN$ 表示假阴性）：

$Accuracy=(\frac{TP+TN}{TP+TN+FP+FN}) \tag{6}$

$Precision=(\frac{TP}{TP+FP}) \tag{7}$

$Recall=(\frac{TP}{TP+FN}) \tag{8}$

$F_1 = 2\times\frac{Precision \times Recall}{Precision + Recall} \tag{9}$

在处理不平衡数据集时，准确性可能不是一个可靠的指标，其中一些类别或类别的代表性不足或过高。在这种情况下，准确性可能会受到优势阶级的偏见，导致忽视少数阶级。

例如，如果一个数据集有90%的正实例和10%的负实例，那么总是预测正实例的分类器将有90%的准确率，但它将无法检测到任何负实例。

因此选择使用 $F_1$ 分数（精度和召回率的调和平均值）作为不平衡数据集的替代指标。 $F_1$ 分数同时考虑了精度和回忆：当精度和召回率都很高时，它会给出更高的值，这意味着分类器可以正确识别正面和负面实例。

代码克隆检测

代码克隆检测涉及识别在语法或语义方面相似或相同的代码片段。代码克隆可以分为四种主要类型：

除了空白、布局和注释的变化之外，其他代码片段都是相同的。这些也被称为精确或文本克隆
除了标识符名称和文字值的变化之外，代码片段是相同的。这些克隆也称为重命名克隆或参数化克隆
语法相似但在语句级别不同的代码片段。这些也被称为缺口克隆或未遂克隆
语法不同但实现相同功能的代码片段。这些也被称为语义或功能克隆

实验中利用基于成对相似性的POJ-104中的代码对，构建了OJClone数据集。从前15个POJ-104问题中的每个问题中选择500个程序，得到180万个克隆对和2620万个非克隆对。所有对的比较将非常耗时，因此随机选择50000个克隆对和50000个非克隆对进行代码检测评估。BigCloneBench是一个广泛使用的基准数据集，包括来自25000个Java存储库的项目：它涵盖了十个功能，包括6000000个真克隆对和260000 个假克隆对。OJClone和BigCloneBench数据集都用于评估无监督和有监督环境中的性能，将本文提出的方法与最先进的预训练代码表示模型CodeBERT和无监督代码表示学习模型InferCode进行比较。

无监督学习

在无监督学习实验中，通过计算两个AST的余弦相似性，并根据阈值将代码对分类为相似或不相似。将阈值T设置为0.75。代码克隆检测中的这个阈值是一个超参数，它基于相似和不相似代码对的余弦相似性来定义它们之间的截止点。由于最佳阈值可能会根据数据特征、模型结构和GPU资源而变化。在过程中还使用阈值T=0.5进行了实验，如表1所示。当把阈值降低到0.5时，性能会恶化，训练时间会增加。这是因为编码器具有简单的架构（2或4层），并且批大小受到GPU资源的限制。本文的模型很难用较低的阈值来区分克隆和非克隆。

表1. BigCloneBench Java上不同体系结构模型的度量

在这里插入图片描述

实验证明，多头注意力的维度固定在64，增加模型深度（nlayers）或隐藏大小（dmodel）可以提高性能。当使用相对位置编码和MLP、两个编码层和1024的隐藏大小时，本文的框架实现了最佳性能和收敛速度。有了这些设置和128的批量大小，TransformCode框架获得了82.36%的F1分数和87.50%的准确率。对比结果如表2所示，TransformCode框架实现的精度略低于InferCode，但在 $F_1$ 分数和召回率方面超过了InferCode 。InferCode 具有较低的召回率和较高的精度，这意味着它在识别代码克隆对时不太一致。

表2. BigCloneBench Java上不同模型的度量（无监督）

在这里插入图片描述

将此配置应用于OJClone数据集上，TransformCode框架在隐藏大小为1024的四层结构中，以67.69%的精度和67.10%的F1得分获得了最佳性能。与其他无监督模型的对比结果如表3所示，在三种不同的参数设置下（两个编码层，隐藏大小为1024；四个编码层，隐藏大小为1024；八个编码层，隐藏大小为128），TransformCode框架都优于其他模型，证明了TransformCode框架在学习代码表示和检测代码克隆方面的有效性。

表3. OJClone上不同模型的度量（无监督）

在这里插入图片描述

根据图8的实验结果可视化分析，TransformCode框架可以学习将相似的代码分组到同一个集群中，即使使用的是相对较小批量（批量大小为86）的对比学习。这表明TransformCode框架可以有效地从代码片段中学习，而不需要大量的数据。

在这里插入图片描述

（a）代码嵌入OJClone C数据集的可视化（b）Java函数数据集中嵌入代码的可视化

图8. 两个不同数据集的TSNE可视化

监督学习

在监督学习实验中，OJClone C数据集的标记数据用于训练模型。TransformCode框架由具有相对位置编码的变换编码器和线性分类器组成，以交叉熵损失作为监督损失进行联合训练，编码器采用对比损耗，使用可训练参数α=0.2作为初始值使用，通过以下公式平衡两种损耗：

$\alpha * Loss_{Contrastive} + (1 - \alpha) * Loss_{Supervised} \tag{10}$

本文所提出的框架与几个基线进行了比较，包括传统方法，如SVM、神经网络模型（TextCNN、LSTM、TBCNN和LSCNN）、基于程序依赖图（基于PDG）的方法和CodeBERT。此外，实验对SVM使用线性核，并从TF-IDF、N-gram和LDA中提取统计特征，将N-gram大小设置为2，并将最大特征数限制为20000。同样为了使基于编码器的预训练模型CodeBERT适应该实验，在CodeBERT之后添加了前馈层和分类器层。实验结果如表4所示，TransformCode框架实现了与TBCNN（表现最好的）相当的性能，但参数较少，证明了其有效性。

表4. 在OJCLONE C数据集上训练的模型的度量（监督）

在这里插入图片描述

同样，结果显示TransformCode框架也取得了比CodeBERT稍好的结果，具有更小的模型大小和更少的参数：只使用了两个编码器层和128个隐藏单元作为编码器，而CodeBERT使用了12个编码器层，768个隐藏单元。这证明了TransformCode框架在从代码数据中学习方面的效率。

方法名称预测

本实验使用Java Small数据集进行训练的，该数据集是从开源项目中提取的Java方法的集合，通过在Java Small数据集上比较了TransformCode框架与InferCode、Code2vec和Code2seq的性能。

关于方法名称预测任务的要求如下：

实验的过程按照三个步骤执行：训练；测试；以及验证。
方法名称预测任务是一个SE任务，旨在为给定的代码片段生成描述性和简明的名称。任务可以公式化为序列生成问题，将输出序列长度限制为五个子字。
TransformCode框架采用与Code2vec和InferCode中相同的评估指标，它们使用了生成的方法名称的子单词的 $F_1$ 分数。（子单词是单词中有意义的部分，由大小写、下划线或连字符分隔）

示例

单词computeMax有两个子词：compute和Max。将预测方法名称的子词与基本事实方法名称的个子词进行比较，并相应地计算Precision、Recall和F1 Score。

如果预测的方法名称是compute max，它将被视为与地面实况方法名称computeMax完全匹配，F1分数将为100%。如果预测结果为getMax，则子词Max将获得完全精度，但Recall的得分仅为50%，F1得分为67%。如果预测结果是getMaxResult，它将获得完全Recall，但Precision只有67%，F1得分将在80%左右。

实验结果如表5所示，就 $F_1$ 得分而言，TransformCode框架实现了与InferCode相当的性能。

表5. JAVA-SMALL数据集F1评分中方法名称预测的结果（监督）

在这里插入图片描述

在这个实验中，尽管在参数、训练数据集以及字典上存在差异，例如CodeBERT有一个更大的字典，而TransformCode框架的字典是专门为这项任务定制的，只使用训练数据，而不访问测试数据，但TransformCode框架比CodeBERT更高效，因为本文提出的框架使用了更小的编码器，并且具有更快的训练速度。在相同的硬件配置上，TransformCode框架可以比CodeBERT快四倍左右。

代码分类

本实验使用POJ-104数据集，该数据集是编程问题和解决方案的大规模集合。实验将数据集分为三个子集（比例为7:2:1）：训练；测试；以及验证。训练集用于训练模型参数，测试集用于测量模型的可推广性，验证集用于调整超参数并防止过拟合。

TransformCode框架的目标是学习一种代码向量表示，为了实现这一目标，实验使用了对比损失和交叉熵损失的组合作为训练损失函数，并用一个可训练的参数 $\alpha$ （初始化为0.1）来平衡这两个损失，由此损失函数定义如下：

$\alpha * Loss_{Contrastive} + (1 - \alpha) * Loss_{Supervised} + 0.5*Loss_{Anchor} \tag{11}$

实验结果如表6所示，可知TransformCode框架获得了与两个最先进的模型（InferCode和ASTNN）相当的结果，但TransformCode框架还具有易于训练和需要较少GPU功率的优点。同样相较于CodeBERT，由于加入了数据扩充，TransformCode框架比transformer编码器略有改进，在使用相同硬件设置的有效性方面超过了CodeBERT。

表6. OJ-C数据集上代码分类的准确性结果（监督）

在这里插入图片描述

TransformCode新框架在本文中由于实验设备硬件问题，而无法使用大批量进行训练
由于编程语言机器语法规则的多样性，没能为代码转换提供更多选项

启发

将树网络和基于令牌的方法相结合，探索两者在代码表示学习中的协同作用
在编译代码时选择不同的优化级别，将其作为变换方法的一种，为模型提供更可靠的变体
将大语言模型（LLM）集成到框架中，解决更实际和更具挑战性的SE问题，如大规模代码检索和代码缺陷检测

图学习小组

关注

9
点赞
踩
27

收藏

觉得还不错? 一键收藏
0
评论
TransformCode: A Contrastive Learning Framework for Code Embedding via Subtree Transformation

基于现有的利用迁移学习的预训练模型（PTM），本文提出了TransformCode框架，这是一种以对比学习方式学习代码嵌入的新框架。该框架是与编码器和语言无关的，同时还提出了一种新的数据扩充技术——抽象语法树（AST）转换，该技术将语法和语义转换应用于原始代码片段，以生成更多样、更稳健的样本用于对比学习。在本文中，证明了它在几个与代码相关的任务上相对于最先进的方法（如SourcererCC、Code2vec和InferCode）的有效性和优越性。
复制链接

扫一扫