Technical Summary: COSMOS

Problem Statement

Federated Learning (FL) faces significant challenges in heterogeneous environments where clients differ not only in data distribution (statistical heterogeneity) but also in model architecture (architectural heterogeneity). While Personalized Federated Learning (PFL) methods have been developed to tailor models to individual clients, most existing approaches assume structural compatibility or require clients to disclose proprietary model architectures. This creates a practical barrier, as clients often wish to utilize arbitrary or proprietary models without revealing parameters or gradients.

Current model-agnostic FL approaches, such as FedMD and COMET, rely on output-level communication (e.g., soft or hard labels) using a shared unlabeled dataset. However, these methods suffer from specific limitations:

1. Passive Server Coordination: They often treat the server as a passive coordinator, failing to leverage the server's computational resources for active personalization.
2. Heuristic Clustering: Methods like COMET rely on K-means clustering, which requires prior knowledge of the cluster count ($K$) and lacks formal mechanisms for personalization guarantees when client diversity is high.
3. Lack of Theoretical Guarantees: Existing model-agnostic PFL literature generally lacks rigorous theoretical guarantees regarding risk contraction, often relying on simplified assumptions (e.g., linear models, Gaussian distributions) or providing only standard convergence-to-stationarity results.

Methodology: COSMOS Framework

The authors propose COSMOS (Clustered Output-based Server Models), a model-agnostic PFL framework that enables server-side personalization using only pseudo-label communication. The framework operates under the assumption that clients possess private labeled data and a shared, unlabeled public dataset $U$ (e.g., web-scraped or synthetic data) is available.

COSMOS proceeds in two distinct phases:

1. Pre-Training and Clustering (PTC) Phase

This phase establishes the initial client-server mapping and cluster-specific models:

• Local Pre-Training: Clients train their local models $f_i$ on their private data $D_i$ and generate pseudo-labels on the public dataset $U$.
• Greedy Clustering: The server clusters clients based on the similarity of their pseudo-labels on $U$. Unlike K-means, COSMOS employs a greedy clustering algorithm that does not require $K$ to be specified in advance. It iteratively selects a client with the most "close" neighbors (within a distance threshold $B_0$) to form a cluster until all clients are assigned. This process is endogenous to the training procedure.
• Server-Side Distillation: For each cluster $G_k$, the server aggregates the pseudo-labels from its member clients to create a mean pseudo-label distribution. It then trains a dedicated "teacher" model $h_k$ for that cluster using these aggregated labels and a specific regularization term to ensure local smoothness and robustness.
• Client-Side Distillation: Clients receive the pseudo-labels from their assigned cluster's teacher model and fine-tune their local models using these labels.

2. Iterative Federated Fine-Tuning (IFFT) Phase

Following the PTC phase, the system enters an iterative loop:

• Local Fine-Tuning: Clients refine their models on their private labeled data.
• Server Distillation: Clients send updated pseudo-labels to the server. The server updates the cluster-specific teacher models based on the new aggregated pseudo-labels.
• Client Distillation: Clients further refine their local models using the updated teacher pseudo-labels.

Crucially, the server models are trained using a regularization term derived from semi-supervised learning (SSL) theory, which penalizes prediction inconsistencies under data transformations (e.g., rotations, translations) to ensure robustness.

Theoretical Contributions

The paper provides the first end-to-end theoretical analysis of personalization risk contraction for model-agnostic PFL. By leveraging tools from SSL theory (specifically expansion-based connectivity and bounded pseudo-label error), the authors derive sufficient conditions for exponential risk contraction.

Key theoretical components include:

• Assumptions: The analysis assumes bounded importance weights between the public and private distributions, locally well-conditioned objectives, and a "confidence margin" where clients are confident in their predictions on the public data. It also relies on an expansion property of the data distribution, ensuring that pseudo-labels carry non-trivial information about ground truth.
• Risk Contraction: The authors prove that under these conditions, the personalization risk for both server models and client models contracts exponentially with the number of iterations. Specifically, the risk bound decreases at a rate determined by the cluster quality and the coverage of the public dataset, eventually dropping below the risk of local training alone.

Experimental Results

The authors evaluated COSMOS across four image classification benchmarks: EMNIST-balanced, CIFAR-10, CIFAR-100, and Tiny ImageNet. The experiments utilized a mix of client architectures (AlexNet, ResNet, MobileNet, SqueezeNet) and a VGG16 server model.

• Performance: COSMOS consistently outperformed all model-agnostic baselines (FedMD, FedCT, COMET) across all benchmarks. The performance gap widened as task complexity increased (e.g., from EMNIST to Tiny ImageNet).
• Heterogeneity: COSMOS remained robust to varying degrees of data heterogeneity (controlled by the Dirichlet parameter $\alpha$) and different mixtures of client architectures. It maintained a significant performance edge even when clients used weaker models.
• Communication Efficiency: By communicating only pseudo-labels (hard or soft) instead of model parameters or gradients, COSMOS reduced communication bandwidth by 1–2 orders of magnitude compared to parameter-sharing methods.
• Homogeneous Settings: Even in homogeneous settings where all clients share the same architecture, COSMOS remained competitive with state-of-the-art PFL methods, demonstrating that its model-agnostic design does not incur a performance penalty when homogeneity is present.

Significance and Claims

The paper claims that COSMOS addresses a critical bottleneck in FL by enabling high-performance personalization without requiring structural knowledge of client models. Its significance lies in three main areas:

1. Practical Deployment: It removes the structural barriers limiting PFL in diverse, real-world ecosystems where clients use proprietary or arbitrary architectures.
2. Theoretical Advancement: It provides the first general risk contraction guarantees for model-agnostic PFL, moving beyond standard convergence results to prove exponential risk reduction under realistic assumptions.
3. Scalability: It demonstrates that personalized server-side learning using pseudo-labels is a viable and scalable paradigm for highly heterogeneous environments, offering substantial communication efficiency gains.

The authors conclude that while the current work focuses on a client-server structure, the framework lays the groundwork for future decentralized learning systems, further enhancing the practical relevance of federated learning.