Layer-wise Adaptive Model Aggregation for Scalable Federated Learning

Sunwoo Lee, Tuo Zhang, Chaoyang He, Salman Avestimehr

Keywords: federated learning, model aggregation, neural network

Abstract: In Federated Learning, a common approach for aggregating local models across clients is periodic averaging of the full model parameters. It is, however, known that different layers of neural networks can have a different degree of model discrepancy across the clients. The conventional full aggregation scheme does not consider such a difference and synchronizes the whole model parameters at once, resulting in inefficient network bandwidth consumption. Aggregating the parameters that are similar across the clients does not make meaningful training progress while increasing the communication cost. We propose FedLAMA, a layer-wise model aggregation scheme for scalable Federated Learning. FedLAMA adaptively adjusts the aggregation interval in a layer-wise manner, jointly considering the model discrepancy and the communication cost. The layer-wise aggregation method enables to finely control the aggregation interval to relax the aggregation frequency without a significant impact on the model accuracy. Our empirical study shows that FedLAMA reduces the communication cost by up to $60\%$ for IID data and $70\%$ for non-IID data while achieving a comparable accuracy to FedAvg.

One-sentence Summary: Layer-wise adaptive model aggregation scheme for communication-efficient Federated Learning.