R2Vul: Learning to Reason about Software Vulnerabilities with Reinforcement Learning and Structured Reasoning Distillation
This is the replication package accompanying our paper "R2Vul: Learning to Reason about Software Vulnerabilities with Reinforcement Learning and Structured Reasoning Distillation".
You can find the paper appendix in the supplementary material (appendix.pdf).
The project is structured as follows.
.
├── scripts/ # bash scripts to run specific training and inference
├── src/ # source code of the project
├── data_annotation/ # code to generate data annotations, positive/negative reasoning, and LLM-as-a-judge rankings.
├── inference/ # inference scripts
├── training/ # training scripts for CLS, SFT, and ORPO
├── Dockerfile # Dockerfile to setup the docker container
├── requirements.txt # required Python libraries
├── tgi_serve.sh # script to serve a HuggingFace model using TGI.
├── tgi_serve_r2vul.sh # script to serve a local model using TGI.
We provide a Dockerfile to setup a docker image to run our code.
The image is based on nvidia/cuda:12.4.0 for Ubuntu. Depending on your machine, you can look for an appropriate base image that runs cuda 12.4.0 on dockerhub.
- Build the docker image
docker build -t r2vul-image .This builds the docker image and ensures Python 3 is properly installed.
- Create the docker container
Next, you can instantiate a new docker container based on the image we just created.
docker run -it --name r2vul -d -v R2Vul:/r2vul --gpus all r2vul-imageNote that if you plan to run inference with a local TGI model, add --network=r2vul-inference-net to have both containers on the same network.
You can then start the container and attach to it:
docker start r2vul
docker exec r2vul -it bash
cd r2vul # go to the source code directory- Setup the virtual environment
Create a new virtual environment and install the required Python libraries.
python -m venv venv
pip install -r requirements.txt
source venv/bin/activate # activate the venvNote that if you do not wish to use Docker, you can simply rely on the Python venv, but we cannot guarantee that everything will run smoothly.
We share our raw results for all experiments discussed in the paper under the /results folder:
/main- runs related to the main experiments (Section 5.1). Includes acompute_metrics.pyscript to compute paired bootstrap tests./external_test_set- runs related to experiments on the external test set (Section 5.2). Includes acompute_metrics.pyscript to compute paired bootstrap tests./data_ablation- runs related to the data ablation (Section 5.3). Includes arender_plot.pyscript to reproduce Figure 2./class_imbalance- runs related to class imbalance experiments (Section 5.4). Includes arender_plot.pyscript to reproduce Figure 3./calibration- runs related to model calibration experiments (Section 5.5). Includes arender_plot.pyscript to reproduce Figure 4.
We make our datasets and models checkpoints available on Zenodo: https://zenodo.org/records/16741648.
raw_dataset.json- our raw data mined from NVD.paired_dataset.json- pre/post-commit function pairs./r2vul_dataset- our dataset for training and inference./external_java_test- external Java test set in Huggingface format.
Create a folder /data containing each dataset.
orpo.zip- Qwen2.5-Coder-Instruct models fine-tuned using R2Vul (with ORPO).sft.zip- Qwen2.5-Coder-Instruct models fine-tuned using SFT.cls.zip- Models fine-tuned using CLSMSIVD.zip- MSIVD model checkpointVulLLM.zip- VulLLM model checkpoint
If you want to replicate experiments with existing checkpoints, you need to download them and place them in a /runs folder.
For SFT / R2Vul models, the adapter needs to be merged with the base model before serving them using TGI (see merge.py script).
We provide bash scripts to run specific experiments. You simply have to run them and change some variables depending on what you want to run.
MSIVD:/scripts/main/run_inference_msivd.shVulLLM:/scripts/main/run_inference_vulllm.shCLS:/scripts/main/run_inference_cls.shCommercial LLMs:/scripts/main/run_inference_oai.shCoT:/scripts/main/run_inference_tgi.sh(runtgi_serve.shfirst)SFTandR2Vul:/scripts/main/run_inference_tgi.sh(runtgi_serve_r2vul.shfirst)
CodeBERT (CLS):/scripts/main/run_training_cls_codebert.shQwen2.5-Coder-Instruct (CLS):/scripts/main/run_training_cls_qwen.shSFT:/scripts/main/run_training_sft.shR2Vul:/scripts/main/run_training_r2vul.sh
CLS:/scripts/external_test_set/run_inference_cls.shSFT:/scripts/external_test_set/run_inference_sft.sh(runtgi_serve_r2vul.shfirst)R2Vul:/scripts/external_test_set/run_inference_r2vul.sh(runtgi_serve_r2vul.shfirst)
SFT:/scripts/data_ablation/run_training_sft.shR2Vul:/scripts/data_ablation/run_training_r2vul.sh
For inference, run tgi_serve_r2vul.sh with a specific checkpoint, then run inference using run/main/run_inference_tgi.sh.
SFT:/scripts/class_imbalance/run_inference_sft.shR2Vul:/scripts/class_imbalance/run_inference_r2vul.sh
R2Vul:/scripts/calibration/run_inference_r2vul.sh