[论文解读] Imbalanced Malware Images Classification: a CNN based Approach

Songqing Yue, Wang, Tianyang|arXiv (Cornell University)|Aug 27, 2017

Advanced Malware Detection Techniques参考文献 13被引用 108

一句话总结

本论文提出一种加权 softmax 损失以解决恶意软件图像分类中的类别不平衡问题，并通过微调 VGG-19 并在多种 CNN 上验证其有效性。

ABSTRACT

Deep convolutional neural networks (CNNs) can be applied to malware binary detection via image classification. The performance, however, is degraded due to the imbalance of malware families (classes). To mitigate this issue, we propose a simple yet effective weighted softmax loss which can be employed as the final layer of deep CNNs. The original softmax loss is weighted, and the weight value can be determined according to class size. A scaling parameter is also included in computing the weight. Proper selection of this parameter is studied and an empirical option is suggested. The weighted loss aims at alleviating the impact of data imbalance in an end-to-end learning fashion. To validate the efficacy, we deploy the proposed weighted loss in a pre-trained deep CNN model and fine-tune it to achieve promising results on malware images classification. Extensive experiments also demonstrate that the new loss function can well fit other typical CNNs, yielding an improved classification performance.

研究动机与目标

• 解决基于 CNN 的分类中恶意软件图像数据集高度不平衡的挑战。
• 引入一种加权 softmax 损失，其中类别权重取决于类别大小以平衡学习。
• 通过对预训练的 CNN（VGG-19）进行微调并在多个模型上验证来演示该方法。

提出的方法

• 定义一个加权 softmax 损失 J0 = -(1/m) sum_i sum_j w_k 1(y(i)=j) log p_j^(i)。
• 计算权重 w_k = 1 + (S_max - S_k) / (beta * S_max) 其中 beta 控制缩放（经验值 beta=20）。
• 对预训练的 VGG-19 模型进行微调，添加 dropout 和 batch normalization 层；将加权损失作为最后一层。
• 使用 25 类不平衡的恶意软件图像数据集进行 Top-1 验证错误的评估；在若干 VGG 变体上比较有无加权损失的效果。
• 在 Nvidia TITAN X 上使用 MatConvNet，最终架构为 60 层且包含新增层。

实验结果

研究问题

• RQ1 一个针对类别不平衡的损失是否能提升在不平衡的恶意软件图像数据集上进行的 CNN 性能？
• RQ2 提出的加权 softmax 损失是否在除所测试模型外的不同 CNN 架构中具有泛化性？
• RQ3 缩放参数 beta 对训练稳定性和准确性的影响？
• RQ4 用带有加权损失的预训练 CNN 进行微调与使用原始损失相比有何差异？
• RQ5 使用加权损失时，特征表示发生了哪些定性变化？

主要发现

• 加权 softmax 损失在多个模型上提升了 Top-1 准确率（例如 VGG-19 从 97.32% 提升到 98.63%）。
• 在 VGG 家族变体（VGG-19、VGG-F、VGG-M、VGG-S）中，加权损失始终优于原始损失的测试准确度。
• 对 VGG-19 进行加权损失微调可获得更高的验证准确性和更稳定的训练曲线。
• Beta = 20 提供比其他测试值更平滑的收敛性和更好的最终误差。
• 特征图在加权损失下呈现类别特异的判别性表征。