VulBERTa: Simplified Source Code Pre-Training for Vulnerability Detection
Hazim Hanif, Sergio Maffeis
Computer Science and Information Technology,University of Malaya, Malaysia
Department of Computing Imperial College London, UK
引用
Hanif H, Maffeis S. Vulberta: Simplified source code pre-training for vulnerability detection[C]//2022 International joint conference on neural networks (IJCNN). IEEE, 2022: 1-8.
论文：https://arxiv.org/pdf/2205.12424.pdf
摘要
本文介绍了VulBERTa，一种用于检测源代码中安全漏洞的深度学习方法。本文的方法在开源C/C++项目的实际代码上预训练了一个具有自定义标记化管道的RoBERTa模型。该模型学习了代码语法和语义的深层知识表示，本文利用这一点来训练漏洞检测分类器。本文在几个数据集（Vuldeepecker、Draper、REVEAL和muVuldeepecker）和基准测试（CodeXGLUE和D2A）上评估了本文的方法，针对二元和多类漏洞检测任务。评估结果显示，VulBERTa实现了最先进的性能，并且在不同的数据集上超越了现有的方法，尽管在概念上简单，并且在训练数据的大小和模型参数的数量上成本有限。
1 引言
MITRE报告称，自2016年以来，每年提交的软件CVE（公共漏洞和曝光）数量呈上升趋势，反映了软件生态系统整体安全性面临的威胁增加。相应地，多年来在软件漏洞检测领域的研究稳步增长，涵盖了诸如静态分析、动态分析以及基于机器学习的检测模型等各种漏洞检测方法。深度学习在软件漏洞方面取得了鼓舞人心的结果，特别是使用基于序列和图的技术，如双向LSTM和图神经网络。这些技术试图显式地嵌入代码的语法和语义信息，例如通过使用各种依赖和数据流分析来预处理源代码并提取各种工件，如代码小工具、控制流图和依赖图，最终将其输入到相应的神经网络中。在本文中，我们采用了一种不同的方法，该方法受到了基于Transformer的神经架构的最近成功的启发，这些架构能够学习自然语言文本数据的深层表示知识。本文的目标是构建一个C/C++的表示模型，该模型嵌入了关于语言的语法和语义信息，而无需我们的直接干预，然后使用该模型作为基础，使用标准神经架构构建漏洞检测模型。本文介绍了VulBERTa，旨在从包含不同软件项目的大型代码库中学习C/C++源代码的深层表示。为了促进学习内部代码表示，我们需要创建一个语言感知的可靠标记化管道，来解析和标记源代码，同时确保向模型提供基本的、关键的语法和语义信息。在我们的标记化管道中，本文实现了几个新颖的特性，以增强现有字节对编码（BPE）标记器的标记化能力。本文引入了一个自定义的标记器，它将BPE与自定义的预定义代码令牌相结合，以构建我们标记器的词汇表。这些预定义的令牌基于Clang的AST节点类型（标准关键字和标点符号）和常见C/C++ API函数名的列表。本文的自定义标记化管道在保持其原始语法结构的同时，创建了更好的代码表示，即使在编码后也是如此。
VulBERTa实现了一种基于Transformer的深度学习架构，称为RoBERTa。通过遮蔽语言模型（MLM）的方法，VulBERTa构建了其代码表示知识。本文在较少的样本（228万）和模型参数（1.25亿）上预训练VulBERTa，相较于现有的方法。然后，本文将预训练的VulBERTa模型分别与多层感知器（VulBERTa-MLP）和卷积神经网络（VulBERTa-CNN）相连接，以便对漏洞检测模型进行微调。
本文在不同的数据集上评估VulBERTa，以测试本文预训练模型的性能和可迁移性。评估结果显示，本文在这些数据集上达到了最先进的检测性能。特别是，在被认为是具有挑战性的Draper数据集上，我们达到了57.92%的F1分数，超越了现有的方法，这是令人鼓舞的，因为这个数据集由于不平衡并且混合了真实世界和合成样本而被认为是具有挑战性的。VulBERTa在较为简单的muVuldeepecker数据集上的多类分类任务上也表现良好，达到了99.59%的加权F1分数。本文还在两个软件漏洞检测基准（CodeXGLUE和D2A）上测试了VulBERTa，在使用更少的训练数据和更少的模型参数的情况下，它超越了大多数现有的方法。这些结果显示，VulBERTa在学习源代码的深层表示以及使用该知识来检测软件漏洞方面是有效的。
总结来说，本文的主要贡献包括：
- 一个自定义的标记化管道，它将字节对编码（BPE）算法与新颖的预定义代码令牌（标准C/C++关键词、标点符号和库API调用）相结合，提供了更好的代码编码，同时保持了源代码的语法结构。
- 一个小型化且简化的预训练模型（VulBERTa），提供了预训练嵌入权重的可迁移性，这些权重可以在较为简单的架构如多层感知器（MLP）和卷积神经网络（CNN）中重用。
- 软件漏洞检测模型VulBERTa-MLP和VulBERTa-CNN，在不同数据集上实现了最先进（SOTA）的检测性能，并在两个基准测试（CodeXGLUE和D2A）中位列前三。
2 技术介绍
在这一部分中，我们介绍了VulBERTa，我们用于检测C/C++源代码中功能级粒度漏洞的预训练架构。该架构分为三个关键组件：一个标记化技术，它使用自定义词汇解析和标记代码；一个预训练会话，用于构建代码的表示模型；以及一个微调会话，用于精炼模型以针对具体的分类任务。图1展示了VulBERTa训练流程的8个主要步骤。
图1：VulBERTa训练流程
2.1 Tokenizer
我们的Tokenizer管道旨在保留语法结构和选定的语义标识符。它由解析器、标记器和编码器组成。这些组件叠加在一起，将原始源代码转换为神经网络能够理解的结构。图2显示了VulBERTa从原始代码到编码输出序列的过程。下面，我们描述了管道的每个步骤：
图2：Tokenizer工作流程
1）解析器：我们使用几个正则表达式从每个函数的源代码中移除注释。然后，我们使用Clang，一个健壮的C/C++解析器，解析源代码，该解析器可以在不包含任何库或外部依赖的情况下解析代码。Clang允许我们在将源代码分解成一系列代码标记的同时，保留源代码的语法结构。
2）标记器：Clang解析器产生的标记经过BPE算法的进一步处理，该算法经过修改以考虑我们的预定义标记，进一步将解析输入细分为细粒度的标记以进行编码。字节对编码（BPE）是一种子词标记化算法，它通过用不出现在数据中的字节替换一对相似的连续字节。子词标记化还减少了遇到词汇表外标记的可能性，因为大多数子词标记都在词汇表中。我们定义了一个词汇表大小为50000，作为词汇表中条目的最大数量。预定义标记是我们在词汇表中明确包含的标记，因此将它们从子词标记化过程中排除。我们的目标是保留它们的语法或语义含义。通过预定义C/C++关键词、标点符号和标准API名称，我们在预训练期间保留了更多关于源代码含义的信息。
表1 自定义预定义token
3）编码器：编码是将代码标记转换为张量的过程。对于预训练，我们将最大序列长度设置为512，以最大化和泛化数据的学习，因为预训练数据集包含不同的代码库。同时，对于微调任务，我们将最大序列长度增加到1024，以便它可以在不截断的情况下平均包含超过90%的实际样本。我们用一个特殊的填充标记（<pad>）向右填充较短的序列。
2.2 预训练
预训练是初始训练阶段，我们在此阶段训练一个标准的RoBERTa模型，采用MLM（遮蔽语言模型）目标，以学习C/C++代码在不同软件项目中的信息性通用表示。我们将从这个预训练阶段获得的模型称为VulBERTa模型。结合不同的软件项目对学习过程有益，因为它进一步泛化了代码的表示知识，覆盖了不同的编码风格。这也有助于增加模型在预训练期间的鲁棒性。我们设置嵌入大小为768维，遵循RoBERTa-base。这是模型的核心嵌入知识，将对下游微调任务有用。
2.3 微调
在微调过程中，我们在一个特定的下游任务上进一步训练预训练模型，在我们的案例中是软件漏洞检测。图3展示了用于漏洞检测案例的微调流程。我们为微调实现了两种不同的分类方法。第一种方法是在预训练模型之上使用一个标准的多层感知器（MLP），第二种方法使用文本卷积神经网络（TextCNN），由于CNN架构的鲁棒性，这种方法更便宜、更快速地微调。
图3 微调（漏洞检测任务）流程
VulBERTa-MLP：我们实现了一个包含768个神经元的全连接层，以及一个基于微调数据集是二分类或多分类数据集的输出层，拥有2或41个神经元。在微调过程中，我们重用VulBERTa的预训练权重，并继续训练几个周期。这种方法是微调预训练模型最常见的方法，因为它利用了整个VulBERTa架构，几乎不需要修改。
VulBERTa-CNN：我们提取预训练VulBERTa模型的嵌入权重，并将它们作为混合文本CNN的嵌入权重。TextCNN架构包括三个一维CNN，每个都有其最大池化层。输出被连接和展平，然后输入到两个全连接层（256和128个神经元），带有一个用于分类的输出层。由于嵌入已经在大型语言模型上预训练，我们在训练任务中冻结它们。这种技术允许TextCNN模型继承和使用嵌入的表示知识，并专注于调整特定任务CNN层的权重。
3 实验评估
3.1 实验设置
1) 硬件和软件：我们在所有微调实验中使用PyTorch 1.7与CUDA 10.2，并基于Python 3.7进行操作。对于预训练，我们使用了配备48 vCPUs、240GB RAM和2个NVIDIA Tesla A100 40GB GPUs的Google Compute Engine（GCP）虚拟机。对于微调，我们使用一台配备48核Intel Xeon Silver CPU、292GB RAM和2个NVIDIA GTX TITAN Xp GPU的机器。每个GPU拥有12GB的显存，以适应不同的模型配置。
2) 性能标准：对于每个实验，我们报告了几个评估指标，包括每个数据集初始工作中使用的指标。这样，我们可以进行更公平的比较。这些指标包括真阴性（TN）、假阴性（FN）、真阳性（TP）、假阳性（FP）、准确率、精确度、召回率、F1分数、接收器操作特征曲线下的面积（ROC-AUC）、精确度-召回率AUC（PR-AUC）和Matthews相关系数（MCC）。
3) 基准方法：在性能评估中，我们将VulBERTa与每个数据集现有方法的两种基准技术进行比较。这两个基准在分析用于漏洞检测的序列化输入时被广泛使用，并且已经在例如文献[6]和[17]中使用。
(i) 基准-BiLSTM：这项技术是LSTM的一种变体，它实现了一个双层双向LSTM和几个全连接层，以从源代码的序列中学习漏洞检测。双向LSTM同时学习代码序列的前向和后向关系。
(ii) 基准-TextCNN：这是CNN的一个变体，其中输入数据是自然语言文本而不是图像。在这种情况下，我们使用源代码作为输入数据，并将其输入到CNN中。这种技术部署了三个具有池化的卷积层，并在传递结果到几个全连接层之前将它们连接成一个单一层。
4) 模型预训练：我们使用Draper和GitHub数据集通过MLM预训练VulBERTa模型。我们尝试了不同的RoBERTa配置（即小型、中型和基础型），以观察模型参数数量如何影响模型的预训练性能。每次预训练会话的持续时间根据模型配置介于72到96小时之间。训练会话进行到500,000步，并使用学习率调度器随着训练损失趋于平稳而降低学习率。基于结果，我们发现基础型模型在500,000步后的训练中给出了最低的损失。因此，我们选择基础配置的VulBERTa模型（约125M参数）作为我们用于微调的参考预训练模型。
5) 模型微调：我们分别对预训练的VulBERTa-MLP和VulBERTa-CNN模型进行微调，使用漏洞检测作为微调目标。我们将最大迭代次数设置为10，这是足够的，因为模型在4-5个周期后开始对训练集过拟合。我们将学习率设置为0.00003，并使用学习率调度器随着训练损失趋于平稳而降低学习率。我们遵循每个数据集的原始划分，但如果划分信息不可用，我们将数据集分割为80/10/10（训练/验证/测试）。每次微调会话持续时间在5到10小时之间，具体取决于数据集和模型的大小。
3.2 在选定数据集上的评估
漏洞检测实验按数据集分开进行。对于每个数据集，我们选择了原始论文中用于比较的首选评估指标（PEM）。表II展示了评估结果，其中我们突出显示了每个数据集的最高PEM分数。
1) Vuldeepecker：Vuldeepecker报告的精确度得分为91.9%。然而，使用VulBERTa-MLP，我们实现了95.76%的精确度得分，比原始得分高出3.86%。此外，我们的模型获得了更高的F1分数，93.03%，相比于Vuldeepecker的92.9%，这表明我们的模型在分类易受攻击和非易受攻击样本时更为平衡。低比率的假阳性（0.39%）和假阴性（9.14%）表明VulBERTa-MLP能够在合成和现实世界代码中以低误分类率检测出漏洞。
2) Draper：[18]的作者使用几个评估指标来评估他们在Draper数据集上的模型。我们选择MCC作为比较的基础，因为它适用于不平衡数据集，而Draper数据集的易受攻击类和非易受攻击类是不平衡的。VulBERTa-CNN在这个数据集上获得了55.86%的MCC得分。这比[18]报告的性能提高了2.26%，这是显著的，因为我们可以在保持类别平衡的同时提供更好的检测。
3) REVEAL：VulBERTa-MLP模型实现了45.27%的F1分数，高于[25]报告的41.25%，尽管没有使用那里提出的数据再平衡技术。相反，我们在微调过程中为每个类（易受攻击和非易受攻击）分配了权重，以减少类别不平衡问题，而不改变数据集。我们的方法还获得了比原来高2.57%的真阳性率（TPR）。这表明VulBERTa-MLP在正确检测易受攻击样本的同时，也保持了与非易受攻击类的平衡。
4) muVuldeepecker：与前面的实验不同，这是一个多类分类任务。muVuldeepecker数据集包含40个CWE的独立类别，所以每个阳性样本可以映射到一个特定的安全漏洞。VulBERTa-MLP实现了非常高的加权F1分数，99.59%，相比于[5]报告的96.28%。它还将假阴性率（FNR）从5.53%降低到0.41%，这是显著的，因为作为易受攻击的附加样本需要被分配到正确的类别。此外，查看特定类别如CWE-190和CWE-191（整数溢出和下溢）的预测，我们可以看到VulBERTa-MLP能够正确地将90%以上的易受攻击样本分配到它们各自的类别中。
我们的模型设法超越了使用后者数据集和首选评估指标的比较模型的性能。这发生在包含合成和现实世界数据、平衡和不平衡类别、二类和多类分类任务的各种数据集上。
3.3 在选定数据集上的评估
我们还在两个公开可用的基准测试上评估了VulBERTa-MLP和VulBERTa-CNN，这两个基准测试被社区用作比较不同源代码模型在标准化任务上的基础。在两种情况下，漏洞检测任务的首选评估指标（PEM）是准确性。
表2 数据集评估结果
1) CodeXGLUE：CodeXGLUE基准测试是由微软研究院引入的第一个也是最受欢迎的编程语言理解基准测试。我们关注的是缺陷检测任务，该任务包括在第IV-B5节描述的Devign数据集上的漏洞检测。表III-A展示了发表时CodeXGLUE基准测试排行榜的情况，包括我们的结果。VulBERTa-MLP实现了64.75%的准确率，排名第三，低于CoText和C-BERT。值得注意的是，VulBERTa-MLP的模型参数数量显著少于CoTexT（55.07%），并且是用前2个模型的一小部分数据进行训练的。令人印象深刻的是，VulBERTa-CNN的性能略低，但模型大小不到CoTexT的1%，C-BERT的2%。
2) D2A：D2A是由IBM研究院引入的漏洞检测基准测试，基于第IV-B6节描述的D2A数据集。表III-B展示了发表时D2A基准测试排行榜的情况，包括我们的模型。对于“函数”任务（与我们的方法唯一相关的任务），VulBERTa-MLP以62.30%的准确率领先排行榜，VulBERTa-CNN以60.68%排在第二。有趣的是，在这个数据集上，我们的两个模型都超过了C-BERT，尽管它有更大的预训练数据集和更多的参数（与VulBERTa-CNN相比）。
表3性能评估结果
3.4 讨论
表II显示，VulBERTa-MLP在4个数据集中的3个上拥有最佳的PEM分数，而VulBERTa-CNN在剩下的数据集上有最佳的PEM分数。然而，如果我们在微调结果的全谱上比较我们的两个模型，我们会注意到两者之间的差异可以忽略不计。同样，在基准评估上，VulBERTa-MLP和VulBERTa-CNN之间也显示出紧密的关系。两个模型在两个基准测试中都一致地排名相邻，并且在两者中都达到了最先进的性能。一个可能的解释是，从预训练的VulBERTa模型继承的代码表示知识对于漏洞检测起着关键作用。
模型大小和预训练数据大小通常在学习更好的代码表示中扮演着重要的角色。然而，VulBERTa只包含125M参数，我们的预训练数据仅包含2.28M个C/C++函数，例如，与CoTexT（375M）和C-Bert（8.5M）相比。我们通过使用较小的模型和预训练数据提供了与这些方法相当的SOTA检测性能。
基于我们对测试和实施不同标记化技术的初步分析，我们相信我们的标记化方法在使语法和语义信息易于被用于微调的简单神经架构可用方面发挥了重要作用。事实上，模型的简单性似乎是解决方案的一部分。
VulBERTa-CNN，我们基于简单的TextCNN方法（仅有2M参数），在Draper数据集上的MCC得分为55.86%，高于最近在[40]中提出的复杂的3GNN模型所获得的52%，后者使用了晶体图卷积网络和自注意力池化。
3.5 限制
尽管结果令人鼓舞，我们还是发现了我们方法的一些限制。与漏洞检测数据集相关的一个共同且尚未解决的问题是，人工检查偶尔会发现标签不准确的情况。虽然深度学习应该对训练中的标签噪声具有弹性，但测试过程中噪声的存在在某种程度上削弱了量化性能结果。尽管我们的模型相对较小，但训练它们仍然代价昂贵。由于资源有限，我们无法探索显著更大的模型配置和组合，包括执行超参数调整。因此，不同的VulBERTa配置可能会达到更高的性能。
最后，我们工作的主要限制是缺乏系统地尝试在野外的开源项目中检测新的0Day漏洞。这是由于手动审核假阳性的挑战所致，我们希望在未来的工作中解决这个问题，利用可解释性技术。
转述：王越
