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