--- library_name: transformers tags: - security - cyber-security - CWE - vulnerability-classification - cve license: apache-2.0 datasets: - zefang-liu/cve-and-cwe-mapping-dataset language: - en metrics: - accuracy - f1 base_model: - distilbert/distilbert-base-uncased pipeline_tag: text-classification model-index: - name: cwe-predictor results: - task: type: text-classification name: CWE Classification metrics: - type: accuracy value: 0.727207 name: Validation Accuracy - type: f1 value: 0.251264 name: Macro F1 Score --- # CWE Predictor - Vulnerability Classification Model This model classifies vulnerability descriptions into Common Weakness Enumeration (CWE) categories. It's designed to help security professionals and developers quickly identify the type of vulnerability based on textual descriptions. ## Model Details ### Model Description This is a fine-tuned DistilBERT model that predicts CWE (Common Weakness Enumeration) categories from vulnerability descriptions. The model was trained on a comprehensive dataset of CVE descriptions mapped to their corresponding CWE identifiers. **Key Features:** - Classifies vulnerabilities into 232 distinct CWE categories - Trained on 111,640 vulnerability descriptions - Achieves 72.72% accuracy on validation set - Macro F1 score of 0.251 demonstrating balanced performance across classes - Lightweight and fast inference using DistilBERT architecture - **Developed by:** [mulliken](https://huggingface.co/mulliken) - **Model type:** DistilBERT (Transformer-based classifier) - **Language(s) (NLP):** English - **License:** Apache 2.0 - **Finetuned from model:** [distilbert/distilbert-base-uncased](https://huggingface.co/distilbert/distilbert-base-uncased) ### Model Sources - **Hugging Face Model:** [mulliken/cwe-predictor](https://huggingface.co/mulliken/cwe-predictor) - **Dataset:** [CVE and CWE Mapping Dataset](https://huggingface.co/datasets/zefang-liu/cve-and-cwe-mapping-dataset) ## Uses ### Direct Use This model can be used directly for: - **Vulnerability Triage:** Automatically classify security vulnerabilities reported in bug bounty programs or security audits - **Security Analysis:** Categorize CVE descriptions to understand vulnerability patterns - **Automated Security Reporting:** Generate CWE classifications for vulnerability reports - **Security Research:** Analyze trends in vulnerability types across codebases ### Downstream Use The model can be integrated into: - Security scanning tools and SAST/DAST platforms - Vulnerability management systems - Security information and event management (SIEM) systems - DevSecOps pipelines for automated vulnerability classification ### Out-of-Scope Use This model should NOT be used for: - Medical or safety-critical systems without additional validation - As the sole method for security assessment (should complement human expertise) - Classifying non-English vulnerability descriptions - Real-time security detection (model is designed for post-discovery classification) ## Bias, Risks, and Limitations ### Known Limitations - **Class Imbalance:** Some CWE categories are underrepresented in the training data, which may lead to lower accuracy for rare vulnerability types - **Temporal Bias:** Model trained on historical CVE data may not recognize newer vulnerability patterns - **Language Limitation:** Only trained on English descriptions - **Context Loss:** Limited to 512 tokens, longer descriptions are truncated ### Risks - False negatives could lead to unidentified security vulnerabilities - Should not replace human security expertise - May not generalize well to proprietary or domain-specific vulnerability descriptions ### Recommendations - Always use this model as a supplementary tool alongside human security expertise - Validate predictions for critical security decisions - Consider retraining or fine-tuning for domain-specific applications - Monitor model performance over time as new vulnerability types emerge ## How to Get Started with the Model ### Installation ```bash pip install transformers torch ``` ### Quick Start ```python from transformers import AutoTokenizer, AutoModelForSequenceClassification import torch # Load model and tokenizer model = AutoModelForSequenceClassification.from_pretrained("mulliken/cwe-predictor") tokenizer = AutoTokenizer.from_pretrained("mulliken/cwe-predictor") # Prediction function def predict_cwe(text: str) -> str: encoded = tokenizer(text, return_tensors="pt", truncation=True, max_length=512) with torch.no_grad(): logits = model(**encoded).logits pred_id = torch.argmax(logits, dim=-1).item() return model.config.id2label[pred_id] # Example usage vuln_description = "Buffer overflow in the authentication module allows remote attackers to execute arbitrary code." cwe_prediction = predict_cwe(vuln_description) print(f"Predicted CWE: {cwe_prediction}") ``` ### Example Predictions ```python examples = [ "SQL injection vulnerability in login form allows attackers to bypass authentication", "Cross-site scripting (XSS) vulnerability in comment section", "Path traversal vulnerability allows reading arbitrary files", "Integer overflow in image processing library causes memory corruption" ] for desc in examples: print(f"Description: {desc}") print(f"Predicted CWE: {predict_cwe(desc)}\n") ``` ## Training Details ### Training Data The model was trained on the [CVE and CWE Mapping Dataset](https://huggingface.co/datasets/zefang-liu/cve-and-cwe-mapping-dataset), which contains: - CVE descriptions from the National Vulnerability Database (NVD) - Corresponding CWE classifications - Dataset size: 124,045 examples after filtering - Training set: 111,640 examples - Validation set: 12,405 examples - Number of CWE classes: 232 (after removing generic categories like "NVD-CWE-Other" and "NVD-CWE-noinfo") ### Training Procedure #### Preprocessing 1. **Data Cleaning:** - Removed entries with missing descriptions or CWE IDs - Filtered out generic CWE categories ("NVD-CWE-Other", "NVD-CWE-noinfo") - Removed CWE categories with only 1 example to ensure stratified splitting 2. **Tokenization:** - Used DistilBERT tokenizer with max_length=512 - Applied truncation for longer descriptions #### Training Hyperparameters - **Learning rate:** 2e-5 - **Batch size:** 2 per device with gradient accumulation of 8 (effective batch size: 16) - **Number of epochs:** 1 - **Weight decay:** 0.01 - **Optimizer:** AdamW - **Training regime:** fp32 with gradient checkpointing - **Evaluation strategy:** Every 1000 steps #### Training Performance - **Total training time:** ~78 minutes (4712 seconds) (per epoch) - **Training steps:** 13,956 - **Training samples per second:** 23.691 - **Final training loss:** 1.134700 - **Best validation loss:** 1.082806 (at step 6000) - **Model size:** ~268MB ## Evaluation ### Testing Data, Factors & Metrics #### Testing Data Validation set of 12,405 examples (10% stratified split from the training data) #### Metrics - **Accuracy:** Overall correctness of predictions - **Macro F1 Score:** Unweighted mean of F1 scores for each class (ensures balanced performance across all CWE types) ### Results | Step | Training Loss | Validation Loss | Accuracy | Macro F1 | |------|--------------|-----------------|----------|----------| | 1000 | 1.044600 | 1.252940 | 0.704716 | 0.220344 | | 2000 | 1.158700 | 1.188677 | 0.711326 | 0.229855 | | 3000 | 1.119900 | 1.159229 | 0.719226 | 0.235295 | | 4000 | 1.112600 | 1.119924 | 0.720193 | 0.242404 | | 5000 | 1.110300 | 1.111053 | 0.722934 | 0.244389 | | 6000 | 1.134700 | 1.082806 | 0.727207 | 0.251264 | #### Summary The model achieves 72.72% accuracy on the validation set with a macro F1 score of 0.251. The relatively lower F1 score reflects the challenge of classifying across 232 different CWE categories with varying representation in the dataset. ## Model Examination The model uses standard DistilBERT attention mechanisms to process vulnerability descriptions. Key observations: - The model learns to identify security-related keywords and patterns - Attention weights typically focus on vulnerability-specific terms (e.g., "overflow", "injection", "traversal") - Performance varies by CWE category based on training data representation ## Environmental Impact Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700). - **Hardware Type:** Apple Silicon (M-series chip) - **Hours used:** ~1.3 hours - **Cloud Provider:** Local training (no cloud provider) - **Compute Region:** N/A (local) - **Carbon Emitted:** Minimal (Apple Silicon is energy efficient, ~15W TDP) ## Technical Specifications [optional] ### Model Architecture and Objective - **Base Architecture:** DistilBERT (distilbert-base-uncased) - **Task:** Multi-class text classification - **Number of labels:** 232 CWE categories - **Objective:** Cross-entropy loss for sequence classification - **Architecture modifications:** Added classification head with 232 output classes ### Compute Infrastructure Local machine with Apple Silicon processor #### Hardware - **Device:** Apple Silicon (MPS backend) - **Memory management:** PYTORCH_MPS_HIGH_WATERMARK_RATIO set to 0.0 #### Software - **Framework:** PyTorch with Hugging Face Transformers - **Python version:** 3.x - **Key libraries:** transformers, torch, datasets, scikit-learn, pandas, numpy ## Citation If you use this model in your research, please cite: ```bibtex @misc{mulliken2024cwepredictcr, author = {mulliken}, title = {CWE Predictor: A DistilBERT Model for Vulnerability Classification}, year = {2024}, publisher = {Hugging Face}, howpublished = {\url{https://huggingface.co/mulliken/cwe-predictor}} } ``` ## Glossary - **CWE (Common Weakness Enumeration):** A community-developed list of software and hardware weakness types - **CVE (Common Vulnerabilities and Exposures):** A list of publicly disclosed cybersecurity vulnerabilities - **NVD (National Vulnerability Database):** U.S. government repository of vulnerability management data - **Macro F1:** The unweighted mean of F1 scores calculated for each class independently - **SAST/DAST:** Static/Dynamic Application Security Testing ## More Information For questions, issues, or contributions, please visit the [Hugging Face model page](https://huggingface.co/mulliken/cwe-predictor). ## Model Card Authors - [mulliken](https://huggingface.co/mulliken) ## Model Card Contact Please use the Hugging Face model repository's discussion section for questions and feedback: [mulliken/cwe-predictor](https://huggingface.co/mulliken/cwe-predictor/discussions)