Pre-training by Predicting Program Dependencies for Vulnerability Analysis Tasks

全文摘要

本篇论文探讨了如何利用预训练技术提高软件漏洞分析任务的理解能力。作者提出了两个新的预训练目标：控制依赖预测和数据依赖预测，旨在仅根据源代码预测代码片段中的语句级控制依赖和标记级数据依赖。通过预训练，模型可以学习到分析代码中细粒度依赖关系所需的知识，并在微调时提升对漏洞代码的理解能力。实验结果表明，使用CDP和DDP进行预训练的Transformer模型PDBERT在三个下游任务上取得了最先进的性能，同时在部分函数和完整函数中预测控制和数据依赖性的F1得分均超过了99%和94%。

论文速读

论文方法

方法描述

该论文提出了一种名为PDBERT的预训练模型，用于程序依赖关系分析任务。PDBERT通过将源代码作为输入，并使用多层Transformer编码上下文嵌入来捕捉程序结构信息。它使用三个预训练任务：Masked Language Model（MLM）、Statement-Level Control Dependency Prediction（CDP）和Token-Level Data Dependency Prediction（DDP），以帮助模型学习关于程序结构的知识。其中，MLM用于捕捉自然性和语法结构，CDP用于预测语句级别的控制依赖关系，DDP用于预测标记级别的数据依赖关系。

方法改进

与现有的基于自回归模型或序列到序列模型的解决方案相比，PDBERT不需要解析输入程序或构建其PDG，这使得它的应用更加方便和高效。此外，PDBERT还采用了掩码策略来处理标记类型的差异，以帮助模型更好地理解程序结构。

解决的问题

PDBERT旨在解决程序依赖关系分析问题，这是一个重要的编程挑战，因为程序中的依赖关系可以影响程序的行为、性能和安全性。PDBERT的目标是为程序依赖关系分析提供一个强大而高效的工具，可以帮助开发人员更轻松地理解和调试程序。

论文实验

本文主要介绍了对基于代码的预训练模型PDBERT在三个方面的应用进行了实验比较：intrinsic evaluation（内省评价）、extrinsic evaluation（外推评价）和vulnerability assessment（漏洞评估）。具体来说，在内省评价中，作者使用了两个数据集来测试PDBERT对于程序依赖分析的能力，并通过F1分数来衡量其性能。结果表明，PDBERT在控制依赖和数据依赖方面都表现出了很高的准确性。此外，作者还比较了PDBERT与一些常见的预训练模型和其他非预训练模型的性能差异，发现PDBERT在所有任务上都表现出更好的性能。

在外推评价中，作者将PDBERT应用于函数级漏洞检测、分类和评估等任务，并与其他预训练模型和非预训练模型进行了比较。结果显示，PDBERT在所有任务上都比其他模型表现更好，这进一步证明了PDBERT的有效性。

在漏洞评估中，作者使用了一个数据集来测试PDBERT对于CVSS指标预测的能力，并使用了宏平均F1分数和多类Matthews相关系数来衡量其性能。结果表明，PDBERT在所有指标上都比最好的基线模型表现更好，这也进一步证明了PDBERT的有效性。

综上所述，本文通过多个实验验证了PDBERT在不同任务上的有效性，并与其他模型进行了比较，得出了PDBERT在软件安全领域具有很好的应用前景。

论文总结

文章优点

• 本文提出了一种新的预训练技术PDBERT，通过控制依赖性和数据依赖性的预测来帮助神经模型学习程序依赖关系知识，并在漏洞检测、分类和评估等下游任务中取得了显著提升。
• PDBERT采用了CodeBERT作为基础模型，并提出了两种新的预训练目标CDP和DDP，能够有效地捕捉代码中的依赖关系信息。
• 实验结果表明，PDBERT比现有基线模型更有效，并且具有较高的可扩展性和通用性。

方法创新点

• 本文提出的CDP和DDP是基于控制流和数据流的依赖关系预测，能够更好地捕捉代码中的依赖关系信息。
• PDBERT是一种针对程序依赖关系的预训练技术，可以应用于多种编程语言和下游任务中。

未来展望

• 未来可以进一步研究如何将CDP和DDP应用到其他类型的神经网络结构中，以提高其效率和准确性。
• 可以考虑将PDBERT与其他静态分析工具相结合，以提高其性能和适用范围。
• 可以探索如何使用多语言语料库来训练PDBERT，以支持更多的编程语言和应用场景。