[Paper Review] Overcoming Forgetting in Federated Learning on Non-IID Data

Neta Shoham, Tomer Avidor+5 more|arXiv (Cornell University)|Oct 17, 2019

Privacy-Preserving Technologies in Data17 references143 citations

TL;DR

This paper adapts Elastic Weight Consolidation to Federated Learning, creating FedCurv, to prevent forgetting and improve convergence on non-IID data with full participation and efficient bandwidth usage.

ABSTRACT

We tackle the problem of Federated Learning in the non i.i.d. case, in which local models drift apart, inhibiting learning. Building on an analogy with Lifelong Learning, we adapt a solution for catastrophic forgetting to Federated Learning. We add a penalty term to the loss function, compelling all local models to converge to a shared optimum. We show that this can be done efficiently for communication (adding no further privacy risks), scaling with the number of nodes in the distributed setting. Our experiments show that this method is superior to competing ones for image recognition on the MNIST dataset.

Motivation & Objective

• Motivate the challenge of Federated Learning with non-IID data where local models drift apart.
• Adopt a Lifelong Learning-inspired penalty to preserve important parameters across rounds.
• Develop FedCurv to share curvature information while preserving privacy and bandwidth.
• Show that FedCurv improves convergence and scalability compared to baselines on MNIST.
• Explore how large local epochs per round (E) interact with forgetting mitigation.

Proposed method

• Adapt Elastic Weight Consolidation (EWC) to Federated Learning to form FedCurv.
• In each round, each node optimizes a local loss augmented with a penalty involving the diagonal Fisher information from other nodes (diag(I)).
• Share only aggregated gradient-related information across nodes to preserve privacy and keep bandwidth low.
• Express the round loss as L_t,s(θ) plus a sum over other nodes of (θ−θ̂_{t-1,j})^T diag(Ĩ_{t-1,j}) (θ−θ̂_{t-1,j}).
• Initialize θ̂_t as the average of node parameters from the previous round and perform E local SGD epochs.
• Maintain two aggregated vectors u_t and v_t to reconstruct the penalty term with minimal communication and allow secure aggregation.

Experimental results

Research questions

• RQ1 Does FedCurv outperform FedAvg and FedProx on non-IID data in terms of convergence speed and final accuracy?
• RQ2 How does the number of local epochs E per round affect performance under non-IID distributions?
• RQ3 Can the curvature-based penalty be implemented with low bandwidth and preserve privacy, akin to FedAvg?
• RQ4 What is the impact of using full participation versus partial participation in FedCurv?
• RQ5 How does FedCurv scale with the number of nodes in a simulated federated setup?

Key findings

• FedCurv achieves faster convergence than FedAvg, especially for larger E (e.g., 50), attaining 90% accuracy in fewer rounds.
• For E=50, FedCurv reaches 90% in 9 rounds and 95% in 38 rounds, outperforming FedAvg and FedProx on the non-IID MNIST setup.
• FedCurv also demonstrates superiority for E=10, with 35 rounds to 90% and 99 rounds to 95%, outperforming alternatives.
• FedProx shows improvements over FedAvg but is less effective than FedCurv at higher E values in achieving top accuracy.
• The method preserves privacy by sharing only aggregated gradient-related information and can use secure aggregation similarly to FedAvg.
• Bandwidth can be further reduced by leveraging a sparse diagonal Fisher information without sacrificing performance.