# pFedAFM: Adaptive Feature Mixture for Batch-Level Personalization in Heterogeneous Federated Learning

arXiv: https://arxiv.org/abs/2404.17847 · HTML: https://mlsanity.com/paper/2404.17847/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 58 references that are in the map.
- First posted on arXiv 27 Apr 2024.

## Builds on (most cited first, 30 of 33)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [An overview of gradient descent optimization algorithms](https://mlsanity.com/paper/1609.04747/) (2016)
- [Generalized Cross Entropy Loss for Training Deep Neural Networks with Noisy Labels](https://mlsanity.com/paper/1805.07836/) (2018)
- [Ensemble Distillation for Robust Model Fusion in Federated Learning](https://mlsanity.com/paper/2006.07242/) (2020)
- [Exploiting Shared Representations for Personalized Federated Learning](https://mlsanity.com/paper/2102.07078/) (2021)
- [FedMD: Heterogenous Federated Learning via Model Distillation](https://mlsanity.com/paper/1910.03581/) (2019)
- [Data-Free Knowledge Distillation for Heterogeneous Federated Learning](https://mlsanity.com/paper/2105.10056/) (2021)
- [FedProto: Federated Prototype Learning across Heterogeneous Clients](https://mlsanity.com/paper/2105.00243/) (2021)
- [Federated Learning on Non-IID Data: A Survey](https://mlsanity.com/paper/2106.06843/) (2021)
- [HeteroFL: Computation and Communication Efficient Federated Learning for Heterogeneous Clients](https://mlsanity.com/paper/2010.01264/) (2020)
- [Think Locally, Act Globally: Federated Learning with Local and Global Representations](https://mlsanity.com/paper/2001.01523/) (2020)
- [Communication-Efficient On-Device Machine Learning: Federated Distillation and Augmentation under Non-IID Private Data](https://mlsanity.com/paper/1811.11479/) (2018)
- [FjORD: Fair and Accurate Federated Learning under heterogeneous targets with Ordered Dropout](https://mlsanity.com/paper/2102.13451/) (2021)
- [FedKD: Communication Efficient Federated Learning via Knowledge Distillation](https://mlsanity.com/paper/2108.13323/) (2021)
- [Distillation-Based Semi-Supervised Federated Learning for Communication-Efficient Collaborative Training with Non-IID Private Data](https://mlsanity.com/paper/2008.06180/) (2020)
- [Parameterized Knowledge Transfer for Personalized Federated Learning](https://mlsanity.com/paper/2111.02862/) (2021)
- [FedRolex: Model-Heterogeneous Federated Learning with Rolling Sub-Model Extraction](https://mlsanity.com/paper/2212.01548/) (2022)
- [Cronus: Robust and Heterogeneous Collaborative Learning with Black-Box Knowledge Transfer](https://mlsanity.com/paper/1912.11279/) (2019)
- [Heterogeneous Ensemble Knowledge Transfer for Training Large Models in Federated Learning](https://mlsanity.com/paper/2204.12703/) (2022)
- [Federated Mutual Learning](https://mlsanity.com/paper/2006.16765/) (2020)
- [FedAUX: Leveraging Unlabeled Auxiliary Data in Federated Learning](https://mlsanity.com/paper/2102.02514/) (2021)
- [FedGH: Heterogeneous Federated Learning with Generalized Global Header](https://mlsanity.com/paper/2303.13137/) (2023)
- [FedGEMS: Federated Learning of Larger Server Models via Selective Knowledge Fusion](https://mlsanity.com/paper/2110.11027/) (2021)
- [Fed2: Feature-Aligned Federated Learning](https://mlsanity.com/paper/2111.14248/) (2021)
- [Wireless Federated Distillation for Distributed Edge Learning with Heterogeneous Data](https://mlsanity.com/paper/1907.02745/) (2019)
- [FedZKT: Zero-Shot Knowledge Transfer towards Resource-Constrained Federated Learning with Heterogeneous On-Device Models](https://mlsanity.com/paper/2109.03775/) (2021)
- [Towards Personalized Federated Learning via Heterogeneous Model Reassembly](https://mlsanity.com/paper/2308.08643/) (2023)
- [FedMatch: Federated Learning Over Heterogeneous Question Answering Data](https://mlsanity.com/paper/2108.05069/) (2021)
- [FedClassAvg: Local Representation Learning for Personalized Federated Learning on Heterogeneous Neural Networks](https://mlsanity.com/paper/2210.14226/) (2022)

## Built on by (most cited first, 3 of 3)

- [Federated Model Heterogeneous Matryoshka Representation Learning](https://mlsanity.com/paper/2406.00488/) (2024)
- [FedSKD: Aggregation-free Model-heterogeneous Federated Learning via Multi-dimensional Similarity Knowledge Distillation for Medical Image Classification](https://mlsanity.com/paper/2503.18981/) (2025)
- [CRIP: Channel Level Representation Injection for Personalized One-Shot Federated Learning](https://mlsanity.com/paper/2608.02222/) (2026)

## Abstract (arXiv)

Model-heterogeneous personalized federated learning (MHPFL) enables FL clients to train structurally different personalized models on non-independent and identically distributed (non-IID) local data. Existing MHPFL methods focus on achieving client-level personalization, but cannot address batch-level data heterogeneity. To bridge this important gap, we propose a model-heterogeneous personalized Federated learning approach with Adaptive Feature Mixture (pFedAFM) for supervised learning tasks. It consists of three novel designs: 1) A sharing global homogeneous small feature extractor is assigned alongside each client's local heterogeneous model (consisting of a heterogeneous feature extractor and a prediction header) to facilitate cross-client knowledge fusion. The two feature extractors share the local heterogeneous model's prediction header containing rich personalized prediction knowledge to retain personalized prediction capabilities. 2) An iterative training strategy is designed to alternately train the global homogeneous small feature extractor and the local heterogeneous large model for effective global-local knowledge exchange. 3) A trainable weight vector is designed to dynamically mix the features extracted by both feature extractors to adapt to batch-level data heterogeneity. Theoretical analysis proves that pFedAFM can converge over time. Extensive experiments on 2 benchmark datasets demonstrate that it significantly outperforms 7 state-of-the-art MHPFL methods, achieving up to 7.93% accuracy improvement while incurring low communication and computation costs.