Federated Learning for Malware Image Classification under Data Heterogeneity

Victor Taiwo, Cemal Nisan, Muhammad Athallah, M. Emre Gursoy, Oznur Ozkasap

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,…