Source-linked AI summary

Heterogeneity for the Win: One-Shot Federated Clustering

Don Kurian Dennis, Tian Li, Virginia Smith

arXiv:2103.00697v2cs.LG

TL;DR

Unsupervised federated learning, especially clustering across heterogeneous devices, is relatively unexplored. The paper develops and analyzes one-shot k-FED, a Lloyd-style federated clustering scheme, and finds that specific heterogeneity regimes can weaken separation requirements while providing practical communication and failure-handling benefits.

  • Problem

    Unsupervised federated learning and clustering across heterogeneous device networks have received relatively limited attention.

  • Method

    k-FED runs a variant of Lloyd’s k-means locally on devices and aggregates local cluster means in one communication round.

  • Results

    The analysis shows measurable benefits from the paper’s specific heterogeneity notion, including weaker separation requirements for many cluster pairs than in centralized clustering.

  • Takeaways & Limitations

    One-shot federated clustering can address communication costs, stragglers, and device failures while supporting applications such as personalization and device sampling.

  • Takeaways & Limitations

    The analysis assumes center separation and a specific heterogeneous partition in which each device contains data from only a small number of target clusters.

Abstract

from arXiv · show

In this work, we explore the unique challenges -- and opportunities -- of unsupervised federated learning (FL). We develop and analyze a one-shot federated clustering scheme, $k$-FED, based on the widely-used Lloyd's method for $k$-means clustering. In contrast to many supervised problems, we show that the issue of statistical heterogeneity in federated networks can in fact benefit our analysis. We analyse $k$-FED under a center separation assumption and compare it to the best known requirements of its centralized counterpart. Our analysis shows that in heterogeneous regimes where the number of clusters per device $(k')$ is smaller than the total number of clusters over the network $k$, $(k'\le \sqrt{k})$, we can use heterogeneity to our advantage -- significantly weakening the cluster separation requirements for $k$-FED. From a practical viewpoint, $k$-FED also has many desirable properties: it requires only round of communication, can run asynchronously, and can handle partial participation or node/network failures. We motivate our analysis with experiments on common FL benchmarks, and highlight the practical utility of one-shot clustering through use-cases in personalized FL and device sampling.

1. Introduction

Unsupervised federated clustering is relatively unexplored despite FL’s heterogeneous device networks and practical clustering applications. The paper proposes one-shot k-FED and shows that certain forms of heterogeneity can weaken separation requirements while addressing communication and reliability constraints.

  • Motivation: Unsupervised federated learning remains relatively unexplored, motivating clustering methods for data distributed across heterogeneous devices.Clustering in FL supports client selection, personalization, and exploratory data analysis.
  • Method: k-FED applies a one-shot variant of Lloyd’s method, with devices solving local k-means problems and sending local cluster means to a central server.The method requires one communication round and messages of size O(dk'_z).
  • Practical properties: One communication round lets k-FED tolerate device failures when enough available devices collectively contain the k target clusters.Previously unavailable devices can be clustered through recomputation at the central server.
  • Main result: Heterogeneity can provide measurable benefits for k-FED under the paper’s specific clustering-oriented notion, contrasting with commonly reported harms in supervised FL.The supervised-FL harms described include poor convergence, unfair models, and the need for personalization.
  • Main result: When each device contains data from at most k' target clusters, the analysis compares k-FED with centralized clustering under center-separation assumptions.The paper considers heterogeneous partitions, including mixture-of-Gaussians examples, and reports a regime where requirements are similar to the centralized counterpart.
  • Main result: A large number of cluster pairs can satisfy a weaker separation requirement in heterogeneous networks, broadening the class of problems solvable relative to centralized clustering.This is presented as the first result the authors are aware of analyzing heterogeneity’s benefits for federated clustering.

2. Background and Related Work

The paper situates k-FED among centralized, parallel, distributed, and federated clustering methods. Unlike typical distributed implementations requiring numerous communication rounds, k-FED performs local Lloyd-style clustering followed by one aggregation round.

  • Centralized clustering: Lloyd’s heuristic iteratively assigns points to their nearest centers and recomputes each center as the mean of assigned points.The method starts with k initial centers and repeats these steps until termination.
  • k-FED: k-FED is a communication-efficient distributed variant that runs a Lloyd-style k-means procedure locally before aggregating and assigning clusters once.Its analysis follows deterministic frameworks for clustering mixture distributions without generative assumptions on the data.
  • Parallel and distributed clustering: Typical parallel and distributed clustering methods directly parallelize algorithms such as Lloyd’s heuristic or DBSCAN and require numerous communication rounds.This distinguishes them from the one-shot communication scheme studied here.
  • Federated clustering: Prior federated clustering work often clusters devices for supervised learning and uses iterative or centralized schemes rather than one-shot unsupervised clustering.These approaches target modeling non-IID data for downstream supervised tasks.
  • Federated clustering: A distributed matrix-factorization clustering approach considered statistical heterogeneity but did not focus on one-shot clustering or distinct heterogeneity benefits.The paper’s distinction is therefore both methodological and analytical.

3. k-FED: Preliminaries and Main Results

k-FED combines local Lloyd-style clustering with centralized aggregation under heterogeneous data partitions. Its analysis shows that, when devices contain only a small number of target clusters, inactive cluster pairs can require weaker separation while the method retains one-round recovery guarantees.

  • Method: k-FED applies a local k-means procedure on each device, then aggregates local cluster centers to form the final clustering.The framework builds on centralized Lloyd-type analyses and returns local cluster assignments before server-side aggregation.
  • Assumptions: The analysis assumes center separation and heterogeneous devices containing data from only a limited number of the total target clusters.This heterogeneity is modeled through local cluster counts and motivates separate active and inactive separation conditions.
  • Main results: Inactive cluster pairs have a weaker separation requirement than active pairs in the federated analysis.The paper explicitly distinguishes active and inactive separation and states that inactive pairs require less separation.
  • Main results: k-FED correctly classifies all but O(1/c^2)n points at termination, and classifies all points correctly when every local point satisfies the proximity condition.The theorem gives an approximate guarantee generally and an exact guarantee under the stronger local proximity condition.
  • Main results: When k' is sufficiently smaller than k, active separation approaches the centralized requirement while inactive separation weakens to a k' dependence.The paper identifies this regime as a benefit of heterogeneity and states that k-FED recovers target partitions in one communication round.
  • Practical implications: After server-side aggregation, new local centers from an unseen device can be assigned in O(k' · k) distance computations.The method is also presented as an inexpensive heuristic for exploration or preprocessing even when formal separation requirements are not satisfied.

4. Applications and Experiments

Experiments validate k-FED’s theoretical and practical claims across synthetic and real federated data. Structured heterogeneity improves clustering, while one-shot clustering supports personalized learning and client selection.

  • Separating Mixture of Gaussians: k-FED recovers the target clustering with high probability for the specialized mixture-of-Gaussians setup.The experiments instantiate the theoretical construction and evaluate clustering accuracy over repeated runs.
  • Separating Mixture of Gaussians: Even when the separation constant c is much smaller than the prescribed c ≥ 100, Figure 1 shows accurate clustering.The figure also reports decreasing variance across runs as separation increases.
  • Benefits of Heterogeneity: Structured partitions achieve k-means costs closer to oracle clustering than IID random partitions, with benefits increasing as k1 decreases.This pattern appears on FEMNIST and Shakespeare, despite real-data separation being below the theoretical requirement.
  • Communication-Efficiency: Using one communication round, k-FED can provide clustering quality similar to naive distributed k-means.The comparison evaluates k-FED against a more communication-intensive distributed baseline.
  • Client Selection: Combining k-FED with client selection boosts convergence and reduces final test-accuracy variance on realistic federated benchmarks.On FEMNIST, k-FED combined with pow-d reduces variance by 35% versus random selection.

5. Conclusion

The conclusion presents k-FED as an example where federated heterogeneity benefits unsupervised clustering. It also identifies practical advantages and broader research directions.

  • Conclusion: k-FED rigorously demonstrates that heterogeneity can benefit a one-shot federated clustering problem.The analysis concerns a one-shot variant of Lloyd’s algorithm for distributed clustering.
  • Conclusion: The method addresses high communication costs, stragglers, and device failures in federated settings.These practical concerns motivate the design of k-FED.
  • Conclusion: Other notions of heterogeneity may benefit additional federated learning problems, motivating future analysis.The paper frames this as an open direction requiring problem-specific analysis.

A.1. Proving Theorem 3.1 (Main Theorem)

The proof establishes correctness by showing that local centers approximate true cluster means, global initialization selects one representative per target cluster, and final assignments preserve the target partition.

  • Local clustering: The proof bounds each device’s local approximation error using the global matrix error and the structure of the device’s at-most-k1 clusters.The argument uses local subsets of target clusters and an operator-norm comparison.
  • Final assignment: The induced clustering satisfies the proximity condition, implying that no points are misclassified and the target clustering is recovered up to permutation.This completes the main correctness argument.
  • Local clustering: Local Lloyd clustering succeeds because each device contains only active cluster pairs, so inactive separation constraints do not affect the local problem.The resulting local centers remain close to the corresponding true cluster means.
  • Initialization: Under the separation assumptions and the nmin ≥ 4c^2k1nmax condition, k-FED selects exactly one local center for every target cluster.The proof maintains an invariant that selected centers correspond to distinct target clusters.
  • Final assignment: Each local center is closer to its corresponding initial center than to any other initial center, so centers from the same target cluster form one partition.This partition excludes centers associated with other target clusters.

A.2. Running Time of k-FED and Handling New Devices

After local clustering, k-FED’s central processing has polynomial distance-computation cost and can assign centers from previously unseen devices without rerunning global clustering.

  • Running time: The runtime analysis excludes local Algorithm-1 execution, which converges in polynomial time under the separation assumptions and typically terminates after few iterations in practice.The paper distinguishes this local computation from the analyzed central steps.
  • Running time: Steps 2–8 require O(Zk1 · k^2) pairwise distance computations.The bound sums the distances computed while selecting k initial centers across at most Zk1 local centers per iteration.
  • Handling New Devices: Once the set M is computed, local centers from a new device can be correctly assigned in O(k1 · k) distance computations.Assignment uses nearest-center matching to the already selected initial centers.

A.3. Separating Data from Mixture of Gaussian

The proof analyzes active and inactive cluster-pair separation under k1 ≤ k, using concentration and mixture-weight properties to establish the required conditions with high probability.

  • The analysis assumes k1 ≤ k and builds on prior results from Lemma 6.3.
  • The proof considers an active cluster pair under the stated separation requirement.
  • Each mixture component contains nearly its expected number of points with very high probability, where the expectation depends on its mixing weight.
  • The resulting separation condition holds with high probability.
  • The inactive separation condition is handled similarly, while the proximity condition follows from Gaussian concentration.

B.1. Datasets

Real-data experiments use EMNIST, FEMNIST, and Shakespeare through the LEAF benchmark, with a manually structured FEMNIST partition for client selection.

  • Real-data experiments use the EMNIST, FEMNIST, and Shakespeare datasets and corresponding models from the LEAF benchmark.
  • For client selection experiments, a subset of FEMNIST’s first 10 classes is partitioned by assigning 2 classes to each of 500 devices.

B.2. Choosing k Based on Separation

The paper selects k by examining relative center separation across candidate clusterings on MNIST. Most cluster pairs have small relative separation, even under oracle clustering.

  • The oracle procedure computes crs for every cluster pair and candidate k, then plots its distribution on MNIST.
  • The oracle clustering chooses k by maximizing the fraction of cluster pairs with crs > c0, using a fixed c0 such as 0.5.
  • For most cluster pairs and candidate k values, crs < 1, indicating relatively small separation.
Loading 2103.00697v2…