# VISION: Robust and Interpretable Code Vulnerability Detection Leveraging Counterfactual Augmentation

arXiv: https://arxiv.org/abs/2508.18933 · HTML: https://mlsanity.com/paper/2508.18933/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 42 references that are in the map.
- First posted on arXiv 26 Aug 2025.

## Builds on (most cited first, 25 of 25)

- [A Unified Approach to Interpreting Model Predictions](https://mlsanity.com/paper/1705.07874/) (2017)
- ["Why Should I Trust You?": Explaining the Predictions of Any Classifier](https://mlsanity.com/paper/1602.04938/) (2016)
- [CodeBERT: A Pre-Trained Model for Programming and Natural Languages](https://mlsanity.com/paper/2002.08155/) (2020)
- [Counterfactual Explanations without Opening the Black Box: Automated Decisions and the GDPR](https://mlsanity.com/paper/1711.00399/) (2017)
- [GNNExplainer: Generating Explanations for Graph Neural Networks](https://mlsanity.com/paper/1903.03894/) (2019)
- [GraphCodeBERT: Pre-training Code Representations with Data Flow](https://mlsanity.com/paper/2009.08366/) (2020)
- [Learning the Difference that Makes a Difference with Counterfactually-Augmented Data](https://mlsanity.com/paper/1909.12434/) (2019)
- [Devign: Effective Vulnerability Identification by Learning Comprehensive Program Semantics via Graph Neural Networks](https://mlsanity.com/paper/1909.03496/) (2019)
- [PGM-Explainer: Probabilistic Graphical Model Explanations for Graph Neural Networks](https://mlsanity.com/paper/2010.05788/) (2020)
- [Simple data balancing achieves competitive worst-group-accuracy](https://mlsanity.com/paper/2110.14503/) (2021)
- [Vulnerability Detection with Code Language Models: How Far Are We?](https://mlsanity.com/paper/2403.18624/) (2024)
- [Explaining NLP Models via Minimal Contrastive Editing (MiCE)](https://mlsanity.com/paper/2012.13985/) (2020)
- [DiverseVul: A New Vulnerable Source Code Dataset for Deep Learning Based Vulnerability Detection](https://mlsanity.com/paper/2304.00409/) (2023)
- [Graph Rationalization with Environment-based Augmentations](https://mlsanity.com/paper/2206.02886/) (2022)
- [Vulnerability Detection with Fine-grained Interpretations](https://mlsanity.com/paper/2106.10478/) (2021)
- [Data Quality for Software Vulnerability Datasets](https://mlsanity.com/paper/2301.05456/) (2023)
- [Robust Counterfactual Explanations for Tree-Based Ensembles](https://mlsanity.com/paper/2207.02739/) (2022)
- [Robust Counterfactual Explanations for Neural Networks With Probabilistic Guarantees](https://mlsanity.com/paper/2305.11997/) (2023)
- [Navigating Shortcuts, Spurious Correlations, and Confounders: From Origins via Detection to Mitigation](https://mlsanity.com/paper/2412.05152/) (2024)
- [Security Vulnerability Detection Using Deep Learning Natural Language Processing](https://mlsanity.com/paper/2105.02388/) (2021)
- [Graph Neural Networks for Vulnerability Detection: A Counterfactual Explanation](https://mlsanity.com/paper/2404.15687/) (2024)
- [Solving the Class Imbalance Problem Using a Counterfactual Method for Data Augmentation](https://mlsanity.com/paper/2111.03516/) (2021)
- [Towards Formalizing Spuriousness of Biased Datasets Using Partial Information Decomposition](https://mlsanity.com/paper/2407.00482/) (2024)
- [Enhancing Code Vulnerability Detection via Vulnerability-Preserving Data Augmentation](https://mlsanity.com/paper/2404.09599/) (2024)
- [Data Quality Issues in Vulnerability Detection Datasets](https://mlsanity.com/paper/2410.06030/) (2024)

## Built on by (most cited first, 3 of 3)

- [MulVul: Retrieval-augmented Multi-Agent Code Vulnerability Detection via Cross-Model Prompt Evolution](https://mlsanity.com/paper/2601.18847/) (2026)
- [Few-Shot Knowledge Distillation of LLMs With Counterfactual Explanations](https://mlsanity.com/paper/2510.21631/) (2025)
- [COFFEQA: Counterfactual-Guided Few-Shot Distillation for LLM Question Answering](https://mlsanity.com/paper/2610.12837/) (2026)

## Abstract (arXiv)

Automated detection of vulnerabilities in source code is an essential cybersecurity challenge, underpinning trust in digital systems and services. Graph Neural Networks (GNNs) have emerged as a promising approach as they can learn structural and logical code relationships in a data-driven manner. However, their performance is severely constrained by training data imbalances and label noise. GNNs often learn 'spurious' correlations from superficial code similarities, producing detectors that fail to generalize well to unseen real-world data. In this work, we propose a unified framework for robust and interpretable vulnerability detection, called VISION, to mitigate spurious correlations by systematically augmenting a counterfactual training dataset. Counterfactuals are samples with minimal semantic modifications but opposite labels. Our framework includes: (i) generating counterfactuals by prompting a Large Language Model (LLM); (ii) targeted GNN training on paired code examples with opposite labels; and (iii) graph-based interpretability to identify the crucial code statements relevant for vulnerability predictions while ignoring spurious ones. We find that VISION reduces spurious learning and enables more robust, generalizable detection, improving overall accuracy (from 51.8% to 97.8%), pairwise contrast accuracy (from 4.5% to 95.8%), and worst-group accuracy (from 0.7% to 85.5%) on the Common Weakness Enumeration (CWE)-20 vulnerability. We further demonstrate gains using proposed metrics: intra-class attribution variance, inter-class attribution distance, and node score dependency. We also release CWE-20-CFA, a benchmark of 27,556 functions (real and counterfactual) from the high-impact CWE-20 category. Finally, VISION advances transparent and trustworthy AI-based cybersecurity systems through interactive visualization for human-in-the-loop analysis.