Overcoming Forgetting in Federated Learning on Non-IID Data

TL;DR

This paper proposes a novel method for federated learning on non-IID data by adding a penalty term to the loss function, which helps local models converge to a shared optimum, improving performance without extra privacy risks.

Contribution

It introduces a penalty-based approach inspired by Lifelong Learning to address model drift in federated learning with non-IID data, scalable and communication-efficient.

Findings

Outperforms existing methods on MNIST image recognition.

Efficiently scales with the number of nodes in the network.

Adds no additional privacy risks during communication.

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.