基于ViT的轻量级恶意代码检测架构
Lightweight Malicious Code Detection Architecture Based on Vision Transformer
黄保华 1杨婵娟 1熊宇 2庞飔1
作者信息
- 1. 广西大学计算机与电子信息学院,南宁 530004
- 2. 武汉数字工程研究所,武汉 430070
- 折叠
摘要
随着信息社会的快速发展,恶意代码变体日益增多,给现有的检测方法带来了挑战.为了提高恶意代码变体的检测准确率和效率,文章提出一种新的混合架构FasterMalViT.该架构通过融合部分卷积结构改进ViT,显著提升其在恶意代码检测领域的性能.为了解决引入卷积操作导致参数量增加的问题,文章采用可分离自注意力机制替代传统的多头注意力,有效减少了参数量,降低了计算成本.针对恶意代码数据集中各类样本分布不均衡的问题,文章引入类别平衡焦点损失函数,引导模型在训练过程中更关注样本数量较少的类别,从而提高难分类别的性能.在Microsoft BIG、Malimg数据集和MalwareBazaar数据集上的实验结果表明,FasterMalViT具有较好的检测性能和泛化能力.
Abstract
With the rapid development of the information society,the number of malware variants is increasing,posing challenges to existing detection methods.To improve the accuracy and efficiency of detecting malware variants,this paper proposed a new hybrid architecture called FasterMalViT.This architecture enhanced the Vision Transformer(ViT)by integrating partial convolutional structures,significantly improving its performance in malware detection.To address the issue of increased parameter count due to the introduction of convolutional operations,the paper employed a separable self-attention mechanism instead of traditional multi-head attention,effectively reducing the number of parameters and computational cost.To tackle the problem of imbalanced sample distribution in malware datasets,the paper introduced a class-balanced focal loss function,guiding the model to pay more attention to categories with fewer samples during training,thus improving performance on hard-to-classify categories.Experimental results on the Microsoft BIG,Malimg,and MalwareBazaar datasets demonstrate that FasterMalViT exhibits good detection performance and generalization ability.
关键词
恶意代码/ViT/部分卷积/可分离自注意力Key words
malicious code/ViT/partial convolution/separable self-attention引用本文复制引用
基金项目
国家自然科学基金(61962005)
中国高校产学研创新基金-新一代信息技术创新项目(2021ITA11003)
出版年
2024
NETL
NSTL
万方数据