[论文解读] An Ensemble of Pre-trained Transformer Models For Imbalanced Multiclass Malware Classification

Ferhat Demirkıran, Aykut Çayır+2 more|arXiv (Cornell University)|Dec 25, 2021

Advanced Malware Detection Techniques被引用 4

一句话总结

该论文提出了一种基于袋装法的随机变换器森林（RTF）集成模型，结合了预训练的 BERT 和 CANINE 变换器架构，通过使用 API 调用序列，在高度不平衡的数据集上提升了多类恶意软件分类的性能。RTF 模型在 Catak 数据集上实现了最先进的 F1 分数 0.6149，通过注意力机制和集成学习捕捉恶意软件行为中的复杂序列关系，优于单个模型和先前方法。

ABSTRACT

Classification of malware families is crucial for a comprehensive understanding of how they can infect devices, computers, or systems. Thus, malware identification enables security researchers and incident responders to take precautions against malware and accelerate mitigation. API call sequences made by malware are widely utilized features by machine and deep learning models for malware classification as these sequences represent the behavior of malware. However, traditional machine and deep learning models remain incapable of capturing sequence relationships between API calls. On the other hand, the transformer-based models process sequences as a whole and learn relationships between API calls due to multi-head attention mechanisms and positional embeddings. Our experiments demonstrate that the transformer model with one transformer block layer surpassed the widely used base architecture, LSTM. Moreover, BERT or CANINE, pre-trained transformer models, outperformed in classifying highly imbalanced malware families according to evaluation metrics, F1-score, and AUC score. Furthermore, the proposed bagging-based random transformer forest (RTF), an ensemble of BERT or CANINE, has reached the state-of-the-art evaluation scores on three out of four datasets, particularly state-of-the-art F1-score of 0.6149 on one of the commonly used benchmark dataset.

研究动机与目标

• 解决利用 API 调用序列的行为特征进行高度不平衡多类恶意软件分类的挑战。
• 评估预训练变换器模型（如 BERT 和 CANINE）在捕捉 API 调用序列中长距离依赖关系以实现恶意软件家族识别方面的有效性。
• 开发一种集成模型，以提升在不平衡数据集中稀有恶意软件家族上的泛化能力和性能。
• 对比基于变换器的模型与传统 LSTM 和单层变换器模型在推理和训练效率方面的表现。
• 展示 CANINE 在恶意软件分类中的首次应用，并通过集成学习验证其性能。

提出的方法

• 该方法采用一种基于袋装法的集成框架，称为随机变换器森林（RTF），其中多个预训练 BERT 或 CANINE 模型在训练数据的随机子集上进行训练。
• 每个 RTF 模型使用无需分词的输入表示方式来表示 API 调用序列，保留序列顺序，使注意力机制能够建模调用之间的依赖关系。
• 通过多数投票方式聚合预测结果，以提高鲁棒性并降低方差，尤其在少数类恶意软件家族上表现更优。
• 应用了预处理步骤，包括序列长度归一化和数据增强，以提升在包含动态 API 调用的 Catak 数据集上的性能。
• 评估使用 F1 分数和 AUC 作为主要指标，以应对类别不平衡问题，同时测量训练和推理时间以进行效率分析。
• 模型利用 BERT 和 CANINE 中固有的多头注意力和位置嵌入，更有效地建模 API 调用序列中的长距离关系，优于 RNN 或单层变换器。

Figure 1 : Transformer model architecture.

实验结果

研究问题

• RQ1 RQ.1：在高度不平衡的多类恶意软件分类中，合适的分类指标是什么？F1 分数和 AUC 在评估模型性能方面如何比较？
• RQ2 RQ.2：与传统模型（如 LSTM）和单层变换器相比，预训练变换器模型（如 BERT 和 CANINE）在不平衡恶意软件数据集上的表现如何？
• RQ3 RQ.3：预训练变换器的集成（如随机变换器森林，RTF）是否能在基准恶意软件数据集中实现最先进性能？
• RQ4 RQ.4：在使用 BERT 和 CANINE 等大型预训练模型时，推理速度与分类准确率之间的权衡如何？
• RQ5 RQ.5：预处理技术如何影响变换器模型在动态 API 调用序列上的性能，特别是在 Catak 数据集上？

主要发现

• RTF 集成模型在 Catak 数据集上实现了最先进的 F1 分数 0.6149，优于所有基线模型和先前方法。
• BERT 和 CANINE 预训练模型在所有四个数据集上的 F1 分数和 AUC 表现显著优于单层变换器和 LSTM 模型，尤其在稀有恶意软件家族上表现更优。
• 首次成功将 CANINE 模型应用于恶意软件分类，并在 VirusShare 和 VirusSample 数据集上实现了最先进结果，证明其在静态 API 调用序列上的有效性。
• BERT 和 CANINE 的训练时间显著长于 LSTM 和单层变换器，BERT 在 Oliveira 数据集上耗时 141.02 分钟，但由于模型为预训练，不影响实时响应。
• RTF 的平均推理时间约为每样本 4.5 秒，对杀毒扫描器而言可接受，且尽管使用了大型预训练模型，仍保持在合理范围内。
• RTF 模型在四个数据集中的三个实现了最先进性能，证实了使用预训练变换器进行集成学习在不平衡恶意软件分类中的有效性。

Figure 2 : The outputs of pre-processing steps.