🇧🇷 Português

GPT-2 Fine-Tuning for Malware Classification

Fine-tuning GPT-2 for malware classification using raw assembly (.asm) files. This repository contains the full pipeline and experiments conducted during my Master's research at IME (Instituto Militar de Engenharia).

---

Abstract

Malware classification and identification are critical tasks in cybersecurity. This work proposes a method based on fine-tuning GPT-2, a transformer-based language model, to classify malware directly from assembly code. The assembly files are preprocessed to extract only the pure code segment, tokenized, and fed into GPT-2 for sequence classification. Experiments on two datasets - Microsoft BIG 2015 (9-class classification) and ESET2022 (binary identification) - achieve results near the state of the art, with 98.96% accuracy on BIG 2015 and 96.18% accuracy on ESET2022.

Academic Reference

• Master's Thesis: Dissertação.pdf (Portuguese)
• Published Paper: Malware Classification Using the GPT-2 Model - Paper.pdf
• LinkedIn Article: Malware Classification using the GPT-2 Model
• Institution: IME - Instituto Militar de Engenharia

Key Features

• Full pipeline: preprocessing → tokenization/chunking → GPT-2 fine-tuning → evaluation
• Multi-dataset support: Microsoft BIG 2015 (9 classes) and ESET2022 (binary)
• 10-fold cross-validation for robust evaluation
• Binary and multi-class classification modes
• CSV-based dataset caching for fast reload across experiments
• WhatsApp notifications during long training runs (opt-in via env vars)
• TF-IDF baseline archived separately from the main pipeline

Architecture

Results

Microsoft BIG 2015 - Classification (9 classes)

Epochs	Accuracy (%)	F1 Macro (%)	Training Time
1	98.62 ± 0.50	96.87 ± 0.76	135.29 h
2	98.82 ± 0.33	97.35 ± 0.81	242.43 h
3	98.96 ± 0.27	97.79 ± 1.04	457.09 h

ESET2022 - Binary Identification

Epochs	Accuracy (%)	F1 Macro (%)	Training Time
1	95.63 ± 0.56	95.40 ± 0.59	7.19 h
2	95.85 ± 0.43	95.65 ± 0.45	14.92 h
3	96.01 ± 0.46	95.82 ± 0.48	23.41 h
4	96.18 ± 0.47	96.00 ± 0.50	29.74 h

Results reported as mean ± standard deviation across 10 folds.

Datasets

Microsoft BIG 2015

• Source: Microsoft Malware Classification Challenge (BIG 2015)
• Size: ~20,000 .asm files
• Classes: 9 malware families (Ramnit, Lollipop, Kelihos_ver3, Vundo, Simda, Tracur, Kelihos_ver1, Obfuscator.ACY, Gatak)
• Task: Multi-class classification

ESET2022

• Source: Private dataset provided by ESET
• Size: ~1,000 .asm files
• Classes: 2 (malware / benign)
• Task: Binary identification

Quick Start

1. Setup

git clone https://github.com/matheusvanzan/gpt-malware-class.git
cd gpt-malware-class
python -m venv venv
source venv/bin/activate   # Linux
venv\Scripts\activate      # Windows
pip install -r requirements.txt

2. Preprocessing

Process the raw .asm files:

python main.py --process

This filters code segments, removes noise, and saves to proc-1/.

3. Dataset Cache

Tokenize and cache the dataset as CSV:

python main.py --cache --fold=1 --version=3

4. Training

# Binary mode (ESET2022)
python main.py --train --model=gpt2 --limit=all --fold=1 --epochs=4

# Multi-class mode (BIG 2015)
python main.py --train --multi --model=gpt2 --limit=all --fold=1 --epochs=3

Run all 10 folds with the provided script:

# Linux
bash src/scripts/run.sh

# Windows
src/scripts/run.bat

5. Evaluation

python main.py --test --multi --fold=1
python main.py --metrics --multi

Project Structure

├── main.py              # CLI entry point (thin wrapper)
├── src/                 # Source code
│   ├── cli.py           # CLI argument parsing and orchestration
│   ├── dataset.py       # Tokenization, chunking, CSV caching
│   ├── model.py         # GPT-2 fine-tuning (train / test / metrics)
│   ├── processor.py     # Assembly preprocessing & cleaning
│   ├── settings.py      # Configuration (paths, hyperparameters)
│   ├── metrics.py       # Accuracy, F1, confusion matrix
│   ├── utils/
│   │   └── whats.py     # WhatsApp notification helper
│   ├── tests/
│   │   └── test_settings.py
│   ├── scripts/         # Experiment automation scripts
│   │   ├── run.bat
│   │   ├── run.sh
│   │   ├── dataset_cache.bat
│   │   ├── tensorboard.bat
│   │   ├── zip_partials.py
│   │   └── dataset_stats.py
│   └── __init__.py
├── papers/              # Academic publications
│   ├── Paper.pdf        # Published paper
│   └── Dissertação.pdf  # Master's thesis
├── assets/              # Images and diagrams
├── requirements.txt
├── pyproject.toml
├── ruff.toml
├── .gitignore
└── LICENSE

Configuration

Set the data path via environment variable:

export GPT_MALWARE_PATH=/path/to/gpt-malware   # Linux
set GPT_MALWARE_PATH=D:\path\to\gpt-malware    # Windows

Or edit src/settings.py to set the default paths.

WhatsApp notifications are disabled by default. To enable:

export WHATSAPP_PHONE=+5511999999999
export WHATSAPP_API_KEY=your_api_key

Citation

If you use this work in your research, please cite:

@mastersthesis{vanzan2024,
    author  = {Matheus Vanzan},
    title   = {Malware Classification Using the GPT-2 Model},
    school  = {Instituto Militar de Engenharia (IME)},
    year    = {2024},
}

License

This project is licensed under the MIT License - see LICENSE.

Acknowledgments

• IME (Instituto Militar de Engenharia) - Master's program support
• ESET - Dataset and collaboration
• HuggingFace - Transformers library