PaperlayerView exact revision on arXiv

Source-linked AI summary

VulBERTa: Simplified Source Code Pre-Training for Vulnerability Detection

Hazim Hanif, Sergio Maffeis

arXiv:2205.12424v1cs.CRcs.AIcs.LG

TL;DR

Software vulnerability detection is challenging, while existing approaches often rely on complex representations or substantial resources. VulBERTa uses custom tokenisation and RoBERTa pre-training for C/C++ code, then fine-tunes simple classifiers; it reports state-of-the-art performance across datasets and benchmarks despite smaller training data and models.

  • Problem

    Software vulnerability detection remains challenging, motivating methods that can learn useful source-code representations for reliable detection.

  • Method

    VulBERTa combines custom syntax-preserving C/C++ tokenisation with RoBERTa masked-language-model pre-training and MLP or CNN fine-tuning.

  • Results

    VulBERTa achieves state-of-the-art detection performance across datasets and benchmarks, including 57.92% F1 on Draper and 99.59% weighted F1 on muVuldeepecker.

  • Takeaways & Limitations

    The results support using learned C/C++ representations with simple downstream architectures for software vulnerability detection.

  • Takeaways & Limitations

    Testing includes occasional label inaccuracies, training remains expensive, larger configurations and hyperparameter sweeps were not explored, and systematic detection analysis was incomplete.

Abstract

from arXiv · show

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.

I. INTRODUCTION

VulBERTa uses custom code tokenisation and RoBERTa pre-training to learn C/C++ representations for vulnerability detection. It achieves strong performance across datasets and benchmarks with a simple, compact approach.

  • Rising software CVEs have driven continued research into vulnerability detection using static analysis, dynamic analysis, and machine-learning approaches.
  • VulBERTa learns C/C++ syntactic and semantic representations with a Transformer-based architecture, without manually extracting code artefacts.
  • Its custom tokeniser combines BPE with predefined C/C++ keywords, punctuation, and API tokens while preserving source-code syntax.
  • The model uses RoBERTa and masked language modelling, then transfers its pre-trained weights to MLP- and CNN-based vulnerability classifiers.
  • 57.92% F1 on Draper and 99.59% weighted F1 on muVuldeepecker demonstrate performance on binary and multi-class vulnerability detection tasks.
  • VulBERTa-MLP and VulBERTa-CNN achieve state-of-the-art detection performance across datasets and top-three positions on CodeXGLUE and D2A.

II. RELATED WORK

Related work spans raw-source, gadget-based, graph-based, and dataset-focused approaches to vulnerability detection. These methods differ in their representations, preprocessing requirements, and treatment of dataset quality.

  • Software vulnerability detection remains challenging, motivating deep-learning methods alongside established static and dynamic analysis approaches.
  • An early raw-source approach extracted deep-learning features and classified them with Random Forest, achieving 90.4% AUC on a real-world dataset.
  • VulDeePecker uses code gadgets centered on library/API calls, while SySeVR extends this representation with syntax, data-dependency, and control-dependency information.
  • Graph-based methods such as Devign and DeepWukong learn from data- and control-dependency graphs, while DeepTective targets PHP vulnerabilities.
  • REVEAL addresses class imbalance and duplicates, whereas D2A constructs a curated benchmark through differential analysis of source-code version pairs.

B. Pre-trained models of source code

Pre-trained source-code models adapt Transformer-based language modelling ideas to programming languages and increasingly incorporate syntax or graph structure. VulBERTa organizes this process into tokenisation, pre-training, and task-specific fine-tuning.

  • Pre-training applies Transformer architectures to learn general representations of textual data for language-specific tasks.
  • C-BERT learns Abstract Syntax Tree features during pre-training and outperforms existing approaches on AST node tagging and vulnerability detection.
  • CodeBERT pre-trains on code and natural-language pairs as well as unimodal programming-language data, achieving high BLEU scores relative to RoBERTa models.
  • DOBF uses code deobfuscation to learn structural information, while GraphCodeBERT incorporates graph structure through graph-guided masked attention.
  • GraphCodeBERT shows strong performance against RoBERTa and CodeBERT on code-clone detection, translation, and refinement tasks.
  • VulBERTa's architecture combines custom-vocabulary tokenisation, representation-building pre-training, and function-level fine-tuning for vulnerability classification.

1) Parser:

VulBERTa converts C/C++ source into model-ready representations through parsing, customized BPE tokenization, and encoding, then uses the pre-trained model for vulnerability classification.

  • 1) Parser:: Clang parses each function into code tokens while preserving syntactic structure, after comments are removed.
  • 2) Tokenisation:: Modified BPE further splits parsed input, while predefined C/C++ keywords, punctuation, and API names preserve selected syntactic and semantic meaning.The vocabulary contains 451 predefined tokens excluded from BPE; literals and identifiers undergo BPE.
  • 2) Tokenisation:: Encoding converts code tokens into tensors, using maximum sequence lengths of 512 for pre-training and 1024 for fine-tuning.The longer fine-tuning limit contains more than 90% of actual samples on average without truncation.
  • Pre-training: Pre-training learns C/C++ representations with masked language modelling across different software projects.Combining projects generalizes representation knowledge across coding styles and increases pre-training robustness.
  • C. Fine-tuning: Fine-tuning reuses the pre-trained model for vulnerability detection through either an MLP or a TextCNN classifier.The MLP uses the full architecture, whereas TextCNN uses extracted embedding weights and freezes them during training.

IV. DATASETS

The study uses public function-level C/C++ datasets from open-source repositories and synthetic samples, separating pre-training data from vulnerability-detection evaluation data.

  • IV. DATASETS: The datasets contain function-level C/C++ source code from varied codebases, including open-source repositories and synthetic samples.
  • IV. DATASETS: Datasets are divided into pre-training and fine-tuning categories according to their prevalent use.
  • IV. DATASETS: Masked language modelling on the GitHub and Draper datasets provides the pre-training task.

1) GitHub:

The GitHub pre-training dataset provides over one million extracted C/C++ functions from more than one thousand open-source repositories, while the associated fine-tuning datasets support vulnerability evaluation.

  • 1) GitHub:: 1,101,075 C/C++ functions were extracted from 1060 open-source GitHub repositories for the GitHub dataset.
  • 1) GitHub:: Joern extracts individual functions from downloaded files because VulBERTa targets function-level security vulnerabilities.
  • 2) Draper:: The Draper dataset contains 1,274,366 C/C++ functions from Debian, GitHub repositories, and the Juliet test suite.Its contents range from highly documented production code to synthetic test samples.
  • Fine-tuning datasets: Fine-tuning uses selected datasets compiled by security researchers to evaluate vulnerability-detection approaches.
  • Vuldeepecker: The Vuldeepecker dataset combines real-world NVD samples with synthetic SARD samples and is frequently used as a C/C++ vulnerability-detection benchmark.
  • Vuldeepecker: Vuldeepecker fine-tuning uses binary labels checked by three static analysers and assigned by security experts.

3) REVEAL:

The evaluation spans real-world vulnerability datasets from Linux, Chromium, QEMU, FFmpeg, and other open-source projects, using varied data formats and labeling procedures.

  • 3) REVEAL:: The REVEAL dataset is a binary dataset drawn from the Linux Debian kernel and Chromium.It was introduced to address duplication and unrealistic vulnerable-class distributions in existing datasets.
  • 4) muVuldeepecker:: The muVuldeepecker dataset is multiclass and uses code gadgets from NVD and SARD rather than ordinary function-level source code.
  • 5) Devign:: The Devign dataset contains function-level C/C++ code from QEMU and FFmpeg, manually labeled and verified by security researchers in two rounds.
  • 6) D2A:: The D2A dataset covers FFmpeg, httpd, Libav, LibTIFF, Nginx, and OpenSSL, with issues labeled through differential analysis of static-analyser reports.
  • Evaluation: The study describes pre-training, fine-tuned VulBERTa-MLP and VulBERTa-CNN models, and evaluation across vulnerability datasets and benchmarks.
  • Experimental setup: Fine-tuning experiments use PyTorch with CUDA, while pre-training uses Google Compute Engine virtual machines equipped with two NVIDIA Tesla A100 GPUs.

2) Performance criteria:

The evaluation uses multiple standard metrics and dataset-specific preferred metrics to compare VulBERTa with established sequence-based baselines. Pre-training and fine-tuning follow defined model configurations, splits, and optimization schedules.

  • Performance criteria: Experiments report dataset-specific metrics alongside TN, FN, TP, FP, accuracy, precision, recall, F1-score, ROC-AUC, PR-AUC, and MCC.Reporting the metrics used in each dataset’s original work supports fairer comparison.
  • Performance criteria: VulBERTa is compared with Baseline-BiLSTM and Baseline-TextCNN, two widely used sequence-based vulnerability-detection techniques.The BiLSTM learns forward and backward code-sequence relationships, while TextCNN applies convolution and pooling to source-code input.
  • Training setup: Pre-training uses Masked Language Modelling on Draper and GitHub datasets across small, medium, and base RoBERTa configurations.Sessions run for up to 500,000 steps and last 72–96 hours depending on configuration.
  • Training setup: VulBERTa-MLP and VulBERTa-CNN are fine-tuned separately on each dataset for vulnerability detection.Fine-tuning uses a maximum of 10 epochs, a learning rate of 0.00003, schedulers, and original or 80/10/10 dataset splits.

B. Evaluation on selected datasets

Across selected datasets, VulBERTa achieves strong binary and multi-class vulnerability-detection results, exceeding compared approaches across varied data compositions and class-balance conditions.

  • Dataset results: 95.76% precision and 93.03% F1 score make VulBERTa-MLP higher than Vuldeepecker’s reported 91.9% precision and 92.9% F1 score.The reported false-positive and false-negative rates are 0.39% and 9.14%, respectively.
  • Dataset results: 55.86% MCC gives VulBERTa-CNN a 2.26% increase over the performance reported in the Draper comparison.MCC is used because Draper’s vulnerable and non-vulnerable classes are imbalanced.
  • Dataset results: 45.27% F1 score exceeds the 41.25% reported in the comparison, despite omitting the proposed data-rebalancing techniques.Instead, class weights are assigned during fine-tuning, and the reported TPR is 2.57% higher.
  • Dataset results: 99.59% weighted F1 exceeds the 96.28% reported for muVuldeepecker’s 40-CWE multi-class task.The false-negative rate falls from 5.53% to 0.41%.
  • Dataset results: VulBERTa-MLP correctly assigns more than 90% of vulnerable samples to their respective CWE-190 and CWE-191 classes.These classes represent integer overflow and underflow.
  • Overall comparison: The models surpass compared approaches across synthetic and real-world data, balanced and imbalanced classes, and binary and multi-class tasks.The comparisons use each dataset’s preferred evaluation metrics.

C. Evaluation on benchmarks

VulBERTa achieves leading or near-leading benchmark performance while using comparatively small models and pre-training data. The results also support the importance of its learned code representation and tokenisation design.

  • CodeXGLUE: 64.75% accuracy places VulBERTa-MLP third on CodeXGLUE’s defect-detection leaderboard.VulBERTa-MLP uses 55.07% fewer model parameters than CoTexT, while VulBERTa-CNN’s model size is below 1% of CoTexT and 2% of C-BERT.
  • D2A: 62.30% accuracy places VulBERTa-MLP first on D2A’s relevant function task, while VulBERTa-CNN ranks second with 60.68%.Both models outperform C-BERT on D2A despite its larger pre-training dataset and parameter count relative to VulBERTa-CNN.
  • Overall benchmark comparison: VulBERTa-MLP has the best preferred evaluation metric on 3 of 4 datasets, while VulBERTa-CNN has the best score on the remaining dataset.The two models differ negligibly across the full fine-tuning results and rank next to each other on both benchmarks.
  • Discussion: The pre-trained VulBERTa model’s inherited code-representation knowledge is identified as a possible explanation for the close performance of the two fine-tuning architectures.This interpretation accompanies the observation that VulBERTa-MLP and VulBERTa-CNN perform similarly across datasets and benchmarks.
  • Efficiency: 125M parameters and 2.28M C/C++ functions yield state-of-the-art detection performance with smaller models and pre-training data than cited comparators.The comparison names CoTexT with 375M parameters and C-BERT with 8.5M functions.
  • Discussion: The tokenisation approach is believed to make syntactic and semantic information more accessible to the simple neural architectures used for fine-tuning.This claim is based on preliminary analysis of different tokenisation techniques.
  • Model simplicity: 55.86% MCC from the 2M-parameter VulBERTa-CNN exceeds the 52% reported for the more complex 3GNN model on Draper.The cited 3GNN combines Crystal Graph Convolution and Self Attention Pooling.

E. Limitations

VulBERTa's evaluation is constrained by label inaccuracies, training-resource limits, and the lack of systematic in-the-wild novel 0-day detection. The authors nevertheless report state-of-the-art performance with relatively low model complexity.

  • Dataset and evaluation constraints: Occasional label inaccuracies in vulnerability datasets undermine quantitative test results when noise is present during evaluation.The authors note that deep learning may tolerate training noise, but testing noise remains problematic.
  • Training constraints: Limited resources prevented exploring significantly larger VulBERTa configurations, model combinations, and hyper-parameter sweeps.Different configurations could therefore achieve higher performance.
  • Deployment scope: The main limitation is the lack of a systematic attempt to detect novel 0-day vulnerabilities in open-source projects in-the-wild.The authors attribute this to the challenge of manually reviewing false positives and propose explainability techniques for future work.
  • Reported strengths: VulBERTa achieves state-of-the-art performance while standing out for conceptual simplicity and low parameter complexity.The conclusion characterizes the model as relatively small compared with more complicated alternatives.
Loading 2205.12424v1…