FLEX-MoE: Federated Mixture-of-Experts with Load-balanced Expert Assignment
This addresses scalability and performance issues for federated learning on edge devices with heterogeneous data, representing an incremental improvement over existing methods.
The paper tackles the challenges of deploying Mixture-of-Experts models in federated learning, such as resource constraints and load imbalance from non-IID data, by proposing FLEX-MoE, which optimizes expert assignment and load balancing, achieving superior performance and balanced utilization across three datasets.
Mixture-of-Experts (MoE) models enable scalable neural networks through conditional computation. However, their deployment with federated learning (FL) faces two critical challenges: 1) resource-constrained edge devices cannot store full expert sets, and 2) non-IID data distributions cause severe expert load imbalance that degrades model performance. To this end, we propose \textbf{FLEX-MoE}, a novel federated MoE framework that jointly optimizes expert assignment and load balancing under limited client capacity. Specifically, our approach introduces client-expert fitness scores that quantify the expert suitability for local datasets through training feedback, and employs an optimization-based algorithm to maximize client-expert specialization while enforcing balanced expert utilization system-wide. Unlike existing greedy methods that focus solely on personalization while ignoring load imbalance, our FLEX-MoE is capable of addressing the expert utilization skew, which is particularly severe in FL settings with heterogeneous data. Our comprehensive experiments on three different datasets demonstrate the superior performance of the proposed FLEX-MoE, together with its ability to maintain balanced expert utilization across diverse resource-constrained scenarios.