Lay Importance on the Layer: Federated Learning for Non-IID data with Layer-based Regularization

We propose Federated Learning with Layer-Adaptive Proximation (FedLap), a new class of algorithms to tackle non-IID data. The novelty is a new regularization term in the local (client) loss function which generalizes that of FedProx and captures divergence between global and local model weights of each client, at the level of Deep Neural Network (DNN) layers. Thus, the weights of different layers of the DNN are treated differently in the regularization function. Divergence between the global and local models is captured through a dissimilarity metric and a distance metric, both applied to each DNN layer. Since regularization is applied per layer and not to all weights as in FedProx, during local updates, only the weights of those layers that drift away from the global model change, while the other weights are not affected. Compared to FedAvg and FedProx, FedLap achieves 3–5 % higher accuracy in the first 20 communication rounds, and up to 10% higher accuracy in cases of unstable client participation. Thus, FedLap paves the way for a novel layer-aware class of algorithms in distributed learning.