Server assigns experts to improve language model training on edge devices
Server-Side Sparse Assignment for Federated Mixture-of-Experts
Machine LearningInformation Theory
Summary
Training large language models on user devices while keeping their data private is hard because different devices have different abilities and needs. The authors propose a new way for a central server to decide which parts of the model each device should use, based on guesses that get better over time. This method helps devices use fewer resources while still learning well from their own data. Tests showed their approach gets better accuracy than letting devices use all parts, while loading fewer model pieces.
What this means in practice
- •For edge device developers: Reduce memory and compute load on edge devices by assigning only the most relevant model components for local training.
- •For federated learning engineers: Improve personalized model fine-tuning by server-driven expert selection that matches client data patterns better than uniform access.
Authors
Xue Zheng, Aylin Yener
Abstract
Large language models (LLMs) are increasingly recognized for their remarkable capabilities; federated learning (FL) offers a promising framework for collaboratively fine-tuning LLMs on edge devices while achieving data privacy. However, edge devices are inherently heterogeneous: clients may exhibit diverse data distributions and resource constraints, including limited memory and computation budgets. Recent federated Mixture-of-Experts (MoE) approaches address these challenges by using sparse expert activation or partial expert loading. However, existing partial-expert methods rely on heuristic full-pool observations or prioritize system-wide load balance, without guarantees for recovering client-specific relevant experts from partial observations. We propose $S^3$-MoE, a server-side sparse expert assignment framework that estimates client-specific expert relevance from partial observations and refines assignments through an exploit--explore mechanism. We also establish relevance-estimation error and exact Top-$K$ recovery guarantees. Experiments using a CNN-based MoE on CIFAR-100 show that $S^3$-MoE recovers task-dependent client--expert preferences and reduces conflicts caused by correlated expert selection across clients. Compared with Full-20 FedAvg, which gives each client access to all 20 experts for local Top-5 routing, $S^3$-MoE achieves 1.31 percentage points higher mean final accuracy while loading only $5$ experts per client.