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

• URL: http://arxiv.org/abs/2605.16690v1
• Date: Fri, 15 May 2026 23:06:59 GMT
• Title: UB-SMoE: Universally Balanced Sparse Mixture-of-Experts for Resource-adaptive Federated Fine-tuning of Foundation Models
• Abstract summary: We propose Universally Balanced Sparse Mixture-of-Experts (UB-SMoE) to address expert utilization imbalance and Top-K routing problems.<n>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.
• Score: 2.3099144596725565
• License: http://creativecommons.org/licenses/by-nc-sa/4.0/
• 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 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.

Related papers

• REFLEX: Rethinking MoE Inference as Refinement-Aware Compute Allocation in Diffusion Language Models [15.630466998208261]
  Mixture-of-experts (MoE) models increase parameter capacity by activating only a small subset of experts for each token.<n>In diffusion language models (DLMs), however, each denoising forward jointly revisits all token positions despite their sharply different refinement demands.<n>We argue that MoE inference in DLMs should therefore be viewed as refinement-aware compute allocation across heterogeneous token refinement states.
  arXiv  Detail & Related papers  (2026-08-03T06:59:06Z)
• Federated Foundation Models Fine-Tuning with Heterogeneous Compressed Clients [22.724544289853558]
  FedSLM is a parameter-centric framework for federated fine-tuning with heterogeneous compressed clients.<n>We provide theoretical guarantees for adapter-level aggregation, subspace-alignment bounds for cross-group fusion, and a characterization of how the confidence loss mitigates weak-supervision noise.
  arXiv  Detail & Related papers  (2026-07-31T06:45:56Z)
• $φ$-Balancing for Mixture-of-Experts Training [20.061328126956028]
  Mixture-of-Experts (MoE) models rely on balanced expert utilization to fully realize their scalability.<n>We propose $$-balancing, a principled framework that directly targets population-level expert balance by minimizing a strictly convex, symmetric, and differentiable potential of the expected routing distribution.
  arXiv  Detail & Related papers  (2026-05-14T20:39:28Z)
• A Replicate-and-Quantize Strategy for Plug-and-Play Load Balancing of Sparse Mixture-of-Experts LLMs [64.8510381475827]
  Sparse Mixture-of-Experts (SMoE) architectures are increasingly used to scale large language models efficiently.<n>SMoE models often suffer from severe load imbalance across experts, where a small subset of experts receives most tokens while others are underutilized.<n>We present a systematic analysis of expert routing during inference and identify three findings: (i) load imbalance persists and worsens with larger batch sizes, (ii) selection frequency does not reliably reflect expert importance, and (iii) overall expert workload and importance can be estimated using a small calibration set.
  arXiv  Detail & Related papers  (2026-02-23T15:11:16Z)
• Mitigating Reward Hacking in RLHF via Bayesian Non-negative Reward Modeling [49.41422138354821]
  We propose a principled reward modeling framework that integrates non-negative factor analysis into the Bradley-Terry preference model.<n>BNRM represents rewards through a sparse, non-negative latent factor generative process.<n>We show that BNRM substantially mitigates reward over-optimization, improves robustness under distribution shifts, and yields more interpretable reward decompositions than strong baselines.
  arXiv  Detail & Related papers  (2026-02-11T08:14:11Z)
• Outcome-Grounded Advantage Reshaping for Fine-Grained Credit Assignment in Mathematical Reasoning [60.00161035836637]
  Group Relative Policy Optimization has emerged as a promising critic-free reinforcement learning paradigm for reasoning tasks.<n>We introduce Outcome-grounded Advantage Reshaping (OAR), a fine-grained credit assignment mechanism that redistributes advantages based on how much each token influences the model's final answer.<n>OAR-G achieves comparable gains with negligible computational overhead, both significantly outperforming a strong GRPO baseline.
  arXiv  Detail & Related papers  (2026-01-12T10:48:02Z)
• Enhanced-FQL($λ$), an Efficient and Interpretable RL with novel Fuzzy Eligibility Traces and Segmented Experience Replay [0.0]
  Enhanced-FQL($$) achieves superior sample efficiency and reduced variance compared to n-step fuzzy TD and fuzzyA($$) baselines.<n>The framework's inherent interpretability, combined with its computational efficiency and theoretical convergence guarantees, makes it suitable for safety-critical applications.
  arXiv  Detail & Related papers  (2026-01-07T20:59:18Z)
• Mixture of Ranks with Degradation-Aware Routing for One-Step Real-World Image Super-Resolution [76.66229730098759]
  In real-world image super-resolution (Real-ISR), existing approaches mainly rely on fine-tuning pre-trained diffusion models.<n>We propose a Mixture-of-Ranks (MoR) architecture for single-step image super-resolution.<n>We introduce a fine-grained expert partitioning strategy that treats each rank in LoRA as an independent expert.
  arXiv  Detail & Related papers  (2025-11-20T04:11:44Z)
• FLAME: Towards Federated Fine-Tuning Large Language Models Through Adaptive SMoE [21.860699562235776]
  FLAME is a novel federated learning framework based on the Sparse Mixture-of-Experts (SMoE) architecture.<n>It retains full (uncompressed) global LoRA matrices and achieves client-side adaptability by varying the number of activated experts per client.<n>It tackles these challenges through a lightweight rescaling mechanism and an activation-aware aggregation scheme.
  arXiv  Detail & Related papers  (2025-06-19T21:02:19Z)
• NDCG-Consistent Softmax Approximation with Accelerated Convergence [67.10365329542365]
  We propose novel loss formulations that align directly with ranking metrics.<n>We integrate the proposed RG losses with the highly efficient Alternating Least Squares (ALS) optimization method.<n> Empirical evaluations on real-world datasets demonstrate that our approach achieves comparable or superior ranking performance.
  arXiv  Detail & Related papers  (2025-06-11T06:59:17Z)
• Distributional Robustness and Regularization in Reinforcement Learning [62.23012916708608]
  We introduce a new regularizer for empirical value functions and show that it lower bounds the Wasserstein distributionally robust value function. It suggests using regularization as a practical tool for dealing with $textitexternal uncertainty$ in reinforcement learning.
  arXiv  Detail & Related papers  (2020-03-05T19:56:23Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.