Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Michael Ibrahim, Nagi Gebraeel, Weijun Xie

TL;DR

FedGEM tackles federated clustering without prior knowledge of the global number of components $K$ by having clients run local EM and form uncertainty sets around local maximizers; the central server merges overlapping uncertainty sets to infer global overlaps and the total cluster count while preserving data privacy. The method provides probabilistic convergence guarantees under standard GEM assumptions and specializes to isotropic GMMs with a tractable bi-convex reformulation for the client radius problem. Empirical results show FedGEM can match or exceed centralized EM performance and outperform several federated baselines across diverse datasets, even when clusters are overlapping or non-ideal. The work demonstrates a practical, scalable approach to unsupervised federated clustering with unknown $K$, offering theoretical and empirical validation and highlighting directions for privacy enhancements and broader model classes.

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.

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

TL;DR

FedGEM tackles federated clustering without prior knowledge of the global number of components by having clients run local EM and form uncertainty sets around local maximizers; the central server merges overlapping uncertainty sets to infer global overlaps and the total cluster count while preserving data privacy. The method provides probabilistic convergence guarantees under standard GEM assumptions and specializes to isotropic GMMs with a tractable bi-convex reformulation for the client radius problem. Empirical results show FedGEM can match or exceed centralized EM performance and outperform several federated baselines across diverse datasets, even when clusters are overlapping or non-ideal. The work demonstrates a practical, scalable approach to unsupervised federated clustering with unknown , offering theoretical and empirical validation and highlighting directions for privacy enhancements and broader model classes.

Abstract

We study the problem of federated clustering when the total number of clusters 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.
Paper Structure (45 sections, 11 theorems, 84 equations, 4 figures, 7 tables, 4 algorithms)

This paper contains 45 sections, 11 theorems, 84 equations, 4 figures, 7 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Problem Setting
  4. Overview of FedGEM Algorithm
  5. Client Computations
  6. Server Computations
  7. Convergence Analysis
  8. FedGEM for Multi-Component Isotropic GMMs
  9. Numerical Experiments
  10. Benchmarking
  11. Sensitivity
  12. Conclusions
  13. Supplementary Pseudo-code
  14. Detailed FedGEM Pseuo-code
  15. Detailed Server Computations Pseudo-code
  16. ...and 30 more sections

Key Result

Proposition 1

Suppose Assumption assump:strong_conc holds. Then, there must exist a unique solution $\varepsilon_{k_g} \geq 0$ to the optimization problem $J_{k_g}(\boldsymbol{\theta}_g^{(t-1)})$ for all components $k_g \in [K_g]$ and all clients $g \in [G]$. (Proof in Appendix app:prop1_proof).

Figures (4)

  • Figure 1: Results of the sensitivity study.
  • Figure 2: Supplementary results of the sensitivity study.
  • Figure 3: Results on the sensitivity of our algorithm to its hyperparameter.
  • Figure 4: Results of the scalability experiment for all experimental settings and benchmark models.

Theorems & Definitions (30)

  • Remark 1
  • Proposition 1: Local Uncertainty Set Radius Problem
  • Definition 1: First-Order Stability
  • Remark 2
  • Theorem 1: Single-Point EM Convergence
  • Theorem 2: Local Convergence of Finite-Sample GEM
  • Theorem 3: Number of Clusters Inference
  • Theorem 4: Radius Problem Reformulation
  • Theorem 5: Local Convergence of Population GEM
  • proof
  • ...and 20 more