Papers for

cloud ml platform teams

Papers whose findings have a practical use for this group, as judged from the abstract. Open a paper to read what it means in practice.

Improving federated learning by routing data to specialized expert models

From Task Mixtures to Specialized Experts

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.

Wed 30 SeptMachine Learning
The gist
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.
Open → 2610.00580v1