# COSMOS: Model-Agnostic Personalized Federated Learning with Clustered Server Models and Pseudo-Label-Only Communication

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 24 of its 39 references that are in the map.
- First posted on arXiv 11 May 2026.

## Builds on (most cited first, 24 of 24)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Ensemble Distillation for Robust Model Fusion in Federated Learning](https://mlsanity.com/paper/2006.07242/) (2020)
- [Ditto: Fair and Robust Federated Learning Through Personalization](https://mlsanity.com/paper/2012.04221/) (2020)
- [Personalized Federated Learning with Moreau Envelopes](https://mlsanity.com/paper/2006.08848/) (2020)
- [An Efficient Framework for Clustered Federated Learning](https://mlsanity.com/paper/2006.04088/) (2020)
- [Towards Personalized Federated Learning](https://mlsanity.com/paper/2103.00710/) (2021)
- [FedMD: Heterogenous Federated Learning via Model Distillation](https://mlsanity.com/paper/1910.03581/) (2019)
- [Bayesian Nonparametric Federated Learning of Neural Networks](https://mlsanity.com/paper/1905.12022/) (2019)
- [Data-Free Knowledge Distillation for Heterogeneous Federated Learning](https://mlsanity.com/paper/2105.10056/) (2021)
- [Adaptive Personalized Federated Learning](https://mlsanity.com/paper/2003.13461/) (2020)
- [FedProto: Federated Prototype Learning across Heterogeneous Clients](https://mlsanity.com/paper/2105.00243/) (2021)
- [Personalized Federated Learning using Hypernetworks](https://mlsanity.com/paper/2103.04628/) (2021)
- [FedBABU: Towards Enhanced Representation for Federated Image Classification](https://mlsanity.com/paper/2106.06042/) (2021)
- [Theoretical Analysis of Self-Training with Deep Networks on Unlabeled Data](https://mlsanity.com/paper/2010.03622/) (2020)
- [Theoretical Analysis of Weak-to-Strong Generalization](https://mlsanity.com/paper/2405.16043/) (2024)
- [FedGroup: Efficient Clustered Federated Learning via Decomposed Data-Driven Measure](https://mlsanity.com/paper/2010.06870/) (2020)
- [Communication Efficiency in Federated Learning: Achievements and Challenges](https://mlsanity.com/paper/2107.10996/) (2021)
- [Towards Model Agnostic Federated Learning Using Knowledge Distillation](https://mlsanity.com/paper/2110.15210/) (2021)
- [Knowledge Distillation for Federated Learning: a Practical Guide](https://mlsanity.com/paper/2211.04742/) (2022)
- [Architecture Agnostic Federated Learning for Neural Networks](https://mlsanity.com/paper/2202.07757/) (2022)
- [FedSelect: Personalized Federated Learning with Customized Selection of Parameters for Fine-Tuning](https://mlsanity.com/paper/2404.02478/) (2024)
- [Little is Enough: Boosting Privacy by Sharing Only Hard Labels in Federated Semi-Supervised Learning](https://mlsanity.com/paper/2310.05696/) (2023)
- [Learning Policy Committees for Effective Personalization in MDPs with Diverse Tasks](https://mlsanity.com/paper/2503.01885/) (2025)

## Abstract (arXiv)

Federated learning (FL) in heterogeneous environments remains challenging because client models often differ in both architecture and data distribution. While recent approaches attempt to address this challenge through client clustering and knowledge distillation, simultaneously handling architectural and statistical heterogeneity remains difficult. We introduce COSMOS, a model-agnostic framework that enables server-side personalization using only pseudo-label communication. Clients train local models and predict on the public data; the server clusters clients by prediction similarity, trains a cluster-specific model for each group using its own compute, and distills the resulting models back to clients. We provide the first theoretical analysis showing that distillation from the learned cluster models can yield exponential personalization risk contraction, going beyond the convergence-to-stationarity guarantees typically provided in model-agnostic FL. Experiments across benchmarks demonstrate that COSMOS consistently outperforms all model-agnostic FL baselines while remaining competitive with state-of-the-art personalized FL methods. More broadly, our results highlight personalized server-side learning with pseudo-labels as a promising paradigm for scalable and model-agnostic federated learning in highly heterogeneous environments.