README.md

---
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