Applying Word Embeddings and Graph Neural Networks for Effective Malware Classification

The significance of word embeddings in natural language processing for capturing semantic relationships between words is widely acknowledged. This study aims to explore the efficacy of word embedding techniques in classifying malware. Specifically, we evaluate the effectiveness of applying Graph Neural Networks (GNNs) to weighted graphs formed from word embeddings generated by analyzing opcode sequences in malware files. In the initial experiments, we employ the Graph Convolution Network (GCN) on weighted graphs generated using different word embedding techniques, including Bag-of-words, TF-IDF, and Word2Vec. The results indicate that Word2Vec provides the most effective word embeddings, serving as the baseline for comparison with three GNN models, namely Graph Convolution Network, Graph Attention Network (GAT), and GraphSAGE Network. Subsequently, we conduct further experiments, generating vector embeddings of varying lengths using Word2Vec, and utilizing these embeddings as node features for constructing weighted graphs. Through performance comparison of the GNN models, we demonstrate that larger vector embeddings significantly enhance the models’ ability to classify malware files into their respective families. Furthermore, we compare the result achieved using Word2Vec embeddings against those obtained through contextualized embeddings from BERT. Overall, our experiments show the potential of word embeddings as node features for GNN classification, with an increase in accuracy from 71.6 to 91.91% when Word2Vec embeddings were used in combination with GCN.