Optimizing Hierarchical Federated Learning: A Reinforcement Learning Approach

Federated learning (FL) represents a decentralized machine learning paradigm where multiple clients cooperatively train a global model over local data, under the coordination of a central server. FL clients exhibit two-fold heterogeneity. First, data samples are unbalanced cross clients, and are not independent and identically distributed (non-IID). Second, clients possess different computation and communication capabilities. Such heterogeneity affects FL model convergence and training cost. In addition, the design of FL is complicated by the large communication overhead caused by frequent model parameter updates from a potentially enormous number of clients. To reduce communication overhead, we present a three-layer hierarchical FL architecture, where a middle layer with cluster servers is introduced for intra-cluster aggregation and global model updates exchange. We address the performance impact from client heterogeneity, and aim to balance two objectives, i.e., maximizing FL model accuracy and minimizing the training cost including computation cost and communication cost. Leveraging a deep reinforcement learning approach, we design a dynamic aggregation weight allocation framework, DaWa, to effectively select clients and obtain the optimal aggregation weights for the selected clients. Extensive experiments under empirical configurations are conducted to verify the efficacy of DaWa. The results show that DaWa can improve convergence rate by up to 21 while reducing training cost by up to 55 state-of-the-art FL frameworks.