Exploiting Personalized Invariance for Better Out-of-distribution Generalization in Federated Learning

TL;DR

This paper introduces a dual-regularized personalized federated learning framework that leverages invariant features to improve out-of-distribution generalization across clients with Non-IID data.

Contribution

It proposes a novel dual-regularized learning approach that explicitly explores personalized invariance, outperforming existing methods in OOD generalization for federated learning.

Findings

Theoretically proven convergence and OOD generalization benefits.

Outperforms existing federated and invariant learning methods in diverse scenarios.

Effective in handling Non-IID data and distribution shifts.

Abstract

Recently, data heterogeneity among the training datasets on the local clients (a.k.a., Non-IID data) has attracted intense interest in Federated Learning (FL), and many personalized federated learning methods have been proposed to handle it. However, the distribution shift between the training dataset and testing dataset on each client is never considered in FL, despite it being general in real-world scenarios. We notice that the distribution shift (a.k.a., out-of-distribution generalization) problem under Non-IID federated setting becomes rather challenging due to the entanglement between personalized and spurious information. To tackle the above problem, we elaborate a general dual-regularized learning framework to explore the personalized invariance, compared with the exsiting personalized federated learning methods which are regularized by a single baseline (usually the global model). Utilizing the personalized invariant features, the developed personalized models can efficiently exploit the most relevant information and meanwhile eliminate spurious information so as to enhance the out-of-distribution generalization performance for each client. Both the theoretical analysis on convergence and OOD generalization performance and the results of extensive experiments demonstrate the superiority of our method over the existing federated learning and invariant learning methods, in diverse out-of-distribution and Non-IID data cases.