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

We formally analyze the convergence of heterogeneous federated SMoE and show that the two discordances are not merely empirical nuisances — they impose a fundamental limit on accuracy.

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 β₁–β₄, 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}
}