Distilling A Universal Expert from Clustered Federated Learning

1 Introduction

Federated learning (FL) has emerged as a promising paradigm for privacy-preserving distributed training, enabling collaborative model development without exposing sensitive personal data [Nguyen et al., 2021; Beltr´an et al., 2023; Zhang et al., 2023]. The primary goal of FL is to train a high-quality consensus model through joint optimization. However, deploying FL across distributed clients is nontrivial due to the statistical heterogeneity of local data (i.e., non- IID). Standard averaging-based aggregation method [McMahan et al., 2017] tends to result in biased local updates, degrading the performance of the consensus model and slowing convergence. Prior works [Karimireddy et al., 2020; Li et al., 2020] have introduced correction mechanisms that align local and global update directions. While effective in reducing update bias, these approaches compromise local model personalization, a key factor in ensuring high performance under data heterogeneity.

∗Corresponding authors.

Clustered Federated learning (CFL) effectively compensates for local personalization by training multiple clusterlevel consensus models for heterogeneous devices [Tu et al., 2025; Huang et al., 2024; Ma et al., 2023]. Unlike traditional FL, CFL inherently preserves individual data characteristics through its clustering mechanism. Yet even state-of-the-art CFL approaches still face challenges in various practical settings. A key problem lies in the limited information flow across clusters: small clusters are prone to overfitting, and large clusters typically converge to local optima. As a result, CFL consensus models tend to overfit local personalization, limiting their generalization capability.

In fact, neither traditional FL nor CFL methods effectively balance personalized knowledge with global consensus. An ideal solution aims to build a consensus model that preserves client-specific features. Yet, this is a significantly challenge, as naive aggregation of personalized knowledge tends to induce cluster-level model bias. These observations prompt us to explore a new question:

Can the consensus model be successfully conducted without undermining personalization?

Transferring knowledge across heterogeneous clusters offers a potential solution. Recent advances in knowledge distillation (KD) have shown promise in facilitating such knowledge migration [Wang and Yoon, 2021]. However, a major limitation of traditional KD methods is their reliance on auxiliary datasets. The choice of auxiliary data significantly impacts model performance and introduces additional computational overhead [Wang et al., 2024]. These limitations motivate us to adopt a data-free knowledge distillation (DFKD) approach [Lin et al., 2020a], in which a generator is trained through model inversion to synthesize pseudo-data on the fly, enabling knowledge transfer from the multiple teacher networks to the student network.

To this end, we face several key challenges: C1. How to design a training method with multiple teacher networks such that heterogeneous knowledge remains non-conflicting? C2. How to effectively extract and ensemble cross-cluster knowledge, given that dynamic clustering induces random intercluster distribution shifts? C3. How to mitigate privacy risks in clustering, as the direct exposure of similarity values in existing CFL methods increases the risk of re-identification.

To simultaneously address all the aforementioned challenges, we propose a novel FL framework, DisUE (Distilled Universal Expert)), which enhances the consensus model while preserving local personalization. DisUE follows a three-stage learning paradigm: local model training, iterative clustering, and universal expert distillation. In particular, to tackle C1, we construct a universal expert model from CFL and propose a category cluster-level knowledge migration mechanism to enable single expert-to-student distillation. To tackle C2, we design two adaptive components at the category cluster-level to handle dynamic inter-cluster heterogeneity. By leveraging category statistics, these components guide both model inversion and ensemble. To tackle C3, we introduce a lightweight similarity encryption protocol that prevents direct exposure of similarity values.

In a nutshell, our main contributions are as follows:

• Novel Framework for FL: We propose DisUE, a distillation-enhanced framework from CFL, which can maintain both intra- and inter-group knowledge in FL.
• Adaptive Category Cluster-Distillation Mechanism: We devise an adaptive Group Label Sampler (GLS) and Group Weighting Factors (GWF) to process heterogeneous knowledge across groups. It incorporates groupbased distillation to reduce distillation aggregation overhead and globally shares a model to improve the performance of minority groups.
• Flexible plugin. As a flexible algorithm that is orthogonal to existing CFL optimizers (IFCA, CFL, CFL-GP, PACFL). Our DisUE can easily enhance existing CFL methods, demonstrating the generality and compatibility of our inter-group aggregation mechanism.
• Superior performance. Extensive experiments on three real-world datasets against diverse advancing baselines prove the consistently state-of-the-art performance of our proposed DisUE.

2 RelatedWork

2.1 Clustered Federated Learning

CFL addresses statistical heterogeneity by grouping clients into clusters with homogeneous data distributions [Long et al., 2023; Ma et al., 2022]. Existing CFL methods [Ghosh et al., 2020; Sattler et al., 2020; Duan et al., 2021] primarily focus on client partitioning to derive accurate personalized expert models. Prior work has explored clustering under various assumptions, such as temporal domain shifts between training and test data [Fan et al., 2024], joint label and feature skewness [Guo et al., 2024], and multiple types of data distribution shifts [Ruan and Joe-Wong, 2022]. Other approaches optimize aggregation schedules [Kim et al., 2024]. However, these methods restrict inter-cluster communication, causing expert models to overfit limited client data and suffer in generalizability. We bridge this gap by designing an expert model based on CFL, and improving its generalization via clusterlevel DFKD transfer.

2.2 Data-free Knowledge Distillation

DFKD removes the need for auxiliary datasets in KD by transferring knowledge through model output exchange [Yoo et al., 2019]. Current approaches fall into two categories based on source networks: single-teacher DFKD [Shin and Choi, 2024; Lin et al., 2020b] and multi-teacher DFKD [Hao et al., 2021; Ye et al., 2020]. While most DFKD methods operate at the client level, they require more training rounds than aggregation methods [McMahan et al., 2017; Karimireddy et al., 2020] due to the larger scale of pseudodata generated relative to original samples. To address this limitation, we introduce a cluster-level DFKD approach that lowers communication costs via group knowledge fusion.

3 Problem Statement

CFL Pipeline. CFL follows an iterative three-stage framework [Liu et al., 2024]:

• Local training (L-phase). Each client performs local training and transmits model updates to the server. Let there be N clients, where the i-th client holds a private dataset Di = {(x(l) i , y(l) i )}ni l=1 with ni samples. Client i updates its local parameters ωi by minimizing the clientlevel empirical risk:

where f(·) denotes the per-sample loss. All clients then send their parameters {ωi}Ni=1 to the server.

• Clustering (C-phase). The server partitions clients into K disjoint clusters C = {C1, . . . , CK} using similarity metrics (e.g., cosine distance).
• Group model aggregation (G-phase). CFL maintains K group models for N clients. Within cluster Ck, m clients with identically distributed data collaboratively train a group model ωCk . We formulate the cluster-level optimization objective as:

Where L denotes the loss function for intra-group clients. The server subsequently distributes the updated cluster models {ωCk }Kk=1 to their respective clusters.

Emerging Requirements from CFL. The core limitation arises from strict inter-cluster communication constraints. As prior analyses demonstrate, overcoming this challenge requires global information sharing - a critical need that drives our architectural redesign of the Group model aggregation:

• G-phase: Directly applying existing methods to intercluster knowledge transfer encounters two fundamental barriers: (1) Distribution Discrepancy. The inherent heterogeneity across clusters in CFL leads to systematic

bias when using traditional aggregation methods (e.g., FedAvg). (2) Adaptation rigidity. Static knowledge transfer mechanisms fail to respond to dynamic distribution shifts between clusters. These challenges demand a new CFL paradigm that supports adaptive transfer and dynamically adjusts to distributional changes induced by shifting cluster structures.

4 Methodology

In this section, we introduce DisUE, a cross-cluster federated learning framework that overcomes the limitations discussed in Section Introduction (Figure 1). The framework operates through three core components: (1) Clustering and Intra-group Aggregation (CGA): Partitions clients using cosine similarity-based clustering and extracts group-specific knowledge through intra-group model aggregation. (2) Intergroup Aggregation (IGA): Distills heterogeneous cluster expertise into a global model via data-free adversarial learning. (3) Secure Similarity Computing Protocol (SSC): Encrypts client parameters using PCSC, ensuring security against parameter leakage during similarity computation.

4.1 Clustering and Intra-group Aggregation

Principles. This stage achieves client dynamic clustering through angular similarity measurement, hyperparameterfree operation, and non-IID robustness. We quantify pairwise client relationships using cosine similarity between gradient updates. Unlike density-based methods [McInnes et al., 2017; Ester et al., 1996] that require predefined neighborhood thresholds, our design eliminates hyperparameters through message passing. Additionally, these algorithms exhibit unstable cluster structure under data heterogeneity. Our goal is to adopt a dynamic clustering mechanism that maintains robustness in extreme non-IID scenarios while enabling seamless integration with FL frameworks-critical for realizing a universal expert model. These requirements drive our adoption of Affinity Propagation (AP) [Frey and Dueck, 2007], which automatically identifies exemplar clients through iterative responsibility/availability updates.

Clustering. We compute pairwise client similarities using cosine distance between parameter vectors ωu and ωv:

where ∥ωu∥ denotes the length of the parameter vector ωu. Let Sim be the cosine similarity matrix. The AP clustering obtains K disjoint groups as follows:

Intra-group Aggregation. For each cluster Cj ∈ C, we aggregate member models through Eq. (2), denotes as {ωC1 , . . . , ωCK }. Yielding specialist models {ωCj }Kj=1 that capture cluster-specific knowledge patterns.

4.2 Inter-group Aggregation

Principles. This stage addresses two co-existing challenges: adaptive cluster-wise label distribution shifts and crosscluster distillation aggregation. Some aggregation methods based on subgroups are constrained by inter-group node symmetry or cluster topology symmetry [Jahani-Nezhad et al., 2022; So et al., 2021], while global averaging suppresses cluster-specific knowledge. Our solution breaks this stalemate through adaptive category group-level DFKD aggregation. Specifically, (1) Dynamically adjusting cluster contributions through adaptive GLS and GWF, and (2) establishing cross-heterogeneous cluster knowledge transfer via prediction logit alignment between cluster experts and the global model. We fuse GLS and GWF into the cluster-level DFKD aggregation through a distillation-driven ensemble, creating a flexible mechanism for cross-cluster knowledge transfer.

Handing Cluster-level Distribution Imbalance. Label distribution imbalance constitutes a primary manifestation of distributional discrepancies across heterogeneous clusters. The dynamic cluster assignment mechanism in CFL induces stochastic cluster-level label distribution variations across training rounds. Within cluster-level DFKD methods, this bias manifests through two critical pathways: (1) leveraging cluster-specific distribution characteristics to optimize generator training, and (2) aggregating cluster models according to their category-specific contributions. Inspired by [Zhang et al., 2022] and the underlying principle of counters, we proposed GLS and GWF.

The GLS samples representative labels guided by intracluster category statistics:

where I(·) denotes the indicator function, and nky represents the cardinality of category y samples in cluster k. The GW dynamically determines category importance weights per cluster:

This adptaive weighting scheme ensures equitable integration of category-specific knowledge across clusters. These adaptive components regulate cluster-level knowledge transfer through accumulated category distribution statistics, enabling effective handling of non-IID data scenarios.

Cluster-level DFKD Aggregation Mechanism. Our framework aims to facilitate cross-cluster knowledge transfer through data-free distillation, enabling the global model ωG to assimilate heterogeneous cluster characteristics. We preserve inter-group knowledge foundations by maintaining cluster-averaged feature representations through aggregated cluster models serving as student networks. These cluster models subsequently function as teacher networks to supervise student model distillation.

The teacher-student knowledge transfer occurs through data-free distillation where teacher models guide student models to capture their specific data distributions. However, client-side generator deployment imposes prohibitive computational overhead on edge devices. To overcome this limitation, we design a generator G at the server that synthesizes pseudo-samples ˆx conditioned on category labels y:

where z ∼ N(0, 1) represents random Gaussian noise vectors, and θG denotes the generator's parameters.

The cross-cluster distillation method achieves knowledge transfer through prediction discrepancy minimization between ensemble teachers and student network on generated samples, effectively capturing teacher-specific distribution patterns. We formulate the cluster distillation objective as:

where ˆ p(y) denotes the cluster-level label distribution from Eq. (5), and αk y represents the adaptive cluster weights defined in Eq. (6) . To enable precise cluster knowledge transfer, we leverage the global model baseline while applying clusterspecific regulation through:

Here C(·) denotes the classifier output layer, σ(·) the softmax function, and KL the Kullback-Leibler divergence measuring prediction distribution discrepancies.

To maintain semantic coherence in cross-cluster distillation while mitigating computational overhead, we optimize the centralized generator through cluster-conditioned crossentropy minimization. The objective function for cluster fidelity, denoted as Lcf , is given by:

The per-cluster fidelity loss is computed as:

In CFL environments, cluster-specific knowledge exhibits multi-category characteristics. We adopt the diversity-aware regularization term Ldiv [Zhu et al., 2021] to enhance knowledge transfer through maximizing inter-sample dissimilarity.

where Q is the number of pseudo-samples, and zi denotes the noise of the i-th pseudo-sample.

4.3 Secure Similarity Computation

Principles. Similarity matrix exposure introduces reidentification risks, as adversaries could infer sensitive client attributes through parametric similarity analysis. We address this challenge through a lightweight Privacy-Preserving Cosine Similarity Computing (PCSC) protocol protocol for CFL [Lu et al., 2014]. The protocol prevents server access to raw similarity values through encrypted operations.

4.4 Optimization and Algorithm of DisUE

Optimization Objective. Under CFL assumptions, intracluster data homogeneity allows cluster-optimal solutions via FedAvg, whereas inter-group label shifts degrade expert model performance. We seek a globally optimal solution ωG that harmonizes k cluster-specific distributions.

The DisUE framework establishes an adversarial game between the global model ωG and a generator G (parameterized by θG), governed by a composite loss (Lcd, Lcf and Ldiv) that coordinates cross-cluster knowledge distillation through alternating optimization phases: (1)Maximization Phasewhere G synthesizes hard-classifiable samples from cluster to amplify prediction conflicts, and (2) Minimization Phase where ωG aligns predictions with cluster models using boundary samples. This process progressively sharpens global decision boundaries while absorbing cluster-specific knowledge.

Algorithm 1 Workflow of DisUE

Require: Initial global model ω(0) G , client set Ns, rounds T,

security parameter=sec params Ensure: Final global model ω(T) G , generator parameters θ(T) G Initialization: 1: Broadcast ω(0) G to all clients in Ns 2: for round t = 1 to T do L-Phase: Local Training 3: Client Update 4: for all client i ∈ Ns in parallel do 5: Local training: ω(t) i ← (ω(t-1) G ) 6: E(ω(t) i ) ← SSC.Encrypt(sec params, ω(t) i ) 7: Upload E(ω(t) i ) to server 8: end for 9: Initialize similarity matrix Sim(t) ← 0|Ns|×|Ns| 10: for all pairs (u, v) ∈ Ns × Ns where u < v do 11: Compute encrypted similarity: 12: Sim(t) u,v ← SSC.Compute(E(ω(t) u ), E(ω(t) v )) 13: Update Sim(t) [u, v] ← Sim(t) u,v 14: end for C-Phase: Clustering 15: {ω(t) 1 , . . . , ω(t) K } ← CGA-Clustering(Sim(t) ) 16: Intra-group aggregation: 17: {ωC(t) 1 , . . . , ωC(t) K } ← CGA-Intra({C (t) 1 , . . . ,C (t) K }) G-Phase: Group Model Aggregation 18: Global aggregation: 19: ω(t) G ← FedAvg(ωC(t) 1 , . . . , ωC(t) K ) 20: Inter-group aggregation: 21: ω(t+1) G ← IGA({ωC(t) 1 , . . . , ωC(t) K }, ω(t) G , θG ) 22: ServerUpdate(ω(t+1) G , θG) 23: end for 24: return ω(T) G , θ(T) G

The unified optimization objective combines three loss components from Eqs. (8), (10), and (12):

where βcf and βdiv balance the loss terms. This minimax optimization maximizes cross-cluster knowledge transfer, ultimately producing a universal model with consistent performance across clusters.

Algorithm. Algorithm 1 illustrates the workflow of DisUE. We begin by broadcasting an initial global model ω(0) G to all clients. In each communication round t, clients first perform L-phase on their private data. Subsequently, each client encrypts its updated parameters using the SSC protocol and transmits these encrypted updates to the server. The server obtains a similarity matrix Sim(t) , which it leverages to execute CGA-Clustering, partitioning clients into K groups. Within group, parameters are aggregated via CGA-Intra, conducing clustering structure {ωCt1 , . . . , ωCtK }.

Next, the global model ωtG is obtained via a federated aggregation (FedAvg) of these group models. To enable effective cross-cluster distillation, the IGA aggregates distilled knowledge from group models by refining ωtG into ωt+1 G without direct access to raw client data. Finally, the server updates the global parameters (and the generator parameters θG before proceeding to the next communication round. After T rounds, DisUE returns the ω(T) G and θ(T) G .

5 Experiments

5.1 Implementation Details

Datasets. We evaluate DisUE on three standard benchmarks: SVHN [Netzer et al., 2011], CIFAR-10 [Krizhevsky et al., 2009], and CIFAR-100 [Krizhevsky et al., 2009]. To simulate federated data heterogeneity, we partition datasets using the Dirichlet distribution Dir(ϵ) [Hsu et al., 2019], with concentration parameters ϵ ∈ [0.01, 0.06] controlling non-IID difficulty levels. Lower ϵ values create more skewed client data distributions, representing challenging FL scenarios.

Baselines. We benchmark DisUE against 11 state-of-the-art FL methods spanning three key categories: (1) CFL methods including IFCA [Ghosh et al., 2020], CFL [Sattler et al., 2020], CFL-GP [Kim et al., 2024], and PACFL [Vahidian et al., 2023]; (2) DFKD approaches FedDF [Lin et al., 2020b] and FedFTG [Zhang et al., 2022]; and (3) conventional FL methods covering FedAvg [McMahan et al., 2017], FedProx [Li et al., 2020], FedDyn [Acar et al., 2021], MOON [Li et al., 2021], and SCAFFOLD [Karimireddy et al., 2020].

Configurations. For all methods, we set the communication rounds T = 500, the number of clients N = 100, with an active fraction Act = 0.15. For local training, we set the number of local epochs localE = 5, batch size = 50, and the weight decay to 1×10-3 . The learning rates for the classifier and generator are initialized to 0.1 and 0.01, respectively. The dimension z is set to 100 for CIFAR-10 and SVHN, and 256 for CIFAR-100. Unless otherwise specified, we adopt βcf = 1.0 and βdiv = 1.0. All our experimental results represent the average over five random seeds. For the classifier, we adopt the network architecture from [He et al., 2016]. The generator architecture from [Zhang et al., 2022] is employed for both FedFTG and FedDF.

5.2 Performance Comparison

Test Accuracy. Table 1 reports the test accuracy achieved by all methods on the SVHN, CIFAR10, and CIFAR100 datasets. (1) Among all scenarios, our method consistently outperforms the three baseline algorithm types on both IID and non-IID settings. (2) Compared to CFL algorithms, DisUE achieves significantly better performance, primarily due to cross-group knowledge transfer, which enables its global model to leverage more comprehensive information. (3) Compared to the DFKD algorithms, performance is further improved. This is attributed to the group-based distillation protocol having a pre-trained component, clustering, thereby mitigating the non-IID problem in advance. (4) Experimental observations reveal distinct performance patterns across algorithms under varying data distributions. Global FL algorithm achieves superior accuracy in IID and Dir(ϵ = 0.1). This phenomenon arises from the relatively balanced category distributions across clients, which induces aligned

(b) Test accuracy on act(·), dirϵ = 0.01

(c) Test accuracy on SVHN dir(ϵ = 0.01)

Figure 2: (a) Test accuracy versus data heterogeneity measured by ϵ(·). (b) Test accuracy versus the fraction C of active clients per round (ϵ = 0.01). (c) Test accuracy over 500 communication rounds (learning curve) with ϵ = 0.01. All experiments are conducted on SVHN.

parameter update directions that undermine clustering efficacy. Conversely, under Dir(ϵ = 0.01) non-IID conditions, CFL demonstrates better performance over global FL.

Communication Overhead. We assessed the communication rounds needed for convergence by DFKD SOTA methods on SVHN with on-IID degrees set to 0.01 in Figure. 2c. (1) Our method achieves accelerated convergence rates through cluster-level adversarial distillation. (2) During the first 100 iterations, our methods demonstrate greater performance gains compared to other DFKD approaches, evidencing that integration clustering effectively mitigates data heterogeneity issues and prevents early-stage local noise.

Data Heterogeneity. We evaluate the robustness of algorithms under varying degrees of non-IID data by measuring the test accuracy of all methods with non-IID degrees set to (0.02, 0.03, 0.04, 0.05, 0.06). Figure 2a illustrates the test accuracy as a function of varying Dir(·), respectively showing the best-performing algorithm for each type of method. (1)DisUE outperforms baseline methods across diverse configurations, demonstrating the effectiveness of CFL's clustering mechanism in (i) preserving heterogeneous data characteristics through clustering and (ii) maximizing knowledge utility via inter-group distillation. (2) Global FL algorithms exhibit performance improvements with increasing Dirichlet coefficient ϵ. (3) CFL performance degrades progressively with higher values due to accumulating clustering errors.

Client Participation Scale. We evaluated all approaches across varying client scales to observe how the impact of each algorithm responds to an increasing number of participants. Thus, we set act= (0.1, 0.15, 0.2, and 0.25) per training round. (1) As shown in Figure 2b, DisUE demonstrates superior performance. (2) Increasing act(·) improves the accuracy of all methods, as more local information becomes available per round. (3) Compared with FL and DFKD, CFL and DisUE demonstrate a more stable increase in accuracy. This indicates that, under a fixed number of clients, the CFL algorithm should avoid selecting a high client participation rate when adapting to a Dirichlet distribution.

5.3 Compatibility Study

To validate the compatibility of the proposed aggregation mechanism within CFL frameworks, we integrated it into several existing CFL optimizers and evaluated their accuracy. As shown in Table 2, our cross-cluster distillation module consistently enhances the global performance of each CFL method, with the combination involving PACFL achieving the highest test accuracy. Essentially, the clustering algorithm acts as a fundamental component in CFL. Therefore, these results also demonstrate our method's compatibility with various clustering strategies, as it improves their performance while remaining orthogonal to the underlying clustering approach.

5.4 Ablation Study

To assess the necessity of each component, we test accuracy after removing different key modules and loss functions. In all experiments, we use Dir(ϵ = 0.01) on the SVHN task. For the modules, we remove GLS, GWF, and IGA, while for the key loss functions, we eliminate Lcf and Ldiv. In Table 3, after removing the IGA, DisUE reverts to FedAvg, resulting in the most performance degradation. This underscores the critical importance of extracting and migrating crosscluster information. GLS and GWF modules cause the model to average samples and integrate identical weights, respectively. Consequently, the global model becomes incapable of addressing differences in data distributions across groups, leading to further performance declines. Lcf and Ldiv play key roles in pseudo data generation. Dropping Lcf leads to blurred or less realistic outputs, whereas dropping Ldiv results in less diverse categories.

5.5 Hyperparameters Sensitivity Analysis

We evaluate the sensitivity of key hyperparameters: βcf, βdiv, and noise dimension. The βcf and βdiv weights are assessed at (0.5, 0.75, 1, 1.25, and 1.5), while the noise dimensions are set to (50, 100, 150, 200, and 250), following the settings in [Zhang et al., 2022]. (1) As illustrated in Figure 3a and Figure 3b, DisUE achieves optimal performance when βcf = 1 and βdiv = 1. (2) An unsuitable βcf may lead to capturing the group model semantic imbalances in the pseudo-data. Similarly, the improper ratio of βdiv could reduce the diversity of pseudo-data. (3) As shown in Figure 3c, a noise dimension of 100 produces the best performance, indicating that excessively high noise dimensions should be avoided in this setup.

6 Conclusion

Training a consensus model that effectively balances personalized and global knowledge under statistical heterogeneity remains a challenge in FL. We address this by introducing a novel framework that dynamically clusters clients to capture personalization and leverages a cluster-level DFKD strategy to distill and transfer global knowledge across groups.

Acknowledgements

Zeqi Leng, Chunxu Zhang and Bo Yang are supported by the National Natural Science Foundation of China underGrant Nos.U22A2098,62172185,62206105, and 62202200; the Major Science and Technology Development Plan of Jilin Province under Grant No. 20240302078GX; the Major Science and Technology Development Plan of Changchun under Grant No.2024WX05. Riting Xia is supported by the Inner Mon-golia University High-level Talent Project under Grant No. 10000-23112101/2861.