LineVul: A Transformer-based Line-Level Vulnerability Prediction

Michael Fu Chakkrit ,Tantithamthavorn

引用

Fu M, Tantithamthavorn C. Linevul: A transformer-based line-level vulnerability prediction[C]//Proceedings of the 19th International Conference on Mining Software Repositories. 2022: 608-620.

论文：LineVul: A Transformer-based Line-Level Vulnerability Prediction | IEEE Conference Publication | IEEE Xplore

摘要

这篇文章介绍了一种基于Transformer的漏洞预测方法，称为LineVul。作者指出现有的漏洞预测方法如IVDetect（一种基于图的神经网络）存在一些限制，因此提出了LineVul以解决最先进的 IVDetect方法的几个局限性。通过对一个大规模的真实世界数据集进行实证评估，作者展示了LineVul在功能级别和行级别漏洞预测方面的准确性和效果，突出了LineVul在更准确和更具成本效益的线级漏洞预测方面的重大进步。

1 引言

软件漏洞是信息系统、安全程序、内部控制或实施中可能被威胁源利用或触发的弱点。那些与软件系统相关的未解决的弱点可能会导致极端的安全或隐私风险。

为了缓解这一挑战，人们引入了程序分析(PA)工具，利用预定义的漏洞模式来分析源代码。例如，Gupta等人发现，有一些静态分析方法被用来检测SQL注入和跨站脚本漏洞。另一方面，机器学习(ML) /深度学习(DL)方法已经被提出。具体来说，这些基于ML/ DL的方法首先生成源代码的表示，以便学习漏洞模式。最后，这样的方法将学习源代码的表示和ground truth之间的关系。尽管基于ML/ DL的方法具有动态学习漏洞模式而无需手动预定义漏洞模式的优点，但以前的方法仍然侧重于粗粒度的漏洞预测，模型只指出文件级或功能级的漏洞， 仍然是粗粒度的。

最近，Li等提出了一种IVDetect方法来解决细粒度漏洞预测的需求。IVDetect利用FA-GCN(即特征注意图卷积网络) 方法来预测功能级漏洞，并利用gnnex -plain来定位漏洞的细粒度位置。Li等人发现IVDetect方法的f值为0.35，优于最先进的方法。然而，IVDetect有以下三个限制：

首先，IVDetect的训练过程仅限于项目特定的数据集。 由于IVDetect使用的训练数据数量有限，语言模型可能无法捕获标记与其周围标记之间最准确的关系。因此，IVDetect对源代码的向量表示仍然是次优的。其次，IVDetect方法的基于rnn的体系结构仍然不能有 效地捕获源代码的有意义的长期依赖关系和语义。IVDetect依赖于基于rnn的模型来生成向量表示，供其图形模型在过程中使用预测步骤。然而，基于rnn的模型在学习长序列的源代码时往往存在困难。这样的限制可能会使生成的向量表示变得不那么有意义，从而导致预测不准确。第三，IVDetect的子图解释仍然是粗粒度的。IVDetect利用gnexplainer来确定哪个子图对预测贡献最大。虽然这样的子图解释可以帮助开发人员缩小范围来定位脆弱的行，但这样的子图仍然包含许多代码行。因此，安全分析人员仍然需要手动定位这些子图中的哪些行实际上是脆弱的。

在本文中，我们提出了LineVul，一种基于transformer的细粒度漏洞预测方法，以解决IVDetect的三个重要限制。首先，我们没有使用基于rnn的模型来生成代码表示，而是利用具有自关注层的BERT架构，该架构能够使用点积操作在长序列中捕获长期依赖关系。其次，我们利用CodeBERT预训练语言模型来生成源代码的向量表示，而不是使用特定项目的训练数据。 第三，我们没有使用gnexplainer来识别有助于预测的子图，而是利用BERT架构的注意力机制来定位易受攻击的线，这比 IVDetect方法更细粒度。最后，我们通过实验将我们的LineVul 方法与7种基线方法(IVDetect、Reveal、SySeVR、 Devign、Russell等、VulDeePecker、BoW+RF)进行比较，并对粗粒度(即功能级)和细粒度(即行级)漏洞预测场景进行评估。

我们得出结论，LineVul比现有的漏洞预测方法更准确、更经济、更细粒度。因此，我们期望我们的LineVul可以帮助安全分析人员以经济有效的方式找到易受攻击的线路。此外，我们建议在未来的研究中使用注意机制来提高基于Transformer的模型的可解释性，因为本文已经展示了使用注意机制进行线级漏洞定位的实质性好处，它优于其他与模型无关的技术(例如DeepLift, LIG和SHAP)。

新颖性&贡献。据我们所知，本文的主要贡献如下:

(1) 我们的LineVul，一种基于变压器的线路级脆弱性预测方法，解决了现有脆弱性预测方法的各种局限性;

(2) 实验结果证实了我们的 LineVul比现有的漏洞预测方法更准确、更经济、更细粒度。

2 技术介绍

我们将LineVul方法设计为两步方法：预测有漏洞的函数和定位有漏洞的行

1. 函数级漏洞预测

函数级漏洞预测包括两个主要步骤：

①BPE 子词标记化.：在步骤1中，我们利用字节对编码(Byte Pair Encoding,BPE)方法通过两个主要步骤构建我们的tokenizer。1a生成合并操作以确定单词应该如何被分割，1b根据子单词词汇表应用合并操作。具体来说，BPE将把所有单词分解成字符序列，并识别出应该合并成一个新符号的最常见符号对(例如，两个连续字符对)。BPE是一种算法，它可以将罕见词分割成有意义的子词，同时保留常见词(即不会将常见词分割成更小的子词)。例如，在图1中，函数名 将被分割成一个子单词列表，即["un"， "Prem"， "ul"， "Sk"， "Image"， "To"， "Prem"， "ul"]。常用词“Image”被保留下来，其他生僻词被拆分。在对各种函数名进行标记时，使用BPE子词标记将有助于减少词汇表的大小，因为它将罕见的函数名拆分为多个子组件，而不是直接将完整的函数名添加到字典中。本文将BPE方法应用于CodeSearchNet语料库上，生成了一个适合于源代码语料库预训练的语言模型的子词分词器（subword tokenizer）。

②LineVul 模型构建：在步骤2中，我们构建了一个基于BERT架构的LineVul模型，并利用了Feng等人预训练的的初始权值。在步骤2a中，LineVul对subword-tokenized的函数执行单词和位置编码，以生成每个单词及其在函数中的位置的嵌入向量。然后，在步骤2b中，向量被输入到BERT体系结构中，该体系结构是由12个Transformer编码器块组成的堆栈。每个编码器由一个多头自注意力层和一个全连接的前馈神经网络组成。最后，在步骤2c中，为了对给定的函数进行二分类，将输出向量送入单个线性层。我们将在下面描述每个步骤。

2a 字和位置编码： 源代码由多个token组成，其中每个token的含义严重依赖上下文(即围绕token)以及每个token在函数中的位置。因此，捕获代码上下文及其在函数中的位置非常重要，特别是对于函数级漏洞预测来说。这一步的目的是生成编码向量，捕获代码tokens的语义含义及其在输入序列中的位置。为此，对于每个subword-tokenized的token，我们生成两个向量：

(1) 单词编码向量，表示给定代码标记与其他代码标记之间的有意义关系;

(2) 位置编码向量，表示给定标记在输入序列中的位置。记号编码向量根据词嵌入矩阵生成，其中为词汇表大小，d为嵌入大小。根据位置嵌入矩阵生成位置编码向量，其中c为上下文大小，d为嵌入大小。最后，将单词编码向量和位置编码向量连接起来，以产生Transformerd编码器块的输入向量。

2b 具有双向自注意力的12个Transformer编码器堆栈： 在这个步骤中，编码向量被输入到12个仅使用编码器的Transformer块的堆栈中(即BERT体系结构)。每个编码器块由两部分组成，即双向多头自注意力层和全连接前馈神经网络。下面，我们简要介绍了多头自我注意和前馈神经网络。

利用多头自注意力层计算每个编码token的注意力权值，得到一个注意力向量。双向自注意力的使用允许每个token分别关注其左右两侧的上下文。生成的注意力权重用于指示Transformer模型应该注意哪些代码语句。通常，自注意力机制用于获取全局依赖关系，其中注意权重表示序列中的每个代码token是如何受到序列中所有其他单词的影响的，从而允许我们的LineVul方法捕获每个代码token之间的依赖关系，以生成更有意义的表示。

自注意力机制采用了信息检索的概念，该概念使用点积操作计算每个代码token的相关得分，其中每个token与其他token进行一次交互。自注意力机制依赖于三个主要组件:Query( )、Key( )和Value( )。Query是当前代码token的表示形式，用于根据存储在Key向量中的所有其他token的keys对它们进行评分。通过对查询向量和关键向量进行点积，得到每个token的注意力分数。然后使用Softmax函数将注意力得分归一化为概率，以获得注意力权重。最后，通过对值向量与注意力权重向量进行点积来更新值向量。

在我们的LiNEVuL中使用的自注意是一个缩放的点积自注意，其中注意分数除以。

我们采用的自注意机制可以用以下公式来概括:

为了获取更丰富的输入序列语义，我们采用多头机制来实现自注意，使得模型能够在不同位置上联合关注来自不同代码表示子空间的信息。对于 d-dimension Q , K和V ，我们将这些向量切分成h个头，每个头具有的dimension。在所有的自注意力操作之后，每个头将被再次连接到一个全连接的前馈神经网络中，包括两个线性转换，中间有一个ReLU激活。多头机制可归纳为:

其中

用于拼接后线性投影到预期维度。

2c 单线性层：多重线性变换与非线性激活函数(即ReLU)构建在叠加Transformer编码器块中，因此，最终编码器输出的表示是有意义的。我们只需要一个线性层将代码表示映射到二进制标签中，即 ˆ= + 。

2 行级漏洞定位

给定一个被LineVul预测为有漏洞的函数，我们通过利用Transformer体系结构中的自注意力机制来定位有漏洞的行，来执行行级的漏洞定位。直觉上，对预测贡献最大的tokens很可能是有漏洞的tokens。

对于函数中的每个subword token，在步骤3中，我们总结了12个Transformer编码器块中的每个块的自注意力分数。在获得subword-token注意力分数之后，我们将这些分数集成为行分数。我们通过换行控制字符(即\n)将整个函数分解为许多tokens列表(每个token列表代表一行)。最后，对于每个token分数列表，我们将其总结为一个注意力行分数，并按降序排列行分数。

3 实验评估

在本节中，我们将介绍我们研究问题的三个动机、我们研究的数据集和我们的实验设置。

3.1 研究问题

为了评估我们的LineVul方法，我们制定了以下三个研究问题。

RQ1 函数级漏洞预测的LineVul有多准确?

最近，Li等人提出了IVDetect，一种最先进的细粒度漏洞预测方法。然而，正如在上文中提到的，IVDetect有三个关键的限制，导致不准确和粗粒度的预测。因此，我们提出我们的LineVul方法来应对这些挑战。因此，我们调查我们的LineVul的准确性是否优于最先进的函数级漏洞预测方法。

RQ2 行级漏洞定位的精确度有多高?

需要行级漏洞预测来帮助开发人员识别漏洞行的细粒度位置，而不是浪费时间检查非漏洞行。虽然IVDetect可以识别有漏洞的函数的子图，但这些子图仍然包含许多安全分析师需要检查的代码行，这些代码行仍然是粗粒度的。因此，我们研究了LineeVul用于行级漏洞预测的准确性。

RQ3 我们的LineVul在行级漏洞定位方面的成本效益是什么?

漏洞预测的主要目标之一是，通过最少的工作量发现最大数量的漏洞，帮助安全分析师以经济有效的方式定位漏洞行。因此，当安全分析人员决定在实践中部署高级方法时，检查源代码所需的工作量就成为一个关键问题。因此，我们研究了LineVul用于行级脆弱性预测的成本效益。

3.2 实验设置

数据分割：与Li等人的实验类似，我们使用了相同的数据分割方法，即整个数据集被分割成80%的训练数据、10%的验证数据和10%的测试数据。

函数级模型实现：为了实现用于函数级漏洞预测的LineVul方法，我们主要使用两个Python库，即transformer和Pytorch。transformer库提供API访问基于transformer模型的架构和预训练的权值，而PyTorch库支持训练过程(例如反向传播和参数优化)中的计算。我们下载了由Feng等人预训练的CodeBERT tokenizer和CodeBERT模型。我们使用我们的训练数据集对预训练模型进行微调，为我们的漏洞预测任务获得合适的权重。该模型在NVIDIA RTX 3090显卡上进行了微调，训练时间约为7小时10分钟。如公式1所示，利用交叉熵损失对函数级预测值与现实真值标签之间的模型进行更新和优化，其中 为类的数量， 为现实概率分布(one-hot)， 是预测的概率分布。为了获得最佳的微调权重，我们使用验证集来按epoch监控训练过程，并根据针对验证集(而不是测试集)的最佳F1-score来选择最佳模型。( , ) = − Σ ( ) ( ) (1)

行级模型实现： 为了实现用于行级漏洞预测的LineVul方法，我们在推理期间使用由微调模型返回的自注意力矩阵。我们总结每个token的点积注意力分数，因此，每个token都有一个注意力分数。我们将tokens集成到行中，并汇总每个token的得分，以获得行得分。然后使用行得分对行进行优先排序，其中较高的行得分表示该行很可能是有漏洞的行。

微调的超参数设置： 对于我们的LineVul方法的模型架构，我们使用CodeBERT的默认设置，即12个Transformer Encoder块，768的隐藏size和12个注意头。我们遵循Feng等人提供的相同的微调策略。在训练过程中，学习率被设置为2e-5，采用线性计划，在整个训练过程中学习率线性递减。我们使用反向传播与AdamW优化广泛采用微调的基于Transformer的模型，以更新模型和最小化损失函数。

RQ1 函数级漏洞预测的LineVul有多准确?

方法： 为了回答这个RQ，我们关注于函数级漏洞预测，并将我们的LineVul与其他七个baseline模型进行比较，描述如下:

IVDetect利用特征注意图卷积网络(GCN)进行漏洞预测，使用了源代码中的五种类型的特征表示(即sub-tokens序列、AST子树、变量名称和类型、数据依赖上下文和控制依赖上下文);

ReVeal利用门控图神经网络(GGNN)来学习图的源代码属性;

Devign利用门控图神经网络(GGNN)自动学习源代码的图属性(即AST、CFG、DFG和代码序列);

SySeVR使用代码语句、程序依赖关系和程序切片作为特征，对几种基于RNN的模型(LR、MLP、DBN、CNN、LSTM等)进行分类;

VulDeePecker利用双向LSTM网络进行语句级漏洞预测;

Russell等利用基于基于RNN的模型对漏洞进行预测。

BoW+RF (LineDP/JITLine)使用单词的bag作为特征，并使用随机森林模型进行软件缺陷预测。

与Li等人的类似，我们使用三种二元分类度量来评估我们的LineVul，即精确度、召回率和F1-score。准确度衡量的是模型正确预测有漏洞的函数占预测有漏洞的的函数数量的比例，其计算公式为 / + 。召回率衡量被正确预测为有漏洞的函数占实际有漏洞的的函数的数量的比例，其计算公式为 /( + )。F-measure是精度和召回率的调和均值，其计算公式为2×Precision×Recall/(precision + recall)。我们使用0.5的概率阈值作为截断值来表示预测标签。这些测量值的范围从0到1，其中1表示最高的精度。

结果： 图3展示了我们的LineVul的实验结果，以及根据我们的三个评价度量(即F1、精确度、召回率)的七个baseline方法。

图3 (RQ1)我们的LineVul实验结果与七个baseline比较，用于函数级漏洞预测。(↗)F1值越高，精确度越高，召回率越高=越好。

在F-measure方面，图3显示LineVul达到了最高的F-measure 0.91，而最先进的方法达到了0.19-0.35的F-measure。这一发现表明，LineVul大大提高了目前的技术水平，提高了160%-379%，中值提高了250%。在精确度方面，图3显示LineVul的精确度最高为0.97，而最先进的方法的精确度为0.12-0.48。这一发现表明，LineVul显著提高了state-of-the-art水平102%-708%，中位数提高了439%。在召回率方面，图3显示LineVul的最高召回率为0.86，而最先进的方法的召回率为0.17-0.86。这一发现表明，LineVul大幅提高了state-of-the-art水平的16%-406%，中位数提高了65%。

换句话说，我们的结果表明，在Transformer体系结构中使用语义和语法特性比使用源代码图形属性的现有工作表现更好。我们的发现与最近许多研究的发现不同，这些研究发现，在漏洞预测中，使用图数据属性(如数据依赖图、抽象语法树、控制流图和数据流图)往往优于使用句法和语义特征。这是因为先前研究中作为baseline的基于RNN的方法(如RNN、LSTM、GRU)(1)是在特定项目数据集上训练的;以及(2)捕获源代码的长期依赖关系有困难。与之前的研究不同，我们的结果证实了我们的LineVul方法比state-of-the-art方法更准确，强调使用CodeBERT预训练的语言模型(在百万GitHub仓库上进行训练)和使用Transformer架构捕获源代码的长期依赖关系的巨大好处，从而显著改进了函数级的漏洞预测方法。

RQ2 行级漏洞定位的精确度有多高?

方法： 为了回答这个问题，我们着重于评估行级漏洞定位的准确性。因此，我们从我们的方法正确预测的有漏洞的函数开始。然后，我们执行步骤3，以确定对于预测为有漏洞的给定函数，哪些行可能是有漏洞的。因此，给定函数中的每一行都有自己的分数(即行级分数)。由于RQ1中的其他方法不是为行级本地化而设计的，所以我们不将我们的方法与它们进行比较。相反，我们将我们的LineVul方法与其他5种通常用于深度学习模型的模型无关技术进行了如下比较:

层综合梯度(Layer Integrated Gradient, LIG)是一种公理性的路径归因方法，通过近似模型输出沿给定基线到输入的路径(行)上的梯度与输入的积分，为每个输入特征赋予一个重要分数；

显著性(Saliency)在输入端对网络进行一阶泰勒展开，梯度只是模型线性表示中每个特征的系数。这些系数的绝对值可以表示特征的重要性；

DeepLift将每个神经元的激活与其参考激活进行比较，并根据差异分配贡献分数;

DeepLiftSHAP扩展了DeepLift算法，并使用DeepLift方法逼近SHAP值;

GradientSHAP通过从基线分布中随机抽样计算期望梯度来近似SHAP值;

CppCheck是一个用于C和C++编程语言的静态代码分析工具。

为了评估我们的行级漏洞定位LineVul方法，我们使用以下两个度量:

Top-10 Accuracy衡量有漏洞的函数的百分比，在这些有漏洞的函数中至少出现一条排名在top-10的实际有漏洞的行。直觉告诉我们，如果有漏洞的行没有出现在前10名中，安全分析师可能会忽略它们，类似于任何推荐系统。因此，top-10 accuracy将帮助安全分析师更好地理解行级漏洞定位方法的准确性。

Initial False Alarm(IFA, 初始假警报)衡量安全分析人员直到找到给定函数的第一个实际有漏洞的行而需要检查的错误预测行(即没有漏洞的行被错误预测为有漏洞的或错误警报)的数量。IFA计算为安全分析人员在发现第一个实际有漏洞的行之前必须检查的假警报总数。低的IFA值表明安全分析师在检查假警报上花费的精力更少。

结果： 图4展示了我们的LineVul的实验结果，以及根据我们的两个评价衡量指标(即Top-10 Accuracy和IFA)的五种baseline方法。

图4 (RQ2)我们的自注意力方法和其他五种方法的Top-10 Accuracy和IFA。(↗) Top10准确度越高=更好，(↘)IFA越低=更好

我们的LineVul达到了0.65的Top-10 Accuracy，比其他baseline方法的精度高12%-25%。在Top-10 accuracy方面，由图4可知，LineVul的Top-10 Accuracy最高，为0.65，而state-of-the-art方法的Top-10 Accuracy为0.52-0.58。在IFA方面，图4显示LineVul的IFA中值最低，为1，而baseline方法的IFA中值为3-4。这一结果表明，LineVul相比baseline方法显著提高了67%-75%，中位数提高了67%。这些结果证实了我们的LineVul方法比行级漏洞定位的baseline方法更准确。

换句话说，我们的结果表明，注意力机制优于其他模型无关技术。在行级缺陷预测文献中，先前的研究经常利用一种LIME model-agnostic的技术来解释基于DL/ML的缺陷预测模型的预测结果。然而，这种model-agnostic技术被认为是一种外部的model-agnostic技术(例如一个model-agnostic在模型被训练之后被应用于解释一个黑盒DL/ML模型)，而不是一个内在的model-agnostic技术(例如，一个DL/ML模型本身是可解释的，因此不需要在之后应用外部的model-agnostic技术)。虽然众所周知，深度学习是复杂且难以解释的，但本文是第一次尝试利用注意力机制来解释基于Transformer的模型的预测，强调行级漏洞定位的注意力机制的实质好处。

RQ3 我们的LineVul在行级漏洞定位方面的成本效益是什么?

方法： 为了回答这个问题，我们着重于评估我们的LineVul方法对于行级漏洞定位的成本效率。在实际场景中，最经济有效的行级漏洞预测方法应该帮助安全分析人员用最少的功夫最大数量地找到实际有漏洞的行。因此，让我们假设测试数据集中的18,864个函数(即504,886 LOC)是安全分析师必须检查的函数。为了衡量我们的方法的成本效益，我们首先从我们的LineVul方法获得预测。然后根据预测概率对预测的函数进行排序。在预测为有漏洞的函数中，我们根据我们的方法步骤3得到的行得分，按降序对行进行排序。因此，得分最高的行(即可能是有漏洞的)将排在最前面。然后，我们使用以下衡量指标来评估成本效益:

Effort@20%Recall衡量安全分析人员为了找出实际的20%有漏洞的行而花费的工作量(以LOC度量)。它的计算方式为：测试集中全部的用于定位实际有漏洞行的20%的LOC除以总LOC。Effort@20%Recall的值很低，这表明安全分析人员可能会花费较少的精力来寻找20%的实际有漏洞的行。

Recall@1%LOC衡量给定一定工作量(即给定测试数据集的LOC的前1%)可以找到的真实有漏洞的行的比例(即正确预测)。Recall@1%LOC的计算方法是:测试集中前1%的行中正确定位的有漏洞的行的总数除以测试集中实际有漏洞的行的总数。Recall@1%LOC的值很高表明一种方法可以将许多实际有漏洞的行排列在顶部。

结果： 图5展示了我们的LineVul的实验结果，以及根据我们的两个评估衡量指标(即Effort@20%Recall和Recall@1%LOC)的五种baseline方法。

图5 (RQ3)我们的自注意力方法和其他五种方法的Effort@20%Recall和Recall@1%LOC。(↘)低Effort@20%Recall=更好，(↗)高Recall@1%LOC=更好

我们的LineVul Effort@20%Recall值为0.75，比其他baseline方法低29%-53%。在Effort@20%Recall方面，图5显示LineVul的Effort@20%Recall最低为0.75%，而baseline方法的Effort@20%Recall最低为1.06%-1.60%。这意味着安全分析师可能会在测试数据集中检查0.75%的LOC (0.75%*504,886= 3786 LOC)，以便找到20%的实际有漏洞的行(即20%的Recall)。因此，这一发现表明，我们的方法可以帮助安全分析师花费更少的资金，以找到相同数量的实际有漏洞的行。

在Recall@1%LOC中，图5显示LineVul获得了最高的Recall@1%LOC值0.24，而baseline方法获得了Recall@1%LOC值0.13-0.19，这表明LineVul比baseline方法显著地提高了26%-85%，中值提高了85%。这一发现意味着我们的方法可以帮助安全分析人员发现比其他baseline方法更多的实际有漏洞的行，尽管他们需要检查的工作量是一样的。

转述：韩廷旭