A Cascaded Machine Learning Approach using the EMBER2024 Dataset
| Name | Role |
|---|---|
| Lorena Sarasua-Fernandez | Collaborator |
| Sophie Liu | Collaborator |
| Chuks Ofojuah | Collaborator |
| Sharon Owusu | Collaborator |
| Atharva Sasankar | Collaborator |
| Brandon Jaipersaud | Collaborator |
Malware detection is a cybersecurity process that aims to identify whether a file is malicious or benign before it can damage a system. Despite the effectiveness of traditional detection methods, the consistent emergence of new malware creates a bottleneck — which is where machine learning comes in.
This project uses a portion of the EMBER2024 dataset to investigate whether an ML model can accurately classify files as malicious or benign. For files identified as malicious, we further explore whether their behavioral tags can be extracted. Through responsible AI practices and a cascaded model design, we evaluate how well a model can detect modern malware and characterize its behavior.
As cyber threats evolve, traditional signature-based detection struggles against newly emerging or modified malware. As individuals, businesses, and governments grow increasingly dependent on digital systems, improving automated detection has become a critical cybersecurity challenge.
Machine learning can learn patterns from known malware samples and apply that knowledge to previously unseen examples — helping to strengthen current defenses. Improving automated classification also helps security professionals respond faster and more consistently, and our findings may shed light on the strengths and limitations of different ML approaches for malware analysis.
EMBER2024 — a large-scale benchmark dataset for holistic malware classification.
Our project uses a 20% deterministic sample of the WIN32 file types, the largest subset of file types, from the EMBER2024 dataset to account for file distribution bias and storage limitations.
The first model uses LightGBM to determine whether a given file is malicious or benign.
LightGBM is a tree-based supervised learning algorithm that builds an ensemble of decision trees sequentially — each tree correcting the errors of the previous one. Compared to a single decision tree (prone to overfitting) or a random forest (independent trees), gradient boosting builds trees in a directed, more effective way on structured data. LightGBM is also optimized for speed and memory efficiency at scale, making it well-suited for a dataset as large as EMBER2024.
Files flagged as malicious by Stage 1 are passed to a second LightGBM model that identifies what the malware does — whether it encrypts files, exfiltrates data, logs keystrokes, etc.
This is a multiclass problem where each file is assigned a behavioral tag from the EMBER2024 set. LightGBM was chosen for this stage because it handles class imbalance well and performs strongly on structured and tabular data.
| Model | Accuracy | ROC AUC | Train Samples | Test Samples |
|---|---|---|---|---|
| LGBMClassifier | 87.92% | 0.9574 | 312,125 | 359,994 |
- Correctly detected 150,870 out of 180,000 malware samples.
- 14,356 false positives (benign files flagged as malware)
- 29,130 false negatives (actual malware missed)
Figure 2: This is the feature importance.
Figure 3: These are the API categories.
Figure 4: These are the imported API counts.
- Anderson, H.S. et al. — EMBER2024 KDD Paper
- CrowdStrike — EMBER2024: Advancing Cybersecurity ML Training on Evasive Malware
- Google Scholar: Papers citing EMBER
- Ke, G. et al. — LightGBM: A Highly Efficient Gradient Boosting Decision Tree (NeurIPS 2017)
- LightGBM Documentation
- XGBoost vs LightGBM (Medium)
- RandomForestClassifier — scikit-learn