FLEx: Personalized Federated Learning for Mixture-of-Experts LLMs via Expert Grafting

Fan Liu, Bikang Pan, Zhongyi Wang, Xi Yao, Zhichao Sheng, Jingya Wang, Xiaoying Tang, Ye Shi

Keywords: Large Language Models; Federated Learning; Mixture of Experts; Personalization

TL;DR: We introduce FLEx, a federated learning framework for MoE-based LLMs that freezes pretrained experts to preserve general knowledge and locally trains a new, personalized expert "grafted" from existing ones to efficiently handle data heterogeneity.

Abstract: Federated instruction tuning of large language models (LLMs) is challenged by significant data heterogeneity across clients, demanding robust personalization. The Mixture of Experts (MoE) architecture, where experts can specialize in distinct data patterns, presents a natural architectural solution to this challenge. The inherent sparsity of the MoE architecture, achieved by selectively activating experts, poses a significant challenge to its integration with federated learning (FL). Conventional FL frameworks, designed for dense models, naively aggregate all expert parameters irrespective of their local activation patterns. This naive approach not only undermines MoE's dynamic sparsity but also risks corrupting the world knowledge within pretrained experts. To address this, we propose FLEx (Federated LLMs with Personalized Experts), a novel framework that leverages pretrained MoE-based LLMs for efficient personalization. By aggregating only the shared non-expert parameters, FLEx significantly reduces communication overhead and preserves the world knowledge stored within the frozen pretrained experts. For personalization, we introduce a novel expert grafting mechanism that leverages dynamic sparsity to construct a client-specific expert from selected components of pretrained experts, tailored to local data. This grafted expert is then fine-tuned locally alongside the gating mechanism. This joint training enables the model to learn when to leverage the shared knowledge from frozen experts and when to employ the personalized one. Evaluations on diverse, non-IID instruction tuning datasets show that FLEx consistently outperforms federated baselines on average, while demonstrating strong knowledge preservation on the knowledge-driven benchmark MMLU. Our code is available at \href{https://anonymous.4open.science/r/FLEx-8F12}{\texttt{https://anonymous.4open.science/r/FLEx-8F12}}.

Primary Area: foundation or frontier models, including LLMs

Submission Number: 5190