Microsoft Malware Detection

A 9-class malware family classifier built on the Microsoft Malware Classification Challenge (Kaggle, 2015). Final XGBoost model achieves a multi-class log-loss of 0.011 — a ≈ 200× improvement over a random baseline (2.486).

Problem

Microsoft runs anti-malware utilities over 150 million computers worldwide, generating tens of millions of daily data points. Given a malware sample (.asm + .bytes representation), identify which of nine malware families it belongs to.

Original Kaggle competition: https://www.kaggle.com/c/malware-classification

Why it's hard

• 9 classes — Ramnit, Lollipop, Kelihos_ver3, Vundo, Simda, Tracur, Kelihos_ver1, Obfuscator.ACY, Gatak
• 200 GB raw data — 10,868 .asm files (~150 GB) + 10,868 .bytes files (~50 GB)
• Constraint: classification must be fast (seconds, not hours)
• Metric: multi-class log-loss — penalises mis-confident predictions, not just wrong ones

Approach

1. Feature engineering (two file types, three feature families)

Byte files — raw hexadecimal content of the binary (no PE header):

• File size as a feature
• Byte-frequency features
• Byte bigrams — built a 66 k vocabulary of byte-pairs, then selected top 1000 most discriminative via Chi-square (SelectKBest)
• Pixel-intensity image features — interpret the first 800 bytes of each file as a grayscale image (technique from the 1st-place Kaggle solution)

ASM files — disassembled x86 with segments / opcodes / registers / API calls:

• File size
• 52 handpicked features (segment counts, opcode frequencies, register usage, etc.) — selected after reading top-solution write-ups and the Nataraj et al. malware-as-images paper
• Pixel-intensity image features from .asm content
• Parallel processing used to make 150 GB of ASM extraction tractable

2. Models compared

Five algorithm families were trained and hyperparameter-tuned across the three feature sets (bytes only, asm only, bytes + asm + bigrams + image):

• Random model (baseline)
• KNN
• Logistic Regression
• Random Forest
• XGBoost (with RandomizedSearchCV for hyperparameter tuning)

3. Train / CV / Test split

Stratified 64 / 16 / 20 split. All hyperparameters chosen on CV, never on test.

Results

Stage	Model	Features	Test log-loss
Baseline	Random model	bytes	2.486
1	KNN	bytes	0.242
2	Logistic Regression	bytes	0.528
3	Random Forest	bytes	0.086
4	XGBoost	bytes	0.079
5	KNN	asm (52 feat.)	0.089
6	Random Forest	asm	0.057
7	XGBoost	asm	0.049
8	Random Forest	bytes + asm	0.040
9	XGBoost	bytes + asm	0.032
10	XGBoost	bytes + asm + bigrams + image	0.011

The combined-feature XGBoost model is the final winner. Reading down the table is the project's story: each row is either a different algorithm on the same features, or the same algorithm on richer features, and the log-loss falls steadily as both axes are explored.

Takeaways

• Feature engineering > algorithm choice on this problem. The biggest single jumps came from adding feature families (asm → asm + bytes; flat features → bigrams + images), not from changing models within a feature set.
• Tree ensembles dominate for tabular malware features — Random Forest and XGBoost consistently top KNN and Logistic Regression by an order of magnitude.
• Malware-as-images works — converting the first 800 bytes of a binary into a pixel-intensity vector adds real signal, validating the Nataraj et al. approach.
• Chi-square feature selection scales — reducing 66 k bigrams → 1 k retained discriminative power while making the model trainable on a single workstation.

Tech stack

Python 3 · NumPy · pandas · scikit-learn · XGBoost · joblib (parallel feature extraction) · matplotlib · seaborn · Jupyter

Run

pip install numpy pandas scikit-learn xgboost matplotlib seaborn jupyter
jupyter notebook MicrosoftMalwareDetection.ipynb

Note on data: the raw competition data is ~200 GB and is not included in this repo. Download from the original Kaggle competition: https://www.kaggle.com/c/malware-classification/data. The notebook contains preserved outputs throughout, so the full analysis, model comparisons, and final results are visible end-to-end without needing to re-execute.

References

• Microsoft Malware Classification Kaggle competition: https://www.kaggle.com/c/malware-classification
• 1st-place solution write-up: http://blog.kaggle.com/2015/05/26/microsoft-malware-winners-interview-1st-place-no-to-overfitting/
• Nataraj et al., "Malware Images: Visualisation and Automatic Classification" (VizSec 2011): http://vizsec.org/files/2011/Nataraj.pdf
• 1st-place solution video walkthrough: https://www.youtube.com/watch?v=VLQTRlLGz5Y

License

MIT