Addressing Non-IID with Data Quantity Skew in Federated Learning

Abstract

Non-IID is one of the key challenges in federated learning. Data heterogeneity may lead to slower convergence, reduced accuracy, and more training rounds. To address the common Non-IID data distribution problem in federated learning, we propose a comprehensive dynamic optimization approach based on existing methods. It leverages MAP estimation of the Dirichlet parameter $\beta$ to dynamically adjust the regularization coefficient $\mu$ and introduces orthogonal gradient coefficients ${\mathbf{\Delta }}_i$ to mitigate gradient interference among different classes. The approach is compatible with existing federated learning frameworks and can be easily integrated. Achieves significant accuracy improvements in both mildly and severely Non-IID scenarios while maintaining a strong performance lower bound.

1. Introduction

Federated Learning (FL) [1,2] is an emerging distributed machine learning framework that enables knowledge sharing and user privacy protection during model training without requiring the upload of raw data. However, the statistical heterogeneity across clients, referred to as non-independent and identically distributed (Non-IID) characteristics, hinders the development of federated learning. In federated settings, datasets from a large number of clients exhibit unique attributes and usage patterns, characterized by different distributions, such as feature skew and label skew, which result in training datasets with statistical heterogeneity. Datasets with Non-IID characteristics slow down the convergence, reduce accuracy, and even result in model divergence [3]. In scenarios with high safety requirements, such as autonomous driving, model divergence caused by Non-IID data may lead to serious traffic accidents and even result in loss of life and property.

Previous works have proposed many normalization-based approaches to address the heterogeneity in federated learning. For example, FedProx [4], FedBN [5], SCAFFOLD [6], and MOON [7] have achieved relatively good performance compared to FedAvg [1]. However, these approaches struggle to adapt to datasets with varying distributions. A study [8] found that a method performing well on an IID dataset often performs worse on Non-IID data. Conversely, a method that performs better on Non-IID data often falls short on IID data, sometimes even underperforming vanilla approaches such as FedAvg, as shown in Figure 1.

Brain-inspired neural networks enhance performance under Non-IID. Mechanisms incorporating synaptic plasticity [9,10] can efficiently alleviate the catastrophic forgetting that exists in sequential learning. Moreover, brain-inspired neural networks have a vast exploration space due to synaptic plasticity and neural dynamics. In artificial neural networks (ANNs), synaptic plasticity can be interpreted as the adjustable connections between the neurons in the different layers. Synapses can be strengthened or weakened to consolidate learning and memory [11]. This process is very similar to the adjustment of model weight in the training. Research on synaptic plasticity in neural networks has made significant progress in mitigating the catastrophic forgetting [12,13,14]. Similarly, spike-timing-dependent plasticity (STDP) [15] has been applied in spiking neural networks (SNNs), achieving significant progress in reducing power consumption and also demonstrating their effectiveness in mitigating catastrophic forgetting [9,16]. Unfortunately, the integration of brain-inspired algorithms and loss function regularization remains underexplored in FL.

In this paper, we propose an algorithm that jointly incorporates a brain-inspired approach and regularization to tackle the Non-IID problem in FL. The proposed method is expected to achieve performance comparable to advanced approaches such as FedProx under IID settings, and to brain-inspired approaches such as FedNACA under Non-IID settings, as illustrated by the blue dashed line in Figure 1.

Our contributions are listed as follows:

• We propose an approach introducing MAP-based dynamic adjustment for regularization coefficient $\mu$ in the objective function and orthogonal gradient modulation to adapt to varying degrees of Non-IID distributions. The approach is compatible with existing federated learning frameworks and can be easily integrated.
• We propose a method to evaluate the degree of Non-IID in local datasets and mitigate model drift by adjusting the coefficient $\mu$ of the regularization term in the objective function according to the data heterogeneity.
• We generated three datasets with different degrees of Non-IID: IID, mild Non-IID, and severe Non-IID. Extensive experimental results demonstrate that our approach achieves significant accuracy improvements in both mildly and severely Non-IID scenarios while maintaining a strong performance lower bound.

2. Related Works

2.1. Regularization-Based Approaches

Regarding Non-IID issues in FL, prior researchers have proposed various solutions, with regularization-based and brain-inspired approaches showing particularly strong performance. In machine learning, regularization-based methods were initially applied to avoid overfitting by adding a term to the objective function, thereby improving model generalization. FedProx [4] incorporated a proximal term into the loss function to alleviate statistical heterogeneity. FedDyn [17] introduced a dynamic regularization term, adding two components—a norm of the model and an inner product between the gradient and the model parameters—to enhance robustness in heterogeneous FL. Similar to these methods, SCAFFOLD [6] and MOON [7] essentially modify the loss function of neural networks to mitigate performance degradation under Non-IID settings. In HFMDS [18], the authors proposed class-relevant feature matching data synthesis by sharing partial features between the server and clients. Similarly, in the study [3], the authors also advocated sharing data between the server and clients to address Non-IID issues in FL. However, such approaches deviate from the original intent of FL. To avoid violating privacy constraints in FL, CCVR [19] shares virtual representations from higher layers of deep networks near the classifier instead of raw data to calibrate the bias of the global model caused by Non-IID. Our approach is regularization-based but differs from the above methods in that it can adapt to varying degrees of Non-IID and can be integrated with other approaches.

Adaptive/personalized FL represents another research focus for addressing the challenges posed by Non-IID in FL. Adaptive FL improves the adaptability of the global model by adjusting optimization algorithms, learning rates, aggregation strategies, and the ratio of client participation. For instance, FedNova [20] mitigates the issue of “bias accumulation” caused by inconsistent numbers of local training steps across clients, while FedOpt [21] employs adaptive optimizers at the server side to enhance the convergence speed and stability of the global model. In contrast, personalized FL aims to tackle the difficulty that a single global model cannot meet the diverse needs of clients under heterogeneous data distributions. It provides each client with a personalized model while maintaining collaborative learning. pFedMe [22] adopts a regularization-based approach such that each client learns its personalized model on top of the global model.

Non-IID problems can be interpreted as a form of catastrophic forgetting [23] in FL, and approaches designed to address catastrophic forgetting can be adapted to solve Non-IID problems in federated learning.

2.2. Biological-Inspired Approaches

Due to the intrinsic complexity of the Non-IID issues, FL has yet to establish effective and reliable solutions to address them. In contrast, intelligent biological systems do not exhibit catastrophic forgetting when learning multiple tasks, and the underlying mechanisms are continuously being explored. Biologically inspired approaches, which are capable of simulating intelligent behavior, adapting to complex and dynamic environments, and demonstrating strong robustness, as well as broad applicability within artificial intelligence, may provide valuable insights for addressing the challenges associated with Non-IID.

In recent years, researchers have proposed various heuristic algorithms inspired by biological intelligence, including mechanisms that mimic memory and synaptic consolidation, to address catastrophic forgetting [24,25], attracting significant attention. For example, ANPyC [10], an adversarial neural pruning and synaptic consolidation algorithm inspired by the memory mechanisms of mammals, provides a long-term solution to forgetting in multitask learning. Synaptic plasticity during learning is governed by principles such as Hebbian theory [26] and STDP [15]. SNNs mimic the processing and encoding of sensory information in the brain [27], and techniques developed for SNNs have been transferred to ANNs, achieving strong performance. For instance, ANNs with global neuromodulation [9] effectively mitigate forgetting issues in incremental learning. The study [28] confirmed that feedback connections in the brain deliver error signals. Similarly, the study [29] proposed synaptic intelligence, incorporating biological complexity to address forgetting of learned knowledge by introducing an importance measure for each synapse, which estimates the sensitivity of the loss function to that synapse. Synapses with lower importance to previous knowledge are adjusted to prevent interference with existing knowledge during the training of a new task.

Studies have shown that biological neuromodulation, when combined with traditional approaches, can also reduce forgetting during the acquisition of new knowledge. Inspired by the memory system of Drosophila, the study [30] proposed a solution that incorporates two additional modules compared to traditional artificial intelligence: stability protection and active forgetting. These modules enable the system to actively regulate the forgetting rate of learned knowledge while maintaining compatibility with different architectures. In the work [31], focusing on spatial behavior and drawing inspiration from neuromodulation, proposed a framework with three distinct granularities: plastic neurons based on Hebbian learning (fine-grained), layers with dropout (medium-grained), and network layers with self-regulating learning rates (coarse-grained). The authors [32] divided the neural network into multiple modules based on connection costs, where neurons within each module employ biological neural regulation mechanisms for learning, and control gates are opened to allocate a new module for acquiring new knowledge. Similarly, the authors [33] combined meta-learning with neuromodulation, dividing neural networks into multiple modules with module connections controlled by gates, as in [32]. Other works, including diffusion-based neural modulation modularization [34], differentiable neural modulation [35], neural modulation with local learning [36], and dynamic learning with neuromodulation [37], have introduced biological neuron mechanisms into neural networks to enhance performance.

Our approach differs from the aforementioned methods. Regularization typically involves adding a term to the loss function, while biological neural modulation is usually applied during weight updates; thus, the two processes operate without mutual interference. The main idea of our method is to effectively integrate regularization with biological neural modulation and validate its effectiveness in mitigating the Non-IID problem in a federated learning framework.

3. System Model

In this section, we provide a detailed description of the system model, including the federated learning architecture, The estimation of the degree of Non-IID using the MAP is based on the number of classes in the local dataset, the coefficient of the regularization term $\mu$ in the objective function, and the orthogonal tensor coefficients in the local updating stage.

3.1. Federated Learning

Let $\mathcal{N}=\{1,2,\dots ,N\}$ be a set of indices that describe N clients. In each communication round, a fraction of the clients $\mathcal{C}\subseteq \mathcal{N}$ is selected to train the model. The sampling ratio $c={\displaystyle \frac{\left|\mathcal{C}\right|}{\left|\mathcal{N}\right|}}$ indicates the fraction of clients chosen per communication round due to resource constraints. Each participant i holds a local dataset ${\mathcal{D}}_i\subseteq \mathcal{D}$ containing $\left|\mathcal{D}i\right|$ samples in the form of $(x_k,y_k)$, where $\mathcal{D}$ denotes the whole dataset, $x_k$ denotes the input vector, and $y_k$ denotes the ground truth. Let $q_i=\{q_i^1,q_i^2,\dots ,q_i^r\}$ denote the sample quantity of each class in the client i, r denote the number of classes, and $|{\mathcal{D}}_i|={\sum }_{n=1}^r{q}_i^n$.

Let $w_i^g$ denote the global model and $w_i^j$ denote the local model on the client j in round i. Every round comprised five iterative steps: broadcasting the global model, selecting participants, the local update, uploading local model, and aggregating the model [38], as shown in Figure 2. The coefficient $\mu$ is adjusted adaptively for heterogeneous data distributions, and the normalized loss function [4] is adopted in this paper, addressing the dataset with Non-IID characteristics better. The loss function follows as

$${\mathcal{L}}^{*}\left(w_i^j\right)=\mathcal{L}\left(w_i^j\right)+{\displaystyle \frac{\mu }{2}}\parallel w_i^g-w_i^j,\parallel$$

(1)

where $\parallel •\parallel$ denotes the distance norm between two models. The coefficient of the proximal term $\mu \ge 0$ is related to the degree of Non-IID in the local datasets. When $\mu =0$, the method reduces to FedAvg. A larger value of $\mu$ is assigned when the local datasets exhibit a higher degree of Non-IID. The coefficient $\mu$ is adaptively adjusted according to the measured degree of Non-IID. Details of the assignment strategy are provided in Section 3.2.

3.2. Evaluation for the Degree of Non-IID of Local Dataset

In Section 3.2, we first describe the characteristics of Non-IID from the perspective of class sample distribution. Then, we formulate the degree of Non-IID in the local datasets using MAP estimation with a Dirichlet prior $Dir\left(\overrightarrow{\mathit{\beta }}\right)$. Finally, we discuss the relationship between the proximal term coefficient $\mu$ in the objective function and the parameter $\overrightarrow{\mathit{\beta }}$ in the Dirichlet distribution.

In real-world scenarios, the sample distribution across clients is typically Non-IID. In this paper, we focus on the perspective of class sample quantities, as illustrated in Figure 3. To illustrate the mechanism for quantifying the degree of Non-IID, we present an example containing a three-class dataset, as shown in Figure 4. Each dimension corresponds to a category, and each data point in the figure represents the class-wise sample distribution for a client. In Figure 4a, $x_1$ exhibits a nearly uniform distribution—34%, 31%, and 35% for classes $c_1$$,c_2$, and $c_3$, respectively, reflecting typical IID settings. Figure 4b illustrates a mild Non-IID, in which certain clients maintain balanced sample distributions, whereas others display pronounced sample count imbalances. In contrast, Figure 4c depicts a severe Non-IID, wherein client $x_2$ exhibits an exceptionally skewed class distribution, with 89% of its samples concentrated in class $c_3$. Consequently, classes $c_1$ and $c_2$ are markedly underrepresented, leading to insufficient feature learning for these classes during local training. These scenarios can be quantitatively modeled using a Dirichlet distribution, as demonstrated in Figure 5. Figure 5 shows heatmaps of client distribution counts obtained by projecting the data from Figure 4 onto the x–y plane. In the IID setting, as Figure 5a, most clients cluster around the center area, indicating similar sample proportions across all classes. In the mild non-IID setting, as Figure 5b, the clients are more evenly dispersed, reflecting the coexistence of balanced and imbalanced distributions. In the severe non-IID setting, as Figure 5c, the majority of clients concentrate near the corners, corresponding to highly skewed distributions. These patterns are parameterized by the Dirichlet distribution with $\overrightarrow{\mathit{\beta }}$, specifically, larger $\overrightarrow{\mathit{\beta }}$ values produce centralized distributions, as in Figure 5a, while smaller $\overrightarrow{\mathit{\beta }}$ values yield corner-concentrated distributions, as in Figure 5c. It should be noted that the parameter $\overrightarrow{\mathit{\beta }}$ of the Dirichlet distribution is r-dimensional. Assuming isotropic Dirichlet sampling, where all components of $\overrightarrow{\mathit{\beta }}$ take identical values, most studies consequently employ a single scalar parameter to characterize the degree of Non-IID.

We assume that the class quantities on each client follow a Dirichlet distribution $Dir\left(\overrightarrow{\mathit{\beta }}\right)$, where $\overrightarrow{\mathit{\beta }}$ is an r-dimensional vector corresponding to the number of classes. All components of $\overrightarrow{\mathit{\beta }}$ are assumed equal, ${\beta }_1={\beta }_2=\cdots ={\beta }_r=\beta$, indicating an isotropic Dirichlet prior. The size of each local dataset is sampled from a uniform distribution, $|{\mathcal{D}}_i|\sim U({\displaystyle \frac{\left|\mathcal{D}\right|}{N}},\lambda )$, where $0\le \lambda \le {\displaystyle \frac{\left|\mathcal{D}\right|}{N}}$ is a variance parameter controlling the degree of Non-IID. The distribution becomes IID when $\lambda =0$ and becomes increasingly Non-IID as $\lambda$ increases.

In local training, the sample quantities are known to each client, while the parameter $\overrightarrow{\mathit{\beta }}$ of the Dirichlet distribution is unknown. The proximal term $\mu$ in the objective function significantly affects the performance of federated learning, with larger values typically assigned under severe Non-IID conditions and smaller values under near-IID conditions. We use the parameter $\overrightarrow{\mathit{\beta }}$ of the Dirichlet distribution to measure the degree of Non-IID. To this end, we employ MAP estimation to infer the posterior distribution of $\overrightarrow{\mathit{\beta }}$ given the observed class quantities as prior information.

The reason for using MAP is that the sample size provided as prior knowledge is too small during local training. Under such limited prior knowledge, MAP tends to outperform maximum likelihood estimation.

To measure the degree of Non-IID of the dataset, we treat sample quantity as a random variable. Assumption that the sample quantity ${\mathit{q}}_i$ for each participant follows $Dir\left({\mathit{x}}_i\right|\overrightarrow{\mathit{\beta }})$ distribution, as shown below.

$${\mathit{q}}_i\sim Dir\left({\mathit{x}}_i\right|\overrightarrow{\mathit{\beta }})={\displaystyle \frac{1}{B\left(\overrightarrow{\mathit{\beta }}\right)}}\prod _{j=1}^r{x_i^j}^{{\beta }_j-1},$$

(2)

where $B\left(\overrightarrow{\mathit{\beta }}\right)$ denotes the Beta function, $\overrightarrow{\mathit{\beta }}=({\beta }_1,{\beta }_2,\dots ,{\beta }_r)$ denotes the parameter of the Beta function, ${\beta }_i>0$ and ${\beta }_i=\beta$, for all $i\in \{1,2,\dots ,r\}$. $x_i^j$ is a prior probability representing the percentage of class j in the client i, and $x_i^j={\displaystyle \frac{q_i^j}{|{\mathcal{D}}_i|}}$, for all $j\in \{1,2,\dots ,r\}$. Parameter ${\beta }_i$ of prior distribution follows Gamma distribution,

$${\beta }_i\sim Gamma(\beta ;a,b)={\displaystyle \frac{b^a}{\Gamma \left(a\right)}}{\beta }^{a-1}{e}^{-b\beta },$$

(3)

where $a,b$ are hyperparameters and $\beta$ is a random variable with $a>0, b>0, \beta >0$. To better fit for the Non-IID setting, we choose $a={\displaystyle \frac{1}{2}}, b=1$.

Posterior estimation of parameter $\overrightarrow{\mathit{\beta }}$ can be expressed as follows:

$$P\left(\overrightarrow{\mathit{\beta }}\right|{\mathit{x}}_i)={\displaystyle \frac{P\left({\mathit{q}}_i|\overrightarrow{\mathit{\beta }}\right)P\left(\overrightarrow{\mathit{\beta }}\right)}{P\left({\mathit{q}}_i\right)}}.$$

(4)

The computation of the denominator $P\left({\mathit{q}}_i\right)$ involves a complex integral, which makes it difficult to solve directly. Considering that $P\left({\mathit{q}}_i\right)$ is the constant with respect to $\overrightarrow{\mathit{\beta }}$, and based on the relationship between posterior $P\left(\overrightarrow{\mathit{\beta }}\right|{\mathit{q}}_i)$ and the product $P\left({\mathit{q}}_i|\overrightarrow{\mathit{\beta }}\right)P\left(\overrightarrow{\mathit{\beta }}\right)$, we adopt the MAP to estimate the parameter $\overrightarrow{\mathit{\beta }}$. The derivation is as follows:

$$\begin{array}{cc}\hfill P\left(\overrightarrow{\mathit{\beta }}\right|{\mathit{x}}_i)\propto P\left({\mathit{x}}_i|\overrightarrow{\mathit{\beta }}\right)P\left(\overrightarrow{\mathit{\beta }}\right)& ={\displaystyle \frac{1}{B\left(\overrightarrow{\mathit{\beta }}\right)}}{\left(\prod _{j=1}^r{x}_j\right)}^{\beta -1}{\left({\displaystyle \frac{b^a}{\Gamma \left(a\right)}}{\beta }^{a-1}{e}^{-b\beta }\right)}^r\hfill \end{array}$$

(5)

$$\begin{array}{cc}\hfill & ={\displaystyle \frac{\Gamma \left(r\beta \right)}{\Gamma {\left(\beta \right)}^r}}{\left(\prod _{j=1}^r{x}_j\right)}^{\beta -1}{\left({\displaystyle \frac{1}{\Gamma \left(0.5\right)}}{\beta }^{-0.5}{e}^{-\beta }\right)}^r.\hfill \end{array}$$

(6)

Let

$$\begin{array}{cc}\hfill \mathbf{L}\left(\beta \right)& =log\left({\displaystyle \frac{\Gamma \left(r\beta \right)}{\Gamma {\left(\beta \right)}^r}}{\left(\prod _{j=1}^r{x}_j\right)}^{\beta -1}{\left({\displaystyle \frac{1}{\Gamma \left(0.5\right)}}{\beta }^{-0.5}{e}^{-\beta }\right)}^r\right)\hfill \end{array}$$

(7)

$$\begin{array}{cc}\hfill & =log\Gamma \left(r\beta \right)-rlog\Gamma \left(\beta \right)+(\beta -1)\sum _{j=1}^{r}log{x}_j-r\left(log\Gamma \left(0.5\right)+{\displaystyle \frac{1}{2}}log\beta +\beta \right).\hfill \end{array}$$

(8)

Since the MAP estimate lacks a closed-form solution, the Newton–Raphson method is used to obtain a numerical estimate of $\beta$. To evaluate the effectiveness of this numerical solution, we run the program 10 times under each setting. The results are illustrated with error bars in Figure 6. The numerical estimates perform well under Non-IID settings, while noticeable deviations from the true value are observed under IID conditions. This is because the final MAP estimate is influenced by the hyperparameters $a,b$ in the Gamma prior distribution.

In MAP estimation, the bias and variance are influenced by the prior knowledge and the sample size, and are generally smaller than those of the MLE. In this work, since the dataset exhibits the more severe Non-IID characteristics, the parameters of the prior distribution are set to favor these severe Non-IID scenarios. As a result, when the dataset exhibits IID characteristics, the bias and variance of the MAP estimator tend to increase. This issue can be addressed via a $\mu$–$\beta$ mapping. Specifically, when the MAP estimate of $\beta$ exceeds a certain threshold, the dataset is considered IID, and an IID-specific approach is employed for training.

3.3. Relationship Between $\mu$ and $\beta$

The degree of Non-IID in a dataset significantly influences both the performance and convergence speed of the model. Introducing a proximal term into the objective function can improve performance when training on heterogeneous data. However, this regularization term may also impair the model’s generalization ability. Therefore, datasets with different levels of Non-IID are not suitable for using the same value of $\mu$. For instance, the proximal term $\mu$ should be set to 0 for IID datasets. In contrast, for datasets with severe Non-IID, a larger $\mu$ is desirable to enhance generalization.

In this paper, we adopt the $\beta$ parameter introduced in the previous section, which not only quantifies the degree of Non-IID in the dataset but also serves as a basis for estimating a suitable value of $\mu$. This enables the objective function to adapt automatically to varying levels of data heterogeneity. A simple piecewise nonlinear mapping between $\mu$ and $\beta$ is designed, as shown below.

$$\mu =\left\{\begin{array}{cc}{A}_0\hfill & ,\beta \le {\beta }_0\hfill \\ A_0{\displaystyle \frac{{\beta }_0}{10\beta }}\hfill & ,{\beta }_0<\beta <{\beta }_1\hfill \\ 0\hfill & ,\beta \ge {\beta }_1\hfill \end{array},\right.$$

(9)

where $\beta$ is the variable, and $A_0>0$ is a real-valued constant to control the magnitude of the proximal term. ${\beta }_0$ and ${\beta }_1$ are fixed thresholds. The mappings are illustrated in Figure 7. A smaller $\beta$ indicates a more severe Non-IID distribution, which corresponds to a larger value of $\mu$, and vice versa. When the estimated $\beta$ exceeds ${\beta }_1$, the dataset is considered to be IID, and $\mu$ is set to zero. In this case, the proximal term vanishes, and the objective reduces to that of standard federated learning. When $\beta <{\beta }_0$, $\mu$ is assigned a constant value $A_0$. If $A_0$ is set too large or tends toward infinity, the objective function may degenerate into a mean squared error (MSE), overwhelming the original loss function. Therefore, it is advisable to avoid setting $A_0$ excessively large. Based on experimental comparisons, the setting ${\beta }_0=0.01$, ${\beta }_1=1$ and $1\le A_0\le 4$ yields relatively better performance [4].

The output metrics are largely insensitive to variations in the $\mu -\beta$ mapping and the latent variables, such as the parameters a and b of the Gamma function in the prior distribution. The effect of parameter changes on the outputs is predictable and exhibits a clear pattern. Moreover, the outputs are primarily governed by the characteristics of the dataset distribution, indicating a high degree of robustness.

3.4. Controlling Gradient Coefficient ${\mathbf{\Delta }}_{\mathit{i}}$

Given that regularization-based objective functions offer limited improvement under Non-IID conditions and may, to some extent, hinder convergence. To this end, we introduce a mutually orthogonal coefficient, ${\mathbf{\Delta }}^{\left(j\right)},j=1,2,\dots ,r$, to control the gradient and isolate the interference among the knowledge. The coefficient ${\mathbf{\Delta }}^{\left(j\right)}$ is a tensor with the same shape as the model, and each class is assigned a coefficient. These coefficients are mutually orthogonal under the Frobenius inner product, ensuring minimal interference among classes and satisfying orthogonality under the Frobenius inner product, and as follows:

$$\left\{\begin{array}{cc}Tr({{\mathbf{\Delta }}^{\left(i\right)}}^T,{\mathbf{\Delta }}^{\left(j\right)})=0\hfill & ,i\ne j\hfill \\ Tr({{\mathbf{\Delta }}^{\left(i\right)}}^T,{\mathbf{\Delta }}^{\left(j\right)})=1\hfill & ,i=j\hfill \end{array},\right.$$

(10)

where $Tr(·,·)$ represents the sum of the main diagonal elements after the multiplication of two tensors with the same dimensions.

The intuitive rationale behind introducing inter-class orthogonal coefficients is based on the assumption that feature representation spaces corresponding to different classes are mutually exclusive. This assumption is well-justified in practice, as objects from distinct categories typically exhibit unique, non-overlapping characteristics.

Next, client k generates its coefficients based on the local data distribution using the following formula:

$${\mathbf{\Delta }}_{\mathit{k}}=\sum _{i=1}^r{x}_k^i{\mathbf{\Delta }}^{\mathbf{\left(}\mathit{i}\mathbf{\right)}},$$

(11)

where $x_k^i={\displaystyle \frac{q_k^i}{{\sum }_j{q}_k^j}}$ is a scalar that performs element-wise multiplication with the tensor ${\mathbf{\Delta }}^{\mathbf{\left(}\mathit{i}\mathbf{\right)}}$. For details on the generation of ${\mathbf{\Delta }}^{\mathbf{\left(}\mathit{i}\mathbf{\right)}}$, please see Appendix A.

In the round $(i+1)-th$, the model parameters on client k are updated according to the following formula:

$${\mathit{w}}_{\mathit{i}+\mathbf{1}}^{\mathit{k}}={\mathit{w}}_{\mathit{i}}^{\mathit{k}}-lr\ast ({\mathbf{\Delta }}_{\mathit{k}}\odot {\displaystyle \frac{\partial {\mathcal{L}}^{*}}{\partial \mathit{w}}}),$$

(12)

where $lr$ denotes the learning rate, and ⊙ denotes element-wise multiplication. It should be noted that ${\mathbf{\Delta }}^{\left(i\right)}$ denotes a class-specific constant tensor, whereas ${\mathbf{\Delta }}_k$ represents a coefficient computed by client k according to the characteristics of its local data distribution.

This idea is inspired by the concept of global synaptic plasticity in computational neuroscience [9]. In the brain, neuromodulatory substances such as dopamine and serotonin regulate synaptic plasticity. These neuromodulators play a key role in modulating the plasticity mechanisms themselves—a phenomenon often referred to as the "plasticity of plasticity"—which enables the brain to accelerate learning and enhance memory consolidation, as illustrated in Figure 8.

This modulation process can be mathematically formulated as shown in Equation (12), where ${\mathbf{\Delta }}_{\mathit{k}}$ represents the global neuromodulator, and ${\displaystyle \frac{\partial {\mathcal{L}}^{*}}{\partial w}}$ corresponds to synaptic plasticity governed by Hebbian learning rules [26], as depicted in Figure 9.

To clarify the construction procedure of our approach to improve reproducibility and implementation guidance, the pseudocode for the implementation is provided in Algorithm 1.

Algorithm 1 The pseudocode of the proposal

1: Generate ${\mathbf{\Delta }}^{\left(i\right)}$ For class i according to Appendix A at the server side 2: Broadcast ${\mathbf{\Delta }}^{\left(i\right)}$ to all clients 3: Client k calculates ${\mathbf{\Delta }}_k$ according to ${\mathbf{\Delta }}_{\mathit{k}}={\sum }_{i=1}^r{x}_k^i{\mathbf{\Delta }}^{\mathbf{\left(}\mathit{i}\mathbf{\right)}}$ 4: repeat 5:    Broadcasting global model to all clients 6:    Local update according to ${\mathit{w}}_{\mathit{i}+\mathbf{1}}^{\mathit{k}}={\mathit{w}}_{\mathit{i}}^{\mathit{k}}-lr\ast ({\mathbf{\Delta }}_{\mathit{k}}\odot {\displaystyle \frac{\partial {\mathcal{L}}^{*}}{\partial w}})$ at client side 7:    Aggregate local model 8: until Global_epoch > Predefined_value

3.5. Convergence Analysis

The convergence analysis of the proposed approach in this paper is divided into two parts: one concerns the convergence of the objective function with the regularization term, and the other concerns the convergence of the weight updates with orthogonal coefficients. The convergence of the objective function with the regularization term has been thoroughly examined in work [4]. The MAP-based estimation of the regularization coefficient employed in our approach does not affect the convergence property. Therefore, we focus on analyzing the convergence of the weight updates with orthogonal coefficients.

Lemma 1.

Let ${\parallel {\mathbf{\Delta }}_i\parallel }_F=1$, $i=1,2,\dots ,r$, and $<{\mathbf{\Delta }}_i,{\mathbf{\Delta }}_j{>}_F=0,i\ne j$, where ${\parallel ·\parallel }_F$ is the Frobenius norm, $<·,·{>}_F$ is Frobenius inner product, $\mathbf{\Delta }={\sum }_{i=1}^r{w}_i{\mathbf{\Delta }}_i$, where $0\le w_i\le 1$ is a scalar and ${\sum }_{i=1}^r{w}_i=1$, then ${\parallel \Delta \parallel }_F\le 1$.

Proof of Lemma 1.

 Since $<{\mathbf{\Delta }}_i,{\mathbf{\Delta }}_j{>}_F=0$ for $i\ne j$ and ${\parallel {\mathbf{\Delta }}_i\parallel }_F=1$, we have

$${\parallel \mathbf{\Delta }\parallel }_F^2=<\sum _{i=1}^r{w}_i{\mathbf{\Delta }}_i,\sum _{j=1}^r{w}_j{\mathbf{\Delta }}_j{>}_F=\sum _{i=1}^r\sum _{j=1}^r{w}_i{w}_j<{\mathbf{\Delta }}_i,{\mathbf{\Delta }}_j{>}_F=\sum _{i=1}^r{w}_i^2$$

because all cross terms vanish.

$0\le w_i\le 1$, then

$${\parallel \mathbf{\Delta }\parallel }_F=\sum _{i=1}^r{w}_i^2\le \sum _{i=1}^r{w}_i=1.$$

Taking square roots yields ${\parallel \mathbf{\Delta }\parallel }_F=1$.    □

We provide a convergence analysis under the assumption that the loss function is convex. The extension to the non-convex setting has been extensively studied in the prior works [4,6,39,40], and the corresponding conclusions can be straightforwardly generalized.

Assumption that the objective function ${\mathcal{L}}^{*}\left({\mathit{w}}_i\right)$ satisfied following conditions:

• L-smoothness: $\parallel \nabla {\mathcal{L}}^{*}\left({\mathit{w}}_i\right)-\nabla {\mathcal{L}}^{*}\left({\mathit{w}}_j\right)\parallel \le {\displaystyle \frac{L}{2}}{\parallel {\mathit{w}}_i-{\mathit{w}}_j\parallel }^2$.
• Strong convex: ${\mathcal{L}}^{*}\left({\mathit{w}}_i\right)\ge {\mathcal{L}}^{*}\left({\mathit{w}}_j\right)+{({\mathit{w}}_i-{\mathit{w}}_j)}^T\nabla {\mathcal{L}}^{*}\left({\mathit{w}}_j\right)+{\displaystyle \frac{\mu }{2}}\parallel {\mathit{w}}_i-{\mathit{w}}_j\parallel$.
• Bounded gradient: $\mathbb{E}\parallel \nabla {\mathcal{L}}^{*}\left({\mathit{w}}_i\right)-\nabla {\mathcal{L}}^{*}\left({\mathit{w}}^{*}\right){\parallel }^2 \le {\sigma }^2$

Theorem 1.

The proposed approach exhibits the same convergence properties as the FedProx [4], specifically, $\parallel {\mathit{w}}^{*}-{\mathit{w}}_{i+1}^{ours}\parallel \le \parallel {\mathit{w}}^{*}-{\mathit{w}}_{i+1}^{fedprox}\parallel$.

Proof of Theorem 1.

$$\parallel {\mathit{w}}^{*}-{\mathit{w}}_{i+1}^{ours}\parallel = \parallel {\mathit{w}}^{*}-{\mathit{w}}_i^{ours}-lr*(\Delta \odot \nabla {\mathcal{L}}^{*}\left({\mathit{w}}_i\right))\parallel$$

According to Lemma 1,

$$\begin{array}{cc}\hfill \parallel {\mathit{w}}^{*}-{\mathit{w}}_i^{ours}-lr*(\Delta \odot \nabla {\mathcal{L}}^{*}\left({\mathit{w}}_i\right))\parallel & \le \parallel {\mathit{w}}^{*}-{\mathit{w}}_i^{ours}+lr*\nabla {\mathcal{L}}^{*}\left({\mathit{w}}_i\right)\parallel \hfill \\ \hfill & =\parallel {\mathit{w}}^{*}-{\mathit{w}}_{i+1}^{fedprox}\parallel \hfill \end{array}$$

where $\parallel ·\parallel$ represents the distance norm.

The convergence of $\parallel {\mathit{w}}^{*}-{\mathit{w}}_{i+1}^{fedprox}\parallel$ has been proved in the work [4]. Thus, our approach also achieves convergence.    □

4. Experiments

In our experiments, we compare the proposed method with several classical federated learning baselines, including FedAvg [1], FedProx with $\mu =0.2$ [4], SCAFFOLD [6], MOON [7], and FedNACA [9]. To evaluate and compare the convergence speed of these approaches, we employ a two-layer neural network with 1,000 neurons in the hidden layer.

The experiments are conducted on two widely used benchmark datasets: MNIST [41] and CIFAR10 [42]. To simulate varying degrees of data heterogeneity, we synthesize three different non-IID datasets using the Dirichlet distribution [43] with concentration parameters $\beta =1000, \beta =0.5$, and $\beta =0.01$, which correspond to IID, mild Non-IID, and severe Non-IID scenarios, respectively. The IID dataset contains all classes of samples, and the quantity of the class is equal to each other. Mild and severe Non-IID datasets are generated by Dirichlet distribution with the setting $\beta =0.5$ and $\beta =0.01$, respectively. In Non-IID datasets, both the number of classes and the number of samples are imbalanced across clients. During the experiments, the client can only access the number of samples for each class in its local dataset, and the value of the parameter $\beta$ is unknown, which is estimated by MAP. The dataset is partitioned into 10 non-overlapping subsets, and each subset is assigned to a distinct client.

4.1. Metrics

To evaluate the convergence speed of different methods, we employ two metrics: (1) the number of communication rounds required to reach a predefined target accuracy and (2) the top-1 accuracy. In each communication round, a fraction of clients, referred to as the sampling ratio $0<C\le 1$, is randomly selected to participate in model training; while both metrics capture aspects of convergence, the number of rounds to reach the target accuracy provides a more intuitive and direct measure of convergence speed across different algorithms compared to tracking accuracy curves. All experiments are conducted over 100 communication rounds.

4.2. Results on Non-IID Degree and Client Sampling Ratio C

We conduct experiments with varying degrees of Non-IID and client sampling ratios on the MNIST and CIFAR-10 datasets. The results are presented in Figure 10 and Figure 11. In both figures, the x-axis denotes the target accuracy, while the y-axis indicates the number of communication rounds required to achieve target accuracy. A steeper slope corresponds to a greater number of rounds needed, reflecting slower convergence, and vice versa.

Figure 10 and Figure 11 illustrate the effects of two parameters: the degree of Non-IID and the client sampling ratio. The degree of Non-IID is categorized into three levels: IID, mild Non-IID, and severe Non-IID. The client sampling ratio C takes values in $0.2,0.5,1.0$.

From the figures, several observations can be made. First, under a fixed Non-IID setting (i.e., subplots within the same row, from left to right), all algorithms exhibit faster convergence as C increases. This implies that increasing the number of participating clients leads to improved model generalization. Second, as the data distribution shifts from IID to severe Non-IID (i.e., subplots within the same column, from bottom to top), the performance of all algorithms degrades significantly, indicating that data heterogeneity substantially impacts convergence behavior. Third, regarding convergence speed, our proposed method consistently outperforms other baselines in both mild and severe Non-IID scenarios, while maintaining comparable performance under IID settings.

These improvements primarily arise from two key aspects. First, unlike the fixed proximal term in FedProx, we estimate the hyperparameter $\mu$ via MAP estimation, enabling it to adapt dynamically to different degrees of Non-IID. Consequently, our method surpasses both FedProx and FedAvg in performance. Second, during the parameter update phase, we introduce an orthogonalized coefficient tensor that is element-wise multiplied with the gradient. This operation effectively isolates the representations of different classes or knowledge components, thereby mitigating mutual interference—a feature absent in FedNACA.

Some approaches, such as FedAvg, exhibit poor performance under conditions of low client sampling ratios and severe Non-IID data distributions, as illustrated in the first subplot of the first row in Figure 10. After 100 rounds of training, FedAvg achieves an accuracy of only 20%. This suboptimal performance can be partly attributed to the relatively small model size. Nevertheless, increasing the model capacity leads to some improvement in accuracy. Additionally, the limited number of training rounds suggests that, under the current experimental settings, FedAvg’s convergence speed is comparatively slow.

4.3. Results on Top-1 Accuracy

The number of communication rounds required to reach the target accuracy serves as a measure of the convergence speed of each approach, whereas the top-1 accuracy reflects the model’s reliability. The top-1 accuracies of all baseline methods after 100 rounds are summarized in

Table 1. Top-1 accuracy. Bold values indicate the highest performance among all methods under the same experimental setting.

		MNIST			CIFAR10
		C = 0.2	C = 0.5	C = 1	C = 0.2	C = 0.5	C = 1
Severe Non-IID	FedNACA	84.09 ± 0.68	84.69 ± 0.41	84.48 ± 0.22	42.88 ± 0.83	57.76 ± 0.63	58.16 ± 0.45
	Proposal	84.65 ± 0.78	91.99 ± 0.38	91.54 ± 0.2	57.42 ± 0.68	58.22 ± 0.61	59.21 ± 0.32
	FedAvg	71.02 ± 0.76	89.96 ± 0.69	90.86 ± 0.38	22.38 ± 1.21	25.33 ± 0.98	26.37 ± 0.55
	FedProx	78.17 ± 1.02	64.53 ± 0.88	62.88 ± 0.65	25.96 ± 1.12	30.35 ± 0.94	30.61  ± 0.89
	MOON	45.59 ± 1.41	53.17 ± 0.9	53.69 ± 0.89	23.89 ± 1.66	27.54 ± 1.32	29.44 ± 1.2
	SCAFFOLD	88.84 ± 0.84	88.69 ± 0.7	89.62 ± 0.66	31.35 ± 1.31	33.14 ± 0.97	36.39 ± 0.71
Mild Non-IID	FedNACA	82.35 ± 0.51	84.65 ± 0.31	85.66 ± 0.15	32.72 ± 0.68	33.85 ± 0.45	34.48 ± 0.33
	Proposal	96.86 ± 0.44	96.46 ± 0.3	96.34 ± 0.16	44.23 ± 0.53	47.28 ± 0.4	48.2 ± 0.28
	FedAvg	96.89 ± 0.52	97.21 ± 0.29	96.28 ± 0.15	23.21 ± 0.66	24.01 ± 0.48	26.3 ± 0.25
	FedProx	93.86 ± 0.52	93.96 ± 0.3	92.61 ± 0.16	43.54 ± 0.7	44.73 ± 0.42	43.9 ± 0.29
	MOON	94.03 ± 0.61	94.58 ± 0.34	92.44 ± 0.19	41.27 ± 0.73	43.16 ± 0.51	43.16 ±  0.36
	SCAFFOLD	88.46 ± 0.48	90.52 ± 0.31	91.52 ± 0.16	36.3 ± 0.68	40.11 ± 0.46	41.53 ±  0.31
IID	FedNACA	92.72 ± 0.22	92.82 ± 0.12	92.75 ± 0.09	28.47 ± 0.56	28.61 ± 0.29	34.48 ± 0.19
	Proposal	97.51 ± 0.2	97.14 ± 0.12	96.97 ± 0.11	41.37 ± 0.57	42.88 ± 0.28	43.57 ± 0.21
	FedAvg	97.57 ± 0.2	97.76 ± 0.11	97.7 ± 0.1	33.01 ± 0.58	34.76 ± 0.3	35.03 ± 0.22
	FedProx	96.34 ± 0.19	96.39 ± 0.11	96.42 ± 0.11	48.77 ± 0.52	49.43 ± 0.26	49.59 ± 0.22
	MOON	96.46 ± 0.25	95.4 ± 0.14	94.28 ± 0.13	46.6 ± 0.61	47.52 ± 0.33	46.81 ± 0.24
	SCAFFOLD	89.12 ± 0.21	90.73 ± 0.12	92.01 ± 0.11	38.87 ± 0.57	40.14 ± 0.31	41.77 ± 0.22

. As shown in Table 1, our proposed method achieves higher top-1 accuracy than most other approaches under both mild and severe Non-IID settings. Under the IID setting, although the accuracy difference between our method and the highest top-1 accuracy reaches up to 6.55%, our approach is still not the worst-performing approaches.

Two main factors contribute to these observations. First, during the estimation of $\beta$ via MAP, the parameters of the prior Gamma distribution $(a={\displaystyle \frac{1}{2}}, b=1)$ are better suited for Non-IID scenarios. When the data distribution is IID, this prior introduces some bias in the estimation of $\beta$. Second, the mapping function between $\beta$ and $\mu$ does not fully capture the true variation, leading to certain estimation errors, as illustrated in Figure 7 and Equation (9).

5. Conclusions

In this paper, we propose a novel approach to address the challenges posed by Non-IID data distributions in federated learning. The proposed approach comprises two key components. First, a new proximal term coefficient is incorporated into the objective function, where the coefficient is derived via Maximum A Posteriori (MAP) estimation based on the local sample statistics to quantitatively reflect the degree of data heterogeneity. Second, to mitigate model drift and reduce inter-class interference during updates, an orthogonal tensor coefficient is assigned to each class. This coefficient modulates the gradient through element-wise multiplication, thereby preserving class-specific knowledge. Extensive experiments conducted on benchmark datasets demonstrate the effectiveness and robustness of the proposed method in comparison with classical baselines under varying degrees of data heterogeneity.