Aggregation Alignment for Federated Learning with Mixture-of-Experts under Data Heterogeneity
This work solves aggregation challenges for fine-tuning MoE-based LLMs in federated learning, which is important for applications with distributed, privacy-sensitive data, though it appears incremental as it builds on existing FL and MoE paradigms.
The paper tackles the problem of fine-tuning Mixture-of-Experts (MoE) large language models in federated learning under data heterogeneity, which causes aggregation issues like poor global gating and expert semantic blurring. The proposed FedAlign-MoE framework addresses these by aligning routing distributions and selectively aggregating semantically aligned experts, achieving faster convergence and superior accuracy compared to state-of-the-art benchmarks.
Large language models (LLMs) increasingly adopt Mixture-of-Experts (MoE) architectures to scale model capacity while reducing computation. Fine-tuning these MoE-based LLMs often requires access to distributed and privacy-sensitive data, making centralized fine-tuning impractical. Federated learning (FL) therefore provides a paradigm to collaboratively fine-tune MoE-based LLMs, enabling each client to integrate diverse knowledge without compromising data privacy. However, the integration of MoE-based LLM fine-tuning into FL encounters two critical aggregation challenges due to inherent data heterogeneity across clients: (i) divergent local data distributions drive clients to develop distinct gating preference for localized expert selection, causing direct parameter aggregation to produce a ``one-size-fits-none'' global gating network, and (ii) same-indexed experts develop disparate semantic roles across clients, leading to expert semantic blurring and the degradation of expert specialization. To address these challenges, we propose FedAlign-MoE, a federated aggregation alignment framework that jointly enforces routing consistency and expert semantic alignment. Specifically, FedAlign-MoE aggregates gating behaviors by aligning routing distributions through consistency weighting and optimizes local gating networks through distribution regularization, maintaining cross-client stability without overriding discriminative local preferences. Meanwhile, FedAlign-MoE explicitly quantifies semantic consistency among same-indexed experts across clients and selectively aggregates updates from semantically aligned clients, ensuring stable and specialized functional roles for global experts. Extensive experiments demonstrate that FedAlign-MoE outperforms state-of-the-art benchmarks, achieving faster convergence and superior accuracy in non-IID federated environments.