Source-linked AI summary
Federated learning with hierarchical clustering of local updates to improve training on non-IID data
Christopher Briggs, Zhong Fan, Peter Andras
TL;DR
Non-IID client data can make a single federated model suboptimal. The paper introduces hierarchical clustering of client updates to train specialised models, finding faster convergence and more clients reaching target accuracy in many non-IID settings.
Problem
Non-IID data makes a single joint federated model unreliable or suboptimal for clients with differing distributions.
Method
FL+HC clusters clients by similarity of their local model updates and trains the resulting groups independently on specialised models.
Results
In 2 of 3 non-IID settings, FL+HC converges more quickly and enables up to 2x more clients to reach a target accuracy than standard FL.
Takeaways & Limitations
Clustering client updates can improve average accuracy and the percentage of clients performing well in most evaluated circumstances.
Abstract
from arXiv · showhide
Federated learning (FL) is a well established method for performing machine learning tasks over massively distributed data. However in settings where data is distributed in a non-iid (not independent and identically distributed) fashion -- as is typical in real world situations -- the joint model produced by FL suffers in terms of test set accuracy and/or communication costs compared to training on iid data. We show that learning a single joint model is often not optimal in the presence of certain types of non-iid data. In this work we present a modification to FL by introducing a hierarchical clustering step (FL+HC) to separate clusters of clients by the similarity of their local updates to the global joint model. Once separated, the clusters are trained independently and in parallel on specialised models. We present a robust empirical analysis of the hyperparameters for FL+HC for several iid and non-iid settings. We show how FL+HC allows model training to converge in fewer communication rounds (significantly so under some non-iid settings) compared to FL without clustering. Additionally, FL+HC allows for a greater percentage of clients to reach a target accuracy compared to standard FL. Finally we make suggestions for good default hyperparameters to promote superior performing specialised models without modifying the the underlying federated learning communication protocol.
I. INTRODUCTION & BACKGROUND
Federated learning trains a joint model from distributed client data, but non-IID data can make a single model suboptimal. FL+HC clusters clients by update similarity to train specialised models, with empirical guidance on its hyperparameters.
- Federated learning: Federated learning aggregates local client training to learn a joint statistical model over massively distributed data.Clients train locally and return model updates for central aggregation.
- Motivation: Non-IID data can reduce the suitability of a single joint model for clients with different data distributions.The paper motivates separate models for client groups with similar update behaviour.
- FL+HC: FL+HC clusters clients by similarity of their updates to the global model before training specialised models for client subsets.The clustering step is inserted after a set number of communication rounds.
- Expected benefits: FL+HC is designed to increase test accuracy while reducing communication rounds needed to reach convergence.Its underlying federated-learning communication protocol is not modified.
- Evaluation: The study characterises hierarchical-clustering effects across IID and non-IID settings and evaluates FL and clustering hyperparameters.The contribution includes recommendations for default hyperparameters when the data distribution is unknown.
B. Client statistical heterogeneity
Client data may differ in feature, label, or conditional distributions, making local objectives unreliable approximations of a joint objective. Clustering model updates provides a privacy-preserving proxy for identifying related clients.
- Challenge: Non-IID client data violates the assumption that all clients sample from the same underlying distribution.The paper denotes this heterogeneity as P_i ≠ P_j across client pairs.
- Challenge: Under non-IID data, a local model trained on client data can be an unreliable approximation of the joint model.This creates a fundamental challenge for training one joint model.
- Forms of heterogeneity: Non-IIDness includes feature-distribution skew, label-distribution skew, and concept shifts involving labels or features.Clients may also have different numbers of training examples.
- Rationale: When client distributions differ substantially, multiple models targeted to groups of similar clients may better meet their objectives.The central entity cannot inspect raw data, so local model updates serve as the similarity signal.
C. Hierarchical clustering
FL+HC uses agglomerative hierarchical clustering on vectorised client updates to form specialised client groups without requiring a preset cluster count. The study evaluates distance metrics, linkage mechanisms, and resulting model accuracy.
- Choice of method: Hierarchical clustering is used because it can determine the number of clusters independently while assigning every client to a relevant cluster.It is also described as scalable and reasonably interpretable.
- Procedure: Agglomerative clustering begins with singleton client-update clusters and repeatedly merges the most similar pair until one cluster remains.Each sample is a vectorised local model update.
- Distance metrics: The evaluated distance metrics are L1 Manhattan, L2 Euclidean, and cosine distance.Cosine distance is invariant to scaling and compares update directions.
- Linkage: Linkage determines inter-cluster similarity using single, complete, average, or Ward’s linkage.Ward’s linkage is restricted to Euclidean distance and minimises intra-cluster variance when merging.
- Evaluation: Experiments assess whether averaged clustered updates achieve higher test accuracy than a single joint model.The comparison uses the normal federated-learning setting as the baseline.
D. Related work
Prior work addresses non-IID federated learning through data sharing, multi-task learning, and clustering approaches. This paper is positioned as closely resembling work that clusters clients for specialised federated models.
- Non-IID federated learning: Prior federated-learning research has proposed methods for mitigating non-IID data, including sharing a small data subset.The cited discussion contrasts these approaches with standard FederatedAveraging.
- Related approaches: Multi-task learning treats local learning objectives, or groups of related clients, as separate tasks in distributed learning.The paper also cites K-means clustering of client updates for isolating Byzantine clusters.
II. FEDERATED LEARNING WITH HIERARCHICAL
FL+HC introduces a clustering step during federated training to separate clients with non-iid objectives and train specialised models for similar clients. The clustering operation is computationally manageable because it runs once on the server.
- FL+HC adds a clustering step at communication round n after clients update the global joint model.The resulting client groupings are used to distinguish non-iid objectives during training.
- Each discovered cluster is trained independently on its own specialised model rather than sharing one joint model.The specialised model for cluster c is denoted fc(w).
- O(n3) server-side clustering is performed once, so its impact on the overall training operation is described as not significant.
A. Experiment setup
The experiments evaluate FL and FL+HC on MNIST and FEMNIST under iid and several non-iid data partitions. They vary client participation and clustering-related settings while using a common client training procedure and reporting accuracy-based outcomes.
- MNIST provides a 60,000-example training set and a 10,000-example test set for evaluating handwritten-digit classification.
- The study compares iid, pathological non-iid, and label-swapped non-iid partitions, including clients with restricted labels or swapped digit meanings.
- The client update procedure uses mini-batch gradient descent across local epochs before returning the updated model to the server.
- The label-swapped partition creates four groups distributed evenly across 25 clients, yielding 100 clients with 600 training examples each.This construction produces four natural clusters under concept shift.
- FEMNIST uses 367 clients formed from users with 12 to 386 samples each, testing feature distribution skew and uneven client sample sizes.
- Each client's test data is drawn from the same distribution as its training data across all partitioning schemes.
- FederatedAveraging trains the baseline for 50 rounds with client fractions 0.1, 0.2, 0.5, and 1.0.Client training uses mini-batch SGD with batch size 10, 3 local epochs, and learning rate 0.1.
- FL+HC first trains a global model for n rounds, then clusters client update vectors produced after 3 additional local epochs.The resulting clusters continue training independently with FederatedAveraging.
III. RESULTS & DISCUSSION
The evaluation measures average client test accuracy and the percentage of clients reaching target accuracy after clustering and at round 50. The reported statistics indicate that similarity-based clustering is beneficial in most circumstances.
- Average test set accuracy is reported immediately after clustering and at round 50, when unclustered FederatedAveraging begins to plateau.
- The percentage of clients reaching 99% test accuracy, or 80% in the FEMNIST non-iid setting, is also reported.
- Clustering clients by similarity beneficially affects average client test accuracy and the percentage of clients performing well in most circumstances.
A. Effect of varying client fraction and number of rounds prior to clustering
FL+HC’s sensitivity to client fraction and pre-clustering rounds depends on the data distribution. More pre-clustering rounds generally help identify useful clusters, while clustering can harm performance when client similarity is poorly estimated.
- Experimental setup: The experiments vary client fraction and the number of FL rounds before clustering to assess their effects on cluster quality and training performance.Client fractions tested include 0.1, 0.2, 0.5, and 1.0.
- Iid setting: In iid data, clustering usually falls back to the single joint model, while clustering after one round slightly reduces average test accuracy.The result is attributed to stochastic client optimization hindering similarity estimation before sufficient training iterations.
- Pathological non-iid setting: 1.3x–1.9x: pathological non-iid FL+HC’s jump in test accuracy in round 2 after clustering, although round-50 final accuracy matches FL.At round 50, the percentage of clients reaching 99% accuracy improves by 1.2x–2.0x over FL.
- Label-swapped non-iid setting: In label-swapped non-iid data, FL+HC substantially outperforms FL, reaching near-iid average accuracy and training up to 80% of clients to 99% accuracy.Increasing pre-clustering rounds improves the number of clients reaching the target, while client fraction has negligible effects on recorded metrics.
- FEMNIST non-iid setting: In FEMNIST, clustering is usually slightly worse than FL and can drastically reduce accuracy when it creates clusters of dissimilar clients.Higher client fractions and more pre-clustering rounds generally benefit FL+HC in this setting.
- Overall pattern: Across non-iid settings, increasing pre-clustering rounds can produce greater gains over FL, while client fraction generally has only a small positive effect.The clearest benefit of more pre-clustering rounds occurs in the label-swapped setting, where FL+HC also lets more clients reach target accuracy per communication round.
B. Effect of varying hierarchical clustering hyperparameters
FL+HC’s distance metric strongly affects clustering outcomes across non-IID settings, while linkage choice generally has little effect. Manhattan distance is strongest for pathological non-IID data, cosine for label-swapped data, and Euclidean for FEMNIST.
- Experimental setup: FL+HC hyperparameter experiments fix 10 communication rounds before clustering and a 0.2 client fraction, then compare distance thresholds, metrics, and linkage methods.The experiments report the best-performing threshold for each hyperparameter combination and data partitioning scheme.
- IID setting: In iid data, Euclidean and cosine fail to split clients, reproducing standard FL, whereas Manhattan clustering reduces test accuracy and target-accuracy coverage.The paper notes that perfectly iid data is unlikely in real-world distributed datasets.
- Pathological non-IID setting: A 1.2x–2.1x increase in clients reaching 99% target accuracy occurs with FL+HC versus FL by round 50 in pathological non-IID data.All tested hyperparameter combinations achieve final test accuracy similar to or better than a single joint model.
- Pathological non-IID setting: Manhattan distance performs best in pathological non-IID data, with test accuracy after clustering at round 10 exceeding FL’s round-50 iid accuracy and communication reduced by >5x.The paper links this result to Manhattan distance’s suitability for sparse high-dimensional update vectors.
- Label-swapped non-IID setting: In label-swapped non-IID data, cosine distance performs best because update direction matters more than magnitude, while linkage choice has little effect.Only cosine-based experiments reach the performance of FL on iid data.
- FEMNIST non-IID setting: FEMNIST shows little overall improvement over FL; cosine degrades performance, Euclidean gives the best final test accuracy, and Manhattan lets 1.1x more clients reach target accuracy by round 50.The linkage mechanism again has less effect than the distance metric.
- Overall findings: Across non-IID settings, the best distance metric depends on the data distribution, whereas linkage has comparatively little influence on measured performance.The paper recommends adjusting hierarchical-clustering hyperparameters when the non-IID structure is unknown.
C. Future work
The authors identify privacy noise, update compression, larger models and datasets, and adversarial clients as directions requiring further study. They also limit the demonstrated evidence to a simple CNN classifying handwritten digits on simulated distributed data.
- Privacy and communication: Differential-privacy noise may affect FL+HC’s ability to find good client clusterings, motivating future evaluation with noisy updates.The current analysis uses full and clean client weight updates.
- Privacy and communication: The effect of compression methods designed to reduce client-update payloads on FL+HC remains unexplored.Compression could alter the updates used for clustering.
- Scope and scalability: The demonstrated method uses a simple CNN for handwritten-digit classification on simulated distributed data, so scaling to larger networks and datasets still requires confirmation.The authors express confidence about scaling but state that further work is required.
- Adversarial settings: FL+HC was not evaluated in the presence of adversaries, including whether clustering could help identify malicious clients.The paper presents adversarial-client analysis as a promising research avenue.
IV. CONCLUSION
The paper introduces FL+HC, which clusters clients by the similarity of their model updates to train specialised models. Across simulated iid and non-IID image-classification settings, it can accelerate convergence and increase target-accuracy coverage, with performance depending on the non-IID structure.
- Contribution: FL+HC adds a clustering step to the federated-learning protocol to group clients by the similarity of their model weight updates and train specialised models.The method targets distributed datasets without changing the underlying federated-learning communication protocol.
- Empirical findings: In 2 of 3 non-IID settings, FL+HC converges faster and enables up to 2x more clients to reach a target accuracy by the end of training.Under iid data, FL+HC matches FL except when clustering occurs after one round.
- Empirical findings: FL+HC reduces communication rounds by >5x in some non-IID settings when using the Manhattan distance metric.Different distance metrics perform best for different non-IID data structures.
- Practical guidance: The paper provides recommendations for default FL and hierarchical-clustering hyperparameters when the non-IID structure of the data is unknown.The recommendations are based on experiments varying client participation, pre-clustering rounds, distance metrics, and linkage mechanisms.