[论文解读] VulBERTa: Simplified Source Code Pre-Training for Vulnerability Detection

Hazim Hanif, Sergio Maffeis|arXiv (Cornell University)|May 25, 2022

Software Engineering Research被引用 9

一句话总结

VulBERTa 提出了一种针对 C/C++ 代码漏洞检测的简化、轻量级预训练方法，采用自定义分词流程，结合 BPE 与预定义的代码标记（关键字、标点符号和 API 调用），以保留语法和语义结构。该模型在 228 万个函数（含 1.25 亿参数）上进行预训练，在多个数据集和基准测试中达到最先进性能，包括在 Draper 数据集上取得 57.92% 的 F1 分数，在 muVuldeepecker 数据集上达到 99.59% 的 F1 分数，尽管模型复杂度极低，仍优于更大规模的模型。

ABSTRACT

This paper presents VulBERTa, a deep learning approach to detect security vulnerabilities in source code. Our approach pre-trains a RoBERTa model with a custom tokenisation pipeline on real-world code from open-source C/C++ projects. The model learns a deep knowledge representation of the code syntax and semantics, which we leverage to train vulnerability detection classifiers. We evaluate our approach on binary and multi-class vulnerability detection tasks across several datasets (Vuldeepecker, Draper, REVEAL and muVuldeepecker) and benchmarks (CodeXGLUE and D2A). The evaluation results show that VulBERTa achieves state-of-the-art performance and outperforms existing approaches across different datasets, despite its conceptual simplicity, and limited cost in terms of size of training data and number of model parameters.

研究动机与目标

• 开发一种简化、低复杂度的 C/C++ 源代码漏洞检测预训练方法。
• 通过自定义分词流程整合语法和语义信息，提升代码表征学习能力。
• 在参数更少、训练数据更少的情况下，实现漏洞检测的最先进性能。
• 评估预训练模型在多样化数据集和基准测试中的可迁移性与鲁棒性。
• 在保持高检测准确率的同时，实现高效微调，采用轻量级分类器（MLP 和 CNN）

提出的方法

• 自定义分词流程将字节对编码（BPE）与 C/C++ 关键字、标点符号和常见 API 函数名的预定义标记相结合，以保留语法和语义结构。
• 模型采用 RoBERTa 作为主干架构，基于来自开源项目的 228 万个 C/C++ 函数，通过掩码语言建模（MLM）进行预训练。
• 预训练在经过筛选的真实世界 C/C++ 代码数据集上进行，专注于语法和语义知识的建模，不依赖显式的依赖关系或数据流分析。
• 微调采用两种轻量级分类器：多层感知机（VulBERTa-MLP）和卷积神经网络（VulBERTa-CNN）。
• 该方法避免使用复杂的图结构或序列结构架构，转而依赖分词质量和预训练过程来编码深层代码表征。
• 评估在包括 Vuldeepecker、Draper、REVEAL、muVuldeepecker 以及基准测试 CodeXGLUE 和 D2A 在内的多个数据集上，涵盖二分类和多分类漏洞检测任务。

实验结果

研究问题

• RQ1 一个简化、小规模的预训练模型是否能在不依赖复杂架构设计的情况下，实现漏洞检测的最先进性能？
• RQ2 一种能保留语法和语义结构的自定义分词流程，在漏洞检测的代码表征学习中能提升多大程度？
• RQ3 在参数更少、预训练数据更少的情况下，小型模型是否仍能在基准数据集上超越大型模型？
• RQ4 预训练表征在多样化数据集和漏洞检测任务中的可迁移性如何？
• RQ5 当与强大的预训练代码表征结合时，轻量级微调头（MLP 和 CNN）是否能实现高性能？

主要发现

• VulBERTa 在具有挑战性的 Draper 数据集上取得了 57.92% 的最先进 F1 分数，尽管模型规模小且训练数据有限，仍优于现有方法。
• 在多分类的 muVuldeepecker 数据集上，VulBERTa-MLP 达到 99.59% 的加权 F1 分数，显示出在更简单、更平衡数据上的强大性能。
• 在 CodeXGLUE 基准测试中，VulBERTa-MLP 以 64.75% 的准确率排名第一，VulBERTa-CNN 以 60.68% 的准确率排名第二，两者均优于更大的模型（如 C-BERT）。
• 在 D2A 基准测试中，VulBERTa-MLP 在 'function' 任务上达到 62.30% 的准确率，排名第一；VulBERTa-CNN 达到 60.68%，排名第二。
• 尽管仅使用 1.25 亿参数和 228 万个预训练样本，VulBERTa 仍优于 C-BERT（810 万个样本，1.1 亿参数）和 CodeBERT（850 万个样本，1.25 亿参数）等模型。
• 该模型在不同微调头上的表现稳健：VulBERTa-MLP 在 4 个数据集中的 3 个上达到最佳 PEM 分数，VulBERTa-CNN 在剩余 1 个上表现最佳。