Explainable Deep Learning for Windows Malware Detection Using Static PE Features

Overview

A comprehensive, explainability-first malware detection framework trained on the EMBER 2018 benchmark dataset (799,876 labeled Windows PE file samples). This project systematically compares four model families and applies SHAP and LIME to produce actionable security explanations.

Results Summary

Model	Accuracy	ROC-AUC	TPR@1%FPR
XGBoost	97.21%	0.9952	94.82%
LightGBM	96.84%	0.9948	94.20%
MLP	96.06%	0.9883	89.64%
RandomForest	94.20%	0.9870	81.80%

Key XAI finding: Import table features account for ~38% of total SHAP importance, followed by byte histogram patterns (~21%).

Project Structure

malware-detection/
├── data/
│   ├── raw/                    # EMBER .npy files
│   ├── processed/              # cleaned & processed
│   └── splits/                 # train/val/test splits
├── notebooks/
│   ├── 01_EDA.ipynb            # Exploratory Data Analysis
│   ├── 02_preprocessing.ipynb  # Feature Engineering
│   ├── 03_baseline.ipynb       # Baseline model results
│   ├── 04_deep_learning.ipynb  # MLP analysis
│   ├── 05_explainability.ipynb # SHAP & LIME
│   └── 06_results.ipynb        # Final results & paper figures
├── src/
│   ├── data_loader.py          # Memory-safe EMBER data loading
│   ├── features.py             # PE feature metadata (2381 dims)
│   ├── models/
│   │   ├── baseline.py         # LightGBM, XGBoost, RandomForest
│   │   ├── mlp.py              # Deep MLP (PyTorch)
│   │   └── cnn.py              # 1D CNN (PyTorch)
│   ├── explainability/
│   │   ├── shap_analysis.py    # SHAP TreeExplainer + group attribution
│   │   └── lime_analysis.py    # LIME local explanations
│   ├── train.py                # Full training pipeline
│   ├── evaluate.py             # Unified evaluation metrics
│   └── visualize.py            # Publication-ready figures
├── results/
│   ├── figures/                # All generated figures
│   ├── metrics/                # JSON results files
│   └── models/                 # Saved model files (.pkl, .pt)
├── paper/
│   ├── draft.md                # Research paper draft
│   └── figures/                # Paper-quality figures
├── config.yaml                 # All hyperparameters
├── run_explainability.py       # XAI pipeline entry point
└── requirements.txt

Quick Start

1. Install Dependencies

pip install -r requirements.txt

2. Prepare Data

Place EMBER 2018 .npy files in data/raw/:

data/raw/X_train_final.npy
data/raw/y_train_final.npy
data/raw/X_test_final.npy
data/raw/y_test_final.npy

3. Train Models

# Train all baseline models (LightGBM, XGBoost, RandomForest)
py -m src.train --model baselines

# Train MLP
py -m src.train --model mlp

# Debug mode (fast, 50k samples)
py -m src.train --debug --model baselines

4. Run Explainability

# SHAP + LIME for all models
py run_explainability.py --model all

# Single model
py run_explainability.py --model lightgbm

5. Run Notebooks

jupyter notebook notebooks/

Dataset

EMBER 2018 — Endgame Malware BEnchmark for Research

Split	Benign	Malicious	Total
Train	254,992	254,940	509,932
Val	44,999	44,989	89,988
Test	99,985	99,971	199,956

Feature space: 2,381 static PE features across 10 groups:

• ByteHistogram (256), ByteEntropy (256), StringFeatures (100)
• GeneralInfo (10), SectionInfo (92), ImportsInfo (384)
• ExportsInfo (128), DataDirectories (30), HeaderInfo (62), Extra (1063)

Key Findings

1. Import table features dominate (~38% of SHAP importance) — consistent with malware using suspicious API calls
2. Byte patterns contribute ~31% (histogram + entropy) — packing/obfuscation detection
3. False positives are mainly legitimately packed software with high byte entropy
4. False negatives typically use dynamic API loading to hide imports

Citation

If you use this code, please cite:

@article{mariam2025explainable,
  title   = {Explainable Static Malware Detection on Windows PE Files Using Ensemble Learning},
  author  = {Mariam Zakaria},
  journal = {arXiv preprint},
  year    = {2025}
}

License

MIT License — see LICENSE file.