Improving federated learning by routing data to specialized expert models
From Task Mixtures to Specialized Experts
Machine Learning
Summary
When many devices work together to learn from their data, the data each device has can come from many different tasks mixed together. The authors show that instead of trying to train one model per device, it’s better to route each input to specialized models tailored to each task. They develop a way to find these task-specific models using some public examples and show that this specialized routing improves overall performance and helps the worst-performing devices.
What this means in practice
- •For mobile device engineers: Design federated learning systems that assign user data inputs to specialized models per task, improving personalized model accuracy on mixed-task data.
- •For cloud ml platform teams: Implement expert specialized routing to better handle diverse client task mixtures in collaborative model fine-tuning services.
Authors
Hojat Allah Salehi, Mehrdad Mahdavi, Andrew Arash Mahyari, M. Hadi Amini
Abstract
In collaborative foundation model fine-tuning, client data is rarely homogeneous. Instead, clients typically possess unknown mixtures of distinct data distributions, or tasks. Conventional federated learning primarily addresses heterogeneity across clients without explicitly resolving latent task mixtures within each client. We study this setting as compound heterogeneity, where data is heterogeneous both across and within clients. We study adaptation over a common frozen representation and show that, when tasks share the same feature geometry, the optimal model for a client's task mixture under squared loss is a convex combination of the optimal models for its underlying tasks. Thus, a single locally trained model represents the client's overall task mixture, while individual inputs may be drawn from different underlying task distributions. This motivates routing inputs to specialized experts, and we show that, when the task optima form a simplex, task-aligned routing achieves lower risk than any single adapted model for genuinely mixed clients. With access to a small set of task-labeled public samples, we derive a convex program to recover task experts and match them to their corresponding tasks. Our routing analysis shows that effective specialization requires input-dependent expert selection aligned with each client's task mixture. Motivated by this analysis, we propose FedSEE. Across our experiments, FedSEE avoids the negative transfer observed in the evaluated baselines and improves performance by 2.9 points overall and 3.7 points for the worst-served quartile.