pFedMoE: Data-Level Personalization with Mixture of Experts for Model-Heterogeneous Personalized Federated Learning

TL;DR

This work tackles data-level personalization in model-heterogeneous federated learning (MHPFL) by introducing pFedMoE, a mixture-of-experts framework that combines a global, small, homogeneous feature extractor with local, heterogeneous large extractors via a per-sample gating network. The method trains three MoE components end-to-end and only exchanges compact homogeneous extractors, balancing generalization and personalization while preserving privacy. Theoretical analysis shows a non-convex convergence rate of $\mathcal{O}(1/T)$ under standard assumptions, and extensive experiments on CIFAR-10/100 demonstrate state-of-the-art accuracy with lower computation and acceptable communication costs, robust to both pathological and practical non-IIDness. Overall, pFedMoE provides a data-level, model-heterogeneous FL solution that improves personalization, reduces communication, and scales across diverse client settings, with potential for extension to continual learning under distribution drift.

Abstract

Federated learning (FL) has been widely adopted for collaborative training on decentralized data. However, it faces the challenges of data, system, and model heterogeneity. This has inspired the emergence of model-heterogeneous personalized federated learning (MHPFL). Nevertheless, the problem of ensuring data and model privacy, while achieving good model performance and keeping communication and computation costs low remains open in MHPFL. To address this problem, we propose a model-heterogeneous personalized Federated learning with Mixture of Experts (pFedMoE) method. It assigns a shared homogeneous small feature extractor and a local gating network for each client's local heterogeneous large model. Firstly, during local training, the local heterogeneous model's feature extractor acts as a local expert for personalized feature (representation) extraction, while the shared homogeneous small feature extractor serves as a global expert for generalized feature extraction. The local gating network produces personalized weights for extracted representations from both experts on each data sample. The three models form a local heterogeneous MoE. The weighted mixed representation fuses generalized and personalized features and is processed by the local heterogeneous large model's header with personalized prediction information. The MoE and prediction header are updated simultaneously. Secondly, the trained local homogeneous small feature extractors are sent to the server for cross-client information fusion via aggregation. Overall, pFedMoE enhances local model personalization at a fine-grained data level, while supporting model heterogeneity.

Figures (11)

• Figure 1: Workflow of MoE.
• Figure 2: Workflow of pFedMoE.
• Figure 3: Gating network structure.
• Figure 4: Accuracy distribution for individual clients.
• Figure 5: T-SNE representation visualization results for FedProto and pFedMoE on CIFAR-10 (non-IID: 2/10).

Theorems & Definitions (4)

• Lemma 1
• Lemma 2
• Theorem 1
• Theorem 2