Federated learning has emerged as a pivotal approach to safeguard data privacy while effectively integrating disparate data silos. FedAvg, as a classic federated learning framework, can enhance the performance of edge device models through an algorithm that ensures user privacy [4]. In this algorithm, the central server first selects a subset of clients to participate in the training and distributes the global model to these selected clients. Each client independently trains the model using their local data and the central server aggregates the model updates from the participating clients to form a new global model. However, due to differences in user demographics or usage habits, the local data on clients may exhibit non-IID characteristics. Li et al. further confirmed that traditional FedAvg faces challenges such as slow convergence of the global model and deviation from optimal solutions under non-IID conditions [6].

Many researchers have proposed to mitigate the biases that arise during local training on non-IID data, aiming to alleviate the adverse effects in the federated averaging process and enhance the performance of the global model. For instance, FedProx introduces a regularization term during local training that constrains updates based on the distance between the local model and the global model, thereby reducing overfitting of the local models [20]. MOON normalizes local training by leveraging the similarity between the representations of local and global models, incorporating a contrastive learning approach that maximizes the similarity to improve model generalization [21]. FedLC further calibrates at the logit level to reduce updates for minority classes, thereby enhancing model performance when dealing with imbalanced data [22]. FedNova refines the global aggregation phase by adjusting the contribution weights of each client, making the global model more aligned with the global optimum [23]. These existing methods have addressed label shift issues to some extent, but further research and improvements are necessary to enhance their effectiveness in real-world scenarios.

A comprehensive survey categorizes heterogeneous federated learning into data space, statistical, system, and model heterogeneity, suggesting further research to improve model generalization and performance across diverse clients [24]. Researchers have proposed several strategies to mitigate these issues. For example, HeteroFL addresses computational heterogeneity by enabling the training of heterogeneous local models with varying complexities [25]. Exploiting Model and Data Heterogeneity in FL (MDH-FL) employs knowledge distillation and symmetric loss to tackle both data and model heterogeneity [26]. Recent approaches also explore adaptive data distribution, regularization terms, contrastive learning, and multi-task learning to address heterogeneity issues [27]. These methods aim to optimize algorithms and model structures to cope with heterogeneity, but they do not address the diversity of client data distributions. Consequently, some researchers propose clustered federated learning (CFL) to group clients with similar data distributions, thereby improving model performance and enhancing privacy protection level.

CFL is a model-agnostic distributed multi-task optimization framework that enhances both model performance and privacy protection. In recent years, various CFL frameworks have been proposed, including confidential aggregation techniques for securing individual updates and customized systems tailored for human activity recognition applications [28]. The latest advancements focus on efficiently identifying the distributional similarity between client data subspaces using principal angles, which accelerates cluster formation and provides convergence guarantees for non-convex objectives [29]. Furthermore, some researchers have adopted non-convex pairwise fusion, enabling autonomous estimation of cluster structures without prior knowledge [30]. As specific examples in the CFL domain, ClusterFL presents a multi-task federated learning framework that automatically captures inherent clustering relationships among nodes, thereby improving accuracy and reducing communication overhead, particularly in human activity recognition applications [28]. ACFL introduces a mean-shift clustering algorithm and an auction-based client selection strategy aimed at mitigating data heterogeneity and balancing energy consumption in mobile edge computing systems [31]. By leveraging the geometric properties of the federated learning loss surface, clients are grouped into clusters with jointly trainable data distributions, suitable for general non-convex objectives, while performing multi-task optimization under privacy preservation [32]. IFCA addresses non-IID data through iterative clustering and model updates, partitioning clients with similar data distributions into several clusters where each cluster independently conducts model training and updates followed by global model aggregation [33]. FedAC effectively integrates global knowledge into cluster learning by decoupling neural networks and employing different aggregation methods for each submodule, thus significantly enhancing performance [22]. FeSEM introduces an expectation-maximization algorithm for client clustering, ensuring that clients within each cluster share similar data distributions [34].

The CFL methods described above address the issue of model parameter inconsistency by incorporating clustering steps during local training, reducing the global model performance decline caused by label shift. They excel at increasing model accuracy, convergence rate, and efficiency. However, there are two major issues that most studies have not addressed. The first issue is model consistency; model consistency across different clusters can have an impact on overall performance, especially when data distribution is complex. The second issue is the ability to adapt dynamically. Most existing methods’ adaptive adjustment capacity may be restricted for extremely uneven data distributions. To address data heterogeneity in federated learning distributed training, this study introduces an intra-cluster correction method based on the CFL framework. We use a specific sampling strategy to ensure that labels are dynamically balanced with each round of sampling, allowing the cluster model to update towards the optimal global solution.

In traditional federated learning algorithms, each iteration typically involves either selecting all clients to participate in the aggregation of the global model or choosing a subset of clients based on a specific sampling strategy. However, when the private datasets on client devices exhibit non-IID characteristics, the performance of the aggregated global model can significantly deteriorate. The purpose of CFIC is to mitigate the issue of client model drift caused by label distribution shifts, thereby obtaining a globally aggregated model with guaranteed performance. The symbols used in this study are shown in Table 1.

In the context of federated learning across devices, consider a system composed of K clients, denoted as {C1,C2,⋯,CK}. All clients are dedicated to integrating their respective data {D1,D2,⋯,DK} in order to obtain a better-performing machine learning model. The central server interacts with these clients to collaboratively find model parameters w that minimize loss through Eq. (1). (1)minwf(w)=∑i=1KpiFi(w)=Ei[Fi(w)]

In this setup, pi represents the probability that participant user i is chosen to contribute to the global model training process, and this probability pi satisfies pi≥0 and ∑i=1Kpi=1. Fi(w) indicates the loss error prediction for model parameters w by the client j on its local dataset (xi,yi), typically calculated using function Fi(w)=ℓ(xi,yi,w), where ℓ(⋅) is a predefined loss function relevant to the specific task. Typically, the FedAvg method is used to minimize model parameters w through unbiased sampling. The probability pi follows a uniform distribution, implying that all users have an equal random chance of being selected to participate in the training.

Existing research primarily focuses on clustering client-trained local models based on their similarities, which poses a significant challenge for large-scale distributed applications in federated learning. In iterative clustering methods, calculating model similarity is a resource-intensive process that requires substantial computational resources. This is especially true in highly heterogeneous environments where the dataset may contain only a few labels. To address this issue, we can introduce a simulated IID data distribution as a reference. By comparing the local client’s data distribution with this reference distribution, we can obtain the differences between them. We map these distributional differences to a value that represents the deviation of different clients’ data from the ideal distribution. By performing calculations based on these values, we can significantly reduce the required computational scale. Following the work [35], we assume a uniform global class distribution when computing local discrepancies to prevent global label information leakage. Simultaneously, only the maximum categorical label is retained during constructed a feature extraction function Ψ(⋅), thereby mitigating unnecessary local label information exposure, as shown in Eq. (2). (2)Ψ(D(k))=arg⁡mini(log⁡|Di(k)−Qi(k)|∑jlog⁡|Dj(k)−Qj(k)|)

In this equation, Di(k) represents the actual proportion of samples with label i in the data set ∑i=1CDi(k)=1 of client k, C denotes the total number of labels, and |Di(k)−Qi(k)| indicates the distributional variance of label i within client k. Clients then upload their extracted label distribution features to the server. The server uses these uploaded features to perform client clustering without needing to predefine the number of clusters. The number of clusters is determined by the distribution characteristics of the clients. Client clustering is initiated only during the first round of communication and when new clients join, thereby conserving significant computational resources. This approach involves clients locally mapping feature extraction functions to numerical values. Clients then only upload these mapped values, ensuring that their own label distributions are not exposed. This method provides a level of privacy protection, as the raw data remains on the clients’ local machines and only aggregated, mapped values are shared.

In this step, we mitigate the issue of client model drift caused by extreme label distribution imbalance. In such scenarios, most clients contain only a few or even just one type of label. For clustering groups, the data within each group is more likely to follow an independent and identically distributed distribution. Traditionally, scholars often adopt aggregation functions such as FedAvg, Krum, and Trimmed Mean, which can effectively protect user data privacy [36]. These aggregation functions are used to summarize local model updates from clients into a global model. The choice of aggregation function directly affects the robustness of the model. However, the results of ablation experiments show that using the FedAvg aggregation function decreases model performance. While the Krum and Trimmed Mean aggregation functions have advantages in terms of robustness and simplicity, they require the use of local model gradients from the previous round, making it vulnerable to adversarial samples and limiting their applications in certain dynamically changing scenarios [37]. In this paper, instead of using traditional aggregation functions, we aggregate each cluster into a cluster model, which exhibits higher accuracy for the labels within its respective cluster. However, these cluster models might deviate from the direction of model aggregation in terms of loss-minimizing model parameters w. Therefore, this paper introduces an intra-cluster model correction strategy to update the cluster models towards the direction of the global model’s optimal solution.

To ensure model stability, we have improved the sampling strategy to make the number of samples from each cluster relatively uniform across iterations. Specifically, K clients are clustered into m clusters represented as {𝒢1,𝒢2,𝒢3,⋯,𝒢m}, with a sampling ratio of γ. Within each cluster, each client’s opportunity to participate in training is determined by uniform random sampling, denoted as ∑i∈𝒢pi=1. The sampling number for each cluster is min(⌊Kγm⌋,1), and then additional Kγ−∑i∈𝒢mini(⌊Kγm⌋,1) clients are sampled from the remaining clients to form a final set of clients S(t) as the sampling result for the t-th communication round. If it is the first round of communication, Kγ clients are randomly and uniformly sampled to form client set S(t).

At the beginning of the t-th communication, the server obtains the current global model parameters wt and distributes them to each local client to get local model parameters wkt+1. This paper aggregates m cluster models {g1,g2,g3,⋯,gm} from m clusters. The cluster model update formula is given in Eq. (3). (3)git+1=∑k∈𝒢inknwkt+1

For the standard gradient descent formula w=w−ηΔw, where η represents the learning rate and Δw is the step size for gradient adjustment at each time step. As it approaches the optimal value, the gradient becomes smaller. Since the learning rate is fixed, the standard gradient descent method converges slowly and may even fall into local optima. This paper introduces momentum h to correct the direction of the global model. Specifically, this is done by comparing the deviation between the cluster model and the direction of the previous round’s global model to correct the direction of the global model aggregation in the current round. We also consider that the cluster model might deviate from the direction of the optimal solution of the global model in some rounds. Therefore, we add the historical correction gradient direction to improve the stability of the model, shown in Eq. (4). (4)ht+1=(α⋅ht−β⋅∑i=1mningit+1−wt‖git+1−wt‖2)

In this formula, α is the momentum term, representing the influence of historical gradients. The larger α is, the greater the impact of the historical correction gradient direction on the current round. The incorporation of the momentum term offers significant advantages in the high non-IID scenario, which is the focus of this study. On one hand, it accelerates the escape from local minima by leveraging historical accumulation. On the other hand, it mitigates the update oscillations caused by client sampling fluctuations. This design is inspired by the classical convergence theory of distributed optimization, where the core idea is to reduce client gradient variance through exponential smoothing, thereby enhancing convergence efficiency. Therefore, the server can smooth the direction of historical updates to alleviate the impact of heterogeneity among client updates. Finally, we update the global model w by averaging the local models wkt+1 trained by clients in this round and adding the correction from the cluster model, shown in Eq. (5). (5)wt+1=∑k∈S(t)nknwkt+1−ht+1

Based on the above theoretical analysis, we provide the following pseudocode for the CFIC algorithm as follows (Algorithm 1):

In the provided pseudocode, the algorithm takes the number of clients k, the number of communication rounds t, the sampling ratio γ, the server-side client list 𝒜, gradient information h0, initial model parameters w0, and the learning rate η as input. On the server side, the algorithm initializes w0 with all client values Lk set to empty (Line 1). For the t-th communication round, the algorithm randomly and uniformly samples clients within clusters based on the sampling ratio γ, forming the client set for this round (Line 3). The algorithm then iterates over the client set for this round and updates the client model parameters (Lines 4–5). If a client is not in the central server’s-maintained client table, the algorithm updates the client in the central server’s client table and re-clusters the client to form a new cluster (Lines 6–8). After completing these steps, the algorithm updates the cluster model at Line 9, corrects the gradient at Line 10, and updates the global model parameters at Line 11.

The client-side update process is as follows: The client first obtains the model parameters at Line 13 and then processes the client set in parallel, updating the model parameters (Lines 14–15). If the client is not within the label distribution characteristics, the algorithm uses a label distribution feature extraction function for client clustering (Lines 16–17). Finally, the client returns the model parameters and label distribution characteristics to the server side (Line 18).

The cold start problem typically arises when new clients join a federated learning system for the first time. Due to the lack of sufficient historical data to train their local models, these new clients are unable to make meaningful contributions to the global model during the initial phase. This issue not only affects the performance of the new clients’ own models but also potentially slows down the overall performance and stability of the entire system. To address this challenge, we propose the CFIC algorithm, which can meet the needs of new clients joining the federated learning system by mitigating the cold start issue. Specifically, we maintain a client table on the server-side that records all participating clients in the federated learning system. When a new client attempts to join the system, the server detects its absence from the maintained client table 𝒜 and subsequently triggers a re-clustering process. During this re-clustering phase, the new client is incorporated into an appropriate cluster based on certain criteria, thereby facilitating its integration and enabling it to contribute more effectively to the global model.

We conducted experiments on seven real-world datasets to validate the accuracy and fairness of our algorithm. The purpose of using these real-world datasets is to test the performance of the proposed algorithm in a realistic setting when compared with other algorithms, as well as to evaluate the accuracy of various algorithms on these datasets under heterogenous conditions.

MNIST dataset [38]. This dataset consists of images of handwritten digits from 0 to 9. The input for the dataset is 784-dimensional (28 × 28) flattened images, and the output is class labels ranging from 0 to 9.

We compared our proposed algorithm with the following six federated learning algorithms to test its performance. 1.

FedAvg algorithm. This algorithm is a widely used aggregation method in federated learning. In each round, the server sends the global model parameters to a randomly selected group of clients. Each client trains the model on its local dataset and then sends the local updates back to the server for aggregation. The updated global model is subsequently disseminated to the clients for the next round of training.

The model structure used in this study is as follows: MNIST adopts dual 5 × 5 convolutional layers (32 → 64 channels) and dual fully connected layers (3136 → 512 → 10); CIFAR-10/100 and SVHN are both dual 5 × 5 convolutions (64 channels) combined with three fully connected layers (1600 → 384 → 192 → output), with CIFAR-100 adjusted to 192 → 100 through the final linear layer; EMNIST uses a three convolutional layer (32 → 32 → 64 channels) and a double fully connected layer (576 → 256 → 62), with a unique pooling connection sequence and custom dimension transformation layers. All models use ReLU activation and 2 × 2 max pooling, and non-linear mapping between fully connected layers is achieved through ReLU. The training and testing ratio is 0.9: 0.1.

The experimental setup utilized an Intel(R) Xeon(R) Gold 5218 CPU operating at 2.30 GHz and CentOS Linux release 7.9.2009 (Core) as the operating system. The platform was built using the Python library PyTorch and equipped with three NVIDIA Tesla V100S PCIe 32 GB graphics cards. We assumed a scenario involving 100 devices participating in federated learning with a sampling rate of 0.3. The batch size was set to 64, and the learning rate was uniformly set to 0.01 across all experiments. Additionally, the number of local iterations was fixed at five for each round of communication.

To simulate the class distribution of client data, we employed the Dirichlet distribution to partition the datasets. Sampling was conducted based on corresponding probability values. For the five real-world datasets, different alpha values (0.05, 0.1, 0.3, 0.5) were selected to conduct experiments aimed at evaluating the testing accuracy under various Dir partitions. Due to variations in dataset sizes and partitioning methods, different algorithms required varying numbers of communications to achieve convergence in terms of testing accuracy and training loss across the datasets. Consequently, the number of communications used in the experiments varied accordingly among the five real-world datasets. Fig. 2 illustrates the distribution of classes owned by clients for the CIFAR10 dataset under different alpha values. The distributions of the other datasets are omitted due to space constraints.

(1) Testing accuracy

Table 2 illustrates the testing accuracy of various algorithms under different data distribution scenarios. FedAvg exhibits a significant disadvantage in handling highly heterogeneous data, particularly on the CIFAR-100 dataset where its accuracy is only 13.50%, compared to 44.94% on the CIFAR-10 dataset. This demonstrates that FedAvg’s global model generalization ability is insufficient when addressing label imbalance across clients. Conversely, algorithms that incorporate regularization and optimization strategies, such as FedProx and FedDyn, achieve performance improvements. FedProx reduces inter-client model update discrepancies through regularization, while FedDyn mitigates the impact of label shift with dynamic regularization, performing particularly well on large-scale datasets like CIFAR-10 and CIFAR-100.

Comparative experiments demonstrate that the proposed CFIC algorithm significantly improves testing accuracy under non-IID conditions through its intra-cluster calibration strategy. While demonstrating suboptimal performance on FMNIST, the proposed algorithm consistently outperformed state-of-the-art methods across all six remaining benchmark datasets. Specifically, CFIC achieved accuracy rates of 73.65%, 33.05%, and 83.45% on CIFAR-10, CIFAR-100, and FEMNIST benchmarks, surpassing all baseline algorithms. Notably, CFIC maintains efficient convergence and superior accuracy even under extreme label imbalance conditions. Although FedFa and FedMGDA+ demonstrate strong generalization capabilities in certain scenarios, their effectiveness diminishes significantly when handling severe data skewness, particularly under small alpha values where substantial performance degradation is observed.

(2) Communication efficiency

To evaluate the communication efficiency of the CFIC algorithm, we measured the number of communication rounds required for each algorithm to reach the target accuracy on various datasets. The accuracy achieved by FedAvg in the last round was used as the benchmark, where “0” indicates that the algorithm failed to reach the target accuracy within a limited number of communication rounds. Fig. 3 illustrates the number of communication rounds needed for different algorithms to achieve this benchmark accuracy across various datasets. By incorporating a unique intra-cluster correction mechanism, the CFIC algorithm optimizes the direction of global model updates, significantly reducing the number of communication rounds. This advantage has been validated across seven different datasets and has significantly lowered the data transmission frequency during federated learning, effectively reducing communication costs and latency. These results demonstrate that CFIC improves both model performance and overall system efficiency in federated learning.

(3) Client experiment

To further validate the effectiveness of the CFIC algorithm in handling data heterogeneity, we conducted a series of experiments on the CIFAR-10 dataset, examining the impact of varying client numbers (20, 500) in a Dir(0.05) heterogeneous environment. The results, shown in Fig. 4, indicate that as the number of clients increases, the CFIC algorithm not only adapts to large-scale federated learning environments but also maintains high accuracy and stability across all scales. Specifically, in an experiment with 20 clients, CFIC achieved an accuracy of approximately 37.19%; with 500 clients, the accuracy reached 52.85%, surpassing other algorithms. These outcomes robustly demonstrate CFIC’s superior performance in non-IID data settings, highlighting its reliability and scalability under different scales and extreme data heterogeneity conditions.

(4) Ablation experiment

The CFIC algorithm is divided into two main parts, i.e., inter-cluster sampling and intra-cluster correction. During the inter-cluster sampling phase, it involves sampling members from each clustered group. This effectively reduces the impact caused by non-IID data among clients, thereby optimizing communication efficiency and the consistency of model training. In the model aggregation phase, the global model is adjusted through an intra-cluster correction strategy, enhancing the generalization performance of the global model. This phase critically affects the model’s ability to handle heterogeneous data. To validate the effectiveness of the CFIC, we remove either of the two core components to evaluate their specific impact on model performance. The experimental results, as shown in Fig. 5, indicate that the removal of either component leads to a significant performance drop on both the CIFAR-100 and SVHN datasets.