AGProto: Adaptive Graph ProtoNet towards Sample Adaption for Few-Shot Malware Classification

Abstract

Traditional malware-classification methods reliant on large pre-labeled datasets falter when encountering new or evolving malware types, particularly when only a few samples are available. And most current models utilize a fixed architecture; however, the characteristics of the various types of malware differ significantly. This discrepancy results in notably inferior classification performance for certain categories or samples with uncommon features, but the threats of these malware samples are of equivalent significance. In this paper, we introduce Adaptive Graph ProtoNet (AGProto), a novel approach for classifying malware in the field of Few-Shot Learning. AGProto leverages Graph Neural Networks (GNNs) to propagate sample features and generate multiple prototypes. It employs an attention mechanism to calculate the relevance of each prototype to individual samples, resulting in a customized prototype for each case. Our approach achieved optimal performance on two few-shot malware classification datasets, surpassing other competitive models with an accuracy improvement of over 2%. In extremely challenging scenarios—specifically, 20-class classification tasks with only five samples per class—our method notably excelled, achieving over 70% accuracy, significantly outperforming existing advanced techniques.