Technical Summary: MARGIN – Margin-Aware Regularized Geometry for Imbalanced Vulnerability Detection

1. Problem Statement

The paper addresses the critical challenge of software vulnerability detection in real-world scenarios, where deep learning models face two distinct but compounding forms of data imbalance:

• Frequency Imbalance: Vulnerability datasets exhibit a severe long-tailed distribution. A small number of Common Weakness Enumeration (CWE) categories (e.g., CWE-79, CWE-89) dominate the data, while many others (e.g., CWE-400, CWE-98) are rare. Standard loss functions treat all samples equally, causing models to bias toward majority classes and neglect minority vulnerabilities.
• Difficulty Imbalance: Beyond frequency, different CWE categories possess varying levels of intrinsic learning difficulty. Vulnerabilities with high structural diversity (e.g., CWE-119, which includes stack/heap overflows and off-by-one errors) exhibit high intra-class variance in feature representations, making them harder to learn than those with repetitive patterns (e.g., CWE-415).

The authors argue that existing methods primarily address frequency imbalance (via resampling or re-weighting) but fail to account for difficulty imbalance. They reinterpret these issues from an embedding geometry perspective, observing that both imbalances induce geometric distortions in the hyperspherical representation space. Specifically, minority or difficult classes fail to converge tightly around their prototypes, causing their probability mass regions to intrude into the Voronoi cells of other classes. This "distortion region" leads to misclassification and unstable decision boundaries.

2. Methodology: MARGIN

The proposed framework, MARGIN (Margin-Aware Regularized Geometry for Imbalanced Vulnerability DetectioN), is a metric-based approach that operates on a unit hypersphere to learn discriminative vulnerability representations. It integrates adaptive margin learning with hyperspherical prototype modeling.

Core Components:

1. Hyperspherical Embedding Space:
   The method utilizes a Cosine Softmax Loss where features and classifier weights are $\ell_2$-normalized, placing them on a unit hypersphere $S^{d-1}$. Classes are represented by prototypes, and the goal is to align sample embeddings with their corresponding prototypes while maintaining angular separation.

2. Geometric Modeling with vMF Distribution:
   To quantify the concentration of class distributions, the authors model embeddings using the von Mises–Fisher (vMF) distribution. The concentration parameter $\kappa_i$ estimates how tightly samples cluster around a class mean.

   • Voronoi Cells: Under ideal balanced conditions (Neural Collapse), classes form an Equiangular Tight Frame (ETF) structure, partitioning the sphere into equal Voronoi cells.
   • Distortion Region: In imbalanced settings, the probability mass of a class (defined by its vMF apex angle $\theta^{vMF}_i$) often exceeds its corresponding Voronoi cell ($\theta^{cell}$). The area where this mass intrudes into other classes' decision regions constitutes the "distortion region," directly causing False Positives and False Negatives.

3. Adaptive Margin Mechanism:
   To reduce the distortion region, MARGIN introduces a dynamic angular margin $m_i$ for each class.

   • The margin is calculated based on the difference between the class's vMF apex angle and the ideal Voronoi apex angle: $m_i = \max(\frac{\theta^{vMF}_i - \theta^{cell}}{2}, \dots)$.
   • This mechanism dynamically adjusts the constraint: classes with high dispersion (high difficulty or low frequency) receive larger margins to force their probability mass back within their Voronoi cells, while compact classes receive smaller margins.

4. Concentration-Aware Logits Scaling:
   The framework includes a scaling factor $s_i$ for the logits, adapted to the class concentration $\kappa_i$.

   • Classes with low concentration (high variance, difficult to learn) are assigned larger scaling factors to sharpen their decision boundaries and amplify gradient signals.
   • Classes with high concentration are assigned smaller factors to prevent them from dominating the training process.
   • This is achieved via a soft ranking strategy in the log-concentration space.

5. Inference with Geometric Medians:
   During inference, MARGIN constructs data-driven prototypes by computing the geometric median of the normalized training embeddings for each class, rather than relying solely on learned weight vectors. This provides more robust prototypes that better reflect the empirical feature distribution.

3. Key Contributions

The paper outlines three primary contributions:

1. Geometric Reinterpretation of Imbalance: The authors revisit frequency and difficulty imbalances, analyzing their effects through the lens of embedding geometry. They demonstrate that these imbalances manifest as a misalignment between class-conditional probability mass regions and their ideal Voronoi decision regions.
2. The MARGIN Framework: They propose a unified metric-learning framework that simultaneously mitigates both imbalances. By integrating adaptive margins (to align probability mass with Voronoi cells) and concentration-aware scaling (to balance gradient contributions), MARGIN stabilizes decision boundaries without requiring complex resampling.
3. Empirical Validation: Extensive experiments on public datasets (BigVul, MegaVul, ReposVul) show that MARGIN consistently outperforms strong baselines, including graph-based methods (ReVeal, LIVABLE) and recent representation learning models (HCL-VC, CLeVeR), particularly in fine-grained CWE classification and on rare classes.

4. Experimental Results

The evaluation was conducted on three datasets with varying scales and long-tailed distributions. Key findings include:

• Performance Gains: MARGIN achieved state-of-the-art results across all datasets. On BigVul, it reached a binary F1 of 86.57% and a macro-F1 of 70.72% for CWE classification, outperforming the next best baseline (CLeVeR) by over 7 percentage points in macro-F1.
• Robustness on Rare Classes: In the MegaVul dataset, which presents severe imbalance, MARGIN achieved a binary F1 of 51.06% and a macro-F1 of 34.65%, significantly outperforming baselines that struggled with rare CWE types.
• Ablation Studies: Removing the adaptive margin or the adaptive logits scaling resulted in performance drops, confirming that both mechanisms are essential. The margin specifically improved the balance between precision and recall, while scaling stabilized learning dynamics for minority classes.
• Geometric Analysis: Case studies demonstrated that MARGIN successfully reduces the "distortion region" (measured via FPR/FNR dynamics) and drives the embedding space toward a stable, well-separated ETF structure, as evidenced by improved clustering metrics (NMI, ARI) and a reduced Gram matrix condition number.

5. Significance and Claims

The paper claims that MARGIN offers a principled geometric interpretation of imbalance in vulnerability detection. By explicitly modeling the embedding space and regularizing the geometry to align probability mass with decision regions, the method improves not only classification accuracy but also the robustness and interpretability of the model.

The authors emphasize that their approach addresses the root cause of bias in imbalanced settings—the geometric distortion of the embedding space—rather than just treating the symptom via data re-weighting. They note that while the current evaluation is limited to C/C++ datasets due to data availability, the geometric framework is backbone-agnostic and theoretically applicable to other languages. The work bridges the gap between practical imbalance handling and theoretical manifold learning, providing a more stable foundation for security-critical vulnerability detection systems.