francesco-univr/malware-visual-analysis

# 恶意软件可视化分析 在 Malimg 数据集上的家族分类与 Zero-Day 检测 计算机视觉与深度学习 — 维罗纳大学，人工智能硕士。 作者：Francesco Simbola (VR545665) Nataraj 等人 (2011) 将恶意软件二进制文件读取为一维字节流，并将其重塑 为灰度图像。本项目在该想法的基础上构建了两个 pipeline，均在 Malimg 数据集上运行： - 任务 1 — 家族分类（闭集）：将恶意软件图像分配到 25 个家族中的一个。 - 任务 2 — Zero-Day 检测（开集）：使用基于重构的异常检测，标记训练期间从未见过的家族的样本。 ## 结果 ### 任务 1 — 家族分类（测试集，1,267 个样本） | 模型 | Accuracy | Macro F1 | F1 Autorun.K | F1 Yuner.A | |-------|----------|----------|--------------|------------| | ResNet-18 (scratch) | 97.63% | ~0.926 | 0.000 | 0.741 | | ResNet-18 (ImageNet-1K) | 98.50% | ~0.941 | 0.000 | 0.741 | | ViT-B/16 (SWAG) | 98.18% | 0.961 | 1.000 | 1.000 | 在总体准确率方面，这三个模型处于彼此的抽样误差范围内（精确的 McNemar 检验，预训练的 ResNet vs ViT：*p* = 0.62）。整个 macro-F1 的差距仅来自一对 样本：**Autorun.K / Yuner.A**：ResNet 变体无法区分它们，而 ViT 将其区分开来。结果 表明这对样本在像素级别上具有 100% 的训练/测试重叠（参见*已知局限性*），因此 ViT 在这里的分数反映的是对已见像素的记忆，而不是对未见像素的泛化能力。 ### 任务 2 — Zero-Day 检测（149 个 zero-day vs 1,118 个已知） | 模型 | AUROC | AUPRC | F1 (calibrated) | |-------|-------|-------|-----------------| | ConvAutoencoder (AE) | 0.648 | 0.253 | 0.097 | | ConvVAE | **0.732** | **0.502** | **0.550** | Monte Carlo 样本计数消融实验（N ∈ {1, 2, 5, 10}，AUROC 波动为 0.0031）表明，VAE 相对于 AE 的提升 来自于变分公式，而不是 MC 平均。 每个家族的 AUROC 分为两种情况。VAE 能够清晰地区分 **Fakerean** (0.952) 和 **VB.AT** (0.928)，但在 **Rbot!gen** (0.438)、**Yuner.A** (0.402) 和 **Autorun.K** (0.388) 上表现极差，全部低于 0.5 的随机基线。原因是表征 本身：它将这三个家族映射到了已知家族所占据的相同区域，因此没有任何 重构分数能将它们区分开来。 ## 项目结构 ``` . ├── configs/ │ └── config.yaml # Hyperparameters and paths ├── src/ │ ├── dataset.py # Dataloaders, transforms, weighted sampling │ ├── models/ │ │ ├── resnet_classifier.py # ResNet-18 head for 25 classes │ │ ├── vit_classifier.py # ViT-B/16 SWAG head │ │ ├── autoencoder.py # Convolutional AE + shared encoder/decoder │ │ └── vae.py # Convolutional VAE + ELBO loss │ ├── prepare_dataset.py # MD5 deduplication + stratified split │ ├── train_classifier.py # Task 1 training │ ├── train_vae.py # Task 2 training (AE / VAE) │ ├── evaluate_classifier.py # Task 1 evaluation + confusion matrices │ ├── evaluate_zeroday.py # Task 2 evaluation (AUROC, AUPRC, thresholds) │ ├── analyze_zeroday_per_family.py # Per-family AUROC decomposition │ ├── ablate_mc_samples.py # Monte Carlo sample-count ablation │ ├── significance_test.py # Exact McNemar test between classifiers │ ├── visualize.py # Grad-CAM, t-SNE, VAE reconstructions │ └── utils.py # Seeding, checkpointing, optimizer, early stopping ├── tests/ │ ├── smoke_test.py # Structural validation of all models │ ├── check_duplicates.py # Byte-level (MD5) duplicate analysis │ └── check_pixel_duplicates.py # Pixel-level leakage analysis ├── main.pdf # Full project report └── requirements.txt ``` ## 安装说明 需要 Python 3.10+ 以及支持 CUDA 的 GPU 来进行训练（CPU 可用于推理）。 ``` pip install -r requirements.txt ``` ## 数据集 由于体积（约 1.2 GB）和许可协议的限制，Malimg 数据集（Nataraj 等人，2011）未包含在此 仓库中。它可以在 Kaggle 上公开获取： 这个镜像文件已经整理成了 `train/`、`val/` 和 `test/` 子文件夹，这正是 `prepare_dataset.py` 期望的输入布局。下载后，请复现本工作中使用的精确的去重后的 70/15/15 划分： ``` python src/prepare_dataset.py --data_root --output_dir clean_dataset ``` 这会通过 MD5 对每个样本的原始 PNG 字节进行哈希处理，移除了 895 个完全相同的字节 重复项（9,339 → 8,444 个样本），并将全新的、无泄漏的分层划分写入 `clean_dataset/`（去重和重新划分是在合并原始文件夹*之后*运行的， 因此没有文件会出现在多个划分中）。五个 zero-day 家族（VB.AT、Fakerean、 Rbot!gen、Autorun.K、Yuner.A）被排除在 autoencoder 的训练集之外，并且配置在 `configs/config.yaml` 中。 ## 复现结果 训练会将 checkpoint 写入 `checkpoints/`，并将图表/指标写入 `outputs/`。 ### 任务 1 — 分类器 ``` python src/train_classifier.py --model resnet18 # from scratch python src/train_classifier.py --model resnet18 --pretrained # ImageNet-1K python src/train_classifier.py --model vit --pretrained # ViT-B/16 SWAG (384x384) python src/evaluate_classifier.py --model resnet18 --checkpoint checkpoints/resnet18_best.pth python src/evaluate_classifier.py --model resnet18 --pretrained --checkpoint checkpoints/resnet18_pretrained_best.pth python src/evaluate_classifier.py --model vit --pretrained --checkpoint checkpoints/vit_pretrained_best.pth --image_size 384 ``` ### 任务 2 — autoencoder ``` python src/train_vae.py --model ae python src/train_vae.py --model vae python src/evaluate_zeroday.py --model ae --checkpoint checkpoints/ae_latent128_best.pth python src/evaluate_zeroday.py --model vae --checkpoint checkpoints/vae_latent128_beta1.0_best.pth python src/analyze_zeroday_per_family.py --model vae --checkpoint checkpoints/vae_latent128_beta1.0_best.pth ``` ### 分析与图表 ``` python src/ablate_mc_samples.py # Monte Carlo ablation python src/significance_test.py # McNemar ResNet vs ViT python src/visualize.py --task gradcam --model resnet18 --checkpoint checkpoints/resnet18_pretrained_best.pth python src/visualize.py --task tsne --model vit --checkpoint checkpoints/vit_pretrained_best.pth --image_size 384 python src/visualize.py --task reconstructions --vae_model vae --vae_checkpoint checkpoints/vae_latent128_beta1.0_best.pth ``` ### 快速验证（无需训练，无需 checkpoint；需要 `clean_dataset/`） ``` python tests/smoke_test.py # builds every model, checks shapes and losses python tests/check_pixel_duplicates.py # reproduces the 425 / 47 / 71 leakage figures ``` ## 关键设计选择 - **类别不平衡**（高达 37:1）由 `WeightedRandomSampler` 处理，而不是通过 加权损失，从而避免了双重补偿。 - **不使用几何增强**：翻转会反转字节顺序，而旋转会破坏 定义每个家族视觉特征的内存线性结构。 - **AE/VAE pipeline 不使用 padding**：常量 padding 像素很容易被 重构，从而稀释异常分数。 - **AdamW**，权重衰减仅应用于 conv/linear 权重（bias 和 norm 参数 不进行衰减）。 - **VAE β-warmup**：β 在前 10 个 epoch 中从 0 逐渐增加到 1，以防止后验坍缩。 - **Monte Carlo 异常分数**（N = 10），遵循 An & Cho (2015)。 所有随机种子（Python、NumPy、PyTorch、CUDA）均固定为 42。 ## 已知局限性 MD5 去重在文件字节级别进行，而不是像素级别。事后 像素级别分析（`tests/check_pixel_duplicates.py`）发现了 425 个像素重复副本 以及 71 个测试样本（5.6%）在训练集中存在像素级孪生样本，这些集中在四个 家族上（Autorun.K 和 Yuner.A 为 100%，Fakerean 为 97%，Adialer.C 为 57%）。因此，任务 1 中关于 Autorun.K/Yuner.A 这一对样本的核心结果应被解读为对 近乎退化的碰撞的鲁棒性，而不是对未见样本的泛化能力。自然的下一步工作应当是进行像素级别的重新去重，随后进行重新训练。 ## 参考文献 - L. Nataraj 等人，*Malware Images: Visualization and Automatic Classification*，VizSec，2011。 - D. P. Kingma, M. Welling，*Auto-Encoding Variational Bayes*，ICLR，2014。 - J. An, S. Cho，*Variational Autoencoder based Anomaly Detection using Reconstruction Probability*，2015。 - A. Dosovitskiy 等人，*An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale*，ICLR，2021。 - M. Singh 等人，*Revisiting Weakly Supervised Pre-Training of Visual Perception Models* (SWAG)，CVPR，2022。 完整的讨论和方法论请见 [`main.pdf`](main.pdf)。

标签：DAST, ResNet, Vectored Exception Handling, ViT, 凭据扫描, 图像分类, 异常检测, 恶意软件分析, 深度学习, 计算机视觉, 逆向工具