MalImgDA: Diffusion-based Data Augmentation for Long-tailed Malware Family Classification

Abstract

With the rapid improvement of machine learning technology, leveraging machine learning methods for malware classification has emerged as a viable approach. However, under real-world circumstance, the imbalanced or long-tailed distribution among various malware families, poses a critical challenge to classify such few-shot malware families, eliminating the effectiveness of trained classifier. In this paper, we propose MalImgDA, a novel data augmentation framework that leverages diffusion model-based approach to tackle the long-tailed malware family classification problem. By fine-tuning the pre-trained diffusion model on few-shot data, we synthesize analogous malware images from existing samples with high resemblance to target family. And then we can mingle synthetic samples with existing data to build a re-balanced dataset for classifier training. Particularly, by utilizing MalImgDA, we can substantially enhance the diversity of data and generate plausible malware variants proactively while persevering the characteristics of target family. Experiments conducted on two publicly available datasets demonstrate the effectiveness of our proposed method in comparison to other commonly-used approaches.