Triplet-trained graph transformer with control flow graph for few-shot malware classification

Abstract

The exponential proliferation of malware requires robust detection mechanisms for the security of global enterprises and national infrastructures. Conventional malware classification methods primarily depend on extensive datasets of curated malware samples, rendering them suboptimal for detecting novel strains exploiting contemporary vulnerabilities. In this paper, we reformulate malware detection as a few-shot learning task, and propose a new distance-based classification method that harnesses the innate functional attributes of malware to mitigate the dependency on sample volume. A disentangled representation of the malware's control flow graph is exploited, and a specialized transformer architecture is trained with a triplet-loss function, aiming to finetune the representation of malicious attributes. An attention mechanism of the transformer judiciously discerns functional signatures from intricate control flow graphs. Empirical evaluations on real-world malware datasets underscore the efficacy of the proposed method, achieving an outstanding recall rate of 83.37% with mere 2,000 training samples. As a result, our method outperforms the state-of-the-art methods with an accuracy of 99.45% and a recall of 97.89%.