FedTSA: A Cluster-based Two-Stage Aggregation Method for Model-heterogeneous Federated Learning
This addresses a practical problem in federated learning for scenarios with varied client hardware, though it is incremental as it builds on existing methods.
The paper tackles the challenge of system heterogeneity in federated learning by proposing FedTSA, a cluster-based two-stage aggregation method that allows clients to train different models based on their resource capabilities, and it outperforms baselines in experiments.
Despite extensive research into data heterogeneity in federated learning (FL), system heterogeneity remains a significant yet often overlooked challenge. Traditional FL approaches typically assume homogeneous hardware resources across FL clients, implying that clients can train a global model within a comparable time frame. However, in practical FL systems, clients often have heterogeneous resources, which impacts their training capacity. This discrepancy underscores the importance of exploring model-heterogeneous FL, a paradigm allowing clients to train different models based on their resource capabilities. To address this challenge, we introduce FedTSA, a cluster-based two-stage aggregation method tailored for system heterogeneity in FL. FedTSA begins by clustering clients based on their capabilities, then performs a two-stage aggregation: conventional weight averaging for homogeneous models in Stage 1, and deep mutual learning with a diffusion model for aggregating heterogeneous models in Stage 2. Extensive experiments demonstrate that FedTSA not only outperforms the baselines but also explores various factors influencing model performance, validating FedTSA as a promising approach for model-heterogeneous FL.