Margin-Based Supervised Contrastive Learning for Robust Malware Detection

Yawar Abbas Abid, Jinsong Wu, Said Jadid Abdulkadir, Ayed Alwadain, Muhammad Waqas, Muhammad Farhan

Abstract: The growth of network technologies has led to a rise in network traffic diversity, which has resulted in increased complexity within each class. However, distinguishing between malicious and benign actions has become more challenging as the two have started to overlap. Consequently, false alarms have become a common issue, which can cause problems in detection models. We proposed a novel approach to enhance intrusion detection based on supervised contrastive learning that combines a maximum-margin-based approach for malware detection in an end-to-end process. This approach differs from traditional methods, where feature selection and detection are considered separate stages. Our model guarantees robustness and correlation between these two processes. To handle the challenges posed by Variation within a class and similarities between different classes, we use the margin-based supervised contrastive learning (MbSCL) technique. This technique involves a two-stage learning process that enables the network to learn embeddings that bring closer the similar data points and separate dissimilar ones. We employ a specialized loss function that combines supervised contrastive loss and marginal loss to improve the accuracy of the classification. This process balances the effectiveness of embedding separation and classification to improve overall performance. Mathematical theorems and proofs support the optimization process, which provides insights into contrastive loss in similarity learning. This two-stage learning approach and loss function design are particularly useful for malware detection, as they enable the identification of subtle data patterns and anomalies in network traffic.

External IDs:dblp:journals/access/AbidWAAWF25