---
title: KLD-Weighted Clustered Federated Learning
url: https://www.emergentmind.com/topics/kld-weighted-clustered-federated-learning
type: topic
---

# KLD-Weighted Clustered Federated Learning

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 [2403.04144] [2302.10413] [2204.13619].

## 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 [2403.04144]. 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 [2403.04144]. 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 [2302.10413]. “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 [2204.13619].

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 [2403.04144]. 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 [2403.04144].

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 [2403.04144].

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 [2403.04144].

## 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 [2302.10413]. 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 \(j\)-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 \(W_g^t\), defining the round-\(t\) similarity between clients \(C_i\) and \(C_j\) as
\[
s^t_{ij} = \frac{(W^{t}_i -W_g^{t})^T \cdot (W^{t}_j - W_g^{t})}{\|W^{t}_i - W_g^{t} \|\|W^{t}_j-W_g^{t}\|}.
\]
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 \(Q\)-matrix, and forms clusters using the binary indicator
\[
u_{ij} = \begin{cases} 1, &\text{if } q_{ij} \geq \varepsilon; \\ 0, &\text{otherwise}, \end{cases}
\]
where \(\varepsilon\) increases over rounds [2302.10413].

CADIS also proposes **transitive learning** to accelerate similarity estimation via the triangle inequality for cosine similarity:
\[
s_{a,b}s_{b,c} - \sqrt{\left(1-s_{a,b}^{2}\right)\left(1-s_{b,c}^{2}\right)} \leq s_{a,c} \leq s_{a,b}s_{b,c} + \sqrt{\left(1-s_{a,b}^{2}\right)\left(1-s_{b,c}^{2}\right)}.
\]
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 [2302.10413].

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 \(\mathcal L_{KD}\). The final local objective is
\[
\mathcal{L} = \mathcal{L}_{CC} + \lambda\mathcal{L}_{KD},
\]
with \(\lambda \ge 0\) controlling the regularization strength [2302.10413]. 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 **\(3:3:2:1:1\)**, 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 [2302.10413].

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 \(\mathcal{L}_{FedAvg} - \mathcal{L}_{CADIS} \geq 0\). 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 [2302.10413].

## 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 [2204.13619]. Its objective is not client clustering but hierarchical personalization through a global model, cluster-specific models, and client-specific models.

The main objective is
\[
\min_{\{\theta_i\}_{i=1}^n} F(\{\theta_i\}_{i=1}^n) \coloneq \sum_{j=1}^k\sum_{i\in\mathcal I_j} \left( f_i(\theta_i) + \frac{(1-\alpha_j)\gamma_i}{2}\|\theta_i-\bar\theta_j\|^2 + \frac{\alpha_j\gamma_i}{2}\|\theta_i-\bar\theta\|^2 \right),
\]
where
\[
\bar\theta_j=\frac{\sum_{i\in\mathcal I_j}\gamma_i\theta_i}{\sum_{i\in\mathcal I_j}\gamma_i}, \qquad
\bar\theta= \frac{\sum_{j=1}^k\alpha_j\bar\theta_j}{\sum_{j=1}^k \alpha_j}.
\]
Here \(\gamma_i \ge 0\) is a client-specific regularization strength and \(\alpha_j \ge 0\) 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 \(w_j\) and a global model \(\bar w\), establishing a direct link between hierarchical personalization and weighted quadratic penalties [2204.13619].

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: \(\gamma_i\) corresponds to inverse within-cluster variance and \(\lambda_j\) 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 \(p_j\), and global communication occurs with probability \(p_0\). The paper derives an unbiased stochastic gradient oracle and proves linear convergence up to a noise floor under the assumption that each \(f_i\) is \(\mu\)-strongly convex and \(L\)-smooth:
\[
\mathbb E\|\theta^t-\hat\theta\|^2 \le (1-\eta\mu)^t\|\theta^0-\hat\theta\|^2 + \frac{2\eta\sigma_{\hat\theta}^2}{\mu}.
\]
It further chooses
\[
\tau_j=\frac{p_0}{p_0+2(1-p_0)p_j}
\]
to minimize the variance contribution of the oracle [2204.13619].

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 \(\lambda_j = \bar\sigma^{-2}\), \(\gamma_i = \bar\sigma_j^{-2}\), and \(X_i^TX_i=\beta_i I_d\), the minimizer of the hierarchical objective is the **best linear unbiased estimator** of each \(\theta_i^*\) 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 [2204.13619].

Empirically, the paper reports simulation results with **20 clusters**, **20 clients per cluster**, and varying local sample size \(m\in\{1,5,25,50,100,200\}\), 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 \(\ell_2\) 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 [2204.13619].

## 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 [2403.04144] | Euclidean distance among final-layer weights; agglomerative hierarchical clustering | None in clustering or weighting |
| CADIS [2302.10413] | Cosine similarity of penultimate-layer improvements; cumulative similarity matrix and thresholding | Used only in local knowledge-distillation regularization |
| Multiple Known Clusters [2204.13619] | 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 [2302.10413].

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 [2403.04144]. 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 [2302.10413]. 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 \(\alpha_j\) and \(\gamma_i\) [2204.13619].

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 [2403.04144]. 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 [2302.10413]. 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 [2204.13619].

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.

Source: https://www.emergentmind.com/topics/kld-weighted-clustered-federated-learning