# MARGIN: Margin-Aware Regularized Geometry for Imbalanced Vulnerability Detection

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 38 references that are in the map.
- First posted on arXiv 11 May 2026.

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

- [CodeBERT: A Pre-Trained Model for Programming and Natural Languages](https://mlsanity.com/paper/2002.08155/) (2020)
- [Class-Balanced Loss Based on Effective Number of Samples](https://mlsanity.com/paper/1901.05555/) (2019)
- [CosFace: Large Margin Cosine Loss for Deep Face Recognition](https://mlsanity.com/paper/1801.09414/) (2018)
- [CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation](https://mlsanity.com/paper/2109.00859/) (2021)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [GraphCodeBERT: Pre-training Code Representations with Data Flow](https://mlsanity.com/paper/2009.08366/) (2020)
- [UniXcoder: Unified Cross-Modal Pre-training for Code Representation](https://mlsanity.com/paper/2203.03850/) (2022)
- [Devign: Effective Vulnerability Identification by Learning Comprehensive Program Semantics via Graph Neural Networks](https://mlsanity.com/paper/1909.03496/) (2019)
- [code2vec: Learning Distributed Representations of Code](https://mlsanity.com/paper/1803.09473/) (2018)
- [Deep Learning based Vulnerability Detection: Are We There Yet?](https://mlsanity.com/paper/2009.07235/) (2020)
- [L2-constrained Softmax Loss for Discriminative Face Verification](https://mlsanity.com/paper/1703.09507/) (2017)
- [SySeVR: A Framework for Using Deep Learning to Detect Software Vulnerabilities](https://mlsanity.com/paper/1807.06756/) (2018)
- [Hyperspherical Prototype Networks](https://mlsanity.com/paper/1901.10514/) (2019)
- [Deep Neural Collapse Is Provably Optimal for the Deep Unconstrained Features Model](https://mlsanity.com/paper/2305.13165/) (2023)
- [Loss Functions and Metrics in Deep Learning](https://mlsanity.com/paper/2307.02694/) (2023)
- [ReposVul: A Repository-Level High-Quality Vulnerability Dataset](https://mlsanity.com/paper/2401.13169/) (2024)
- [MegaVul: A C/C++ Vulnerability Dataset with Comprehensive Code Representation](https://mlsanity.com/paper/2406.12415/) (2024)
- [Guiding Neural Collapse: Optimising Towards the Nearest Simplex Equiangular Tight Frame](https://mlsanity.com/paper/2411.01248/) (2024)
- [KappaFace: Adaptive Additive Angular Margin Loss for Deep Face Recognition](https://mlsanity.com/paper/2201.07394/) (2022)
- [Data Quality Issues in Vulnerability Detection Datasets](https://mlsanity.com/paper/2410.06030/) (2024)
- [One-for-All Does Not Work! Enhancing Vulnerability Detection by Mixture-of-Experts (MoE)](https://mlsanity.com/paper/2501.16454/) (2025)
- [LIVABLE: Exploring Long-Tailed Classification of Software Vulnerability Types](https://mlsanity.com/paper/2306.06935/) (2023)
- [An Empirical Study of the Imbalance Issue in Software Vulnerability Detection](https://mlsanity.com/paper/2602.12038/) (2026)

## Abstract (arXiv)

Software vulnerability detection is critical for ensuring software security and reliability. Despite recent advances in deep learning, real-world vulnerability datasets suffer from two severe challenges: frequency imbalance and difficulty imbalance. We reinterpret these challenges from an embedding geometry perspective, observing that such imbalances induce geometric distortions in hyperspherical representation space. To address this issue, we propose MARGIN, a metric-based framework that learns discriminative vulnerability representations through adaptive margin metric learning and hyperspherical prototype modeling. MARGIN dynamically adjusts geometric regularization according to the distribution structure estimated by the von Mises-Fisher concentration, aligning the probability mass of embedding distributions with their corresponding Voronoi cells, thereby reducing geometric distortion and yielding more stable decision boundaries. Extensive experiments on public vulnerability datasets show that MARGIN consistently outperforms strong baselines, achieving notable improvements in classification and detection, especially on challenging, imbalanced datasets. Further analysis demonstrates that MARGIN produces more structured embedding geometries, improving robustness, interpretability, and generalization.