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

Jul 16, 2026

Upload by Patrick

Federated learning in heterogeneous settings remains difficult when clients differ in both model architecture and data distribution. COSMOS introduces a model-agnostic framework enabling server-side personalization using pseudo-label-only communication. Clients train local models and generate predictions on shared public unlabeled data. The server clusters clients by prediction similarity, trains a dedicated cluster-specific teacher model, and distills it back. Theoretical results show exponential personalization risk contraction, and experiments demonstrate consistent gains over model-agnostic baselines.

LanguageEnglish

CategoryResearch & Report

FormatPDF

Pages17

Reading Time43 min

1 Introduction

Model-agnostic federated learning problem

Prior work limitations (COMET)

2 COSMOS overview

arXiv :2605 . 1 1 165v 3 [ cs .LG] 11 Jul 2026 COSMOS: Model-Agnostic Personalized Federated Learning with Clustered Server Models and Pseudo-Label-Only Communication Ben Rachmut 1 , Luise Ge 1 (􀀀), Ning Zhang1 , William Yeoh 1 , and Yevgeniy Vorobeychik 1 Washington University in [St. Louis](St. Louis g.luise@wustl.edu)[ g.luise@wustl.edu](St. Louis g.luise@wustl.edu) Abstract. 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-tostationarity 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. Keywords: Federated Learning · Personalized Federated Learning. 1 Introduction Federated learning (FL) is a distributed training paradigm where clients collaboratively train one or more server-side models without disclosing local data [12, 18,22] . Motivated by the heterogeneity of client data distributions, a host of personalized FL (PFL) methods has been developed to tailor models to individual clients [1,3,16,29] . However, most existing FL schemes still presume some structural knowledge or compatibility of client model architectures. This imposes a significant practical barrier, as clients may often wish to use whatever model architecture best fits their needs, or make use of proprietary architectures that they would not wish to disclose. Consequently, an important practical need in FL is to be model agnostic, allowing clients using whatever models they choose 2 B. Rachmut et al. to simply “plug in” to the FL scheme. The implication of model-agnostic FL is that it disallows any communication about model parameters or gradients. Despite the clear practical need, the problem of model-agnostic FL remains underexplored, particularly when clients simultaneously differ in both data distributions and model architectures. Existing model-agnostic approaches therefore rely on output-level communication, typically utilizing a shared unlabeled dataset [1,3,16] . Such datasets provide a common reference set on which heterogeneous models can exchange predictive signals without revealing parameters or private data. In many practical deployments, such datasets are readily available through public corpora (e.g., web-scraped images or text), synthetic generation, or institutionally shared benchmark pools. As a result, prediction-based communication has emerged as one of the most practical mechanisms for enabling collaboration across heterogeneous models in FL. Moreover, by using pseudo-labels instead of parameters or gradients, communication efficiency can be improved significantly, which has been widely recognized as a critical concern in FL as the wireless and other end-user connections are typically slower, more expensive and less reliable [26] . To the best of our knowledge, COMET [3] is the closest prior wor

FAQs

What problem does COSMOS address in federated learning?

COSMOS targets personalization in federated learning when clients are heterogeneous in both data distributions and model architectures, where prior model-agnostic methods struggle to handle both aspects jointly.

How does COSMOS enable communication without sharing model parameters or gradients?

COSMOS relies on pseudo-label communication: clients generate pseudo-labels using their local models on shared public unlabeled data, and the server uses only these predictive signals for clustering and distillation.

How does COSMOS personalize on the server side?

The server clusters clients by prediction similarity, trains a cluster-specific teacher model using server compute for each group, and distills the learned cluster models back to clients for personalization.

Jul 16, 2026

Upload by Patrick