Communication-Efficient Loss Minimization over Heterogeneous Data with Federated Hierarchical Ensemble Aggregation via Distillation

Distributed optimization through federated learning (FL) suffers from data heterogeneity particularly when the client datasets are highly imbalanced. Model aggregation via ensemble distillation is an effective solution to address this issue. However, there is no previous work on ensemble distillation in FL that considers hierarchical model aggregation, which is important for reducing communication overhead over a large network. In this work, we propose new methods to enable ensemble distillation for a hierarchical FL system. We develop a Federated Hierarchical Ensemble Aggregation via Distillation (FedHEAD) algorithm that performs ensemble distillation by reusing the clients' local data within each network sector of the hierarchy. We also extend it to FedHEAD+ so as to take advantage of reference data when it is available at the server. We provide theoretical analysis on FedHEAD and FedHEAD+, showing that under a wide range of conditions, our proposed schemes achieve faster convergence than existing non-hierarchical alternatives. Furthermore, extensive experiments over computer vision, natural language processing, and network traffic classification datasets show that the proposed schemes are robust towards hierarchical model aggregation in the network.