Shift Happens: Mixture of Experts based Continual Adaptation in Federated Learning

TL;DR

This paper presents ShiftEx, a mixture of experts framework for federated learning that adaptively handles distribution shifts in streaming data, significantly improving accuracy and adaptation speed in non-stationary environments.

Contribution

It introduces a novel shift-aware mixture of experts model with a latent memory and optimization strategy for dynamic adaptation to distributional shifts in federated learning.

Findings

Achieves 5.5-12.9% accuracy improvement over baselines.

Adapts 22-95% faster to distribution shifts.

Effectively handles covariate and label shifts in FL.

Abstract

Federated Learning (FL) enables collaborative model training across decentralized clients without sharing raw data, yet faces significant challenges in real-world settings where client data distributions evolve dynamically over time. This paper tackles the critical problem of covariate and label shifts in streaming FL environments, where non-stationary data distributions degrade model performance and necessitate a middleware layer that adapts FL to distributional shifts. We introduce ShiftEx, a shift-aware mixture of experts framework that dynamically creates and trains specialized global models in response to detected distribution shifts using Maximum Mean Discrepancy for covariate shifts. The framework employs a latent memory mechanism for expert reuse and implements facility location-based optimization to jointly minimize covariate mismatch, expert creation costs, and label imbalance. Through theoretical analysis and comprehensive experiments on benchmark datasets, we demonstrate 5.5-12.9 percentage point accuracy improvements and 22-95 % faster adaptation compared to state-of-the-art FL baselines across diverse shift scenarios. The proposed approach offers a scalable, privacy-preserving middleware solution for FL systems operating in non-stationary, real-world conditions while minimizing communication and computational overhead.