Malware detection with MaxViT transformer

Detect PE malicious file with Multi-Axis Vision Transformer (MaxViT model)

Getting started

clone repo

git clone https://github.com/prozip/malware-detection_MaxViT-Transformer

Prepare data

• Modify location path in config/train_config.yml
• Dataset contains folders with name as label

Train

python3 train.py -t S
python3 train.py -t S -b 64

Parameters

• edit para in config/*.yml

Architecture

1. Convert PE file to color image

python2 bin2rgb.py

• PE file -> entropy (Shannon)
• entropy -> rgb color (Binvis)
• rgb color -> image (Hilbert curve)
  

2. MaxViT: Multi-Axis Vision Transformer

• MaxT-0 (31M): S1 hidden_size 64; num_blocks [2, 2, 5, 2]
• MaxT-1 (69M): S1 hidden_size 96; num_blocks [2, 2, 5, 2]
• MaxT-2 (119M): S1 hidden_size 96; num_blocks [2, 6, 14, 2]
• MaxT-3 (212M): S1 hidden_size 128; num_blocks [2, 6, 14, 2]
• MaxT-4 (476M): S1 hidden_size 192; num_blocks [2, 6, 14, 2]

Dataset: DikeDataset

• Benign: 982 samples

• Malware: 8970 samples

  RGB image samples:

Benign	Malicious

Train result:

• Best accuracy: 99.6%

• Epochs: 200

• Time: 6 hours