FLEX-MoE: Federated Mixture-of-Experts with Load-balanced Expert Assignment for Edge Computing

TL;DR

FLEX-MoE is a federated mixture-of-experts framework designed for edge computing that balances expert load and improves model accuracy under resource constraints and non-IID data.

Contribution

It introduces a novel client-expert fitness scoring and an optimization algorithm to jointly enhance expert specialization and load balancing in federated MoE models.

Findings

Achieves superior accuracy compared to baseline methods.

Ensures balanced expert utilization across heterogeneous edge devices.

Performs well under resource constraints and diverse data distributions.

Abstract

Mixture-of-Experts (MoE) models enable scalable neural networks through conditional computation, offering enhanced effectiveness and efficiency for next-generation wireless communications. However, deploying MoE with federated learning (FL) over wireless and IoT edge networks faces two critical challenges: 1) resource-constrained clients cannot store large AI models with full expert sets, and 2) non-IID data distributions cause severe expert load imbalance that degrades model performance. To this end, we propose FLEX-MoE, a 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 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 greedy methods that focus solely on personalization while ignoring load imbalance, FLEX-MoE addresses expert utilization skew, which is particularly severe in heterogeneous edge FL. Our experimental results demonstrate superior accuracy and consistently balanced expert utilization across diverse resource-constrained scenarios for edge computing.