Federated learning improves malware detection with diverse client data
Federated Learning for Malware Image Classification under Data Heterogeneity
Cryptography and SecurityDistributed, Parallel, and Cluster ComputingMachine Learning
Summary
Malware can be detected by turning its code into pictures and teaching computers to recognize patterns. Sometimes many computers work together to train these detection programs without sharing sensitive data, called federated learning. The authors found that if the data on each computer is too different, the detection gets worse. They discovered a way to pick which computers should join each training round based on how varied their data is, which helps the program learn better. This makes collaborative malware detection more accurate and stable.
What this means in practice
- •For cybersecurity teams: Build federated malware detection models that maintain accuracy when client data varies widely.
- •For federated learning engineers: Improve model training by selecting client subsets with diverse malware classes to stabilize training outcomes.
Authors
Victor Taiwo, Cemal Nisan, Muhammad Athallah, M. Emre Gursoy, Oznur Ozkasap
Abstract
Malware image classification (MIC), in which binary files are converted into visual representations and deep neural networks are trained to identify malware families, has recently emerged as an effective approach for malware detection. Federated learning (FL) enables malware image classifiers to be trained collaboratively by multiple clients without sharing raw data. However, a fundamental challenge in FL is data heterogeneity (non-IID), which is particularly relevant but understudied in the malware context. In this paper, we address this gap via three contributions. First, we present an empirical study of federated MIC under varying degrees of data heterogeneity across three popular malware image datasets (MalImg, Virus-MNIST, and MalNet). We show that while FL achieves strong accuracy under IID conditions, both accuracy and training stability degrade as data becomes increasingly non-IID. Second, we introduce a set of round-level statistics that characterize the data contributed by participating clients in each FL round, and demonstrate that class entropy (diversity of class labels present in a given round) exhibits strong positive correlation with global model accuracy. Third, leveraging this insight, we propose an entropy-aware client selection strategy that filters candidate client subsets based on class entropy, and show that our strategy outperforms standard random selection. Our results highlight the role of client composition in federated MIC and underscore the importance of distribution-aware classifier training under data heterogeneity.