MoEVD: Enhancing Vulnerability Detection by Mixture-of-Experts (MoE)

2025-01-27

AI Summary

Existing deep learning–based vulnerability detection (DLVD) methods adopt a “one-size-fits-all” paradigm, exhibiting poor adaptability to diverse Common Weakness Enumeration (CWE) types and particularly weak performance on long-tail CWEs with scarce training data. Method: This paper proposes the first two-stage Mixture-of-Experts (MoE) framework tailored for vulnerability detection: Stage I precisely classifies the CWE type of a given code snippet; Stage II dispatches the input to a dedicated expert network—trained exclusively on that CWE—for fine-grained vulnerability identification. This decouples CWE classification from vulnerability judgment, mitigating the generalization bottleneck inherent in monolithic models. Contribution/Results: The framework significantly enhances detection accuracy for long-tail CWEs. Experiments show an overall F1-score of 0.44—surpassing the best baseline by ≥12.8%. For long-tail CWEs specifically, F1-score improves by ≥7.3%, and recall increases by 9.0%–77.8%.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: Multi-modal Vision
• Data Mining & Knowledge Management: Anomaly/Outlier Detection

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Web Mining and Content Analysis: Large pretrained models with web data
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Deep Learning-based Vulnerability Detection (DLVD) techniques have garnered significant interest due to their ability to automatically learn vulnerability patterns from previously compromised code. Despite the notable accuracy demonstrated by pioneering tools, the broader application of DLVD methods in real-world scenarios is hindered by significant challenges. A primary issue is the"one-for-all"design, where a single model is trained to handle all types of vulnerabilities. This approach fails to capture the patterns of different vulnerability types, resulting in suboptimal performance, particularly for less common vulnerabilities that are often underrepresented in training datasets. To address these challenges, we propose MoEVD, which adopts the Mixture-of-Experts (MoE) framework for vulnerability detection. MoEVD decomposes vulnerability detection into two tasks, CWE type classification and CWE-specific vulnerability detection. By splitting the task, in vulnerability detection, MoEVD allows specific experts to handle distinct types of vulnerabilities instead of handling all vulnerabilities within one model. Our results show that MoEVD achieves an F1-score of 0.44, significantly outperforming all studied state-of-the-art (SOTA) baselines by at least 12.8%. MoEVD excels across almost all CWE types, improving recall over the best SOTA baseline by 9% to 77.8%. Notably, MoEVD does not sacrifice performance on long-tailed CWE types; instead, its MoE design enhances performance (F1-score) on these by at least 7.3%, addressing long-tailed issues effectively.

Problem

• Deep Learning Vulnerability Detection
• Rare Vulnerability
• Performance Improvement

Innovation

• MoEVD
• Mixture-of-Experts
• Deep Learning-based Vulnerability Detection