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

Abstract

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 a 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× compared to existing heterogeneous LoRA-rank methods.

Motivation

Real-world federated networks comprise devices with vastly different computational budgets. A single model configuration for all devices is ineffective: we cannot exploit high-end clients if the global model is shrunk for edge devices, and a large model cannot run on low-resource clients. We need resource-adaptive fine-tuning, where models scale their capacity to match client capabilities.

The Two Critical Discordances

Theoretical Foundation

The convergence analysis traces the accuracy limit to two effects: biased sparse-MoE gradients and an error floor that grows as computation budgets shrink.

UB-SMoE Framework

UB-SMoE resolves both discordances with two complementary mechanisms that form a self-reinforcing cycle: PG maintains expert viability through approximate gradients, enabling DMR to balance utilization, which generates real gradients that further refine all experts.

Main Results

We evaluate UB-SMoE on Commonsense-15K (8 commonsense-reasoning datasets) and the TeleQuAD telecommunications question-answering benchmark, using OLMoE-1B-7B and OLMo-1B base models. Baselines span heterogeneous LoRA-rank methods (HetLoRA, FlexLoRA, FLoRA, FLoRIST) and heterogeneous sparsity methods (A$^3$SMoE, SMoE-LLB).

Commonsense-15K: SMoE Backbone (OLMoE-1B-7B)

Commonsense-15K: Dense Backbone (OLMo-1B)

Performance Across Resource Budgets (the 8.7× story)

Across four computation budgets, from β₁ to β₄, heterogeneous LoRA-rank methods collapse under tight budgets while sparsity-based methods stay effective.

TeleQuAD: Telecommunications QA (non-IID)

BERTScore F1 on the domain-specific TeleQuAD benchmark under non-IID client data, where UB-SMoE is best at every budget.

Ablation: Contribution of Each Mechanism

Scalability to 32 Clients

Key Findings

Citation

If you find this work useful in your research, please consider citing:

@inproceedings{tran2026ubsmoe,
  title     = {UB-SMoE: Universally Balanced Sparse
               Mixture-of-Experts for Resource-adaptive
               Federated Fine-tuning of Foundation Models},
  author    = {Tran, Van-Tuan and Nguyen-Le, Hong-Hanh
               and Ruffini, Marco and Dzaferagic, Merim},
  booktitle = {Proceedings of the International Conference
               on Machine Learning (ICML)},
  year      = {2026}
}