Source-linked AI summary
Finetuning Large Language Models for Vulnerability Detection
Alexey Shestov, Rodion Levichev, Ravil Mussabayev, Evgeny Maslov, Anton Cheshkov, Pavel Zadorozhny
TL;DR
Vulnerability detection is complex, and this paper investigates whether finetuned large language models can improve function-level classification, including under real-world class imbalance. It finetunes WizardCoder and explores training procedures and loss choices; the reported results improve on ContraBERT in specified dataset settings. The authors also note that quality remains below business needs and that function-level data limits coverage of vulnerabilities spanning functions.
Problem
Function-level vulnerability detection is complex, and real-world data imbalance can make standard training objectives favor negative examples over positive ones.
Source: §1 · p. 2Method
The study finetunes WizardCoder and experiments with training configurations and loss formulations for balanced and imbalanced vulnerability datasets.
Source: §1 · p. 2, §8 · p. 14Results
ROC AUC improved from 0.66 to 0.69 over ContraBERT on X1 without P3, and F1 improved from 0.22 to 0.27 on the dataset with P3.
Source: §8 · p. 14Takeaways & Limitations
Finetuning WizardCoder improves vulnerability-detection quality, though the reported quality remains below what business tasks require.
Source: §8 · p. 15Function-level dataset granularity limits the study's coverage of vulnerabilities that span multiple functions.
Source: §7 · p. 14
Abstract
from arXiv · showhide
This paper presents the results of finetuning large language models (LLMs) for the task of detecting vulnerabilities in source code. We leverage WizardCoder, a recent improvement of the state-of-the-art LLM StarCoder, and adapt it for vulnerability detection through further finetuning. To accelerate training, we modify WizardCoder's training procedure, also we investigate optimal training regimes. For the imbalanced dataset with many more negative examples than positive, we also explore different techniques to improve classification performance. The finetuned WizardCoder model achieves improvement in ROC AUC and F1 measures on balanced and imbalanced vulnerability datasets over CodeBERT-like model, demonstrating the effectiveness of adapting pretrained LLMs for vulnerability detection in source code. The key contributions are finetuning the state-of-the-art code LLM, WizardCoder, increasing its training speed without the performance harm, optimizing the training procedure and regimes, handling class imbalance, and improving performance on difficult vulnerability detection datasets. This demonstrates the potential for transfer learning by finetuning large pretrained language models for specialized source code analysis tasks.
1 INTRODUCTION
The introduction frames vulnerability detection as a complex task with separate representation and classification stages, and focuses on improving the classification model. It motivates finetuning LLMs and investigating training regimes for balanced and imbalanced datasets.
1 INTRODUCTION
- The study investigates whether larger LLMs can improve on CodeBERT-style vulnerability detection and whether limited model capacity constrains performance.The authors compare text-represented functions with CodeBERT-based models.Source: §1 · pp. 1–2
- Vulnerability detection has two stages—representing functions and classifying those representations—and improving either can improve overall solution quality.The paper focuses on the classification stage, leaving function representation unchanged.Source: §1 · p. 2
- The paper treats finetuning as a promising direction after prior work found simple ChatGPT prompting no noticeably better than random prediction on vulnerability detection.Source: §1 · p. 2
- Real-world vulnerable-function data is highly imbalanced, and standard cross-entropy may favor the dominant negative class over positive examples.The introduction identifies focal loss and sample weighting as techniques for emphasizing the minority class.Source: §1 · p. 2
- The experiments seek effective model architectures, training hyperparameters, and loss formulations for finetuning on balanced and imbalanced datasets.The project also examines model choice, limited GPU resources, decoder-model classification, and underused sequence length that slows training.Source: §1 · p. 2
2 RELATED WORK
Prior work covers prompting, retrieval-augmented generation, and finetuning LLMs for vulnerability detection, with reported challenges including false positives and limited Java-specific research. The paper situates its Java WizardCoder study within that developing area.
2 RELATED WORK
- A review reports that 82% of surveyed LLM vulnerability studies used fine-tuning, while class imbalance, noisy data, and scarce labeled datasets remained challenges.CodeBERT appeared in 38.5% of reviewed studies; GPT-3.5 and GPT-4 appeared in 10.3% and 2.6%.Source: §2 · p. 3
- ChatGPT and GPT-3 without additional tuning performed no better than random guessing in one study, while another found ChatGPT lagging specialized vulnerability models.Source: §2 · p. 3
- A GPT model using BM-25 and TF-IDF retrieval for few-shot examples improved precision, recall, and F1 on Devign over traditional methods.Source: §2 · p. 3
- Prior C/C++ LLM studies reported high recall but high false-positive rates, while Java vulnerability detection remained underexplored.The related-work passage distinguishes this paper's Java WizardCoder finetuning from primarily prompt-based prior approaches.Source: §2 · p. 3
3 IMPLEMENTATION DETAILS
The implementation adapts a decoder LLM to function-level vulnerability classification while addressing model choice, limited hardware, and wasted computation from short functions. The authors use WizardCoder, a classification-focused loss, and batch packing, which raised training throughput substantially.
3.1 CHALLENGES
- The implementation challenges were selecting a transferable, usable model, fitting it to limited GPU resources, aligning decoder training with classification, and speeding training for short inputs.The standard fixed-length approach wastes computation because most functions are under 50 tokens; packing multiple functions into a sequence mitigates this issue.Source: §3.1 · p. 4, §3.5 · p. 5
3.2 CHOOSING AN LLM TO FINETUNE
- WizardCoder was selected for its stronger YES/NO responses than CodeGen, although experiments found no significant performance difference between WizardCoder and StarCoder.The paper uses “StarCoder” to refer specifically to the WizardCoder variant in its subsequent writing; more elaborate experiments are needed to establish differences among the family models.Source: §3.2 · pp. 4–5
3.3 ADAPTING LLMS FOR LIMITED COMPUTATIONAL RESOURCES
- LoRA enabled adaptation of the 13-billion-parameter WizardCoder under hardware constraints by reducing the trainable parameters and memory requirements.The authors report reducing the model to 25 million parameters with r = 8, alpha = 32, and dropout = 0.05.Source: §3.3 · p. 5
3.4 ADAPTING LLMS FOR CLASSIFICATION
- The classification loss uses the final-token probability matched to the class label, rather than the generative next-token objective.This formulation is intended to make weight updates follow the classification objective without influence from next-token generation.Source: §3.4 · p. 5
3.5 SPEEDING UP THE CLASSIFICATION TRAINING
- Batch packing concatenates complete function examples into training sequences, with separate next-token-prediction and binary-classification regimes.StarCoder-based sequences are 2,048 tokens and CodeBERT-based sequences are 512 tokens; prediction uses one test function without packing.Source: §3.5 · pp. 6–7
- 13x speedup raised training throughput from 1,500 to 20,000 samples per hour through batch packing of short functions.Packing reduces computation spent on padding; the authors say it makes practical finetuning possible and is generally applicable to decoder models with short inputs.Source: §3.5 · p. 7
- Dynamic gradient accumulation could stabilize batch size while retaining packing efficiency, but requires non-trivial HuggingFace Transformers training-loop modifications.Training used one input sequence per batch because of hardware memory limits.Source: §3.5 · p. 7
4 EXPERIMENTAL SETUP
The experimental setup builds and filters function-level Java vulnerability datasets, evaluates checkpoints using validation ROC AUC, and compares WizardCoder with ContraBERT. The datasets include balanced and highly imbalanced variants, and the authors provide code and data for reproduction.
4.1 DATASETS
- CVEfixes may contain inconsistencies and errors because it is automatically constructed, while the Manually-Curated Dataset uses expert-found fix commits.VCMatch covers only ten popular repositories and contains only fix commits.Source: §4.1 · p. 7
- Filtering to commits that change only one function produces X1, intended to reduce false positive labels from cleanup, refactoring, or irrelevant changes.The source datasets include CVEfixes, a Manually-Curated Dataset, and VCMatch; function labels are derived from vulnerability-fixing commits, with pre-change versions labeled vulnerable and post-change versions non-vulnerable.Source: §4.1 · pp. 7–8
- X1 without P3 has 1,334 samples and an approximately 1:1 positive-to-negative balance; adding P3 yields 22,945 samples with an approximately 1:34 balance.P3 consists of unchanged functions from files associated with patches; the majority of negatives in the larger dataset come from P3.Source: §4.1 · p. 8
4.2 EVALUATION METRICS
- The best checkpoint is selected by validation ROC AUC, then evaluated on the test set; reported metrics include ROC AUC and vulnerable-class F1.The F1 threshold is determined on validation data, and most cases report both validation and test metrics.Source: §4.2 · p. 8
4.3 PRETRAINED MODELS
- WizardCoder is compared with ContraBERT, a CodeBERT-family model pretrained with masked-language modeling and contrastive tasks.The study identifies ContraBERT as the state-of-the-art CodeBERT-family model for its experiments.Source: §4.3 · p. 8
4.4 FULL REPRODUCIBILITY PACKAGE
- The authors state that all code and data needed to reproduce the experiments are available in their GitHub repository.The repository is identified in the paper as vul-llm-finetune.Source: §4.4 · p. 9
5 RESEARCH QUESTIONS
The research questions examine whether StarCoder-based models outperform CodeBERT-based models on balanced and imbalanced vulnerability detection, and how decoder training choices affect performance. They also investigate context length and techniques for handling class imbalance.
5.1 COMPARISON TO CODEBERT-BASED MODELS
- RQ1 and RQ2 ask whether StarCoder-based models outperform CodeBERT-based models on balanced and imbalanced vulnerability datasets, respectively.The experiments vary training recipes, including epochs and loss functions, across datasets with and without P3 to examine hyperparameter transferability.Source: §5 · p. 9
- RQ3 tests whether next-token prediction, the standard decoder-LLM training approach, is effective for vulnerability classification.The approach predicts the next token alongside a binary vulnerability label and trains with cross-entropy for both objectives.Source: §5 · p. 9
- RQ4 asks whether dynamic batch sizes from batch packing harm model quality; RQ5 compares mean and sum loss reduction, and RQ6 tests increased context size.The context-size experiment compares 512-token and 2,048-token inputs on the dataset without P3.Source: §5 · p. 9
5.3 IMBALANCED CLASSIFICATION IMPROVEMENT
- RQ7 tests whether focal loss and sample weighting effectively address class imbalance by emphasizing minority positive examples on the full dataset.The experiments systematically vary the focal-loss gamma parameter and sample weights; ROC AUC, F1, accuracy, and the optimal classification threshold are evaluated.Source: §5 · p. 9
6 EXPERIMENTAL RESULTS
Across balanced and imbalanced X1 settings, finetuned WizardCoder is compared with ContraBERT and evaluated under training and imbalance-focused variants. WizardCoder leads on both datasets, while ablations identify useful training choices and imbalance methods yield modest gains.
6.1 COMPARISON TO CODEBERT-BASED MODELS
- WizardCoder surpassed ContraBERT in ROC AUC and F1 on balanced X1 without P3 and imbalanced X1 with P3.The reported gains in ROC AUC were smaller on the imbalanced dataset.Source: Table 2 · p. 10, Table 3 · p. 10, §6.1 · p. 10
6.2 ABLATION STUDY
- Next-token prediction reached ROC AUC 0.75 on imbalanced X1 with P3, below classification-only CodeBERT at 0.85 and WizardCoder at 0.86.It nevertheless exceeded random performance.Source: Table 4 · p. 11, §6.2 · p. 11
- Dynamic batch sizing reduced training time by up to 13 times, and limiting batch size did not improve model quality.The study reports no significant influence from limiting the maximum number of functions per batch.Source: §6.2 · p. 11, Table 5 · p. 11
- Sum loss reduction outperformed mean reduction by better handling uneven packed-sequence lengths, with each function contributing equally to the gradient.The comparison used the validation portion of X1 without P3.Source: §6.2 · p. 12, Table 6 · p. 12
- Reducing WizardCoder’s context from 2048 to 512 tokens left its validation ROC AUC unchanged at 0.69.The authors suggest gains mainly reflect improved code representations rather than increased context size.Source: Table 7 · p. 12, §6.2 · p. 12
6.3 IMBALANCED CLASSIFICATION IMPROVEMENT
- Focal loss with γ = 1 modestly improved ROC AUC over standard cross-entropy with γ = 0, and focal loss produced better-calibrated models.Quality peaked at γ = 1 before declining at higher values; the authors suggest noisy labels may impede identifying hard cases.Source: §6.3 · p. 13
- A 3x weight on P1 and P2 produced marginal ROC AUC and F1 gains, while excessive weighting such as 30x hurt validation and test metrics.The authors identify moderate weighting as a possible optimum, with further tuning potentially yielding additional gains.Source: Table 9 · p. 13, §6.3 · p. 13
- Combining focal loss and sample weighting yielded only minor gains, with best ROC AUC 0.877 versus 0.86 for baseline training.The authors conclude that more advanced methods are needed to account explicitly for severe class imbalance.Source: Table 10 · p. 14, §6.3 · p. 14
7 THREATS TO VALIDITY
The study identifies threats tied to dataset construction, training inputs, and function-level scope. Dataset size may also limit how well the sample represents the real vulnerability distribution.
7 THREATS TO VALIDITY
- The dataset is not split by projects, and project-based splitting might have significantly decreased quality.The authors identify this as a threat to the objectiveness of the study.Source: §7 · p. 14
- Batch packing lets the model see preceding functions during classification-regime finetuning, potentially limiting training effectiveness.The authors say random placement should mitigate negative impact, and test inference uses one function per input sequence.Source: §7 · p. 14
- Function-level data limits coverage of vulnerabilities spanning multiple functions.Source: §7 · p. 14
- The relatively small dataset may limit representativeness of the true vulnerability distribution.Source: §7 · p. 14
8 CONCLUSION
Finetuned WizardCoder improves vulnerability detection over previous CodeBERT-based models, with a reported ROC AUC of 0.69 on the dataset without easy negative examples. The study also reports faster training and benchmark contributions, while identifying unresolved quality and data challenges.
8 CONCLUSION
- ROC AUC 0.69 on the dataset without easy negative examples is reported as state-of-the-art and an improvement over previous CodeBERT-based models.The authors attribute the improvement likely to WizardCoder’s larger model capacity and pretraining corpus.Source: §8 · p. 14
- Over 13x speedup in training time came from batch packing designed to mitigate small sequence lengths.The strategy enables faster iteration and tuning.Source: §8 · p. 14
- ROC AUC improved from 0.66 to 0.69 over ContraBERT on X1 without P3, while F1 improved from 0.22 to 0.27 on the dataset with P3.Source: §8 · p. 14
- Focal loss with sample weighting improved results from 0.86 to 0.878 on the highly imbalanced dataset, but stronger methods are still needed to emphasize minority positives.Source: §8 · p. 14
- The study collected a smaller, more precise vulnerability benchmark with quality labels free from errors stemming from irrelevant code and cleanup changes.The passage contrasts its size with most available datasets.Source: §8 · p. 15
- Further work is needed especially for training with P3, scarce minority data, and the gap between achieved and business-task quality.The authors point to curriculum learning, active sampling, data augmentation, vulnerability representation, and broader project-level context as research directions.Source: §8 · p. 15
A DISTRIBUTION OF CWES
The appendix presents distributions of the top 19 CWE categories and other categories in the test set, shown separately with and without P3.
A DISTRIBUTION OF CWES
- Figure 1 presents the distribution of the top 19 CWE categories and others in the test set with P3.Source: Figure 1 · p. 17
- The two figures distinguish test-set CWE distributions by whether P3 is included.Source: Figure 1 · p. 17, Figure 2 · p. 18
- Figure 2 presents the distribution of the top 19 CWE categories and others in the test set without P3.Source: Figure 2 · p. 18