VulRepair: A T5-Based Automated Software Vulnerability Repair

Michael Fu yeh.fu@monash.edu Monash University Australia

Chakkrit Tantithamthavorn∗ chakkrit@monash.edu Monash University Australia

Trung Le trunglm@monash.edu Monash University Australia

Van Nguyen khacvan.nguyen@adelaide.edu.au The University of Adelaide Australia

Dinh Phung dinh.phung@monash.edu Monash University Australia

引用

Fu M, Tantithamthavorn C, Le T, et al. Vulrepair: a t5-based automated software vulnerability repair[C]//Proceedings of the 30th ACM joint european software engineering conference and symposium on the foundations of software engineering. 2022: 935-947.

论文：VulRepair: a T5-based automated software vulnerability repair | Proceedings of the 30th ACM Joint European Software Engineering Conference and Symposium on the Foundations of Software Engineering

摘要

在本文中，我们提出了VulRepair，这是一种基于T5的自动化软件漏洞修复方法，它利用预训练和 BPE 组件来解决先前工作的各种技术限制。通过对来自真实世界软件项目的 8,482 个漏洞修复的广泛实验，我们的 VulRepair 实现了 44% 的完美预测，比竞争基线方法准确13%-21%。根据这些结果我们得出结论，我们的 VulRepair 比两种基线方法更准确，突出了基于 NMT 的自动漏洞修复的巨大进步。我们的额外调查还表明，我们的 VulRepair 可以准确地修复现实世界众所周知的1,706个漏洞中的745个，证明了我们的 VulRepair 在生成漏洞修复方面的实用性和意义，帮助资源不足的安全分析师修复漏洞。

1 引言

软件漏洞是攻击者可以利用的软件系统中发现的安全漏洞、故障或弱点来进行恶意活动。特别是，犯罪群体可以利用软件中未解决的安全漏洞来攻击和破坏系统窃取机密信息或极端资产，造成严重的经济损害。

最近，研究人员提出了各种基于人工智能的方法来帮助资源不足的安全分析师更好地理解漏洞的特征和快速发现漏洞。例如，提出了基于 AI 的漏洞预测方法，用于预测给定函数是否容易受到各种粒度级别和使用的各种类型信息影响。这种漏洞预测方法只帮助安全分析师找到、检测和定位漏洞的位置。然而，安全分析师仍然需要花费大量的精力来手动修复或修复漏洞。

最近，Chen等人提出了VRepair，这是一种利用基于Transformer的神经机器翻译（Neural Machine Translation，NMT）来进行自动漏洞修复的方法。VRepair 旨在解决自动化程序修复领域先前工作的各种挑战然而，由于以下三个限制，VRepair 仍然不准确。

VRepair 在 23,607 个C或C++ 函数的小型错误修复语料库上进行训练，这可能会生成次优的向量表示。VRepair利用单词级别的标记化和复制机制来处理词汇外（Out-Of-Vocabulary，简称OOV）的情况，这限制了它生成新token的能力。这些新token可能从未在易受攻击的函数中出现过，但在进行漏洞修复时被引入了进去。VRepair 利用 Vanilla Transformer（即基本的 Encoder-Decoder Transformer 架构），它使用绝对位置编码，限制了其自注意力机制学习输入序列中代码token的相对位置信息的能力。

在本文中，我们提出了VulRepair，这是一种基于T5的漏洞修复方法，旨在解决上述 VRepair的局限性。首先，我们的 VulRepair 使用来自大型代码库（即CodeSearchNet+C/C#的预训练CodeT5组件，其中包含来自8个不同编程语言的共计8.35亿个函数）来生成更有意义的向量表示，使用 BPE 标记化来处理词汇外（OOV）问题，并采用T5架构来考虑自注意力机制中的相对位置信息。通过对我们的VulRepair在CVEFixes和BigVul数据集上的广泛评估，我们解决了以下四个研究问题：

RQ1 VulRepair生成软件漏洞修复的准确性是多少？

我们的VulRepair实现了44%的完美预测，比VRepair准确度提升21%，比CodeBERT准确度提升13%。

RQ2 使用预训练组件进行漏洞修复的好处是什么？

PL/NL预训练语料库在漏洞修复方法中将完美预测的百分比提高了30%-38%，突出了使用预训练组件进行漏洞修复方法的巨大好处。

RQ3 使用BPE标记化进行漏洞修复的好处是什么？

BPE子词标记化在漏洞修复方法中将完美预测的百分比提高了9%-14%，突出了使用BPE标记化进行漏洞修复方法的巨大好处。

RQ4 我们的VulRepair组件的贡献是什么？

我们的VulRepair的预训练组件是最重要的组件。如果没有为VulRepair正确设计T5架构，性能可以从44%降低到 1%。这一发现表明，设计基于NMT的自动化漏洞修复方法仍然是一项具有挑战性的任务，这需要深入了解现代Transformer架构以实现最高可能的完美预测。

本文的新颖性和贡献如下：

VulRepair是一种基于T5的漏洞修复方法，旨在解决VRepair的各种限制。用两种竞争基线方法（即VRepair、CodeBERT）对我们的VulRepair进行了广泛的评估。 对预训练组件对软件漏洞修复的影响进行实证评估。对标记化技术对软件漏洞修复的影响进行实证评估。消融研究，以研究VulRepair方法每个组件的贡献。

2 背景

以前，循环神经网络(RNN)被广泛用作各种软件工程任务，例如SequenceR和Tufano等人的基于RNN的自动程序修复方法。随着源代码长度的增加，基于RNN的模型在输入token之间存在长期依赖关系，使得基于RNN的模型忘记了长序列标记的一些过去信息（这在源代码中很常见）。

Google Brain引入的基于Transformer的NMT模型是一种具有自注意力机制的编码器-解码器架构。与RNN不同，Transformers不一定按顺序处理标记序列。相反，自注意力机制为输入序列中的任何位置提供上下文向量（即上下文向量用于提供模型应该关注的标记的权重），从而允许基于Transformer的NMT架构比基于RNN的NMT架构更准确。因此，chen等人提出了VRepair，这是一个基于Transformer的NMT自动漏洞修复方法，旨在解决基于RNN的NMT自动程序修复方法的各种限制。具体来说，VRepair 包括以下三个步骤：

第1步：代码表示。对于输入易受攻击的函数，VRepair 首先利用具有复制机制的词级Clang分词器将C函数标记为标记序列。然后，使用词嵌入层生成序列中每个标记的向量表示，以捕获输入标记之间的语义信息。然后，使用绝对位置编码层来生成考虑输入标记之间位置信息的相同序列的另一个向量表示。最后，将两个向量表示相加以形成最终的代码表示，该表示将用作编码器-解码器模型的输入向量。

第2步：编码器-解码器Transformer。VRepair通过采用基于编码器-解码器的Transformer模型来生成修复代码。具体来说，首先将输入向量通过六层Transformer编码器进行处理。接着，来自最后一层编码器的输出向量会分别输入到六层Transformer解码器中。在所有解码器处理完成后，由最后一层解码器产生的输出向量将通过一个带有softmax激活函数的线性层，以获得最终的词汇表概率分布，这个分布用于生成修复代码。

第 3 步：用于修复生成的Beam搜索。给定Transformer解码器输出向量，VRepair使用Beam搜索算法生成50个易受攻击的修复候选。

但是，VRepair也存在上文所述的三个主要限制，我们提出了VulRepair，这是一种基于T5的漏洞修复方法，旨在解决上述VRepair的局限性。

3 VulRepair方法

针对一个有漏洞的函数，我们在第一步中使用基于CodeT5预训练语言模型的字节对编码（Byte-Pairs Encoding, BPE）方法进行子词标记化，以产生子词标记化的函数（即，每个函数的子词代码标记列表）。在第二步中，我们基于T5架构构建了一个VulRepair模型。对于每一个子词标记化的函数，在第2a步，VulRepair执行单词嵌入来为每个标记生成一个嵌入向量，并将其组合成一个矩阵。然后，在第2b步，这个矩阵被输入到T5编码器堆栈中，并且最后一个T5编码器的输出被输入到每个T5解码器中，在第2c步进行。在第2d步，T5解码器堆栈的输出被输入到一个具有softmax激活的线性层中，以生成词汇的概率分布。最后，在第三步中，我们利用Beam搜索在词汇的概率分布之上生成最终的候选预测。

3.1 代码表示

图1 VulRepair架构概述

我们所研究的数据集中的每个局部有漏洞的函数的代码表示由两个主要步骤组成：

1 BPE子词标记化。我们利用字节对编码（BPE）方法进行子词级别的标记化，包括两个主要步骤。首先，生成合并操作以确定如何分割单词；然后，根据子词词汇应用这些合并操作。具体来说，BPE会将所有代码标记分割为字符序列，并识别最常见的符号对（例如，两个连续字符的对），将其合并成一个新的符号。BPE算法会将罕见的标记分割成有意义的子词，同时保留常见的标记（即，不将常见单词分割为更小的子词）。举例来说，函数名称"IsValidSize"会被分割成子词列表，如["IsValid", "Size"]。罕见词"IsValidSize"被分割为两个常见词"IsValid"和"Size"。

BPE子词标记化的使用有助于减小在对各种标记进行标记化时的词汇量，因为它会将罕见的标记分割成多个子词，而不是直接将完整的标记添加到词汇表中。在本文中，我们应用了一个在CodeSearchNet（CSN）和Wang等人提取的C/C#语料库上预训练的BPE标记器。该标记器在八种不同的编程语言上预训练（即Ruby, JavaScript, Go, Python, Java, PHP, C, C#），适用于对源代码进行标记化。为了适应代码生成任务，我们添加了"<s>"和"</s>"标记来表示序列的开始（BOS）和结束（EOS）。如果需要，使用"<pad>"标记将输入序列填充到相同的长度。此外，我们还在词汇表中额外添加了四个特殊标记（即"<StartLoc>", "<EndLoc>", "<ModStart>", "<ModEnd>"）作为额外的词汇ID，因此在标记化过程中它们不会被分割成子组件。使用源代码的预训练语言模型将确保生成的向量表示比VRepair更有意义，因为它是在更大的代码语料库（即CodeSearchNet+C/C#）上预训练的。此外，BPE的使用将确保在漏洞修复中可以生成从未在易受攻击函数中出现过的新标识符。

2a 单词嵌入。源代码由多个标记组成，其中每个标记的含义严重依赖于上下文（即周围的标记）和每个标记在函数中的位置。因此，捕捉代码上下文及其在函数内的位置非常重要。这一步的目的是生成嵌入向量，捕捉代码标记的语义含义以及它们在函数内的位置。对于每个子词标记化的函数，在第2a步中，我们为每个子词标记生成一个 [1x768] 的嵌入向量，并将其组合成矩阵，以代表给定代码标记与其他代码标记之间的有意义关系。为了捕捉代码标记的语义含义，我们利用在与上文提到的预训练分词器相同语料库上预训练的词嵌入向量。为了捕捉每个代码标记在函数中的位置，我们的VulRepair利用相对位置嵌入，在自注意力计算过程中计算并添加到关键矩阵和值矩阵中。

3.2 VulRepair模型架构

2b 编码器堆栈。在第2b步中，实现了十二层编码器块的堆栈，用于生成解码器使用的编码器隐藏状态。类似于原始的Transformer编码器，每个编码器块以一个层归一化开始，其中激活仅进行重新缩放，不应用任何添加偏差。每个编码器块包括两个子组件：一个多头自注意力层，带有相对位置编码，接着是一个前馈神经网络。每个编码器中的每个子组件（即自注意力和前馈神经网络）都围绕它有一个残差连接，并在其后进行一个层归一化步骤。

不同于 VRepair 使用绝对位置编码层和词嵌入层来捕捉输入序列中的位置信息，我们使用相对位置编码来高效考虑输入序列（即关系感知的自注意力机制）内标记之间的相对位置和距离表示。我们在 VulRepair 中使用的自注意力是带有相对位置编码的缩放点积自注意力。自注意力操作使用四个矩阵计算，即Q（查询）, K（键）, V（值）和 P（相对位置信息）。相对位置信息P 作为额外组件提供给模型，添加到键矩阵和值矩阵中，如下所示：Attention(Q, K, V) = softmax(Q(K+P)^T/√dk)(V + P)，其中 P 是点积操作中两个输入的边缘表示，用于确定标记之间的位置信息。与使用固定嵌入为每个位置编码的绝对位置编码不同，成对的位置编码根据自注意力操作中 K 和 Q 之间的偏移产生一个不同的学习嵌入。因此，它可以有效捕捉标记之间的相对信息。

为了捕获输入序列的更丰富语义含义，我们使用多头机制实现自注意力，这允许模型同时关注来自不同位置的不同代码表示子空间的信息。对于维度为d的Q、K和 V，我们将这些向量分割成ℎ个头，每个头具有d/ℎ维度。所有自注意力操作完成后，每个头将被重新连接起来，输入到一个包括两次线性变换和一个ReLU激活函数的全连接前馈神经网络中。

2c 解码器堆栈。在第2c步中，实现了十二层解码器块的堆栈，以根据最后一个编码器块提供的隐藏状态生成漏洞修复。与原始变换器解码器（Transformer Decoder）相似，每个解码器块都以与编码器块中相同的层归一化开始。每个解码器块由三个子组件组成：一个带有相对位置编码的掩码多头自注意力层、一个带有相对位置编码的多头编码器-解码器自注意力，以及一个前馈神经网络。与编码器块一样，解码器中的每个子组件都围绕它具有残差连接，并且之后是一个层归一化步骤。在我们的生成模型的训练阶段使用掩码自注意力，以限制模型在不关注后续上下文的情况下预测下一个标记。因此，模型在生成期间只能关注之前的标记。这遵循了推理阶段的情况，在该阶段模型将不会有任何后续上下文，并且在生成期间只能关注先前的标记。

2d Linear层和Softmax层。线性层是一个完全连接的神经网络，它将解码器堆栈产生的向量投影到更大的logits向量中，单元的数量等于词汇表中唯一标记的数量。然后，以下Softmax层将值转换为最多一个的概率分布，用于在第3步中生成最终输出。

3.3 漏洞修复生成

给出softmax概率的输出，在步骤3中，我们利用Beam搜索根据条件概率在每个时间步为输入序列选择多个漏洞修复候选者。修复候选的数量依赖于称为Beam Width的参数b设置。具体来说，Beam搜索在每个时间步使用最佳优先搜索策略选择具有最高概率的最佳修复候选者。当生成EOS令牌（即“</s>）时，Beam搜索将终止。

4 实验设置

4.1 研究数据集

在我们的实验中，我们使用漏洞修复数据集CVEFixes和Big-Vul，其中包含8,482个漏洞修复（易受攻击的C函数和对其修复）。表1显示了实验数据集的描述性统计。

表1 研究数据集的描述性统计

为了确保与VRepair进行公平的比较，我们严格遵循Chen等人提供的复制包对实验数据集进行预处理。每个输入序列都包含一个指定序列的CWE类型的特殊标签。输入序列中每个易受攻击的代码片段都使用特殊标签“<StartLoc>”和“<EndLoc>”进行标记，其中“<StartLoc>表示易受攻击的代码片段的开始，该代码片段将以特殊标签结尾“<EndLoc>”。对于输出标签，每个修复代码片段都表示为特殊标记“<ModStart>和<ModEnd>，其中 "<ModStart>表示易受攻击修复和非易受攻击上下文的开头，以特殊标签结尾<ModEnd>。将此类特殊标签添加到标记器的主要目的是确保此类特殊标签不会被视为常规代码标记，并且不会由标记器拆分。同样，这种特殊的标签将有助于模型关注易受攻击的代码片段和漏洞修复的区域。

4.2 实验步骤

分割。与Chen等人相同，我们将实验数据集分为70%的训练、10%的验证和20%的测试数据。

漏洞修复模型实现。我们在两个深度学习Python库（即 Transformers和 PyTorch）之上构建了我们的VulRepair 方法。Transformer库提供对基于Transformer的模型架构和预训练权重的API访问，而PyTorch库支持训练过程中的计算（例如，反向传播和参数优化）。

VulRepair的模型训练。我们从Transformers库的API中获得Wang等人预训练的CodeT5 tokenizer和模型。我们使用我们的训练数据集来微调预训练模型，以获得适合我们的漏洞修复任务的权重。该模型在NVIDIA RTX 3090显卡上进行微调，训练时间约为5小时。

在训练过程中，我们使用交叉熵损失来更新模型并在预测序列中的每个位置和实况序列中的每个位置之间进行优化。为了获得最佳微调的权重，我们使用验证集逐个epoch监控训练过程，并根据验证集（不是测试集）的最佳损失值选择最佳模型。

微调的超参数设置。对于我们的 VulRepair 方法的模型架构，我们使用 CodeT5 的默认设置，即 12个 Transformer 编码器块，12个 Transformer 解码器块，768的隐藏层大小，以及 12个注意力头。在微调期间，学习率设置为2e−5，并采用线性计划，其中学习率在整个训练过程中线性衰减。我们使用带有AdamW 优化器的反向传播，这是广泛用于微调基于 Transformer 的模型以更新模型和最小化损失函数的方法。

5 实验结果

RQ1 VulRepair 生成软件漏洞修复的准确性是多少？

图2显示了我们VulRepair和两种基线方法根据我们的评估措施的实验结果。我们的VulRepair实现了44%的完美预测，比基线方法准确13%-21%。

图2 VulRepair的实验结果和漏洞修复的两个基线比较

RQ2 使用预训练组件进行漏洞修复的好处是什么？

图3展示了使用大型预训练语料库进行漏洞修复的好处的实验结果。无论模型架构如何，基于PL/NL的预训练语料库都将漏洞修复方法的完美预测百分比提高了30%-38%。

图3 六种不同模型消融研究的实验结果

RQ3 使用BPE标记化进行漏洞修复的好处是什么？

为了回答这个 RQ，我们旨在调查标记化组件对漏洞修复的影响。为此，对于每种方法，我们只更改每个漏洞修复方法的标记器。具体来说，我们扩展了我们的实验，系统地评估了以下六种基于DL的漏洞修复方法的变体，即2个分词器（子词分词器和词级分词器）×3 个模型架构（VulRepair、CodeBERT、VRepair）：

Subword Tokenizer + CodeT5（我们的 VulRepair）：带有 CodeT5 模型的 BPE 分词器。Word-level Tokenizer + CodeT5：带有 CodeT5 模型的词级标记器。Subword Tokenizer + Vanilla Transformer：带有 Encoder-Decoder Transformer 模型的 BPE 分词器。Word-level Tokenizer + Vanilla Transformer (VRepair)：带有 Encoder-Decoder Transformer 模型的词级标记器和 OOV 问题的复制机制。Subword Tokenizer + CodeBERT（原始 CodeBERT）：BPE 分词器与 CodeBERT 模型。Word-level Tokenizer + CodeBERT：带有 CodeBERT 模型的词级标记器。

最后，我们使用相同的评估措施评估这些变体的准确性。图4展示了使用BPE标记化进行漏洞修复的好处的实验结果。无论模型架构如何，BPE子词标记化在漏洞修复方法中将完美预测的百分比提高了9%-14%。

图4 不同标记化技术用于漏洞修复的各种方法的实验结果

RQ4 我们的 VulRepair 组件的贡献是什么？

为了回答这个RQ，我们的目标是通过检查VulRepair （Pre-training+BPE+T5）中每个组件在变化时的模型精度来研究VulRepair中每个组件的贡献，并与基本的T5相比（无预训练+Wordlevel+T5）。具体来说，我们扩展了我们的实验，系统地评估了以下四种基于T5的漏洞修复方法的变体，即2种预训练策略（预训练，没有预训练）×2 标记器（子词级别，词级）：

预训练 + BPE + T5 (VulRepair)：一个带有 BPE 分词器的预训练 T5 模型。预训练 + 词级 + T5：一个带有词级标记器的预训练 T5 模型。没有预训练 + BPE + T5：一个带有 BPE 分词器的非预训练 T5 模型。没有预训练 + 词级 + T5：一个带有词级标记器的非预训练 T5 模型。

类似地，我们使用相同的评估措施评估这些变体的准确性。结果。图5显示了评估VulRepair组件贡献的消融研究。

图5 VulRepair的消融研究结果

在我们的VulRepair中，预训练组件贡献了完美预测的14%。在比较消除预训练组件的 Pre+BPE+T5 和 NoPre+BPE+T5 时，我们观察到性能下降 44% 到 30%，占 14%。在我们的 VulRepair 中，BPE 组件贡献了完美预测的 9%。在比较 BPE 组件更改为词级的 Pre+BPE+T5 和 Pre+Word+T5 时，我们观察到性能下降从 44% 到 35%，占 9%。然而，如果没有为 VulRepair 正确设计 T5 架构，性能从 44% 下降到 1%。这一发现表明，设计基于 NMT 的自动化漏洞修复方法是一项具有挑战性的任务，这需要深入了解现代 Transformer 架构以实现最高可能百分比的完美预测。

6 结论

本文提出了一种基于T5的自动软件漏洞修复方法VulRepair。通过广泛评估，我们得出结论，VulRepair比VRepair和CodeBERT要准确13%-21%，突显了基于NMT的自动漏洞修复的重大进展。更重要的是，我们的VulRepair能够准确修复1,706个真实世界知名漏洞中的745个。此外，我们发现我们的VulRepair能够正确修复与前十个最危险CWE（如CWE-416（使用后释放）、CWE-20（不当输入验证）和CWE78（OS命令注入））相关的38%的易受攻击功能，这展示了我们的VulRepair对于生成漏洞修复的实用性和重要性，有助于帮助资源不足的安全分析人员修复漏洞。我们的额外分析的重要发现，引出了许多未来研究人员应该探索的开放研究挑战（例如，处理更大的功能、处理罕见类型的漏洞、处理包含许多新标记的困难修复）。

转述：杨襄龙