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

# COSMOS: Model-Agnostic Personalized Federated Learning with Clustered Server Models and Pseudo-Label-Only Communication
Source: [https://arxiv.org/html/2605.11165](https://arxiv.org/html/2605.11165)
Luise GeWilliam YeohNing ZhangYevgeniy Vorobeychik Washington University in St\. Louis

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

Keywords:Federated Learning Personalized Federated Learning\.

## 1Introduction

Federated learning \(FL\) is a distributed training paradigm where clients collaboratively train one or more server\-side models without disclosing local data\[[22](https://arxiv.org/html/2605.11165#bib.bib30),[18](https://arxiv.org/html/2605.11165#bib.bib15),[12](https://arxiv.org/html/2605.11165#bib.bib4)\]\. Motivated by the heterogeneity of client data distributions, a host of personalized FL \(PFL\) methods has been developed to tailor models to individual clients\[[16](https://arxiv.org/html/2605.11165#bib.bib26),[3](https://arxiv.org/html/2605.11165#bib.bib20),[1](https://arxiv.org/html/2605.11165#bib.bib19),[29](https://arxiv.org/html/2605.11165#bib.bib42)\]\. 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 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\[[16](https://arxiv.org/html/2605.11165#bib.bib26),[3](https://arxiv.org/html/2605.11165#bib.bib20),[1](https://arxiv.org/html/2605.11165#bib.bib19)\]\. 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 \(see our comparison in Table[4](https://arxiv.org/html/2605.11165#A3.T4)\), 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](https://arxiv.org/html/2605.11165#bib.bib64)\]\.

To the best of our knowledge, COMET\[[3](https://arxiv.org/html/2605.11165#bib.bib20)\]is the closest prior work that targets the crossover\. However, COMET has three major limitations\. First, the server acts solely as a passive coordinator, neglecting the significant computational resources potentially available at the server side that could be leveraged to facilitate personalization\. Second, its reliance on heuristic K\-means clustering necessitates prior knowledge of the cluster count K and lacks a formal mechanism for personalization when the underlying client diversity is high\. Consequently, COMET’s theoretical framework is restricted to standard convergence\-to\-stationarity for non\-convex objectives, providing no formal guarantees regarding personalization performance or risk reduction\.

We address all these limitations with COSMOS \(ClusteredOutput\-basedServerModels\)\. While a high\-capacity server has the potential to assist clients, it is highly non\-trivial whether the server can effectively learn from the noisy, heterogeneous pseudo\-labels provided by the clients in the first place\. COSMOS explicitly overcomes this bottleneck through careful algorithmic design\. Specifically, COSMOS enables clients to train arbitrary local models on their private data and use them to generate pseudo\-labels on a shared unlabeled dataset\. The server then performs distance\-controlled clustering to group clients with similar data distributions and trains a dedicated teacher model for each cluster\. Notably, COSMOS does not require the dataset to match client distributions exactly, only that it provides broad coverage of the input space\. As we show in the experiments \(Section[6](https://arxiv.org/html/2605.11165#S6)\), COSMOS remains effective even when the public dataset constitutes only a small fraction of the overall training data\.

Another critical gap in all existing model\-agnostic PFL literature is the absence of risk contraction guarantees\. To address this, we provide the first end\-to\-end analysis of personalization risk contraction\. To achieve this level of rigor, we leverage standard tools from semi\-supervised learning \(SSL\) theory\[[14](https://arxiv.org/html/2605.11165#bib.bib54),[31](https://arxiv.org/html/2605.11165#bib.bib53)\], including expansion\-based connectivity and bounded pseudo\-label error\. They allow us to derive general sufficient conditions for exponential risk contraction without restricting the model class\. Importantly, these assumptions are used only for analysis and do not impose constraints on the practical implementation of COSMOS\.

Our main contributions are:

1. 1\.Algorithmic Framework\.We present COSMOS, the first model\-agnostic PFL framework where the server actively trains cluster\-specific models using clients’ pseudo\-labels\.
2. 2\.Theoretical Guarantee\.We establish an end\-to\-end exponential contraction of personalization risk bounds for COSMOS under sufficient conditions, providing the first general risk contraction guarantee in model\-agnostic PFL\.
3. 3\.Empirical Evaluation\.We demonstrate that COSMOS not only consistently outperforms existing model\-agnostic FL methods but also maintains competitive performance in homogeneous settings, while reducing communication from parameter sharing by 1\-2 orders of magnitude\.

## 2Related Work

Classical Federated Learning\.Federated learning was introduced through FedAvg\[[22](https://arxiv.org/html/2605.11165#bib.bib30)\], which trains a single global model by aggregating client weight updates\. While simple and communication\-efficient, FedAvg suffers with non\-IID data, motivating methods such as FedProx\[[18](https://arxiv.org/html/2605.11165#bib.bib15)\]and SCAFFOLD\[[12](https://arxiv.org/html/2605.11165#bib.bib4)\]that stabilize optimization via proximal or control\-variance corrections\. Nonetheless these methods converge to a single global model and do not offer personalization\.

Model\-Agnostic, Model\-Heterogeneous, and Knowledge\-Distillation–Based Federated Learning\.Since classical federated learning communicates parameters or their updates, it requires every model on the clients and the server to share the same architecture\. To provide greater flexibility, a number of model\-heterogeneous approaches, such as communicating instance\-level representations as in FedHeNN\[[21](https://arxiv.org/html/2605.11165#bib.bib39)\], or abstract class prototypes as in FedProto\[[30](https://arxiv.org/html/2605.11165#bib.bib40)\], have been proposed to relax the architectural homogeneity assumption, as have many personalized FL methods \(see below\)\. However, while model\-agnostic approaches are necessarily model\-heterogeneous, most model\-heterogeneous methods are*not*model\-agnostic, since they still impose some architectural constraints\. Additionally, model\-agnostic FL is necessarily knowledge\-distillation–based FL \(KD\-FL\)\[[23](https://arxiv.org/html/2605.11165#bib.bib38)\]\. Nevertheless, many KD\-FL methods still rely on parameter aggregation at certain stages\[[19](https://arxiv.org/html/2605.11165#bib.bib9),[25](https://arxiv.org/html/2605.11165#bib.bib46),[34](https://arxiv.org/html/2605.11165#bib.bib44),[2](https://arxiv.org/html/2605.11165#bib.bib2)\]\. To our knowledge, only FedMD\[[16](https://arxiv.org/html/2605.11165#bib.bib26)\]and COMET\[[3](https://arxiv.org/html/2605.11165#bib.bib20)\]explore purely model\-agnostic FL using soft labels, while FedCT\[[1](https://arxiv.org/html/2605.11165#bib.bib19)\]relies on hard labels\. Recently, the communication efficiency of pseudo\-labels has also been leveraged in federated multi\-view clustering \(e\.g\., CeFMC\[[20](https://arxiv.org/html/2605.11165#bib.bib63)\]\) although its objective is orthogonal to ours\.

Personalized Federated Learning \(PFL\)\.Almost all practical federated learning settings exhibit statistical heterogeneity, where different clients’ local distributions can vary substantially\. Personalized federated learning \(PFL\) addresses this by learning client\-adapted models\[[29](https://arxiv.org/html/2605.11165#bib.bib42)\]\. Existing PFL approaches can be broadly grouped by whether they maintain a single shared server model or a small number of server\-side models\. In the first group, a single global model is adapted to each client via meta\-learning \(Per\-FedAvg\[[8](https://arxiv.org/html/2605.11165#bib.bib3)\]\), regularization \(pFedMe\[[6](https://arxiv.org/html/2605.11165#bib.bib61)\], Ditto\[[17](https://arxiv.org/html/2605.11165#bib.bib59)\]\), adaptive mixing of local and global models \(APFL\[[5](https://arxiv.org/html/2605.11165#bib.bib28)\]\), representation refinement \(FedBABU\[[24](https://arxiv.org/html/2605.11165#bib.bib62)\]\), or hypernetwork\-based parameter generation \(pFedHN\[[27](https://arxiv.org/html/2605.11165#bib.bib60)\], FedSelect\[[28](https://arxiv.org/html/2605.11165#bib.bib65)\]\)\. In the second group, clustered PFL methods explicitly maintain multiple server\-side models and assign clients to them, as in IFCA\[[10](https://arxiv.org/html/2605.11165#bib.bib10)\], FedGroup\[[7](https://arxiv.org/html/2605.11165#bib.bib8)\], AutoCFL\[[11](https://arxiv.org/html/2605.11165#bib.bib5)\], and pFedCK\[[33](https://arxiv.org/html/2605.11165#bib.bib27)\]\. To the best of our knowledge, only COMET is also a model\-agnostic PFL approach\[[3](https://arxiv.org/html/2605.11165#bib.bib20)\]\. However, it requires one to specify the number of clustersKKin advance, and applies standardKK\-means clustering, which is heuristic and not easily amenable to theoretical personalization guarantees\.

Theory for Model\-Agnostic FL\.While theoretical convergence results abound for conventional FL schemes, model\-agnostic settings as well as personalization make such results significantly more challenging\. The earliest model\-agnostic FL approach, FedMD\[[16](https://arxiv.org/html/2605.11165#bib.bib26)\]does not provide any theoretical guarantees\. A recent FedCT method\[[1](https://arxiv.org/html/2605.11165#bib.bib19)\]requires an oversimplified assumption directly that the training algorithms*always*yield monotone increasing accuracy to achieve convergence\. The theoretical analysis for COMET\[[3](https://arxiv.org/html/2605.11165#bib.bib20)\], on the other hand, requires linear models and a Gaussian data distribution to obtain generalization results\. Thus, there are no general sufficient conditions on risk bound contraction for model\-agnostic personalized FL\. We bridge this gap by adopting the analysis tools from the semi\-supervised learning literature\[[31](https://arxiv.org/html/2605.11165#bib.bib53)\]\.

## 3Model

We consider a federated learning scenario forMM\-class classification withNNclients and a server\. Each clienti∈\[N\]i\\in\[N\]has a private labeled datasetDi=\{\(xi​j,y​\(xi​j\)\)\}jD\_\{i\}=\\\{\(x\_\{ij\},y\(x\_\{ij\}\)\)\\\}\_\{j\}withxi​jx\_\{ij\}drawn i\.i\.d\. from its local distribution𝒟i\\mathcal\{D\}\_\{i\}over the input space𝒳\\mathcal\{X\}, andy​\(xi​j\)y\(x\_\{ij\}\)the true label ofxi​jx\_\{ij\}\. We assume that each𝒟i\\mathcal\{D\}\_\{i\}admits a density functionpi​\(x\)p\_\{i\}\(x\)\. A clientiitrains a modelfi:𝒳→\[0,1\]Mf\_\{i\}:\\mathcal\{X\}\\rightarrow\[0,1\]^\{M\}with‖fi​\(x\)‖1=1\\\|f\_\{i\}\(x\)\\\|\_\{1\}=1from a hypothesis classℋi\\mathcal\{H\}\_\{i\}that can be distinct for eachii, representing, for example, distinct neural network architectures for different clients\. A server, in turn, has a hypothesis classℋS\\mathcal\{H\}\_\{S\}and can train a*collection*of modelsH=\{h1,…,hK\}⊂ℋSH=\\\{h\_\{1\},\\ldots,h\_\{K\}\\\}\\subset\\mathcal\{H\}\_\{S\}, wherehk:𝒳→\[0,1\]Mh\_\{k\}:\\mathcal\{X\}\\rightarrow\[0,1\]^\{M\}and‖hk​\(x\)‖1=1\\\|h\_\{k\}\(x\)\\\|\_\{1\}=1for eachkk\. In our setting, the value ofKKis obtained*endogenously as part of the training procedure*\. Furthermore, letπ:\[N\]→\[K\]\\pi:\[N\]\\rightarrow\[K\]be a mapping \(also obtained during training\) which assigns each clientiito a corresponding server modelhπ​\(i\)h\_\{\\pi\(i\)\}\. Finally, following prior work on PFL, we assume availability of an*unlabeled public dataset*U=\{xj\}U=\\\{x\_\{j\}\\\}of sizen=\|U\|n=\|U\|, with eachxjx\_\{j\}drawn i\.i\.d\. from a global distribution𝒬\\mathcal\{Q\}over𝒳\\mathcal\{X\}\. We assume that𝒬\\mathcal\{Q\}admits a density functionq​\(x\)q\(x\)\. It is not difficult to obtain unlabeled datasets of this kind, for example, scraping \(open\-license\) images or text from the internet, or generating it synthetically\.

To formalize the learning objective, we introduce some additional notation\. Letggbe a classification model with outputs a distribution overMMclasses \(i\.e\., pseudo\-labels\)\. We useA∘g\(x\)=argmaxm∈\[M\]\[g\(x\)\]mA\\circ g\(x\)=\\arg\\max\_\{m\\in\[M\]\}\[g\(x\)\]\_\{m\}to denote the predicted class \(i\.e\., the class with the highest probability undergg\)\. Further, letR𝒟​\(g\)=𝔼x∼𝒟​\[ℓ0−1​\{A∘g​\(x\),y​\(x\)\}\]R\_\{\\mathcal\{D\}\}\(g\)=\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\[\\ell\_\{0\-1\}\\\{A\\circ g\(x\),y\(x\)\\\}\]denote the risk \(probability of a mistake\) ofggunder distribution𝒟\\mathcal\{D\}, whereℓ0−1\\ell\_\{0\-1\}is the0−10\-1loss, andErrx∈D​\(g\)=1\|D\|​∑xℓ0−1​\{A∘g​\(x\),y​\(x\)\}\\mathrm\{Err\}\_\{x\\in D\}\(g\)=\\frac\{1\}\{\|D\|\}\\sum\_\{x\}\\ell\_\{0\-1\}\\\{A\\circ g\(x\),y\(x\)\\\}denote the hard label prediction error over a datasetDD\.

Learning Objective:Our goal is to train a collection\{hk\}\\\{h\_\{k\}\\\}ofKKpersonalized server models, along with a client\-to\-model mappingπ\\pi, which minimizes the*personalization risk*, defined as the average risk over theNNclients:R†​\(h1,…,hK\)=1N​∑i=1NR𝒟i​\(hπ​\(i\)\)\.R^\{\\dagger\}\(h\_\{1\},\\dots,h\_\{K\}\)=\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}R\_\{\\mathcal\{D\}\_\{i\}\}\(h\_\{\\pi\(i\)\}\)\.A full notation summary appears in Table[1](https://arxiv.org/html/2605.11165#A1.T1)in Appendix[A](https://arxiv.org/html/2605.11165#A1)\.

## 4Algorithmic Approach

![Refer to caption](https://arxiv.org/html/2605.11165v1/figures_ijcai/workflow.png)Figure 1:Overview of the four steps in the COSMOS workflow\.At the high level, the proposed COSMOS framework involves iteratively fine\-tuning the client and server models using pseudo\-labels collected from one another\. As shown in Figure[1](https://arxiv.org/html/2605.11165#S4.F1), clients generate pseudo\-labels for global data using locally trained models, which the server clusters and aggregates to train cluster\-specific models\. The resulting pseudo\-labels are iteratively returned to clients for local model refinement\. We enable personalization by creating and training a small collectionHHofKKserver models, with each server modelhkh\_\{k\}specialized to a subset of similar clients obtained through clustering\. More precisely, COSMOS proceeds in two phases: 1\) the*pre\-training and clustering \(PTC\)*phase and 2\) the*iterative federated fine\-tuning \(IFFT\)*phase\.

### 4\.1The Pre\-Training and Clustering Phase

The pre\-training and clustering phase involves four steps: 1\) local pre\-training, 2\) clustering, 3\) server\-side distillation, and 4\) client\-side distillation\.

Step 1: Local Pre\-Training\.We pre\-train each clientii’s modelfif\_\{i\}locally onDiD\_\{i\}forE0E\_\{0\}epochs using conventional \(e\.g\., cross\-entropy\) loss\. We refer to the resulting client models asfi\(1,1\)f\_\{i\}^\{\(1,1\)\}, where the first superscript refers to the time step \(t=1t=1\) and the second means that the client models are trained \(or fine\-tuned, in the IFFT phase\) on*local*data \(as opposed to global dataUUwith server\-provided pseudo\-labels, as in Step 4 below\)\. Each clientiithen generates pseudo\-labels over the global setUU,fi\(1,1\)​\(U\)f\_\{i\}^\{\(1,1\)\}\(U\), and sends these to the server\.

Step 2: Client Clustering\.Upon receiving pseudo\-labelsfi\(1,1\)​\(U\)f\_\{i\}^\{\(1,1\)\}\(U\)from all clientsii, the server proceeds to cluster the clients intoKKclusters\{𝒢k\}k=1K\\\{\\mathcal\{G\}\_\{k\}\\\}\_\{k=1\}^\{K\}based on pseudo\-label similarity\. Our goal is to find the minimum number of clustersKKwith respect to a given distance function\. However, this problem is NP\-hard and the best known approximation has computational complexity exponential in the data dimension\[[9](https://arxiv.org/html/2605.11165#bib.bib52)\]\. On the other hand, while numerous heuristic clustering approaches have been proposed, they cannot be easily used to provide*risk convergence*guarantees\.

To address this challenge, we propose the following greedy approach, which*is*amenable to theoretical analysis \(see Section[5\.1](https://arxiv.org/html/2605.11165#S5.SS1)\)\. We define the distance between two clientsi,ji,jas their pseudolabels’ℓ1\\ell\_\{1\}distance, i\.e\.d\(t\)​\(i,j\)=\|fi\(t,1\)​\(U\)−fj\(t,1\)​\(U\)\|1d^\{\(t\)\}\(i,j\)=\|f\_\{i\}^\{\(t,1\)\}\(U\)\-f\_\{j\}^\{\(t,1\)\}\(U\)\|\_\{1\}\. Define the setNi​\(B0,C\)N\_\{i\}\(B\_\{0\},C\)of*B0B\_\{0\}\-close neighbors*of a clientiiamong clients inCCasNi​\(B0,C\)=\{j\|j∈C∖\{i\},d\(1\)​\(i,j\)≤B0\}N\_\{i\}\(B\_\{0\},C\)=\\\{j\|j\\in C\\setminus\\\{i\\\},d^\{\(1\)\}\(i,j\)\\leq B\_\{0\}\\\}\. At the high level, our algorithm greedily selects a client with the mostB0B\_\{0\}\-close neighbors among those not previously selected, and defines a new cluster associated with this client and its neighbors, proceeding until all clients belong to some cluster\. This is made more precise in Algorithm[1](https://arxiv.org/html/2605.11165#algorithm1)\. The resulting clustering scheme induces a cluster assignment functionπ\\piwhich maps each clientiito a clusterkk\. Notably,*we need not know the number of clusters upfront*, as this is endogenous to our clustering approach\. WhileB0B\_\{0\}can be treated as another hyperparameter for clustering, this hyperparameter is theoretically grounded: it is directly connected to the personalization guarantees of our framework\.

Input:Clients

1,…,N1,\\dots,Nwith associated hard label vectors

\{yi\}\\\{y\_\{i\}\\\}; clustering parameter

B0B\_\{0\}\.

Output:The number of clusters

KKand the associated collection of client clusters

G=\{𝒢k\}k=1KG=\\\{\\mathcal\{G\}\_\{k\}\\\}\_\{k=1\}^\{K\}\.

Initialization:

C=\[1,…,N\]C=\[1,\\dots,N\],

G=\{∅\}G=\\\{\\emptyset\\\},

K=1K=1
while*C≠∅C\\neq\\emptyset*do

i∗∈arg⁡maxi∈Ci^\{\*\}\\in\\arg\\max\_\{i\\in C\}\|Ni​\(B0,C\)\|\|N\_\{i\}\(B\_\{0\},C\)\|
𝒢K=\{i∗∪Ni∗​\(B0,C\)\}\\mathcal\{G\}\_\{K\}=\\\{i^\{\*\}\\cup N\_\{i^\{\*\}\}\(B\_\{0\},C\)\\\}
G=G∪𝒢KG=G\\cup\\mathcal\{G\}\_\{K\};

C=C∖𝒢KC=C\\setminus\\mathcal\{G\}\_\{K\};

K=K\+1K=K\+1

Algorithm 1Greedy Clustering\.Step 3: Server\-Side Distillation\.Our next step is to train the server\-side model in each clusterkkbased on the pseudolabels received from all clients\. Specifically, fix a cluster𝒢k\\mathcal\{G\}\_\{k\}and a datapointx∈Ux\\in U\. The server update for each cluster modelhk\(1\)h\_\{k\}^\{\(1\)\}\(where the superscript references iteration\) andx∈Ux\\in Uthen uses the aggregate pseudo\-labels over all clients in this cluster which we define asf¯k\(1\)​\(x\)=1\|𝒢k\|​∑i∈𝒢kfi\(1,1\)​\(x\)\\bar\{f\}^\{\(1\)\}\_\{k\}\(x\)=\\frac\{1\}\{\|\\mathcal\{G\}\_\{k\}\|\}\\sum\_\{i\\in\\mathcal\{G\}\_\{k\}\}f\_\{i\}^\{\(1,1\)\}\(x\)for clientsi∈𝒢ki\\in\\mathcal\{G\}\_\{k\}, i\.e\., the average of the pseudo\-labels of the clients in thekk\-th cluster\. Finally, the server performs gradient\-based training for each cluster\-specific modelkkto minimize the following objective:

J^​\(hk\(1\);ℓ\)=∑x∈U\[ℓ​\(hk​\(x\),f¯k\(1\)​\(x\)\)\+λ​rℬ​\(hk\(1\);ℓ\)​\(x\)\],\\displaystyle\\\!\\\!\\\!\\\!\\\!\\\!\\hat\{J\}\(h\_\{k\}^\{\(1\)\};\\ell\)=\\sum\_\{x\\in U\}\\Big\[\\ell\(h\_\{k\}\(x\),\\bar\{f\}\_\{k\}^\{\(1\)\}\(x\)\)\+\\lambda\\,r\_\{\\mathcal\{B\}\}\(h\_\{k\}^\{\(1\)\};\\ell\)\(x\)\\Big\],\(1\)whereℓ​\(⋅,⋅\)\\ell\(\\cdot,\\cdot\)is a loss function andrℬ​\(⋅,ℓ\)r\_\{\\mathcal\{B\}\}\(\\cdot,\\ell\)an input\-specific regularization term\. The role of the regularization term in our setting is to train the models which are locally smooth \(related to the consistency and robustness properties central to our theoretical analysis in Section[5](https://arxiv.org/html/2605.11165#S5)\)\. Specifically, let𝒯\\mathcal\{T\}denote a set of permissible data augmentations \(e\.g\., image translations and rotations\) and fixddto be a small radius capturing minor perturbations around augmented data points \(in particular, we assume thatddis significantly smaller than the typical norm ofxx\)\. For an inputx∈𝒳x\\in\\mathcal\{X\}, define itstransformation ballas:ℬ​\(x\)=\{x′∈𝒳:∃T∈𝒯​s\.t\.​‖x′−T​\(x\)‖≤d\}\\mathcal\{B\}\(x\)=\\\{x^\{\\prime\}\\in\\mathcal\{X\}:\\exists T\\in\\mathcal\{T\}\\text\{ s\.t\. \}\\\|x^\{\\prime\}\-T\(x\)\\\|\\leq d\\\}\. Next, let the neighborhood𝒩​\(x\)\\mathcal\{\\mathcal\{N\}\}\(x\)of a pointxxbe the set of inputs whose transformation balls intersect withℬ​\(x\)\\mathcal\{B\}\(x\):𝒩​\(x\)=\{x′:ℬ​\(x\)∩ℬ​\(x′\)≠∅\}\\mathcal\{N\}\(x\)=\\\{x^\{\\prime\}:\\mathcal\{B\}\(x\)\\cap\\mathcal\{B\}\(x^\{\\prime\}\)\\neq\\emptyset\\\}\. LetS⊂𝒩​\(x\)S\\subset\\mathcal\{N\}\(x\)be a finite sample of the neighborhood ofxx\(e\.g\., obtained by rejection sampling\)\. We now define a regularization term for any modelggasrℬ​\(g;ℓ\)​\(x\):=∑x′∈S\{ℓ​\(g​\(x′\),g​\(x\)\)\}r\_\{\\mathcal\{B\}\}\(g;\\ell\)\(x\):=\\sum\_\{x^\{\\prime\}\\in S\}\\bigl\\\{\\ell\(g\(x^\{\\prime\}\),g\(x\)\)\\bigr\\\}\. We use the same loss function for the supervised term and regularized term, but allowing different loss functions is also admissible\. We motivate this form of regularization more precisely in Section[5](https://arxiv.org/html/2605.11165#S5)\.

After gradient\-based training for each server\-side modelhk\(1\)h\_\{k\}^\{\(1\)\}using the objective \([1](https://arxiv.org/html/2605.11165#S4.E1)\), we send the resulting pseudo\-labelshk\(1\)​\(U\)h\_\{k\}^\{\(1\)\}\(U\)to clients in cluster𝒢k\\mathcal\{G\}\_\{k\}\.

Step 4: Client\-Side Distillation\.Finally, each clientiifine\-tunes its local modelfi\(1,2\)f\_\{i\}^\{\(1,2\)\}, where the second superscript refers to the second round of local fine\-tuning based on the pseudo\-labelshπ​\(i\)\(1\)​\(U\)h\_\{\\pi\(i\)\}^\{\(1\)\}\(U\)\. We use the following objective:

F^​\(fi\(1,2\);ℓ\)=∑x∈Uwi​\(x\)​\[ℓ​\(fi\(1,2\)​\(x\),hπ​\(i\)\(1\)​\(x\)\)\+λ​rℬ​\(fi\(1,2\);ℓ\)​\(x\)\],\\hat\{F\}\(f\_\{i\}^\{\(1,2\)\};\\ell\)=\\sum\_\{x\\in U\}w\_\{i\}\(x\)\\Big\[\\ell\(f\_\{i\}^\{\(1,2\)\}\(x\),h\_\{\\pi\(i\)\}^\{\(1\)\}\(x\)\)\\\\ \+\\lambda\\,r\_\{\\mathcal\{B\}\}\(f\_\{i\}^\{\(1,2\)\};\\ell\)\(x\)\\Big\],\(2\)
wherewi​\(x\)w\_\{i\}\(x\)are the client\-specific weights of data inUUthat allow for importance sampling techniques \(see Section[5](https://arxiv.org/html/2605.11165#S5)\)\. In practice, we find that settingwi​\(x\)=1w\_\{i\}\(x\)=1for alliiandxxis effective in practice, and avoids the need to make assumptions about the local and global distribution differences\.

### 4\.2The Iterative Federated Fine\-Tuning Phase

After the PTC phase, we enter an iterative FL phase in which we alternate 1\) local fine\-tuning, 2\) server\-side distillation, and 3\) client\-side distillation over a fixed number of iterations\. We now describe each step for a fixed iterationtt\.

Step 1: Local Fine\-Tuning\.The PTC phase effectively serves as iterationt=1t=1\. In any IFFT iterationt≥2t\\geq 2, the local fine\-tuning step for clientiistarts from the modelfi\(t−1,2\)f\_\{i\}^\{\(t\-1,2\)\}\(obtained by fine\-tuning on the public dataUUannotated with the server’s pseudo\-labels, as described in Step 3 below\), and performs supervised training on its labeled local datasetDi=\{\(x,y\)\}D\_\{i\}=\\\{\(x,y\)\\\}to obtainfi\(t,1\)f\_\{i\}^\{\(t,1\)\}\. Concretely, as before, clientiiapproximately minimizes the empirical risk with a calibrated surrogate lossℓ\\ellsuch as cross\-entropy lossR^ℓ,Di\(f\)=1\|Di\|∑\(x,y\)∈Diℓ\(f\(x\),y\)\)\.\\hat\{R\}\_\{\\ell,D\_\{i\}\}\(f\)\\;=\\;\\frac\{1\}\{\|D\_\{i\}\|\}\\sum\_\{\(x,y\)\\in D\_\{i\}\}\\ell\\\!\\big\(f\(x\),y\\big\)\)\.We require this supervised update to be*classification\-safe*in the sense that, with high probability over the local sample/optimization noise, it does not increase the true0–11risk on𝒟i\\mathcal\{D\}\_\{i\}beyond a small generalization slack for analysis purpose\. After fine\-tuning, the pseudo\-labelsfi\(t,1\)​\(U\)f\_\{i\}^\{\(t,1\)\}\(U\)are then sent to the server to fine\-tune the modelhπ​\(i\)th\_\{\\pi\(i\)\}^\{t\}\.

Step 2: Client\-Side Distillation\.In iterationtt, each server modelk∈\[1,…,K\]k\\in\[1,\\ldots,K\]is fine\-tuned on the pseudo\-labelsfi\(t,1\)​\(U\)f\_\{i\}^\{\(t,1\)\}\(U\)received from clientsi∈𝒢ki\\in\\mathcal\{G\}\_\{k\}in a way that closely mirrors Step 3 of the PTC phase\. That is, we definef¯k\(t\)​\(x\)=1\|𝒢k\|​∑i∈𝒢kfi\(t,1\)​\(x\)\\bar\{f\}^\{\(t\)\}\_\{k\}\(x\)=\\frac\{1\}\{\|\\mathcal\{G\}\_\{k\}\|\}\\sum\_\{i\\in\\mathcal\{G\}\_\{k\}\}f\_\{i\}^\{\(t,1\)\}\(x\)fori∈𝒢ki\\in\\mathcal\{G\}\_\{k\}\. The server modelhk\(t\)h\_\{k\}^\{\(t\)\}is then fine\-tuned using the objective \([1](https://arxiv.org/html/2605.11165#S4.E1)\), withhk\(1\)h\_\{k\}^\{\(1\)\}replaced byhk\(t\)h\_\{k\}^\{\(t\)\}andf¯k\(1\)​\(x\)\\bar\{f\}^\{\(1\)\}\_\{k\}\(x\)replaced byf¯k\(t\)​\(x\)\\bar\{f\}^\{\(t\)\}\_\{k\}\(x\)\. The server then sends pseudo\-labelshk\(t\)​\(U\)h\_\{k\}^\{\(t\)\}\(U\)to clientsi∈𝒢ki\\in\\mathcal\{G\}\_\{k\}for allkk\.

Step 3: Client\-Side Distillation\.Finally, we fine\-tune each clientiifrom the pseudo\-labelshπ​\(i\)\(t\)​\(U\)h\_\{\\pi\(i\)\}^\{\(t\)\}\(U\)using the objective from Equation[2](https://arxiv.org/html/2605.11165#S4.E2)\(as in Step 4 of PTC\), in which we replacefi\(1,2\)f\_\{i\}^\{\(1,2\)\}withfi\(t,2\)f\_\{i\}^\{\(t,2\)\}andhπ​\(i\)\(1\)h\_\{\\pi\(i\)\}^\{\(1\)\}withhπ​\(i\)\(t\)h\_\{\\pi\(i\)\}^\{\(t\)\}\.

## 5Personalization Guarantees

In this section, we analyze the population risk of COSMOS with respect to the client distributions\{𝒟i\}i=1N\\\{\\mathcal\{D\}\_\{i\}\\\}\_\{i=1\}^\{N\}\. We provide such guarantees both for the individual client models, as well as for the overall*personalization risk*on the server\. This analysis entails several technical challenges\. First, as clustering is NP\-hard, heuristic approaches are typically used, and it is therefore not evident how to achieve convergence guarantees for personalization risk for algorithms that leverage clustering \(as COSMOS does\)\. Second, it is clear that arbitrary distributions𝒟i\\mathcal\{D\}\_\{i\}\(of the client data\) and𝒬\\mathcal\{Q\}\(of global data\) cannot achieve convergence\. For example, if𝒬\\mathcal\{Q\}puts all mass on a single datapoint, the clustering step will fail\. Prior approaches deal with this issue by making strong distributional assumptions, such as assuming a Gaussian data distribution\[[3](https://arxiv.org/html/2605.11165#bib.bib20)\]\. We aim to obtain*general*sufficient conditions on the data distributions\. Third, we wish to make minimal assumptions on the nature of client and server models, unlike prior approaches that require linearity for generalization bounds\[[3](https://arxiv.org/html/2605.11165#bib.bib20),[27](https://arxiv.org/html/2605.11165#bib.bib60)\]\.

Notation:Before diving into the technical details, we define some useful notation\. For a clientii, let𝒮i=supp​\(𝒟i\)⊆𝒳\\mathcal\{S\}\_\{i\}=\\mathrm\{supp\}\(\\mathcal\{D\}\_\{i\}\)\\subseteq\\mathcal\{X\}denote the support ofii’s local data distribution and𝒮ic=𝒳∖𝒮i\\mathcal\{S\}\_\{i\}^\{c\}=\\mathcal\{X\}\\setminus\\mathcal\{S\}\_\{i\}be its complement\. We writeUi=U∩𝒮iU\_\{i\}=U\\cap\\mathcal\{S\}\_\{i\}for the client\-supported portion of public data andUic=U∖UiU\_\{i\}^\{c\}=U\\setminus U\_\{i\}for its complement\. Let𝒬​\(𝒮i\)\\mathcal\{Q\}\(\\mathcal\{S\}\_\{i\}\)denote the probability of𝒮i\\mathcal\{S\}\_\{i\}under𝒬\\mathcal\{Q\}\. Let𝒬i:=𝒬\(⋅∣𝒮i\)\\mathcal\{Q\}\_\{i\}:=\\mathcal\{Q\}\(\\cdot\\mid\\mathcal\{S\}\_\{i\}\)denote the restriction of𝒬\\mathcal\{Q\}to𝒮i\\mathcal\{S\}\_\{i\}, with densityqi​\(x\)=q​\(x\)/𝒬​\(𝒮i\)q\_\{i\}\(x\)=q\(x\)/\\mathcal\{Q\}\(\\mathcal\{S\}\_\{i\}\)\. Let𝒬ic:=𝒬\(⋅∣𝒮ic\)\\mathcal\{Q\}\_\{i\}^\{c\}:=\\mathcal\{Q\}\(\\cdot\\mid\\mathcal\{S\}\_\{i\}^\{c\}\)\. For any functionsggandg′g^\{\\prime\}, and inputxx, defineG​\(g,g′,ℓ\)=ℓ​\(g,g′\)\+λ​rℬ​\(g;ℓ\)\.G\(g,g^\{\\prime\},\\ell\)=\\ell\(g,g^\{\\prime\}\)\+\\lambda r\_\{\\mathcal\{B\}\}\(g;\\ell\)\.Then, we can write the server\-side clusterkk’s objectiveJ^​\(hk\)=∑x∈UG​\(hk,f¯k,ℓ\)​\(x\)\\hat\{J\}\(h\_\{k\}\)=\\sum\_\{x\\in U\}G\(h\_\{k\},\\bar\{f\}\_\{k\},\\ell\)\(x\), wherehkh\_\{k\}is the server\-side model andf¯k\\bar\{f\}\_\{k\}the aggregate model over clients in thekkth cluster\. Similarly, we can writeF^​\(fi;ℓ\)=∑x∈Uwi​\(x\)​G​\(fi,hπ​\(i\),ℓ\)​\(x\)\\hat\{F\}\(f\_\{i\};\\ell\)=\\sum\_\{x\\in U\}w\_\{i\}\(x\)G\(f\_\{i\},h\_\{\\pi\(i\)\},\\ell\)\(x\)for a clientii’s modelfif\_\{i\}\. It will also be useful to defineJ^i​\(hk;ℓ\)=∑x∈UiG​\(hk,f¯k,ℓ\)​\(x\)\\hat\{J\}\_\{i\}\(h\_\{k\};\\ell\)=\\sum\_\{x\\in U\_\{i\}\}G\(h\_\{k\},\\bar\{f\}\_\{k\},\\ell\)\(x\)andJ^ic​\(hk;ℓ\)=∑x∈UicG​\(hk,f¯k,ℓ\)​\(x\)\\hat\{J\}\_\{i\}^\{c\}\(h\_\{k\};\\ell\)=\\sum\_\{x\\in U\_\{i\}^\{c\}\}G\(h\_\{k\},\\bar\{f\}\_\{k\},\\ell\)\(x\)for the server, and similarly,F^i​\(fi;ℓ\)=∑x∈Uiwi​\(x\)​G​\(fi,hπ​\(i\),ℓ\)​\(x\)\\hat\{F\}\_\{i\}\(f\_\{i\};\\ell\)=\\sum\_\{x\\in U\_\{i\}\}w\_\{i\}\(x\)G\(f\_\{i\},h\_\{\\pi\(i\)\},\\ell\)\(x\)andF^ic​\(fi;ℓ\)=∑x∈Uicwi​\(x\)​G​\(fi,hπ​\(i\),ℓ\)​\(x\)\\hat\{F\}\_\{i\}^\{c\}\(f\_\{i\};\\ell\)=\\sum\_\{x\\in U\_\{i\}^\{c\}\}w\_\{i\}\(x\)G\(f\_\{i\},h\_\{\\pi\(i\)\},\\ell\)\(x\)for each clientii\.J^i\\hat\{J\}\_\{i\}andF^i\\hat\{F\}\_\{i\}are the server and client objectives with respect to the public dataUiU\_\{i\}restricted to support ofii’s data distribution, whereasJ^ic\\hat\{J\}\_\{i\}^\{c\}andF^ic\\hat\{F\}\_\{i\}^\{c\}are these objectives on the portion ofUUoutside clientii’s support\. SinceU=Ui∪UicU=U\_\{i\}\\cup U\_\{i\}^\{c\}for eachii, we can also note thatJ^​\(hk;ℓ\)=J^i​\(hk;ℓ\)\+J^ic​\(hk;ℓ\)\\hat\{J\}\(h\_\{k\};\\ell\)=\\hat\{J\}\_\{i\}\(h\_\{k\};\\ell\)\+\\hat\{J\}\_\{i\}^\{c\}\(h\_\{k\};\\ell\)andF^​\(fi;ℓ\)=F^i​\(fi;ℓ\)\+F^ic​\(fi;ℓ\)\\hat\{F\}\(f\_\{i\};\\ell\)=\\hat\{F\}\_\{i\}\(f\_\{i\};\\ell\)\+\\hat\{F\}\_\{i\}^\{c\}\(f\_\{i\};\\ell\)\.

Moreover, insofar as our interest is in convergence \(in terms of*true*risk\), the objectives for the clients and server defined above are only estimates thereof\. The true objective of each server\-side cluster isJ​\(hk;ℓ\)=𝔼x∼𝒬​\[G​\(hk,f¯k,ℓ\)​\(x\)\]J\(h\_\{k\};\\ell\)=\\mathbb\{E\}\_\{x\\sim\\mathcal\{Q\}\}\[G\(h\_\{k\},\\bar\{f\}\_\{k\},\\ell\)\(x\)\], withJi​\(hk;ℓ\)=𝒬​\(𝒮i\)​𝔼x∼𝒬i​\[G​\(hk,f¯k,ℓ\)​\(x\)\];Jic​\(hk\)=𝒬​\(𝒮ic\)​𝔼x∼𝒬ic​\[G​\(hk,f¯k,ℓ\)​\(x\)\]J\_\{i\}\(h\_\{k\};\\ell\)=\\mathcal\{Q\}\(\\mathcal\{S\}\_\{i\}\)\\mathbb\{E\}\_\{x\\sim\\mathcal\{Q\}\_\{i\}\}\[G\(h\_\{k\},\\bar\{f\}\_\{k\},\\ell\)\(x\)\];J\_\{i\}^\{c\}\(h\_\{k\}\)=\\mathcal\{Q\}\(\\mathcal\{S\}\_\{i\}^\{c\}\)\\mathbb\{E\}\_\{x\\sim\\mathcal\{Q\}\_\{i\}^\{c\}\}\[G\(h\_\{k\},\\bar\{f\}\_\{k\},\\ell\)\(x\)\]; and we defineF​\(fi;ℓ\)F\(f\_\{i\};\\ell\),Fi​\(fi;ℓ\)F\_\{i\}\(f\_\{i\};\\ell\), andFic​\(fi;ℓ\)F\_\{i\}^\{c\}\(f\_\{i\};\\ell\)analogously for each clientii\. While cumbersome, the key consideration that this notation enables us to deal with is how well\-behaved the objective functions of the server and clients are on the supported and unsupported parts of the public data\.

### 5\.1Clustering

Our risk analysis relies on a controlled within\-cluster pseudo\-label distanceBBover the shared unlabeled setUU\. In particular, the greedy clustering procedure in Section[4\.1](https://arxiv.org/html/2605.11165#S4.SS1)enforces this property in the first round by constructing clusters with thresholdB0B\_\{0\}, whereB0≤BB\_\{0\}\\leq B\. In subsequent rounds, bounded in\-cluster disagreement is expected to continue to hold because all clients are anchored to the same public set and distill from the same cluster\-level teacher in each round\. This condition is the mechanism that lets us control how much a cluster\-level aggregation step can degrade clientii’s pseudo\-labels, which leads to Lemma[5\.2](https://arxiv.org/html/2605.11165#S5.Thmtheorem2)\. Without an explicit bound, similarity\-based clustering may work well empirically but offers no convergence guarantees: aggregation within a cluster can introduce uncontrolled label error\.

We additionally require clients to have confidence margin onUiU\_\{i\}to prevent small within\-cluster discrepancies from flipping argmax labels\. This is reasonable becausefi\(t,1\)f\_\{i\}^\{\(t,1\)\}is obtained by fine\-tuning on clientii’s labeled data, so its predictions on the in\-support public subsetUiU\_\{i\}are expected to be sufficiently decisive\. Formally, for anyg​\(x\)g\(x\)mapping inputs to a distribution over labels\[M\]\[M\], define the marginΔg​\(x\)=g​\(x\)\(1\)−g​\(x\)\(2\)\\Delta\_\{g\}\(x\)=g\(x\)\_\{\(1\)\}\-g\(x\)\_\{\(2\)\}, whereg​\(x\)\(m\)g\(x\)\_\{\(m\)\}denotes themm\-th largest coordinate ofg​\(x\)g\(x\)\.

###### Assumption 5\.1\.

For each iterationttand clusterkk, the within\-cluster pseudo\-label distance is bounded:maxi,j∈𝒢k⁡d\(t\)​\(i,j\)≤B\.\\max\_\{i,j\\in\\mathcal\{G\}\_\{k\}\}d^\{\(t\)\}\(i,j\)\\leq B\.Moreover, for each clientiithere existsγ\>0\\gamma\>0such that for everyx∈Uix\\in U\_\{i\},Δfi\(t,1\)​\(x\)≥γ\.\\Delta\_\{f\_\{i\}^\{\(t,1\)\}\}\(x\)\\geq\\gamma\.

###### Lemma 5\.2\.

For each clientii,Errx∈Ui​\(f¯k\(t\)\)≤Errx∈Ui​\(fi\(t,1\)\)\+2​Bγ​\|Ui\|\.\\mathrm\{Err\}\_\{x\\in U\_\{i\}\}\(\\bar\{f\}\_\{k\}^\{\(t\)\}\)\\;\\leq\\;\\mathrm\{Err\}\_\{x\\in U\_\{i\}\}\(f\_\{i\}^\{\(t,1\)\}\)\\;\+\\;\\frac\{2B\}\{\\gamma\\,\|U\_\{i\}\|\}\.

### 5\.2Data and Objective Conditions

Our next challenge comes from the server’s reliance on the public unlabeled poolUU: since it never observes clients’ labeled data, any guarantee must ensure that the distributionQQinduced byUUprovides adequate coverage of each local distribution𝒟i\\mathcal\{D\}\_\{i\}\. In particular, a key requirement is that the public distributionQQdoes not under\-cover nor over\-concentrate on regions where clientiihas support\. We formalize this as the following condition\.

###### Assumption 5\.3\(Distributional coverage\)\.

For each clientiiwith local data distribution𝒟i\\mathcal\{D\}\_\{i\}that has densitypip\_\{i\},wi,1=supx∈𝒮ipi​\(x\)qi​\(x\)w\_\{i,1\}\\;=\\;\\sup\_\{x\\in\\mathcal\{S\}\_\{i\}\}\\frac\{p\_\{i\}\(x\)\}\{q\_\{i\}\(x\)\}andwi,2=supx∈𝒮iqi​\(x\)pi​\(x\)w\_\{i,2\}\\;=\\;\\sup\_\{x\\in\\mathcal\{S\}\_\{i\}\}\\frac\{q\_\{i\}\(x\)\}\{p\_\{i\}\(x\)\}are bounded\.

We note that we can in principle*construct*UUto be sufficiently diverse, or define a sampling distribution𝒬\\mathcal\{Q\}to have a strictly positive density over𝒳\\mathcal\{X\}, such as a Gaussian distribution\. Then, if𝒳\\mathcal\{X\}is bounded \(for example,𝒳=\[0,1\]m\\mathcal\{X\}=\[0,1\]^\{m\}, as for normalized image data\), the condition will hold for reasonable𝒟i\\mathcal\{D\}\_\{i\}\(since the sup is over its support\)\.

Furthermore, we must ensure convexity for a well\-conditioned optimization landscape, and that data points outside the support of each client’s distribution do not excessively interfere with learning from client\-supported data, where pseudo\-labels are expected to be more accurate\. To capture this, we require key data\-dependent regularity conditions on the server and client objectives,J,J^,FJ,\\hat\{J\},F, andF^\\hat\{F\}, captured by the following definition\.

###### Definition 5\.4\(Locally well\-conditioned Objective\)\.

We say a differentiable functionΦ​\(θ;ℓ\)\\Phi\(\\theta;\\ell\)is\(μ,L\)\(\\mu,L\)\-*locally well\-conditioned*if

1. 1\.Φi\\Phi\_\{i\}isμ\\mu\-strongly convex:∇2Φi​\(θ\)⪰μ​I​for​μ\>0\.\\nabla^\{2\}\\Phi\_\{i\}\(\\theta\)\\succeq\\mu I\\ \\text\{for \}\\mu\>0\.
2. 2\.∇Φic\\nabla\\Phi^\{c\}\_\{i\}isL2L\_\{2\}\-Lipschitz continuous:‖∇Φic​\(θ\)−∇Φic​\(θ′\)‖≤L​‖θ−θ′‖​∀θ,θ′\\\|\\nabla\\Phi\_\{i\}^\{c\}\(\\theta\)\-\\nabla\\Phi\_\{i\}^\{c\}\(\\theta^\{\\prime\}\)\\\|\\leq L\\\|\\theta\-\\theta^\{\\prime\}\\\|\\ \\forall\\theta,\\theta^\{\\prime\}for some0<L<μ0<L<\\mu\.
3. 3\.‖∇Φ^ic​\(θ∗\)‖≤τ\\\|\\nabla\\hat\{\\Phi\}\_\{i\}^\{c\}\(\\theta^\{\*\}\)\\\|\\leq\\tauwhereθ∗=arg⁡minθ⁡Φ^i​\(θ\)\\theta^\{\*\}=\\arg\\min\_\{\\theta\}\\hat\{\\Phi\}\_\{i\}\(\\theta\)\.

In the following, we note that the functionsFF,F^\\hat\{F\},JJ, andJ^\\hat\{J\}are, effectively, functions of*parameters*of the models they are constructed around, whether these are the server\-side modelshkh\_\{k\}or the client\-side modelsfif\_\{i\}\.

###### Assumption 5\.5\.

For all clientsiiand COSMOS iterationstt:

1. 1\.F​\(θi\(t\);ℓ\)F\(\\theta\_\{i\}^\{\(t\)\};\\ell\),F^​\(θi\(t\);ℓ\)\\hat\{F\}\(\\theta\_\{i\}^\{\(t\)\};\\ell\),J​\(θk\(t\);ℓ\)J\(\\theta\_\{k\}^\{\(t\)\};\\ell\), andJ^​\(θk\(t\);ℓ\)\\hat\{J\}\(\\theta\_\{k\}^\{\(t\)\};\\ell\)are locally well\-conditioned, and
2. 2\.The pseudolabels generated by any pairs of parametersθ,θ′\\theta,\\theta^\{\\prime\}from the same hypothesis class areL1L\_\{1\}\-Lipschitz continuous:supx∈𝒳‖θ​\(x\)−θ′​\(x\)‖∞≤L1​‖θ−θ′‖\.\\sup\_\{x\\in\\mathcal\{X\}\}\\big\\\|\\theta\(x\)\-\{\\theta^\{\\prime\}\}\(x\)\\big\\\|\_\{\\infty\}\\;\\leq\\;L\_\{1\}\\,\\\|\\theta\-\\theta^\{\\prime\}\\\|\.

The final regularity condition bounds how often inputs fall arbitrarily close to the decision boundary, so that a small change in the objective does not frequently flip the predicted label\. Concretely, we adopt a standard Tsybakov\-style margin condition for the optimal classifiers\.

###### Assumption 5\.6\(Confidence margin condition\)\.

For each clientii, letfi∗∈arg⁡ming∈ℋi⁡Fi​\(g;ℓ\)f\_\{i\}^\{\*\}\\in\\arg\\min\_\{g\\in\\mathcal\{H\}\_\{i\}\}F\_\{i\}\(g;\\ell\)and lethi∗∈arg⁡ming∈ℋS⁡Ji​\(g;ℓ\)h\_\{i\}^\{\*\}\\in\\arg\\min\_\{g\\in\\mathcal\{H\}\_\{S\}\}J\_\{i\}\(g;\\ell\)\. Then for allii, there exist constantsC\>0C\>0andα\>0\\alpha\>0such that∀t≥0\\forall t\\geq 0,Pr⁡\(Δfi∗​\(x\)≤t\)≤C​tα\\Pr\\big\(\\Delta\_\{f\_\{i\}^\{\*\}\}\(x\)\\leq t\\big\)\\;\\leq\\;C\\,t^\{\\alpha\}andPr⁡\(Δhi∗​\(x\)≤t\)≤C​tα\.\\Pr\\big\(\\Delta\_\{h\_\{i\}^\{\*\}\}\(x\)\\leq t\\big\)\\;\\leq\\;C\\,t^\{\\alpha\}\.

### 5\.3Pseudolabel and Label Conditions

While our clustering controls the additional error introduced by aggregation, and our data/objective conditions support stable gradient\-based optimization, a purely pseudo\-label\-based method still requires additional structure to be provably effective\. In particular, the pseudo\-labels must carry nontrivial information about the ground truth, and the ground truth labels should satisfy some structural connectivity so that the learner can generalize beyond the pseudo\-labeled points rather than merely memorizing arbitrary functions\. These conditions are what enable learning from pseudo\-labels, inducing the weak\-to\-strong generalization we observe at the server’s end\[[14](https://arxiv.org/html/2605.11165#bib.bib54)\]\.

Specifically, we adopt the expansion\-and\-robustness framework of\[[31](https://arxiv.org/html/2605.11165#bib.bib53)\], which expresses label connectivity via input\-space transformations𝒯\\mathcal\{T\}and associated neighborhoods𝒩​\(x\)\\mathcal\{N\}\(x\)of a given inputxx\(see Section[4\.1](https://arxiv.org/html/2605.11165#S4.SS1), Step 3\)\. For a subsetV⊆𝒳V\\subseteq\\mathcal\{X\}, define𝒩​\(V\)\\mathcal\{N\}\(V\)as the union of neighborhoods of its points:𝒩​\(V\)=⋃x∈V𝒩​\(x\)\\mathcal\{N\}\(V\)=\\bigcup\_\{x\\in V\}\\mathcal\{N\}\(x\)\. Expansion ensures that if a setVVcontains a small fraction of a given class, then closingVVunder the neighborhood operator𝒩​\(⋅\)\\mathcal\{N\}\(\\cdot\)captures a larger fraction of that same class\.

###### Definition 5\.7\(\(b,c\)\(b,c\)\-expansion\)\.

A distribution𝒟\\mathcal\{D\}satisfies\(b,c\)\(b,c\)\-expansion if for every classmmand∀V⊆𝒳\\forall V\\subseteq\\mathcal\{X\}withPrx∼𝒟\|x∈V⁡\[y​\(x\)=m\]≤b\\Pr\_\{x\\sim\\mathcal\{D\}\|x\\in V\}\[y\(x\)=m\]\\leq b, we havePrx∼𝒟\|x∈𝒩​\(V\)⁡\[y​\(x\)=m\]≥min⁡\{c​Prx∼𝒟⁡\[y​\(x\)=m\|c∈V\],1\}\\Pr\_\{x\\sim\\mathcal\{D\}\|x\\in\\mathcal\{N\}\(V\)\}\[y\(x\)=m\]\\geq\\min\\\{c\\Pr\_\{x\\sim\\mathcal\{D\}\}\[y\(x\)=m\\ \|\\ c\\in V\],1\\\}\.

To quantify the probability that pseudo\-labelers make mistakes for any clientiiand iterationtt, define

b¯=supt,imax\{𝒬i\(ℳ\(hπ​\(i\)\(t\)\)\),𝒬i\(ℳ\(f¯π​\(i\)\(t\)\)\),\\displaystyle\\overline\{b\}=\\sup\_\{t,i\}\\max\\Bigl\\\{\\mathcal\{Q\}\_\{i\}\\\!\\left\(\\mathcal\{M\}\(h\_\{\\pi\(i\)\}^\{\(t\)\}\)\\right\),\\;\\mathcal\{Q\}\_\{i\}\\\!\\left\(\\mathcal\{M\}\(\\bar\{f\}\_\{\\pi\(i\)\}^\{\(t\)\}\)\\right\),\\;𝒟i\(ℳ\(hπ​\(i\)\(t\)\)\),𝒟i\(ℳ\(f¯π​\(i\)\(t\)\)\)\},\\displaystyle\\mathcal\{D\}\_\{i\}\\\!\\left\(\\mathcal\{M\}\(h\_\{\\pi\(i\)\}^\{\(t\)\}\)\\right\),\\;\\mathcal\{D\}\_\{i\}\\\!\\left\(\\mathcal\{M\}\(\\bar\{f\}\_\{\\pi\(i\)\}^\{\(t\)\}\)\\right\)\\Bigr\\\},whereℳ​\(f\):=\{x∈𝒳:A∘f​\(x\)≠y​\(x\)\}\\mathcal\{M\}\(f\):=\\\{x\\in\\mathcal\{X\}:A\\circ f\(x\)\\neq y\(x\)\\\}\.

The following natural condition requires that the probabilityb¯\\overline\{b\}of pseudo\-labeling errors is not too large, as well as inputs with a given class are sufficiently well\-connected\.

###### Assumption 5\.8\(Effective pseudo\-labelers\)\.

b¯≤13\\overline\{b\}\\leq\\tfrac\{1\}\{3\}\.

Complementing expansion, robustness formalizes the idea that inputs close in the input space share the same labels\. We characterize this property in terms of robustness loss\.

###### Definition 5\.9\.

For a functionffand distribution𝒟\\mathcal\{D\},*robustness loss*is the fraction of examples that are not robust to input transformations:

Rℬ​\(f,𝒟\)=𝔼x∼𝒟​\[𝟏​\{∃x′∈𝒩​\(x\):f​\(x′\)≠f​\(x\)\}\]\.R\_\{\\mathcal\{B\}\}\(f,\\mathcal\{D\}\)\\;=\\;\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\Bigl\[\\mathbf\{1\}\\bigl\\\{\\exists\\,x^\{\\prime\}\\in\\mathcal\{N\}\(x\)\\colon f\(x^\{\\prime\}\)\\neq f\(x\)\\bigr\\\}\\Bigr\]\.

###### Assumption 5\.10\(Label connectivity\)\.

For every clientiiand true classifiery​\(x\)y\(x\),𝒬i\\mathcal\{Q\}\_\{i\}and𝒟i\\mathcal\{D\}\_\{i\}satisfy\(b¯,c¯\)\(\\overline\{b\},\\overline\{c\}\)\-expansion for somec¯\>3\\overline\{c\}\>3andmax\{Rℬ\(y,𝒬i\),\\max\\\{R\_\{\\mathcal\{B\}\}\(y,\\mathcal\{Q\}\_\{i\}\),Rℬ\(y,𝒟i\)\}≤ρR\_\{\\mathcal\{B\}\}\(y,\\mathcal\{D\}\_\{i\}\)\\\}\\leq\\;\\rho\.

### 5\.4Risk Contraction

We now put these tools together to provide bounds for the client\-side and server\-side risk\. The proofs are provided in the Appendix \([B\.2](https://arxiv.org/html/2605.11165#A2.SS2)\)\. Definec=min⁡\{c¯,1b¯\}c=\\min\\left\\\{\\overline\{c\},\\,\\frac\{1\}\{\\overline\{b\}\}\\right\\\}\. Suppose that 1\)c\>2​w1​w2\+1,c\>2w\_\{1\}w\_\{2\}\+1,wherew1=maxi⁡\{wi,1\}w\_\{1\}=\\max\_\{i\}\\\{w\_\{i,1\}\\\}andw2=maxi⁡\{wi,2\}w\_\{2\}=\\max\_\{i\}\\\{w\_\{i,2\}\\\}, 2\)λ=2​cc\+1\\lambda=\\tfrac\{2c\}\{c\+1\}, and 3\) Conditions[5\.1](https://arxiv.org/html/2605.11165#S5.Thmtheorem1)–[5\.10](https://arxiv.org/html/2605.11165#S5.Thmtheorem10)hold\.

###### Theorem 5\.11\.

For any clientiiat any iterationtt, with probability at least1−δ1\-\\delta, its server modelhπ​\(i\)\(t\)h\_\{\\pi\(i\)\}^\{\(t\)\}’s risk is contracting:R𝒟i​\(hπ​\(i\)\(t\)\)≤κ1​R𝒟i​\(fi\(t,1\)\)\+O~​\(n−1/2\),R\_\{\\mathcal\{D\}\_\{i\}\}\(h\_\{\\pi\(i\)\}^\{\(t\)\}\)\\;\\leq\\;\\kappa\_\{1\}\\,R\_\{\\mathcal\{D\}\_\{i\}\}\(f\_\{i\}^\{\(t,1\)\}\)\+\\tilde\{O\}\(n^\{\-1/2\}\),whereκ1=2​w1​w2c−1<1\\kappa\_\{1\}=\\tfrac\{2w\_\{1\}w\_\{2\}\}\{c\-1\}<1and the lower order term hides constants and poly\-log terms innn\.

###### Theorem 5\.12\.

Letκ2=4​w1​w2\(c−1\)2<1\\kappa\_\{2\}=\\tfrac\{4w\_\{1\}w\_\{2\}\}\{\(c\-1\)^\{2\}\}<1\. For each clientiiand iterationtt, with probability at least1−δ1\-\\delta,R𝒟i​\(fi\(t,2\)\)≤κ2​R𝒟i​\(fi\(t−1,2\)\)\+O~​\(log⁡\(1/δ\)n\)\.R\_\{\\mathcal\{D\}\_\{i\}\}\\\!\\left\(f\_\{i\}^\{\(t,2\)\}\\right\)\\;\\leq\\;\\kappa\_\{2\}\\,R\_\{\\mathcal\{D\}\_\{i\}\}\\\!\\left\(f\_\{i\}^\{\(t\-1,2\)\}\\right\)\+\\;\\tilde\{O\}\\\!\\left\(\\sqrt\{\\frac\{\\log\(1/\\delta\)\}\{n\}\}\\right\)\.

Applying a union bound, we obtain a contraction ove personalization risk bound at exponential rate\.

###### Corollary 5\.13\(Personalization Risk Bound\)\.

With probability at least1−δ1\-\\deltapersonalization risk afterTTiterations is

R†​\(h1\(T\),…,hK\(T\)\)≤1N​∑i=1Nκ2T​R𝒟i​\(fi\(1,1\)\)\+O~​\(log⁡\(T/δ\)n\)\.R^\{\\dagger\}\\\!\\left\(h^\{\(T\)\}\_\{1\},\\ldots,h^\{\(T\)\}\_\{K\}\\right\)\\leq\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\kappa\_\{2\}^\{\\,T\}R\_\{\\mathcal\{D\}\_\{i\}\}\\\!\\left\(f\_\{i\}^\{\(1,1\)\}\\right\)\+\\tilde\{O\}\\\!\\left\(\\sqrt\{\\frac\{\\log\(T/\\delta\)\}\{n\}\}\\right\)\.

Moreover, for sufficiently largennandTT, personalization risk drops below the risk from local training alone\.

![Refer to caption](https://arxiv.org/html/2605.11165v1/x1.png)Figure 2:Comparison of COSMOS and model\-agnostic baselines on four benchmarks\. Client models consist of MobileNet and SqueezeNet architectures\. Curves show mean client accuracy over rounds \(Top\-1 for CIFAR\-10/100 and EMNIST; Top\-5 for Tiny ImageNet\)\. A statistical comparison of final\-round performance is provided in the Appendix \(Appendix[D](https://arxiv.org/html/2605.11165#A4), Table[2](https://arxiv.org/html/2605.11165#A3.T2)\)\.![Refer to caption](https://arxiv.org/html/2605.11165v1/x2.png)Figure 3:Performance of COSMOS and heterogeneous\-model baselines on CIFAR\-100 under different mixtures of client architectures\. Each subfigure corresponds to a distinct architecture combination\. A statistical summary of the final communication round appears in Appendix[D](https://arxiv.org/html/2605.11165#A4), Table[3](https://arxiv.org/html/2605.11165#A3.T3)\.

## 6Experiments

Experiment Setup:We construct client datasets using Dirichlet non\-IID partitioning with concentration parameterα\\alpha\[[32](https://arxiv.org/html/2605.11165#bib.bib6)\]\. Specifically, the label space is partitioned into five disjoint groups, each containing20%20\\%of the classes, and clients are assigned to one group\. For each class, its samples are distributed among the clients within the same group according to a Dirichlet draw\. To introduce limited cross\-group exposure and better reflect realistic data sharing, a fraction of each client’s local samples \(10%\) is pooled, randomly shuffled, and uniformly redistributed across all clients\. Our unlabeled datasetUUconsists of 20% of all training data\. We use accuracy as the efficacy measure\.

We selected client architectures among 4 options: AlexNet, ResNet18, MobileNet, and SqueezeNet\. We use VGG16 for the server\. Our focus is on the model\-agnostic setup, where we compare COSMOS with the only three truly model\-agnostic baselines \(all others require knowledge of client model architecture; see Section[2](https://arxiv.org/html/2605.11165#S2)\): FedMD\[[16](https://arxiv.org/html/2605.11165#bib.bib26)\], FedCT\[[1](https://arxiv.org/html/2605.11165#bib.bib19)\], and COMET\[[3](https://arxiv.org/html/2605.11165#bib.bib20)\]\. In the Appendix, we also evaluate COSMOS in*homogeneous model*settings, where we compare to*eight state\-of\-the\-art FL baselines*, focusing primarily on PFL approaches\.

We use four image classification benchmarks: EMNIST\-balanced\[[4](https://arxiv.org/html/2605.11165#bib.bib57)\], CIFAR\-10 and CIFAR\-100\[[13](https://arxiv.org/html/2605.11165#bib.bib58)\], and Tiny ImageNet\[[15](https://arxiv.org/html/2605.11165#bib.bib56)\]\. The main experiments use 25 clients over 10 communication rounds, using Dirichlet non\-IID sampling \(α=5\\alpha=5\)\. Each data point in the figures reflects the*mean client accuracy*over three random seeds\. We present further details about the*experiment setup, confidence intervals, scalability evaluation of COSMOS as we increase the number of clients, and other parametric ablations, in the Appendix\.*

Results:Figure[2](https://arxiv.org/html/2605.11165#S5.F2)presents the main results across all four benchmarks in the model\-agnostic setting withα=5\\alpha=5\. In all cases, COSMOS outperforms all model\-agnostic baselines by a large margin\. Moreover, its advantage increases with increasing task complexity \(from EMNIST to TinyImageNet\)\. Results with varyingα\\alpha\(which are largely consistent with this observation\) as well as confidence bound information are in the Appendix\.

Figure[3](https://arxiv.org/html/2605.11165#S5.F3)evaluates COSMOS under different mixtures of client architectures\. Across all combinations, COSMOS remains robust to heterogeneity in client capabilities, including settings with a high proportion of weaker models \(AlexNet\), maintaining substantial performance edge over baselines\. Additionally, our evaluation in the homogeneous architecture setting \(see the Appendix, Appendix[D](https://arxiv.org/html/2605.11165#A4), Figure[6](https://arxiv.org/html/2605.11165#A4.F6)\) demonstrates that COSMOS remains competitive in such environments with the best\-performing PFL and model\-agnostic baselines\. This is particularly notable, as the best\-performing baseline*changes for different model architectures*, further demonstrating the value of our model\-agnostic design\. Moreover, our analysis of the communication cost demonstrates the bandwidth advantage of COSMOS over prior art \(see the Appendix, Table[4](https://arxiv.org/html/2605.11165#A3.T4)\)\.

Finally, we analyzed the sensitivity of COSMOS to key design choices and hyperparameters \(including temperatureTTand weight of the regularization termλ\\lambda\)\. We find that COSMOS is quite robust to small changes inλ\\lambda, withT=1T=1andλ=5\\lambda=5yielding the best performance\. Additionally, we observe improvement in COSMOS performance as the number of clients or the size ofUUincrease, as well as with the use of a more powerful server model architecture; see the Appendix for further details\.

## 7Conclusion

In this work, we introduce COSMOS, a framework that addresses a critical bottleneck in federated learning: achieving high\-performance personalization while remaining fully model\-agnostic\. 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\. We establish a general end\-to\-end theoretical analysis of risk \(generalization\) bound contraction of our framework that significantly generalizes past theoretical results in this setting\. And we validate our algorithm’s effectiveness against state\-of\-art baselines across multiple benchmarks\. An exciting next step will be to move beyond a client\-server structure to a fully decentralized learning system, further enhancing federated learning’s practical relevance\.

## References

- \[1\]A\. Abourayya, J\. Kleesiek, K\. Rao, E\. Ayday, B\. Rao, G\. I\. Webb, and M\. Kamp\(2025\)Little Is Enough: Boosting Privacy by Sharing Only Hard Labels in Federated Semi\-Supervised Learning\.InProceedings of the AAAI Conference on Artificial Intelligence,Cited by:[§1](https://arxiv.org/html/2605.11165#S1.p1.1),[§1](https://arxiv.org/html/2605.11165#S1.p2.1),[§2](https://arxiv.org/html/2605.11165#S2.p2.1),[§2](https://arxiv.org/html/2605.11165#S2.p4.1),[§6](https://arxiv.org/html/2605.11165#S6.p2.1)\.
- \[2\]A\. Afonin and S\. P\. Karimireddy\(2022\)Towards Model Agnostic Federated Learning Using Knowledge Distillation\.InInternational Conference on Learning Representations,Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p2.1)\.
- \[3\]Y\. J\. Cho, J\. Wang, T\. Chirvolu, and G\. Joshi\(2023\)Communication\-Efficient and Model\-Heterogeneous Personalized Federated Learning via Clustered Knowledge Transfer\.IEEE Journal of Selected Topics in Signal Processing\.Cited by:[§1](https://arxiv.org/html/2605.11165#S1.p1.1),[§1](https://arxiv.org/html/2605.11165#S1.p2.1),[§1](https://arxiv.org/html/2605.11165#S1.p3.1),[§2](https://arxiv.org/html/2605.11165#S2.p2.1),[§2](https://arxiv.org/html/2605.11165#S2.p3.2),[§2](https://arxiv.org/html/2605.11165#S2.p4.1),[§5](https://arxiv.org/html/2605.11165#S5.p1.4),[§6](https://arxiv.org/html/2605.11165#S6.p2.1)\.
- \[4\]G\. Cohen, S\. Afshar, J\. Tapson, and A\. Van Schaik\(2017\)EMNIST: Extending MNIST to Handwritten Letters\.InInternational Joint Conference on Neural Networks \(IJCNN\),Cited by:[§6](https://arxiv.org/html/2605.11165#S6.p3.1)\.
- \[5\]Y\. Deng, M\. M\. Kamani, and M\. Mahdavi\(2020\)Adaptive Personalized Federated Learning\.arXiv preprint arXiv:2003\.13461\.Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p3.2)\.
- \[6\]C\. T\. Dinh, N\. Tran, and J\. Nguyen\(2020\)Personalized Federated Learning With Moreau Envelopes\.Advances in Neural Information Processing Systems\.Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p3.2)\.
- \[7\]M\. Duan, D\. Liu, X\. Ji, R\. Liu, L\. Liang, X\. Chen, and Y\. Tan\(2021\)FedGroup: efficient federated learning via decomposed similarity\-based clustering\.InIEEE International Conference on Parallel & Distributed Processing with Applications \(ISPA/BDCloud/SocialCom/SustainCom\),Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p3.2)\.
- \[8\]A\. Fallah, A\. Mokhtari, and A\. Ozdaglar\(2020\)Personalized federated learning with theoretical guarantees: a model\-agnostic meta\-learning approach\.InAdvances in Neural Information Processing Systems,Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p3.2)\.
- \[9\]L\. Ge, M\. Lanier, A\. Sarkar, B\. Guresti, Y\. Vorobeychik, and C\. Zhang\(2025\)Learning Policy Committees for Effective Personalization in MDPs With Diverse Tasks\.Proceedings of the 42nd International Conference on Machine Learning\.Cited by:[§4\.1](https://arxiv.org/html/2605.11165#S4.SS1.p3.5)\.
- \[10\]A\. Ghosh, J\. Chung, D\. Yin, and K\. Ramchandran\(2020\)An Efficient Framework for Clustered Federated Learning\.Advances in Neural Information Processing Systems\.Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p3.2)\.
- \[11\]B\. Gong, T\. Xing, Z\. Liu, W\. Xi, and X\. Chen\(2022\)Adaptive Client Clustering for Efficient Federated Learning Over Non\-IID and Imbalanced Data\.IEEE Transactions on Big Data\.Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p3.2)\.
- \[12\]S\. P\. Karimireddy, S\. Kale, M\. Mohri, S\. Reddi, S\. Stich, and A\. T\. Suresh\(2020\)Scaffold: Stochastic Controlled Averaging for Federated Learning\.InInternational Conference on Machine Learning,Cited by:[§1](https://arxiv.org/html/2605.11165#S1.p1.1),[§2](https://arxiv.org/html/2605.11165#S2.p1.1)\.
- \[13\]A\. Krizhevsky, G\. Hinton,et al\.\(2009\)Learning Multiple Layers of Features From Tiny Images\.Toronto, ON, Canada\.Cited by:[§6](https://arxiv.org/html/2605.11165#S6.p3.1)\.
- \[14\]H\. Lang, D\. Sontag, and A\. Vijayaraghavan\(2024\)Theoretical Analysis of Weak\-to\-Strong Generalization\.Advances in Neural Information Processing Systems\.Cited by:[§1](https://arxiv.org/html/2605.11165#S1.p5.1),[§5\.3](https://arxiv.org/html/2605.11165#S5.SS3.p1.1)\.
- \[15\]Y\. Le and X\. Yang\(2015\)Tiny ImageNet Visual Recognition Challenge\.Note:Stanford CS231N: Convolutional Neural Networks for Visual RecognitionCited by:[§6](https://arxiv.org/html/2605.11165#S6.p3.1)\.
- \[16\]D\. Li and J\. Wang\(2019\)FedMD: Heterogenous Federated Learning via Model Distillation\.arXiv preprint arXiv:1910\.03581\.Cited by:[§1](https://arxiv.org/html/2605.11165#S1.p1.1),[§1](https://arxiv.org/html/2605.11165#S1.p2.1),[§2](https://arxiv.org/html/2605.11165#S2.p2.1),[§2](https://arxiv.org/html/2605.11165#S2.p4.1),[§6](https://arxiv.org/html/2605.11165#S6.p2.1)\.
- \[17\]T\. Li, S\. Hu, A\. Beirami, and V\. Smith\(2021\)Ditto: Fair and Robust Federated Learning Through Personalization\.InInternational Conference on Machine Learning,Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p3.2)\.
- \[18\]T\. Li, A\. K\. Sahu, M\. Zaheer, M\. Sanjabi, A\. Talwalkar, and V\. Smith\(2020\)Federated Optimization in Heterogeneous Networks\.Proceedings of Machine Learning and Systems\.Cited by:[§1](https://arxiv.org/html/2605.11165#S1.p1.1),[§2](https://arxiv.org/html/2605.11165#S2.p1.1)\.
- \[19\]T\. Lin, L\. Kong, S\. U\. Stich, and M\. Jaggi\(2020\)Ensemble Distillation for Robust Model Fusion in Federated Learning\.Advances in Neural Information Processing Systems\.Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p2.1)\.
- \[20\]J\. Liu, X\. Liu, S\. Wang, X\. Wan, D\. Li, K\. Lu, and K\. He\(2025\)Communication\-efficient federated multi\-view clustering\.IEEE Transactions on Pattern Analysis and Machine Intelligence48,pp\. 17–32\.External Links:[Link](https://api.semanticscholar.org/CorpusID:280767120)Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p2.1)\.
- \[21\]D\. Makhija, X\. Han, N\. Ho, and J\. Ghosh\(2022\)Architecture Agnostic Federated Learning for Neural Networks\.InInternational Conference on Machine Learning,Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p2.1)\.
- \[22\]B\. McMahan, E\. Moore, D\. Ramage, S\. Hampson, and B\. A\. y Arcas\(2017\)Communication\-Efficient Learning of Deep Networks from Decentralized Data\.InInternational Conference on Artificial Intelligence and Statistics,Cited by:[§1](https://arxiv.org/html/2605.11165#S1.p1.1),[§2](https://arxiv.org/html/2605.11165#S2.p1.1)\.
- \[23\]A\. Mora, I\. Tenison, P\. Bellavista, and I\. Rish\(2024\)Knowledge Distillation in Federated Learning: A Practical Guide\.InProceedings of the Thirty\-Third International Joint Conference on Artificial Intelligence,Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p2.1)\.
- \[24\]J\. Oh, S\. Kim, and S\. Yun\(2022\)FedBABU: toward enhanced representation for federated image classification\.InInternational Conference on Learning Representations,Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p3.2)\.
- \[25\]F\. Sattler, A\. Marban, R\. Rischke, and W\. Samek\(2022\)CFD: communication\-efficient federated distillation via soft\-label quantization and delta coding\.InIEEE Transactions on Network Science and Engineering,Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p2.1)\.
- \[26\]O\. Shahid, S\. Pouriyeh, R\. M\. Parizi, Q\. Z\. Sheng, G\. Srivastava, and L\. Zhao\(2021\)Communication efficiency in federated learning: achievements and challenges\.arXiv preprint arXiv:2107\.10996\.Cited by:[§1](https://arxiv.org/html/2605.11165#S1.p2.1)\.
- \[27\]A\. Shamsian, A\. Navon, E\. Fetaya, and G\. Chechik\(2021\)Personalized Federated Learning Using Hypernetworks\.InInternational Conference on Machine Learning,Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p3.2),[§5](https://arxiv.org/html/2605.11165#S5.p1.4)\.
- \[28\]R\. Tamirisa, C\. Xie, W\. Bao, A\. Zhou, R\. Arel, and A\. Shamsian\(2024\)FedSelect: Personalized Federated Learning With Customized Selection of Parameters for Fine\-Tuning\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p3.2)\.
- \[29\]A\. Z\. Tan, H\. Yu, L\. Cui, and Q\. Yang\(2022\)Towards Personalized Federated Learning\.IEEE Transactions on Neural Networks and Learning Systems\.Cited by:[§1](https://arxiv.org/html/2605.11165#S1.p1.1),[§2](https://arxiv.org/html/2605.11165#S2.p3.2)\.
- \[30\]Y\. Tan, G\. Long, L\. Liu, T\. Zhou, Q\. Lu, J\. Jiang, and C\. Zhang\(2022\)FedProto: Federated Prototype Learning Across Heterogeneous Clients\.InProceedings of the AAAI Conference on Artificial Intelligence,Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p2.1)\.
- \[31\]C\. Wei, K\. Shen, Y\. Chen, and T\. Ma\(2020\)Theoretical Analysis of Self\-Training With Deep Networks on Unlabeled Data\.InInternational Conference on Learning Representations,Cited by:[§B\.2](https://arxiv.org/html/2605.11165#A2.SS2.12.p3.2),[§B\.2](https://arxiv.org/html/2605.11165#A2.SS2.18.p2.1),[§1](https://arxiv.org/html/2605.11165#S1.p5.1),[§2](https://arxiv.org/html/2605.11165#S2.p4.1),[§5\.3](https://arxiv.org/html/2605.11165#S5.SS3.p2.9)\.
- \[32\]M\. Yurochkin, M\. Agarwal, S\. Ghosh, K\. Greenewald, N\. Hoang, and Y\. Khazaeni\(2019\)Bayesian Nonparametric Federated Learning of Neural Networks\.InInternational Conference on Machine Learning,Cited by:[§6](https://arxiv.org/html/2605.11165#S6.p1.3)\.
- \[33\]J\. Zhang and Y\. Shi\(2024\)A Personalized Federated Learning Method Based on Clustering and Knowledge Distillation\.Electronics\.Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p3.2)\.
- \[34\]Z\. Zhu, J\. Hong, and J\. Zhou\(2021\)Data\-Free Knowledge Distillation for Heterogeneous Federated Learning\.InInternational Conference on Machine Learning,Cited by:[§2](https://arxiv.org/html/2605.11165#S2.p2.1)\.

## Appendix ANotation Table

Table 1:Summary of the main notation used throughout this paper\.
## Appendix BOmitted Proofs

### B\.1Proof of Lemma[5\.2](https://arxiv.org/html/2605.11165#S5.Thmtheorem2)

###### Proof\.

Letc^i​\(x\):=arg⁡maxc∈\[M\]⁡fi\(t,1\)​\(x\)c\\hat\{c\}\_\{i\}\(x\):=\\arg\\max\_\{c\\in\[M\]\}f\_\{i\}^\{\(t,1\)\}\(x\)\_\{c\}andc^k​\(x\):=arg⁡maxc∈\[M\]⁡f¯k\(t\)​\(x\)c\\hat\{c\}\_\{k\}\(x\):=\\arg\\max\_\{c\\in\[M\]\}\\bar\{f\}\_\{k\}^\{\(t\)\}\(x\)\_\{c\}\. For anyx∈Uix\\in U\_\{i\}we have the pointwise implication

𝟏​\{c^k​\(x\)≠y​\(x\)\}≤1​\{c^i​\(x\)≠y​\(x\)\}\+𝟏​\{c^k​\(x\)≠c^i​\(x\)\},\\mathbf\{1\}\\\{\\hat\{c\}\_\{k\}\(x\)\\neq y\(x\)\\\}\\;\\leq\\;\\mathbf\{1\}\\\{\\hat\{c\}\_\{i\}\(x\)\\neq y\(x\)\\\}\+\\mathbf\{1\}\\\{\\hat\{c\}\_\{k\}\(x\)\\neq\\hat\{c\}\_\{i\}\(x\)\\\},because ifc^k​\(x\)=c^i​\(x\)\\hat\{c\}\_\{k\}\(x\)=\\hat\{c\}\_\{i\}\(x\)thenc^k​\(x\)≠y​\(x\)\\hat\{c\}\_\{k\}\(x\)\\neq y\(x\)impliesc^i​\(x\)≠y​\(x\)\\hat\{c\}\_\{i\}\(x\)\\neq y\(x\)\. Now consider the set of public supported points where the averaged predictor disagrees with clientiion the hard classification label\. Call itBad:=\{x∈Ui:c^k​\(x\)≠c^i​\(x\)\}\\mathrm\{Bad\}:=\\\{x\\in U\_\{i\}:\\hat\{c\}\_\{k\}\(x\)\\neq\\hat\{c\}\_\{i\}\(x\)\\\}\.

Averaging overUiU\_\{i\}gives

Errx∈Ui​\(f¯k\(t\)\)≤Errx∈Ui​\(fi\(t,1\)\)\+1\|Ui\|​\|Bad\|\.\\mathrm\{Err\}\_\{x\\in U\_\{i\}\}\\\!\\bigl\(\\bar\{f\}\_\{k\}^\{\(t\)\}\\bigr\)\\;\\leq\\;\\mathrm\{Err\}\_\{x\\in U\_\{i\}\}\\\!\\bigl\(f\_\{i\}^\{\(t,1\)\}\\bigr\)\+\\frac\{1\}\{\|U\_\{i\}\|\}\\,\|\\mathrm\{Bad\}\|\.\(3\)
It remains to bound\|Bad\|\|\\mathrm\{Bad\}\|\. Fix anyx∈Badx\\in\\mathrm\{Bad\}, and denotec⋆:=c^i​\(x\)c^\{\\star\}:=\\hat\{c\}\_\{i\}\(x\)andc~:=c^k​\(x\)\\tilde\{c\}:=\\hat\{c\}\_\{k\}\(x\)\. By definition ofBad\\mathrm\{Bad\}we havec~≠c⋆\\tilde\{c\}\\neq c^\{\\star\}and

f¯k\(t\)​\(x\)c~≥f¯k\(t\)​\(x\)c⋆\.\\bar\{f\}\_\{k\}^\{\(t\)\}\(x\)\_\{\\tilde\{c\}\}\\;\\geq\\;\\bar\{f\}\_\{k\}^\{\(t\)\}\(x\)\_\{c^\{\\star\}\}\.Define the coordinate deviationsδc:=f¯k\(t\)​\(x\)c−fi\(t,1\)​\(x\)c\\delta\_\{c\}:=\\bar\{f\}\_\{k\}^\{\(t\)\}\(x\)\_\{c\}\-f\_\{i\}^\{\(t,1\)\}\(x\)\_\{c\}\. Rearranging the inequality above yieldsδc~−δc⋆≥fi\(t,1\)​\(x\)c⋆−fi\(t,1\)​\(x\)c~\.\\delta\_\{\\tilde\{c\}\}\-\\delta\_\{c^\{\\star\}\}\\;\\geq\\;f\_\{i\}^\{\(t,1\)\}\(x\)\_\{c^\{\\star\}\}\-f\_\{i\}^\{\(t,1\)\}\(x\)\_\{\\tilde\{c\}\}\.By the margin assumption \(applied with competitor classc~\\tilde\{c\}\),fi\(t,1\)​\(x\)c⋆−fi\(t,1\)​\(x\)c~≥γf\_\{i\}^\{\(t,1\)\}\(x\)\_\{c^\{\\star\}\}\-f\_\{i\}^\{\(t,1\)\}\(x\)\_\{\\tilde\{c\}\}\\geq\\gamma, henceδc~−δc⋆≥γ\.\\delta\_\{\\tilde\{c\}\}\-\\delta\_\{c^\{\\star\}\}\\;\\geq\\;\\gamma\.Therefore, by the triangle inequality,

‖f¯k\(t\)​\(x\)−fi\(t,1\)​\(x\)‖1=\\displaystyle\\bigl\\\|\\bar\{f\}\_\{k\}^\{\(t\)\}\(x\)\-f\_\{i\}^\{\(t,1\)\}\(x\)\\bigr\\\|\_\{1\}=∑c∈\[M\]\|δc\|\\displaystyle\\sum\_\{c\\in\[M\]\}\|\\delta\_\{c\}\|\\;≥\\displaystyle\\geq\|δc~\|\+\|δc⋆\|\\displaystyle\\;\|\\delta\_\{\\tilde\{c\}\}\|\+\|\\delta\_\{c^\{\\star\}\}\|\\;≥\\displaystyle\\geq\|δc~−δc⋆\|≥γ\.\\displaystyle\\;\|\\delta\_\{\\tilde\{c\}\}\-\\delta\_\{c^\{\\star\}\}\|\\;\\geq\\;\\gamma\.
Summing this lower bound overx∈Badx\\in\\mathrm\{Bad\}gives

∑x∈Ui‖f¯k\(t\)​\(x\)−fi\(t,1\)​\(x\)‖1\\displaystyle\\sum\_\{x\\in U\_\{i\}\}\\bigl\\\|\\bar\{f\}\_\{k\}^\{\(t\)\}\(x\)\-f\_\{i\}^\{\(t,1\)\}\(x\)\\bigr\\\|\_\{1\}≥∑x∈Bad‖f¯k\(t\)​\(x\)−fi\(t,1\)​\(x\)‖1\\displaystyle\\;\\geq\\;\\sum\_\{x\\in\\mathrm\{Bad\}\}\\bigl\\\|\\bar\{f\}\_\{k\}^\{\(t\)\}\(x\)\-f\_\{i\}^\{\(t,1\)\}\(x\)\\bigr\\\|\_\{1\}≥γ​\|Bad\|\.\\displaystyle\\;\\geq\\;\\gamma\|\\mathrm\{Bad\}\|\.
On the other hand, let us denote clientii’s correspond cluster as\|𝒢k\|\|\\mathcal\{G\}\_\{k\}\|\.

Usingf¯k\(t\)​\(x\)−fi\(t,1\)​\(x\)=1\|𝒢k\|​∑j∈𝒢k\(fj\(t,1\)​\(x\)−fi\(t,1\)​\(x\)\)\\bar\{f\}\_\{k\}^\{\(t\)\}\(x\)\-f\_\{i\}^\{\(t,1\)\}\(x\)=\\frac\{1\}\{\|\\mathcal\{G\}\_\{k\}\|\}\\sum\_\{j\\in\\mathcal\{G\}\_\{k\}\}\\bigl\(f\_\{j\}^\{\(t,1\)\}\(x\)\-f\_\{i\}^\{\(t,1\)\}\(x\)\\bigr\)and the triangle inequality, we get

‖f¯k\(t\)​\(x\)−fi\(t,1\)​\(x\)‖1≤1\|𝒢k\|​∑j∈𝒢k‖fj\(t,1\)​\(x\)−fi\(t,1\)​\(x\)‖1\.\\displaystyle\\bigl\\\|\\bar\{f\}\_\{k\}^\{\(t\)\}\(x\)\-f\_\{i\}^\{\(t,1\)\}\(x\)\\bigr\\\|\_\{1\}\\;\\leq\\;\\frac\{1\}\{\|\\mathcal\{G\}\_\{k\}\|\}\\sum\_\{j\\in\\mathcal\{G\}\_\{k\}\}\\bigl\\\|f\_\{j\}^\{\(t,1\)\}\(x\)\-f\_\{i\}^\{\(t,1\)\}\(x\)\\bigr\\\|\_\{1\}\.
As the greedy clustering algorithm proceeds by putting all the clients that areBB\-close to the center into one cluster, two clients from the same cluster’s distance is bounded by2​B2B\. Summing overx∈Uix\\in U\_\{i\}yields

∑x∈Ui‖f¯k\(t\)​\(x\)−fi\(t,1\)​\(x\)‖1\\displaystyle\\sum\_\{x\\in U\_\{i\}\}\\bigl\\\|\\bar\{f\}\_\{k\}^\{\(t\)\}\(x\)\-f\_\{i\}^\{\(t,1\)\}\(x\)\\bigr\\\|\_\{1\}≤1\|𝒢k\|​∑j∈𝒢k∑x∈Ui‖fj\(t,1\)​\(x\)−fi\(t,1\)​\(x\)‖1\\displaystyle\\leq\\frac\{1\}\{\|\\mathcal\{G\}\_\{k\}\|\}\\sum\_\{j\\in\\mathcal\{G\}\_\{k\}\}\\sum\_\{x\\in U\_\{i\}\}\\bigl\\\|f\_\{j\}^\{\(t,1\)\}\(x\)\-f\_\{i\}^\{\(t,1\)\}\(x\)\\bigr\\\|\_\{1\}≤1\|𝒢k\|​∑j∈𝒢k2​B=2​B\.\\displaystyle\\leq\\frac\{1\}\{\|\\mathcal\{G\}\_\{k\}\|\}\\sum\_\{j\\in\\mathcal\{G\}\_\{k\}\}2B\\;=\\;2B\.Thus\|Bad\|≤2​B/γ\|\\mathrm\{Bad\}\|\\leq 2B/\\gamma\. Plugging into Equation \([3](https://arxiv.org/html/2605.11165#A2.E3)\) yields

Errx∈Ui​\(f¯k\(t\)\)≤Errx∈Ui​\(fi\(t,1\)\)\+1\|Ui\|⋅2​Bγ,\\mathrm\{Err\}\_\{x\\in U\_\{i\}\}\\\!\\bigl\(\\bar\{f\}\_\{k\}^\{\(t\)\}\\bigr\)\\;\\leq\\;\\mathrm\{Err\}\_\{x\\in U\_\{i\}\}\\\!\\bigl\(f\_\{i\}^\{\(t,1\)\}\\bigr\)\+\\frac\{1\}\{\|U\_\{i\}\|\}\\cdot\\frac\{2B\}\{\\gamma\},which establishes the desired bound\. ∎

### B\.2Proofs of Risk Bounds

To prove our main results, which are the clients’ and server’s risk bounds, we begin with the following key lemma\.

###### Lemma B\.1\(Stability under objective shift\)\.

Fix an iterationtt, and a clientii, letΦ\\Phidenote either the client\-side objectiveFiF\_\{i\}or the server\-side objectiveJiJ\_\{i\}in the surrogate version\. Letθ^\\hat\{\\theta\}represent the corresponding model minimizingΦ^\\hat\{\\Phi\}\(i\.e\.,fi\(t,2\)f\_\{i\}^\{\(t,2\)\}orhπ​\(i\)\(t\)h\_\{\\pi\(i\)\}^\{\(t\)\}\), andθ\\thetarepresentarg⁡minΘ⁡Φi​\(θ\)\\arg\\min\_\{\\Theta\}\\Phi\_\{i\}\(\\theta\)then

R𝒬i​\(θ^\)−R𝒬i​\(θ∗\)≤C​\(2​L1​\(τμ−L2\+4​εnμ\)\)α\.R\_\{\\mathcal\{Q\}\_\{i\}\}\(\\hat\{\\theta\}\)\-R\_\{\\mathcal\{Q\}\_\{i\}\}\(\\theta^\{\*\}\)\\;\\leq\\;C\\left\(2L\_\{1\}\(\\frac\{\\tau\}\{\\mu\-L\_\{2\}\}\+\\sqrt\{\\frac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\)\\right\)^\{\\alpha\}\.and

Ri​\(θ^\)−Ri​\(θ∗\)≤C​\(2​L1​\(τμ−L2\+4​εnμ\)\)α\.R\_\{i\}\(\\hat\{\\theta\}\)\-R\_\{i\}\(\\theta^\{\*\}\)\\;\\leq\\;C\\left\(2L\_\{1\}\(\\frac\{\\tau\}\{\\mu\-L\_\{2\}\}\+\\sqrt\{\\frac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\)\\right\)^\{\\alpha\}\.

###### Proof\.

Recall the empirical objective can be decomposed into supported and unsupported parts:Φ^​\(θ\)=Φ^i​\(θ\)\+Φ^ic​\(θ\)\\hat\{\\Phi\}\(\\theta\)=\\hat\{\\Phi\}\_\{i\}\(\\theta\)\+\\hat\{\\Phi\}\_\{i\}^\{c\}\(\\theta\)\. Letθ^1:=arg⁡minΘ⁡Φ^i​\(θ\)\\hat\{\\theta\}\_\{1\}:=\\arg\\min\_\{\\Theta\}\\hat\{\\Phi\}\_\{i\}\(\\theta\)\. Sinceθ^\\hat\{\\theta\}minimizesΦ^\\hat\{\\Phi\}, we have∇Φ^i​\(θ^\)\+∇Φ^ic​\(θ^\)=0\\nabla\\hat\{\\Phi\}\_\{i\}\(\\hat\{\\theta\}\)\+\\nabla\\hat\{\\Phi\}\_\{i\}^\{c\}\(\\hat\{\\theta\}\)=0\. Using the mean value theorem on∇Φ^i\\nabla\\hat\{\\Phi\}\_\{i\}betweenθ^\\hat\{\\theta\}andθ^1\\hat\{\\theta\}\_\{1\}:

∇Φ^i​\(θ^\)=∇2Φ^i​\(θ¯\)​\(θ^−θ^1\),\\nabla\\hat\{\\Phi\}\_\{i\}\(\\hat\{\\theta\}\)=\\nabla^\{2\}\\hat\{\\Phi\}\_\{i\}\(\\bar\{\\theta\}\)\(\\hat\{\\theta\}\-\\hat\{\\theta\}\_\{1\}\),for someθ¯\\bar\{\\theta\}on the line segment betweenθ^1\\hat\{\\theta\}\_\{1\}andθ^\\hat\{\\theta\}\.

Because∇Φ^i​\(θ^1\)=0\\nabla\\hat\{\\Phi\}\_\{i\}\(\\hat\{\\theta\}\_\{1\}\)=0andΦ^i\\hat\{\\Phi\}\_\{i\}isμ\\mu\-strongly convex, we obtain

‖θ^−θ^1‖\\displaystyle\\\|\\hat\{\\theta\}\-\\hat\{\\theta\}\_\{1\}\\\|=‖∇2Φ^i​\(θ¯\)−1​∇Φ^i​\(θ^\)‖\\displaystyle=\\\|\\nabla^\{2\}\\hat\{\\Phi\}\_\{i\}\(\\bar\{\\theta\}\)^\{\-1\}\\nabla\\hat\{\\Phi\}\_\{i\}\(\\hat\{\\theta\}\)\\\|≤1μ​‖∇Φ^i​\(θ^\)‖=1μ​‖∇Φ^ic​\(θ^\)‖\.\\displaystyle\\leq\\frac\{1\}\{\\mu\}\\\|\\nabla\\hat\{\\Phi\}\_\{i\}\(\\hat\{\\theta\}\)\\\|=\\frac\{1\}\{\\mu\}\\\|\\nabla\\hat\{\\Phi\}\_\{i\}^\{c\}\(\\hat\{\\theta\}\)\\\|\.
UsingL2L\_\{2\}\-smoothness ofΦ^ic\\hat\{\\Phi\}\_\{i\}^\{c\}, we bound

‖∇Φ^ic​\(θ^\)‖≤‖∇Φ^ic​\(θ^1\)‖\+L2​‖θ^−θ^1‖\.\\\|\\nabla\\hat\{\\Phi\}\_\{i\}^\{c\}\(\\hat\{\\theta\}\)\\\|\\leq\\\|\\nabla\\hat\{\\Phi\}\_\{i\}^\{c\}\(\\hat\{\\theta\}\_\{1\}\)\\\|\+L\_\{2\}\\\|\\hat\{\\theta\}\-\\hat\{\\theta\}\_\{1\}\\\|\.Rearranging,

‖θ^−θ^1‖≤1μ−L2​‖∇Φ^ic​\(θ^1\)‖≤τμ−L2\.\\\|\\hat\{\\theta\}\-\\hat\{\\theta\}\_\{1\}\\\|\\leq\\frac\{1\}\{\\mu\-L\_\{2\}\}\\\|\\nabla\\hat\{\\Phi\}\_\{i\}^\{c\}\(\\hat\{\\theta\}\_\{1\}\)\\\|\\leq\\frac\{\\tau\}\{\\mu\-L\_\{2\}\}\.
We now relateθ^1\\hat\{\\theta\}\_\{1\}toθ∗\\theta^\{\*\}\. By standard uniform convergence arguments \(e\.g\., Rademacher complexity\), with high probability1−δ1\-\\deltaover the training sample,

supθ\|Φ^i​\(θ\)−Φi​\(θ\)\|≤εn,where​εn∈O~​\(log⁡\(1/δ\)n\)\.\\sup\_\{\\theta\}\\left\|\\hat\{\\Phi\}\_\{i\}\(\\theta\)\-\\Phi\_\{i\}\(\\theta\)\\right\|\\leq\\varepsilon\_\{n\},\\quad\\text\{where \}\\varepsilon\_\{n\}\\in\\tilde\{O\}\(\\sqrt\{\\frac\{\\log\(1/\\delta\)\}\{n\}\}\)\.Using this, and the fact thatθ^1\\hat\{\\theta\}\_\{1\}minimizesΦ^i\\hat\{\\Phi\}\_\{i\}, we get:

Φi​\(θ^1\)≤Φ^i​\(θ^1\)\+εn≤Φ^i​\(θ∗\)\+εn≤Φi​\(θ∗\)\+2​εn\.\\Phi\_\{i\}\(\\hat\{\\theta\}\_\{1\}\)\\leq\\hat\{\\Phi\}\_\{i\}\(\\hat\{\\theta\}\_\{1\}\)\+\\varepsilon\_\{n\}\\leq\\hat\{\\Phi\}\_\{i\}\(\\theta^\{\*\}\)\+\\varepsilon\_\{n\}\\leq\\Phi\_\{i\}\(\\theta^\{\*\}\)\+2\\varepsilon\_\{n\}\.By strong convexity ofΦi\\Phi\_\{i\}, we obtain

‖θ^1−θ∗‖≤4​εnμ\.\\\|\\hat\{\\theta\}\_\{1\}\-\\theta^\{\*\}\\\|\\leq\\sqrt\{\\frac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\.Then by triangle inequality:

‖θ^−θ∗‖≤‖θ^−θ^1‖\+‖θ^1−θ∗‖≤τμ−L2\+4​εnμ\.\\\|\\hat\{\\theta\}\-\\theta^\{\*\}\\\|\\leq\\\|\\hat\{\\theta\}\-\\hat\{\\theta\}\_\{1\}\\\|\+\\\|\\hat\{\\theta\}\_\{1\}\-\\theta^\{\*\}\\\|\\leq\\frac\{\\tau\}\{\\mu\-L\_\{2\}\}\+\\sqrt\{\\frac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\.
Now consider the pseudolabels generated byθ^\\hat\{\\theta\}andθ∗\\theta^\{\*\}ass^,s\\hat\{s\},srespectively\. Letc⋆c^\{\\star\}be the unique maximizer ofss\. If‖s^−s‖∞≤ε\\\|\\hat\{s\}\-s\\\|\_\{\\infty\}\\leq\\varepsilon, then\|s^j−sj\|≤ε\|\\hat\{s\}\_\{j\}\-s\_\{j\}\|\\leq\\varepsilonfor allj∈\[M\]j\\in\[M\]\. In particular,s^c⋆≥sc⋆−ε\\hat\{s\}\_\{c^\{\\star\}\}\\geq s\_\{c^\{\\star\}\}\-\\varepsilonand, for anyj≠c⋆j\\neq c^\{\\star\},s^j≤sj\+ε\\hat\{s\}\_\{j\}\\leq s\_\{j\}\+\\varepsilon\. Hence for eachj≠c⋆j\\neq c^\{\\star\},

s^c⋆−s^j≥\(sc⋆−ε\)−\(sj\+ε\)=\(sc⋆−sj\)−2​ε\.\\hat\{s\}\_\{c^\{\\star\}\}\-\\hat\{s\}\_\{j\}\\;\\geq\\;\(s\_\{c^\{\\star\}\}\-\\varepsilon\)\-\(s\_\{j\}\+\\varepsilon\)\\;=\\;\(s\_\{c^\{\\star\}\}\-s\_\{j\}\)\-2\\varepsilon\.
We can rank theMMscores for inputxxbyssass\(1\)\>s\(2\)\>⋯\>s\(M\)\.s\_\{\(1\)\}\>s\_\{\(2\)\}\>\\dots\>s\_\{\(M\)\}\.Sincesj≤s\(2\)s\_\{j\}\\leq s\_\{\(2\)\}for allj≠c⋆j\\neq c^\{\\star\}andsc⋆=s\(1\)s\_\{c^\{\\star\}\}=s\_\{\(1\)\}, we havesc⋆−sj≥s\(1\)−s\(2\)=Δ​\(s\)s\_\{c^\{\\star\}\}\-s\_\{j\}\\geq s\_\{\(1\)\}\-s\_\{\(2\)\}=\\Delta\(s\), and thus

s^c⋆−s^j≥Δ​\(s\)−2​ε\.\\hat\{s\}\_\{c^\{\\star\}\}\-\\hat\{s\}\_\{j\}\\;\\geq\\;\\Delta\(s\)\-2\\varepsilon\.IfΔ​\(s\)\>2​ε\\Delta\(s\)\>2\\varepsilon, thens^c⋆−s^j\>0\\hat\{s\}\_\{c^\{\\star\}\}\-\\hat\{s\}\_\{j\}\>0for allj≠c⋆j\\neq c^\{\\star\}, soc⋆c^\{\\star\}remains the unique maximizer ofs^\\hat\{s\}\.

Then applying the condition

supx∈𝒳‖sθ^​\(x\)−sθ∗​\(x\)‖∞≤L1​‖θ^−θ∗‖\.\\sup\_\{x\\in\\mathcal\{X\}\}\\big\\\|s\_\{\\hat\{\\theta\}\}\(x\)\-s\_\{\\theta^\{\*\}\}\(x\)\\big\\\|\_\{\\infty\}\\;\\leq\\;L\_\{1\}\\,\\\|\\hat\{\\theta\}\-\\theta^\{\*\}\\\|\.to the boundε\\varepsilon, we have

\|R​\(θ^\)−R​\(θ⋆\)\|≤\\displaystyle\|R\(\\hat\{\\theta\}\)\-R\(\\theta^\{\\star\}\)\|\\;\\leqℙ​\(hθ^​\(x\)≠hθ⋆​\(x\)\)\\displaystyle\\;\\mathbb\{P\}\\big\(h\_\{\\hat\{\\theta\}\}\(x\)\\neq h\_\{\\theta^\{\\star\}\}\(x\)\\big\)≤\\displaystyle\\;\\leqℙ​\(Δ​\(s\)≤2​L1​\(τμ−L2\+4​εnμ\)\)\\displaystyle\\;\\mathbb\{P\}\\\!\\left\(\\Delta\(s\)\\leq 2L\_\{1\}\(\\frac\{\\tau\}\{\\mu\-L\_\{2\}\}\+\\sqrt\{\\frac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\)\\right\)≤\\displaystyle\\;\\leqC​\(2​L1​\(τμ−L2\+4​εnμ\)\)α\.\\displaystyle\\;C\\left\(2L\_\{1\}\(\\frac\{\\tau\}\{\\mu\-L\_\{2\}\}\+\\sqrt\{\\frac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\)\\right\)^\{\\alpha\}\.
∎

Next we prove the risk bounds, restating the theorems for completeness\.

###### of Thm\.[5\.11](https://arxiv.org/html/2605.11165#S5.Thmtheorem11)\.

SinceQiQ\_\{i\}andDiD\_\{i\}has bounded importance weights,

1=∫Sipi​\(x\)​𝑑x≤supx∈Sipi​\(x\)q​\(x\)​∫Siq​\(x\)​𝑑x=wi,1​Q​\(Si\),1\\;=\\;\\int\_\{S\_\{i\}\}p\_\{i\}\(x\)\\,dx\\;\\leq\\;\\sup\_\{x\\in S\_\{i\}\}\\frac\{p\_\{i\}\(x\)\}\{q\(x\)\}\\int\_\{S\_\{i\}\}q\(x\)\\,dx\\;=\\;w\_\{i,1\}\\,Q\(S\_\{i\}\),soQ​\(Si\)≥1/wi,1\>0Q\(S\_\{i\}\)\\geq 1/w\_\{i,1\}\>0\. Applying a multiplicative Chernoff bound on data drawn i\.i\.d\. fromQQ, we have\|Ui\|∈O​\(Q​\(Si\)​n\)=O​\(n\)\|U\_\{i\}\|\\in O\(Q\(S\_\{i\}\)n\)=O\(n\)with probability at least1−exp⁡\(−n/8​wi,1\)1\-\\exp\(\-n/8w\_\{i,1\}\)\. And Lemma[5\.2](https://arxiv.org/html/2605.11165#S5.Thmtheorem2)the aggregated pseudolabel modelf¯k\(t\)\\bar\{f\}\_\{k\}^\{\(t\)\}yields:

ErrUi​\(f¯k\(t\)\)≤ErrUi​\(fi\(t,1\)\)\+ε1,where​ε1∈O~​\(B​n−1\)\.\\mathrm\{Err\}\_\{U\_\{i\}\}\(\\bar\{f\}\_\{k\}^\{\(t\)\}\)\\;\\leq\\;\\mathrm\{Err\}\_\{\{U\_\{i\}\}\}\(f\_\{i\}^\{\(t,1\)\}\)\+\\varepsilon\_\{1\},\\text\{where \}\\varepsilon\_\{1\}\\in\\tilde\{O\}\(Bn^\{\-1\}\)\.
A standard Rademacher complexity bound then gives, with high probability:

R𝒬i​\(f¯k\(t\)\)≤R𝒬i​\(fi\(t,1\)\)\+ε1\+ε2,where​ε2∈O~​\(n−1/2\)\.R\_\{\\mathcal\{Q\}\_\{i\}\}\(\\bar\{f\}\_\{k\}^\{\(t\)\}\)\\;\\leq\\;R\_\{\\mathcal\{Q\}\_\{i\}\}\(f\_\{i\}^\{\(t,1\)\}\)\+\\varepsilon\_\{1\}\+\\varepsilon\_\{2\},\\quad\\text\{where \}\\varepsilon\_\{2\}\\in\\tilde\{O\}\(n^\{\-1/2\}\)\.
Leth∗h^\{\*\}be the optimizer forJiJ\_\{i\}\. Applying Theorem 4\.3 of\[[31](https://arxiv.org/html/2605.11165#bib.bib53)\], we obtain:

R𝒬i​\(h∗\)\\displaystyle R\_\{\\mathcal\{Q\}\_\{i\}\}\(h^\{\*\}\)≤2c−1​R𝒬i​\(f¯k\(t\)\)\+2​c​ρc−1\\displaystyle\\leq\\frac\{2\}\{c\-1\}\\,R\_\{\\mathcal\{Q\}\_\{i\}\}\(\\bar\{f\}\_\{k\}^\{\(t\)\}\)\+\\frac\{2c\\rho\}\{c\-1\}≤2c−1​\(R𝒬i​\(fi\(t,1\)\)\+ε1\+ε2\)\+2​c​ρc−1\.\\displaystyle\\leq\\frac\{2\}\{c\-1\}\\left\(R\_\{\\mathcal\{Q\}\_\{i\}\}\(f\_\{i\}^\{\(t,1\)\}\)\+\\varepsilon\_\{1\}\+\\varepsilon\_\{2\}\\right\)\+\\frac\{2c\\rho\}\{c\-1\}\.
Using the importance weighting trick and bounded density ratios, we have:

Ri​\(hk\(t\)\)\\displaystyle R\_\{i\}\\\!\\left\(h\_\{k\}^\{\(t\)\}\\right\)=𝔼x∼𝒟i​\[ℓ0​\-​1​\(hk\(t\)​\(x\),y​\(x\)\)\]\\displaystyle=\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\_\{i\}\}\\\!\\left\[\\ell\_\{0\\text\{\-\}1\}\\\!\\left\(h\_\{k\}^\{\(t\)\}\(x\),\\,y\(x\)\\right\)\\right\]=𝔼x∼𝒬i​\[p​\(x\)q​\(x\)​ℓ0​\-​1​\(hk\(t\)​\(x\),y​\(x\)\)\]\\displaystyle=\\mathbb\{E\}\_\{x\\sim\\mathcal\{Q\}\_\{i\}\}\\\!\\left\[\\tfrac\{p\(x\)\}\{q\(x\)\}\\,\\ell\_\{0\\text\{\-\}1\}\\\!\\left\(h\_\{k\}^\{\(t\)\}\(x\),\\,y\(x\)\\right\)\\right\]≤w1​R𝒬i​\(hk\(t\)\)\\displaystyle\\leq w\_\{1\}\\,R\_\{\\mathcal\{Q\}\_\{i\}\}\\\!\\left\(h\_\{k\}^\{\(t\)\}\\right\)
By Lemma[B\.1](https://arxiv.org/html/2605.11165#A2.Thmtheorem1),

Ri​\(hk\(t\)\)≤w1​\(R𝒬i​\(h∗\)\+C​\(2​L1​\(τμ−L2\+4​εnμ\)\)α\)\\displaystyle R\_\{i\}\\\!\\left\(h\_\{k\}^\{\(t\)\}\\right\)\\leq w\_\{1\}\\biggl\(R\_\{\\mathcal\{Q\}\_\{i\}\}\(h^\{\*\}\)\+C\\Bigl\(2L\_\{1\}\\Bigl\(\\tfrac\{\\tau\}\{\\mu\-L\_\{2\}\}\+\\sqrt\{\\tfrac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\\Bigr\)\\Bigr\)^\{\\alpha\}\\biggr\)≤w1​\(2c−1​R𝒬i​\(fi\(t,1\)\)\+2​\(ε1\+ε2\+c​γ\)c−1\)\\displaystyle\\leq w\_\{1\}\\left\(\\tfrac\{2\}\{c\-1\}\\,R\_\{\\mathcal\{Q\}\_\{i\}\}\\\!\\left\(f\_\{i\}^\{\(t,1\)\}\\right\)\+\\tfrac\{2\(\\varepsilon\_\{1\}\+\\varepsilon\_\{2\}\+c\\gamma\)\}\{c\-1\}\\right\)\+w1​C​\(2​L1​\(τμ−L2\+4​εnμ\)\)α\\displaystyle\\qquad\+w\_\{1\}\\,C\\Bigl\(2L\_\{1\}\\Bigl\(\\tfrac\{\\tau\}\{\\mu\-L\_\{2\}\}\+\\sqrt\{\\tfrac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\\Bigr\)\\Bigr\)^\{\\alpha\}=2​w1c−1​𝔼x∼𝒟i​\[q​\(x\)p​\(x\)​ℓ​\(fi\(t,1\)​\(x\),y​\(x\)\)\]\\displaystyle=\\tfrac\{2w\_\{1\}\}\{c\-1\}\\,\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\_\{i\}\}\\\!\\left\[\\tfrac\{q\(x\)\}\{p\(x\)\}\\,\\ell\\\!\\left\(f\_\{i\}^\{\(t,1\)\}\(x\),\\,y\(x\)\\right\)\\right\]\+2​w1​\(ε1\+ε2\+c​γ\)c−1\+w1​C​\(2​L1​\(τμ−L2\+4​εnμ\)\)α\\displaystyle\+\\tfrac\{2w\_\{1\}\(\\varepsilon\_\{1\}\+\\varepsilon\_\{2\}\+c\\gamma\)\}\{c\-1\}\+w\_\{1\}\\,C\\Bigl\(2L\_\{1\}\\Bigl\(\\tfrac\{\\tau\}\{\\mu\-L\_\{2\}\}\+\\sqrt\{\\tfrac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\\Bigr\)\\Bigr\)^\{\\alpha\}≤2​w1​w2c−1​Ri​\(fi\(t,1\)\)\+2​w1​\(ε1\+ε2\+c​γ\)c−1\\displaystyle\\leq\\tfrac\{2w\_\{1\}w\_\{2\}\}\{c\-1\}\\,R\_\{i\}\\\!\\left\(f\_\{i\}^\{\(t,1\)\}\\right\)\+\\tfrac\{2w\_\{1\}\(\\varepsilon\_\{1\}\+\\varepsilon\_\{2\}\+c\\gamma\)\}\{c\-1\}\+w1​C​\(2​L1​\(τμ−L2\+4​εnμ\)\)α\.\\displaystyle\+w\_\{1\}\\,C\\Bigl\(2L\_\{1\}\\Bigl\(\\tfrac\{\\tau\}\{\\mu\-L\_\{2\}\}\+\\sqrt\{\\tfrac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\\Bigr\)\\Bigr\)^\{\\alpha\}\.
By assumption,κ1:=2​w1​w2c−1<1\\kappa\_\{1\}:=\\tfrac\{2w\_\{1\}w\_\{2\}\}\{c\-1\}<1and the claim follows\. ∎

###### of Thm\.[5\.12](https://arxiv.org/html/2605.11165#S5.Thmtheorem12)\.

By Lemma[B\.1](https://arxiv.org/html/2605.11165#A2.Thmtheorem1), with probability at least1−δ1\-\\delta, the difference of the true risk between the obtained model, and the minimizerf∗f^\{\*\}forℱi,supp\\mathcal\{F\}\_\{i,\\mathrm\{supp\}\}is bounded as:

Ri​\(fi\(t,2\)\)−Ri​\(f∗\)≤C​\(2​L1​\(τμ−L2\+4​εnμ\)\)α\.R\_\{i\}\(f\_\{i\}^\{\(t,2\)\}\)\-R\_\{i\}\(f^\{\*\}\)\\;\\leq\\;C\\left\(2L\_\{1\}\(\\frac\{\\tau\}\{\\mu\-L\_\{2\}\}\+\\sqrt\{\\frac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\)\\right\)^\{\\alpha\}\.
Then by\[[31](https://arxiv.org/html/2605.11165#bib.bib53), Theorem 4\.3\],

Ri​\(fi\(t,2\)\)\\displaystyle R\_\{i\}\\\!\\left\(f\_\{i\}^\{\(t,2\)\}\\right\)≤Ri​\(f∗\)\+C​\(2​L1​\(τμ−L2\+4​εnμ\)\)α\\displaystyle\\leq R\_\{i\}\(f^\{\*\}\)\+C\\Bigl\(2L\_\{1\}\\Bigl\(\\tfrac\{\\tau\}\{\\mu\-L\_\{2\}\}\+\\sqrt\{\\tfrac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\\Bigr\)\\Bigr\)^\{\\alpha\}≤2c−1​Ri​\(hk\(t\)\)\+2​c​ρc−1\+C​\(2​L1​\(τμ−L2\+4​εnμ\)\)α\\displaystyle\\leq\\tfrac\{2\}\{c\-1\}\\,R\_\{i\}\\\!\\left\(h\_\{k\}^\{\(t\)\}\\right\)\+\\tfrac\{2c\\rho\}\{c\-1\}\+C\\Bigl\(2L\_\{1\}\\Bigl\(\\tfrac\{\\tau\}\{\\mu\-L\_\{2\}\}\+\\sqrt\{\\tfrac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\\Bigr\)\\Bigr\)^\{\\alpha\}≤2c−1​\(κ​Ri​\(fi\(t,1\)\)\+2​w1​\(ε1\+ε2\+c​ρ\)c−1\)\+2​c​ρc−1\\displaystyle\\leq\\tfrac\{2\}\{c\-1\}\\Bigl\(\\kappa\\,R\_\{i\}\\\!\\left\(f\_\{i\}^\{\(t,1\)\}\\right\)\+\\tfrac\{2w\_\{1\}\(\\varepsilon\_\{1\}\+\\varepsilon\_\{2\}\+c\\rho\)\}\{c\-1\}\\Bigr\)\+\\tfrac\{2c\\rho\}\{c\-1\}\+C​\(2​L1​\(τμ−L2\+4​εnμ\)\)α\(Theorem[5\.11](https://arxiv.org/html/2605.11165#S5.Thmtheorem11)\)\\displaystyle\\qquad\{\}\+C\\Bigl\(2L\_\{1\}\\Bigl\(\\tfrac\{\\tau\}\{\\mu\-L\_\{2\}\}\+\\sqrt\{\\tfrac\{4\\varepsilon\_\{n\}\}\{\\mu\}\}\\Bigr\)\\Bigr\)^\{\\alpha\}\\qquad\\text\{\(Theorem~\\ref\{thm:server\_improvement\}\)\}=4​w1​w2\(c−1\)2​Ri​\(fi\(t,1\)\)\+O~​\(n−1/2\)\\displaystyle=\\tfrac\{4w\_\{1\}w\_\{2\}\}\{\(c\-1\)^\{2\}\}\\,R\_\{i\}\\\!\\left\(f\_\{i\}^\{\(t,1\)\}\\right\)\+\\tilde\{O\}\\\!\\left\(n^\{\-1/2\}\\right\)≤4​w1​w2\(c−1\)2​Ri​\(fi\(t−1,2\)\)\+O~​\(n−1/2\)\.\\displaystyle\\leq\\tfrac\{4w\_\{1\}w\_\{2\}\}\{\(c\-1\)^\{2\}\}\\,R\_\{i\}\\\!\\left\(f\_\{i\}^\{\(t\-1,2\)\}\\right\)\+\\tilde\{O\}\\\!\\left\(n^\{\-1/2\}\\right\)\.The last step is justified as the only difference betweenfi\(t−1,2\)f\_\{i\}^\{\(t\-1,2\)\}andfi\(t,1\)f\_\{i\}^\{\(t,1\)\}is that the clientiicontinues to be trained over accurate real data, and by the uniform convergence via Rademacher complexity, and that the local training doesn’t worsen the empirical risk\.

Finally, defineκ2=4​w1​w2\(c−1\)2\\kappa\_\{2\}=\\tfrac\{4w\_\{1\}w\_\{2\}\}\{\(c\-1\)^\{2\}\}and letζ∈O~​\(n−1/2\)\\zeta\\in\\tilde\{O\}\(n^\{\-1/2\}\)\. By the union bound,Ri​\(fi\(T,2\)\)≤κ2T−1​Ri​\(fi\(1,2\)\)\+ζ​∑m=0T−1κ2m≤κ2T​Ri​\(fi\(1,1\)\)\+ζ​∑m=0Tκ2m=κ2T​Ri​\(fi\(1,2\)\)\+ζ​\(1−κ2T\)1−κ2\.R\_\{i\}\(f\_\{i\}^\{\(T,2\)\}\)\\;\\leq\\;\\kappa\_\{2\}^\{T\-1\}\\,R\_\{i\}\(f\_\{i\}^\{\(1,2\)\}\)\\;\+\\;\\zeta\\sum\_\{m=0\}^\{T\-1\}\\kappa\_\{2\}^\{m\}\\leq\\kappa\_\{2\}^\{T\}R\_\{i\}\(f\_\{i\}^\{\(1,1\)\}\)\+\\zeta\\sum\_\{m=0\}^\{T\}\\kappa\_\{2\}^\{m\}\\;=\\;\\kappa\_\{2\}^\{T\}\\,R\_\{i\}\(f\_\{i\}^\{\(1,2\)\}\)\\;\+\\;\\frac\{\\zeta\\bigl\(1\-\\kappa\_\{2\}^\{T\}\\bigr\)\}\{1\-\\kappa\_\{2\}\}\.Hiding the constants, we complete the proof\. ∎

## Appendix CAdditional Experimental Details

#### Overall Setup:

Our development environment was based on Python \(version 3\.10\) and PyTorch \(version 2\.5\.0\+cu121\), with GPU acceleration provided by two NVIDIA GeForce RTX 4090 GPUs\.111The source code is available at:[https://anonymous\.4open\.science/r/cluster\_fed\_learning\-5D71/](https://anonymous.4open.science/r/cluster_fed_learning-5D71/)All experiments follow a standard cross\-device federated learning setting, where a central server coordinates training across many clients with disjoint local datasets\.

We measure performance using the classification accuracy of each client on a test set sampled from its own local data distribution, ensuring that evaluation aligns with the personalized nature of the task\. For COSMOS, the prediction model used for clientiiis the*cluster\-specific server model*hπ​\(i\)h\_\{\\pi\(i\)\}, which represents the distilled knowledge from its assigned cluster\. For baseline methods that do not employ a server\-side model \(e\.g\., FedMD, FedCT, COMET, and homogeneous PFL methods\), accuracy is computed directly using the client’s own trained model\. This evaluation protocol ensures that all methods are compared according to the model actually deployed to each client\.

#### Resource Requirements:

Our experiments were conducted on a server running Ubuntu20\.04\.620\.04\.6LTS, with the development environment based on Python3\.103\.10and PyTorch2\.5\.0\+c​u​1212\.5\.0\+cu121\. The system was equipped with22x NVIDIA GeForce RTX40904090GPUs,9494GB of RAM, and2\.02\.0GB of swap memory\.

#### Time Requirements:

When running COSMOS with a VGG architecture using55clusters, training models for one communication round took approximately2020minutes\. A full experiment \(spanning2020rounds and2525clients\) typically required77hours\.

For benchmarks with lower complexity \(e\.g\., AlexNets\), training models for one communication round took approximately55minutes, and a full experiment \(spanning2020rounds and2525clients\) typically required1\.51\.5hours\.

#### Data Sources and Details:

1. 1\.Tiny ImageNet:The Tiny ImageNet dataset is downloaded from the official Stanford CS231n website\. It is provided as a ZIP file containing 200 categories of images, used for classification tasks\. The dataset is extracted and reorganized into appropriate folders for training and validation\. The dataset is available at:[http://cs231n\.stanford\.edu/tiny\-imagenet\-200\.zip](http://cs231n.stanford.edu/tiny-imagenet-200.zip)\.
2. 2\.CIFAR\-10 & CIFAR\-100:These datasets consist of 60,000 32x32 color images in 10 and 100 classes, respectively\. They are commonly used for training machine learning models for image classification\. The datasets are available for download directly via the torchvision library\. More information on these datasets can be found at:[https://www\.cs\.toronto\.edu/˜kriz/cifar\.html](https://www.cs.toronto.edu/~kriz/cifar.html)\.
3. 3\.Extended MNIST \(EMNIST\):The EMNIST dataset extends the original MNIST dataset to include handwritten letters\. We specifically use theBalancedversion, which includes 131,600 characters across 47 balanced classes\. The dataset is available for download via the torchvision library\. Additional details can be found on the homepage:[https://www\.nist\.gov/itl/products\-and\-services/emnist\-dataset](https://www.nist.gov/itl/products-and-services/emnist-dataset)\.

#### Hyper Parameters and Optimizers:

1. 1\.Loss Functions:During training using pseudo\-labels, we use the Kullback\-Leibler Divergence \(nn\.KLDivLoss\) to align model outputs with mean pseudo\-labels, with reduction=’batchmean’\. During client fine\-tuning \(using their local data true labels\), we use Cross\-Entropy Loss \(nn\.CrossEntropyLoss\)\.
2. 2\.Optimizers and Learning Rate:We use the Adam optimizer for both training stages\. After performing hyperparameter optimization, we use a higher learning rate of0\.0010\.001when training on client\-side data with true labels\. This enables the model to adapt quickly and effectively to reliable supervision\. In contrast, for pseudo\-labeled data, we reduce the learning rate to0\.00010\.0001to account for the inherent noise and uncertainty in these labels\.
3. 3\.Weights Initialization:For consistency across experiments, we initialize weights using Kaiming He initialization for both convolutional and linear layers\. This method, which is well\-suited for ReLU activations, helps prevent vanishing or exploding gradients\. The biases are initialized to zeros\. The seed is updated and applied to ensure reproducibility across runs\.
4. 4\.COSMOS Clustering Hyperparameter \(BB\):COSMOS uses a neighborhood threshold hyperparameterBBto determine client similarity during the greedy\-elimination clustering step\. We tuneBBsuch that the resulting number of clusters is approximatelyK=5K=5, matching the 20% class\-grouping structure used to induce heterogeneous label distributions across clients\.

Table 2:Final\-round performance on benchmark datasets \(mean±\\pmSEM\) under Dirichlet data heterogeneity withα=5\\alpha=5\. Rows marked with†denote that COSMOS’s SEM error bar does not overlap with any baseline \(COMET, FedCT, or FedMD\)\. Figure[2](https://arxiv.org/html/2605.11165#S5.F2)presents performance across communication rounds\.Table 3:Final\-round performance under different client architecture mixtures \(mean±\\pmSEM\)\. Rows marked with†indicate COSMOS’s SEM interval does not overlap with*any*baseline\. Figure[3](https://arxiv.org/html/2605.11165#S5.F3)presents performance across communication rounds\.Table 4:Comparison of Baselines by Communication Type \(Bandwidth Measured Per Client in MB\)\.

## Appendix DAdditional Results

This appendix provides supplementary empirical results that support and extend the main findings in Section[6](https://arxiv.org/html/2605.11165#S6)\. The results offer additional evidence regarding robustness, scalability, and design choices in COSMOS\.

#### Effect of Data Heterogeneity:

Table 5:Final\-round performance on benchmark datasets \(mean±\\pmSEM\) under Dirichlet data heterogeneity withα=1\\alpha=1\. Rows marked with†denote that COSMOS’s SEM error bar does not overlap with any baseline \(COMET, FedCT, or FedMD\)\. Figure[4](https://arxiv.org/html/2605.11165#A4.F4)presents performance across communication rounds\.Table 6:Final\-round performance on benchmark datasets \(mean±\\pmSEM\) under Dirichlet data heterogeneity withα=100\\alpha=100\. Rows marked with†denote that COSMOS’s SEM error bar does not overlap with any baseline \(COMET, FedCT, or FedMD\)\. Figure[5](https://arxiv.org/html/2605.11165#A4.F5)\. presents performance across communication rounds\.![Refer to caption](https://arxiv.org/html/2605.11165v1/x3.png)Figure 4:Comparison of COSMOS and heterogeneous\-model baselines across four benchmarks under Dirichlet non\-IID sampling \(α=1\\alpha=1\)\. Client models consist of MobileNet and SqueezeNet architectures\. Curves show mean client accuracy over rounds \(Top\-1 for CIFAR\-10/100 and EMNIST; Top\-5 for Tiny ImageNet\)\. A statistical comparison of final\-round performance is provided in Table[5](https://arxiv.org/html/2605.11165#A4.T5)\.![Refer to caption](https://arxiv.org/html/2605.11165v1/x4.png)Figure 5:Comparison of COSMOS and heterogeneous\-model baselines across four benchmarks under Dirichlet non\-IID sampling \(α=100\\alpha=100\)\. Client models consist of MobileNet and SqueezeNet architectures\. Curves show mean client accuracy over rounds \(Top\-1 for CIFAR\-10/100 and EMNIST; Top\-5 for Tiny ImageNet\)\. A statistical comparison of final\-round performance is provided in Table[6](https://arxiv.org/html/2605.11165#A4.T6)\.In Figures[4](https://arxiv.org/html/2605.11165#A4.F4)and[5](https://arxiv.org/html/2605.11165#A4.F5), as well as Tables[5](https://arxiv.org/html/2605.11165#A4.T5)and[6](https://arxiv.org/html/2605.11165#A4.T6), we examine the impact of data heterogeneity by varying the Dirichlet parameterα\\alpha\. Across all datasets, COSMOS consistently outperforms model\-agnostic baselines for all evaluated values ofα\\alpha\. While the absolute performance of all methods degrades as heterogeneity increases, the relative advantage of COSMOS becomes more pronounced in highly non\-IID regimes\.

#### Comparison in Setting with Homogeneous Clients:

Table 7:Final\-round performance under homogeneous client architectures \(mean±\\pmSEM\)\. Best result per architecture in bold\.![Refer to caption](https://arxiv.org/html/2605.11165v1/x5.png)Figure 6:Top\-1 accuracy over communication rounds for homogeneous client architectures \(AlexNet, ResNet, MobileNet, and SqueezeNet\) on CIFAR\-100 with Dirichlet non\-IID partitioning \(α=5\\alpha=5\)\. Each curve shows the mean accuracy across seeds for a variety of algorithms that do not require clients to be model\-agnostic: COMET, Ditto, FedBABU, FedMD, FedSelect, COSMOS, pFedCK, and pFedHN\. A statistical summary of the final communication round appears in Table[7](https://arxiv.org/html/2605.11165#A4.T7)\.In Table[7](https://arxiv.org/html/2605.11165#A4.T7)and Figure[6](https://arxiv.org/html/2605.11165#A4.F6), we relax the model\-agnostic constraint and compare COSMOS to additional baselines in a homogeneous setting where all clients share the same architecture\. COSMOS remains competitive with state\-of\-the\-art personalized FL methods, while no single baseline dominates across all datasets and architectures\. This shows that COSMOS’s model\-agnostic design does not come at the cost of performance when homogeneity is available\.

#### Additional Comparisons across Various Settings:

![Refer to caption](https://arxiv.org/html/2605.11165v1/x6.png)Figure 7:Top\-1 accuracy over communication rounds for COSMOS on CIFAR\-100 with Dirichlet non\-IID sampling \(α=5\\alpha=5\): \(a\) varying numbers of participating clients, \(b\) different ratios \(fractions\) of unlabeled global data, and \(c\) different server\-side architectures\. Client models are MobileNet and SqueezeNet\.Table 8:Final\-round performance of COSMOS under: \(a\) varying numbers of participating clients, \(b\) different ratios \(fractions\) of unlabeled global data, and \(c\) different server\-side architectures\.We next evaluate the performance of COSMOS under different settings in Figure[7](https://arxiv.org/html/2605.11165#A4.F7)and Table[8](https://arxiv.org/html/2605.11165#A4.T8)by varying \(a\) the number of participating clients, \(b\) the size of the unlabeled public datasetUU, and \(c\) the server\-side model architecture\. As the number of clients increases, COSMOS continues to improve, benefiting from increased pseudo\-label diversity at the server\. Similarly, increasing\|U\|\|U\|results in consistent performance gains\. Figure[7](https://arxiv.org/html/2605.11165#A4.F7)\(c\) shows that COSMOS further benefits from stronger server\-side models\.

#### Comparison with Different Numbers of Clusters:

![Refer to caption](https://arxiv.org/html/2605.11165v1/x7.png)Figure 8:Top\-1 accuracy over communication rounds for COSMOS showing the performance on CIFAR\-100 with Dirichlet non\-IID sampling \(α=5\\alpha=5\) under varying number of clusters \(i\.e\., values ofβ\\betashown in the legend\), Client models are AlexNet\.In Figure[8](https://arxiv.org/html/2605.11165#A4.F8), we study the effect of tuningβ\\beta, where different values correspond to different numbers of clusters\. We observe that performance saturates once the number of clusters reaches approximately five, with no further improvement beyond this point\. In contrast, larger values ofβ\\beta\(resulting in fewer clusters\) lead to degraded performance, likely because grouping too many clients together reduces data similarity within each cluster\.

![Refer to caption](https://arxiv.org/html/2605.11165v1/x8.png)Figure 9:Top\-1 accuracy over communication rounds for COSMOS\. Subfigure \(a\) shows performance under different values ofλ\\lambda\. Subfigure \(b\) shows performance on CIFAR\-100 with Dirichlet non\-IID sampling \(α=5\\alpha=5\) under varying distillation temperatures and different statistical aggregation rules \(mean, median, max\)\. Client models are MobileNet and SqueezeNet\.
#### Comparison with Different Hyperparameter Values:

Lastly, in Figure[9](https://arxiv.org/html/2605.11165#A4.F9), we examine the performance of COSMOS under different parameter settings\. In subfigure \(a\), we vary the value ofλ\\lambda, while in subfigure \(b\) we evaluate different pseudo\-label aggregation functions within each cluster, as well as the effect of pseudo\-label temperature \(with temperature0corresponding to hard labels\)\. COSMOS is robust to moderate variations inλ\\lambda, with performance remaining stable across a wide range of values andλ=5\\lambda=5yielding the best overall results in our experiments\. We also observe thatT=1T=1performs best across datasets, while larger temperatures tend to oversmooth pseudo\-labels and slightly degrade performance\. This indicates that preserving confident predictions during pseudo\-label exchange is beneficial\. Additionally, mean aggregation consistently achieves the strongest performance among the evaluated aggregation strategies\.

For completeness, we provide statistical comparisons of final\-round performance using mean accuracy±\\pmSEM across three random seeds\. In most evaluated settings, COSMOS exhibits non\-overlapping error bars with model\-agnostic baselines, providing strong evidence of consistent performance improvements\.