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

2026-05-11

AI Summary

This work addresses the challenge of personalization in federated learning under dual heterogeneity—both in client model architectures and data distributions—by proposing a model-agnostic framework. The approach clusters clients based on prediction similarity and trains a dedicated global model for each cluster at the server. Notably, it achieves cross-device knowledge distillation using only pseudo-label communication, eliminating the need to share model parameters or raw data. Theoretical analysis demonstrates that model distillation within clusters yields an exponential reduction in personalized risk. Empirical evaluations show that the proposed method significantly outperforms existing model-agnostic baselines across multiple benchmarks and attains performance comparable to state-of-the-art personalized federated learning algorithms.

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

Problem

• federated learning
• heterogeneous environments
• model-agnostic
• personalization
• architectural heterogeneity

Innovation

• model-agnostic
• personalized federated learning
• pseudo-label communication
• client clustering
• knowledge distillation