UB-SMoE: Universally Balanced Sparse Mixture-of-Experts for Resource-adaptive Federated Fine-tuning of Foundation Models

Heterogeneous LoRA-rank methods address system heterogeneity in federated fine-tuning of foundation models by assigning client-specific ranks based on computational capabilities. However, these methods achieve only marginal computational savings, as dense feed-forward computations dominate. Sparse Mixture-of-Experts (SMoE) provides a promising alternative through conditional computation, yet we identify that its naive application to heterogeneous federated settings introduces two critical discordances: (i) expert utilization imbalance and (ii) non-differentiability of Top-K routing. Our convergence analysis demonstrates that these discordances lead to degraded convergence, particularly for resource-constrained clients. To address these challenges, we propose Universally Balanced Sparse Mixture-of-Experts (UB-SMoE), which introduces Dynamic Modulated Routing (DMR) to rebalance expert utilization, and Universal Pseudo-Gradient (PG) to reconstruct learning signals for non-activated experts. These mechanisms form a self-reinforcing cycle that maintains expert viability across heterogeneous clients. Experiments on benchmarks show that UB-SMoE achieves up to $45.0\%$ computational reduction on low-resource clients while improving their performance by $8.7 \times$ compared to existing heterogeneous LoRA-rank methods.

Modern AI foundation models can be specialized for private domains, such as medical records across hospitals, through federated learning - where devices collaborate on training without sharing raw data. But real-world devices range from powerful servers to tiny edge devices, forcing a trade-off: either everyone uses a small model the weakest device can run, or weak devices barely benefit at all. A promising fix for this is giving the model a set of "experts" and lets each device call only as many as it can afford. Unfortunately, this creates a rich-get-richer dynamic: experts favored by powerful devices get all the "practice", while experts that weaker devices depend on are trained the least, creating a loop that causes extreme experts imbalance. We introduce UB-SMoE, which solves this with two cooperating ideas: a router that promotes neglected-but-relevant experts using usage statistics aggregated across devices, and a pseudo training signal that keeps neglected experts from falling behind. Together they cut the computation low-resource devices need by a half while improving their results several times. This work helps everyday devices, not just data centers, to enable usage of powerful foundation models.