KLD-Weighted Clustered Federated Learning

• The paper clarifies that weight-driven clustering, rather than KL divergence, forms the basis for grouping non-IID client models.
• It compares three techniques: FedClust uses Euclidean distance of final-layer weights, CADIS incorporates KL in local regularization, and known-cluster FL applies quadratic regularization.
• Practical insights reveal that one-shot weight-driven clustering reduces communication rounds while boosting accuracy on benchmarks like CIFAR-10 and FMNIST.

Searching arXiv for the cited papers and nearby clustered-FL / KL-related work to ground the article. Clustered federated learning (CFL) is a family of federated learning methods for non-IID client populations in which clients are grouped according to shared structure, either because similar data distributions are inferred from model behavior or because cluster membership is known a priori. In the strict technical sense implied by the cited works, “KLD-weighted clustered federated learning” would require KL divergence to determine clustering or server-side aggregation weights. The representative methods most closely adjacent to that label do not do so. FedClust is a weight-driven CFL method based on Euclidean proximity among final-layer weights; CADIS is a cluster-aware FL framework whose clustering is based on penultimate-layer cosine similarity and whose KL divergence term appears only in local knowledge-distillation regularization; and the known-cluster personalized FL formulation is entirely KL-free and instead uses weighted quadratic regularization across client, cluster, and global models (Islam et al., 2024, Nguyen et al., 2023, Lyu et al., 2022).

1. Problem formulation and scope

Federated learning enables collaborative model training on decentralized devices without exposing local data, but its central difficulty is uneven data distribution across clients, which violates the IID assumption used in conventional machine learning (Islam et al., 2024). CFL addresses this by exploiting client structure rather than averaging all updates as though they were sampled from a single homogeneous population.

The cited literature distinguishes at least three heterogeneity regimes. FedClust addresses generic non-IID data by grouping clients whose local training induces similar model parameters (Islam et al., 2024). CADIS introduces cluster-skewed non-IID, a setting in which clients can be partitioned into clusters with similar data distributions and the cluster sizes are unequal; this is more structured than generic label skew because it explicitly includes inter-client correlation (Nguyen et al., 2023). “Personalized Federated Learning with Multiple Known Clusters” studies the case in which the cluster partition is already known, denoted $\mathcal I_1,\dots,\mathcal I_k$, and the question is how to regularize learning within and across those clusters (Lyu et al., 2022).

Across these settings, the operational question is not merely whether clients differ, but how that difference should enter the algorithm. The three principal design choices in the cited works are: inferring similarity from model weights, equalizing or regularizing information flow across clusters, and deciding whether the relevant weighting is geometric, cardinality-based, or probabilistic. A central conclusion of the cited papers is that these are distinct mechanisms and should not be conflated with KL-based weighting.

2. FedClust: weight-driven clustering from final-layer parameters

FedClust is explicitly described as a weight-driven clustered federated learning method rather than a KL-divergence-based method (Islam et al., 2024). Its premise is an “implicit relationship between the local model weights and the underlying data distribution on a client device,” with the specific empirical observation that the final layer or classifier layer reflects data-induced differences more clearly than earlier layers. Clients with similar local label distributions are reported to learn similar final-layer weights; those weights are therefore treated as proxies for client data distributions.

The workflow is one-shot. The server broadcasts an initial global model; clients train locally; clients upload updated final-layer weights; the server interprets those weights as a representation of client data distribution; and clustering is performed from those weights. The server computes a proximity matrix using the Euclidean distance among final layer weights and then applies agglomerative hierarchical clustering (HC). The clustering mechanism is therefore geometric rather than information-theoretic.

The only explicit optimization equations reproduced in the paper are the standard FedAvg objective and local objective,

$$\min_{\theta} F(\theta) \overset{\scriptscriptstyle\Delta}{=} \sum_{i=1}^{m}\frac{n_i}{N}F_i(\theta)$$

and

$$F_i(\theta)=\frac{1}{n_i}\sum_{j_i=1}^{n_i}f_{j_i} (\theta; x_{j_i},y_{j_i}),$$

which are not clustering objectives. For clustering, the paper states in words that Euclidean distance is computed among final-layer weights and used to form the matrix $\mathcal M$, but it does not provide a KL-divergence formula, weighted-KL formula, or any other information-theoretic clustering objective (Islam et al., 2024).

FedClust contrasts itself with earlier CFL methods that require many communication rounds for stable cluster formation, a predefined number of clusters, or full-model weight comparisons. Its one-shot strategy completes cluster formation in a single communication round and then proceeds with intra-cluster training in a workflow described as similar to FedAvg. This suggests a specific tradeoff: faster cluster formation and lower communication cost are obtained by relying heavily on the quality of the initial representation encoded in the final layer.

The reported evaluation uses CIFAR-10, Fashion MNIST (FMNIST), and SVHN under Non-IID Dir (0.1). VGG16 appears in the motivation experiment demonstrating the relevance of final-layer weights, while LeNet-5 is used in the main performance evaluation. Baselines are FedAvg, FedProx, CFL, IFCA, and PACFL. The main metric is test accuracy, reported as mean $\pm$ standard deviation. FedClust achieves the best reported accuracies in the excerpted table: 60.25 $\pm$ 0.58 on CIFAR-10, 95.51 $\pm$ 0.17 on FMNIST, and 78.23 $\pm$ 0.30 on SVHN; the paper also emphasizes improved communication cost, although the excerpt does not provide a separate numeric communication-efficiency table (Islam et al., 2024).

The limitations stated or implied in the excerpt are equally important. FedClust is not a KLD-weighted method; it relies on the heuristic that classifier-layer weights are reliable proxies for data distribution; one-shot clustering may be fragile if the initial local training is noisy or insufficient; broader validation is left as future work; and the dynamic accommodation of newcomers is claimed but not fully formalized in the excerpt (Islam et al., 2024).

3. CADIS: cluster-aware aggregation with KL-regularized local training

CADIS is a framework for cluster-skewed non-IID data rather than a KLD-weighted clustering method in the usual sense (Nguyen et al., 2023). Its starting point is that clients can be grouped into clusters with similar data distributions and that those clusters may have very different sizes. Under such skew, FedAvg can bias the global model toward majority clusters because it averages client updates using weights based only on client sample counts, $p_i = \frac{n_i}{\sum_i n_i}$, and therefore ignores cluster cardinality and inter-client correlation.

The clustering signal in CADIS comes from the penultimate layer. The paper argues that training on samples from class $j$ tends to increase the $$\min_{\theta} F(\theta) \overset{\scriptscriptstyle\Delta}{=} \sum_{i=1}^{m}\frac{n_i}{N}F_i(\theta)$$0-th row of the penultimate-layer weight matrix and decrease others, so changes in that layer provide class-dependent evidence about local data. The server compares the improvement of clients’ penultimate layers relative to the global model’s penultimate layer $$\min_{\theta} F(\theta) \overset{\scriptscriptstyle\Delta}{=} \sum_{i=1}^{m}\frac{n_i}{N}F_i(\theta)$$1, defining the round-$$\min_{\theta} F(\theta) \overset{\scriptscriptstyle\Delta}{=} \sum_{i=1}^{m}\frac{n_i}{N}F_i(\theta)$$2 similarity between clients $$\min_{\theta} F(\theta) \overset{\scriptscriptstyle\Delta}{=} \sum_{i=1}^{m}\frac{n_i}{N}F_i(\theta)$$3 and $$\min_{\theta} F(\theta) \overset{\scriptscriptstyle\Delta}{=} \sum_{i=1}^{m}\frac{n_i}{N}F_i(\theta)$$4 as

$$\min_{\theta} F(\theta) \overset{\scriptscriptstyle\Delta}{=} \sum_{i=1}^{m}\frac{n_i}{N}F_i(\theta)$$5

Because only a subset of clients participates in each round, the server maintains a cumulative similarity matrix, rescales it by min-max normalization into a $$\min_{\theta} F(\theta) \overset{\scriptscriptstyle\Delta}{=} \sum_{i=1}^{m}\frac{n_i}{N}F_i(\theta)$$6-matrix, and forms clusters using the binary indicator

$$\min_{\theta} F(\theta) \overset{\scriptscriptstyle\Delta}{=} \sum_{i=1}^{m}\frac{n_i}{N}F_i(\theta)$$7

where $$\min_{\theta} F(\theta) \overset{\scriptscriptstyle\Delta}{=} \sum_{i=1}^{m}\frac{n_i}{N}F_i(\theta)$$8 increases over rounds (Nguyen et al., 2023).

CADIS also proposes transitive learning to accelerate similarity estimation via the triangle inequality for cosine similarity: $$\min_{\theta} F(\theta) \overset{\scriptscriptstyle\Delta}{=} \sum_{i=1}^{m}\frac{n_i}{N}F_i(\theta)$$9 This lets the server infer unseen pairwise similarities from observed ones. In the reported ablation, transitive learning speeds up similarity-matrix convergence to about 40 rounds versus 100 rounds for the standard method (Nguyen et al., 2023).

The aggregation rule is cluster-equalized weighted averaging. The excerpt states that the weight assigned to a client depends on both inverse cluster cardinality and FedAvg-style sample weighting, with the explicit intention that large clusters do not dominate while clients with more data still matter more within their clusters. The global model is then aggregated from the weighted client models. The design objective is equality between clusters rather than divergence-based weighting.

The KL term appears elsewhere: in local training regularization. CADIS treats the global model as a teacher and the client model as a student, uses KDE-based conditional probabilities to model pairwise sample interactions, and defines a KL-based distillation loss $$F_i(\theta)=\frac{1}{n_i}\sum_{j_i=1}^{n_i}f_{j_i} (\theta; x_{j_i},y_{j_i}),$$0. The final local objective is

$$F_i(\theta)=\frac{1}{n_i}\sum_{j_i=1}^{n_i}f_{j_i} (\theta; x_{j_i},y_{j_i}),$$1

with $$F_i(\theta)=\frac{1}{n_i}\sum_{j_i=1}^{n_i}f_{j_i} (\theta; x_{j_i},y_{j_i}),$$2 controlling the regularization strength (Nguyen et al., 2023). The paper is explicit that this KL divergence is used to reduce local overfitting and improve training, not to compute cluster assignments or server aggregation weights.

The experiments cover PILL, MNIST, CIFAR-10, and CIFAR-100, including a multi-cluster (MC) cluster-skewed non-IID setting with 5 clusters and client ratios $$F_i(\theta)=\frac{1}{n_i}\sum_{j_i=1}^{n_i}f_{j_i} (\theta; x_{j_i},y_{j_i}),$$3, as well as PA, BC, and UC variants. Representative top-1 results include PILL: CADIS 79.71\% versus FedAvg 73.36\%, a gain of 8.7\%; CIFAR-10: CADIS 50.09\% versus FedAvg 48.74\%; and CIFAR-100: CADIS improves over FedAvg by up to 16.5\%–20.2\% depending on client participation level. Convergence is also faster: on PILL, CADIS reaches 60\% accuracy in 45 rounds versus 74 for FedAvg/FedProx, and on CIFAR-10 it reaches 48\% in 537 rounds versus 960 for FedAvg (Nguyen et al., 2023).

The theoretical claims are correspondingly scoped. The paper argues that the penultimate layer contains enough information to infer class presence in local data and gives a simplified three-client convergence comparison in which $$F_i(\theta)=\frac{1}{n_i}\sum_{j_i=1}^{n_i}f_{j_i} (\theta; x_{j_i},y_{j_i}),$$4. The excerpt explicitly notes that this is not a general convex optimization theorem for all FL settings but a proof sketch supporting the design intuition that cluster-aware weighting is preferable to naive averaging when clusters exist (Nguyen et al., 2023).

4. Known-cluster personalized federated learning: weighted regularization without KL divergence

“Personalized Federated Learning with Multiple Known Clusters” studies a different regime: the cluster structure is assumed known in advance rather than discovered from client updates (Lyu et al., 2022). Its objective is not client clustering but hierarchical personalization through a global model, cluster-specific models, and client-specific models.

The main objective is

$$F_i(\theta)=\frac{1}{n_i}\sum_{j_i=1}^{n_i}f_{j_i} (\theta; x_{j_i},y_{j_i}),$$5

where

$$F_i(\theta)=\frac{1}{n_i}\sum_{j_i=1}^{n_i}f_{j_i} (\theta; x_{j_i},y_{j_i}),$$6

Here $$F_i(\theta)=\frac{1}{n_i}\sum_{j_i=1}^{n_i}f_{j_i} (\theta; x_{j_i},y_{j_i}),$$7 is a client-specific regularization strength and $$F_i(\theta)=\frac{1}{n_i}\sum_{j_i=1}^{n_i}f_{j_i} (\theta; x_{j_i},y_{j_i}),$$8 is a cluster-specific mixing weight between cluster-level and global-level sharing. The paper also gives an equivalent multi-task formulation with cluster models $$F_i(\theta)=\frac{1}{n_i}\sum_{j_i=1}^{n_i}f_{j_i} (\theta; x_{j_i},y_{j_i}),$$9 and a global model $\mathcal M$0, establishing a direct link between hierarchical personalization and weighted quadratic penalties (Lyu et al., 2022).

The paper explicitly states that it does not use KLD or KL-divergence at all. The weighting is instead quadratic and can be interpreted through a hierarchical Gaussian model: $\mathcal M$1 corresponds to inverse within-cluster variance and $\mathcal M$2 to inverse between-cluster variance. In that sense, the weights are precision weights or shrinkage parameters rather than divergence weights.

Algorithmically, the method introduces Async-L2GD, a loopless asynchronous method that randomly alternates among local descent, within-cluster communication, and between-cluster communication. Clusters communicate independently, each cluster can have its own communication probability $\mathcal M$3, and global communication occurs with probability $\mathcal M$4. The paper derives an unbiased stochastic gradient oracle and proves linear convergence up to a noise floor under the assumption that each $\mathcal M$5 is $\mathcal M$6-strongly convex and $\mathcal M$7-smooth: $\mathcal M$8 It further chooses

$\mathcal M$9

to minimize the variance contribution of the oracle (Lyu et al., 2022).

The paper’s statistical analysis uses a hierarchical linear model in which cluster centers are drawn around a global mean and client parameters are drawn around their cluster center. Under $\pm$0, $\pm$1, and $\pm$2, the minimizer of the hierarchical objective is the best linear unbiased estimator of each $\pm$3 given the data and the other clients’ local estimates. It also proves that the method improves over both purely local estimation and a single shared global estimator in the modeled regime (Lyu et al., 2022).

Empirically, the paper reports simulation results with 20 clusters, 20 clients per cluster, and varying local sample size $\pm$4, as well as a real-world donor-response dataset in which ZIP codes are clustered by median household income into 5 groups. The proposed model achieves the lowest average $\pm$5 estimation error in simulation, gives competitive or better average cross-entropy than baselines on donor-response prediction, reduces worst-case client loss relative to the global model, and substantially outperforms purely local learning (Lyu et al., 2022).

5. Comparative interpretation of weighting, clustering, and KL divergence

The three papers occupy adjacent but distinct positions in the clustered FL design space.

Method	Cluster signal or structure	Role of KL divergence
FedClust (Islam et al., 2024)	Euclidean distance among final-layer weights; agglomerative hierarchical clustering	None in clustering or weighting
CADIS (Nguyen et al., 2023)	Cosine similarity of penultimate-layer improvements; cumulative similarity matrix and thresholding	Used only in local knowledge-distillation regularization
Multiple Known Clusters (Lyu et al., 2022)	Known clusters; quadratic regularization between client, cluster, and global models	Not used

This comparison makes two distinctions unavoidable. First, clustering and weighting are separate operations. FedClust clusters from weight-space geometry and then trains within clusters in a FedAvg-like manner; CADIS clusters from penultimate-layer similarity and aggregates with cluster-aware equalization; the known-cluster method does not infer clusters at all but regularizes models according to a fixed hierarchy. Second, the presence of KL divergence somewhere in the training pipeline does not imply KLD-weighted clustering or KLD-weighted aggregation. Among these papers, only CADIS contains a KL term, and that term regularizes local client training rather than determining cluster assignments or server aggregation weights (Nguyen et al., 2023).

A plausible implication is that the phrase “KLD-weighted clustered federated learning” is technically meaningful only if the role of KL divergence is specified. The cited literature shows at least three non-equivalent possibilities: KL in clustering, KL in server aggregation weighting, and KL in local regularization. The surveyed methods instantiate only the third, and only in CADIS.

6. Misconceptions, limitations, and boundaries of the term

A common misconception is to equate any distribution-aware CFL method with a KL-based method. FedClust directly contradicts that reading: it is distribution-aware only in the indirect sense that final-layer weights are treated as proxies for local data distributions, and its actual clustering criterion is Euclidean distance in weight space (Islam et al., 2024). Another misconception is that any use of KL divergence makes a method “KLD-weighted.” CADIS uses KL divergence in a knowledge-distillation loss, but its clustering is based on penultimate-layer cosine similarity and its aggregation is based on cluster size and sample size, not KL-derived weights (Nguyen et al., 2023). A third misconception is that weighted clustered FL necessarily refers to probabilistic divergence weighting; the known-cluster personalized FL paper shows a purely quadratic, KL-free alternative based on hierarchical regularization with weights $\pm$6 and $\pm$7 (Lyu et al., 2022).

The limitations reported in the papers define the present boundary of the topic. FedClust depends on the heuristic that final-layer weights reliably encode client data distributions, may be fragile because of one-shot clustering, evaluates only certain datasets under Non-IID Dir (0.1) in the excerpt, and does not fully formalize newcomer handling there (Islam et al., 2024). CADIS derives its convergence comparison from a simplified three-client analysis, reports its strongest advantage in clustered settings, and uses KL only as an auxiliary local regularizer rather than as the core cluster-weighting mechanism (Nguyen et al., 2023). The known-cluster personalized method assumes cluster labels are available a priori, derives optimality only within the class of linear unbiased estimators under a hierarchical linear model, and notes that better biased estimators may exist but are harder to implement efficiently in FL (Lyu et al., 2022).

Taken together, these works delimit the phrase rather than fully instantiate it. They show that contemporary clustered FL adjacent to “KLD-weighted” is more accurately partitioned into weight-driven clustering, cluster-aware aggregation with KL-regularized local training, and precision-weighted hierarchical personalization. This suggests that any strict use of the label “KLD-weighted clustered federated learning” should be reserved for methods in which KL divergence explicitly governs cluster formation or aggregation weights, a condition not satisfied by the representative papers surveyed here.