Explaining the Contributing Factors for Vulnerability Detection in Machine Learning

TL;DR

The paper investigates what factors contribute to effective vulnerability detection in machine learning applied to software code. By systematically varying tokenization, embedding methods, architectural metrics, and learning models across 17 Java projects and three labelled datasets, it identifies a practical baseline: token-based bag-of-words representations with a Random Forest classifier, achieving robust within-domain performance and a notable baseline against deep models. Architectural metrics offer little improvement, while token-based approaches provide the strongest signal for vulnerability signatures; cross-domain transferability is limited by domain heterogeneity, highlighting the need for transfer learning to realize cross-project applicability. Overall, the study provides a rigorous, data-driven baseline and a clear set of directions for advancing transferable vulnerability detection in real-world software engineering contexts.

Abstract

There is an increasing trend to mine vulnerabilities from software repositories and use machine learning techniques to automatically detect software vulnerabilities. A fundamental but unresolved research question is: how do different factors in the mining and learning process impact the accuracy of identifying vulnerabilities in software projects of varying characteristics? Substantial research has been dedicated in this area, including source code static analysis, software repository mining, and NLP-based machine learning. However, practitioners lack experience regarding the key factors for building a baseline model of the state-of-the-art. In addition, there lacks of experience regarding the transferability of the vulnerability signatures from project to project. This study investigates how the combination of different vulnerability features and three representative machine learning models impact the accuracy of vulnerability detection in 17 real-world projects. We examine two types of vulnerability representations: 1) code features extracted through NLP with varying tokenization strategies and three different embedding techniques (bag-of-words, word2vec, and fastText) and 2) a set of eight architectural metrics that capture the abstract design of the software systems. The three machine learning algorithms include a random forest model, a support vector machines model, and a residual neural network model. The analysis shows a recommended baseline model with signatures extracted through bag-of-words embedding, combined with the random forest, consistently increases the detection accuracy by about 4% compared to other combinations in all 17 projects. Furthermore, we observe the limitation of transferring vulnerability signatures across domains based on our experiments.

Figures (5)

• Figure 1: Learning software vulnerability as a classification task
• Figure 2: The data flow of the feature engineering and learning. The feature engineering is consistent for all the classification models. The modeling part illustrates the structure of the revised ResNet model that consists of one convolutional layer, one dense layer and 7 ResNet blocks. Each ResNet block is composed of 16 layers.
• Figure 3: OWASP token distribution. Most of the tokens have fewer than 20 occurrences.
• Figure 4: Juliet token distribution. Most of the tokens are have fewer than 25 occurrences.
• Figure 5: All Android projects token distribution. The majority of the tokens have fewer than 500 occurrences.