FedPAGR: Federated Prototype Alignment via Geometric Refinement for Heterogeneous Architectures
Paper · 2026 · Kris Prasad and Md Abdullah Al Hafiz Khan
Summary
The paper addresses federated learning when clients run different model architectures, so their parameters cannot be averaged. Prototype-exchange methods such as FedProto avoid this by sharing per-class mean embeddings. Under strongly skewed label distributions, however, plain averaging produces noisy prototypes that are not well separated. FedPAGR changes the pipeline in three ways:
- Shared space: each client maps its features through a private projection head into a common normalized d-dimensional space.
- Server-side refinement: instead of a trainable server network (as in FedTGP), the server runs a few SGD steps directly on the prototype matrix. The objective balances closeness to client submissions against an angular margin between classes.
- Client training: clients copy the refined prototypes into their classifier weights. They then train with cross-entropy, a temperature-scaled prototype-alignment loss, and an entropy term.
The method was tested on four image benchmarks and on MIMIC-IV with ClinicalBERT clients fine-tuned to different depths, using 20 clients, 5 architectures and Dirichlet α=0.1. FedPAGR gives the best ensemble accuracy on every image dataset and the best local accuracy on CIFAR-10, Fashion-MNIST and MIMIC-IV. This matters for privacy-sensitive NLP, where institutions use different encoders and cannot share data, public corpora or models.
Key points
- MIMIC-IV: clients used ClinicalBERT variants that were frozen, or had 2 or 4 layers fine-tuned. FedPAGR reached 76.38% local accuracy, compared with 71.39% for FedTGP (+4.99). FedProto (11.98%) and FedProc (7.89%) essentially failed.
- CIFAR-10: local / ensemble accuracy was 83.05% / 40.76%, versus 74.98% / 38.40% for FedTGP.
- High-class datasets: FedTGP keeps higher local accuracy on CIFAR-100 (44.15 vs. 41.86) and Tiny ImageNet (24.18 vs. 23.07). FedPAGR still wins on ensemble accuracy there (23.10 and 10.48).
- Ablations on CIFAR-10:
- Dropping the alignment terms (β=0) reduces accuracy to about chance (9.97%).
- Removing separation lowers ensemble accuracy by 2.73 points; removing entropy lowers it by 3.32 points.
- More refinement steps mostly help ensemble accuracy, not local accuracy.
- Communication cost: each client exchanges only C×d prototype floats, about 20KB each way per round for CIFAR-10 with d=512. The separation loss grows quadratically with the number of classes, and the evaluation uses only synthetic Dirichlet label skew.
Institutions
Fields, topics and concepts
📚 Shelves 2The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 works
Cite
Cite as Kris Prasad and Md Abdullah Al Hafiz Khan. 2026. FedPAGR: Federated Prototype Alignment via Geometric Refinement for Heterogeneous Architectures. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 4: Student Research Workshop).
@inproceedings{prasad2026fedpagr,
title = {{FedPAGR: Federated Prototype Alignment via Geometric Refinement for Heterogeneous Architectures}},
author = {Kris Prasad and Md Abdullah {Al Hafiz Khan}},
booktitle = {Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 4: Student Research Workshop)},
year = {2026},
doi = {10.18653/v1/2026.acl-srw.37},
url = {https://doi.org/10.18653/v1/2026.acl-srw.37}
}References
- Asad Aali, Dave Van Veen, Yamin Arefeen, Jason Hom, Christian Bluethgen, Eduardo Pontes Reis, Sergios Gatidis, Namuun Clifford, Joseph Daws, Arash Tehrani, Jangwon Kim, and Akshay Chaudhari. 2025. MIMIC-IV-Ext-BHC: Labeled Clinical Notes Dataset for Hospital Course Summarization. PhysioNet. Version… find
- Asad Aali, Dave Van Veen, Yamin Ishraq Arefeen, Jason Hom, Christian Bluethgen, Eduardo Pontes Reis, Sergios Gatidis, Namuun Clifford, Joseph Daws, Arash S Tehrani, Jangwon Kim, and Akshay S Chaudhari. 2024. A dataset and benchmark for hospital course summarization with adapted large language models… find
- Samiul Alam, Luyang Liu, Ming Yan, and Mi Zhang. 2022. Fedrolex: Model-heterogeneous federated learning with rolling sub-model extraction. In Advances in Neural Information Processing Systems. find
- Yutong Dai, Zeyuan Chen, Junnan Li, Shelby Heinecke, Lichao Sun, and Ran Xu. 2023. Tackling data heterogeneity in federated learning with class prototypes. In Proceedings of the AAAI Conference on Artificial Intelligence. find
- Enmao Diao, Jie Ding, and Vahid Tarokh. 2021. Heterofl: Computation and communication efficient federated learning for heterogeneous clients. In International Conference on Learning Representations. find
- Boyu Fan, Siyang Jiang, Xiang Su, Sasu Tarkoma, and Pan Hui. 2024. A survey on model-heterogeneous federated learning: Problems, methods, and prospects. IEEE International Conference on Big Data. find
- Samuele Fonio, Mirko Polato, and Roberto Esposito. 2024. Fedhp: Federated learning with hyperspherical prototypical regularization. In European Symposium on Artificial Neural Networks. find
- A. E. W. Johnson, L. Bulgarelli, L. Shen, and et al. 2023. Mimic-iv, a freely accessible electronic health record dataset. Scientific Data, 10.
- Alistair Johnson, Lucas Bulgarelli, Tom Pollard, Brian Gow, Benjamin Moody, Steven Horng, Leo Anthony Celi, and Roger Mark. 2024. MIMIC-IV. PhysioNet. Version 3.1. find
- Peter Kairouz and H. Brendan McMahan. 2021. Advances and open problems in federated learning. Foundations and Trends in Machine Learning, 14(1-2):1–210. find
- Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J. Reddi, Sebastian U. Stich, and Ananda Theertha Suresh. 2020. Scaffold: Stochastic controlled averaging for federated learning. In Proceedings of the 37th International Conference on Machine Learning. PMLR. find
- Alex Krizhevsky. 2009. Learning multiple layers of features from tiny images. Technical report, University of Toronto. find
- Ya Le and Xuan S. Yang. 2015. Tiny imagenet visual recognition challenge. find
- Daliang Li and Junpu Wang. 2019. Fedmd: Heterogenous federated learning via model distillation. In NeurIPS Workshop on Federated Learning for Data Privacy and Confidentiality. find
- Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. 2020. Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems, volume 2, pages 429–450. find
- Mang Ye, Xiuwen Fang, Bo Du, Pong C Yuen, and Dacheng Tao. 2023. Heterogeneous federated learning: State-of-the-art and research challenges. ACM Computing Surveys, 56(3):1–44. find
- Jianqing Zhang, Yang Liu, Yang Hua, and Jian Cao. 2024. Fedtgp: Trainable global prototypes with adaptive-margin-enhanced contrastive learning for data and model heterogeneity in federated learning. Preprint, arXiv:2401.03230. — [arXiv:2401.03230](https://arxiv.org/abs/2401.03230) find
- Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. 2018. Federated learning with non-iid data. find
- Xiang Li, Yue Zhang, Junpu Wang, Jian Chen, and 1 others. 2024. Federated learning via knowledge distillation: A survey. arXiv preprint, arXiv:2404.08564. — [arXiv:2404.08564](https://arxiv.org/abs/2404.08564) find
- Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. 2020. Ensemble distillation for robust model fusion in federated learning. In Advances in Neural Information Processing Systems. find
- H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas. 2017. Communication-efficient learning of deep networks from decentralized data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (AISTATS). PMLR. find
- Xutong Mu, Yulong Shen, Ke Cheng, Xueli Geng, Jiaxuan Fu, Tao Zhang, and Zhiwei Zhang. 2023. Fedproc: Prototypical contrastive federated learning on non-iid data. Future Generation Computer Systems, 143:93–104. find
- Felix Sattler, Tim Korjakow, Roman Rischke, and Wojciech Samek. 2021. Fedaux: Leveraging unlabeled auxiliary data in federated learning. IEEE Transactions on Neural Networks and Learning Systems. find
- Jake Snell, Kevin Swersky, and Richard S. Zemel. 2017. Prototypical networks for few-shot learning. In Advances in Neural Information Processing Systems. find
- Changlin Song, Divya Saxena, Jiannong Cao, and Yuqing Zhao. 2024. Feddistill: Global model distillation for local model de-biasing in non-iid federated learning. arXiv preprint arXiv:2404.09210. — [arXiv:2404.09210](https://arxiv.org/abs/2404.09210) find
- Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang. 2022a. Fedproto: Federated prototype learning across heterogeneous clients. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 8432–8440. find
- Yue Tan, Guodong Long, Jie Ma, Lu Liu, Tianyi Zhou, and Jing Jiang. 2022b. Federated learning from pre-trained models: A contrastive learning approach. arXiv preprint arXiv:2209.10083. — [arXiv:2209.10083](https://arxiv.org/abs/2209.10083) find
- Han Xiao, Kashif Rasul, and Roland Vollgraf. 2017. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. Preprint, arXiv:1708.07747. — [arXiv:1708.07747](https://arxiv.org/abs/1708.07747) find
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/prasad-2026-fedpagr-federated-prototype-alignment-geometric/