[论文解读] MalBERT: Using Transformers for Cybersecurity and Malicious Software Detection

Abir Rahali, Moulay A. Akhloufi|arXiv (Cornell University)|Mar 5, 2021

Advanced Malware Detection Techniques参考文献 31被引用 23

一句话总结

本文提出 MalBERT，一种基于 Transformer 的新型模型，将 Android 应用源代码视为文本以实现静态恶意软件检测。通过在预处理的代码特征（包括权限、意图和活动）上微调 BERT，该模型在二分类任务中达到 97.61% 的准确率，在多分类恶意软件分类任务中达到 91.02%，优于 LSTM 和其他 Transformer 基线模型。

ABSTRACT

In recent years we have witnessed an increase in cyber threats and malicious software attacks on different platforms with important consequences to persons and businesses. It has become critical to find automated machine learning techniques to proactively defend against malware. Transformers, a category of attention-based deep learning techniques, have recently shown impressive results in solving different tasks mainly related to the field of Natural Language Processing (NLP). In this paper, we propose the use of a Transformers' architecture to automatically detect malicious software. We propose a model based on BERT (Bidirectional Encoder Representations from Transformers) which performs a static analysis on the source code of Android applications using preprocessed features to characterize existing malware and classify it into different representative malware categories. The obtained results are promising and show the high performance obtained by Transformer-based models for malicious software detection.

研究动机与目标

• 应对传统检测方法难以规避的复杂恶意软件日益增长的威胁。
• 探索基于 Transformer 的模型（尤其是 BERT）在 Android 恶意软件静态分析中的潜力。
• 通过将源代码建模为文本标记序列（如权限、意图、活动）来开发一种新型特征表示方法，以实现 NLP 风格的分类。
• 证明预训练语言模型在恶意软件检测任务中可优于 RNN 模型（如 LSTM）。
• 发布一个包含 11 个类别的预处理 Android 恶意软件数据集，以支持未来研究。

提出的方法

• 预处理 Android 应用源代码，提取并表示关键特征（如权限、意图、活动和 API 调用）为文本序列。
• 将预处理后的代码特征视为自然语言输入，并对预训练的 BERT 模型进行微调，以实现二分类（恶意/良性）和多分类（11 种恶意软件类别）任务。
• 利用 Transformer 架构中的注意力机制捕捉代码结构中的长距离依赖关系和关键模式。
• 通过在包含 11 种恶意软件类别的公开 Android 恶意软件数据集（Androzoo）上微调 BERT 模型，应用迁移学习。
• 使用保留的测试集上的准确率、马修斯相关系数（MCC）、F1 分数和交叉熵损失来评估性能。
• 将 BERT 与基线模型（包括 LSTM、XLNet、RoBERTa 和 DistilBERT）进行比较，以评估相对性能。

实验结果

研究问题

• RQ1 预训练的 Transformer 模型（如 BERT）是否能有效适应为以源代码作为文本输入的静态恶意软件检测？
• RQ2 基于 BERT 的恶意软件检测在性能上与传统的 RNN 模型（如 LSTM）及其他预训练 NLP 模型相比如何？
• RQ3 不同预训练语言模型（如 XLNet、RoBERTa、DistilBERT）对恶意软件分类准确率和泛化能力有何影响？
• RQ4 将 Android 源代码建模为标记序列（如权限、意图）是否能实现有效的表示学习以用于恶意软件分类？
• RQ5 微调后的 BERT 是否能在真实世界的 Android 恶意软件数据集上，在二分类和多分类恶意软件检测任务中均实现高性能？

主要发现

• BERT 在二分类任务中达到最高的准确率 97.61%，显著优于 LSTM（94.05%）、XLNet（95.79%）、RoBERTa（95.33%）和 DistilBERT（95.42%）。
• 在 11 个类别的多分类恶意软件分类任务中，BERT 达到 91.02% 的准确率，优于所有基线模型，包括 LSTM（85.07%）和 DistilBERT（59.81%）。
• BERT 在二分类任务中还取得了最高的马修斯相关系数（MCC）0.9559，表明其在类别不平衡数据上表现优异。
• 模型仅经过两轮微调后损失开始上升，表明在最优初始学习率下收敛效率高。
• 在二分类任务中，BERT 的交叉熵损失最小（0.1274），表明其模型校准性能优于其他模型。
• 本研究证明，源代码可被有效表示为文本标记序列，以支持基于 NLP 的恶意软件检测，从而实现高性能分类。