An Efficient Framework for Clustered Federated Learning

Avishek Ghosh; Jichan Chung; Dong Yin; Kannan Ramchandran

Source-linked AI summary

An Efficient Framework for Clustered Federated Learning

Avishek Ghosh, Jichan Chung, Dong Yin, Kannan Ramchandran

arXiv:2006.04088v2Jun 7, 2020stat.MLcs.LG

TL;DR

The paper studies federated learning with users partitioned into unknown clusters, where each cluster represents a distinct learning objective. It proposes IFCA, which alternates cluster assignment estimation with cluster-model optimization, and proves convergence under good initialization while demonstrating empirical effectiveness, including with random restarts and neural networks.

Problem
Federated learning must learn from heterogeneous user data while exploiting within-cluster similarity, despite unknown cluster memberships and distinct cluster objectives.
Method
IFCA alternates between assigning users to the lowest-loss cluster model and optimizing the cluster models in a distributed setting.
Results
IFCA has convergence guarantees for linear and strongly convex losses under good initialization, and experiments show effectiveness on clustered FL benchmarks, neural networks, and random initialization with multiple restarts.
Takeaways & Limitations
Clustered FL provides a framework for jointly learning user memberships and personalized cluster models while retaining theoretical guarantees and practical performance across convex and non-convex settings.
Takeaways & Limitations
The theoretical guarantees assume good initialization, while ambiguous cluster structure requires combining IFCA with weight sharing for practical neural-network training.

Abstract

from arXiv · show

We address the problem of federated learning (FL) where users are distributed and partitioned into clusters. This setup captures settings where different groups of users have their own objectives (learning tasks) but by aggregating their data with others in the same cluster (same learning task), they can leverage the strength in numbers in order to perform more efficient federated learning. For this new framework of clustered federated learning, we propose the Iterative Federated Clustering Algorithm (IFCA), which alternately estimates the cluster identities of the users and optimizes model parameters for the user clusters via gradient descent. We analyze the convergence rate of this algorithm first in a linear model with squared loss and then for generic strongly convex and smooth loss functions. We show that in both settings, with good initialization, IFCA is guaranteed to converge, and discuss the optimality of the statistical error rate. In particular, for the linear model with two clusters, we can guarantee that our algorithm converges as long as the initialization is slightly better than random. When the clustering structure is ambiguous, we propose to train the models by combining IFCA with the weight sharing technique in multi-task learning. In the experiments, we show that our algorithm can succeed even if we relax the requirements on initialization with random initialization and multiple restarts. We also present experimental results showing that our algorithm is efficient in non-convex problems such as neural networks. We demonstrate the benefits of IFCA over the baselines on several clustered FL benchmarks.

1 Introduction

The paper addresses heterogeneous, non-i.i.d. data in federated learning by modeling users as unknown clusters with shared learning objectives. It proposes IFCA, which alternates cluster assignment estimation and cluster-model optimization, and reports theoretical and empirical support.

Motivation: Federated learning must handle naturally arising data heterogeneity because users’ devices hold data from different distributions.Treating all users’ data as i.i.d. can undermine personalized predictions, while individual devices may lack enough data.
Problem and approach: Clustered federated learning partitions users into groups and trains a model for each group, but cluster identities are unknown.The resulting problem jointly requires identifying memberships and optimizing distributed cluster models.
Problem and approach: IFCA alternates between estimating users’ cluster identities and minimizing cluster-specific loss functions.The paper views this procedure as alternating minimization in a distributed setting.
Theory and evidence: IFCA has convergence guarantees for linear models and general strongly convex losses under good initialization, with exponential convergence and near-optimal statistical error in certain regimes.For two linear-model clusters, convergence is guaranteed when initialization is slightly better than random.
Theory and evidence: Experiments report success with random initialization and multiple restarts, neural networks, and clustered FL benchmarks based on MNIST, CIFAR-10, and Federated EMNIST.These experiments broaden evaluation beyond the good-initialization theoretical setting.

2 Related Work

The paper situates clustered federated learning among global-model, personalized-model, and latent-variable approaches to non-i.i.d. federated data. Its analysis extends convergence insights for alternating methods to a distributed setting.

Positioning and contribution: Relative to concurrent work by Mansour et al., this paper establishes population-loss convergence rates under good initialization.The authors state that this simultaneously guarantees convergence of training loss and generalization to test data, whereas the concurrent work provided only generalization guarantees.
Federated learning and non-i.i.d. data: Federated learning research addresses non-i.i.d. user data through global-model and personalized-model formulations.MOCHA, for example, formulates multi-task learning with a user-correlation regularizer.
Latent-variable problems: Clustered federated learning treats user cluster identities as latent variables in a distributed statistical estimation problem.This connects the formulation to classical problems including Gaussian mixtures and mixtures of linear regressions.
Latent-variable problems: Expectation maximization and alternating minimization are established approaches for latent-variable problems, but finite-sample convergence analysis is difficult because the optimization landscape is non-convex.Prior centralized results often require a suitable starting point for fast convergence.

3 Problem Formulation

The formulation models users as machines grouped into unknown clusters, with each cluster associated with its own population loss and target model. Learning uses finite-sample empirical losses to estimate these cluster-specific solutions.

Distributed setting: The distributed system contains one center machine and m worker machines, with each worker corresponding to a federated-learning user.The machines communicate through a predefined protocol.
Cluster structure: Workers are partitioned into k clusters, but the partition and each machine’s cluster identity are not revealed to the learning algorithm.The unknown partition creates the latent assignment component of the problem.
Data and distributions: Each worker in cluster j holds n i.i.d. data points drawn from distribution D_j.Each data point consists of a feature and response pair.
Objectives: The paper defines a loss f(θ; z) over parameter space Θ = R^d and seeks to minimize each cluster’s population loss.For theoretical analysis, strongly convex losses provide guarantees for estimating the unique minimizer of each population loss.
Empirical objectives: Because only finite data are available, the formulation uses empirical losses computed from subsets of each worker’s data.F_i(θ) denotes an empirical loss associated with some or all data on worker i.

4 Algorithm

IFCA iteratively estimates worker cluster identities while optimizing cluster-specific models, addressing limitations of one-shot clustering. It also supports weight sharing for ambiguous neural-network clustering structures and can reduce communication after assignments stabilize.

One-Shot Clustering: One-shot clustering can lock in incorrect assignments, burden the center machine, and restrict the initial local models to settings such as convex losses.IFCA addresses these concerns by improving worker clustering over multiple iterations rather than retaining a single initial estimate.
Iterative Federated Clustering Algorithm (IFCA): IFCA alternates cluster-identity estimation with loss minimization through gradient averaging or model averaging.Workers identify the model with the lowest local loss, then return gradients or locally updated models for cluster-specific aggregation.
Iterative Federated Clustering Algorithm (IFCA): IFCA offers gradient averaging and model averaging, with model averaging allowing multiple local gradient steps before workers send models back.The algorithm specification includes a step size, initialization for each cluster model, parallel iterations, and a local-step count for model averaging.
Iterative Federated Clustering Algorithm (IFCA): In each iteration, the center broadcasts cluster models to participating workers, which select memberships, perform local updates, and return updates for within-cluster aggregation.Model averaging follows the server broadcast, local identification and updates, worker-to-server transmission, and averaging within each estimated cluster.
Practical Implementation of IFCA: Weight sharing combines IFCA with shared early neural-network layers when clusters have common data properties but distinct distributions.The shared layers learn a representation from all available data, while IFCA operates on the last or last few layers; this can also reduce communication cost.
Practical Implementation of IFCA: When cluster assignments remain stable for several iterations, the center can send each worker only its assigned model instead of all k models.This communication reduction is enabled by monitoring whether workers’ estimated cluster identities change across parallel iterations.

5 Theoretical Guarantees

Theoretical analysis shows that IFCA converges under suitable initialization and regularity conditions for linear and strongly convex losses, with statistical error rates that are near optimal in stated regimes.

Algorithmic setup: IFCA alternates cluster-identity estimation with gradient descent updates of cluster-specific parameters, using fresh data samples in the theoretical analysis.The analysis assumes all workers participate in every round and uses separate samples for estimating identities and computing gradients; re-sampling is for tractability and is not required in practice.
Linear model: The linear-model analysis assumes Gaussian features, Gaussian noise, squared loss, separated cluster parameters, and an initialization close enough to the ground truth.For two clusters, the initialization requirement can be only slightly better than random; the assumptions also impose conditions on sample size, cluster proportions, dimension, and separation.
Linear model: IFCA iterates are contractive in the linear model, and the convergence accelerates as parameter closeness improves because cluster-identity estimation becomes more accurate.The error recursion contains a term associated with cluster misclassification, while the closeness parameter increases over iterations.
Linear model: After an initial phase that increases parameter closeness, the linear-model guarantee reaches a statistical error whose leading term is near the minimax rate up to logarithmic and p-dependent factors.The analysis states that the second error term decays exponentially in the per-worker sample size, making the final statistical error rate near optimal.
Strongly convex losses: For general strongly convex losses, IFCA also converges under assumptions including good initialization, smoothness, strong convexity, sufficient data, and adequate cluster separation.The resulting accuracy has no explicit dependence on the initial closeness parameter after the initial phase, and the stated rate is near optimal when n is at least on the order of m.
Strongly convex losses: In the strongly convex setting, the statistical error is near optimal in the regime n ≳ m because its 1/√(mn) component matches the optimal rate even when cluster identities are known.Compared with the linear model, the additional term is 1/n rather than exp(−n), reflecting the weaker bounded-variance assumption instead of Gaussian noise.

6 Experiments

Experiments evaluate IFCA on synthetic, rotated-image, and FEMNIST settings. Results show improved success with stronger separation and more data, efficient cluster identification, and better accuracy than global and local baselines.

Synthetic Data: Random initialization with a few restarts can suffice for linear-model experiments.The empirical study relaxes the theoretical initialization requirement.
Synthetic Data: Success probability increases with parameter separation and with more workers and samples per worker.The experiments use R as the separation parameter and vary m and n while holding other quantities fixed.
Rotated MNIST and CIFAR: IFCA gradually identifies worker clusters and achieves higher test accuracy than global and local model baselines.The global model combines heterogeneous distributions, while the local model can overfit individual data.
Rotated MNIST and CIFAR: Cluster identities are identified after about 30 communication rounds on Rotated MNIST and 10 on Rotated CIFAR.Figure 4 tracks cluster-identity estimation accuracy across parallel iterations.
Federated EMNIST: On FEMNIST, IFCA uses weight sharing and shows a clear advantage over global and local model approaches.The experiments vary the number of clusters and compare IFCA with centralized clustering and the two baseline schemes.

7 Conclusions and Future Work

The paper addresses clustered federated learning with an iterative algorithm and establishes convergence guarantees for strongly convex and smooth losses. Experiments extend the approach beyond the convex regime, while future work targets broader loss classes, stochastic gradients, and partial participation.

The paper proposes an iterative algorithm for clustered federated learning.
Convergence guarantees are established for strongly convex and smooth functions.
Experiments use random initialization with multiple restarts and show efficient performance beyond the convex regime.
Future directions include weakly convex and non-convex analysis, stochastic worker gradients, and small subsets of participating devices.

A Proof of Theorem 1

The proof analyzes one IFCA iteration by separating cluster-assignment errors from gradient-update errors. It combines concentration, Gaussian-noise, and probability bounds to establish high-probability control under the stated assumptions.

The proof studies a single iteration using estimated cluster sets and current model parameters.Estimated clusters are denoted S_j, while the true clusters are S*_j.
Cluster identity errors are analyzed by comparing losses under the true and competing cluster parameters.The argument models feature matrices as Gaussian and observation noise as Gaussian with variance σ^2.
The gradient-update error is decomposed into terms involving covariance deviation and feature-noise interaction.The proof concatenates data from correctly assigned workers and bounds ||X^T X − E[X^T X]||op and ||X^Tϵ||.
Gaussian covariance concentration and a noise-concentration lemma provide high-probability bounds for the update terms.The analysis first establishes that the effective sample size is sufficiently large, then controls the resulting terms.
A union bound extends the first-cluster argument to all clusters while retaining success probability 1 − 1/poly(m).The proof chooses a step size and uses the bound k ≤ m to simplify the final probability.

A.2 Proof of Lemma 2

The lemma bounds the feature-noise interaction X^Tϵ for independent Gaussian feature and noise variables. The proof conditions on X, applies Hanson–Wright and chi-squared concentration, and obtains an exponential tail bound.

The proof conditions on X and applies the Hanson–Wright inequality to control ||X^Tϵ||.
Chi-squared concentration and operator-norm control yield probability bounds decaying with min{d, N}.

B Proof of Corollary 1

The proof shows that IFCA’s alignment parameter reaches at least 1/4 after a constant number of iterations under the stated assumptions, after which the final accuracy is O(ε).

The initialization and separation assumptions provide the conditions needed to bound the iteration error and raise the alignment parameter.The proof uses 0 < α_0 < 1 and separates the relevant error terms into two sufficient conditions.
IFCA reaches α_t ≥ 1/4 after a constant number of iterations.The proof establishes this once T′ is at least 2.
With T = log(∆/4ε), the final accuracy is O(ε) with high probability.

C Proof of Theorem 2

The proof of Theorem 2 controls cluster misclassification and gradient-update errors separately, then combines their probability bounds to obtain a high-probability convergence guarantee.

The proof analyzes erroneous cluster identity estimation through pairwise error events and bounds their probabilities using the separation condition.For the target cluster, the objective gap is at least αλ∆^2, yielding probability bounds for incorrect comparisons.
Gradient-descent error is decomposed into separate terms and controlled using strong convexity, smoothness, Markov’s inequality, and concentration bounds.The resulting bounds are combined across cluster-update components and classification events.
The correctly assigned workers provide a sufficiently large intersection with the target cluster with probability at least 1 − 2 exp(−cpm).The proof lower-bounds this intersection by 1/4 p_1m.
The complete iteration succeeds with probability at least 1 − δ after combining the individual failure probabilities and imposing the sample-size condition.The proof aggregates δ_0, kδ_1, δ_2, and 2 exp(−cpm), then concludes the bound under p ≳ 1/m log(mn′).

D Proof of Corollary 2

The proof of Corollary 2 first increases the alignment parameter to at least 1/4, then runs additional iterations to reduce the error floor while controlling the overall failure probability.

The initial phase uses the separation and initialization conditions to ensure that the alignment-update error is sufficiently small.
The alignment sequence {α_0, α_1, …} is non-decreasing, and IFCA reaches α_T′ ≥ 1/4 after T′ iterations.
After α_t ≥ 1/4, the proof runs another T′′ iterations to reduce the error floor according to the stated iteration schedule.The total iteration count is T = T′ + T′′.
The total number of iterations is bounded using the schedules for the initial and refinement phases.
The final failure probability is controlled by selecting per-iteration failure probabilities and applying the sample-size condition.The proof replaces δ_0 with the global choice to obtain the final error floor.