FedStein: Enhancing Multi-Domain Federated Learning Through James-Stein Estimator

TL;DR

FedStein introduces a novel federated learning method that leverages the James-Stein estimator to improve multi-domain data generalization by sharing only specific batch normalization statistics, leading to significant accuracy gains.

Contribution

This paper proposes FedStein, a new federated learning approach that enhances multi-domain performance by sharing James-Stein estimates of BN statistics, addressing domain heterogeneity.

Findings

FedStein outperforms FedAvg and FedBN in accuracy.

Achieves over 14% accuracy improvement in some domains.

Enhances domain generalization across multiple datasets.

Abstract

Federated Learning (FL) facilitates data privacy by enabling collaborative in-situ training across decentralized clients. Despite its inherent advantages, FL faces significant challenges of performance and convergence when dealing with data that is not independently and identically distributed (non-i.i.d.). While previous research has primarily addressed the issue of skewed label distribution across clients, this study focuses on the less explored challenge of multi-domain FL, where client data originates from distinct domains with varying feature distributions. We introduce a novel method designed to address these challenges FedStein: Enhancing Multi-Domain Federated Learning Through the James-Stein Estimator. FedStein uniquely shares only the James-Stein (JS) estimates of batch normalization (BN) statistics across clients, while maintaining local BN parameters. The non-BN layer parameters are exchanged via standard FL techniques. Extensive experiments conducted across three datasets and multiple models demonstrate that FedStein surpasses existing methods such as FedAvg and FedBN, with accuracy improvements exceeding 14% in certain domains leading to enhanced domain generalization. The code is available at https://github.com/sunnyinAI/FedStein