REVIEW 3 major objections 4 minor 17 references
Clustered Federated Learning via Embedding Distributions
T0 review · 3 major / 4 minor · reviewed 2026-08-07 · deepseek-v4-flash
Pith's one-line read A one-shot clustered federated learning method groups clients by the Earth Mover's distance between embedding distributions and matches an oracle that knows the true clusters, without being told how many clusters exist.
desk verdict Novel one-shot FL clustering via embedding EMD, with strong empirical results but a theory section that needs fixing and a calibration heuristic that is not yet derived. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central object is the Earth Mover's distance $W_1(\mu_i^{\mathcal{Z}}, \mu_j^{\mathcal{Z}})$, the minimal transport cost between the distributions of two clients' data induced in the embedding space $\mathcal{Z}$ by their neural feature extractors. The method couples this with a reference-normalized threshold rule: each client $c$ computes $\tau_c = W_1(\mu_c^{\mathcal{Z}}, \mu_{\text{val},c}^{\mathcal{Z}})$ between its own training and validation embeddings, and the server links two clients only if the cross-client EMDs computed under both clients' models are within a fixed tolerance $\epsilon$ of the respective $\tau_c$ values. This turns a privacy-preserving, one-shot comparison of client data into a cluster adjacency graph whose connected components are the output clusters.
What would settle it
Train two clients on data drawn from the same underlying distribution but with deliberately different embedding models (for example, one pretrained on rotated images and one not) and compute the cross-client EMD under each model; if either distance exceeds $\tau_c + \epsilon$ despite identical data distributions, the clustering rule would split same-distribution clients. Alternatively, on Rotated MNIST, use a skewed validation split (for example, validation images from a single class) so that $\tau_c$ becomes abnormally small; EMD-CFL should then fragment the correct clusters and produce an adjusted Rand index well below 1.0.
Extended reading notes
Core claim
EMD-CFL computes, for each pair of clients, the EMD between their randomly projected embedding distributions, evaluated under each client's own embedding model, and normalizes by a per-client reference distance between that client's training and validation embeddings. Two clients are clustered together only if both normalized distances fall below a fixed tolerance. The paper claims this rule recovers the ground-truth cluster structure exactly, reaching an adjusted Rand index of 1.00 on all tested benchmarks, and that it is the only method among 16 baselines to match the Oracle on both Rotated CIFAR10 and Backdoor CIFAR10. The theoretical section proves that the generalisation error of a federated ensemble hypothesis, as well as the expected distance between client gradients and parameters, is bounded by the embedding-space EMD, so clustering by EMD targets an upper bound that prior gradient- and parameter-based clustering methods implicitly try to reduce.
Load-bearing premise
The method assumes that the EMD between a client's own training and validation data is a reliable ruler for judging whether another client's data comes from the same distribution, and that a single fixed tolerance value works across clients and datasets.
Editorial extensions
If this is right
- The cluster structure is identified after the first epoch and stays fixed, so the method is one-shot and does not require iterative refinement.
- The method does not need the number of clusters in advance; the number emerges from the adjacency graph.
- Cluster quality is robust to large random projection reductions (up to 90% dimension reduction) and to approximate Sinkhorn EMDs, so communication and computation can be traded off.
- On backdoor datasets, EMD-CFL separates backdoored clients from clean targets and preserves target accuracy, whereas parameter- and gradient-based methods leave targets vulnerable on CIFAR10.
- The theoretical bounds imply that clustering by embedding EMD also keeps within-cluster gradient and parameter distances small, directly connecting the method to the signals used by earlier clustered FL work.
Reading between the lines
- Because the reference-normalized threshold is a heuristic rather than a derived bound, a natural extension would be to estimate $\epsilon$ or $\tau_c$ adaptively from the within-client distribution of validation EMDs, which the paper does not explore.
- The one-shot EMD signal could be used as a preprocessing step for other clustered FL methods that require the number of clusters, supplying them with an initial partition rather than a final clustering.
- Since the method operates on embedding distributions from arbitrary pretrained feature extractors, it may transfer to text or tabular data, though the paper only tests images.
- On backdoor tasks, the method's ability to isolate subtle feature shifts suggests a potential defense-in-depth mechanism for federated systems, but the paper does not claim security guarantees against adaptive adversaries.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes EMD-CFL, a one-shot clustered federated learning method that groups clients by the Earth Mover's distance between their randomly projected embedding distributions. Clients exchange embedding models in pairs, compute cross-client EMDs under each other's models, and compare these against a per-client training/validation reference distance tau_c; the server then thresholds the normalized statistic with a fixed epsilon and derives clusters from the resulting adjacency matrix. The authors motivate the use of EMDs with domain-adaptation bounds on hypothesis disagreement, gradient distance, and parameter distance, and they report experiments on Rotated MNIST, Rotated CIFAR10, PACS, and two backdoor datasets, comparing against 16 baselines and an Oracle with ground-truth clusters. The central empirical claim is that EMD-CFL matches the Oracle and recovers the true clustering without knowing the number of clusters in advance.
Significance. If the empirical results hold, EMD-CFL is a practically attractive method: it is one-shot, does not require the number of clusters as input, and its evaluation is unusually thorough, with 16 baselines, multiple datasets, three runs, and robustness checks over epsilon, random-projection dimension, Sinkhorn approximation, model architecture, and partial participation. The authors also release code, which strengthens reproducibility. However, the theoretical contribution as written is not sound: Corollary 2.5 contains an inverted learning rate, the proof of Theorem 2.2 swaps mixture coefficients, and the core clustering statistic in Algorithm 1 is not justified by the Section 2 bounds. The paper is therefore best understood at present as an empirical method with heuristic motivation, and the theoretical framing needs substantial repair before the central claims can be accepted as stated.
major comments (3)
- [Section 2.3, Corollary 2.5 and Appendix A.1] The inequality E[||phi^{t+1}_i - phi^{t+1}_j||] <= (M/kappa) W1(mu_i, mu_j) has the learning rate inverted. From the stated update E[phi^{t+1}_i] = phi^t_i - kappa E[grad L_i] and Theorem 2.4, the triangle inequality gives E[||phi^{t+1}_i - phi^{t+1}_j||] <= kappa M W1(mu_i, mu_j) when both clients start from a common phi^t. As written, the bound decreases as the learning rate increases, which is the opposite of the behavior of the update. This matters because the paper lists 'our method targets an upper bound of the parameter and gradient distances' as a contribution; the corollary needs to be corrected or removed.
- [Section 2.2, Theorem 2.2 and Appendix A.1] The proof of Theorem 2.2 does not establish the stated inequality. The theorem defines the target mixture as T = alpha mu_i + (1-alpha) mu_j, so epsilon_T(h) = alpha epsilon_i(h) + (1-alpha) epsilon_j(h), which is the first line of the proof. After applying Lemma 2.1, however, the proof concludes with alpha epsilon_j(h) + (1-alpha) epsilon_i(h), swapping the coefficients. Unless the mixture definition is a typo, the displayed bound does not follow. Since Corollary 2.3 and the subsequent clustering motivation rely on Theorem 2.2, the theoretical chain needs to be re-derived and corrected.
- [Section 3, Algorithm 1 Lines 16-17; Section 4.2] The one-shot adjacency test uses the statistic W1(Z_{c|c}^R, Z_{c'|c}^R) - tau_c and a fixed threshold epsilon=0.025. Section 2 provides bounds on hypothesis disagreement, gradients, and parameters, but none of these results bounds this normalized EMD statistic or says how tau_c calibrates across different embedding models. Section 4.2 explicitly concedes that epsilon 'in general depends on the dataset, the embedding model and the random projection dimension.' Thus the paper's claim to recover the true cluster structure without knowing K rests on an uncalibrated heuristic. The robustness sweeps in Appendices E.1-E.6 vary epsilon and dim(R) on benchmarks whose within-cluster clients have very similar embedding distributions; they do not exercise the calibration mechanism under heterogeneous within-cluster spread or unrepresentative validation splits. I would like either a probabilistic or worst-case bound for the normalized statistic under the same-cluster hypothesis, or a direct calibration study that fixes epsilon while varying within-cluster dispersion and validation split quality.
minor comments (4)
- [Table 1 caption] The caption states that 'EMD-CFL is the only method to match the Oracle and obtains optimal clustering (bold) on both datasets,' but in the Rotated MNIST columns PACFL also reports ARI=1.00, average accuracy 98.86, and worst accuracy 97.73, matching or exceeding the Oracle row. If the intended meaning is 'the only method to match the Oracle on both datasets,' the caption should say so explicitly.
- [Equation (2)] The summation in Equation (2) is typeset as 'C X k=c', which appears to be a typo for the sum over clients c=1 to C; the indexing should be corrected.
- [Section 3, Client Embeddings] The definition of tau_c is written as tau_c = W1(Z_{c|c}, Z_{val|c}^R), but the first argument is missing the projection superscript R. The notation should be consistent with Algorithm 1, where tau_c is used together with projected distributions.
- [Algorithm 1] Line 11 tests 'if W[c][c'] is None', but the loop structure does not introduce c'. The pseudocode should explicitly iterate over c' in P(t) (with c' != c) so that the pairwise update in Lines 16-17 is well defined.
Circularity Check
No circularity: EMD-CFL's clustering statistic is a data-derived heuristic with fixed thresholds; the theoretical bounds are imported from external works and no fitted parameter is renamed as a prediction.
full rationale
EMD-CFL's derivation chain is self-contained against external, non-self-cited results: Lemma 2.1 is a bound due to Shen et al. (2018); Theorem 2.2 and Corollary 2.3 are proved in Appendix A.1 from that lemma and triangle/convexity inequalities; Theorem 2.4 is due to Fallah et al. (2020); and Corollary 2.5 follows algebraically. The algorithm-level statistic in Algorithm 1 (Lines 16-17) subtracts each client's own train-validation EMD tau_c from the cross-client EMD and thresholds with fixed epsilon=0.025. This is an uncalibrated statistical heuristic rather than a consequence of the Section 2 bounds, but it is not circular: tau_c is measured from each client's own held-out validation split, epsilon is a fixed hyperparameter, and neither encodes the ground-truth cluster labels, the Oracle accuracy, or the target ARI. The reported Oracle match is therefore empirically contingent, not true by construction. The paper also reports robustness sweeps over epsilon, random projection dimension, Sinkhorn approximation, architectures, and partial participation, consistent with a non-fitted evaluation. No load-bearing self-citation or imported uniqueness theorem is used. The only weakness is the hand-set global epsilon and the unproven reference-scale assumption, which is a correctness risk, not a circularity.
Assumptions & free parameters
free parameters (3)
- epsilon =
0.025
- dim(R) =
0.9 * dim(Z)
- subsample_size =
min(10% of client data, 512)
assumptions (5)
- domain assumption Existence of a ground-truth clustering of clients into K underlying distributions
- domain assumption Hypotheses h, h' are L-Lipschitz continuous with respect to the embedding space Z
- domain assumption The induced distributions over embeddings are well-defined and the reference distance tau_c is a valid within-client baseline
- standard math Johnson-Lindenstrauss lemma
- domain assumption The loss function is M-Lipschitz with respect to the embedding space
Cite this review
Pith. "Pith review of Clustered Federated Learning via Embedding Distributions." pith.science (2026). https://pith.science/paper/HDQTZJXT
@misc{pith2026250607769,
author = {Pith},
title = {Pith review of: Clustered Federated Learning via Embedding Distributions},
year = {2026},
howpublished = {\url{https://pith.science/paper/HDQTZJXT}},
note = {Machine review of arXiv:2506.07769}
}
read the original abstract
Federated learning (FL) is a widely used framework for machine learning in distributed data environments where clients hold data that cannot be easily centralised, such as for data protection reasons. FL, however, is known to be vulnerable to non-IID data. Clustered FL addresses this issue by finding more homogeneous clusters of clients. We propose a novel one-shot clustering method, EMD-CFL, using the Earth Mover's distance (EMD) between data distributions in embedding space. We theoretically motivate the use of EMDs using results from the domain adaptation literature and demonstrate empirically superior clustering performance in extensive comparisons against 16 baselines and on a range of challenging datasets.
Figures
Figures from the paper (14 more)
Reference graph
Works this paper leans on
-
[5]
We use Python Optimal Transport 0.9.5 [Flamary et al., 2021] for the EMD calculations
We implemented all experiments in PyTorch 2.5 and trained on an HPC cluster with NVIDIA L40S GPUs. We use Python Optimal Transport 0.9.5 [Flamary et al., 2021] for the EMD calculations. Table 5: Native embedding dimensions and learning rates used for training. Model dim(Z)κ Simple CNN 128 0.01 ResNet18 on CIFAR10 512 0.01 ResNet18 on PACS 512 0.001 ResNet...
work page 2021
-
[7]
Jakub Koneˇcn`y, H Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated optimization: Distributed machine learning for on-device intelligence.arXiv preprint arXiv:1610.02527,
-
[11]
Theoretical analysis of domain adaptation with optimal transport
Ievgen Redko, Amaury Habrard, and Marc Sebban. Theoretical analysis of domain adaptation with optimal transport. InMachine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2017, Skopje, Macedonia, September 18–22, 2017, Proceedings, Part II 10, pages 737–753. Springer,
work page 2017
-
[14]
12 A Theory A.1 Proofs Theorem A.1.Under the assumptions from Lemma 2.1, we have ϵT (h)≤αϵ j(h) + (1−α)ϵ i(h) + 2LW1(µZ i , µZ j ) +λ whereλ≡min h∈H ϵi(h) +ϵ j(h)is the sum error of the ideal hypothesis. Proof. ϵT (h) =αϵ i(h) + (1−α)ϵ j(h) ≤α[ϵ i(h, h∗) +ϵ i(h∗)] + (1−α) [ϵ j(h, h∗) +ϵ j(h∗)] (triangle inequality) =α[ϵ i(h, h∗)−ϵ j(h, h∗) +ϵ j(h, h∗) +ϵ ...
work page 1984
-
[15]
is closest in spirit to our approach and also focuses on comparing differences in data directly, but does so on the raw input level. In addition to the clustering identification approach, another point of distinction is if the clustering is soft or hard. Soft clustering effectively allows clients to participate in training multiple cluster models [Marfoq ...
work page 2021
-
[1985]
FedClust: Optimizing Federated Learning on Non-IID Data through Weight-Driven Client Clustering
Md Sirajul Islam, Simin Javaherian, Fei Xu, Xu Yuan, Li Chen, and Nian-Feng Tzeng. Fedclust: Optimizing federated learning on non-iid data through weight-driven client clustering.arXiv preprint arXiv:2403.04144,
-
[2004]
Intriguing properties of neural networks.arXiv preprint arXiv:1312.6199,
C Szegedy. Intriguing properties of neural networks.arXiv preprint arXiv:1312.6199,
-
[2009]
Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning.arXiv preprint arXiv:2002.10619,
arXiv 2002
Show all 17 references
-
[2010]
Towards federated learning at scale: System design.arXiv preprint arXiv:1902.01046,
Keith Bonawitz. Towards federated learning at scale: System design.arXiv preprint arXiv:1902.01046,
1902 arXiv
-
[2013]
An image is worth 16x16 words: Transformers for image recognition at scale
Alexey Dosovitskiy. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929,
2010 arXiv
-
[2017]
Federated adversarial domain adaptation
Xingchao Peng, Zijun Huang, Yizhe Zhu, and Kate Saenko. Federated adversarial domain adaptation. arXiv preprint arXiv:1911.02054,
1911 arXiv
-
[2018]
Balancing similarity and complementarity for federated learning
Kunda Yan, Sen Cui, Abudukelimu Wuerkaixi, Jingfeng Zhang, Bo Han, Gang Niu, Masashi Sugiyama, and Changshui Zhang. Balancing similarity and complementarity for federated learning. arXiv preprint arXiv:2405.09892,
-
[2019]
Fedrc: Tackling diverse distribution shifts challenge in federated learning by robust clustering.arXiv preprint arXiv:2301.12379,
Yongxin Guo, Xiaoying Tang, and Tao Lin. Fedrc: Tackling diverse distribution shifts challenge in federated learning by robust clustering.arXiv preprint arXiv:2301.12379,
-
[2020]
PACFL and FedClust require a threshold similar to our ϵ which we tune to find the closest to the optimal split in the first epoch
for this purpose. PACFL and FedClust require a threshold similar to our ϵ which we tune to find the closest to the optimal split in the first epoch. For PACFL, we used an ϵ of 12 for Rotated MNIST and CIFAR10 experiments, 11 for PACS, 2.5 for Backdoor MNIST and 10.5 for Backdo...
-
[2021]
Privacy via the johnson-lindenstrauss transform.arXiv preprint arXiv:1204.2606,
Krishnaram Kenthapadi, Aleksandra Korolova, Ilya Mironov, and Nina Mishra. Privacy via the johnson-lindenstrauss transform.arXiv preprint arXiv:1204.2606,
-
[2023]
Domain adaptation: Learning bounds and algorithms.arXiv preprint arXiv:0902.3430,
Yishay Mansour, Mehryar Mohri, and Afshin Rostamizadeh. Domain adaptation: Learning bounds and algorithms.arXiv preprint arXiv:0902.3430,
-
[2024]
Extensions of lipshitz mapping into hilbert space
William B Johnson. Extensions of lipshitz mapping into hilbert space. InConference modern analysis and probability, 1984, pages 189–206,
1984
Reviewed August 7, 2026 · model on record in the stance chip above.
Discussion (0). Sign in to comment.