FedStein: Enhancing Multi-Domain Federated Learning Through James-Stein Estimator
Keywords: Distributed Machine Learning, Federated Learning, Domain Generalization, Machine Learning
TL;DR: FedStein uses the James-Stein estimator to improve batch normalization in federated learning, ensuring consistent performance despite non-i.i.d data from diverse domains, addressing data heterogeneity challenges.
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. Non-BN layer parameters are exchanged via standard FL techniques. Extensive experiments conducted across three datasets and multiple models demonstrate that FedStein surpasses the existing methods such as FedAvg and FedBN, with accuracy improvements exceeding 14% in certain domains leading to enhanced domain generalization.
Submission Number: 38
Loading