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

AI summary

Overview

Research area: Federated learning (unsupervised/federated clustering), finite mixture models, expectation-maximization (EM) theory.

Technical level: Advanced — the paper combines a new federated algorithm with convergence proofs built on first-order stability (FOS), strong concavity, and contraction-region analysis.

Scope: The paper introduces FedGEM, a federated generalized expectation-maximization algorithm that trains mixture models across privacy-constrained clients when the global number of clusters is unknown and client cluster sets are heterogeneous but may overlap.

What This Paper Is About

Clients in a federated system each hold data drawn from their own mixture of clusters, some clusters are shared between clients but no client sees all of them, and nobody knows how many distinct clusters exist overall. Existing federated clustering methods generally assume the server already knows the total cluster count, and the one prior method that relaxes this (AFCL) requires clients to share arrays the same size as their raw data, which the authors argue allows data reconstruction at the server. FedGEM instead lets clients run EM locally, send only a short summary per component, and lets the server infer the global cluster count from those summaries.

Key Contributions

1. A federated GEM algorithm for unknown global component counts. FedGEM lets clients with overlapping clusters collaborate on shared cluster parameters while keeping cluster weights local, enabling personalization. Cluster overlaps are detected by the server through intersections of client-constructed uncertainty sets, using closed-form computations, which yields an estimate of the global number of clusters.

2. Probabilistic convergence guarantees. Under Assumptions 1–6 (including strong concavity, first-order stability, continuity, likelihood boundedness, and a finite-sample/population M-step proximity condition), the authors prove that iterates converge to a neighborhood of the ground-truth parameters with a stated probability, and that the inferred cluster count equals the true count with probability at least the product of (1 − δ_g) over all clients and local components.

3. A tractable instantiation for isotropic Gaussian mixture models. For multi-component isotropic GMMs the paper derives a bi-convex, two-dimensional reformulation of the client's uncertainty-set radius problem, provides closed-form population and finite-sample M-steps, and establishes an upper bound on the distance between the population and finite-sample M-step maps.

4. Empirical benchmarking. Experiments over 50 repetitions show FedGEM achieving performance comparable to centralized EM and outperforming existing federated clustering methods, at times even beating methods that are given prior knowledge of the cluster count.

Main Findings

• Uncertainty sets drive overlap detection. Each client solves an optimization problem to find the radius of a Euclidean ball around the maximizer of its local expected complete-data log-likelihood for each component; the server merges components into super-clusters when the distance between two maximizers is at most the sum of the two radii.
• A unique radius exists under strong concavity. Proposition 1 establishes that the local uncertainty-set radius problem has a unique solution ε_{k_g} ≥ 0 for every component and every client when the strong-concavity assumption holds.
• Local GEM convergence mirrors EM. Theorem 2 shows that a client update anywhere inside the uncertainty ball converges to a neighborhood of the ground truth with the same contraction factor β_g/λ_g plus a vanishing term ε̂(t) that goes to zero as t → ∞.
• Cluster count recovery with quantified probability. Theorem 3 states that if the final aggregation radius per component is set to ε_g^unif(N_g, δ_g) and the largest such radius is at most R_min/4, then the inferred number of clusters equals the true K with probability at least ∏g ∏{k_g} (1 − δ_g).
• The final aggregation radius is a user-set hyperparameter. Unlike the iterative stage, where radii are computed by solving an optimization problem, the final aggregation radius is treated as a user-defined hyperparameter.
• Isotropic GMM analysis is tractable. The general semi-infinite radius problem reduces to a bi-convex two-dimensional problem in (ε_{k_g}, α_{k_g}), and the strong-concavity parameter of the population Q_g is π_{min_g} when the current parameters equal the ground truth.
• Convergence for GMMs requires well-separated clusters. The GMM-specific FOS, contractive-radius, and population-to-finite-sample bound results (Theorems 6, 7, 8 in Appendix B.3) hold under a well-separated cluster requirement.
• Competitive empirical performance. Averaged over 50 repetitions, FedGEM is reported as comparable to centralized EM and better than existing federated clustering methods, scaling well with problem size. Specific dataset names and numeric scores are not included in the available content.

Methodology in Plain English

The method runs in two stages. In the collaborative training stage, every client takes some number of ordinary EM steps on its own local data to update its mixture parameters. For each local component, the client then solves an optimization problem to find how far it could move that component's center without making its local expected complete-data log-likelihood worse — this defines a "ball of uncertainty" around the center. The client sends only the center and the radius to the server. The server compares all the balls from all clients: if two balls overlap, it treats those two components as the same global cluster, merges them, and computes a shared center that stays inside both balls. It then sends the updated centers back. Because updates are allowed anywhere inside a ball that does not decrease the local likelihood objective, the procedure is a generalized EM rather than a strict EM.

The second stage is a single final aggregation pass. Clients send their centers along with a final aggregation radius (this time a user-chosen number rather than a solved one), and the server merges any centers that fall within that radius of each other. Counting the resulting super-clusters gives the estimate of the total number of unique clusters, and the assignment of clusters to clients follows from which components were merged. For isotropic Gaussian mixtures with identity covariance, the client-side radius problem is rewritten as a small bi-convex problem in two variables, which makes each client's per-iteration computation cheap.

Why This Matters

Impact on research. FedGEM is presented as the first federated GEM algorithm for training mixture models without prior knowledge of the global number of components. It relaxes two assumptions that prior federated clustering work retained — identical cluster sets across clients, and server-side knowledge of the cluster count — and it addresses the privacy weakness the authors attribute to AFCL, where clients share arrays sized like their raw data and the server can reconstruct data via simple scalar multiplication and subtraction. FedGEM instead shares a maximizer-and-radius tuple per component per iteration, and the paper includes a preliminary differential privacy discussion of those shared finite-sample maximizers in Appendix B.4.

Real-world applications (drawn from the paper's motivating scenario for original equipment manufacturers):

• Fault detection and diagnosis in power generators and other capital-intensive industrial equipment.
• Servicing of medical imaging systems, where the OEM must detect and classify faults without a global labeling standard.
• Long-term service contracts (LTSCs) guaranteeing reliability standards, where missing those guarantees can incur multi-million dollar penalties.
• Any multi-site deployment where clients cannot share raw data because of its size, dimensionality, or privacy concerns, labels are inconsistent across sites, and the full set of fault classes is not known in advance.

Industry relevance. The motivation is explicitly industrial: OEMs of capital-intensive systems need to diagnose faults across many client sites, cannot rely on client-supplied labels because maintenance practices and labeling conventions differ, and cannot centralize training for privacy and data-size reasons. The combination of an unknown number of fault classes, heterogeneous but overlapping fault sets, and no raw-data sharing maps directly onto that setting.

Future Directions

• Differential privacy guarantees. The paper provides only a preliminary differential privacy discussion for the shared finite-sample maximizers; formal privacy guarantees remain open.
• Relaxing the well-separation requirement. The convergence results for isotropic GMMs require the clusters to be well separated, so extending the analysis to overlapping or poorly separated clusters is an open question.
• Robustness beyond the modeling assumptions. The authors note that they highlight performance in problem settings that violate modeling assumptions, suggesting further work on when the assumptions can be dropped.
• Verifying assumptions on more model families. The general convergence results rely on Assumptions 1–6, which the paper verifies for isotropic GMMs; checking them for other component distributions (and other covariance structures) is a natural extension.

Target Audience

Researchers and graduate students working on federated learning, distributed clustering, and mixture-model estimation; practitioners in industrial reliability, prognostics, and fault diagnosis who need multi-site learning under privacy constraints; and theoretically inclined readers interested in convergence analysis of EM-type algorithms built on first-order stability and contraction arguments.

Authors’ abstract

We study the problem of federated clustering when the total number of clusters $K$ across clients is unknown, and the clients have heterogeneous but potentially overlapping cluster sets in their local data. To that end, we develop FedGEM: a federated generalized expectation-maximization algorithm for the training of mixture models with an unknown number of components. Our proposed algorithm relies on each of the clients performing EM steps locally, and constructing an uncertainty set around the maximizer associated with each local component. The central server utilizes the uncertainty sets to learn potential cluster overlaps between clients, and infer the global number of clusters via closed-form computations. We perform a thorough theoretical study of our algorithm, presenting probabilistic convergence guarantees under common assumptions. Subsequently, we study the specific setting of isotropic GMMs, providing tractable, low-complexity computations to be performed by each client during each iteration of the algorithm, as well as rigorously verifying assumptions required for algorithm convergence. We perform various numerical experiments, where we empirically demonstrate that our proposed method achieves comparable performance to centralized EM, and that it outperforms various existing federated clustering methods.