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

arXiv:2605.11165 · cs.LG · Submitted 2026-05-11 · Read on arXiv

Listen

Radio episode about this paper

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

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

Washington University in St. Louis

cs.LG

Submitted: 2026-05-11

Updated: 2026-10-19

Code: https://github.com/YODA-Lab/COSMOS

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 87/100

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

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

Summary

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.

Introduction and Motivation

Federated learning (FL) in heterogeneous environments remains challenging because client models often differ in both architecture and data distribution <ref:2605.11165#pg2>. 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>. 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>. Prediction-based communication has emerged as one of the most practical mechanisms for enabling collaboration across heterogeneous models in FL, and by using pseudo-labels instead of parameters or gradients, communication efficiency can be improved significantly <ref:2605.11165#pg3>. Existing prior work like COMET has limitations, such as its reliance on heuristic K-means clustering and lack of formal guarantees regarding personalization performance <ref:2605.11165#pg4>.

COSMOS Framework Overview

COSMOS enables clients to train arbitrary local models on their private data and use them to generate pseudo-labels on a shared unlabeled dataset <ref:2605.11165#pg4>. The server then performs distance-controlled clustering to group clients with similar data distributions and trains a dedicated teacher model for each cluster <ref:2605.11165#pg5>. This process is divided into two phases:

The pre-training and clustering (PTC) phase

The iterative federated fine-tuning (IFFT) phase.

Pre-Training and Clustering Phase

The PTC phase involves four steps <ref:2605.11165#pg8>:

  1. Local Pre-Training: Each client i pre-trains its model fi locally on Di for E0 epochs using conventional (e.g., cross-entropy) loss, resulting in models f(1,i) <ref:2605.11165#pg8>. Each client then generates pseudo-labels over the global set U, f(1,i)(U), and sends these to the server <ref:2605.11165#pg8>.

  2. Client Clustering: The server clusters clients into K clusters Gk based on pseudo-label similarity using a greedy approach where the distance is defined as d(t)(i, j) = f(t,1)i(U)−f(t,1)j(U)1 <ref:2605.11165#pg8>. This clustering induces a cluster assignment function π which maps each client i to a cluster k <ref:2605.11165#pg8>.

  3. Server-Side Distillation: For each cluster Gk, the server trains a dedicated teacher model hk based on aggregate pseudo-labels ¯f(1)k(x) = 1/Gk P i∈Gk f(1,1)i(x), and performs gradient-based training to minimize an objective Jˆ(h(1)k;l) <ref:2605.11165#pg8>.

  4. Client-Side Distillation: Finally, each client i fine-tunes its local model f(1,2)i using the pseudo-labels h(1)π(i)(U), minimizing an objective Fˆ(f(1,2)i; l) <ref:2605.11165#pg8>.

Iterative Federated Fine-Tuning Phase

The IFFT phase alternates between three steps for a fixed number of iterations t <ref:2605.11165#pg9>:

Step 1: Local Fine-Tuning:

In iteration t ≥ 2, the local fine-tuning step starts from the model f(t-1,2)i and performs supervised training on its labeled local dataset Di to obtain f(t,1)i <ref:2605.11165#pg9>. The resulting pseudo-labels f(t,1)i(U) are then sent to the server to fine-tune the model h tπ(i) <ref:2605.11165#pg9>.

Step 2: Client-Side Distillation:

In iteration t, each server model k is fine-tuned on the pseudo-labels f(t,1)i(U) received from clients i ∈ Gk in a way that closely mirrors Step 3 of the PTC phase <ref:2605.11165#pg9>. The server then sends pseudo-labels h(t)k(U) to clients i ∈ Gk for all k <ref:2605.11165#pg9>.

Step 3: Client-Side Distillation:

Finally, each client i is fine-tuned from the pseudo-labels h(t)π(i)(U) using the objective from Equation 2 <ref:2605.11165#pg9>.

Personalization Guarantees

The analysis establishes an end-to-end exponential contraction of personalization risk bounds for COSMOS under sufficient conditions <ref:2605.11165#pg4>. The key components enabling this are:

Clustering Control:

The greedy clustering procedure enforces the property that the within-cluster pseudo-label distance is bounded by B, and for each client i there exists γ > 0 such that for every x ∈ Ui, ∆f(t,1)i(x) ≥ γ <ref:2605.11165#pg10>.

Data and Objective Conditions:

Condition 2 (Distributional coverage) requires that the public distribution Q does not under-cover nor over-concentrate on regions where client i has support <ref:2605.11165#pg10>. Condition 3 requires that F, Fˆ, J, and Jˆ are locally well-conditioned and the pseudolabels generated by any pairs of parameters are L1-Lipschitz continuous <ref:2605.11165#pg10>.

Effective Pseudo-Labelers:

Condition 5 requires that the probability b of pseudolabeling errors is not too large, specifically b ≤ 1/3 <ref:2605.11165#pg10>. Condition 6 ensures Label connectivity via (b, c)-expansion and robustness loss conditions <ref:2605.11165#pg10>.

Risk Contraction Results

Theorem 1 shows that for any client i at any iteration t, its server model h(t)π(i)’s risk is contracting: RDi(h(t)π(i)) ≤ κ1 RDi(f(t,1)i)+O˜n−1/2), where κ1 = 2w1w2c−1. Theorem 2 shows that for each client i and iteration t, the risk of the client model f(t,2)i is bounded by RDi f(t,2)i ≤ κ2 RDi f(t-1,2)i+O˜q log(1/δ)n <ref:2605.11165#pg10>. Applying a union bound yields a contraction of the personalization risk bound at an exponential rate <ref:2605.11165#pg10>.

Empirical Evaluation

Experiments demonstrate that COSMOS consistently outperforms all model-agnostic FL baselines while remaining competitive with state-of-the-art personalized FL methods <ref:2605.11165#pg4>. Specifically, Figure 2 presents the main results across four benchmarks in the model-agnostic setting with α = 5, where COSMOS outperforms all modelagnostic baselines by a large margin <ref:2605.11165#pg10>. Furthermore, Figure 3 evaluates COSMOS under different mixtures of client architectures and shows that it remains robust to heterogeneity in client capabilities <ref:2605.11165#pg10>. The per-client transmission cost for pseudo-labels is significantly lower than parameter sharing in this setting <ref:2605.11165#pg10>.

Conclusion

COSMOS introduces a framework that addresses a critical bottleneck in federated learning: achieving high-performance personalization while remaining fully model-agnostic <ref:2605.11165#pg8>. By allowing clients to utilize arbitrary (even proprietary) architectures through the communication of predictions on a shared unlabeled pool, COSMOS removes the structural barriers that have historically limited the deployment of PFL in diverse, real-world ecosystems <ref:2605.11165#pg8>.

Improvements for AI systems

  1. Model-Agnostic Personalization Framework COSMOS can enable clients using whatever model architecture best fits their needs to participate in federated learning without disclosing proprietary architectures or gradients.

  2. The system can achieve exponential contraction of personalization risk bounds, providing a general sufficient conditions on risk bound contraction for model-agnostic personalized FL, which ensures that the resulting client models are significantly better adapted to individual data distributions than standard FL methods.

  3. COSMOS reduces communication from parameter sharing by 1-2 orders of magnitude by relying on pseudo-label communication, making it viable for end-user connections that are slower, more expensive and less reliable.

  4. The server can actively train cluster specific models using clients’ pseudo-labels to create a collection H of K server models, where each server model hk is specialized to a subset of similar clients obtained through clustering.

  5. The system allows for the end-to-end theoretical analysis of personalization risk contraction, demonstrating that the personalization risk drops below the risk from local training alone for sufficiently large data and iterations.

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.

Sources

Related papers