Avoid Forgetting by Preserving Global Knowledge Gradients in Federated Learning with Non-IID Data

TL;DR

This paper introduces FedProj, a federated learning framework that preserves the global decision boundary during local training on non-IID data, addressing catastrophic forgetting and improving model performance.

Contribution

The paper presents a novel federated learning method that maintains the global decision boundary by using ensemble knowledge transfer and episodic memory, outperforming existing approaches.

Findings

FedProj significantly reduces forgetting of the global decision boundary.

The ensemble knowledge transfer loss improves global model calibration.

Experimental results show FedProj outperforms state-of-the-art methods on non-IID data.

Abstract

The inevitable presence of data heterogeneity has made federated learning very challenging. There are numerous methods to deal with this issue, such as local regularization, better model fusion techniques, and data sharing. Though effective, they lack a deep understanding of how data heterogeneity can affect the global decision boundary. In this paper, we bridge this gap by performing an experimental analysis of the learned decision boundary using a toy example. Our observations are surprising: (1) we find that the existing methods suffer from forgetting and clients forget the global decision boundary and only learn the perfect local one, and (2) this happens regardless of the initial weights, and clients forget the global decision boundary even starting from pre-trained optimal weights. In this paper, we present FedProj, a federated learning framework that robustly learns the global decision boundary and avoids its forgetting during local training. To achieve better ensemble knowledge fusion, we design a novel server-side ensemble knowledge transfer loss to further calibrate the learned global decision boundary. To alleviate the issue of learned global decision boundary forgetting, we further propose leveraging an episodic memory of average ensemble logits on a public unlabeled dataset to regulate the gradient updates at each step of local training. Experimental results demonstrate that FedProj outperforms state-of-the-art methods by a large margin.