MARGIN: Margin-Aware Regularized Geometry for Imbalanced Vulnerability Detection

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

MARGIN: Margin-Aware Regularized Geometry for Imbalanced Vulnerability Detection

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

cryptography cybersecurity privacy security

cryptography, security, privacy, cybersecurity

ARXIV PREPRINT

MARGIN: Margin-Aware Regularized Geometry for Imbalanced Vulnerability Detection

arXiv:2605.10240v1 [cs.SE] 11 May 2026

Yuteng Zhang, Huifang Ma, Member, IEEE, Jiahui Wei, Qingqing Li, and Yafei Yang

Abstract—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 (Margin-Aware Regularized Geometry for Imbalanced Vulnerability DetectioN), a metricbased 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. Index Terms—Software Vulnerability Detection, Software Security, Imbalance Learning, Metric Learning.

I. I NTRODUCTION S modern software systems keep growing in size and complexity, security vulnerabilities hidden in source code have become a serious threat to software reliability and system safety. These vulnerabilities can cause unexpected program failures or be exploited by attackers to break into systems. Therefore, automatically detecting and classifying vulnerabilities from source code has become an important problem in software engineering. Recent advances in deep learning have significantly improved software vulnerability detection and classification. However, despite these advances, these models still face substantial challenges when applied to real-world vulnerability datasets. Deep learning-based vulnerability analysis aims to learn discriminative class representations by clustering samples of the same category while separating different categories. However, real-world vulnerability data are often highly imbalanced in both class frequency and sample difficulty, which can bias model optimization, hinder reliable convergence, and degrade generalization performance.

Yuteng Zhang, Huifang Ma, Jiahui Wei, Yafei Yang, and Qingqing Li are with the College of Artificial Intelligence and Computer Science, Northwest Normal University, Lanzhou 730070, China. Corresponding author: Huifang Ma (email: [email protected]). Manuscript created May, 2026;

Frequency Imbalance. A key challenge in deep learning for vulnerability detection is the severe frequency imbalance in real-world software data. This issue appears not only as the imbalance between vulnerable and non-vulnerable code, but more importantly as the uneven distribution across different Common Weakness Enumeration (CWE) categories. As shown in Fig. 1, which presents the frequencies of the top 25 CWE types from recent CVE records (2023–2025), the data follows a clear long-tailed distribution with large differences in sample counts. A small number of common categories, such as CWE-79 (Cross-Site Scripting) and CWE89 (SQL Injection), contain tens of thousands of samples and make up a large portion of the dataset. In contrast, many other vulnerabilities, such as CWE-400 (Uncontrolled Resource Consumption), CWE-98 (Use of an Improperly Restricted Directory), and CWE-863 (Incorrect Authorization),appear very rarely and only account for a small part of the data. This large gap in frequency leads to biased metric learning. Standard loss functions like Cross-Entropy and contrastive Loss treat all samples equally [1], so classes with more samples contribute much more to the gradients and dominate the parameter updates. Meanwhile, minority classes provide much weaker signals and have limited influence on the learned representations. In the learned embedding space, majorityclass samples typically cluster tightly, while minority-class samples remain scattered and fail to reliably converge around their prototypes. As illustrated in Fig. 3 (a), even when class difficulty is similar, the embeddings of minority class a shows a more dispersed and blur distribution than class b. Because of this, the model tends to learn biased patterns from high-frequency classes, shifting the decision boundary toward dominant categories while overlooking the subtle characteristics of rare vulnerabilities [2], [3]. Such bias severely limits the reliability of vulnerability detection systems in realworld security scenarios, where accurately identifying rare yet critical vulnerabilities is particularly important. Difficulty imbalance. Most existing methods for handling imbalanced vulnerability detection primarily focus on frequency imbalance. Common solutions include resampling, cost-sensitive learning, focal loss, and class-balanced reweighting [1], [4]. While these approaches mitigate the influence of majority classes, they usually assume that all classes are equally difficult to learn. Difficulty imbalance can significantly hinder model performance even when class frequencies are balanced, leading to slower convergence, lower precision for difficult classes, and biased evaluation metrics in real-world vulnerability datasets. In practice, vulnerability datasets also exhibit clear differ-

ARXIV PREPRINT

7RS&:(LQ&9(5HFRUGV

(

&:(&DWHJRU\

(

&RXQW

Fig. 1. Top 25 CWE Types in CVE Project (2023–2025). Frequency distribution of CWE categories in publicly disclosed vulnerabilities from the CVE Project. The distribution is highly imbalanced, where a small number of CWE types account for a large portion of vulnerabilities while many others appear only rarely. This long-tailed pattern reflects the uneven real-world occurrence of software weaknesses and highlights the inherent class imbalance in vulnerability datasets.

Fig. 2. Illustration of vulnerability diversity and learning difficulty across two CWE types. (a) CWE-119 exhibits high embeddings variance with diverse patterns, including stack overflows, heap overflows, and off-by-one errors. (b) In contrast, CWE-415 follows a highly templated allocate→free→free pattern with limited variation.

Class

Frequency: Difficulty:

Embedding Dispersion:

(a)

(b)

Fig. 3. (a) With comparable class difficulty, lower-frequency classes yield embeddings that are harder to converge around their prototypes on the manifold. (b) With equal class frequencies, higher difficulty (i.e., greater intraclass variance) results in embeddings that are harder to converge around their prototypes on the manifold.

ences in learning difficulty across CWE categories [5]. Although CWE defines a standardized taxonomy, vulnerabilities within the same CWE can vary significantly in syntax, control flow, data flow, and usage context [6], leading to diverse implementation patterns and semantic meanings. As a result, the features extracted by the backbone network from the same class can be highly heterogeneous, especially for difficult classes, where more diverse patterns generate more varied feature vectors, increasing intra-class variance and reflecting higher learning difficulty. For example, as shown in Fig. 2, CWE-119 (Out-of-bounds

Read) can appear as stack overflows, heap overflows, or offby-one loop errors, exhibiting rich structural diversity. In contrast, CWE-415 (Double Free) often follows simpler and repetitive patterns, where the same pointer is freed multiple times. This suggests that class difficulty is closely related to the diversity of feature representations: classes with higher structural diversity produce more high-variance embeddings from the backbone, complicating optimization and slowing convergence. In this work, it’s hard to quantify the accurate difficulty directly so we quantify difficulty imbalance by treating intraclass feature variance as a proxy for class difficulty. Prior studies have confirmed that higher intra-class variance correlates with increased learning difficulty [7], reflecting greater dispersion, noise, and inconsistency within class representations—a manifestation of the well-known Bias–Variance Dilemma. We attribute this difficulty imbalance to the feature geometry learned by the backbone. Under this setting, difficult classes show high feature variance in the embedding space [3], [8]. Easy classes form compact clusters due to simpler and more consistent patterns, whereas difficult classes remain more dispersed. Their semantic heterogeneity causes the backbone to produce overly dispersed features, destabilizing the decision boundary and hindering convergence—an issue that traditional imbalance learning strategy cannot resolve. As illustrated in Fig. 3(b), even when class frequencies are identical, the difficult class a exhibits a more scattered distribution than class b. When frequency imbalance and difficulty imbalance coexist, the learned representations show

ARXIV PREPRINT

varying degrees of compactness across classes [3], both of which can be interpreted geometrically as the concentration of samples around their class prototypes in the embedding space. Lack of Interpretability. Most existing approaches mainly aim to improve classification performance, while paying less attention to the interpretability. Common evaluation metrics such as MCC and F1-score only reflect prediction results, but cannot show whether the embedding space has clear and welldiscriminated geometric structures under severe imbalance [2]. Without explicit constraints on representation geometry, models may depend on locally optimal but unstable decision boundaries, which can be highly sensitive to distribution shifts or rare vulnerability patterns. For security-critical tasks like vulnerability detection, such instability can greatly reduce practical reliability. In essence, both frequency imbalance and difficulty imbalance can be unifiedly viewed as a shared geometric issue in the embedding space, namely the uneven concentration of class distributions. Our work shows that imbalance in vulnerability detection is closely related to distortion in the learned embedding geometry [3], [4]. Frequency and difficulty imbalance cause dominant classes to have much dispersed and distort distribution in the learned decision space. Therefore, explicitly modeling the embedding space from a geometric and manifold perspection in metric learning offers a more principled way to improve both robustness and interpretability. To address these challenges, we propose MARGIN (Margin-Aware Regularized Geometry for Imbalanced Vulnerability DetectioN), a metric-based framework that revisits vulnerability detection from a hyperspherical embedding perspective. It provides a unified geometric view in which both frequency and difficulty imbalance are described by how concentrated each class is on the hypersphere. By explicitly shaping the geometry of the embedding space, MARGIN reduces representation bias through adaptive margin learning and prototype-based regularization. This helps stabilize decision boundaries and improves the separation between classes under severe imbalance. Overall, the framework connects practical imbalance handling with a clear geometric interpretation, leading to better performance and improved interpretability. The contributions of this paper are summarized as follows: 1) We revisit two critical imbalances in vulnerability detection, i.e., frequency imbalance and difficulty imbalance, and analyze their effects in an interpretable manner from the perspective of embedding geometry. (Section III) 2) We propose MARGIN, a metric-learning framework that integrates adaptive mechanisms and hyperspherical prototype modeling to mitigate both imbalances simultaneously. (Section IV) 3) Extensive experiments on multiple public vulnerability datasets demonstrate that MARGIN consistently improves vulnerability detection and classification performance. (Section V, VII)

II. P RELIMINARIES A. Problem Formulation We consider a vulnerability dataset D = {(xi , yi )}N i=1 , where each sample consists of a code snippet xi ∈ X and a class label yi ∈ Y. The label space is defined as Y = {Non-Vul, CWE-a, CWE-b, . . . }, where Non-Vul denotes non-vulnerable code and each CWE-* represents a specific vulnerability category defined in the CWE. This task is formulated as a single-label multi-class classification problem in close-set, where each code sample is associated with exactly one label. In practice, real-world datasets exhibit severe frequency and difficulty imbalance across different labels.

B. Cosine Softmax Loss To address class imbalance in vulnerability detection from a geometric perspective, we adopt a unit hyperspherical prototype-based classification framework. This formulation defines a structured manifold that can be explicitly analyzed and regularized, allowing imbalance to be tackled in a unified geometric manner. Each class is represented by a prototype, serving as a stable and interpretable semantic center in the embedding space [9]. We begin with the standard Cosine Softmax Loss for singlelabel multi-class classification. For each input, the encoder produces a d-dimensional feature, and we apply ℓ2 normalization to both the feature e ∈ Rd and classifier weight wi ∈ Rd . The linear classifier then produces logits for each class, followed by the standard softmax cross-entropy loss:

zi = w̃i⊤ ẽ,

exp(zi ) , ℓi = − log P|Y| j=1 exp(zj )

i∈Y

(1)

The logit zi can be also decomposed as: zi = cos θi

(2)

where θi denotes the angle between w̃i and ẽ. Thus, the decision score depends purely on angular similarity and vector magnitude variations are eliminated. We introduce the scaling factor s0 to sharpen the logits and the resulting cosine softmax loss is defined as: exp(s0 cos θi ) ℓi = − log P|Y| j=1 exp(s0 cos θj )

(3)

The cosine softmax loss makes both features and class weights lie on the unit hypersphere Sd−1 . We interpret w̃i as the Weight Prototype of class i. Training therefore encourages samples to align with their corresponding prototypes, forming compact clusters separated by angular margins. This geometric view enables explicit analysis of imbalance effects in embedding space.

ARXIV PREPRINT

C. Voronoi Cell and Neural Collapse

Recent studies show that deep classifiers without prediction bias and trained on ideal balanced data tend to exhibit the Equiangular Tight Frame (ETF) structure under Neural Collapse phenomenon during the later stage of training [10], all weight prototypes converge to a highly symmetric geometric configuration: the pairwise similarity between weight prototypes converge to a constant, resulting in prototypes that are uniformly distributed on the unit hypersphere: w̃i⊤ w̃j = −

1 , |Y| − 1

∀i ̸= j

(4)

In this setting, the unit hypersphere Sd−1 can be partitioned into |Y| spherical regions centered around each weight prototype, where all regions are of equal size. These regions correspond to the Voronoi Cells. When the number of classes |Y| is sufficiently large, the shape of each Voronoi cell approaches a spherical cap on the hypersphere. Taking each Voronoi cell as the base and the center of the unit hypersphere as the apex, we obtain a cone whose base is the corresponding Voronoi cell. We have the apex angle of hyperspherical cone which base is Voronoi cell, referred as Voronoi Apex Angle: 1 (5) θcell = arccos − |Y| − 1 We also referrer the hyperspherical cone as the Voronoi Cone. At the same time, the probability mass regions of embeddings for each class converge to be equal in size and lie predominantly within the corresponding Voronoi cell defined by its weight prototype. Also, the mean feature prototype of each class aligns with its corresponding weight prototype. Let µ̃i denote the normalized mean feature of class i. Under Neural Collapse, we have µ̃i = w̃i , for all i ∈ 1, . . . , |Y|. Ideally, the ETF structure yields clearly discriminative and unbiased decision boundaries [3]. However, under classimbalanced conditions, such an ETF structure cannot be properly established, resulting in geometrically less separable representations, which in turn gives rise to the distortion region. This observation suggests that, to alleviate the prediction bias induced by class imbalance from a geometric perspective, it is essential to explicitly encourage the embedding distribution to evolve toward an ETF structure. D. vMF Distribution for Hyperspherical Modeling Following the principle of maximum entropy, feature embeddings are represented on the unit hypersphere. The von Mises–Fisher (vMF) distribution is employed as a quantitative tool to characterize class-wise feature concentration. In this framework, the concentration parameter κi measures how tightly samples cluster around their mean direction µi [11]. Formally, a normalized embedding ẽ ∈ Sd−1 follows a vMF distribution ẽ ∼ vMF(µ̃i , κi ) if p(ẽ | µ̃i , κi ) = Cd (κi ) exp κi µ̃⊤ ẽ (6) i where

Cd (κi ) =

κi2

−1

(2π) 2 I d −1 (κi )

(7)

and Iν (·) denotes the modified Bessel function of the first kind. A larger κi indicates stronger concentration of embeddings around µ̃i , leading to more compact class distributions. Under the vMF assumption, the high-density region of each class on the hypersphere can be approximated as a spherical cap centered at µ̃i . Taking this region as the base and the center of the unit hypersphere as the apex, we obtain a hyperspherical cone, refer as the vMF Confidence Cone. For a given confidence level α, the corresponding highprobability region satisfies vMF =α (8) P arccos(µ̃⊤ ẽ) ≤ θ i i where θivMF denotes the apex angle of the hyperspherical cone induced by the vMF distribution. This angle is also referred to as the vMF Apex Angle. Importantly, for a fixed confidence level α, the apex angle θivMF is solely determined by the concentration parameter κi . A larger κi leads to a smaller apex angle, indicating tighter concentration around the prototype, whereas a smaller κi produces a wider apex angle, reflecting more dispersed intra-class distributions. III. M OTIVATION A NALYSIS Under the dual imbalance of class frequency and difficulty, neural networks in the late training stage—upon reaching the Neural Collapse regime—learn hyperspherical embeddings where class prototypes are well separated with large angular distances. As illustrated in Fig. 4, let Pa and Pb denote the prototypes of classes a and b, respectively. According to the argmax decision rule, the angular bisectors of ∠Pa OPb and ∠Pb OPa , denoted as OQ1 and OQ2 , partition the unit hypersphere into regions with equal angular boundaries. These regions define the Voronoi cells of classes a \ and b, namely Q\ 1 Pa Q2 and Q1 Pb Q2 , each with a fixed apex cell angle θ . The Voronoi cells therefore act as ideal decision regions. However, under imbalance, class-conditional embeddings do not always concentrate tightly around their corresponding prototypes. Instead, they may deviate and extend beyond their associated Voronoi cells, leading to misclassification [3]. To characterize this phenomenon, consider the binary classification scenario shown in Fig. 4. The embeddings of classes a and b can be reasonably approximated as following vMF distributions on S1 . Their probability mass regions can be viewed as \ spherical caps, namely A\ 1 OA2 and B1 OB2 , with vMF apex vMF vMF angles θa and θb , where typically θavMF > θbvMF . This indicates that class a exhibits a more dispersed embedding distribution. Crucially, these probability mass regions may exceed their corresponding Voronoi cells, i.e.,θavMF > θcell , θbvMF > θcell . For \ class a, the exceeding portions Q\ 1 OA1 and Q2 OA2 intrude \ into the Voronoi cell Q1 Pb Q2 of class b, leading to false

ARXIV PREPRINT

Class

Probablity Mass Region

Class

Probablity Mass Region

Geometric Interpretation of Misclassification

Distortion Region

All parts of probability mass region exceed its corresponding Voronoi Cell is the distortion region.

Weight Prototype Argmax Boundary Probablity Mass Boundary of Class : Probablity Mass Boundary of Class :

Class samples in , be mistakenly predicted as class .

will

Class samples in , be mistakenly predicted as class .

will

Shrink inward the boundary of each class's probability mass region toward its prototype by an angular margin of , shifting it closer to the argmax decision boundary, to reduce the distortion region. This encourages the probability mass regions to align within their corresponding Voronoi cell.

Voronoi Cell for Class : Voronoi Cell for Class : Probability Mass Region for Class : Probability Mass Region for Class :

b. After adjusting probability mass region boundary

a. Cosine softmax spherical decision boundary with distortion region

Fig. 4. Geometric Interpretation of Misclassification and Heuristic Design of MARGIN Conceptual illustration of spherical angular spans for class decision of Cosine Softmax Loss and MARGIN on S1 .

negatives for class a and false positives for class b. Similar phenomenon applies to class b. We define the union of all such probability masses that lie outside its corresponding Voronoi cell and enter those of other classes as the distortion region, specifically A\ 1 OB1 and A\ OB . This distortion region reflects the geometric 2 2 mismatch between class-conditional embedding distributions and Voronoi decision regions, and directly leads to misclassification and prediction bias under dual imbalance that causes increment of FPR and FNR. Therefore, the core challenge is not merely the class imbalance itself, but the misalignment between the geometry of embedding distributions and the induced decision regions. An effective learning strategy should explicitly constrain the embedding distribution to reduce the distortion region. To this end, we introduce a margin constraint into the cosine softmax loss, enforcing that each sample of class i remains within an angular distance mi from its corresponding prototype. This imposes an upper bound on the angular support of the class-conditional distribution, effectively contracting its probability mass region (i.e., reducing θivMF ). As a result, intraclass embeddings become more compact and better aligned with their Voronoi cells, thereby mitigating the distortion region. However, a fixed margin contradicts the heterogeneous dispersion induced by dual imbalance (e.g., θavMF > θbvMF ). Classes with highly dispersed probability mass region require stronger constraints to sharpen decision boundary, while more concentrated classes benefit from loose margins to maintain stable optimization. This reveals a fundamental limitation of static margin designs [12]. To address this, we propose a dynamic margin mechanism that adapts mi according to the geometric structure of each class during training. By explicitly modeling class-conditional distributions on the hypersphere and adjusting margins based on their concentration, our approach aims to align the probability mass region (e.g., A\ 1 OA2 ) with its corresponding Voronoi cell (e.g., Q\ P Q ), thereby minimizing the distortion region. 1 a 2 IV. M ETHODOLOGY A. Feature Extraction and Normalization Given a code sample X , we employ CodeT5 as the transformer backbone to extract a 768-dimensional semantic vector

ei ∈ R768 from the final [CLS] token. While MARGIN is backbone-agnostic, we adopt CodeT5 due to its widespread adoption in code representation tasks. This choice ensures that observed performance gains are attributable to MARGIN rather than the backbone, as confirmed by our ablation studies; preliminary experiments with alternative encoders, such as CodeBERT and GraphCodeBERT, yield consistent improvements, further validating the method’s backbone-independent nature. To enforce a geometrically interpretable space, we apply ℓ2 -normalization to both feature embeddings and classifier prototypes. B. Approximation of Kappas To quantify each class’s concentration on the hypersphere, we estimate the vMF parameter κi via Maximum Likelihood Estimation (MLE). Given the normalized embeddings of class i, we first compute the mean resultant length from the empirical mean vector. The MLE of κi involves a ratio of modified Bessel functions, which has no closed-form inverse and usually requires iterative optimization. To simplify, we use the high-dimensional approximation of Banerjee et al. [11]. κi ≈

||µi ||(d − µ2i ) 1 − µ2i

(9)

Here, d is the embedding dimensionality. The approximation is accurate in high dimensions; with CodeT5 producing 768dimensional embeddings, it is both empirically justified and widely used in prior work. C. Align vMF Apex Angle with Voronoi Apex Angle To characterize the angular uncertainty of each class on the hypersphere, we estimate the vMF Apex Angle for each class based on its underlying vMF distribution. On the unit hypersphere, the confidence level α under a vMF distribution for class i satisfies R θivMF (sin θ)d−2 eκ cos θ dθ α = R0 π (sin θ)d−2 eκ cos θ dθ 0

(10)

In practice, we set α = 95%, and the corresponding vMF apex angle is given by

ARXIV PREPRINT

θivMF = arccos Fi−1 (1 − α)

(11)

where Fi−1 denotes the inverse marginal cumulative distribution function of the vMF distribution for class i. Since no closed-form solution exists, we adopt an efficient approximation. Given the concentration parameter κi of class i, we define the effective concentration as κi . Based on the effective concentration, the apex angle of the vMF confidence cone is approximated as s χ2α (d − 1) (12) θivMF = κi where χ2α (d − 1) denotes the α-quantile of the chi-square distribution with d − 1 degrees of freedom. This follows from the asymptotic property of vMF distributions that, under high concentration, local angular deviations can be approximated by a Gaussian distribution on the tangent space [8], [13]. To constrain each class probability mass region within its Voronoi cell, we align the vMF apex angle with the Voronoi apex angle via an adaptive angular margin. Considering the symmetry of the hyperspherical cone angle, we divide it by two to define the margin with Non-negativity constraint: vMF θ − θcell ,0 (13) mi = max i 2 However, in the later stages of training, even when θivMF is aligned with θcell , the Voronoi cell is not a perfect hyperspherical circle that covers the vMF probability mass region, in order to maintain strong regularization and further enhance training, we need to continue aligning it with the smallest vMF vMF apex angle θmin . It serves to maintain compactness of class embeddings even after the vMF confidence cone has aligned with the Voronoi cone. This prevents late-stage relaxation of the embedding space and stabilizes prototypebased classification. Therefore, we have: vMF vMF θ − θcell θivMF − θmin mi = max i , (14) 2 2 When the vMF confidence cone exceeds the Voronoi boundary, the excess region indicates potential inter-class overlap. The adaptive margin therefore penalizes excessive dispersion beyond the Voronoi partition, improving hyperspherical separability and reducing class interference. D. Concentration-Aware Logits Scaling To adapt classifier confidence to class compactness, we adjust the logit scaling factors according to the convergence status of each class. The concentration parameter κi reflects the compactness of class representations on the hypersphere. A larger κi indicates that class i is well optimized with low intra-class variance; assigning large logits in this case may lead to over-confident predictions and dominate gradient updates. Therefore, we assign smaller scaling factors to such classes to suppress their influence. In contrast, a smaller κi indicates a more

dispersed distribution and insufficient convergence, for which we increase the scaling factors to sharpen logits and amplify gradient signals. To achieve this behavior while avoiding the instability of explicit inverse mappings, we adopt a soft ranking-based strategy in the log-concentration space. Specifically, we compute: ri = |Y| · softmax

log κi |Y|

(15)

where C is the number of classes, ensuring ri > 0 and unit mean. We then reverse the allocation order according to the ranking of κi , assigning larger weights to lower-concentration classes: r̃ = Reverse(r) (16) Finally, the class-wise scaling factors are: si = s0 · r̃i

(17)

where s0 is a base scaling factor. This formulation enables geometry-aware and difficultyadaptive gradient reallocation, enhancing under-converged classes while preventing well-optimized ones from dominating training. E. Objective Function The resulting difficulty and frequency-aware kappa loss with adaptive margin and scaling factor for a sample of class i is: exp si cos(θi + mi ) P ℓi = − log exp(si cos(θi + mi )) + j̸=i exp sj cos θj (18) F. Inference on Test Instead of directly using classifier weights as prototypes, on evaluation stage, we construct data-driven prototypes from training embeddings to better capture the empirical geometry of the feature space. Specifically, for each class i ∈ Y, we compute the geometric median of its normalized training embeddings and project it onto the unit hypersphere: pi = arg min p

∥p − ẽi ∥2 ,

xi ∈Di

p̃i =

pi ∥pi ∥2

(19)

Compared to weight-based prototypes, geometric medians are more robust to outliers and better reflect class-wise embedding distributions. Inference is performed via nearest-prototype matching in angular space: ŷ = arg max ẽ⊤ (20) test p̃i y∈Y

This induces a Voronoi partition on Sd−1 , where each sample is assigned to the angular region of its closest prototype.

ARXIV PREPRINT

TABLE I S TATISTICS OF VULNERABILITY DETECTION DATASETS . Dataset

Amount

Vul/Non-Vul IR

CWEs

Labels CV

BigVul MegaVul ReposVul

43184 147423 82037

18.57 18.55 31.17

10 15 13

3.13 3.78 3.61

program analysis and code slicing; and HCL-VC(2024, EMNLP) [21], MoEVD(2025, FSE) [22], and CLeVeR(2025, ACL) [23], which represent recent advances in CWE-aware representation learning and vulnerability classification. For fair comparison, we follow the hyperparameter settings recommended in the original papers. When official implementations are unavailable, models are reproduced based on publicly released code or reported configurations.

V. E XPERIMENTS All experiments were conducted on a single NVIDIA RTX5880 Ada (48 GB) with CUDA 12.4. All methods were evaluated under identical settings unless otherwise specified. In both the baseline and ablation experiments, to validate the effectiveness of our method and mitigate random fluctuations, we conducted each experiment five times with different random seeds. To comprehensively evaluate the effectiveness of MARGIN, we design the following research questions: • RQ1: How effective is the proposed method compared with baseline approaches? • RQ2: How do the proposed adaptive mechanisms contribute to performance? • RQ3: How sensitive is MARGIN to different hyperparameter settings? • RQ4: How much time does MARGIN take? A. Datasets We evaluate MARGIN on three widely used public vulnerability detection datasets: BigVul [14], MegaVul [15], and ReposVul [16]. These datasets, collected from open-source repositories, CVE/NVD reports, and vulnerability-fixing commits, vary in scale and CWE distribution, and are highly imbalanced, with non-vulnerable samples dominating, providing realistic benchmarks for imbalanced vulnerability detection. BigVul and MegaVul are large-scale C/C++ datasets, while ReposVul emphasizes higher data quality through automated vulnerability disentanglement and dependency-aware filtering. Following prior work, we perform standard preprocessing, to obtain stable and reliable data splits, we further select approximately the top 10-15 most frequent CWE categories while preserving the original long-tailed distribution [5]. This ensures sufficient samples per class for training and evaluation while preserving the original long-tailed characteristics of the data distribution. All datasets are split into training, validation, and test sets with a ratio of 8:1:1. B. Baselines To comprehensively evaluate the effectiveness of MARGIN for CWE-based single-label multi-class vulnerability classification, we compare our method against representative baselines from three categories, including graph-based methods, program analysis–based approaches, and recent representation learning models. Specifically, the compared baselines include ReVeal(2021, TSE) [17] and LIVABLE(2024, TSE) [18], which leverage graph neural networks for structural vulnerability modeling; µVulDeePecker(2021, DSC) [19] and SySeVR(2022, DSC) [20], which exploit vulnerability-aware

C. Metrics We evaluate MARGIN from two complementary perspectives: (i) Binary Vulnerability Detection and (ii) Finegrained CWE Classification. Given the severe class imbalance, we adopt imbalance-robust metrics. For binary detection, we collapse all CWE categories into a single positive class (Vulnerable) and treat Non-Vul as the negative class. We report Precision, Recall, F1-score, and Matthews Correlation Coefficient (MCC). For multi-CWE classification, we follow a one-vs-all (OvA) scheme over all classes to compute per-class Precision, Recall, F1-score, and MCC. To specifically assess vulnerability-type recognition, we exclude the Non-Vul class and compute macroaveraged metrics over all CWE classes, referred to as CWEMacro metrics. VI. D ISCUSSION A. Overall Performance Comparison (RQ1) Across BigVul, MegaVul, and ReposVul, MARGIN consistently outperforms most baselines in both vulnerability detection and CWE classification, with clear gains in F1 score (Table II). On BigVul, MARGIN achieves the best binary F1 (86.57%) and MCC (85.88%), surpassing HCL-VC. For CWE classification, it reaches a macro-F1 of 70.72% and macroprecision of 73.35%, outperforming CLeVeR by over 7 points, showing strong performance in both detection and fine-grained classification. On the more challenging MegaVul dataset, MARGIN still performs best, with a binary F1 of 51.06% and a high precision of 66.66%, indicating fewer false positives under severe class imbalance. It also achieves a macro-F1 of 34.65%, improving over LIVABLE (32.68%), showing more reliable CWE predictions. On ReposVul, which reflects real-world scenarios, MARGIN achieves a binary F1 of 43.85%, higher than µVulDeePecker (42.01%). For CWE classification, it obtains the best macro-F1 (34.39%) and macro-recall (48.60%), showing good generalization. Other sequence-based methods (e.g., SySeVR, µVulDeePecker) generally achieve high recall but low precision, resulting in more false positives. Structure-aware methods (e.g., HCL-VC, LIVABLE, CLeVeR) are more balanced but underperform on rare CWE types. MARGIN improves the recall–precision balance and boosts performance on rare classes. Statistical tests confirm consistent macro-F1 gains (p < 0.05 in most cases). For binary F1, MARGIN clearly outperforms

ARXIV PREPRINT

BIGVUL

275000 250000 5000 4000 3000 2000 1000 0

260000

255,846

240000

4,278

2,272 1,484 1,260

1,016 986

682 646 622 532

-vul 119 -20 399 125 200 264 189 416 190 362 Non CWE- CWE CWE- CWE- CWE- CWE- CWE- CWE- CWE- CWE-

1500 1250 1000 750 500 250 0

REPOSVUL

MEGAVUL 425000 406,832 400000 375000

244,740

1,256 1,0661,054

3000 690 678

2000

614 558 380 328 314 314

300 300

-vul 125 190 119 787 476 416 -20 399 362 200 415 264 189 NonCWE-CWE-CWE-CWE-CWE-CWE- CWECWE-CWE-CWE-CWE-CWE-CWE-

1000

3,000 2,6482,560 2,468 2,140 1,824 1,1841,024 886 878 830 758

678 542 512

-vul 119 125 787 476 -20 416 190 200 399 264 120 362 400 401 189 NonCWEC- WEC- WEC- WE- CWCEWEC- WEC- WEC- WEC- WEC- WEC- WEC- WEC- WEC- WE-

Fig. 5. Classes distribution across labels. In terms of frequency, non-vulnerable samples dominate the gaps, while the remaining CWE vulnerability samples exhibit a long-tailed distribution. Overall, the label distribution is highly imbalanced.

traditional methods, while differences with recent strong baselines are often insignificant, likely due to near-saturation. Overall, MARGIN mainly improves feature representation, leading to better class-level discrimination rather than just higher overall accuracy. These results confirm its effectiveness, especially for imbalanced data and real-world vulnerability analysis.

B. Contribution of Adaptive Mechanisms (RQ2) To investigate the impact of the adaptive mechanisms in MARGIN, we conducted ablation studies by selectively removing each key component. Table III summarizes the performance results. Removing the adaptive margin significantly affects the binary detection recall, which increases to 87.66% while precision drops to 82.97%, indicating a shift in the decision boundary that favors sensitivity over specificity. For vulnerability classification, the macro-MCC and macro-F1 decrease compared to the full model, highlighting that the adaptive margin contributes to balanced performance across classes by properly calibrating inter-class margins. Disabling the adaptive logits scaling slightly improves binary MCC and F1, but the recall decreases to 85.05% and precision drops to 89.60%, suggesting a less robust separation in the feature space. For classification, the macro-MCC and macro-F1 also drop compared to the full model, confirming that the scaling factor helps stabilize the learning dynamics and enhances discriminative power for minority classes. The full model consistently achieves the best performance. In binary detection, it attains a high F1 of 87.47% with balanced recall and precision. For classification, it achieves the highest macro-MCC (71.29%) and macro-F1 (71.13%), showing that combining both adaptive mechanisms improves detection and classification, especially for rare CWE types. Ablation results also confirm that removing geometric constraints causes a significant drop even with CodeT5 fixed, indicating the gains stem from the geometric constraints rather than the feature extractor. In summary, ablation results confirm that both the adaptive margin and logits scaling are crucial for MARGIN, with the margin calibrating inter-class distances and the scaling factor stabilizing training and enhancing discriminability.

C. Hyperparameter Sensitivity Analysis (RQ3) Despite incorporating adaptive margins and class-specific adjustments, MARGIN introduces only a single explicit hyperparameter, the base scaling factor s0 , which controls the magnitude of geometric modulation in the hyperspherical feature space. We evaluate its sensitivity on BigVul (Table IV). The best binary F1 (87.35%) and macro-MCC (71.22%) are achieved at s0 = 20, with stable performance observed for s0 ∈ [15, 25]. Smaller values (e.g., s0 = 10) reduce inter-class separation, lowering macro-MCC to 63.78%, while excessively large values (e.g., s0 = 60) decrease binary F1 to 85.00% due to over-confident margins. The trend is non-monotonic: increasing s0 initially improves both detection and classification, but beyond s0 ≈ 30, metrics may fluctuate. This reflects a trade-off in modulating hyperspherical decision boundaries: moderate scaling promotes well-separated and stable boundaries, while extreme values can reduce robustness. The optimal s0 varies across datasets. Empirically, noisier or intrinsically harder datasets benefit from slightly larger s0 to strengthen gradients and sharpen boundaries, whereas cleaner or easier datasets favor smaller s0 to prevent overconfident predictions. These insights offer practical guidance for tuning s0 . D. Computation Costs of MARGIN (RQ4) Due to differences in hardware, pretraining strategies, implementation, and data preprocessing, direct efficiency comparisons with prior work may be unfair. We therefore report only the efficiency of the proposed method. We adopt the approximation of kappa given in equation 9, which has constant time complexity and incurs no additional computational overhead. In addition, computing class prototypes using the Weiszfeld algorithm in fixed dimension d has linear time complexity O(n), which matches the per-iteration complexity of training. Therefore, the prototype computation does not increase the overall time complexity or delay convergence, and is considered acceptable for practical use. MARGIN’s computational cost in training and inference is largely determined by the backbone encoder. The adaptive margin and scaling factor mechanisms introduce additional constraints, which empirically improve gradient signals and accelerate convergence. Using CodeT5 as the backbone, average training latency is 167–200 ms per sample, and inference

ARXIV PREPRINT

TABLE II BASELINE COMPARISON . p- VALUES ARE COMPUTED ON F1 SCORES USING T- TEST AGAINST MARGIN. Dataset

B IG V UL

M EGAV UL

R EPOS V UL

Vulnerability Detection Performance (Vul/Non-vul Binary)

Method

Vulnerability Classification (CWE-Macro)

MCC(%)

F1(%) (p-value)

Recall(%)

Precision(%)

Macro-MCC(%)

Macro-F1(%) (p-value)

Macro-R(%)

Macro-P(%)

SySeVR µVulDeePecker ReVeal LIVABLE HCL-VC MoEVD CLeVeR

65.54±2.1 69.30±2.3 78.88±2.5 83.69±2.7 84.64±2.8 81.84±2.6 84.30±2.7

63.94±1.8 (1.6e-06) 68.58±2.2 (5.2e-06) 79.74±2.5 (3.7e-03) 84.48±2.7 (2.6e-01) 85.42±2.8 (5.3e-01) 82.74±2.6 (5.5e-02) 85.06±2.7 (4.1e-01)

95.36±2.8 93.61±3.1 87.37±2.9 88.10±2.9 85.05±2.8 85.92±2.8 88.82±2.9

48.10±1.5 54.11±1.8 73.33±2.4 81.15±2.7 85.80±2.8 79.78±2.6 81.60±2.7

44.03±1.2 41.53±1.1 25.06±0.8 61.91±1.6 66.37±1.7 34.99±0.9 64.09±1.6

42.61±1.1 (2.6e-08) 39.54±1.0 (2.5e-08) 21.76±0.7 (1.9e-08) 60.86±1.5 (1.7e-05) 66.17±1.7 (3.4e-03) 33.93±0.9 (1.9e-08) 63.76±1.6 (2.1e-04)

52.39±1.4 42.91±1.2 30.88±0.9 60.87±1.6 68.62±1.8 38.40±1.0 66.11±1.7

42.28±1.3 46.61±1.4 34.31±1.1 65.95±2.0 65.17±2.0 35.24±1.1 63.48±2.0

MARGIN

85.88±2.8

86.57±2.8 (–)

88.00±2.9

85.19±2.8

70.91±1.8

70.72±1.8 (–)

69.47±1.8

73.35±2.2

SySeVR µVulDeePecker ReVeal LIVABLE HCL-VC MoEVD CLeVeR

27.32±1.8 25.02±1.6 32.80±2.1 47.49±3.5 43.06±1.4 46.53±3.0 48.25±3.2

24.30±1.6 (4.7e-06) 22.94±1.5 (4.5e-06) 31.88±2.1 (1.5e-05) 50.18±3.4 (6.9e-01) 45.64±3.0 (2.7e-02) 49.31±3.2 (4.2e-01) 50.82±1.6 (8.9e-01)

84.78±6.6 80.58±6.3 72.84±5.7 50.32±3.3 56.97±3.7 50.32±3.8 49.59±3.5

14.19±0.9 13.37±0.9 20.40±1.3 50.05±3.3 38.06±2.5 48.34±3.2 52.11±3.4

12.24±0.8 6.56±0.4 17.42±1.1 32.80±2.2 25.74±1.7 28.98±1.9 30.43±2.0

9.53±0.6 (6.3e-06) 5.35±0.4 (5.6e-06) 15.17±1.0 (5.5e-06) 32.68±2.2 (2.0e-01) 25.20±1.7 (1.2e-04) 28.85±1.9 (2.7e-03) 30.21±2.0 (1.2e-02)

30.73±2.0 16.55±1.1 32.21±2.1 32.00±2.1 30.71±2.0 29.44±1.9 29.57±1.9

6.37±0.4 4.95±0.3 10.97±0.7 34.74±2.3 23.15±1.5 29.76±2.0 32.67±2.2

MARGIN

50.59±3.3

51.06±3.3 (–)

41.38±2.7

66.66±4.4

35.95±2.4

34.65±2.3 (–)

20.00±1.3

47.36±3.1

SySeVR µVulDeePecker ReVeal LIVABLE HCL-VC MoEVD CLeVeR

22.03±1.4 40.50±2.1 38.69±2.0 37.15±1.9 34.71±1.8 41.05±2.1 38.95±2.0

16.77±1.1 (6.0e-07) 42.01±2.1 (2.2e-01) 39.88±2.0 (2.0e-02) 37.90±1.9 (2.3e-03) 33.85±1.7 (8.1e-05) 42.66±2.2 (4.3e-01) 40.24±2.1 (3.2e-02)

82.19±4.2 51.15±2.6 52.42±2.7 54.20±2.8 60.81±3.1 39.95±2.0 51.91±2.7

9.34±0.6 35.64±1.8 32.19±1.6 29.14±1.5 23.45±1.2 45.77±2.4 32.85±1.7

19.31±1.0 31.46±1.6 32.13±1.6 31.20±1.6 30.74±1.6 29.80±1.5 32.31±1.6

13.72±0.7 (1.7e-06) 30.84±1.5 (1.0e-02) 31.29±1.6 (2.1e-02) 29.80±1.5 (2.5e-03) 28.56±1.5 (6.0e-04) 29.59±1.5 (2.0e-03) 31.47±1.6 (2.7e-02)

53.71±2.7 36.99±1.9 39.15±2.0 41.48±2.1 44.65±2.3 28.03±1.4 40.16±2.1

8.62±0.4 28.02±1.4 27.76±1.4 25.15±1.3 24.12±1.2 32.62±1.7 27.15±1.4

MARGIN

42.08±2.2

43.85±2.3 (–)

48.60±2.5

39.95±2.1

34.80±1.8

34.39±1.8 (–)

48.60±2.5

39.95±2.1

TABLE III A BLATION STUDY OF KEY COMPONENTS IN MARGIN(C ODE T5 AS BACKBONE ). Vulnerability Detection (Vul/Non-vul Binary)

Variant

Vulnerability Classification (CWE-Macro)

MCC(%)

F1(%)

R(%)

P(%)

Macro-MCC(%)

Macro-F1(%)

Macro-R(%)

Macro-P(%)

Cosine Softmax Only

85.94±2.6

86.56±2.7

83.16±3.2

90.24±2.4

64.03±2.0

63.78±2.1

60.67±2.5

68.57±2.2

w/o Adaptive Margin w/o Adaptive Logits scaling factor

84.47±2.8 86.63±2.5

85.25±2.9 87.27±2.6

87.66±3.0 85.05±3.1

82.97±2.9 89.60±2.5

68.92±1.9 69.66±1.8

68.79±2.0 69.38±1.9

71.01±2.3 66.67±2.4

67.76±2.1 73.83±2.0

MARGIN (Full)

86.87±2.4

87.47±2.5

84.61±3.0

90.52±2.3

71.29±1.7

71.13±1.8

68.39±2.2

75.08±1.9

TABLE IV S ENSITIVITY ANALYSIS OF THE BASE SCALING FACTOR s0 ON B IG V UL . s0

Binary F1 (%)

CWE-Macro F1 (%)

5 10 15 20 25 30 35 40 45 50 55 60

86.56 85.98 87.38 87.35 86.66 87.45 87.37 87.14 86.86 85.98 85.25 85.00

69.50 63.78 70.09 71.22 69.32 67.82 70.12 68.18 70.10 68.40 68.79 69.01

latency is 111–125 ms per sample, measured with a batch size of 1 on a single NVIDIA A100 GPU with 40 GB memory. On BigVul, the model typically converges within 10– 12 epochs, after which gains are marginal. Peak GPU memory usage during training is around 12 GB, well within industrial capacity. Given these measurements, the computational cost of MARGIN is considered acceptable for industrial deployment under typical resource and latency constraints.

VII. C ASE S TUDY We use the BigVul dataset as a case study. Motivated by our analysis, we examine training dynamics and key metrics to show why our model remains effective under class imbalance. Per-CWE Performance To evaluate the classification performance of our method across different classes, we present the per-CWE F1 performance, examining how classes with varying frequencies are classified. This allows us to verify the robustness of our method under extreme class imbalance. Figure. 6 illustrates the per-class F1 scores in relation to the support (number of samples) for various CWE types, providing a detailed view of how each method copes with class imbalance. A clear pattern emerges from the visualization: while baseline methods exhibit considerable volatility that correlates with sample size, the proposed MARGIN method (depicted as a solid green line) demonstrates remarkable stability and robustness across all classes. Unlike the baselines, MARGIN consistently achieves high F1 scores—generally above 60%—showing minimal sensitivity to fluctuations in support, as indicated by the grey bars. The advantages of MARGIN are particularly pronounced in long-tail scenarios, where data scarcity presents a signif-

ARXIV PREPRINT

Per-CWE F1 for Comparison with Baselines

100

200 150

Support

clever LIVABLE

MARGIN

E-3

MoEVD muVulDeePecker ReVeal SySeVR

90 E-1 CW

16 E-4 CW

E-1 CW

E-2 CW

E-3 CW

E-1

0 E-2 CW

E-1 CW

100

F1 (%)

Fig. 6. Per-CWE F1 Performance Comparsion with Baselines Per-class F1 scores and support for MARGIN and baseline methods across different CWE categories. The bar chart indicates the number of samples per class (support), while the line plots represent F1 scores. MARGIN is highlighted with a red solid line to emphasize its performance.

a. Geometric Median Prototypes Similarity

b. Geometric Median - Weight Prototypes Similarity

Fig. 7. Prototypes Convergence During the Training Epoches. Guided by the adaptive margin mechanism, the geometric median prototypes are not only mutually separated but also aligned with the weight prototypes, suggesting that the embedding distributions of different classes gradually form a clear, separable, and stable structure.

Clustering Metrics & Gram Condition Number over Epochs 1.0

1400 1200 1000

0.6

800 600

0.4

Condition Number

Clustering Metrics (0~1)

0.8

400 0.2

200 0

0.0

NMI

ARI

Epoch

ANGULAR_SILHOUETTE_SCORE

GRAM COND NUM

Fig. 8. Clustering and ETF Dynamics During training, NMI, ARI, angular silhouette, and the Gram condition number of geometric median prototypes change rapidly at first and stabilize later, demonstrating our method’s effectiveness in achieving a well-structured, stable prototype configuration. Distortion Region Dynamics MARGIN Cosine Softmax Loss Only

0.7

FNR + FPR

0.6 0.5 0.4 0.3 0

Epoch

Fig. 9. Distortion Region Dynamics Comparison of Macro FNR + Macro FPR across training epochs. MARGIN consistently yields lower error rates than cosine softmax, indicating reduced distortion region effects and improved robustness under class imbalance.

icant challenge. For example, in CWE-416, which exhibits extreme imbalance with very few samples, most competing methods such as reveal, moevd, and mvd suffer dramatic performance degradation, with F1 scores approaching zero. In contrast, MARGIN maintains a robust F1 of approximately

60%, indicating its ability to learn discriminative features even in severely underrepresented classes. Moreover, while methods like clever and livable achieve competitive peaks in high-support classes (e.g., CWE-190), they display notable instability in medium-support categories such as CWE-200 and CWE-189. MARGIN, however, preserves a steady performance across the spectrum, demonstrating that it does not merely overfit to head classes but generalizes effectively to both frequent and rare vulnerability types. Structure Convergence Our method learns embeddings on the unit hypersphere. To examine the geometric structure during training, we conducted a per-epoch dynamic analysis across three dimensions: class concentration (via κ statistics), prototype convergence (toward ETF), and global cluster structure. Figure 7(a) shows that on BigVul, geometric median prototypes initially fluctuate but rapidly diverge and stabilize, indicating effective class separation. Figure 7(b) reveals that class weight prototypes progressively align with their geometric medians, confirming that the adaptive margin mechanism simultaneously enhances inter-class separation and intra-class compactness. Figure 8 tracks the first 11 epochs, with clustering metrics (NMI, ARI, Angular Silhouette) on the left axis and the Gram matrix condition number for geometric median prototypes on the right. All metrics rise sharply in early epochs, then plateau around epoch 4, signaling rapid convergence to a discriminative feature space towards the ETF structure. Notably, ARI reaches the highest value, indicating strong agreement with ground-truth assignments. Concurrently, the condition number plummets from > 1400 to near 1 and remains low, marking a transition from ill-conditioned, highly correlated prototypes to a numerically stable, well-spread configuration with uniform

ARXIV PREPRINT

angular distances—hallmarks of an ETF. These results demonstrate that the learned prototypes rapidly self-organize into a stable, maximally separated structure, providing empirical validation that the adaptive mechanisms effectively enforce embedding space uniformity. Distortion Elimination As discussed in our earlier section III, the distortion region refers to the phenomenon where the probability mass of one class intrudes into the Voronoi cells of other classes, directly leading to increases in global macro FNR and FPR. Our adaptive margin mechanism explicitly aims to reduce these distortion regions and thereby improve classification accuracy. The figure below illustrates the dynamics of the global macro FNR and FPR across training epochs. For comparison, we also show results obtained without the adaptive mechanism, using only the Cosine Softmax Loss, highlighting the effectiveness of our adaptive margin in mitigating distortion regions. Due to computational constraints, precisely quantifying the distortion region on the unit hypersphere is infeasible. However, as discussed earlier, it is closely related to FPR and FNR, allowing us to estimate its size indirectly by analyzing these metrics for each class. We compare MARGIN with a variant that removes the adaptive mechanism, relying solely on the cosine softmax loss (Figure 9). Even at later training stages, the cosine softmax loss exhibits relatively high FPR and FNR, whereas MARGIN significantly reduces both. This suggests the adaptive mechanism reduces distortion across classes and improves alignment within each region, thereby enhancing robustness under class imbalance. VIII. R ELATED W ORKS A. Representation Learning Code representation learning aims to map programs into continuous embedding spaces for downstream tasks such as vulnerability detection and code search. Early approaches, such as Code2Vec [6], rely on token- or path-based representations and largely overlook global structural dependencies. Graph-based methods, including MGVD [24] and Devign [25], incorporate structural information from ASTs, CFGs, and data-flow graphs, improving semantic expressiveness at the cost of higher computational complexity. More recently, Transformer-based pre-trained models, such as CodeBERT [26], GraphCodeBERT [27], CodeT5 [28], and UniXCoder [29], have become the dominant paradigm due to their strong transferability. However, existing work mainly evaluates representations via aggregate performance metrics, with limited analysis of embedding-space geometry and classwise distributional properties. B. Metric Learning Metric learning enhances representations by enforcing intraclass compactness and inter-class separability. Early pairwise or triplet losses are sensitive to sampling [30], while marginbased Softmax variants, such as ArcFace [31], introduce angular margins under hyperspherical constraints to improve open-set recognition [1], [32]. These methods apply fixed,

sample-level margins to enlarge class separation. Moreover, it focus on open-set recognition and do not explicitly address class imbalance, whereas our work targets the closed-set, imbalanced scenario. Unlike margin-based methods, our approach does not treat the margin as a tunable parameter. We model class-conditional distributions on the hypersphere and derive constraints from their geometric misalignment with Voronoi decision regions, directly regularizing the embedding geometry at the class level. While superficially similar to ArcFace-style formulations, the resulting separation emerges naturally from the global class structure rather than manually designed, sample-level margins. C. Imbalanced Classification Learning Class imbalance remains a major challenge in vulnerability detection. Conventional methods, such as Weighted CrossEntropy, Focal Loss, and Class-Balanced Loss [1], [2], mitigate imbalance through re-weighting strategies, but generally assume homogeneous intra-class distributions and lack explicit geometric [2], [4], [30], [33]. Recent work explores prototype learning and metric-space regularization to improve minority discrimination. However, such geometry-aware approaches remain underexplored in software vulnerability detection. Existing studies typically treat representation learning, metric learning, and imbalance handling as separate problems. In contrast, multi-CWE vulnerability detection requires jointly addressing semantic heterogeneity, long-tailed distributions, and embedding geometry. Our work unifies these aspects by integrating hyperspherical modeling, adaptive margins, and prototype-based learning into a single framework. IX. C ONCLUSION AND F UTURE W ORK We study vulnerability detection under realistic, highly imbalanced CWE distributions from an embedding geometry perspective. By regularizing the hyperspherical structure of representations, our method enhances class separability and decision boundary stability, achieving consistent gains over strong baselines while providing a geometrically interpretable, imbalance-aware framework. Evaluation is limited to C/C++ due to dataset availability, and generalization to other languages remains untested. Results may be influenced by stochastic factors, though improvements are consistent. Future work will focus on enhancing representation learning to better handle class imbalance, improve feature discriminability for rare CWEs, and enable more robust detection across diverse distributions. DATA AND C ODE AVAILABILITY The source code, dataset, and full pipelines are available at https://github.com/stratum-dev/MARGIN for reproducibility. ACKNOWLEDGMENT The authors thank their colleagues and reviewers for their valuable feedback. AI-assisted tools were used solely for language polishing, without influencing any technical content or scientific results.

ARXIV PREPRINT

R EFERENCES [1] J. Terven, D.-M. Cordova-Esparza, J.-A. Romero-González, A. Ramı́rezPedraza, and E. Chávez Urbiola, “A comprehensive survey of loss functions and metrics in deep learning,” Artificial Intelligence Review, vol. 58, 04 2025. [2] Y. Guo, Q. Hu, Q. Tang, and Y. L. Traon, “An empirical study of the imbalance issue in software vulnerability detection,” in Computer Security – ESORICS 2023, G. Tsudik, M. Conti, K. Liang, and G. Smaragdakis, Eds. Cham: Springer Nature Switzerland, 2024, pp. 371–390. [3] F. Li, J. Luo, F. Peng, Y. Qian, H. Yan, and Z. Zhu, “Neural collapse to multiple centers for imbalanced data,” in Advances in Neural Information Processing Systems 37, ser. NeurIPS 2024. Neural Information Processing Systems Foundation, Inc. (NeurIPS), 2024, p. 65583–65617. [Online]. Available: https://dx.doi.org/10.52202/ 079017-2095 [4] Y. Cui, M. Jia, T.-Y. Lin, Y. Song, and S. Belongie, “Classbalanced loss based on effective number of samples,” in 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). IEEE, Jun. 2019, p. 9260–9269. [Online]. Available: https://dx.doi.org/10.1109/cvpr.2019.00949 [5] Y. Guo and S. Bettaieb, “An investigation of quality issues in vulnerability detection datasets,” in 2023 IEEE European Symposium on Security and Privacy Workshops (EuroS&;PW). IEEE, Jul. 2023, p. 29–33. [Online]. Available: https://dx.doi.org/10.1109/eurospw59978. 2023.00008 [6] U. Alon, M. Zilberstein, O. Levy, and E. Yahav, “Code2vec: Learning distributed representations of code,” Proc. ACM Program. Lang., vol. 3, no. POPL, pp. 40:1–40:29, Jan. 2019. [Online]. Available: https://doi.acm.org/10.1145/3290353 [7] S. Geman, E. Bienenstock, and R. Doursat, “Neural networks and the bias/variance dilemma,” Neural Computation, vol. 4, no. 1, p. 1–58, Jan. 1992. [Online]. Available: https://dx.doi.org/10.1162/neco.1992.4.1.1 [8] P. Mettes, E. van der Pol, and C. G. M. Snoek, “Hyperspherical prototype networks,” in Advances in Neural Information Processing Systems, 2019. [9] R. Ranjan, C. D. Castillo, and R. Chellappa, “L2-constrained softmax loss for discriminative face verification,” ArXiv, vol. abs/1703.09507, 2017. [Online]. Available: https://doi.org/10.48550/arXiv.1703.09507 [10] E. Markou, T. Ajanthan, and S. Gould, “Guiding neural collapse: Optimising towards the nearest simplex equiangular tight frame,” in NeurIPS, 2024. [11] A. Banerjee, I. S. Dhillon, J. Ghosh, and S. Sra, “Clustering on the unit hypersphere using von mises–fisher distributions,” in Journal of Machine Learning Research, vol. 6, 2005, pp. 1345–1382. [12] C. Lampert, M. Mondelli, and P. Súkenı́k, “Deep neural collapse is provably optimal for the deep unconstrained features model,” in Advances in Neural Information Processing Systems 36, ser. NeurIPS 2023. Neural Information Processing Systems Foundation, Inc. (NeurIPS), 2023, p. 52991–53024. [Online]. Available: https: //dx.doi.org/10.52202/075280-2306 [13] S. Gopal and Y. Yang, “Von mises-fisher clustering models,” in Proceedings of the 31st International Conference on Machine Learning, ser. Proceedings of Machine Learning Research, E. P. Xing and T. Jebara, Eds., vol. 32, no. 1. Bejing, China: PMLR, 22–24 Jun 2014, pp. 154–162. [Online]. Available: https: //proceedings.mlr.press/v32/gopal14.html [14] J. Fan, Y. Li, S. Wang, and T. N. Nguyen, “A c/c++ code vulnerability dataset with code changes and cve summaries,” in 2020 IEEE/ACM 17th International Conference on Mining Software Repositories (MSR), 2020, pp. 508–512. [15] C. Ni, L. Shen, X. Yang, Y. Zhu, and S. Wang, “Megavul: A c/c++ vulnerability dataset with comprehensive code representations,” in Proceedings of the 21st International Conference on Mining Software Repositories, ser. MSR ’24. New York, NY, USA: Association for Computing Machinery, 2024, p. 738–742. [Online]. Available: https://doi.org/10.1145/3643991.3644886 [16] X. Wang, R. Hu, C. Gao, X.-C. Wen, Y. Chen, and Q. Liao, “Reposvul: A repository-level high-quality vulnerability dataset,” in Proceedings of the 2024 IEEE/ACM 46th International Conference on Software Engineering: Companion Proceedings, ser. ICSE-Companion ’24. New York, NY, USA: Association for Computing Machinery, 2024, p. 472–483. [Online]. Available: https://doi.org/10.1145/3639478.3647634 [17] S. Chakraborty, R. Krishna, Y. Ding, and B. Ray, “Deep learning based vulnerability detection: Are we there yet?” IEEE Transactions on Software Engineering, vol. 48, no. 9, p. 3280–3296, Sep. 2022. [Online]. Available: https://dx.doi.org/10.1109/tse.2021.3087402

[18] X.-C. Wen, C. Gao, F. Luo, H. Wang, G. Li, and Q. Liao, “Livable: Exploring long-tailed classification of software vulnerability types,” IEEE Transactions on Software Engineering, vol. 50, no. 6, p. 1325–1339, Jun. 2024. [Online]. Available: https://dx.doi.org/10.1109/ tse.2024.3382361 [19] D. Zou, S. Wang, S. Xu, Z. Li, and H. Jin, “µvuldeepecker: A deep learning-based system for multiclass vulnerability detection,” IEEE Transactions on Dependable and Secure Computing, vol. 18, no. 5, pp. 2224–2236, 2021. [20] Z. Li, D. Zou, S. Xu, H. Jin, Y. Zhu, and Z. Chen, “Sysevr: A framework for using deep learning to detect software vulnerabilities,” IEEE Transactions on Dependable and Secure Computing, vol. 19, no. 4, p. 2244–2258, Jul. 2022. [Online]. Available: https://dx.doi.org/ 10.1109/tdsc.2021.3051525 [21] C. Ji, S. Yang, H. Sun, and Y. Zhang, “Applying contrastive learning to code vulnerability type classification,” in Conference on Empirical Methods in Natural Language Processing, 2024. [Online]. Available: https://api.semanticscholar.org/CorpusID:273901675 [22] X. Yang, S. Wang, J. Zhou, and W. Zhu, “One-for-all does not work! enhancing vulnerability detection by mixture-of-experts (moe),” Proceedings of the ACM on Software Engineering, vol. 2, no. FSE, p. 446–464, Jun. 2025. [Online]. Available: https://dx.doi.org/10.1145/ 3715736 [23] J. Li, L. Cui, S. Zhao, Y. Yang, L. Li, and H. Zhu, “Clever: Multi-modal contrastive learning for vulnerability code representation,” in Findings of the Association for Computational Linguistics: ACL 2025. Association for Computational Linguistics, 2025, p. 7940–7951. [Online]. Available: https://dx.doi.org/10.18653/v1/2025.findings-acl.414 [24] F. Qiu, Z. Liu, X. Hu, X. Xia, G. Chen, and X. Wang, “Vulnerability detection via multiple-graph-based code representation,” IEEE Transactions on Software Engineering, vol. 50, no. 8, p. 2178–2199, Aug. 2024. [Online]. Available: https://dx.doi.org/10.1109/ tse.2024.3427815 [25] Y. Zhou, S. Liu, J. Siow, X. Du, and Y. Liu, Devign: effective vulnerability identification by learning comprehensive program semantics via graph neural networks. Red Hook, NY, USA: Curran Associates Inc., 2019. [26] Z. Feng, D. Guo, D. Tang, N. Duan, X. Feng, M. Gong, and M. Zhou, “Codebert: A pre-trained model for programming and natural languages,” in Findings of the Association for Computational Linguistics: EMNLP 2020, 2020, pp. 1536–1547. [Online]. Available: https://dx.doi.org/10.18653/v1/2020.findings-emnlp.139 [27] D. Guo, S. Ren, S. Lu, Z. Feng, D. Tang, S. Liu, L. Zhou, N. Duan, A. Svyatkovskiy, S. Fu, M. Tufano, S. K. Deng, C. B. Clement, D. Drain, N. Sundaresan, J. Yin, D. Jiang, and M. Zhou, “Graphcodebert: Pre-training code representations with data flow.” in ICLR. OpenReview.net, 2021. [Online]. Available: https: //dblp.uni-trier.de/db/conf/iclr/iclr2021.html#GuoRLFT0ZDSFTDC21 [28] Y. Wang, W. Wang, S. Joty, and S. C. Hoi, “Codet5: Identifier-aware unified pre-trained encoder-decoder models for code understanding and generation,” in Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 2021, p. 8696–8708. [Online]. Available: https://dx.doi.org/10.18653/v1/2021.emnlp-main.685 [29] D. Guo, S. Lu, N. Duan, Y. Wang, M. Zhou, and J. Yin, “Unixcoder: Unified cross-modal pre-training for code representation,” in Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). Association for Computational Linguistics, 2022, p. 7212–7225. [Online]. Available: https://dx.doi.org/10.18653/v1/2022.acl-long.499 [30] X. Yang, S. Wang, Y. Li, and S. Wang, “Does data sampling improve deep learning-based vulnerability detection? yeas! and nays!” in Proceedings of the 45th International Conference on Software Engineering, ser. ICSE ’23. IEEE Press, 2023, p. 2287–2298. [Online]. Available: https://doi.org/10.1109/ICSE48619.2023.00192 [31] J. Deng, J. Guo, N. Xue, and S. Zafeiriou, “Arcface: Additive angular margin loss for deep face recognition,” in 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2019, pp. 4685– 4694. [32] Y. Meng, J. Huang, G. Wang, C. Zhang, H. Zhuang, L. Kaplan, and J. Han, Spherical text embedding. Red Hook, NY, USA: Curran Associates Inc., 2019. [33] Y. He, G. Lin, X. Ma, J. W. Keung, C. Tan, W. Hu, and F. Li, “Enhancing deep learning vulnerability detection through imbalance loss functions: An empirical study,” Proceedings of the 15th Asia-Pacific Symposium on Internetware, 2024. [Online]. Available: https://api.semanticscholar.org/CorpusID:271293860

Record · ID 175133 · SHA-256 5eb5eef292dc57c7
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.