FedCDWA: Decoupled Federated Prototype Distillation with Hierarchical Wasserstein Aggregation

Federated learning enables decentralized clients to collaboratively train models without sharing local data. However, heterogeneous client distributions often induce client drift and hinder convergence. This paper proposes FedCDWA, a decoupled hierarchical federated distillation framework. FedCDWA decouples client-side personalized distillation from server-side mutual distillation to mitigate distillation-induced optimization conflicts. It further adopts Hierarchical Wasserstein Aggregation to aggregate prototypes without restrictive parametric assumptions while preserving intra-class structure and inter-class geometry. To achieve finer-grained feature alignment, Prototype–Variance Dual Alignment matches feature means and variances in the feature space. We prove convergence guarantees for FedCDWA. Experiments on three datasets demonstrate that FedCDWA consistently improves both global and personalized accuracy across heterogeneity levels, with smaller performance degradation under more severe heterogeneity.

Many machine learning applications need to learn from data stored across different users, hospitals, companies, or devices, but these data often cannot be collected in one place because of privacy or practical constraints. Federated learning addresses this by allowing participants to train a shared model while keeping their raw data local. However, the data held by different participants can be very different, which makes the shared model harder to train and can reduce accuracy. This paper proposes FedCDWA, a method that helps federated learning work better when participants have diverse data. Instead of directly sharing private data, each participant shares compact information about what it has learned. The server then combines this information in a more careful way, while also allowing each participant to keep useful local characteristics. Experiments show that this approach improves both the shared model and the participant-specific models, especially when the data distributions differ strongly across participants.