Source-linked AI summary
Flexible Clustered Federated Learning for Client-Level Data Distribution Shift
Moming Duan, Duo Liu, Xinyuan Ji, Yu Wu, Liang Liang, Xianzhang Chen, Yujuan Tan
TL;DR
Federated learning must train across non-IID, imbalanced, and shifting client data without centralizing private data, which can increase divergence between local and global models. FlexCFL clusters clients by optimization-direction similarity, supports newcomer cold starts, and migrates clients under distribution shift; evaluations report higher accuracy and a balance between accuracy and communication efficiency. Its main added cost is pre-training computation and one-time gradient communication.
Problem
Federated learning faces non-IID, imbalanced, and distribution-shifted client data, increasing divergence between local and global models and degrading performance.
Method
FlexCFL clusters client optimization directions with EDC, uses cold-start mechanisms for groups and newcomers, and migrates clients when Wasserstein distance detects distribution shift.
Results
FlexCFL achieved higher classification accuracy than FedAvg, FedProx, FeSEM, and RAC across evaluations, including +40.9% top-1 test accuracy on FEMNIST versus FedSEM.
Takeaways & Limitations
FlexCFL can balance accuracy and communication efficiency in client-level distribution-shift environments.
Takeaways & Limitations
FlexCFL requires additional pre-training computation and communication to transmit group optimization gradients once per client.
Abstract
from arXiv · showhide
Federated Learning (FL) enables the multiple participating devices to collaboratively contribute to a global neural network model while keeping the training data locally. Unlike the centralized training setting, the non-IID, imbalanced (statistical heterogeneity) and distribution shifted training data of FL is distributed in the federated network, which will increase the divergences between the local models and the global model, further degrading performance. In this paper, we propose a flexible clustered federated learning (CFL) framework named FlexCFL, in which we 1) group the training of clients based on the similarities between the clients' optimization directions for lower training divergence; 2) implement an efficient newcomer device cold start mechanism for framework scalability and practicality; 3) flexibly migrate clients to meet the challenge of client-level data distribution shift. FlexCFL can achieve improvements by dividing joint optimization into groups of sub-optimization and can strike a balance between accuracy and communication efficiency in the distribution shift environment. The convergence and complexity are analyzed to demonstrate the efficiency of FlexCFL. We also evaluate FlexCFL on several open datasets and made comparisons with related CFL frameworks. The results show that FlexCFL can significantly improve absolute test accuracy by +10.6% on FEMNIST compared to FedAvg, +3.5% on FashionMNIST compared to FedProx, +8.4% on MNIST compared to FeSEM. The experiment results show that FlexCFL is also communication efficient in the distribution shift environment.
1 INTRODUCTION
FlexCFL addresses heterogeneous and shifting client data by clustering clients, supporting newcomers, and migrating clients when distributions change. It reports improved accuracy and communication efficiency against established FL and CFL methods.
- FlexCFL is an efficient, flexible clustered federated learning framework designed for heterogeneous federated training.
- The framework clusters clients using EDC, maintains an auxiliary server for newcomer cold starts, and uses Wasserstein distance to trigger client migration.
- FlexCFL was evaluated on four open datasets and compared with several FL and CFL frameworks.
- The framework includes an efficient cold start strategy for groups and newcomers to improve scalability and practicality.
- The paper investigates client-level distribution shift and proposes a communication-efficient client migration algorithm.
2 BACKGROUND AND RELATED WORK
Federated learning trains a global model by aggregating local client updates, but statistical heterogeneity makes consensus optimization difficult. Clustered federated learning addresses this challenge by dividing the optimization into client-specific subobjectives.
- FedAvg: FedAvg weights client contributions by local training data and updates each selected client through local mini-batch gradient steps.The algorithm initializes a global model, selects clients, broadcasts the current model, and aggregates their locally computed updates.
- Federated learning: Federated learning coordinates a server and participating devices through repeated communication rounds that broadcast, locally optimize, aggregate, and update the global model.Training commonly requires hundreds of rounds to reach the target accuracy.
- FedAvg: The participation rate and local epoch count are key FedAvg hyperparameters affecting convergence, with smaller participation rates recommended for non-IID data to alleviate straggler effects.A higher participation rate can improve convergence in IID settings, while an appropriately large local epoch count can increase convergence speed.
- Statistical heterogeneity: Statistical heterogeneity—including non-IID, size-imbalanced, and class-imbalanced data—can make consensus global models perform unsatisfactorily.The global objective jointly represents the local client objectives, which may measure risk over differing data distributions.
- Clustered federated learning: Clustered federated learning divides the optimization goal into several subobjectives instead of optimizing one consensus model.The first CFL framework recursively separates client groups with incongruent descent directions, but its bi-partitioning process is computationally inefficient and may require multiple communication rounds.
3 FLEXCFL
FlexCFL clusters clients by optimization-direction similarity to reduce divergence, supports newcomers through cold starts, and migrates clients when local data distributions shift. Its design combines grouped optimization with communication-conscious aggregation and flexible client reassignment.
- Motivation: High statistical heterogeneity makes local models diverge more, slows convergence, and reduces model accuracy under FedAvg.The motivating MNIST experiment varies the number of classes available per client and measures accuracy and model-weight discrepancy.
- Framework Overview: FlexCFL groups client training by local-optimization proximity, using static grouping to avoid rescheduling clients every round.The framework builds on clustered federated learning while changing the grouping strategy and its similarity measure.
- Framework Overview: The framework performs intra-group FedAvg, weighted inter-group aggregation, and gradient maintenance for newcomer cold starts rather than broadcasting server optimization gradients to every client each round.Clients and groups maintain current models and updates, while the auxiliary server supports initial group assignment for newcomers.
- Group Cold Start: FlexCFL uses EDC to cluster client update directions into sub-optimization groups through an unsupervised procedure based on decomposed cosine similarities.EDC addresses limitations of pairwise proximity measures and uses truncated SVD to reduce update directions before clustering.
- Client Cold Start and Migration: FlexCFL detects client-level distribution shifts with Wasserstein distance and schedules cold starts to migrate clients whose shift exceeds the threshold.This flexible migration strategy preserves static clustering while addressing the distribution-shift limitation of static scheduling.
- Key Features: FlexCFL reduces the discrepancy between joint and sub-optimization objectives while supporting unsupervised grouping without assuming incongruent risk functions.The framework’s stated features also include communication reduction through its client scheduling mechanism.
4 CONVERGENCE ANALYSIS
FlexCFL’s convergence analysis bounds divergence between federated group models and virtual centralized-training models under convexity, Lipschitz continuity, and smoothness assumptions. It treats both settings without and with inter-group aggregation, recovering FedAvg as a special case.
- Assumptions and group loss: The analysis assumes each client loss is convex, M-Lipschitz continuous, and L-Lipschitz smooth.These assumptions are stated for client losses and are inherited by the group loss function.
- Assumptions and group loss: The group loss function is convex, M-Lipschitz continuous, and L-Lipschitz smooth because it is a linear combination of local loss functions.
- Divergence analysis: FlexCFL introduces a virtual group model trained centrally on members’ combined data and synchronized with the federated model each communication round.This model measures divergence between federated and centralized-training behavior.
- Divergence analysis: Intra-group gradient divergence measures the gradient difference between a client’s loss and the group loss under a fixed group membership.
- Without inter-group aggregation: Lemma 2 gives an upper bound on divergence between the FlexCFL model and virtual group model under Assumptions 1–3.The proof uses smoothness and the intra-group gradient-divergence definition, then combines the result with the local update and Jensen’s inequality.
- Without inter-group aggregation: Theorem 1 establishes a convergence bound between the federated group model and virtual group model without inter-group aggregation.The analysis first considers ηg = 0 before extending the bound to ηg > 0.
- With inter-group aggregation: With inter-group aggregation, FlexCFL obtains a corresponding convergence bound that reduces to the no-aggregation bound when ηg = 0.When |G| = 1, the framework degrades to FedAvg.
5 EVALUATION
FlexCFL is evaluated on four federated datasets against established FL and clustered FL baselines, including settings with and without client-level distribution shift. Across these experiments, its clustering, cold-start, migration, and inter-group aggregation strategies improve accuracy, convergence, or communication efficiency under the reported conditions.
- Experimental setup: FlexCFL is evaluated on four federated datasets using image-classification and synthetic tasks, with FedAvg, FedProx, IFCA, and FeSEM as baselines.The evaluation uses MNIST, FEMNIST, FashionMNIST, and Synthetic datasets; experiments use K = 20 selected clients per round and E = 10 local epochs.
- Without distribution shift: +8.4% on MNIST, +40.9% on FEMNIST, and +11.3% on FashionMNIST are FlexCFL’s reported absolute test-accuracy improvements over FeSEM without distribution shift.The comparison reports FlexCFL and IFCA as significantly superior to the other evaluated frameworks.
- Without distribution shift: FedProx reduces training divergence and stabilizes accuracy, while FlexCFL converges faster than IFCA and FeSEM in the reported MNIST experiments.The paper links faster convergence to reduced communication consumption in federated systems.
- Ablation studies: RAC decreases accuracy by an average of 5.5% on MNIST and 19.7% on FEMNIST, with final scores worse than FedAvg.RAC randomly assigns newcomer clients while retaining clustered clients in their groups, whereas RCC causes only moderate degradation in most settings.
- Distribution shift: Client-level distribution shifts degrade clustered methods more than FedAvg because changing client data distributions mismatches fixed clustered training assignments.In FEMNIST swap-all, accuracy decreases by 10.2% for IFCA, 11.8% for FeSEM, and 23.3% for FedGroup, while FedAvg decreases by 1.1%.
- Distribution shift: FlexCFL’s migration strategy improves accuracy under distribution shift, including +2.2% on MNIST, +3.9% on FEMNIST, and +1.4% on FashionMNIST over FedAvg in swap-all.In swap-part, FlexCFL improves over IFCA by +8.9% on MNIST, +10.8% on FEMNIST, and +5.3% on FashionMNIST.
- Communication efficiency: FlexCFL reaches a 5.2% accuracy improvement over FedAvg with 1.26× communication consumption, whereas IFCA-MIG needs 5.48× communication for a 1.9% improvement.The comparison concerns client migration under distribution shift.
6 CONCLUSION
FlexCFL improves federated training through efficient clustering and cold-start strategies, while balancing classification accuracy and communication in client-level distribution shifts.
- FlexCFL improves federated model performance through efficient clustering and cold-start strategies.
- +40.9% top-1 test accuracy is achieved on FEMNIST compared to FedSEM.
- Higher classification accuracy is achieved than with FedAvg, FedProx, FeSEM, and the random strategy RAC across eight models.
- FlexCFL balances communication consumption and accuracy in the client-level distribution shift environment.