FLEX-MoE: Federated Mixture-of-Experts with Load-balanced Expert Assignment
Boyang Zhang, Xiaobing Chen, Songyang Zhang, Shuai Zhang, Xiangwei Zhou, Mingxuan Sun
TL;DR
FLEX-MoE tackles federated learning with Mixture-of-Experts under resource constraints and non-IID data by introducing a server-driven framework that jointly optimizes expert assignment and load balancing. It leverages a client-expert fitness score matrix $Q(c,e)$ and an ILP-based assignment (via PuLP) to select $k_c$ experts per client, while maintaining balanced expert utilization through dynamic per-expert load bounds and EMA-based metrics $U(e)$ and $D_e^{(t)}$. The method updates $Q$ using either accuracy or loss-driven indicators and a gating-aware routing during local training, and aggregates updates with weighted contributions from participating clients. Experimental results on CIFAR-10, EMNIST, and GTSRB show FLEX-MoE achieves superior load balance (lower $\mathrm{CV}$ and $\mathrm{Gap}$) with competitive accuracy, notably under severely non-IID data distributions, demonstrating practical viability for edge-device deployments.
Abstract
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.