Case Study: Fine-tuning Small Language Models for Accurate and Private CWE Detection in Python Code
This provides a privacy-preserving, on-premise solution for software developers to integrate accurate vulnerability detection into their workflows, though it is incremental as it applies existing fine-tuning methods to a new security domain.
This work tackled the problem of detecting security vulnerabilities (MITRE Top 25 CWEs) in Python code by fine-tuning a small language model (350M parameters) instead of relying on large cloud-based models, achieving approximately 99% accuracy, 98.08% precision, 100% recall, and 99.04% F1-score on a specialized test set.
Large Language Models (LLMs) have demonstrated significant capabilities in understanding and analyzing code for security vulnerabilities, such as Common Weakness Enumerations (CWEs). However, their reliance on cloud infrastructure and substantial computational requirements pose challenges for analyzing sensitive or proprietary codebases due to privacy concerns and inference costs. This work explores the potential of Small Language Models (SLMs) as a viable alternative for accurate, on-premise vulnerability detection. We investigated whether a 350-million parameter pre-trained code model (codegen-mono) could be effectively fine-tuned to detect the MITRE Top 25 CWEs specifically within Python code. To facilitate this, we developed a targeted dataset of 500 examples using a semi-supervised approach involving LLM-driven synthetic data generation coupled with meticulous human review. Initial tests confirmed that the base codegen-mono model completely failed to identify CWEs in our samples. However, after applying instruction-following fine-tuning, the specialized SLM achieved remarkable performance on our test set, yielding approximately 99% accuracy, 98.08% precision, 100% recall, and a 99.04% F1-score. These results strongly suggest that fine-tuned SLMs can serve as highly accurate and efficient tools for CWE detection, offering a practical and privacy-preserving solution for integrating advanced security analysis directly into development workflows.