FedOPCS: An Optimized Poisoning Countermeasure for Non-IID Federated Learning with Privacy-Preserving Stability

Abstract

Federated learning (FL), as a distributed machine learning framework, enables multiple participants to jointly train models without sharing data, thereby ensuring data privacy and security. However, FL systems still struggle to escape the typical poisoning threat launched by Byzantine nodes. The current defence measures almost all rely on the anomaly detection of local gradients in a plaintext state, which not only weakens privacy protection but also allows malicious clients to upload malicious ciphertext gradients once they are encrypted, which thus easily evade existing screenings. At the same time, mainstream aggregation algorithms are generally based on the premise that “each client’s data satisfy an independent and identically distributed (IID)”, which is obviously difficult to achieve in real scenarios where large-scale terminal devices hold their own data. Symmetry in data distribution and model updates across clients is crucial for achieving robust and fair aggregation, yet non-IID data and adversarial attacks disrupt this balance. To address these challenges, we propose FedOPCS, an optimized poisoning countermeasure for non-IID FL algorithms with privacy-preserving stability by introducing three key innovations: Conditional Generative Adversarial Network (CGAN)-based data augmentation with conditional variables to simulate global distribution, a dynamic weight adjustment mechanism with homomorphic encryption, and two-stage anomaly detection combining gradient analysis and model performance evaluation. Extensive experiments on MNIST and CIFAR-10 show that, in the model poisoning and mixed poisoning environments, FedOPCS outperforms the baseline methods by 11.4% and 4.7%, respectively, while maintaining the same efficiency as FedAvg. FedOPCS therefore offers a privacy-preserving, Byzantine-robust, and communication-efficient solution for future heterogeneous FL deployments.

1. Introduction

In the digital age, data have become a key resource driving the development of machine learning models, bringing service improvements and efficiency gains to businesses and institutions [1,2]. However, data privacy and security issues have become major challenges. FL is an advanced machine learning paradigm that enables multiple participants to jointly train models while maintaining data privacy. FL is particularly suitable for distributed scenarios where data are dispersed across different geographical locations or organizations, such as retail, healthcare, and financial services. With the proliferation of IoT devices, FL is also applied in edge computing, optimizing device performance and user experience without exposing user data [3]. The ideal FL scenario assumes symmetric contributions from all clients, where data distributions and model updates are balanced. However, real-world non-IID data and poisoning attacks break this symmetry, leading to biased aggregation and compromised performance.

However, real-world non-IID data and poisoning attacks break this symmetry, leading to biased aggregation and compromised performance. To better illustrate this issue, Figure 1 provides an overview of the federated learning architecture and highlights its susceptibility to poisoning attacks under non-IID conditions.

Although FL theoretically offers advantages in data privacy protection, the existing research and practices indicate that most FL algorithms lack robustness. In fact, as the application of FL continues to expand, it begins to suffer from various security threats, with poisoning attacks currently being the most significant threat to FL. Poisoning attacks are mainly divided into data poisoning and model poisoning [4]. A data poisoning attack occurs when an attacker intentionally injects false or misleading information into training data to affect the model’s learning and prediction capabilities. Such attacks may occur during the data collection, storage, or uploading stages, where attackers tamper with local datasets to bias the trained model toward specific outputs or reduce its generalization ability. As shown in [3], in FL, due to the reliance on data provided by multiple participants, even a small amount of contaminated data can significantly degrade the performance of the global model. These attacks not only weaken the performance of the model, but also may implant backdoors that allow attackers to activate malicious behaviours under certain conditions. This threatens the model’s accuracy and reliability and may destabilize the entire system, ultimately jeopardizing user privacy and system security. A viable solution to this issue is to introduce a metric criteria algorithm, where the server side blocks clients that are poisoned to mitigate the harm to FL due to poisoning [5,6]. At the same time, some studies adopt encryption during the communication process of FL to enhance security performance [7].

As FL applications continue to expand, ensuring the system’s robustness in the face of various threats is crucial. The existing research on defending against model poisoning and data poisoning attacks has significant shortcomings [8]. One major limitation of current defence methods is that they often employ relatively simple and crude measures, such as directly discarding anomalous data or model updates. While this approach may be feasible for IID data, in a non-independent and identically distributed (non-IID) environment, the diversity and complexity of data may lead to the loss of important information, affecting the model’s generalization ability and overall performance [9]. Since local optimizers on each client minimize local objectives closely related to data distribution, local updates may steer the model in optimal directions for local objectives. However, these local objectives may deviate from the global data’s non-IID distribution, potentially causing instability in the training process and making it difficult for the FL model to converge. Over-reliance on discarding data may also weaken the model’s representativeness for minority groups, exacerbating data bias issues. Although the current research has proposed data augmentation strategies, such as the preprocessing and enhancement of data before training [10] and z-score-based data augmentation [11], their effectiveness is often limited, and they do not simultaneously address both types of attacks effectively. Research has shown that although data generated from Generative Adversarial Networks (GANs) [12] can enhance the robustness of non-IID data at FL nodes, existing studies using GAN have not proposed effective optimal node selection and privacy protection mechanisms. Therefore, it is necessary to develop adaptive defense mechanisms to enhance the robustness of the system [13]. Researchers need to develop more refined and adaptive defense mechanisms to improve the robustness of FL systems.

To address the challenges of model poisoning and data poisoning attacks in federated learning under non-IID data environments and the urgent need to improve model robustness and security, the contributions of this paper are as follows:

• Novel weight constraints and metric criteria: We propose a new set of weight constraints and metrics for evaluating and ensuring the security and robustness of models in FL environments. The dynamic weight adjustment mechanism ensures symmetric contributions from clients by penalizing malicious updates and rewarding trustworthy ones, thus maintaining equilibrium in the aggregation process.
• Conditional Generative Adversarial Network (CGAN) data generation method to enhance model robustness: Our CGAN-based data augmentation restores symmetry in local data distributions by generating synthetic samples that align with the global data profile, mitigating the skewness caused by non-IID conditions. This method allows us to provide a more accurate view of the data for the FL system while protecting data privacy, thereby effectively guiding the adjustment of client model aggregation weights.
• Model aggregation method using adaptive weighting and homomorphic encryption (HE) techniques: We propose an innovative client model aggregation framework that combines adaptive weighting and homomorphic encryption techniques. The adaptive weighting mechanism dynamically adjusts the weight of models in the global model based on their performance and trustworthiness to optimize the aggregation process and resist model poisoning attacks. Homomorphic encryption ensures the security of the entire aggregation process, preventing potential man-in-the-middle attacks and model leakage risks.

Distinctiveness Compared to Classical Algorithms: Compared with FedAVG, which only performs a simple weighted average of gradients for each client and lacks protection against non-IID data and poisoning threats [3], and FedProx, which alleviates statistical heterogeneity through near-end regularization terms but still cannot identify malicious updates [14], FedPNS mainly selects high-contribution nodes based on probability to accelerate convergence, without involving privacy and security [15]. FedOPCS also introduces three links, namely “inner product + test loss” dual anomaly detection, CGAN data enhancement, and Paillier encryption adaptive aggregation, achieving the unity of robustness, privacy, and interpretability. Thanks to these three improvements, FedOPCS can still improve the accuracy of CIFAR-10 by 11.4% (relative to FedAVG) and 4.7% (relative to FedPNS) in a 50% mixed poisoning environment, while maintaining a convergence speed close to FedProx.

2. Related Work

2.1. Security and Defence Mechanisms in Federated Learning

We categorise existing defences along two orthogonal axes—(i) privacy preservation and (ii) Byzantine robustness—and further discuss how each handles non-IID data.

In FL, since model training and parameter updates occur across multiple clients, ensuring the security of the entire model training and aggregation process becomes a critical issue. The model aggregation algorithm is a key component in FL, responsible for synthesizing the model updates provided by multiple clients into a global model. During this process, it is essential to carefully detect and avoid malicious updates that could contaminate the model. Anomaly detection techniques play a crucial role in this aspect, as they can identify and isolate abnormal data or model updates that may be submitted by malicious participants. Portet [5] proposed the Model-Heterogeneous Aggregation Training (MHAT) scheme to address the issues of high communication overheads and slow model convergence in traditional FL, reducing computational demands and optimizing model convergence precision through knowledge distillation. Although some innovative methods have been proposed for evaluating and improving FL aggregation algorithms, there are still shortcomings in terms of their applicability in heterogeneous environments, dynamic data distribution processing, and extensive experimental verification.

Building upon the foundation of efficient model aggregation, Krum [16] and Zeno [17] presented algorithms to study the resilience of the stochastic gradient descent (SGD) algorithm in the face of Byzantine failures and demonstrated its ability to resist Byzantine attacks. These algorithms are pivotal in understanding how to maintain the integrity of the learning process under adversarial conditions. The Bulyan [18] algorithm overcomes the shortcomings of SGD by ensuring the quality of model updates while achieving convergence. FLTrust bootstraps a root dataset on the server to score clients; however, it still requires plaintext gradients and assumes IID data are being used, leaving privacy and heterogeneity issues unsolved [19]. These aggregation-only defences are vulnerable once gradients are encrypted or data distribution is skewed. This advancement is a significant step towards creating more reliable FL systems. Furthermore, the DarkneTZ framework [20] protects model privacy from membership inference attacks by combining the trusted execution environments (TEE) of edge devices and model partitioning. This framework represents a novel approach to safeguarding user data within FL ecosystems. The papers mentioned in this paragraph proposed some innovative methods to improve the robustness of distributed learning systems. However, they still have shortcomings in terms of their computational complexity, communication overheads, ability to adapt to dynamic environments, ability to handle complex attacks, experimental verification, applicability, and universality.

Continuing the discussion on the challenges in distributed machine learning systems, Tomsett [21] analysed the unique challenges posed by computational and communication constraints and adversarial threats in the military domain, emphasizing the need for models to have high flexibility and robustness. This analysis underscores the importance of robust models in high-stakes environments. Although this paper provides a detailed introduction to the mechanism of model poisoning attacks, there is limited discussion on how to effectively defend against these attacks. The ShieldFL method proposed by Ma [22] enhances the robustness of model aggregation by using cryptographic protocols to protect model updates in FL processes, performing gradient similarity measurements without decryption through double trapdoor homomorphic encryption technology. However, its effectiveness may be limited when facing more complex and advanced attack strategies, such as poisoning data generated by Generative Adversarial Networks. Lancelot employs FHE to protect model updates but forgoes non-IID correction, leading to a drop in accuracy under distribution shifts [23]. RFLPA extends SecAgg with a cosine-similarity screening, yet still assumes IID data and omits gradient distribution smoothing [24]. PEAR adopts a more robust Byzantine tolerance aggregation mechanism, in which the trust score generation method that is constructed eliminates the operation of filtering or directly discarding abnormal gradients in previous schemes, which can prevent the misjudgment of outliers and preserve seemingly abnormal benign gradients. However, ensuring the privacy, security, and robustness of the aggregated results in both IID and non-IID data scenarios remains a challenging task [25]. The existing privacy-preserving BRFL schemes either ignore non-IID heterogeneity or generate excessive ciphertext aggregation overheads. FedOPCS simultaneously (i) filters out encrypted toxic updates and (ii) smooths non-IID data through CGAN.

Cao [26] developed the Model Poisoning Attacks based on Fake clients (MPAF) framework, demonstrating how attackers can significantly reduce the accuracy of the global model by generating fake clients, effectively countering traditional defence methods. This framework exposes the limitations of current defence mechanisms and calls for more sophisticated countermeasures. Sun [27] proposed the FL-WBC mechanism, which mitigates the impact of model poisoning attacks by introducing parameter space perturbations during local training and simplifies the process of resetting to an unpolluted state. This mechanism offers a practical solution to counteract the effects of poisoning attacks. However, when attackers can adopt more covert and complex strategies to bypass detection mechanisms, further improvements in defense capabilities are needed. Zhang [28] introduced PoisonGAN, utilizing the GAN framework to generate malicious data samples that mislead FL algorithms. This effectively reduces the accuracy of the global model. However, there is a lack of extensive empirical data and validation of practical application scenarios. Li [29] proposed the LoMar algorithm, which quantifies the deviation between model updates submitted by clients and uses statistical methods to shield malicious updates, enhancing the robustness of the model. This method mainly targets attack detection and processing strategies in static environments, while in actual applications, the distribution of nodes and data may change dynamically. Awan [30] used a cosine similarity-based method to evaluate the credibility of each local model update and designed a reputation mechanism to reward or penalize clients based on their contribution to the model training, thereby improving the performance and accuracy of the global model. To ensure the effectiveness of defense mechanisms, frequent data exchange and communication may be necessary to detect and handle malicious nodes. In conclusion, securing FL requires the strategic detection and mitigation of malicious activities, emphasizing strong model aggregation and anomaly detection. Integrating cryptographic methods and developing defense frameworks is essential for combating advanced attacks.

2.2. Handling Non-IID Data in Federated Learning

Since the local datasets of clients are often highly heterogeneous, this presents challenges to the design and performance of FL algorithms. Strategies for handling non-IID data mainly focus on improving model generalization and optimizing communication efficiency. Li [14] proposed the FedProx framework to address system heterogeneity and statistical heterogeneity in FL. Building on this foundation, Wang [31] introduced the Favor framework, which uses an experience-based control strategy to intelligently select client devices for training, effectively alleviating the bias problem caused by non-independent and identically distributed (non-IID) data and accelerating the convergence of the model. To optimize communication efficiency, Sattler [32] proposed the Sparse Ternary Compression (STC) framework, which reduces the amount of data transmitted, showing significant effects, especially in scenarios with large communication overheads. Concurrently, Li [33] emphasized the impact of different data allocation strategies on the performance of FL algorithms in data island scenarios, highlighting the need for tailored approaches rather than a one-size-fits-all solution. Although the papers mentioned in this paragraph showed good performance in improving model generalization and communication optimization, their methods may affect the effectiveness of the algorithm in the face of diverse demands and uneven data distribution in heterogeneous network environments. More experimental evidence is needed to evaluate their applicability in different environments.

Addressing the challenges from both architectural and resource management perspectives, Briggs [34] used a hierarchical clustering method to cluster clients and train specialized models, while Sery [35] introduced the COTAF algorithm to reduce the impact of noise in wireless transmissions. Furthermore, Qu [36] and Mendieta [37] indicated that self-attention mechanisms, such as transformer architectures, demonstrate greater robustness in handling data distribution heterogeneity, validated through experiments. Wang [38] proposed a method considering the heterogeneity of device communication and computational resources, improving resource utilization efficiency and model accuracy through intelligent device sampling and data offloading strategies. Balakrishnan [39] introduced the DivFL algorithm, optimizing global model updates through submodular client selection, enhancing learning efficiency. Zhang [40] introduced a Semi-Supervised Federated Learning (SSFL) model, significantly improving model accuracy in scenarios rich in unlabelled data. Lastly, Zhang [7] proposed BatchCrypt, a system solution for cross-domain federated learning combining homomorphic encryption algorithms. Although the method proposed in this paragraph performs well in specific datasets and scenarios, its applicability and universality in various types of non-IID data and diverse practical application scenarios have not been fully validated.

3. System Model

Figure 2 illustrates the overall structure of the proposed FedOPCS framework. It consists of four main components: synthetic data generation, local training, node aggregation, and node credibility assessment. Specifically, data augmentation is achieved using CGANs to generate synthetic data based on real and random-labeled inputs. These data are then used by clients for local training. The server collects local model updates, filters malicious or low-quality updates through an expectation-based node exclusion mechanism, and updates the global model accordingly. This process helps mitigate the effects of poisoning attacks and ensures robust aggregation in non-IID environments.

Typically, federated learning approaches aim to handle consensus learning tasks in a decentralized manner, where a central server coordinates the global learning objective and multiple devices train models using locally collected data [14]. Consider a network with K local nodes, where each node i has a local dataset $D_i$ of size $D_i$. Nodes connect to a central server and work together to find a global model parametrised by w that minimizes the empirical risk [14,15].

$$\mathrm{Global}\mathrm{Objective}:\mathcal{L}\left(w\right)={\displaystyle \frac{1}{K}}\sum _{i=1}^K{\mathcal{L}}_i\left(w\right)$$

(1)

Specifically, in the federated learning problem, each training sample contains a feature vector x and is labelled in the feature space X and the label space Y. For each available training sample, a federated learning model parameterized by w is considered, aiming to learn a prediction probability vector with empirical risk.

From the perspective of federated learning [14], the global objective in (1) is replaced by local objectives, which can be further expressed as

$${\mathcal{L}}_i\left(w\right)={\mathbb{E}}_{(x,y)\sim D_i}[\ell (w;x,y)]$$

(2)

For node i, the local empirical risk on the dataset is usually measured (e.g., cross-entropy loss), which is defined as follows [41]:

$${\mathcal{L}}_i\left(w\right)=-\sum _{j=1}^C{p}_{ij}log{p}_{\mathrm{model},w}(y=j|x)$$

(3)

where $p_{ij}$ represents the probability of data samples and $p_{\mathrm{model},w}(y=j|x)$ represents the node i’s data distribution on class $j\in \left[C\right]$.

The most commonly used algorithm to solve (2) is FedAvg, where the training consists of multiple rounds of communication. In each round of communication t, the server selects a small number of nodes c to participate in training. Referring to the global model of the previous round, each participating node performs local SGD to optimize its objective [14]:

$$w_i^{t+1}=w_i^t-\eta \nabla {\mathcal{L}}_i\left(w_i^t\right)$$

(4)

where $\eta$ is the learning rate and $\nabla {\mathcal{L}}_i\left(w_i^t\right)$ is the gradient at node i. Equation (4) provides the general principle of SGD optimization, where $w_i^{t+1}$ is the result of the local updates after t epochs of mini-batch SGD. Then, the participating nodes communicate their model updates back to the server, and the server aggregates them and updates the global model as follows [14]:

$$w^{t+1}={\displaystyle \frac{1}{K}}\sum _{i\in S}{w}_i^{t+1}$$

(5)

When malicious clients exist, Byzantine bus aggregation rules (such as Krum and Zeno) should be used to improve convergence stability [16].

4. Node Credibility Assessment in Federated Learning Based on CGAN Data Augmentation

4.1. Optimal Aggregation with Node Credibility Assessment

Optimal Aggregation with Node Credibility Assessment requires a special node selection strategy to identify nodes that may adversely affect global updates. By excluding potentially detrimental local updates and reducing the probability of selecting such nodes, nodes that contribute significantly to the reduction in global loss are more likely to be selected. We use the inner product as a criterion to check nodes. In each communication round t, the inner product of the local gradient and the global gradient is considered, where nodes with negative inner products will affect the convergence speed of the model. However, it is not easy to exclude nodes that are harmful to the global model because this will affect the relationship between the local gradient and the global gradient.

To find the best subset of local updates to aggregate, we use a test loss to ensure that excluding some local updates enables the global update to perform better in terms of model convergence, thus ensuring that

$$\mathcal{L}\left(w_{\mathrm{agg}}\right)\le \mathcal{L}\left(w\right)$$

(6)

where $\mathcal{L}\left(w_{\mathrm{agg}}\right)$ is the test loss after removing the detrimental local nodes from the node set. We use an aggregation strategy based on the inner product measure to better aggregate the local updates in each round. The algorithm eliminates harmful local updates to find the optimal local update subset [16]. Specifically, for a set of participating nodes in each global communication round t, the server removes each local update one by one to obtain the global gradient after removal. If the expected global gradient after excluding one local update is higher than when all local updates are retained, then this local update can be temporarily considered detrimental. Further test losses are needed to further determine the impact of this local update on the global model. By comparing the test loss with the obtained global weights $w_{\mathrm{agg}}$ and w, if the condition holds, then the node’s removal is beneficial to the global model. If the test loss check fails, which indicates that the current node has a significant influence on the global weights, the server will keep all local updates. This process is repeated until no detrimental local updates are found. To mitigate the influence of harmful local updates and select credible nodes, we design an inner product-based aggregation strategy, as described in the Algorithm 1. This algorithm iteratively removes harmful local updates based on their inner product with the global gradient and evaluates the impact on test loss to ensure model robustness.

Algorithm 1 Aggregation Based on Inner Product Metrics

1: procedure Aggregation($W_{\mathrm{global}},W_{\mathrm{local}},\mathrm{thresh}$) 2:     while $\mathrm{thresh}>0$ do 3:         Select a local update to remove 4:         if The removed update improves the global gradient then 5:            Remove the update and update thresh 6:         end if 7:     end while 8: end procedure

Excluding potentially detrimental local updates can achieve a greater reduction in the expected loss in each round. Furthermore, we can prefer nodes that contribute more significantly to the global weights and are more credible when selecting nodes before weight aggregation. Therefore, we adopt a probabilistic node selection design based on the optimal aggregation algorithm and adjust the probability of each node being selected according to its contribution and credibility in each round of communication.

We reduce the selection probability of nodes that are ultimately removed and marked in the optimal aggregation algorithm process, and correspondingly increase the selection probability of the remaining nodes. The selection probability of nodes is determined using the following formula [42]:

$$P_i^{t+1}=P_i^t·(1-\delta )$$

(7)

where $P_i^t$ represents the probability of node i being selected in the t-th global round and $\delta$ represents the decrease in the probability of the node being selected in the next round. The parameter selection for $\delta$ and $\alpha$ follows $\delta ={\displaystyle \frac{\alpha }{x}}$. The parameter $\alpha$ controls the degree of probability reduction given a ratio x. For example, a large value of $\alpha$ brings about a drastic reduction because the probability reduction occurs over a large range as x increases in a very small interval, causing the node selection probability to drop very quickly as x increases. At the same time, a larger $\alpha$ makes the node selection very sensitive to misidentification. Ultimately, we have

$$P_i^{t+1}={\displaystyle \frac{P_i^t·(1-\delta )}{{\sum }_j{P}_j^t·(1-\delta )}}$$

(8)

To provide stable privacy protection for the entire federated learning process and prevent information exchange between different clients, we introduce HE algorithms into the entire federated learning framework. HE is a cryptographic technique that allows for computation to be performed on encrypted data without the need to decrypt these data. HE has important advantages in protecting data privacy and security, especially in the fields of cloud computing and distributed computing. HE can be divided into two types: partially homomorphic encryption (PHE) and fully homomorphic encryption (FHE). PHE only supports a limited number of operations, while FHE can handle any number of operations. In FedOPCS, the server’s sole cryptographic task is to aggregate (add) encrypted model updates and return a single ciphertext; no ciphertext-level multiplications are needed. Consequently, the Paillier cryptosystem [43], which supports exact additive homomorphism, is sufficient, enabling a one-round aggregation protocol and eliminating the multi-round masking steps required by Secure Aggregation [42]. Relative to lattice-based schemes such as BFV or CKKS [44,45], Paillier produces ciphertexts that are roughly two to three times smaller and delivers over five-fold faster encryption on ARM processors while still providing at least 128-bit security under the decisional composite residuosity assumption. Moreover, because Paillier operates natively on integers, fixed-point encoded gradients can be encrypted without the numerical errors inherent to CKKS. These properties explain why the reported cryptographic overheads—0.23 s per encryption and 17.6 s per decryption (Section 5.2.4)—are negligible compared with local SGD computation and remain acceptable for resource-constrained edge devices. This paper selects the Paillier cryptosystem, a public key encryption scheme with additive homomorphic properties, as the encryption algorithm for federated learning.

In this encryption algorithm, a pair of public and private keys are first generated according to the following formula. Let p and q be two large prime numbers satisfying $gcd(p,q)=1$, where gcd represents the greatest common divisor. Then, take a random number g such that $g\in {\mathbb{Z}}_n^{*}$ and $g^2\not\equiv 1(modn)$. Finally, take n as the public key and $\lambda$ as the private key, where $n=p·q$ and ${\mathbb{Z}}_n^{*}$ denotes the multiplicative group of integers modulo n.

When each round of client training is completed, the client encrypts the model and sends it to the server. The encryption algorithm first takes a random number r, where $r\in {\mathbb{Z}}_n^{*}$ and $r\not\equiv 0(modn)$. The ciphertext is calculated using the Paillier cryptosystem [43]:

$$c=(m+r)·g^k(mod{n}^2)$$

(9)

where c is the ciphertext and m is the client model parameter. After encryption, client i sends the ciphertext $c_i$ to the server, and the server integrates the ciphertexts, resulting in

$$c_{\mathrm{agg}}=\prod _{i=1}^K{c}_i(mod{n}^2)$$

(10)

Then, the aggregated parameters are returned to the client, and the client decrypts and calculates the plaintext as follows:

$$m=\mu ·c^{\lambda -1}(modn)$$

(11)

where $\mu$ is the modular inverse of n modulo $\lambda$ and m refers to the decrypted plaintext. The encryption and decryption processes are all performed locally, which effectively enhances data privacy, and the data exchange behaviour between clients is also strongly curbed.

4.2. Data Augmentation with Conditional GANs

Adopting a probability-based node selection algorithm for optimal aggregation can effectively integrate credible nodes and exclude interference from unfavourable local updates. However, a data augmentation algorithm can also be designed to reduce the damage to local nodes caused by non-IID data and malicious attacks, enhancing the robustness of the model. Conditional Generative Adversarial Networks (CGAN) can generate data that are as close as possible to the actual data distribution based on conditional variables. The power of CGAN lies in its conditional generation mechanism, which can produce images that are visually similar to the original dataset based on given category labels. These synthetic images are highly consistent with the statistical characteristics of the real data in terms of their color distribution, texture features, and object shape. This not only expands the sample size of each category in the original dataset, especially for rarer categories, but also effectively reduces the risk of overfitting when the model faces an unbalanced dataset and improves the model’s ability to learn different category features. Data augmentation with CGAN enhances the model’s adaptability and robustness to different data distributions. CGAN can simulate different data distribution characteristics to train the model to better adapt to changes in data distribution.

The task of the generator network is to produce data that are as realistic as possible, while the discriminator network tries to distinguish between the generated data and real data. In CGAN, conditional information is merged with the noise input of the generator network G to guide the generator to produce data of a specific category; in the discriminator network D, conditional information is also used to assist in judging whether the input data are real or generated. The model description of CGAN is as follows [46]:

For the discriminator,

$$\underset{D}{min}{\mathbb{E}}_{(x,y)\sim \mathrm{Data}}[logD(x,y)]+{\mathbb{E}}_{(z,y)\sim \mathrm{Noise}}[log(1-D(G(z,y),y))].$$

(12)

For the generator,

$$\underset{G}{min}{\mathbb{E}}_{(z,y)\sim \mathrm{Noise}}[log(1-D(G(z,y),y))].$$

(13)

Unlike FedPNS’s node selection based solely on contribution, this section further combines test loss and the inner product threshold to counteract poisoning through dual judgment. The model implies that the parameters of the generator G are minimized and the parameters of the discriminator are maximized to ensure that the data generated by the generator are as realistic as possible and the discriminator is as accurate as possible. ${\mathbb{E}}_D[·]$ is the expected loss function of the discriminator, where Data is the joint distribution of real data and conditions, and the goal of the discriminator is to maximize this expectation, i.e., for real data samples x, the judgment probability given by the discriminator is close to 1. ${\mathbb{E}}_G[·]$ is the expected loss function of the generator, where Noise is the distribution of prior noise z and conditions y, and the goal of the generator is to minimize this expectation, i.e., for the data generated by the generator according to the noise, the discriminator obtains a judgment probability close to 0. After training the generator G, synthetic data can be obtained according to the label. To enhance robustness under non-IID conditions, we further develop a data augmentation-driven node selection process, which is detailed in the Algorithm 2. This algorithm combines probabilistic node selection with CGAN-based data augmentation, adjusting node participation probabilities based on their contribution credibility across training rounds.

Algorithm 2 Node Selection with Data Augmentation

1: procedure NodeSelection($\mathrm{Server},\mathrm{Nodes},\mathrm{T}$) 2:     $\mathrm{Server}\mathrm{initializes}\mathrm{selected}\mathrm{nodes}\mathrm{probability}{P}_i={\displaystyle \frac{1}{K}}\mathrm{for}\mathrm{all}\mathrm{nodes}i$ 3:     for t = 1 to T do 4:         $\mathrm{Sample}\mathrm{the}\mathrm{selected}\mathrm{nodes}\mathrm{based}\mathrm{on}{P}_i$ 5:         $\mathrm{Server}\mathrm{sends}\mathrm{model}\mathrm{parameters}w\mathrm{and}\mathrm{conditions}y\mathrm{to}\mathrm{nodes}\mathrm{with}\mathrm{HE}$ 6:         $\mathrm{Nodes}\mathrm{use}w\mathrm{and}y\mathrm{to}\mathrm{train}\mathrm{local}\mathrm{model},\mathrm{and}\mathrm{send}\nabla {\mathcal{L}}_i\left(w\right)\mathrm{back}\mathrm{to}\mathrm{the}\mathrm{server}$ 7:         $\mathrm{Server}\mathrm{executes}\mathrm{Algorithm}$ 2 8:         $\mathrm{Server}\mathrm{updates}\mathrm{selected}\mathrm{probability}\mathrm{using}\left(7\right)$ 9:     end for 10: end procedure

5. Experiment

5.1. Experimental Setup

We selected two widely used datasets for evaluation: MNIST and CIFAR10. The MNIST dataset consists of 60,000 $28\times 28$ greyscale images of handwritten digits, used as the training set, with an additional 10,000 images as the test set. The CIFAR10 dataset contains 50,000 $32\times 32$ color images, divided into 10 classes, with 5000 images per class for training, and 10,000 images for the test set. We chose FedAVG, Fedprox, and FedPNS as the baseline methods to measure the effectiveness of our proposed method. During local training, we employed the SGD optimization algorithm, setting the momentum to 0.9, the initial learning rate to 0.01, the batch size to 64, and the number of local epochs to 10. To simulate real-world data distribution, we used a Dirichlet distribution to distribute client data and adjusted the proportions based on the number of samples each client already has, ensuring that the number of samples in each class on the client does not exceed the ratio of the total number of samples of that class in the entire dataset to the number of clients.

5.2. Experimental Results and Analysis

5.2.1. Model Performance Evaluation

Under the given experimental setup, we first explored the impact of different numbers of clients and the proportion of random selection on model performance. Through extensive testing on the MNIST and CIFAR10 datasets, we found that, based on the MNIST dataset, the model performance under different configurations did not show significant differences. This may be due to the simplicity and uniformity of the data distribution of the MNIST dataset, which results in a smaller impact of different client selection strategies on the final model performance. In contrast, on the more complex and diverse CIFAR10 dataset, we observed that when 20 clients were selected for training and the proportion of clients randomly selected in each round of communication was 0.5, the model performance reached its optimal state. This finding indicates that appropriate client selection and communication frequency play an important role in enhancing model performance.

The 20-client setting delivers the best overall performance because it strikes a four-way balance. As shown in Figure 3, we analyze the impact of different client numbers and selection ratios on model performance across datasets. First, once the client pool grows much beyond 20, the non-IID data diversity inflates gradient variance and degrades accuracy, whereas with 20 clients, this inflation is still contained [47]. Second, sampling half of those clients per round creates a weak cyclic schedule that converges faster than purely random sampling; adding more clients lowers each participant’s selection frequency and slows learning [48]. Third, our inner-product credibility filter remains reliable only up to about 20–30 clients; larger pools increase the false-rejection rates, and recent Byzantine-robust studies likewise show that “small-yet-reliable” cohorts are easier to defend [49]. Finally, 20 clients impose a computation load that aligns with the 17-second Paillier-2048 encryption overhead; beyond this point, extra local training is masked by encryption and communication delays, so wall-clock convergence actually slows [50]. These combined factors explain why 20 clients yield the global optimum in our experiments.

5.2.2. Robustness Testing Against Adversarial Attacks

To verify the robustness of our proposed method against malicious attacks, we compared it with several baseline methods. In the model poisoning attack test, we simulated a scenario where an attacker introduces malicious changes during the model parameter update phase. Figure 4 presents the comparative performance of various models under model poisoning attacks. The experimental results show that, due to the adoption of a refined node selection mechanism, our method can effectively identify and exclude nodes that negatively impact the global model update, thus maintaining the model’s convergence performance. Moreover, when facing both model poisoning and data poisoning attacks, our method significantly enhances the model’s adaptability to non-IID data by introducing a CGAN for data augmentation, enhancing the system’s anti-interference capabilities.

5.2.3. Analysis of the Impact of the Poisoning Ratio

To further explore the specific impact of poisoning attacks on model performance, we specifically analysed the performance of the model under different poisoning ratios for the CIFAR10 dataset. In the experiment, we set two different poisoning ratios of 0.2 and 0.5 to simulate different attack intensities of the attacker. The performance of various methods under different ratios of model poisoning is summarized in

Table 1. Performance of model poisoning with different ratios of nodes.

Method	Ratio = 0.2		Ratio = 0.5
	Min	Max	Min	Max
FedAVG	34.3	41.6	30.9	33.8
Fedprox	36.5	42.7	31.7	35.6
FedPNS	68.3	68.5	60.1	60.2
Ours	72.7	72.8	71.4	71.6

. The results show that at a lower poisoning ratio (0.2), the performance of our method is slightly better than the baseline methods; when the poisoning ratio is increased to 0.5, the performance advantage of our method becomes more apparent. When both model and data poisoning attacks are present, the comparative results are shown in

Table 2. Performance of model combined with data poisoning with different ratios of nodes.

Method	Ratio = 0.2		Ratio = 0.5
	Min	Max	Min	Max
FedAVG	43.1	44.7	40.2	41.8
Fedprox	45.0	46.9	43.3	44.6
FedPNS	58.4	58.5	55.6	55.7
Ours	61.6	61.8	60.3	60.4

. Table 2 highlights the classification accuracy of different methods under combined poisoning scenarios, showing that our method consistently achieves better robustness even when both types of attacks coexist. The impact of varying poisoning ratios and types on model performance is detailed in Figure 5. This result verifies that our method can still maintain high robustness in the face of more intense attacks, mainly due to the innovative design of our method in node selection and data enhancement.

5.2.4. Time Consumption and Rationale for Paillier HE

In the federated learning strategy adopted in this paper, the encryption and communication time are negligible compared to the model training. The encryption time in the homomorphic encryption process is positively correlated with the size of the plaintext. The experimental code was implemented in Python 3.8 environment, utilizing PyTorch version 1.13.1 for model training and evaluation. The key parameter size of the Paillier encryption algorithm used in this paper is set to 2048. During the experiment, taking the images of the CIFAR10 dataset as an example, the image size is $32\times 32\times 3$, occupying 3 KB of Python memory space, with an average encryption time of 0.23 s. The convolutional neural network used in this experiment occupies about 230 KB of Python memory space, with an average plaintext encryption time of 17.63 s. Compared to the training process, which lasts for several hours, the encryption process is negligible and does not affect the entire training process.

6. Conclusions

In this study, we explore the robustness of FL in facing the heterogeneous threats of model contamination and data contamination and propose a new algorithmic framework to enhance the security and robustness of FL. By introducing weight restrictions and measurement criteria, our method can effectively identify and mitigate the impact of malicious actors, while considering the challenges brought about by non-IID data. In addition, we also adopt a data generation method based on CGAN and a model aggregation method combining adaptive weighting and homomorphic encryption technology to further improve the robustness and privacy protection of the system.

Extensive experiments on MNIST and CIFAR-10 show that FedOPCS outperforms baseline methods in poisoning environments while maintaining its efficiency. By restoring symmetry in gradient distributions through CGAN-based augmentation and homomorphic encryption, our framework ensures both robustness and the preservation of privacy.

Our work addresses the critical challenge of preserving symmetry in FL under adversarial and non-IID conditions. There are still some limitations and directions for future work. First, the performance of our algorithm in handling extremely non-IID data distribution needs further verification. Second, although homomorphic encryption technology improves security, it also increases the computational overhead and communication overhead, so finding the best balance between efficiency and security is an important research direction. Finally, future directions include extending this framework to other asymmetric scenarios, such as cross-device FL with heterogeneous architectures.