| --- |
| title: Malware Detection Pipeline |
| emoji: 🔬 |
| colorFrom: red |
| colorTo: purple |
| sdk: docker |
| app_port: 7860 |
| pinned: false |
| --- |
| |
| # Two-Stage Malware Detection Pipeline |
| ### MLP + Vision Transformer |
|
|
| **Student:** Amine Garaali | **Deep Learning Project** | **2026** |
|
|
| --- |
|
|
| ## How it works |
|
|
| This demo combines two deep learning models into a sequential detection pipeline — the same architecture used in real-world security products. |
|
|
| ``` |
| Upload .exe or .dll |
| ↓ |
| ┌─────────────────────────────────────────┐ |
| │ Stage 1 — MLP (EMBER features) │ |
| │ Extracts 2381 static PE features │ |
| │ Outputs: malware probability │ |
| └─────────────────────────────────────────┘ |
| ↓ |
| probability < 0.5 → ✅ BENIGN — stop |
| probability ≥ 0.5 → 🔴 MALWARE → Stage 2 |
| ↓ |
| ┌─────────────────────────────────────────┐ |
| │ Stage 2 — Vision Transformer (ViT) │ |
| │ Converts binary to grayscale image │ |
| │ Classifies into 1 of 25 malware │ |
| │ families using self-attention │ |
| └─────────────────────────────────────────┘ |
| ↓ |
| Family name + confidence score |
| ``` |
|
|
| No file is executed at any point. Both models perform static analysis only. |
|
|
| --- |
|
|
| ## Stage 1 — MLP Malware Detector |
|
|
| - **Architecture:** 4-layer MLP (2381 → 512 → 256 → 128 → 1) |
| - **Input:** 2381-dimensional EMBER feature vector (imports, byte histogram, sections, strings, header info) |
| - **Task:** Binary classification — malware vs benign |
| - **Dataset:** EMBER 2018 — 800,000 labeled Windows PE files |
| - **Performance:** 95% accuracy, ROC-AUC 0.9878 |
| - **Key components:** BatchNorm, ReLU, Dropout(0.3), Sigmoid output, Adam optimizer |
|
|
| The MLP acts as a fast, high-accuracy first filter. If the file scores below 50% malware probability it is immediately classified as benign and Stage 2 is skipped. |
|
|
| --- |
|
|
| ## Stage 2 — Vision Transformer (ViT) Family Classifier |
|
|
| - **Architecture:** ViT built from scratch — PatchEmbedding + 6 TransformerBlocks + classification head |
| - **Input:** 64×64 grayscale image derived from raw file bytes (each byte = one pixel) |
| - **Task:** 25-class malware family identification |
| - **Dataset:** MalImg — 9,339 grayscale malware images |
| - **Performance:** 98% accuracy across 25 families |
| - **Key components:** Multi-head self-attention (8 heads), GELU activation, CLS token, positional embeddings, AdamW optimizer |
|
|
| The ViT processes the file's binary content as a sequence of 64 image patches, computing attention between all patch pairs simultaneously. This allows it to detect structural relationships across the entire file that convolutional networks cannot capture in a single layer. |
|
|
| --- |
|
|
| ## The 25 malware families |
|
|
| ``` |
| Adialer.C Agent.FYI Allaple.A Allaple.L Alueron.gen!J |
| Autorun.K C2LOP.P C2LOP.gen!g Dialplatform.B Dontovo.A |
| Fakerean Instantaccess Lolyda.AA1 Lolyda.AA2 Lolyda.AA3 |
| Lolyda.AT Malex.gen!J Obfuscator.AD Rbot!gen Skintrim.N |
| Swizzor.gen!E Swizzor.gen!I VB.AT Wintrim.BX Yuner.A |
| ``` |
|
|
| --- |
|
|
| ## Why two stages? |
|
|
| | | Stage 1 (MLP) | Stage 2 (ViT) | |
| |---|---|---| |
| | Question answered | Is this malware? | What kind of malware? | |
| | Input | 2381 engineered features | Raw binary as image | |
| | Feature learning | Human-designed | Learned by model | |
| | Speed | Very fast | Slower | |
| | Can detect benign? | Yes | No | |
|
|
| The MLP is fast and handles the binary decision efficiently. The ViT only runs when needed, providing deeper analysis for confirmed malware. This mirrors how production security tools work — cheap detection first, expensive classification second. |
|
|
| --- |
|
|
| ## Limitations |
|
|
| - Stage 1 was trained on PE files from 2018 — very recent malware using novel evasion may score lower than expected |
| - Stage 2 recognizes only the 25 MalImg families — novel malware will be classified into the closest-looking known family |
| - Neither model executes the file — dynamic behavior (e.g. fileless malware) is not analyzed |
| - For research and educational purposes only |
|
|
| --- |
|
|
| ## References |
|
|
| - Anderson & Roth (2018). *EMBER: An Open Dataset for Training Static PE Malware Machine Learning Models.* arXiv:1804.04637 |
| - Nataraj et al. (2011). *Malware Images: Visualization and Automatic Classification.* VizSec 2011 |
| - Dosovitskiy et al. (2020). *An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.* arXiv:2010.11929 |
| - Khan & Nauman (2024). *Interpretable Detection of Malicious Behavior in Windows PE using Multi-Head 2D Transformers.* Big Data Mining and Analytics |
|
|
