Federated Learning for Malware Image Classification under Data Heterogeneity

Federated Learning for Malware Image Classification under Data Heterogeneity Victor Taiwo1 , Cemal Nişan1 , Muhammad Athallah2 , M. Emre Gürsoy1 , and Öznur Özkasap1 1 Department of Computer Engineering, Koç University, Istanbul, Türkiye 2 Informatics Institute, Istanbul Technical University, Istanbul, Türkiye

arXiv:2610.25264v1 [cs.CR] 18 Oct 2026

{vtaiwo25, cnisan21, emregursoy, oozkasap}@ku.edu.tr, athallah24@itu.edu.tr

Keywords:

Federated learning, malware, malware image classification, data heterogeneity, entropy.

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 (nonIID), 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.

1

Introduction

Malware continues to pose a major threat to modern computing systems, with attacks growing in scale, diversity, and sophistication. Deep learning-based approaches are actively being developed for malware detection and classification, particularly using imagebased representations of binaries (Gibert et al., 2020; Freitas et al., 2022; Brosolo et al., 2024). In this line of work, often termed malware image classification (MIC), binary files are transformed into visual representations (images) and vision models (e.g., convolutional neural networks) are trained to identify structural patterns associated with different malware families (Nataraj et al., 2011; Vasan et al., 2020a; Vasan et al., 2020b; Kumar and Janet, 2022; Tekerek and Yapici, 2022). Despite their effectiveness, conventional deep learning approaches require aggregating data from multiple sources into a single location. In the context of malware classification, such data collection raises significant privacy, security, and regulatory concerns,

since the raw data of a client may reveal sensitive information about installed software, system vulnerabilities, and malware infections. Here, federated learning (FL) (McMahan et al., 2017) provides a compelling solution by enabling clients to collaboratively train a malware classifier without sharing their data. In FL, clients locally train models using their own data and transmit only the updated model weights to the FL server that coordinates the training process, aggregates local updates and computes the weights of the global model. Indeed, successful approaches applying FL to malware image classification have recently been proposed (Abdel-Basset et al., 2022; Ambekar et al., 2025). However, a fundamental challenge in FL is data heterogeneity, where client data is non-independent and identically distributed (non-IID). This issue is particularly pronounced in malware-related applications, as different client devices (e.g., personal devices running different operating systems, enterprise systems, IoT devices) are exposed to distinct threats and malware, resulting in highly imbalanced data and

skewed class distributions. While non-IID data is a known challenge in FL (Li et al., 2020b; Li et al., 2022; Karimireddy et al., 2020; McMahan et al., 2017), there is limited coverage or consideration of it in the MIC context. To address this gap, in this paper, we first present an extensive empirical study of federated MIC under varying degrees of data heterogeneity (non-IID) and datasets. Using three popular malware image datasets (MalImg, Virus-MNIST, and MalNet), we show that FL can achieve strong performance under IID conditions, but its accuracy and stability degrade as data becomes increasingly non-IID. Second, we introduce a set of round-level statistics (metrics) to characterize the data contributed by participating clients in each FL round, and demonstrate that these statistics (in particular, class entropy) are correlated with model accuracy. Third, we leverage this insight to design an entropy-aware client selection strategy that improves the accuracy of federated malware image classifiers over standard random selection. Overall, our methods and results provide new insights into the role of client composition in FL of malware image classifiers, and highlight the importance of distribution-aware training strategies under data heterogeneity. The rest of this paper is organized as follows. We survey related work in Section 2. Preliminaries for MIC, FL, and our problem statement are provided in Section 3. Our methods, including the proposed round-level statistics and entropy-aware client selection strategy, are described in Section 4. Experiment results and related analyses are presented in Section 5, followed by conclusions in Section 6.

2

Related Work

In this section, we review FL approaches for malware detection and classification. We start with a broader scope, including different data modalities (not only image-based methods), and then focus on malware images, which are most directly relevant to our work. Several FL methods have been proposed in the literature for malware detection and classification. An FL-based framework for detecting infected Internet of Things (IoT) devices through network traffic, employing the federated averaging (FedAvg) algorithm to aggregate local model updates, was proposed in (Rey et al., 2022). In (Amjath et al., 2024), application structure and behavior were modeled using function call graphs, and malware detection was performed using two FL-based graph neural network models: Fed-MalGCN and Fed-MalGAT. In (D’Angelo et al., 2023), a Markov chain and association rule-based

classifier was adapted to the FL setting to classify IoT malware. In (Nobakht et al., 2024), a lightweight one-dimensional CNN with optimized hyperparameters was proposed for IoT malware detection. A concept drift-resilient FL framework named FL-MalDrift was proposed in (Patel et al., 2025), which integrates lightweight on-device drift detection. Ransomware detection using FL, which is a specific type of malware, was the focus of studies in (Vehabovic et al., 2023) and (Lan et al., 2025). Another recent work studied the robustness aspect and highlighted that FLbased malware detectors remain vulnerable to label flipping attacks (Eslamnejad et al., 2025). Two attack strategies and two density-based defense mechanisms were introduced to counter label flipping attacks. A few works within the domain of malware detection explored the challenges posed by non-IID data. FedHGCDroid (Jiang et al., 2022) is an FL framework for Android malware classification. It combines HGCDroid, a malware classification model, with FedAdapt, a contribution degree-based adaptive training mechanism that employs attention and metalearning to improve model resilience to data heterogeneity. In addition, FedDRC (Jiang et al., 2024) was proposed to enhance robustness against spatiotemporal heterogeneity of Android malware. This framework pairs an API encoding mechanism based on functional semantics with an information bottleneckbased local model. The works discussed above do not use malware images, but rather the likes of statistical features extracted from network traffic or graph representations of function calls; therefore, their data modality is fundamentally different from ours. Among works focusing on malware images is (Abdel-Basset et al., 2022), which proposed FED-MAL, an FL framework for MIC in the IoT setting. It utilizes a custom lightweight CNN backbone as the local model and achieves an accuracy of 91.33% on the Virus-MNIST dataset. Another CNN-based FL approach was proposed in (Ullah et al., 2023), in which raw malware binaries were converted into color images, and a data augmentation method was utilized to balance malware data during local training. In (Ambekar et al., 2025), Siamese networks were integrated as feature extractors for MIC in the FL setting. This approach was supplemented with adversarial training to improve model robustness against evasion attacks. Finally, in (Ciaramella et al., 2024), the authors evaluated the impact of batch normalization on the FedAvg algorithm and provided a comparison with group normalization. While these works demonstrate the effectiveness of federated MIC, they focus primarily on model ar-

chitecture, accuracy, and adversarial robustness under homogeneous or mildly heterogeneous data distributions. Critically, they do not perform a systematic evaluation and improvement regarding the impacts of data heterogeneity on classification performance. In this work, we address this gap, studying the impact of non-IID data specifically on federated MIC. As with prior works, we utilize the FedAvg algorithm for aggregation, and similar to (Jiang et al., 2022; Jiang et al., 2024), we simulate the non-IID setting using a Dirichlet distribution.

3

Preliminaries and Problem Statement

3.1

System Model and Notation

Consider an FL setup with K clients, where each client has a local dataset of malware image samples nk . Here xi,k ∈ Rd is denoted by Dk = {(xi,k , yi,k )}i=1 an image generated from a program’s binary, yi,k ∈ {1, 2, ..., L} is the ground truth class label of xi,k (malware type or benign), and nk is the number of samples in Dk . We denote by n the total number of samples across all clients, i.e., n = ∑Kk=1 nk . Each image xi,k is a visual representation of a program’s binary (malware or benign). There are multiple ways in which the images can be constructed, e.g., the binary is first processed as a vector of 8-bit unsigned integers, then converted to a 2D array, and then visualized as a grayscale image (Nataraj et al., 2011; Noever and Noever, 2021). Malware image classification (MIC) can be formulated as a supervised multi-class classification problem, where the goal is to learn a model fw : Rd → {1, 2, ..., L} that maps image representations of binaries to their corresponding class labels. The model parameters w can be obtained by minimizing: min w

1 K nk ∑ ∑ ℓ( fw (xi,k ), yi,k ) n k=1 i=1

(1)

where ℓ(·) denotes the loss function (e.g., crossentropy loss). In a traditional centralized learning approach, all clients send their Dk to a central server, who trains fw according to Equation 1. While effective, this centralized approach requires clients to share their sensitive data Dk with a central server, which raises privacy concerns. The sharing of Dk is particularly problematic in the context of MIC since the samples in Dk can enable the server to learn which programs are running and whether there are any malware infections on

client k’s device. The use of FL in MIC eliminates the need to share Dk with a central server, and therefore, it can help address privacy concerns.

3.2

Federated Learning for Malware Image Classification

Federated learning (FL) enables multiple clients to collaboratively train a global model without exchanging their local datasets (McMahan et al., 2017; Kairouz et al., 2021; Li et al., 2020a). FL is a roundby-round, iterative process. We denote by wt the global model at communication round t, and by wtk the local model parameters of client k at round t. At round t, the FL server maintains the current global model wt and performs the following steps: 1. A subset of clients St ⊆ {1, . . . , K} is selected to participate in round t, e.g., by sampling a fraction C of all clients: |St | = C · K. 2. The global model wt is sent to the participating clients. 3. Each participating client k ∈ St locally updates the model by performing E epochs of training on its local dataset Dk , producing updated parameters wt+1 k . 4. The clients send updated models wt+1 to the k server, and the server aggregates received models to obtain the new global model wt+1 . This process is repeated either for a fixed number of T rounds or until model convergence. In step 3, the local update performed on client k’s device minimizes the local objective: Fk (w) =

1 nk ∑ ℓ( fw (xi,k ), yi,k ) nk i=1

(2)

by performing E epochs of stochastic gradient descent (SGD) initialized at wt . For the aggregation step (step 4), we employ Federated Averaging (FedAvg) (McMahan et al., 2017), the de facto aggregation algorithm in FL. In FedAvg, the global model is computed as the weighted average of local models: nk wt+1 = ∑ wt+1 (3) k n ∑ j t k∈S j∈St

where the contribution of each client is proportional to the size of its local dataset.

3.3

Problem Statement

A fundamental challenge in FL is data heterogeneity, where client data is not identically and independently

distributed (non-IID), i.e., Dk of different clients k differ significantly in terms of cardinality, class distribution, and/or sample characteristics. Prior work has shown that such heterogeneity can significantly degrade model performance, and slow or prevent convergence in algorithms such as FedAvg, FedProx, and SCAFFOLD (McMahan et al., 2017; Li et al., 2020a; Li et al., 2022; Jimenez-Gutierrez et al., 2025). Non-IID distributions are particularly relevant in federated learning of malware image classifiers. In realistic environments, clients may span personal electronic devices (laptops, smartphones, tablets), enterprise systems, and IoT devices, each exposed to distinct threat landscapes. For example, enterprise systems may encounter targeted or advanced malware families, whereas personal devices are more likely to contain commodity malware such as adware or spyware. Furthermore, different devices run different operating systems and applications; therefore, they are likely to contain different vulnerabilities and be affected by different malware. Consequently, the data distributions across different clients are likely to be imbalanced and non-IID. Data heterogeneity is further exacerbated by the client sampling process in FL: in each round t, only a subset of clients St participates in training, meaning that the effective training data distribution varies across rounds depending on the participating clients. As a result, the aggregated model update may be biased toward the class distributions present in St , introducing additional variability and instability in the training process. Despite the inherently non-IID nature of data in malware-related applications, its impact on federated MIC remains underexplored. In this paper, we aim to fill this gap. More specifically, we formulate and answer the following research questions (RQs): RQ1: Can FL be used to train accurate malware image classifiers under IID and non-IID data distributions? How does increasing data heterogeneity (nonIID’ness) affect the robustness and accuracy of these models? RQ2: Can certain round-level characteristics of participating clients, such as aggregate statistical properties of their data, influence or correlate with global model accuracy? RQ3: Can these round-level statistics be leveraged to design more effective client selection strategies that improve the accuracy of federated malware image classifiers, compared to random client participation?

4

Methodology

4.1

System Architecture

We utilize an FL setup consisting of an FL server (aggregator) and K = 50 clients, as illustrated in Figure 1. The FL server is responsible for global model initialization, broadcasting the current model to clients, and aggregating the updated models received from clients. Clients are responsible for training local models using their private dataset Dk of malware images, and communicating the updated model weights to the server. We use MobileNetV2 as the model architecture (Sandler et al., 2018). MobileNetV2 is a lightweight convolutional neural network architecture designed for efficient use on resource-constrained devices; therefore, it is well-suited for FL and image data.

4.2

Malware Image Datasets

We used three publicly available malware image datasets in our work: MalImg, Virus-MNIST, and MalNet. Sample images from the MalImg dataset are provided in Figure 2. It can be observed that samples from the same malware family Allaple.A (Figure 2a and 2b) are visually similar, whereas samples from the other families (Figure 2c and 2d) are visually distinct. MalImg (Nataraj et al., 2011) is a widely used malware image dataset consisting of grayscale images generated from malware binaries. The dataset contains a total of 9,339 samples spanning 25 malware families, such as Allaple, Yuner, Instantaccess, and Fakerean. We used an 80%-20% train-test split to ensure sufficient representatives from each family in our test set. Virus-MNIST (Noever and Noever, 2021) is a dataset of 51,880 grayscale images across 10 classes generated from Windows Portable Executable files. It consists of 2,516 samples from benign programs and

Figure 1: Our system architecture for federated learning of malware image classifiers (MICs).

(a)

(b)

inantly observe a small subset of classes, while larger values of α yield more balanced distributions closer to IID. Figure 3 provides visualizations of the data distributions for different values of α on the MalNet dataset. We used E = 3 local epochs, fraction C = 0.1, and T = 100 communication rounds for training. Local models were trained with the SGD optimizer, using learning rate η = 10−2 , and momentum µ = 0.9. We did not utilize weight decay as baseline experiments without it already showed satisfactory performance. A batch size of 10 for MalImg and 32 for MalNet and Virus-MNIST is used, based on our empirical observations regarding model convergence.

4.4 (c) (d) Figure 2: Sample images from the MalImg dataset. (a) and (b) are samples from the Allaple.A family of network worms. (c) is a sample from the Fakerean trojan, and (d) is a sample from the Obfuscator.AD! trojan.

49,364 malware samples divided into 9 malware families. We maintained the 85%-15% train-test split implemented by the authors. MalNet (Freitas et al., 2022) is a large-scale malware image dataset derived from Android applications (APK files). It contains hundreds of thousands of samples spanning numerous malware families, where each application is converted into an image representation based on its bytecode. In addition to grayscale representations, it incorporates an additional semantic coloring of each pixel based on its function. We performed our experiments on a subset of this dataset consisting of 45,385 samples from the ten largest classes. Furthermore, visual inspection revealed that informative patterns were primarily along the vertical dimension, so we cropped the images to 32×256 to improve model convergence speed, allowing us to better focus on the impact of IID versus non-IID. We used a 70%-30% train-test split.

4.3

Non-IID Simulation and Hyperparameters

To simulate non-IID data distributions across clients, we adopt the commonly used Dirichlet partitioning strategy. Specifically, for each class l ∈ {1, . . . , L}, we sample a K-dimensional proportion vector Pl = (pl,1 , . . . , pl,K ) ∼ Dir(α), where α > 0 is the concentration parameter controlling the degree of nonIID’ness. Each pl,k determines the fraction of samples of class l assigned to client k. Smaller values of α lead to highly skewed distributions where clients predom-

Round-Level Statistics

Only a subset of clients St ⊆ {1, . . . , K} participate in each round t of FL. As discussed in Section 3.3, this induces variability in the effective training data amount and distribution across rounds, especially in non-IID scenarios. To analyze the effects of this variability and which aspects of the variability affect model accuracy the most, we propose a set of roundlevel statistics that capture different properties of the data contributed by clients St . Our statistics are designed to consider complementary aspects of the data distribution at each round, enabling us to analyze how client composition influences model accuracy under non-IID data. 4.4.1

Total Samples

We define the total number of samples N t available in round t as: N t = ∑ nk (4) k∈St

This metric captures the overall amount of data used to compute the global update in round t. In general, larger values of N t are expected to yield more stable and accurate models, since the model learns from a larger number of samples. However, under non-IID settings, large N t may still correspond to highly imbalanced data if the majority of the samples come from a few classes. 4.4.2

Median Samples

We compute the median of the set {nk : k ∈ St }, which reflects the median dataset size among participating clients. Unlike N t , this metric is less sensitive to outliers, i.e., few clients with very large or very small nk .

= 0.5

= 0.1

= 0.05

= 0.01

Clients

IID

Number of Samples addisplay

adsware

Number of Samples clicker++trojan

Number of Samples

downloader

exploit

Number of Samples

smssend

smssend++trojan

Number of Samples spr

spyware

troj

Figure 3: Distribution of MalNet dataset among 50 clients under varying degrees of non-IID, controlled by the α parameter. Each class is represented by a different color. It can be observed that α = 0.5 is closer to IID, as clients have more similar numbers of samples, and their samples span multiple different classes. In contrast, when α = 0.01, clients have very different dataset sizes and their samples are predominantly from one or two classes.

4.4.3

Sample Entropy

To quantify how evenly the data is distributed across participating clients, we compute the Shannon entropy of the sample distribution. Let: nk ptk = t , k ∈ St (5) N denote the fraction of samples contributed by client k in round t. The sample entropy is then defined as: t Hsample = − ∑ ptk log2 ptk

(6)

k∈St

t indicate a more balanced Higher values of Hsample contribution across clients, while lower values indicate that a small number of clients dominate the update. Such dominance can bias the aggregated model toward specific client distributions, which can be problematic under non-IID data.

4.4.4

Low class entropy indicates that the round is dominated by a small number of classes, which can lead to biased model updates and reduced generalization performance, as the model will not see sufficiently many samples from other classes. 4.4.5

Average Class Entropy

We also consider the average class entropy across in(l,t) dividual clients. For each client k ∈ St , let nk denote the number of samples of class l in client k’s dataset, and define: (l,t) n qtk,l = k (9) nk Then, the class entropy of client k at round t is: L

k,t Hclass = − ∑ qtk,l log2 qtk,l

(10)

l=1

Finally, the average class entropy in round t is:

Class Entropy

We note that sample entropy considers only the amount of data, not the class the data belongs to. To measure the diversity of class labels in the data available in round t, we compute the Shannon entropy of the class distribution. Let n(l,t) denote the number of samples of class l in round t, and define: n(l,t) , l ∈ {1, . . . , L} Nt The class entropy is given by: qtl =

1 k,t ∑t Hclass |St | k∈S

(11)

This metric reflects how diverse the class distributions are within individual clients. Low values indicate that clients’ local datasets Dk are concentrated on a few classes, which can amplify heterogeneity effects during FL.

(7)

4.4.6

(8)

We define a weighted version of average class entropy that accounts for client dataset sizes: 1 k,t t bclass H = t ∑ nk Hclass (12) N k∈St

L t Hclass = − ∑ qtl log2 qtl

t H̄class =

l=1

The key difference between class entropy and sample entropy is that class entropy captures the diversity of class labels present in the participating clients.

Weighted Average Class Entropy

This metric captures the overall class diversity while giving higher weight to clients with larger datasets,

since their updates are likely to be more dominant in FL. It serves to balance the class distribution and sample size.

4.5

Entropy-Aware Client Selection

t We empirically observed that class entropy Hclass , in particular, exhibits a strong and consistent correlation with model accuracy. To answer RQ3 by turning this observation into an actionable insight, we propose an entropy-aware client selection strategy. The key intuition of our strategy is that rounds with higher class entropy tend to produce more representative global updates; therefore, explicitly favoring such client subsets in FL will improve model accuracy. Our proposed strategy operates as follows. At each round t: 1. The FL server generates a large pool of candidate client subsets, each of size C · K, by sampling combinations of clients uniformly at random among all K clients. 2. For each candidate subset, the statistics defined in Section 4.4 are computed. 3. The candidate subsets are filtered by discarding: (i) those subsets with total number of samples N t below the 25th percentile, and (ii) those with class entropy below the 95th percentile. 4. Among the remaining candidates, one of them is randomly selected to become the St that will participate in round t. The filtering in the third step ensures that: (i) there will be sufficient data per round of training, and (ii) following the intuition from the previous section, high class entropy (among the top 5%) will be enforced among the possible candidates. The selection in step 4 is non-deterministic and without replacement, so that our strategy does not always select the same clients every round.

5

Results and Analysis

In this section, we experimentally analyze and discuss the answers to the three RQs stated in Section 3.3. We dedicate Section 5.1 to RQ1, Section 5.2 to RQ2, and Section 5.3 to RQ3. Lastly, we perform more detailed hyperparameter analyses in Section 5.4.

5.1

Impact of Non-IID on Model Accuracy and Convergence

To answer RQ1, we evaluate the performance of our FL setup under varying degrees of non-IID, simulated

using Dirichlet distributions with concentration parameter α ∈ {0.5, 0.1, 0.05, 0.01}. We also include results under IID data as a point of comparison. (The data distributions are exemplified in Figure 3.) Table 1 summarizes the accuracy values of the global model in each setting, across the MalImg, Virus-MNIST, and MalNet datasets. Results in Table 1 show a consistent degradation in model accuracy as the degree of non-IID’ness increases. Across all datasets, models achieve the highest accuracy under IID conditions, and their performance decreases as α becomes smaller. When α = 0.5, which represents a moderate level of heterogeneity, the reduction in accuracy remains relatively limited (approximately 4.5% on average), indicating that FedAvg can tolerate mild heterogeneity. However, as α decreases further, the impact becomes more pronounced. In the extreme case of α = 0.01, the model accuracy drops to less than half of the IID baseline. We would like to highlight that under IID conditions, highly accurate malware image classifiers can be trained using FL. For example, as provided in Table 1, we achieve accuracies exceeding 99% on MalImg and 91% on Virus-MNIST. These results are comparable to those reported in the original dataset papers and to those obtained under centralized training. Therefore, the primary challenge does not stem from FL itself, but rather the presence of non-IID data distributions. In addition to reduced accuracy, we observe increased instability in the training process as heterogeneity increases. Specifically, the model exhibits greater variability in test accuracy across communication rounds when α is small. This behavior is illustrated in Figure 4, where learning curves become increasingly volatile as α is reduced. Under IID conditions, the training process is smooth and stable, with accuracy steadily increasing and quickly converging to a high level across all datasets. For moderate non-IID (e.g., α = 0.5), the model still converges to relatively high accuracy (although the final accuracy is slightly lower than that of IID), but with noticeable fluctuations during training. However, as α decreases further, the instability becomes significantly more pronounced. For α = 0.1 and α = 0.05, the accuracy curves exhibit frequent and large oscillations, indicating that successive global updates may move the model in inconsistent directions. The final model accuracy at the end of training is also noticeably lower than IID and α = 0.5. In the extreme case of α = 0.01, the training process becomes highly erratic, with accuracy fluctuating substantially and often failing to converge to a stable or accurate model. While some learning occurs on MalImg, the model does not per-

Table 1: Test accuracy (%) under varying degrees of non-iid

IID

α = 0.5

α = 0.1

α = 0.05

α = 0.01

99.04 ± 0.28 91.97 ± 0.42 84.30 ± 0.17

96.83 ± 4.00 86.35 ± 2.19 78.53 ± 1.78

83.92 ± 9.03 65.69 ± 11.74 55.77 ± 13.32

74.04 ± 12.45 59.24 ± 13.44 50.35 ± 5.17

47.63 ± 22.96 22.99 ± 0.12 14.31 ± 11.74

Dataset

(a) MalImg

80 60 40 20 0

0

20

40

60

(b) Virus-MNIST

100

Test Accuracy

Test Accuracy

100

80

Communication Rounds IID

100

80 60 40 20 0

0

20

40

60

(c) MalNet

100

Test Accuracy

MalImg Virus-MNIST MalNet

80

100

Communication Rounds =0.5 =0.1 =0.05

80 60 40 20 0

0

20

40

60

80

Communication Rounds =0.01

100

Figure 4: Accuracy versus rounds on different datasets under varying α

form better than a random guess on Virus-MNIST or MalNet at any point during training. Overall, these behaviors can be attributed to the increasing mismatch between local data distributions and the global objective as α becomes smaller. When client data is highly skewed, local updates computed during each round are strongly biased toward a small subset of classes. As different subsets of clients participate in different rounds, the aggregated updates vary significantly, leading to accuracy oscillations and convergence instability.

5.2

Round-Level Statistics and Model Accuracy

Motivated by the reduced accuracy and increased oscillations observed under non-IID data, we investigate how the composition of participating clients in each round influences model performance (RQ2). More specifically, for each round t, we compute the roundlevel statistics of participating St using the metrics proposed in Section 4.4, and also record the accuracy of the model in that round. Then, we compute the Spearman correlation between each metric and model accuracy, across 100 rounds and 5 independent runs (a total of 500 recordings). Table 2 summarizes the results across different datasets and α values. In each column, the metric that achieves the strongest correlation is indicated in bold. Cells shaded with gray background color indicate strong statistical significance: p-value < 0.05 in at least 4 of 5 runs. Across all datasets and most values of α, class t entropy Hclass exhibits the strongest and most consistent positive correlation with model accuracy. In

particular, for moderate heterogeneity levels (α = 0.5, 0.1, 0.05), the correlation is both strong and statistically significant across all datasets. This indicates that rounds with more diverse class distributions tend to yield better-performing global models. However, under extreme heterogeneity (α = 0.01), this relationship weakens, suggesting that when client data becomes highly skewed, even high-entropy rounds t may become insufficient. Sample entropy Hsample also shows a generally positive correlation with model accuracy, particularly on Virus-MNIST and MalNet. This suggests that a more balanced contribution of data across clients serve to stabilize model performance and mitigate accuracy degradation. However, its correlation is weaker and less significant than that of class entropy. Similarly, while average class ent bt tropy H̄class and its weighted version H class exhibit moderate positive correlations in some cases, they do not outperform class entropy in general. This suggests that the aggregate class distribution across all particit pating clients (captured by Hclass ) is more informative t bt . than client-level H̄class and H class We also study the correlations of Total Samples and Median Samples with model accuracy. Contrary to the common expectation that more data leads to better performance, we often observe a negative correlation between the total number of samples N t and model accuracy in several configurations, particularly on the MalImg dataset. Furthermore, median number of samples also does not demonstrate a significant correlation with accuracy. These do not imply that larger datasets are inherently detrimental; rather, it reflects the interaction between sample size and data heterogeneity. Due to the inherent skew as shown in

Table 2: Correlation of round-level statistics with model accuracy across different datasets and α values MalImg

Virus-MNIST

MalNet

Metric

0.5

0.1

0.05

0.01

0.5

0.1

0.05

0.01

0.5

0.1

0.05

0.01

Total Samples Median Samples Sample Entropy Class Entropy Avg. Class Entropy Wght. Avg. Class Entropy

-0.358 -0.292 0.193 0.578 0.532 0.562

-0.347 -0.132 0.356 0.413 0.296 0.425

-0.154 0.009 0.291 0.375 0.230 0.359

0.062 0.027 -0.001 0.065 -0.022 -0.135

-0.271 -0.162 0.256 0.507 0.430 0.491

0.000 0.172 0.311 0.295 0.045 0.168

0.119 0.260 0.350 0.459 0.061 0.252

-0.124 0.043 0.307 0.313 0.094 0.172

-0.014 -0.094 -0.023 0.289 0.201 0.233

-0.091 0.165 0.400 0.589 0.242 0.377

-0.014 0.205 0.423 0.499 0.001 0.179

0.399 0.245 0.214 0.239 -0.075 -0.019

Table 3: Test accuracy (%) under random client selection versus our entropy-aware client selection strategy

Table 4: Impact of participating client fraction (C) on model accuracy

Dataset

α

Random

Entr-Aware

C

IID

α = 0.1

α = 0.01

MalImg

0.5 0.1 0.05 0.01

96.83 ± 4.00 83.92 ± 9.03 74.04 ± 12.45 47.63 ± 22.96

98.90 ± 0.09 96.51 ± 3.64 92.78 ± 2.95 58.51 ± 12.43

86.35 ± 2.19 65.69 ± 11.74 59.24 ± 13.44 22.99 ± 0.12

85.47 ± 2.97 71.08 ± 9.90 62.51 ± 12.27 26.52 ± 11.28

99.04 ± 0.28 99.20 ± 0.19 99.10 ± 0.06 98.90 ± 0.11 98.94 ± 0.13

83.92 ± 9.03 92.21 ± 4.40 91.65 ± 4.16 93.65 ± 4.55 92.12 ± 4.22

47.63 ± 22.96 51.04 ± 15.77 56.75 ± 19.80 47.88 ± 10.76 41.26 ± 9.74

Virus-MNIST

0.5 0.1 0.05 0.01

0.1 0.2 0.4 0.8 1.0

MalNet

0.5 0.1 0.05 0.01

78.53 ± 1.78 55.77 ± 13.32 50.35 ± 5.17 14.31 ± 11.74

78.82 ± 1.52 67.17 ± 3.87 54.16 ± 4.13 24.61 ± 11.30

Figure 3, rounds with large N t are often dominated by a small number of clients with large but skewed datasets, which reduces both class and sample entropy and leads to worse model updates. Indeed, having sufficiently many samples in a round is often beneficial to improve accuracy, as evidenced by positive correlations especially when α = 0.01 and 0.05.

5.3

Evaluation of Entropy-Aware Client Selection

Recall that we proposed an entropy-aware client selection strategy in Section 4.5. In this section, we evaluate the effectiveness of our strategy by comparing it against standard random client selection. Hyperparameters are kept the same as in previous experiments, each experiment is repeated five times to ensure statistical reliability, and the average results are reported in Table 3. Across almost all datasets and non-IID levels α, the proposed entropy-aware strategy consistently outperforms random selection. The only exception is the Virus-MNIST dataset at α = 0.5, where both methods achieve comparable performance (the difference between them is lower than 1%). The Wilcoxon signedrank test with a significance level of 0.01 also showed

that these results are statistically significant. These results demonstrate that prioritizing class entropy in client selection leads to more effective training and improved model accuracy. Overall, these experiments provide empirical evidence that round-level statistics, particularly class entropy, can be used not only to explain performance variability but also to actively improve the learning of malware image classifiers under non-IID data.

5.4

Hyperparameter Analyses

We study the effects of the following three parameters on model accuracy: the participating fraction C, the number of local training epochs E, and the total number of FL clients K. We maintain the same hyperparameters used in previous experiments, only varying the specific parameter being tested. Experiments are performed on the MalImg dataset, in IID and non-IID settings with α = 0.1 and 0.01. Impact of C. Table 4 summarizes the results under IID and non-IID settings. Under IID conditions, varying C has only a marginal effect on model accuracy, with all configurations achieving consistently high performance (around 99%). This indicates that when data is uniformly distributed, the number of participating clients per round is not a critical factor. In contrast, we observe more significant variation in model performance in the non-IID settings. When α = 0.1, increasing C initially improves accuracy, with accuracy peaking at C = 0.8, suggesting that involving more clients helps incorporate more diverse data and mitigates heterogeneity. However, further increasing C to 1.0 does not yield additional gains and

Table 5: Impact of the number of local epochs (E) on model accuracy

Table 6: Impact of the number of clients (K) on model accuracy

E

IID

α = 0.1

α = 0.01

K

IID

α = 0.1

α = 0.01

1 3 5 7 9 11

98.82 ± 0.16 99.04 ± 0.28 99.22 ± 0.11 98.92 ± 0.09 98.94 ± 0.16 98.86 ± 0.27

81.56 ± 12.97 83.92 ± 9.03 88.93 ± 6.80 90.05 ± 5.10 87.28 ± 9.76 85.53 ± 11.97

39.99 ± 10.36 47.63 ± 22.96 51.57 ± 16.62 32.83 ± 14.30 37.12 ± 5.27 31.14 ± 12.36

10 25 50 100

99.14 ± 0.25 99.26 ± 0.11 99.04 ± 0.28 98.93 ± 0.19

84.94 ± 17.80 88.84 ± 4.72 83.92 ± 9.03 82.91 ± 13.79

63.35 ± 15.60 69.05 ± 15.55 47.63 ± 22.96 58.70 ± 3.68

slightly degrades performance. A similar but more pronounced trend is observed when α = 0.01. Increasing C from 0.1 to 0.4 improves accuracy, but beyond this point, performance deteriorates, with C = 1.0 resulting in the lowest accuracy. Overall, these results indicate that using the majority or all of the clients does not guarantee better accuracy – in fact, increasing C can yield diminishing or even negative returns. It should also be noted that an increased C also requires higher communication and computational overhead for the clients, which is undesirable. Impact of E. Table 5 summarizes the results of varying E on model accuracy under IID and non-IID settings. As shown in Table 5, under IID conditions, varying E has only a minor effect on accuracy, with all configurations achieving consistently high performance around 99%. This observation is similar to the C parameter. In contrast, under non-IID settings, the effect of E is more pronounced and non-monotonic. When α = 0.1, increasing E from 1 to 7 leads to a steady improvement in accuracy, with the best performance achieved at E = 7. However, further increasing E results in a decline in accuracy. This behavior is even more evident when α = 0.01, where performance peaks at E = 5 and then drops sharply as E increases, with E = 11 yielding the worst results. We attribute this behavior to the local models overfitting to their local data at higher values of E. Increasing E allows clients to better optimize their local objectives. However, under non-IID data, larger values of E cause local models to overfit to their skewed local distributions, leading to updates that deviate significantly from each other and the global objective. As a result, the aggregated model generalizes poorly. Hence, our results suggest that moderate values of E should be preferred under non-IID settings. Impact of K. We report the accuracy impact of varying K in Table 6. In general, there is no strong relationship between K and the model performance. Under IID conditions, the accuracies under varying K are close to one another. Under non-IID conditions, the relationship becomes non-trivial and inconsistent. For both α = 0.1 and α = 0.01, accuracy peaks at K = 25, while both smaller and larger K values result in

reduced accuracy. In particular, for larger K, both the occurrence likelihood and the probability of selecting highly skewed clients in each round increases, which negatively affects model accuracy.

6

Conclusion

In this paper, we studied FL for malware image classification under data heterogeneity (non-IID). Through an empirical study across three datasets (MalImg, Virus-MNIST, and MalNet), we showed that both accuracy and training stability degrade as data becomes increasingly non-IID. We introduced a set of roundlevel statistics to characterize the data contributed by participating clients per FL round, and observed that class entropy exhibits the strongest correlation with global model accuracy. Leveraging this insight, we proposed an entropy-aware client selection strategy that enforces high class entropy among participating clients, and showed that it outperforms standard random client selection. The entropy-aware client selection mechanism proposed depends on server access to client data statistics. Future work could investigate privacypreserving estimation of the statistics in section 4.4, such as inference from model updates or utilizing differential private mechanisms. Combining distribution-aware client selection with robust FL aggregation algorithms, and studying their impacts on malware image classification is another promising direction.

ACKNOWLEDGEMENTS This research was supported by the Scientific and Technological Research Council of Türkiye (TÜBİTAK) under 2247-D Award 123C595, the BAGEP Outstanding Young Scientist Award, TÜBİTAK 2247-A Award 121C338, and EU CHISTERA TÜBİTAK Grant 125N724. The authors thank TÜBİTAK and the Science Academy for their support.

REFERENCES Abdel-Basset, M., Hawash, H., Sallam, K. M., Elgendi, I., Munasinghe, K., and Jamalipour, A. (2022). Efficient and lightweight convolutional networks for iot malware detection: A federated learning approach. IEEE Internet of Things Journal, 10(8):7164–7173. Ambekar, N. G., Samal, S., Devi, N. N., and Thokchom, S. (2025). Fasnet: Federated adversarial siamese networks for robust malware image classification. Journal of Parallel and Distributed Computing, 198:105039. Amjath, M., Henna, S., and Rathnayake, U. S. (2024). Graph representation federated learning for malware detection in internet of health things. Results in Engineering. Brosolo, M., Puthuvath, V., Ka, A., Rehiman, R., and Conti, M. (2024). Sok: visualization-based malware detection techniques. In Proceedings of the 19th International Conference on Availability, Reliability and Security, pages 1–13. Ciaramella, G., Martinelli, F., Mercaldo, F., Peluso, C., and Santone, A. (2024). An approach for privacypreserving mobile malware detection through federated machine learning. In International Conference on Enterprise Information Systems. D’Angelo, G., Farsimadan, E., Ficco, M., Palmieri, F., and Robustelli, A. (2023). Privacy-preserving malware detection in android-based iot devices through federated markov chains. Future Generation Computer Systems, 148:93–105. Eslamnejad, M., Taheri, R., Shojafar, M., and BaderEl-Den, M. (2025). Federated learning-based robust android malware detection: label-flipping attacks and defenses. Neural Computing and Applications, 37(32):27057–27082. Freitas, S., Duggal, R., and Chau, D. H. (2022). Malnet: A large-scale image database of malicious software. In Proceedings of the 31st ACM International Conference on Information & Knowledge Management, pages 3948–3952. Gibert, D., Mateu, C., and Planes, J. (2020). The rise of machine learning for detection and classification of malware: Research developments, trends and challenges. Journal of Network and Computer Applications, 153:102526. Jiang, C., Xia, C., Liu, M., Chen, C., Li, H., Wang, T., and Li, P. (2024). Feddrc: A robust federated learning-based android malware classifier under heterogeneous distribution. 2024 27th International Conference on Computer Supported Cooperative Work in Design (CSCWD), pages 1029–1036. Jiang, C., Yin, K., Xia, C., and Huang, W. (2022). Fedhgcdroid: An adaptive multi-dimensional federated learning for privacy-preserving android malware classification. Entropy, 24. Jimenez-Gutierrez, D. M., Hassanzadeh, M., Anagnostopoulos, A., Chatzigiannakis, I., and Vitaletti, A. (2025). A thorough assessment of the non-iid data impact in federated learning. Journal of Industrial Information Integration, page 101052.

Kairouz, P., McMahan, H. B., Avent, B., Bellet, A., Bennis, M., Bhagoji, A. N., Bonawitz, K., Charles, Z., Cormode, G., Cummings, R., et al. (2021). Advances and open problems in federated learning. Foundations and Trends® in Machine Learning, 14(1–2):1–210. Karimireddy, S. P., Kale, S., Mohri, M., Reddi, S., Stich, S., and Suresh, A. T. (2020). Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, pages 5132– 5143. PMLR. Kumar, S. and Janet, B. (2022). Dtmic: Deep transfer learning for malware image classification. Journal of Information Security and Applications, 64:103063. Lan, K., Li, G., Huang, W., and Li, J. (2025). Hfl-rd: Heterogeneous federated learning-empowered ransomware detection via apis and traffic features. IEEE Transactions on Network and Service Management. Li, Q., Diao, Y., Chen, Q., and He, B. (2022). Federated learning on non-iid data silos: An experimental study. In 2022 IEEE 38th International Conference on Data Engineering (ICDE), pages 965–978. IEEE. Li, T., Sahu, A. K., Talwalkar, A., and Smith, V. (2020a). Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50–60. Li, T., Sahu, A. K., Zaheer, M., Sanjabi, M., Talwalkar, A., and Smith, V. (2020b). Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429–450. McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A. (2017). Communication-efficient learning of deep networks from decentralized data. In Artificial Intelligence and Statistics, pages 1273–1282. PMLR. Nataraj, L., Karthikeyan, S., Jacob, G., and Manjunath, B. S. (2011). Malware images: visualization and automatic classification. In Proceedings of the 8th International Symposium on Visualization for Cyber Security, pages 1–7. Nobakht, M., Javidan, R., and Pourebrahimi, A. (2024). Sim-fed: Secure iot malware detection model with federated learning. Computers and Electrical Engineering, 116:109139. Noever, D. and Noever, S. E. M. (2021). Virus-mnist: A benchmark malware dataset. arXiv preprint arXiv:2103.00602. Patel, A., Tomar, D. S., Pateriya, R., and Haripriya, R. (2025). Fl-maldrift: a federated learning framework for malware detection under local concept drift. Scientific Reports. Rey, V., Sánchez, P. M. S., Celdrán, A. H., and Bovet, G. (2022). Federated learning for malware detection in iot devices. Computer Networks, 204:108693. Sandler, M., Howard, A., Zhu, M., Zhmoginov, A., and Chen, L.-C. (2018). Mobilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 4510–4520. Tekerek, A. and Yapici, M. M. (2022). A novel malware classification and augmentation model based on con-

volutional neural network. Computers & Security, 112:102515. Ullah, F., Srivastava, G., Ullah, S., and Mostarda, L. (2023). Privacy-preserving federated learning approach for distributed malware attacks with intermittent clients and image representation. IEEE Transactions on Consumer Electronics, 70(1):4585–4596. Vasan, D., Alazab, M., Wassan, S., Naeem, H., Safaei, B., and Zheng, Q. (2020a). Imcfn: Image-based malware classification using fine-tuned convolutional neural network architecture. Computer Networks, 171:107138. Vasan, D., Alazab, M., Wassan, S., Safaei, B., and Zheng, Q. (2020b). Image-based malware classification using ensemble of cnn architectures (imcec). Computers & Security, 92:101748. Vehabovic, A., Zanddizari, H., Shaikh, F., Ghani, N., Pour, M. S., Bou-Harb, E., and Crichigno, J. (2023). Federated learning approach for distributed ransomware analysis. In International Conference on Applied Cryptography and Network Security, pages 621–641. Springer.

Related documents

• FEDetect: A Federated Learning-Based Malware Detection and Classification Using Deep Neural Network Algorithms#361193
• Federated Learning Models for Privacy-Preserving Medical Image Analysis#215781
• Federated multimodal malware classification under non-IID data#413282
• Federated learning and differential privacy for medical image analysis#177282
• Federated Learning Framework for Privacy-Preserving Kidney Stone Detection#1014328