pFedAFM: Adaptive Feature Mixture for Data-Level Personalization in Heterogeneous Federated Learning on Mobile Edge Devices

Federated learning (FL), an emerging distributed machine learning paradigm, utilizes edge decentralized data from multiple edge nodes (clients) to train a shared model under preserved data privacy. Furthermore, model-heterogeneous personalized federated learning (MHPFL) enables FL clients to train structurally different personalized models on non-independent and identically distributed (non-lID) local data. Existing MHPFL methods focus on data distribution differences among clients, and they propose various client-level personalization approaches to alleviate non-lID issues. However, different data samples in one client may also have different features, which are often ignored, resulting in constrained model performances. To bridge this gap, we propose a novel model-heterogeneous personalized Federated learning approach with Adaptive Feature Mixture (pFedAFM) to achieve data-level personalization while maintaining efficient communication and computation. It consists of three innovative designs: 1) We add a homogeneous small feature extractor alongside each client's local heterogeneous model, and the server aggregates these homogeneous small feature extractors for cross-client knowledge fusion. 2) We design an iterative training strategy to alternately train the global homogeneous small feature extractor and the local heterogeneous client model, for effective bidirectional exchange between global generalized knowledge and local personalized knowledge. 3) During model training, we devise a trainable weight vector to adaptively mix the features (representation) extracted by the global homogeneous and local heterogeneous models for different data samples, i.e., fulfilling data-level personalized feature mixture. Theoretical analysis proves that pFedAFM converges over time. Extensive experiments on 3 computer vision (CV) and 1 natural lan-guage processing (NLP) benchmark datasets demonstrate that pFedAFM significantly outperforms 8 state-of-the-art MHPFL methods, achieving up to 7.93% accuracy improvement while incurring low communication and computation costs.