Elamur Research⌕
☰
home / papers / 2026 / FedPAGR: Federated Prototype Alignment via Geometric Refinement for Heterogeneous Architectures

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.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 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).

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.acl-srw.37

@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

28 references, 1 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. 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
  2. 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
  3. 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
  4. 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
  5. 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
  6. 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
  7. Samuele Fonio, Mirko Polato, and Roberto Esposito. 2024. Fedhp: Federated learning with hyperspherical prototypical regularization. In European Symposium on Artificial Neural Networks. find
  8. A. E. W. Johnson, L. Bulgarelli, L. Shen, and et al. 2023. Mimic-iv, a freely accessible electronic health record dataset. Scientific Data, 10.
  9. 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
  10. 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
  11. 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
  12. Alex Krizhevsky. 2009. Learning multiple layers of features from tiny images. Technical report, University of Toronto. find
  13. Ya Le and Xuan S. Yang. 2015. Tiny imagenet visual recognition challenge. find
  14. 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
  15. 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
  16. 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
  17. 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
  18. Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. 2018. Federated learning with non-iid data. find
  19. 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
  20. 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
  21. 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
  22. 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
  23. 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
  24. Jake Snell, Kevin Swersky, and Richard S. Zemel. 2017. Prototypical networks for few-shot learning. In Advances in Neural Information Processing Systems. find
  25. 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
  26. 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
  27. 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
  28. 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

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/prasad-2026-fedpagr-federated-prototype-alignment-geometric/