Scalable Hierarchical Over-the-Air Federated Learning

• Published in: IEEE transactions on wireless communications 2024-08, Vol.23 (8), p.8480-8496
• Main Authors: Azimi-Abarghouyi, Seyed Mohammad, Fodor, Viktoria
• Format: Article
• Language: English
• Subjects: Algorithms; Cluster analysis; Downlink; Downlinking; Error analysis; Federated learning; Heterogeneity; hierarchical systems; Interference; Machine learning; over-the-air computation; Servers; Spatial data; Task analysis; Uplink; Uplinking; Wireless communication; Wireless networks; Wireless sensor networks
• ISSN: 1536-1276; 1558-2248; 1558-2248
• DOI: 10.1109/TWC.2024.3350923

When implementing hierarchical federated learning over wireless networks, scalability assurance and the ability to handle both interference and device data heterogeneity are crucial. This work introduces a new two-level learning method designed to address these challenges, along with a scalable over-the-air aggregation scheme for the uplink and a bandwidth-limited broadcast scheme for the downlink that efficiently use a single wireless resource. To provide resistance against data heterogeneity, we employ gradient aggregations. Meanwhile, the impact of uplink and downlink interference is minimized through optimized receiver normalizing factors. We present a comprehensive mathematical approach to derive the convergence bound for the proposed algorithm, applicable to a multi-cluster wireless network encompassing any count of collaborating clusters, and provide special cases and design remarks. As a key step to enable a tractable analysis, we develop a spatial model for the setup by modeling devices as a Poisson cluster process over the edge servers and rigorously quantify uplink and downlink error terms due to the interference. Finally, we show that despite the interference and data heterogeneity, the proposed algorithm not only achieves high learning accuracy for a variety of parameters but also significantly outperforms the conventional hierarchical learning algorithm.