Heterogeneity for the Win: One-Shot Federated Clustering

In plain words

Imagine many phones, each holding data that only a few groups of items belong to. You want to group all the data together, but you cannot collect the raw data in one place. This paper's method, k-FED, has each device group its own data once and send only a handful of group centers to a server. The server matches those centers across devices and gives every point a global group.

Each device sees only a few groups, so its local job is simple. The paper argues that this spread of data helps: some pairs of groups can be told apart with a weaker requirement than a single machine would need. The method needs one round of messages, and it still works when some devices drop out.

The benefit is shown under assumptions. The real-data results are measured, not proven.

TL;DR

k-FED clusters data split across devices in one communication round, with no network-wide initialization. Each device runs a Lloyd-style local clustering and sends its few local centers; the server seeds k groups with a farthest-point rule and makes one assignment pass. The analysis needs well-separated cluster means and at most sqrt(k) clusters per device. Under that heterogeneity, pairs of clusters that never share a device need only about k^(1/4) separation, against sqrt(k) for centralized clustering §3.2. The experiments are author-reported and run on synthetic mixtures, FEMNIST, Shakespeare and rotated MNIST §4.

Problem

Lloyd-type k-means converges in polynomial time only under assumptions and with careful initial centers §2. The cited centralized analysis needs center separation of order sqrt(k) for every pair of clusters §1. Parallel and distributed k-means variants usually need many communication rounds §2.

The setting here is federated: data is generated locally, is not identically distributed across devices, and devices can join, leave or fail §1. The analysis assumes that each device holds data from at most k' clusters, with k' = max over devices of k^(z), and that a network is heterogeneous when k' ≤ sqrt(k) §3.2, heterogeneity definition. It also assumes that each device holds at least a 1/m0 share of the points of each target cluster it touches, for a constant m0 > 1 §3.2.

No generative model is assumed. All quantities are defined on the data matrix A and a fixed target clustering T = (T_1, …, T_k) §3.1.

Method

The method has two stages §3.2. In the first, each device solves a local k-means problem with its own cluster count k^(z) using the procedure of alg. 1, and returns its local centers and point groups. In the second, the server collects all device centers, seeds k groups from them, and assigns each center to a group alg. 2. The induced clustering gives every point the group of its device center §3.2, induced-clustering definition.

Notation. Let ||A − C|| be the spectral norm of the difference between the data matrix and the matrix C whose rows are the target means of each point's cluster §3.1. The target-cluster quantities used throughout are:

\phi(\mathcal{T}) = \sum_{j=1}^{k} \sum_{i \in T_j} \lVert A_i - \mu(T_j) \rVert^2 eq. 1

\tilde\Delta_r = \sqrt{k}\,\frac{\lVert A - C \rVert}{\sqrt{n_r}} eq. 2

Two clusters are well separated when their means are at least a constant c times the sum of their Δ̃ values §3.1:

\lVert \mu_r - \mu_s \rVert \ge c\,(\tilde\Delta_r + \tilde\Delta_s) eq. 3

For the federated analysis, define the device-level quantities

\Delta_r = k'\,\frac{\lVert A - C \rVert}{\sqrt{n_r}}, \qquad \lambda = \sqrt{k'}\,\frac{\lVert A - C \rVert}{\sqrt{n_{\min}}} eq. 4

where n_min is the smallest device size §3.2.

Local procedure alg. 1. It projects the device data onto the top k^(z) singular directions, runs a rough clustering there, and forms initial centers from points that are much closer to one rough center than to the others. Lloyd steps on the full data then give the local centers.

# Local k-means on one device  [alg. 1]
# Input:  local data (n, d); local cluster count k_local
# Output: local centers (k_local, d) and point groups

def local_centers(data, k_local):
    # 1. Project onto the top k_local singular directions, then run any
    #    10-approximation k-means there  [alg. 1, step 1]
    projected = project_top_directions(data, k_local)   # (n, k_local)
    rough = approximate_kmeans(projected, k_local)      # (k_local, k_local)

    # 2. Keep the points at least 3 times closer to rough center r than
    #    to any other rough center; their mean is the initial center
    #    [alg. 1, step 2]
    dist = pairwise_distance(projected, rough)         # (n, k_local)
    for r in range(k_local):
        members[r] = points_with_3x_margin(dist, r)
        init[r] = mean(data[members[r]])

    # 3. Run Lloyd steps on the unprojected data until convergence
    #    [alg. 1, step 3]
    centers, groups = lloyd_until_converged(data, init)
    return centers, groups

Federated procedure alg. 2. Each device runs the local procedure once. The server then picks k well-separated seeds by a farthest-point rule and groups all device centers around them with one assignment round.

# k-FED: one-shot clustering across devices  [alg. 2]
# Input:  per-device data (list of (n_z, d) arrays); global count k
# Output: group id per device center; induced label per point

def kfed(devices, k, k_per_device):
    # 1. Each device clusters locally once; nothing else is sent
    #    [alg. 2, step 1]
    device_centers = []
    for z, data in enumerate(devices):
        centers, _ = local_centers(data, k_per_device[z])   # alg. 1
        device_centers.append(centers)
    all_centers = concatenate(device_centers)               # (sum k_z, d)

    # 2. Seed the set M with the centers of one device  [alg. 2, step 2]
    seeds = device_centers[0]
    # 3. Add the center farthest from the seeds until there are k
    #    [alg. 2, steps 3-6]
    while len(seeds) < k:
        far = argmax(distance_to_set(all_centers, seeds))
        seeds = append(seeds, all_centers[far])

    # 4. One assignment round: each center joins its nearest seed
    #    [alg. 2, step 7]
    groups = nearest_seed(all_centers, seeds)               # group per center

    # 5. A point inherits the group of its local center  [§3.2]
    labels = induced_point_labels(devices, device_centers, groups)
    return groups, labels

Hyperparameters

Results

Author-reported throughout. The theory results are statements about the conditions under which the method is guaranteed to work §3.2. The experiments check the theory on synthetic mixtures and measure practical use on real data §4.

Theory §3.2:

• Main result. Under the active and inactive separation conditions, all but O(1/c²)·n points are classified correctly. If each device's points satisfy the proximity condition, all points are correct. The O-constant is not stated.
• Corollary. With k' ≤ sqrt(k), an active pair needs ‖μ_r − μ_s‖ ≥ c·sqrt(m0·k)·(‖A−C‖/sqrt(n_r) + ‖A−C‖/sqrt(n_s)). This matches the centralized requirement up to the sqrt(m0) factor. An inactive pair needs only the same bracket times 10·sqrt(m0)·k^(1/4) §3.2.
• Runtime. Steps 2 to 8 of alg. 2 take O(Z·k'·k²) pairwise distance computations. A new device's local centers can be assigned in O(k'·k) distance computations to the nearest seed.

Synthetic mixtures §4.1, tab. 1. Setup: Gaussian components arranged so that the separation conditions hold by construction; k' = sqrt(k); m0 = 5; c = 100; accuracy averaged over 10 runs. The table caption calls the metric "clustering accuracy" and does not state its scale.

Real data, heterogeneity §4.2, fig. 2. On FEMNIST and Shakespeare, the k-means cost of structured partitions is closer to the oracle cost than the cost of IID random partitions. The gap widens as k' decreases. The text reports this as qualitative; no numeric table is given for these curves.

Real data, communication §4.2, fig. 3. One-shot k-FED gives similar k-means cost to naive distributed k-means, which needs several rounds. The comparison is qualitative in the text.

Personalized federated learning §4.2, tab. 2. Test accuracy on rotated MNIST. The table does not state the metric's scale.

Client selection §4.2, fig. 4. Adding k-FED's cluster filter to the power-of-choice selection (pow-d) speeds up convergence on FEMNIST and Shakespeare, per the text. Per the text, the spread of final accuracies across devices is 35% lower with the combined method than with random selection on FEMNIST. The number of runs behind that figure is not stated.

Paper vs code

Code: official repository https://github.com/metastableB/kfed, commit 39879fcd8edf58f6a6aa48a8abcc99c9c08779e6. The paper links it in §4 §4. The repository has a single package, fkms, with one module README.md algorithms.py. It has no LICENSE file; code is cited by path and line only and nothing is quoted.

• Local rough clustering. The paper uses any 10-approximation [alg. 1, step 1]. The code uses k-means++ seeding and says so in its docstring algorithms.py algorithms.py.
• Projection. The paper uses the top k^(z) singular vectors [alg. 1, step 1]. The code uses randomized SVD with no seed, so the projection is approximate and differs between runs algorithms.py.
• Split rule. The 1/3 distance rule matches alg. 1, step 2 algorithms.py.
• Initial local centers. The paper takes μ(S_r) over points of the data alg. 1, step 2. The code takes the mean of the projected points and maps it back to d dimensions algorithms.py. These differ unless the mean already lies in the projected subspace, which is a difference the paper does not mention.
• Local Lloyd. The paper runs Lloyd steps on the original data until convergence alg. 1, step 3. The code runs sklearn KMeans seeded with those means on the original data and sets no iteration or tolerance limit algorithms.py.
• Devices with few points. When a device has at most k points, the code returns its raw points as centers and skips the local procedure algorithms.py. The paper has no such case.
• Per-device k. The paper uses k^(z) for each device §3.2. The code passes one dev_k to every device algorithms.py.
• Seed set. The paper seeds M with the centers of one device [alg. 2, step 2]. The code takes the first dev_k rows of the concatenated centers algorithms.py. Those rows come from device 0 only when device 0 returns dev_k centers, which breaks when it has few points.
• Farthest-point rule. Matches alg. 2, steps 3-6 algorithms.py.
• Global step. The paper runs one Lloyd round from M alg. 2, step 7. The code runs sklearn KMeans to convergence from the same seeds algorithms.py.
• Outputs. The paper defines point labels through the groups τ §3.2. The code returns the centers and the local estimates, not point labels algorithms.py. Labels must be derived by the caller, for example with nearest-center assignment algorithms.py.
• New devices. The paper's new-device rule §3.2 has no function in the code.
• Experiments. The paper points to the repository for experimental setup §4. The repository has no experiment scripts, data loaders or configs for the reported numbers README.md.

Pitfalls

• Guarantees do not transfer to real data as checked. The separation conditions depend on the unknown target clustering T and cannot be checked from data alone. The real-data experiments build an oracle clustering first §4.2, app. B. The paper says that real-data separation is much smaller than the theory requires §4.2. The real-data results are therefore practical evidence, not a check of the bound.
• Synthetic runs are a best case. The synthetic data is generated so that the separation conditions hold §4.1, so Table 1 does not stress the method.
• Heterogeneity is what helps. When k' is close to k, the active separation requirement is stricter than the centralized one, Ω(k) against Ω(sqrt(k)) §3.2. Use k-FED for heterogeneous partitions and a centralized method for IID-like data.
• Constants. The proofs need c ≥ 100 §3.1, but the O(1/c²) error constant is not stated, so the expected error at a given c cannot be read off the paper.
• k, k^(z) and m0 are inputs. The paper assumes them known §3.2 and does not say how to estimate them. The oracle criterion for k uses an example threshold c_0 = 0.5 and does not state the one used app. B.
• Small devices. If a device holds few points for a cluster, the 1/m0 condition may fail §3.2. The code's few-points branch [code: fkms/algorithms.py#L170-L173] also changes the local output shape, and the seed set may mix devices (see Paper vs code).
• "One round" is two messages per device. The server must send cluster identities back to each device §3.2.
• "One Lloyd round" is ambiguous. The paper does not say whether the seeds in M are replaced by the means of their groups after the assignment. Read it as one nearest-seed assignment alg. 2, step 7, or run to convergence as the code does. The results can differ.
• Farthest-point seeding and outliers. The max rule picks the farthest center as a seed alg. 2. An outlier device center will be chosen. The code's own docstring says farthest-point selection is affected by outliers, although that comment is about a different function [code: fkms/algorithms.py#L81-L83].
• Non-reproducible runs. k-means++ and randomized SVD use unseeded randomness algorithms.py. The paper reports means over 10 runs §4.1. Fix seeds and report variance.
• Selection and personalization details are thin. The number of runs, the rounds, the batch size and the learning rate for the selection study are not stated app. B. The 35% variance figure §4.2 cannot be reproduced from the paper alone.
• Dead code. The n ≤ k branch of kmeans_pp has a typo (np.aray) at algorithms.py. awasthisheffet returns before reaching it, so kfed does not hit it.

Lineage

• Builds on: the deterministic center-separation framework of Kumar and Kannan, and the improved analysis of Awasthi and Sheffet, for a Lloyd variant §2, §3.1. Lloyd's heuristic is the base method §2.
• Compared against: naive distributed k-means, which needs several rounds §4.2; IFCA (cited as ghosh2020), an iterative clustered-FL method §4.2; power-of-choice client selection (cited as cho2020client) §4.2; a single global model tab. 2.
• Superseded by: not known from content in this wiki.