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

summary

Video file (mp4)

The gist

The gist: COSMOS introduces a model-agnostic framework for personalized federated learning that enables server-side personalization using only pseudo-label communication, providing the first

In short

COSMOS is a model-agnostic framework for personalized federated learning that allows clients to use any local model without revealing its architecture or gradients. It achieves this by having clients generate pseudo-labels on a shared dataset, which the server uses to cluster clients and train specialized teacher models for each group. This method provides a theoretical guarantee of exponential risk contraction for personalization.

Key concepts

Model-Agnostic Framework
This means the system works regardless of what type of local machine learning model (like a CNN or RNN) the client uses. Clients can plug in their specific models without needing to know or share any structural details about them, which is crucial for privacy and flexibility in federated learning.
Pseudo-Label Communication
Instead of sending large model parameters or gradients, clients communicate by generating predictions (pseudo-labels) on a shared dataset. This communication is much smaller and more efficient than sharing full models, allowing for collaboration even when client models are completely different.
Clustered Server Models
The server groups clients based on the similarity of their generated pseudo-labels. For each group (cluster), the server trains a dedicated 'teacher' model specifically tailored to that cluster's data distribution. This specialization helps improve personalization for clients within that specific group.
Exponential Risk Contraction
This is a mathematical guarantee showing that the quality of the personalized client models improves at an exponential rate over time. In simple terms, this means the model gets significantly better at fitting individual client data compared to just training on the global average, provided enough data and iterations are used.

Terminology used across episodes

This episode discusses

The paper

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

Ben Rachmut, Luise Ge, Ning Zhang, William Yeoh, Yevgeniy Vorobeychik

Washington University in St. Louis

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.

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "COSMOS: Model-Agnostic Personalized Federated Learning with Clustered Server Models and Pseudo-Label-Only Communication".

Jane: Here's the gist of the paper. The gist: COSMOS introduces a model-agnostic framework for personalized federated learning that enables server-side personalization using only pseudo-label communication, providing the first theoretical analysis showing exponential risk contraction.

Tom: Federated learning (FL) in heterogeneous environments remains challenging because client models often differ in both architecture and data distribution <ref:2605.11165#pg2>.

Jane: Most existing FL schemes still presume some structural knowledge or compatibility of client model architectures, which imposes a significant practical barrier <ref:2605.11165#pg4>.

Tom: The need to be model-agnostic allows clients using whatever models they choose to simply “plug in” to the FL scheme <ref:2605.11165#pg5>.

Title and authors: Tom: So we’re looking at this paper called "COSMOS: Model-Agnostic Personalized Federated Learning with Clustered Server Models and Pseudo-Label-Only Communication". It sounds a bit dense, but basically, it’s about making federated learning work even when the clients all use totally different model architectures.

Jane: That’s the core problem they’re tackling. They are trying to solve that challenge where clients have both different data and different models simultaneously, which is really hard because most existing methods require you to know something about the client's architecture beforehand.

Lu: Exactly. The title highlights that it’s model-agnostic, meaning clients can just plug in whatever model they want without any restrictions on the server side. That removes a huge practical barrier for deployment one <ref:2605.11165#pg1>.

Meng: But what does that mean for me when I'm trying to deploy this at scale? If I can’t know the architecture, how do we even start setting up the initial communication protocol between clients and the server?

Lalam: It suggests using prediction-based communication instead of sending actual model parameters or gradients, which they call pseudo-labels two <ref:2605.11165#pg1>. This way, you don't have to worry about revealing proprietary model details.

Tom: Right. So we’re talking about a system where clients train locally on their private data and then use those models to generate these shared pseudo-labels that they send to the server instead of sending the models themselves.

Jane: And the authors are aiming for something pretty big here, which is getting a theoretical guarantee about how much better these personalized models will be compared to just using standard federated learning one <ref:2605.11165#pg1>.

Lu: That theoretical guarantee is what makes it interesting. They're providing the first formal analysis showing an exponential contraction of personalization risk bounds, going beyond the typical convergence-to-stationarity guarantees we usually see in federated learning one <ref:2605.11165#pg1>.

The paper's summary: Tom: So, to summarize what COSMOS is actually doing, it sets up a process where clients train their models locally and then generate pseudo-labels over a shared unlabeled dataset that they send back to the server.

Jane: And the server doesn’t just average everything; it clusters the clients based on how similar those generated pseudo-labels are to each other. This is key because it groups together clients whose data distributions or model behaviors look alike.

Lu: After clustering, the server trains a dedicated teacher model for each of those clusters using all the aggregated pseudo-labels from that group one <ref:2605.11165#pg1>. This means the server learns specialized models tailored to specific subgroups of clients.

Meng: So, it’s not just one global model anymore; it’s potentially a collection H of K server models, where each h k is specialized for a subset of similar clients based on how the clustering worked. How does that affect deployment complexity?

Tom: Well, the process is split into two phases. First, there's the pre-training and clustering phase where we get those initial cluster assignments and teacher models set up. Then there’s the iterative federated fine-tuning phase where things get refined over time one <ref:2605.11165#pg1>.

Jane: And that second phase involves clients fine-tuning their models using these new cluster-specific pseudo-labels, which is how they adapt to their specific local data distribution in a structured way.

Lalam: The communication efficiency part is also a big deal because they argue that using these pseudo-labels instead of sending full parameters or gradients can improve communication efficiency significantly two <ref:2605.11165#pg1>.

The paper's improvements: Tom: One of the main points they bring up is that COSMOS introduces an algorithmic framework where the server actively trains cluster-specific models using those client pseudo-labels, which is a big step away from just having one central model.

Jane: And the second major contribution is this theoretical guarantee we mentioned earlier: they establish an end-to-end exponential contraction of personalization risk bounds under certain conditions one <ref:2605.11165#pg1>. This means the personalization risk drops below what you’d expect from just local training alone if you give it enough data and iterations.

Lu: What’s compelling about that guarantee is that they derive general sufficient conditions for this exponential risk contraction without putting any restrictions on the model class itself, which is pretty powerful for a model-agnostic framework one <ref:2605.11165#pg1>.

Meng: But those guarantees rely on a lot of assumptions. They list several conditions they need to hold for the analysis to work, like things about how well the pseudo-labels are generated and how much error is allowed in those labels one <ref:2605.11165#pg1>.

Tom: They detail these conditions under "Data and Objective Conditions," which involve things like distributional coverage of the public data and ensuring that the loss functions are locally well-conditioned, plus that the pseudolabels they generate are L1-Lipschitz continuous one <ref:2605.11165#pg1>.

Jane: And for those assumptions, they specify conditions like having a probability b of pseudolabeling errors that is not too large, specifically b one/three one <ref:2605.11165#pg1>. That’s a concrete number you can use when designing your system.

Lalam: The paper also covers the risk contraction results explicitly. Theorem one shows that for any client i at iteration t, its server model h t pi i's risk is contracting, specifically R D i(h t pi(i)) kappa one R D i(f t,1i) + O(-one/two) one.

Conclusion: Tom: So to wrap up, COSMOS is this model-agnostic framework that uses pseudo-label communication and server-side clustering to create specialized models for different client groups, with the key result being an exponential contraction of personalization risk bounds one <ref:2605.11165#pg1>.

Jane: The implication for us is that we can deploy personalized federated learning across diverse ecosystems without needing to know anything about the clients' specific model architectures one <ref:2605.11165#pg1>. It’s a big step because it tackles structural barriers in the field.

Lu: If those sufficient conditions hold, the risk of getting a personalized model is exponentially better than just training locally, which is a strong theoretical result for this type of setup one <ref:2605.11165#pg1>.

Meng: From an engineering standpoint, we still have to worry about how much compute that server-side clustering and training takes. Can we actually run K cluster-specific models efficiently when K gets large?

Tom: That’s the next practical hurdle, definitely. But for now, this paper gives us a solid mathematical foundation on how to achieve personalization in model-agnostic FL one <ref:2605.11165#pg1>.

Jane: It’s a really interesting piece of work because it moves the conversation away from just heuristic methods and towards providing concrete guarantees about performance improvement one <ref:2605.11165#pg1>. That’s where we are headed next.

Lalam: The COSMOS paper shows that with the right setup, you can achieve personalization risk drops below what local training alone provides, which is a powerful result for data efficiency one <ref:2605.11165#pg1>.

More episodes

← Home