Source-linked AI summary
Clustered Federated Learning: Model-Agnostic Distributed Multi-Task Optimization under Privacy Constraints
Felix Sattler, Klaus-Robert Müller, Wojciech Samek
TL;DR
Federated Learning can perform poorly when private client distributions diverge, because a single model may not fit all clients. The paper introduces Clustered Federated Learning, which clusters clients using cosine similarity of gradient or weight updates. Theoretical guarantees and experiments show improved performance in clustered settings, while preserving privacy and the existing communication protocol.
Problem
Federated Learning may be suboptimal when clients have divergent data distributions that a single model cannot jointly fit.
Method
Clustered Federated Learning uses cosine similarity between client updates at an FL stationary point to infer clusters of jointly trainable clients.
Results
CFL achieves drastic classification-accuracy and perplexity improvements over the Federated Learning baseline when client data has a clustering structure.
Takeaways & Limitations
CFL provides mathematically guaranteed clustering quality for arbitrary non-convex objectives while preserving privacy and requiring no FL communication-protocol changes.
Abstract
from arXiv · showhide
Federated Learning (FL) is currently the most widely adopted framework for collaborative training of (deep) machine learning models under privacy constraints. Albeit it's popularity, it has been observed that Federated Learning yields suboptimal results if the local clients' data distributions diverge. To address this issue, we present Clustered Federated Learning (CFL), a novel Federated Multi-Task Learning (FMTL) framework, which exploits geometric properties of the FL loss surface, to group the client population into clusters with jointly trainable data distributions. In contrast to existing FMTL approaches, CFL does not require any modifications to the FL communication protocol to be made, is applicable to general non-convex objectives (in particular deep neural networks) and comes with strong mathematical guarantees on the clustering quality. CFL is flexible enough to handle client populations that vary over time and can be implemented in a privacy preserving way. As clustering is only performed after Federated Learning has converged to a stationary point, CFL can be viewed as a post-processing method that will always achieve greater or equal performance than conventional FL by allowing clients to arrive at more specialized models. We verify our theoretical analysis in experiments with deep convolutional and recurrent neural networks on commonly used Federated Learning datasets.
I. INTRODUCTION
Federated Learning assumes one model can fit all clients, but divergent conditional distributions or limited model expressiveness can violate this assumption. Clustered Federated Learning addresses this by identifying client groups through cosine similarity of gradient updates, with theoretical guarantees for correct separation.
- Federated Learning: Federated Learning jointly trains one model across private client data through iterative client updates and server aggregation.Clients download a master model, perform local training, and upload updates for weighted averaging.
- Motivation: The standard FL assumption fails when clients have divergent conditional distributions or the model cannot express all distributions simultaneously.These violations include differing labeling preferences, heterogeneous text statistics, and insufficient model complexity.
- Motivation: Ordinary FL cannot provide every client with a locally optimal model when incongruent client distributions are treated equally under one global model.The paper frames this as a Federated Multi-Task Learning problem involving more specialized client models.
- Clustered Federated Learning: Clustered Federated Learning assumes the client population can be partitioned so that each cluster satisfies the conventional Federated Learning assumption.This generalizes the single-population assumption to jointly trainable subsets of clients.
- Method: At an FL stationary point, cosine similarity between client gradient updates can reveal whether clients share a data-generating distribution.The approach infers clustering without direct access to client data or metadata.
- Theory: Theorem 1 provides a client bi-partition with bounded cross-cluster update similarity, while the cosine criterion can remain correct with many distributions and noisy empirical risks.The latter result supports practical clustering when empirical risks only loosely approximate true risks.
A. Distinguishing Congruent and Incongruent Clients
CFL distinguishes congruent from incongruent clients using clustering criteria based on client updates after convergence. This avoids splitting clients when separate models would harm performance and retains the conventional FL solution as a fallback.
- Separating clients with varying distributions can degrade performance in congruent non-iid Federated Learning because clusters exchange less knowledge.
- CFL evaluates clustering after Federated Learning reaches a stationary point, when clients’ updates can reveal whether their distributions are incongruent.
- If splitting degrades performance, CFL falls back to the Federated Learning solution, so it performs at least as well as conventional FL.
III. CLUSTERED FEDERATED LEARNING
CFL first obtains a stationary Federated Learning solution, then recursively partitions clients using update similarities and a stopping criterion. The procedure continues until the identified groups satisfy the conditions for mutually congruent clients.
- CFL begins with Federated Learning and uses the resulting stationary solution before evaluating whether client clustering is necessary.
- When clients are incongruent, the server computes pairwise cosine similarities of recent updates and separates clients by minimizing cross-cluster similarity.
- The optimal bipartition can be solved in O(m3), with clustering overhead typically negligible because the server has greater computational power than clients.
- The correctness criterion can be evaluated using estimated intra-cluster similarity and cross-cluster similarity, independently of the number of distributions k.
- CFL recursively reapplies the procedure to separated groups until no sub-cluster violates the stopping criterion, identifying groups of mutually congruent clients.
IV. RELATED WORK
Related work covers Federated Learning under congruent distributions, centralized multi-task learning, and federated clustering methods. CFL differs by supporting arbitrary non-convex objectives and providing theoretical criteria for adaptive clustering.
- Federated Learning commonly assumes that one central model can fit all clients’ distributions simultaneously, with prior work studying congruent iid and non-iid settings.
- Centralized multi-task learning addresses incongruent data, while MOCHA extends multi-task learning to the Federated Learning setting.
- Compared with a prior clustering method, CFL uses cosine similarity rather than l2-distance, avoids convexity and separated-minima requirements, and can distinguish congruent from incongruent settings.
- CFL is adaptive because clustering need not be decided after the first communication round, and it applies to arbitrary Federated Learning problems with non-convex objectives.
V. IMPLEMENTATION CONSIDERATIONS
CFL generalizes gradient-based similarity analysis to the weight-updates already communicated in Federated Learning. Experiments report that update-based cosine similarities can separate clients better than gradient-based similarities.
- Federated Learning communicates weight-updates rather than gradients because of client memory and communication constraints.
- For sufficiently smooth losses and low learning rates, one-epoch weight-updates approximate the direction of the true gradient.
- CFL therefore computes cosine similarities between weight-updates instead of gradients to remain close to the classical Federated Learning algorithm.
- Experiments found that cosine similarities based on weight-updates achieved even better separations than similarities based on gradients.
B. Preserving Privacy
CFL preserves privacy by transforming updates before communication while retaining the unchanged Federated Learning protocol. Its parameter tree supports specialized models, new-client assignment, and client populations that vary over time.
- CFL applies a transformation to client updates before communication, after which the server averages them and clients apply the inverse operation.
- The privacy-preserving protocol can resume unchanged and remains compatible with encryption, unlike multi-task approaches requiring direct access to client data.
- CFL is flexible enough to handle client populations that vary over time.
- A parameter tree stores clusters, their stationary solutions, and pre-split child updates, with the conventional Federated Learning solution at its root.
- New clients are assigned to leaf clusters by traversing the tree along the branch whose cached client updates are most similar.
- The tree provides models with varying specificity, from the general root model to specialized leaf models, and can support model ensembling.
VI. EXPERIMENTS
The experiments test whether cosine similarities can recover client clusters under practical departures from ideal conditions. Across MNIST and CIFAR label-swap settings, weight-updates provide effective separation with limited data and training.
- The theoretical cosine-similarity criterion assumes convergence, sufficient data, and full-gradient computation.
- CFL can correctly infer clustering even when clients have small datasets and training reaches only an approximately stationary solution.
- Experiments use MNIST and CIFAR-10 with 20 clients in 4 clusters, created by swapping two labels according to each client’s cluster.
- Number of Data points: 20 data points per MNIST client already produce a positive separation gap when weight-updates compute similarity.
- Proximity to Stationary Solution: Separation quality increases monotonically with communication rounds, and correct clustering is achieved after around 10 rounds on both MNIST and CIFAR.
- Weight-Updates instead of Gradients: Weight-updates yield better separation with fewer data points and farther from a stationary solution than gradients, while leaving the communication protocol unchanged.
B. Distinguishing Congruent and Incongruent Clients
The experiments show that norm criteria distinguish congruent from incongruent clients after Federated Learning reaches a stationary point, while CFL then separates incongruent clients and improves local performance. Across CIFAR and Ag-News, clustering yields substantially better accuracy or perplexity than conventional Federated Learning.
- Norm criteria: In the congruent case, client-update norms decrease with the server-update norm, whereas in the incongruent case they stagnate or increase.In both settings, the server update norm approaches zero, indicating convergence to a stationary point.
- CIFAR-10 label permutation: After 50 rounds on CIFAR with 20 clients and four label permutations, CFL begins splitting clients as separation gaps emerge.Further splits occur at rounds 100 and 150 until all incongruent distributions are separated.
- CIFAR-10 label permutation: CIFAR validation accuracy more than doubled to close to 60% after CFL finalized clustering, compared with the conventional Federated Learning solution.The first split produced an immediate 25% accuracy increase for the separated purple cluster.
- Ag-News language modeling: On Ag-News, CFL separated clients in rounds 30, 60, and 90, reducing local test perplexity below 36 while Federated Learning stagnated at 42.Lower perplexity indicates better language-model performance in this experiment.
- Conclusion: CFL is presented as a post-processing framework that improves existing Federated Learning by enabling clients to learn more specialized models.The authors report applicability to convolutional and recurrent neural networks, privacy-preserving implementation, and no required communication-protocol modification.
A. Proving the Separation Theorem
The Separation Theorem uses geometric properties of noisy client gradients to construct a binary partition with bounded cross-cluster similarity. Its proof reduces the vector configuration to two dimensions and establishes a separation guarantee under the theorem’s assumptions.
- Proof strategy: The proof combines bounds for noisy approximations of identical and different vectors with a geometric partitioning lemma.These steps bound intra-cluster similarity from below and cross-cluster similarity from above.
- Two-dimensional geometry: For d = 2, partitioning at the largest and second-largest neighboring angular gaps guarantees a cross-cluster angle at least as large as the second-largest gap.The worst case occurs when the largest gap is maximized and the remaining gaps are equal.
- Two-dimensional geometry: The largest neighboring angle cannot exceed π, yielding the worst-case lower bound π/(k − 1) for the separating angle.This follows because a larger gap would place all vectors on one side of a line through the origin, contradicting the required zero-sum condition.
- Higher-dimensional reduction: For d > 2, projection onto the plane spanned by the closest pair of vectors reduces the problem to the proven two-dimensional case without decreasing the relevant separation.The projected configuration preserves an optimal clustering with the same minimum cross-cluster angle.
- Theorem statement: The theorem seeks a client bipartition whose maximum cross-cluster gradient similarity is bounded below the within-cluster similarity.The result is stated for clients sampled from k data-generating distributions at stationary Federated Learning solutions.
- Theorem consequence: Under the theorem’s sampling and empirical-risk assumptions, the resulting clustering is guaranteed to be correct.The guarantee applies at stationary solutions of the Federated Learning objective.