Pre-training by Predicting Program Dependencies for Vulnerability Analysis Tasks论文分享

今天分享的论文是《Pre-training by Predicting Program Dependencies for Vulnerability Analysis Tasks》

原文链接：Pre-training by Predicting Program Dependencies for Vulnerability Analysis Tasks | IEEE Conference Publication | IEEE Xplore

开源代码1：https://github.com/ZJU-CTAG/PDBERT

开源代码2：Pre-training by Predicting Program Dependencies for Vulnerability Analysis Tasks

这篇论文是关于预训练模型应用漏洞检测的，主要技术是微调。

漏洞分析对软件安全至关重要。受预训练模型在软件工程任务中成功应用的启发，本研究聚焦于利用预训练技术增强对漏洞代码的理解并提升漏洞分析能力。预训练模型的代码理解能力与其预训练目标高度相关。代码的语义结构（如控制依赖和数据依赖）对漏洞分析至关重要，但现有预训练目标要么忽略此类结构，要么仅关注学习如何使用这些结构，而对学习语义结构分析知识的可行性和优势尚未展开研究。为此，本研究提出了两个新颖的预训练目标：控制依赖预测（Control Dependency Prediction, CDP）和数据依赖预测（Data Dependency Prediction, DDP），分别用于仅基于代码片段的源代码预测语句级控制依赖和标记级数据依赖。在预训练过程中，CDP和DDP可引导模型学习分析代码中细粒度依赖关系所需的知识。预训练后，该模型在微调阶段可增强对漏洞代码的理解，还能直接用于分析部分函数和完整函数的依赖关系。为验证预训练目标的有效性，本文预训练了一个名为PDBERT的Transformer模型，结合CDP和DDP，并在三个漏洞分析任务（漏洞检测、漏洞分类和漏洞评估）上进行微调，同时对程序依赖分析进行评估。实验结果表明，PDBERT受益于CDP和DDP，在三个下游任务上均取得了最先进的性能。此外，PDBERT在部分函数和完整函数中预测控制依赖和数据依赖的F1分数分别超过99%和94%。 

本研究的主要贡献： 

1. 引入了预训练代码模型以整合端到端程序依赖分析所需知识的思想，并提出了两个新颖的预训练目标（CDP和DDP）来实现这一思想。 

2. 构建了一个名为PDBERT的预训练模型，结合了CDP、DDP和MLM。据本文所知，这是首个能够分析语句级控制依赖和标记级数据依赖的神经模型。 

3. 进行了内在和外在评估，结果表明PDBERT能够准确高效地识别部分和完整函数中的程序依赖关系，并能促进各种漏洞分析任务。 

程序依赖（包括控制依赖和数据依赖）反映了代码的语义结构。程序依赖图（PDG）通过不同类型的边显式表示控制依赖和数据依赖，常被先前的研究用于漏洞分析。下图展示了一个示例函数及其PDG。在PDG中，每个节点表示一条语句或一个谓词，有两种类型的边：控制依赖边和数据依赖边。控制依赖边（图中的绿色虚线）表示谓词和语句之间的控制流关系，可用于推断语句所依赖的控制条件。数据依赖边（图中的红色实线）表示定义-使用关系，每条边都标记有一个变量，该变量在源节点中定义并在目标节点中使用。通过预测程序中的控制依赖和数据依赖，模型可以学习加强程序中计算相关部分之间的连接，从而更好地捕获代码的语义结构。 

PDBERT仅将代码片段作为输入，该代码片段可以是部分函数或完整函数。与现有预训练模型一样，PDBERT自身输出上下文嵌入序列。这些嵌入通过头部（即输出层）转换为任务输出，例如分类标签或标记序列。在预训练期间，不同的头部连接到PDBERT以执行不同的预训练任务，如下图所示。预训练后，本文可以简单地将PDBERT与CDP和DDP的头部连接，以推断代码片段中的语句级控制依赖和标记级数据依赖。对于下游任务，将新的头部（通常是MLP）连接到PDBERT，并在特定任务的数据集上与PDBERT一起微调。每个头部是任务特定的，经过训练以包含基于PDBERT输出的上下文嵌入生成任务输出的信息。此类信息通常对其他任务无用。由于MLM、CDP和DDP的头部是任务特定的，且不用于下游任务，因此本文认为它们没有副作用。请注意，PDBERT在微调期间仅将源代码作为输入，不需要解析输入程序或构建其PDG。本文假设程序依赖的知识已在预训练期间被PDBERT学习和吸收，并且该知识可以提升对程序依赖敏感的下游任务。 

本文在三个广泛使用的C/C++函数级漏洞检测数据集上评估PDBERT，即ReVeal、Big-Vul和Devign。它们的统计数据如下表所示：

其中，#Func、#Vul 和 #Non-Vul 分别指所有函数、易受攻击（存在漏洞）的函数以及不易受攻击（不存在漏洞）的函数的数量。CC 表示圈复杂度。VC 和 VA 分别指用于漏洞分类（vulnerability classification）和漏洞评估（vulnerability assessment）的数据集。

ReVeal由ReVeal论文发布，从Linux Debian内核和Chromium收集。该数据集不平衡且接近现实场景。本文按照先前的工作随机将数据集的70%/10%/20%划分为训练、验证和测试集。考虑到数据不平衡，遵循先前的工作，本文使用F1分数作为评估指标。 

Big-Vul由Fan等人发布，从348个GitHub项目中收集。考虑到该数据集的大规模，遵循先前的工作，本文将其随机划分为训练、验证和测试集，比例为80%/10%/10%。由于该数据集高度不平衡，本文也使用F1分数作为评估指标。 

Devign由Devign论文发布，包含从FFmpeg和Qemu收集的27.2K个函数。该数据集是平衡的，但不如ReVeal和Big-Vul现实。本文使用该数据集是因为它包含在常用的CodeXGLUE基准中。本文还使用CodeXGLUE的训练/验证/测试划分。遵循CodeXGLUE的设计和先前的工作，本文使用准确率作为评估指标。 

讨论与分析

局限性 

由于计算资源的限制，遵循CodeBERT，本文在这项工作中将PDBERT的最大序列长度设置为512。因此，当用于程序依赖分析时，PDBERT无法正确分析标记数超过512的代码片段。然而，如5.1.2节所示，从开源项目构建的测试集中超过83%的C/C++函数符合此长度限制。基于本文的内在评估，本文认为PDBERT在实践中对于分析大多数函数是有效和高效的。当应用于下游任务时，PDBERT在处理前会截断长代码片段。尽管这可能对性能产生负面影响，但本文的外在评估表明，PDBERT仍然有效地提升了下游任务。另一方面，此限制来自Transformer架构，而非本文的预训练技术。可以通过简单地增加PDBERT的最大序列长度来缓解此限制，这需要更多的计算资源和先进的硬件。考虑到本文的预训练技术与底层模型架构正交，还可以采用专门用于长代码片段的模型架构（如LongFormer和LongCoder）来解决此限制。由于这项工作专注于预训练技术，本文选择了最广泛使用的Transformer架构，遵循先前工作的设置，并将处理长代码片段留作未来的工作。 

此外，CDP专注于语句级控制依赖，不处理单个语句内的控制流（如三元运算符）。未来的工作可以通过考虑此类控制流来改进CDP。 

有效性威胁 

PDBERT的泛化性 

PDBERT在C/C++程序上进行预训练，可能不适用于其他编程语言。然而，本文提出的预训练目标与语言无关。从业者可以使用CDP和DDP为其他甚至多种编程语言预训练他们的模型。 

使用Joern的局限性 

为了构建CDP和DDP的真实标签，本文使用Joern提取PDG。因此，期望PDBERT预测由Joern分析的程序依赖关系，并可能与Joern共享相同的局限性。然而，Joern是最先进的基于源代码的程序依赖分析工具，它具有高精度并已广泛用于与漏洞相关的任务。此外，本文的预训练技术可以被视为训练模型学习Joern的程序依赖分析知识，这仍然是有益的（如本文的外在评估所示）。本文的内在评估展示了PDBERT在程序依赖分析中的独特优势，并至少评估了PDBERT在分析完整函数方面与Joern的接近程度。因此，本文认为Joern的局限性不会影响这项工作的结论。 

数据污染 

预训练数据集中的一些函数可能也出现在某些微调数据集中，因为它们都是从开源存储库收集的。然而，预训练数据集包含漏洞函数和良性函数，并且在预训练期间每个函数的任务特定标签是未知的。此外，先前的工作表明，数据污染可能对预训练模型的性能影响很小。因此，本文认为此威胁是有限的。 

相关工作

预训练代码模型 

现有的预训练代码模型通常可分为仅编码器模型、仅解码器模型和编码器-解码器模型。作为概念验证，本研究聚焦于预训练仅编码器模型以用于漏洞分析。对于现有的仅编码器预训练模型，部分模型直接采用为自然语言设计的预训练任务。例如，CodeBERT采用了掩码语言模型（MLM）和替换标记检测（RTD），这些模型能够捕获代码的自然特性。近年来，一些预训练技术被提出以帮助模型学习代码的语法结构，如DISCO通过预训练目标预测输入中被掩码的抽象语法树（AST）节点类型。此外，部分前期研究还利用对比学习（CL）目标辅助模型学习程序间的高级功能相似性。与这些工作不同，本文的预训练目标旨在帮助模型学习捕获代码的语义结构。 

据本文所知，仅有Code-MVP和GraphCodeBERT这两种预训练代码模型考虑了代码的语义结构（如控制流和数据流信息）。Code-MVP在预训练阶段显式将控制流图（CFG）作为输入，目标是从CFG中学习表示；GraphCodeBERT则通过输入变量序列和未被掩码的数据流边来预测部分被掩码的数据流边以完成预训练，且在推理时需要提取程序的变量序列和所有数据流边作为输入，这表明其目标同样是从数据流中学习表示。相比之下，PDBERT具有以下优势：首先，其同时考虑控制依赖和数据依赖，能够学习更全面的程序依赖知识；更重要的是，PDBERT旨在学习编码代码语义结构的表示，并被设计为“吸收”端到端程序依赖分析所需的知识，这是此前未被研究的方向。如第1节所述及第5节实验验证，该设计带来了若干独特且重要的优势：（1）PDBERT仅以源代码为输入，可正确处理“无法解析”的代码；（2）PDBERT学习的知识具有更强的通用性，能更好地提升下游任务；（3）PDBERT可直接用于分析语句级控制依赖和标记级数据依赖，这是现有神经模型无法实现的能力。因此，本文认为PDBERT相较现有预训练代码模型具有显著的创新性贡献。 

深度学习驱动的漏洞分析 

诸多基于深度学习的方法已被提出用于漏洞分析任务。对于漏洞检测，当前最先进的方法通常先借助静态分析工具（如Joern）提取程序依赖图（PDG），再基于PDG将程序表示为不同形式，例如代码片段、基于语法的表示、基于语义的向量表示或图结构表示，随后利用双向长短期记忆网络（Bi-LSTM）、卷积神经网络（CNN）或图神经网络（GNN）等不同神经模型提取输入程序的特征向量以实现漏洞检测。与这些方法相比，PDBERT无需将程序依赖作为输入，且能处理部分代码，具有独特优势。此外，现有方法通常从零开始训练以学习如何在特定任务中使用程序依赖，而PDBERT通过预训练学习程序依赖知识，可轻松应用于多种漏洞分析任务。 

在漏洞分类领域，现有研究大多聚焦于基于专家整理的漏洞描述预测常见弱点枚举（CWE）类别，仅有μVulDeePecker是基于源代码预测CWE类别的工作，但其采用了过程间分析，而本研究聚焦于函数级漏洞分类。对于漏洞评估任务，多数前期工作基于漏洞描述而非源代码预测通用漏洞评分系统（CVSS）指标，DeepCVA利用CNN和多任务学习实现提交级漏洞评估，Le等人提出基于漏洞语句预测CVSS指标，但这类方法需要大量人工标注或依赖已知的漏洞修复信息。

做个总结

本研究提出了两个新颖的预训练目标（CDP和DDP），以帮助将端到端程序依赖分析的知识融入神经模型并提升漏洞分析能力。CDP和DDP分别旨在仅基于源代码预测程序中的语句级控制依赖和标记级数据依赖。作为概念验证，本文构建了预训练模型PDBERT，其在部分函数和完整函数中预测控制依赖和数据依赖的F1分数分别超过99%和94%。此外，本文在漏洞检测、漏洞分类和漏洞评估三个任务上对PDBERT进行微调，实验结果表明，PDBERT受益于CDP和DDP，性能显著优于现有最先进的基线模型。 

未来，本文计划探索以下方向：（1）研究如何结合CDP和DDP预训练编码器-解码器模型；（2）评估PDBERT在更多下游任务中的有效性，并尝试在多语言代码语料库上预训练PDBERT；（3）对于静态分析工具可成功提取输入代码程序依赖的下游任务，探究将本预训练技术与显式提供的程序依赖结合是否能进一步提升性能。