A Federated Generalized Expectation-Maximization Algorithm for Mixture Models with an Unknown Number of Components¶

Conference: ICLR 2026 arXiv: 2601.21160 Code: None Area: Other Keywords: federated clustering, generalized EM algorithm, mixture models, unknown number of clusters, uncertainty sets

TL;DR¶

This paper proposes FedGEM, an algorithm in which clients perform local EM steps and construct uncertainty sets, while the server detects cluster overlap via set intersections and infers the global number of clusters. FedGEM is the first method to achieve federated clustering with an unknown global cluster count \(K\), and comes with probabilistic convergence guarantees.

Background & Motivation¶

Background: Clustering in federated learning typically assumes that all clients share the same set of clusters and that the global number of clusters \(K\) is known in advance.

Limitations of Prior Work: In practice (e.g., fault diagnosis for OEM devices), clients may have heterogeneous yet potentially overlapping cluster sets, and the total number of fault categories is unknown. Existing federated clustering methods (k-FED, FFCM, FedKmeans) all require \(K\) to be specified beforehand. The only method that does not require prior knowledge of \(K\), AFCL, suffers from serious privacy vulnerabilities—data shared by clients can be reconstructed by the server through simple computation.

Key Challenge: Clients cannot share raw data (privacy constraints), there is no unified label standard (heterogeneous labels), and the global number of clusters is unknown. Neither conventional centralized training nor existing federated methods can simultaneously satisfy all these constraints.

Goal: Design a privacy-preserving federated clustering algorithm capable of inferring \(K\) and training mixture models without prior knowledge of the global cluster count.

Key Insight: Combine EM with uncertainty sets—after the M-step, each client computes an uncertainty radius defined as the maximum perturbation range that does not cause a decrease in likelihood; the server then uses the intersection of uncertainty sets to detect cluster overlap.

Core Idea: Use the intersection of likelihood-nondecreasing uncertainty sets around EM maximizers to identify cross-client cluster overlap, thereby enabling federated clustering without prior knowledge of the global \(K\).

Method¶

Overall Architecture¶

FedGEM consists of two phases: (1) an iterative collaborative training phase—clients perform local EM and compute uncertainty sets, while the server aggregates parameters using set intersections; (2) a final aggregation phase—the server merges converged cluster estimates and infers the global \(K\).

Key Designs¶

Client Computation: EM Step + Uncertainty Sets
Function: Each client runs EM on its local data and computes an uncertainty set radius for each mixture component.
Mechanism: The uncertainty set \(\mathcal{U}_{k_g}\) is a Euclidean ball centered at the M-step maximizer with radius \(\sqrt{\varepsilon_{k_g}}\), obtained by solving an optimization problem such that any iterate within the ball does not decrease the expected complete-data log-likelihood. For isotropic GMMs, this radius problem admits a closed-form solution via a doubly convex two-dimensional reformulation.
Design Motivation: The uncertainty set characterizes the range of parameter perturbation that still guarantees algorithmic non-degradation; its size reflects estimation precision. Clients only need to transmit center-radius pairs, which are far smaller than raw data, thereby preserving privacy.
Server Aggregation: Uncertainty Set Intersection Detection
Function: Identify cluster overlap by checking whether uncertainty sets of components from different clients intersect.
Mechanism: If \(\|M_{k_g} - M_{k_{g'}}\|_2 \leq \sqrt{\varepsilon_{k_g}} + \sqrt{\varepsilon_{k_{g'}}}\), the two components are assigned to the same super-cluster. During collaborative training, the optimal vector \(\nu^*\) lying within the intersection is computed to update parameters; during final aggregation, estimates within the same super-cluster are directly merged.
Design Motivation: Set intersection provides a purely geometric, closed-form mechanism for detecting cross-client cluster correspondence, eliminating the need to share raw data or large arrays.
Convergence Guarantees
Function: Prove that under standard assumptions, the algorithm converges in probability to a neighborhood of the true parameters and correctly estimates the global \(K\).
Mechanism: Using the first-order stationarity (FOS) condition, the GEM iterates are shown to converge to a \(\frac{1}{1-\beta/\lambda}\epsilon^{\text{unif}}\) neighborhood of the true parameters \(\theta^*\). It is further proved that when the final aggregation radius is set to \(\epsilon^{\text{unif}}\) and does not exceed \(R_{\min}/4\), the algorithm correctly infers \(K\).
Design Motivation: The theoretical guarantee is one of the core contributions of this work, elevating the algorithm beyond a heuristic to one with rigorous convergence rates.

Loss & Training¶

The framework naturally optimizes the log-likelihood under the EM paradigm. Uncertainty set radii are obtained by solving constrained optimization problems. Under the isotropic GMM setting, low-complexity feasible algorithms are provided.

Key Experimental Results¶

Main Results¶

Dataset	FedGEM (ARI)	AFCL (ARI)	DP-GMM (ARI)	GMM (known K)	Est. K (true K)
MNIST	0.355±.064	0.137±.055	0.326±.080	0.287±.067	14.7 (10)
FMNIST	0.415±.022	0.143±.019	0.332±.069	0.385±.023	12.3 (10)
CIFAR-10	0.058±.008	0.022±.007	0.032±.017	0.402±.022	33.7 (10)
Abalone	0.128±.030	0.091±.022	0.061±.018	0.096±.028	4.5 (3)

Ablation Study¶

Configuration	ARI (Synthetic)	Notes
Centralized EM	~0.95	Upper bound
FedGEM	~0.90	Close to centralized
AFCL	~0.50	Significantly behind
k-FED (known K)	~0.85	Requires known K

Key Findings¶

FedGEM consistently outperforms AFCL—the only competing method that does not require prior knowledge of \(K\)—across all datasets, and in some cases even surpasses methods that assume \(K\) is known.
On non-Gaussian datasets such as CIFAR-10, the number of clusters is substantially overestimated (33.7 vs. 10), revealing limitations of the isotropic GMM assumption on complex data.
FedGEM scales well; its performance degradation with increasing problem size is notably smaller than that of competing methods.
Reasonable performance is maintained on real-world datasets even when Gaussian assumptions are violated.

Highlights & Insights¶

Elegant design of uncertainty sets: The property of GEM that allows inexact M-steps is repurposed as a tool for cross-client cluster matching—theoretically elegant and practically efficient. The radius naturally reflects estimation precision, shrinking as precision improves.
Closed-form aggregation: The server requires only simple distance comparisons and linear operations, with no iterative optimization, yielding high computational efficiency.
Theoretical completeness: The convergence proof chain progresses rigorously from general mixture models to isotropic GMMs, with each step well-justified.

Limitations & Future Work¶

The isotropic GMM assumption limits performance on complex high-dimensional data (e.g., CIFAR-10); extending the framework to full-covariance GMMs or more expressive mixture models is a natural direction.
Cluster count estimation remains noticeably biased (particularly toward overestimation), and better strategies for selecting the final aggregation radius are needed for practical deployment.
Convergence guarantees rely on the assumption of sufficient cluster separation; they may fail when clusters overlap significantly.
The discussion of differential privacy is only preliminary; formal privacy guarantees remain to be established.

vs. k-FED/FFCM: These methods require the global \(K\) to be specified in advance, limiting practical applicability; FedGEM infers \(K\) automatically.
vs. AFCL: The only comparable method that does not require prior knowledge of \(K\), but AFCL has serious privacy vulnerabilities; FedGEM transmits only parameters and radii.
vs. DP-GMM (centralized): FedGEM approaches—and in terms of ARI occasionally surpasses—centralized DP-GMM performance in the federated setting.

Rating¶

Novelty: ⭐⭐⭐⭐ Uncertainty-set-driven federated clustering is a novel idea with originality in both theoretical and algorithmic design.
Experimental Thoroughness: ⭐⭐⭐⭐ Evaluations span 8 datasets, 7 baselines, and scalability analysis, providing comprehensive coverage.
Writing Quality: ⭐⭐⭐⭐ Theoretical derivations are rigorous, though the notation is heavy and readability is moderate.
Value: ⭐⭐⭐⭐ Fills a theoretical and methodological gap in federated clustering with an unknown number of clusters.