Title: Structured Federated Learning through Clustered Additive Modeling
Venue: NeurIPS 2023
Accepted: yes
Keywords: Federated Learning
Abstract: Heterogeneous federated learning without assuming any structure is challenging due to the conflicts among non-identical data distributions of clients. In practice, clients often comprise near-homogeneous clusters so training a server-side model per cluster mitigates the conflicts. However, FL with client clustering often suffers from “clustering collapse'', i.e., one cluster's model excels on increasing clients, and reduces to single-model FL. Moreover, cluster-wise models hinder knowledge sharing between clusters and each model depends on fewer clients. Furthermore, the static clustering assumption on data may not hold for dynamically changing models, which are sensitive to cluster imbalance/initialization or outliers. To address these challenges, we propose ''Clustered Additive Modeling  (CAM)'', which applies a globally shared model $\Theta_g$ on top of the cluster-wise models $\Theta_{1:K}$, i.e., $y=h(x;\Theta_g)+f(x;\Theta_k)$ for clients of cluster-$k$. The global model captures the features shared by all clusters so $\Theta_{1:K}$ are enforced to focus on the difference among clusters. To train CAM, we develop a novel Fed-CAM algorithm that alternates between client clustering and training global/cluster models to predict the residual of each other.  We can easily modify any existing clustered FL methods by CAM and significantly improve their performance without ‘’clustering collapse'' in different non-IID settings. We also provide a convergence analysis of Fed-CAM algorithm.

Structured Federated Learning through Clustered Additive Modeling

Abstract

Heterogeneous federated learning without assuming any structure is challenging due to the conflicts among non-identical data distributions of clients. In practice, clients often comprise near-homogeneous clusters so training a server-side model per cluster mitigates the conflicts. However, FL with client clustering often suffers from 'clustering collapse', i.e., one cluster's model excels on increasing clients, and reduces to single-model FL. Moreover, cluster-wise models hinder knowledge sharing between clusters and each model depends on fewer clients. Furthermore, the static clustering assumption on data may not hold for dynamically changing models, which are sensitive to cluster imbalance/initialization or outliers. To address these challenges, we propose ' C lustered A dditive M odeling ( CAM )', which applies a globally shared model Θ g on top of the cluster-wise models Θ 1: K , i.e., y = h ( x ; Θ g )+ f ( x ; Θ k ) for clients of clusterk . The global model captures the features shared by all clusters so Θ 1: K are enforced to focus on the difference among clusters. To train CAM, we develop a novel Fed-CAM algorithm that alternates between client clustering and training global/cluster models to predict the residual of each other. We can easily modify any existing clustered FL methods by CAM and significantly improve their performance without 'clustering collapse' in different non-IID settings. We also provide a convergence analysis of Fed-CAM algorithm.

1 Introduction

Federated learning (FL) trains a global model over distributed clients and enforces data localization, i.e., data only stay local for model training at the client side while the server periodically averages client models' weights to update a global model and broadcast it to all clients. When the local data distributions are identical across clients, one global model suffices to serve all clients [1] However, non-identical distributions across clients (i.e., statistical heterogeneity) [2] are more common in practical FL scenarios, which leads to conflicts between the global objective and local ones. Instead of applying one model to all the m clients, an ideal case for non-IID settings would be training a local model per client without any interference from others. However, the local data are usually insufficient so a global model trained on heterogeneous clients exploiting all their data can still be helpful to local training. Hence, non-IID FL methods [3-5] usually struggle to find a sweet spot between the global consensus and local personalization. Without any assumptions on the structure among clients, all clients' distributions can be equally different from each other so a global model is impacted by the conflicts of all clients and may provide limited guidance to their local training.

Clustered Federated Learning. That being said, non-IID clients in practice usually have rich structures that have not been explored by most existing FL methods. A common structure is clusters, i.e., heterogeneous clients can be grouped into several near-homogeneous clusters each composed of

IFCA, client-wise non-lID

IFCA-CAM, client-wise non-lID

IFCA, cluster-wise non-lID

IFCA-CAM, cluster-wise non-lID

9(200)

17%

9(35)

17%

(66)6

17%

9(20)

Cluster size

2

8(0)

7(0)

6(0)

5(0)

Cluster size

2

8(22)

6(14)

5(23)

7(33)

Cluster size

15

150

6(14)

(es)e

7(24)

5(6)

Cluster size

15

8(42)

7(0)

6(0)

5(41)

4(0)

4(20)

4(2)

4(40)

(ole

3(25)

(ole

2(0)

2(28)

2(1)

2(37)

1(0)

1(0)

1(0)

1(0)

0(0)

0(0)

0(1)

0(20)

Round

Round

Round

Round

CIFAR-10, Cluster-wise non-lID (Dirichlet)

CIFAR-10, Cluster-wise non-lID (Dirichlet)

CIFAR-10, Client-wise non-lID (Dirichlet)

CIFAR-10, Client-wise non-lID (Dirichlet)

W/o CAM

W/o CAM

W/o CAM

W/o CAM

CAM

CAM

CAM

CAM

(%)

TestAccuracy

Test Macro-F1

TestAccuracy

Test Macro-F1

IFCA(5)

IFCA(10)

FeSEM(5)

FeSEM(10)

IFCA(5)

IFCA(10)

FeSEM(5)

FeSEM(10)

IFCA(5)

IFCA(10)

FeSEM(5)

FeSEM(10)

IFCA(5)

IFCA(10)

FeSEM(5)

FeSEM(10)

clients with similar distributions. In practice, clusters might be associated with geographic/age/income groups, affiliations, etc. Hence, we can train a server-side model for each cluster, hence mitigating the conflicts caused by heterogeneity. Unfortunately, clients' cluster memberships are usually undefined or inaccessible due to sensitive/private information and have to be jointly optimized with cluster-wise models, as recent clustered FL [6-9] approaches do. They maintain K models Θ 1: K for K clusters and assigns one Θ k to each clienti (with local data X i and local model θ i ), e.g., by min-loss ( Θ k with the minimum loss on X i ) or K-means (the nearest Θ k to θ i ) criterion. Hence, 1 ≤ K ≤ m models can accommodate more heterogeneity than single-model FL but also allows knowledge sharing among similar clients, which is lacking if training m client models independently. Hence, it may reach a better trade-off between global consensus and local personalization in non-IID settings.

However, compared to the general non-IID assumption, clustered FL's assumption might be too restrictive since it prohibits inter-cluster knowledge sharing and enforces every cluster-wise model's training to only depend on a few clients . This is contradictory to the widely studied strategy that different tasks or domains can benefit from sharing low-level or partial representations. It is due to the gap between the assumption of 'clustered data distributions' and the algorithms of 'clustering models (represented by loss vectors or model weights)': they are not equal and the latter is more restrictive. In other words, clients of different clusters can still benefit from feature/parameter sharing.

Moreover, clustered FL usually suffers from optimization instability because dynamically changing models can violate the static clustering assumption and lead to imbalanced cluster assignment, which affects Θ 1: K and local training in the future. In particular: (1) Clustering collapse , i.e., the clients assigned to one cluster keeps growing so 'the rich becomes richer (i.e., the cluster-wise model becomes even stronger)' until reducing to single-model FL. This happens because most clients tend to first learn shared features before focusing on client-specific ones; (2) Fragile to outliers such as malicious clients that may dominate some clusters and push all other benign ones to one or a few clusters; (3) Sensitive to initialization. The process highly depends on initial and earlier cluster assignments since they determine which clients' local training starts from the same model.

Main Contributions. To overcome the above problems of clustered FL, we propose a novel structured FL model termed ' C lustered A dditive M odeling ( CAM )'. Compared to clustered FL, CAM trains a global model Θ g on top of the K clusters' models Θ 1: K . Its prediction for clienti combines the outputs of Θ g and the associated cluster c ( i ) 's model, i.e., y = h ( x ; Θ g ) + f ( x ; Θ c ( i ) ) . This simple additive model removes the restriction of clustered FL by letting all clients share a base model Θ g . It enforces Θ 1: K to focus on learning the different features between clusters, hence mitigating 'clustering collapse'. Moreover, CAM tends to learn balanced clusters and determine the number of clusters automatically (by starting from more clusters and then zeroing out some of them), as demonstrated in Fig 1. Furthermore, CAM is less vulnerable to outliers, which can be mainly captured by Θ 1: K and have less impact on the global model Θ g . In addition, interactions between Θ 1: K and Θ g make CAM less sensitive to initial cluster assignments since updating Θ g also changes the clustering.

CAM is a general model-agnostic method that can modify any existing clustered FL methods. As examples, we apply CAM to two representative methods, i.e., IFCA [6] and FeSEM [8]. In the optimization of CAM, Θ 1: K and Θ g aim to fit the residual of each other's prediction. To this end, we propose an efficiently structured FL algorithm ' Fed-CAM ', which alternates between cluster assignment (server), local training (clients), and update of Θ 1: K and Θ g (server). In experiments on several benchmarks in different non-IID settings, CAM significantly improves SOTA clustered FL methods, as shown in Fig 2. Moreover, we provide a convergence analysis of Fed-CAM algorithm.

2 Related Work

Non-IID FL aims to tackle statistical heterogeneity across clients. FedAvg [1] is designed for the IID setting, so it suffers from client drift and slow convergence with non-IID clients [2]. To address this challenge, FedDANE [10] proposed a federated Newton-type optimization method by adapting a method for classical distributed optimization, i.e., DANE, to the FL setting. Instead of synchronizing all clients' models to be the same global model periodically, FedProx [3] only adds a proximal term to the local training objective that discourages the local model from drifting away from the global model and thus preserves the heterogeneity. [11] applies adaptive learning rates to clients and [12] conducts attention-based adaptive weighting to aggregate clients' models. [13] studies the convergence of the FedAvg in non-IID scenarios. Recent work also studies client-wise personalized FL [14-20], which aim to address the non-IID challenge by training a personalized model per client with the help of the shared global model. Their objectives focus on training local models rather than the server-side model.

Clustered FL assumes that non-IID clients can be partitioned into several groups and clients in each group share a cluster-wise model. It jointly optimizes the cluster assignments and the clusters' models. K-means-based methods [8] assign clusters to clients according to their model parameters' distance. CFL [21] divides clients into two partitions based on the cosine similarity between client gradients and then checks whether a partition is congruent according to the gradient norm. IFCA [6] and HypCluster [7] assign to each client the cluster whose model achieves the minimum loss on the client's data. Few-shot clustering has been introduced to clustered FL by [22,23]. FedP2P [24] allows communication between clients in the same cluster. [25] uses cluster-based contexts to enhance the fine-tuning of personalized FL models. [9] proposed the first cluster-wise non-IID setting and a bi-level optimization framework unifying most clustered FL methods.

Additive modeling in FL trains multiple models and adds their outputs together as its prediction. It was introduced to FL very recently. Federated residual learning [26] proposed an FL algorithm to train an additive model for regression tasks. [27] applies additive modeling to combining the outputs of a shared model and a local model in a partial model personalization framework. However, additive modeling has not been studied for clustered FL.

3 Clustered Additive Modeling (CAM)

In this section, we introduce clustered additive modeling (CAM), which combines a global model and cluster-wise model prediction in FL. CAM conducts a joint optimization of the global and cluster-wise models defined by a cluster assignment criterion. In particular, we provide two examples of CAM using different cluster assignment criteria, i.e., min-loss and K-means, which have been adopted respectively by two SOTA clustered-FL methods, i.e., IFCA and FeSEM. For each of them, we derive alternating optimization procedures (i.e., IFCA-CAM and FeSEM-CAM) that can be implemented in FL setting using two parallel threads of local model training. At the end of this section, we unify both algorithms in a structured FL algorithm Fed-CAM.

Notations. We assume that there are m clients and K clusters, where clienti has n i examples and all clients have n = ∑ m i =1 n i examples. On the server side, we have a global model Θ g and K cluster-wise models Θ 1: K . On the client side, we train m cluster models θ 0 1: m used to update the global model Θ g in FL and θ 1: m used to update the cluster-wise model Θ c ( i ) assigned to each clienti , where c ( i ) is its cluster label determined by the cluster assignment criterion c ( · ) . We further define C k ≜ { i ∈ [ m ] : c ( i ) = k } as the set of clients in clusterk . For simplicity, we will use X i and Y i to respectively represent the local training data on clienti and their ground truths, and ℓ ( Y i , F ( X i )) denotes the batch loss of model F ( · ) on ( X i , Y i ) . A CAM model for clienti can be

For classification, F i ( · ) produces logits and we can apply softmax to get the class probabilities.

3.1 IFCA-CAM: model performance-driven clustering

We extend the min-loss criterion used in IFCA [6] to CAM for cluster assignment, i.e., each clienti is assigned to the clusterk whose model Θ k leads to the minimal loss of CAM on clienti 's data, i.e.,

IFCA-CAM optimizes Θ g and Θ 1: K for minimizing the above minimal loss over all the m clients, i.e.,

where the inner minimization performs the min-loss assignment in Eq. (2). We solve Eq. (3) by the following alternating minimization of cluster membership, cluster-wise models, and the global model.

(i) Cluster assignment by applying Eq. (2) to the latest Θ g and Θ 1: K . This yields c ( · ) and C 1: K .

(ii) Fixing Θ g , we can optimize the cluster-wise models Θ 1: K by gradient descent:

In FL, the gradient can be approximated by aggregating the model updates of local models θ i from clients, whose training on the client side is: (1) initializing θ i ← Θ c ( i ) ; (2) starting from the initialization, running E local epochs updating θ i by

and (3) aggregating the local model update θ i -Θ k from client i ∈ C k to update Θ k , i.e.,

(iii) Fixing Θ 1: K , we can optimize the global model Θ g by gradient descent:

In FL, this gradient step can be approximated by aggregating the local models θ 0 i (similar to FedAvg): (1) initializing θ 0 i ← Θ g ; (2) running E local epochs training θ 0 i by

and (3) aggregating the updated local models θ 0 i of all the m clients to update Θ g , i.e.,

We can run two parallel threads of local training for θ 0 i and θ i for each clienti because their training in Eq. (8) and Eq. (5) does not depend on each other (but they both depend on the cluster assignments in (i)). This is analogous to the simultaneous update algorithm (FedSim) in [27]. One may also consider an alternative update algorithm (which may enjoy a slightly faster convergence) that iterates (i) → (ii) → (i) → (iii) in each round. However, it doubles the communication rounds ((i) requires one communication round) and does not allow parallel local training. Since the alternative update does not show a significant empirical improvement over FedSim in [27], we mainly focus on the parallel one in the remainder of this paper.

3.2 FeSEM-CAM: parameter similarity-based clustering

We follow a similar procedure of IFCA-CAM to derive FeSEM-CAM, which applies a K-means style clustering to the client models θ 1: m , whose objective is minimizing the sum of squares of client-cluster distance, i.e.,

Hence, similar to FeSEM [8], FeSEM-CAM assigns the nearest cluster-wise model to each client and updates the cluster-wise models as the cluster centroids (i.e., K-means algorithm), i.e.,

We iterate the above K-means steps for a few times until convergence in practice. FeSEMCAM applies the K-means objective in Eq. (10) as a regularization to the loss of CAM model ℓ ( Y i , h ( X i ; Θ g ) + f ( X i ; θ i )) , i.e.,

where the minimization w.r.t. Θ 1: K (with θ 1: m fixed) recovers the (weighted) K-means objective in Eq. (10). Unlike IFCA-CAM, where client model θ i is an auxiliary/latent variable for FL not showing in the objective of Eq. (3), it is explicitly optimized in Eq. (12). Similar to IFCA-CAM, we solve Eq. (3) by iterating the following alternating minimization steps (i)-(iii).

(i) K-means clustering that iterates Eq. (11) for a few steps until convergence, which yields c ( · ) , C 1: K , and Θ 1: K . The update of Θ 1: K is analogous to Eq. (6).

(ii) Fixing Θ g , we optimize θ 1: m by client-side local gradient descent:

The first two terms in Eq. (13) compute a linear interpolation between θ i and its assigned cluster's model Θ c ( i ) . This is a result of the K-means regularization term in Eq. (12) and keeps θ i close to Θ c ( i ) .

(iii) Fixing θ 1: m , we can optimize Θ g by gradient descent:

In FL, this gradient step can be approximated by aggregating the local models θ 0 i (similar to FedAvg): (1) initializing θ 0 i ← Θ g ; (2) running E local epochs training θ 0 i by

and (3) aggregating the updated local models θ 0 i of all the m clients to update Θ g by Eq. (9).

3.3 Algorithm

In Algorithm 1, we propose a structured FL algorithm for CAM, i.e., Fed-CAM, which can unify the derived optimization procedures for IFCA-CAM and FeSEM-CAM and can be easily extended to other clustered FL and clustering criteria.

Warmup. As an alternating optimization framework, it would be unstable if both Θ g and Θ 1: K are randomly initialized and jointly optimized in parallel since they may capture overlapping information and result in an inefficient competitive game. To encourage them to learn complementary knowledge, warmup training for one of them before the joint optimization is helpful. For example, a few rounds of FedAvg can produce a 'warm' Θ g , whose predictions' residuals are more informative to train Θ 1: K . Another warmup strategy could be to run a few local training epochs and extract warm Θ 1: K by clustering the lightly-trained local models θ 1: m . In Fed-CAM, we can apply the former warmup to IFCA-CAM and the latter to FeSEM-CAM.

Algorithm 1: Fed-CAM

initialize : Randomly initialize Θ 1: K and Θ g . warmup : (1) β rounds of FedAvg to get an initial Θ g (IFCA-CAM) or (2) β epochs of local training only to get a initial θ 1: m (FedSEM-CAM). Broadcast Θ g and Θ 1: K to all clients. 1 while not converge do /* Client (in parallel) */ 2 for every selected client i do 3 Model performance-driven clustering (e.g., IFCA-CAM): cluster assignment by Eq. (2); 4 Initialize θ i ← Θ c ( i ) and θ 0 i ← Θ g ; 5 Local training of θ i for E epochs: e.g., Eq. (5) (IFCA-CAM) or Eq. (13) (FeSEM-CAM); 6 Local training of θ 0 i for E epochs: e.g., Eq. (8) (IFCA-CAM) or Eq. (15) (FeSEM-CAM); 7 Upload θ i and θ 0 i to the server; /* Server */ 8 Update cluster-wise models Θ 1: K : e.g., Eq. (6) (IFCA-CAM) or Eq. (11) (FeSEM-CAM); 9 Update global model Θ g by Eq. (9); 10 Broadcast Θ g and Θ 1: K to all clients; output : Global model Θ g , cluster-wise models Θ 1: K and c ( i ) ∀ i ∈ [ m ] .

4 Convergence Analysis

Based on the convergence analysis presented in [27], which aims to minimize the following objective:

where u represents the shared parameters and V = v 1 , v 2 , · · · , v m denotes the personalized parameters. If we map Θ g to u , and Θ 1: K to V respectively, this appears strikingly similar to our methods as illustrated in Equations 3 and 12. Provided that the clustering remains stable, we can employ the theoretical framework of [27]. And firstly, we make some standard assumptions for the convergence analysis as below.

Assumption 1. (Smoothness). For i = 1 , · · · , m , the loss function l is continuously differentiable, and there exist constants L that ∇ Θ g ℓ (Θ g , Θ k ) is L-Lipschitz with respect to Θ g and Θ k , and ∇ Θ k ℓ (Θ g , Θ k ) is L-Lipschitz with respect to Θ g and Θ k .

Assumption 2. (Unbiased gradients and bounded variance). The stochastic gradients are unbiased and have bounded variance. For all Θ g and Θ k ,

and

Assumption 3. (Partial gradient diversity). There exists a constant for all θ 0 i and Θ g , θ i and Θ k ,

Assumption 4. (Convexity of cluster models). Fix Θ g , assume ℓ (Θ k ) is convex.

Theorem 1. (Convergence of Fed-CAM). Let Assumptions 1, 2, 3 and 4 hold, and learning rates chosen as η = τ/ ( LE ) for a τ depending on the parameters L, σ 2 g , σ 2 k , δ 2 , s, m, T , provided clustering stable, we have (ignoring absolute constants),

where △ ℓ = ℓ 0 -ℓ ⋆ , and we define effective variance terms,

Remark 1 . It is straightforward to prove that the clustering of both IFCA-CAM and FeSEM-CAM converges, as evidenced by Ma et al. (2022). However, proving the stability of these clustering methods is more challenging due to the oscillation phenomenon often seen in K-means. The stability of clustering will be further demonstrated through experimental analysis in Section 5.2.

Remark 2 . Besides the clustering structure, there is a distinct difference between FedSim [27] and Fed-CAM. In Fed-CAM, we need to aggregate both Θ g and Θ 1: K , while in FedAlt [27], only Θ g requires aggregation. The σ 2 sim, 1 and σ 2 sim, 2 reflect the impact of sample number s and local steps E . Larger s or smaller E , better convergence rate. According to the results presented in [27], alternative gradient descent surpasses simultaneous gradient descent in terms of convergence rate. The asymptotic 1 / √ T rate is achieved when each device is seen at least once on average, and the 1 /T term is dominated by the 1 / √ T term, a situation that occurs when (ignoring absolute constants)

5 Experiments

Benchmark datasets and partitions. The proposed methods have been validated using several benchmark datasets. While detailed results for PathMNIST and TissueMNIST from the MedMNIST [28] are provided in the supplementary material, the other datasets used for validation include:

Fashion-MNIST [29] includes 70,000 labeled fashion images (28 × 28 grayscale) in 10 classes, such as T-shirts, Trouser, and Bag, with others.

CIFAR-10 [30] consists of 60,000 images (32 × 32 color) in 10 classes, including airplane, automobile, bird, and truck, among others. The divergence among classes in CIFAR-10 is relatively higher than in other datasets from the MNIST family.

Each dataset is split among 200 clients and we create the following non-IID scenarios:

Client-wise non-IID by Dirichlet distribution ( α = 0 . 1 ) : This approach uses the Dirichlet distribution to control the randomness of non-IID data, as proposed by [31]. This is a standard method used in most personalized FL methods, which are usually client-wise non-IID.

Cluster-wise non-IID by Dirichlet distribution ( α = (0 . 1 , 10) ) : This strategy divides the dataset into K clusters with α = 0 . 1 to generate substantial variance in cluster-wise non-IID. Then, each cluster is divided into m/K clients with α = 10 to control the non-IID across clients.

Client-wise non-IID by n-class (2) : This method randomly selects n -class out of all classes in the dataset for each client, as proposed by [1], and then samples instances from these datasets.

Cluster-wise non-IID by n-class (3, 2) : This approach randomly assigns 3 classes to each cluster, ensuring a relatively balanced number of instances per class. It then assigns 2 classes to each client.

Baselines. We select baseline methods from four categories as follows:

Single model-based FL: We choose FedAvg [1] and FedProx [3] with a coefficient of 230 and a regularization of 0.95 as the baselines.

Ensemble FL: We train FedAvg and FedProx K times and then learn an ensemble model via soft voting to serve all clients, which are named FedAvg+ and FedProx+, respectively.

Clustered FL: We choose FeSEM [8] and IFCA [6], which is similar to HypCluster [7].

Clustered FL with additive modeling: We integrate CAM with IFCA and FeSEM, denoting them as IFCA-CAM and FeSEM-CAM, respectively.

Learning-related hyperparameters. We use the Convolutional Neural Network (CNN) [32] as the basic model architecture for each client, as detailed in the supplementary material. For optimization, we employ SGD with a learning rate of 0.001 and momentum of 0.9 to train the model, and the batch size is 32. We evaluate the performance using both micro accuracy (%) and macro F1-score (%) on the client-wise test datasets to better capture the non-IID nature per client.

FL system settings. We conduct 100 global communication rounds in the FL system, including 30 warm-up rounds if applicable. Each communication involves 10 local steps. For the clustering process of FeSEM-CAM, we measure distance on the flattened parameters of the fully-connected layers, and use K-Means as the clustering algorithm. The coefficient λ is chosen from 0 . 001 , 0 . 01 , 0 . 1 based on the best performance.

5.1 Main Results and Comparisons

Cluster-wise non-IID scenarios make the assumption that there are underlying clustering structures among clients. Table 1 compares the methods using two benchmark datasets, namely Fashion-MNIST and CIFAR-10. Results using two biomedical datasets are presented in the appendix. The following are some notable observations and analyses:

The application of the ensemble mechanism to FedAvg and FedProx yields minor improvements. This is because the server-side model in FedAvg or FedProx is already a relatively strong model, while ensemble mechanisms usually excel with weak models.

The introduction of CAM significantly enhances the performance of IFCA, which typically struggles with clustering collapse in cluster-wise non-IID scenarios. Notably, CAM decomposes the shared components into a global model and personalized parts into cluster models. Thus, the clustering collapse is mitigated by isolating the dominant shared knowledge.

FeSEM generally exhibits robust performance on cluster-wise non-IID without outliers. Implementing CAM in FeSEM further improves the Macro-F1 performance. The clustering process in FeSEM tends to overfit the label distribution (imbalanced classes) of clients to achieve higher accuracy. However, the application of CAM introduces a global model with a balanced label distribution by averaging all clients, thereby boosting the Macro-F1 performance while preserving the cluster-wise non-IID for high accuracy.

With an increase in the number of clusters K , the CAM-based methods show substantial improvements in Macro-F1. The decomposition of shared knowledge and cluster-wise non-IID characteristics benefit from a reasonably larger K , which facilitates fine-grained, cluster-wise personalization.

Client-wise non-IID Table 2 presents comparative results under client-wise non-IID scenarios using two benchmark datasets: Fashion-MNIST and CIFAR-10. Interestingly, IFCA maintains stable performance under client-wise non-IID conditions, primarily because it cannot form a single dominant cluster model - a primary cause of clustering collapse - in a highly heterogeneous environment. The application of CAM to IFCA and FeSEM shows a significant enhancement, particularly on the CIFAR-10 dataset. This improvement is likely due to FeSEM's typical restriction on knowledge sharing across clusters. In contrast, CAM utilizes a global model to capture more useful common

FeSEM, client-wise non-lID

FeSEM-CAM, client-wise non-lID

FeSEM, cluster-wise non-lID

FeSEM-CAM, cluster-wise non-lID

9(48)

17%

9(28)

175

9(35)

17%

9(20)

Cluster size

8(37)

7(30)

6(20)

5(15)

Cluster size

150

25

8(26)

7(23)

6(21)

5(20)

Cluster size

8(27)

7(25)

6(22)

5(20)

Cluster size

t25

8(20)

7(20)

6(20)

5(20)

100

4(12)

4(20)

4(18)

4(20)

3(12)

3(19)

3(18)

3(20)

1(9)

2(10)

2(18)

1(13)

2(15)

1(12)

2(20)

1(20)

0(7)

0(12)

0(8)

0(20)

Round

Round

40

Round

Round

40

knowledge across clusters, thereby substantially enhancing the generalization capability of each cluster. Furthermore, CIFAR-10, being a relatively complex dataset with a diversity of images, underscores the importance of sharing common knowledge.

5.2 Visualization: CAM combats clustering collapse

Figures 1 and 3 demonstrate the effectiveness of applying CAM to mitigate clustering collapse in IFCA and FeSEM under both cluster-wise and client-wise non-IID scenarios, using the CIFAR-10 dataset with K = 10 . Each color represents a cluster, and the X-axis represents the iteration rounds.

In the case of IFCA, we observe a severe clustering collapse issue in cluster-wise non-IID scenarios. A single cluster can encompass all clients in the client-wise non-IID setting and up to 50% of clients in the cluster non-IID setting. Furthermore, the clustering remains unstable throughout the process. However, when CAM is applied in IFCA-CAM, it quickly identifies some clustering structures within a few rounds, and this structure closely approximates the ground truth.

As for FeSEM, while the phenomenon of clustering collapse is not as pronounced, a single cluster can still dominate up to 25% of all clients if there are no outliers. CAM can expedite the clustering convergence, sometimes achieving it in just one round. Moreover, under client-wise non-IID settings, the application of CAM results in lower variance and more uniform cluster size. In the case of cluster-wise non-IID settings, FeSEM-CAM can easily identify the ground truth.

6 Conclusions

We propose a novel structured FL model 'clustered additive modeling (CAM)' and an efficient FL algorithmic framework Fed-CAM to address non-IID FL challenges with clustering structure. CAMis a general mode-agnostic tool that can improve various existing non-IID FL methods. It can capture more general non-IID structures with global knowledge sharing among clients than clustered FL and overcome several weaknesses such as clustering collapse, vulnerability to cluster imbalance/initialization, etc. Theoretically, Fed-CAM is capable of achieving an asymptotic convergence rate of O (1 / √ T ) . Extensive experiments show that CAM brings substantial improvement to existing clustered FL methods, improves cluster balance, and effectively mitigates clustering collapse.

References

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics , pages 1273-1282. PMLR, 2017.

Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and Trends® in Machine Learning , 14(1-2):1-210, 2021.

Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems , 2:429-450, 2020.

Hangyu Zhu, Jinjin Xu, Shiqing Liu, and Yaochu Jin. Federated learning on non-iid data: A survey. arXiv preprint arXiv:2106.06843 , 2021.

Dashan Gao, Xin Yao, and Qiang Yang. A survey on heterogeneous federated learning. arXiv preprint arXiv:2210.04505 , 2022.

Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. Advances in Neural Information Processing Systems , 33:19586-19597, 2020.

Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arXiv preprint arXiv:2002.10619 , 2020.

Ming Xie, Guodong Long, Tao Shen, Tianyi Zhou, Xianzhi Wang, Jing Jiang, and Chengqi Zhang. Multi-center federated learning. arXiv preprint arXiv:2108.08647 , 2021.

Jie Ma, Guodong Long, Tianyi Zhou, Jing Jiang, and Chengqi Zhang. On the convergence of clustered federated learning. arXiv preprint arXiv:2202.06187 , 2022.

Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smithy. Feddane: A federated newton-type method. In 2019 53rd Asilomar Conference on Signals, Systems, and Computers , pages 1227-1231. IEEE, 2019.

Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Koneˇ cn` y, Sanjiv Kumar, and H Brendan . Adaptive federated optimization. arXiv preprint arXiv:2003.00295 , 2020.

Jing Jiang, Shaoxiong Ji, and Guodong Long. Decentralized knowledge acquisition for mobile internet applications. World Wide Web , 23(5):2653-2669, 2020.

Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of fedavg on non-iid data. arXiv preprint arXiv:1907.02189 , 2019.

Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Toward personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems , 2022.

Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Systems , 33:3557-3568, 2020.

Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. Personalized federated learning using hypernetworks. In International Conference on Machine Learning , pages 9489-9502. PMLR, 2021.

Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint arXiv:2003.13461 , 2020.

Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In International Conference on Machine Learning , pages 2089-2099. PMLR, 2021.

Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes. Advances in Neural Information Processing Systems , 33:21394-21405, 2020.

Karan Singhal, Hakim Sidahmed, Zachary Garrett, Shanshan Wu, John Rush, and Sushant Prakash. Federated reconstruction: Partially local federated learning. Advances in Neural Information Processing Systems , 34:11220-11232, 2021.

Felix Sattler, Klaus-Robert Müller, and Wojciech Samek. Clustered federated learning: Model-agnostic distributed multitask optimization under privacy constraints. IEEE transactions on neural networks and learning systems , 2020.

Don Kurian Dennis, Tian Li, and Virginia Smith. Heterogeneity for the win: One-Shot federated clustering. March 2021.

Pranjal Awasthi and Or Sheffet. Improved spectral-norm bounds for clustering. In Approximation, Randomization, and Combinatorial Optimization. Algorithms and Techniques , pages 37-49. Springer, 2012.

Li Chou, Zichang Liu, Zhuang Wang, and Anshumali Shrivastava. Efficient and less centralized federated learning. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases , pages 772-787. Springer, 2021.

Xueyang Tang, Song Guo, and Jingcai Guo. Personalized federated learning with contextualized generalization. In Lud De Raedt, editor, Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence, IJCAI-22 , pages 2241-2247. International Joint Conferences on Artificial Intelligence Organization, 7 2022. Main Track.

Alekh Agarwal, John Langford, and Chen-Yu Wei. Federated residual learning. arXiv preprint arXiv:2003.12880 , 2020.

Krishna Pillutla, Kshitiz Malik, Abdel-Rahman Mohamed, Mike Rabbat, Maziar Sanjabi, and Lin Xiao. Federated learning with partial model personalization. In International Conference on Machine Learning , pages 17716-17758. PMLR, 2022.

Jiancheng Yang, Rui Shi, Donglai Wei, Zequan Liu, Lin Zhao, Bilian Ke, Hanspeter Pfister, and Bingbing Ni. Medmnist v2: A large-scale lightweight benchmark for 2d and 3d biomedical image classification. arXiv preprint arXiv:2110.14795 , 2021.

Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747 , 2017.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.06335 , 2019.

Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature , 521(7553):436-444, 2015.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors study heterogeneous federated learning, an important problem in federated learning, where the goal is to leverage the collective intelligence of multiple clients with diverse data distributions, features, or models to train a global model that can generalize well across all clients' data. In particular, the authors base their method on clustered federated learning and propose Clustered additive modeling to deal with the clustering collapse and dynamically changing models. The proposed method involves training a shared global model on top of cluster-specific models, and prediction results are obtained by combining the outputs from both the global model and the associated cluster model. Empirical results on two well-known datasets show the effectiveness of the proposed method.\n\nSTRENGTHS:\nStrength\n1. The paper writing is clear and easy to follow\n2. The empirical performance seems promising\n\nWEAKNESSES:\nWeakness:\n1. Comparison with other methods under non-iid data setting needs to be improved.\n2. Ablation study on structural clustering needs to be improved. \n3. Fed-CAM introduces several hyper-parameters, including the number of rounds for warm-up, and the number of cluster size K.\n\nQUESTIONS:\nQuestions for authors:\n1. It is unclear which part of the structural clustering leads to the final good performance, as there are several hyper-parameters and assumptions in cluster generation. For example,\n\t1) How does the number of rounds for warm-up affect the quality?\n\t2) Can we dynamically adjust the cluster size K based on some coarse-to-fine heuristic?\n\t3) Can the proposed method be generalized to handle graph data?\n\n2. More existing methods could be compared under more non-iid settings as well as iid settings. For example,\n\t1) more Distribution-based label imbalance setting, feature skew case, and label & feature skew case\n\t2) comparison with knowledge transfer-based methods\n\n3. The authors do not include sufficient discussions about the communication and computational cost, which are suggested to be included.\n\nLIMITATIONS:\nClustering-based FL approaches have a common limitation that is how to determine the cluster size. Some analysis on this would be desirable in the revision.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studied the problem of clustered federated learning. It analyzed the limitations of existing clustered FL algorithms, including clustering collapse and missing shared knowledge among clusters. Then this paper introduced a novel Clustered Additive Modeling (CAM) framework for learning both a globally shared model and cluster-wise models. A Fed-CAM algorithm was proposed to train CAM model parameters iteratively. The experimental results showed that CAM outperformed existing clustered FL algorithms in prediction accuracy.\n\n#######################################\nI appreciate the detailed response from the authors. However, my major concerns regarding the overall contributions still remain. Therefore, I keep my scores unchanged.\n\nSTRENGTHS:\nThe strengths of this paper are summarized below.\n(1) It proposed a novel Clustered Additive Modeling (CAM) framework for clustered federated learning. Compared to previous algorithms, CAM leveraged a global model to discover the shared knowledge among clusters.\n\n(2) It demonstrated the flexibility of CAM framework by incorporating CAM with two popular clustered FL methods, thus leading to IFCA-CAM and FeSEM-CAM models. Then it provided a unified Fed-CAM optimization method for iteratively updating the clustering and global model parameters.\n\n(3) Experimental results on both cluster and client-wise non-IID settings showed that compared to vanilla IFCA and FeSEM, IFCA-CAM and FeSEM-CAM could achieve better prediction performance in local clients.\n\nWEAKNESSES:\n(1) The rationale behind the proposed CAM framework lacks sufficient justification. As illustrated in the introduction section, existing cluttered FL algorithms suffers from clustering collapse, fragility to outliers, and sensitivity to initialization. However, the techniques regarding how CAM addresses these limitations are not explained.\n(1-1) CAM introduced a global model to mitigate the lustering collapse. As shown in Subsection 3.3, it requires that global and cluster-wise models learn complementary knowledge. One concern is whether the global knowledge shared by all clients always exists in any non-IID FL settings. Furthermore, line 49 states that “different tasks or domains can benefit from sharing low-level or partial representations”. A natural question is whether it is more reasonable to learn partial representations (e.g., common low-level model parameters in FedRep [24]) for discovering the globally shared knowledge among clusters. In addition, the sensitivity of complementarity relationship between global and cluster-wise models can be analyzed with respect to warm-up strategies.\n(1-2) This paper stated that the outlier clients have less impact on the global model, thus enabling the less vulnerability of CAM. But the outliers might have less common knowledge with other clients/clusters. Besides, based on the update rules in Eqs. (9)(14), the outliers can also dominate the training of global model, if they have a significantly large number of training samples.\n(1-3) It is more convincing to quantitively evaluate the sensitivity of CAM to initialization, e.g., showing the clustering results with different random initializations in the experiments.\n\n(2) The optimization of IFCA-CAM and FeSEM-CAM is confusing. Algorithm 1 updates the model parameters \\theta_i and \\theta_i^0 (lines 5-6) separately. This might lead to sub-optimal solutions of local objective functions, compared to optimizing \\theta_i and \\theta_i^0 at each training epoch simultaneously. In addition, it is much more efficient to update \\theta_i and \\theta_i^0 together within E training epochs.\n\n(3) Though this paper focuses on clustered FL scenarios, it is better to compare the proposed CAM framework with other related FL frameworks. This is because learning the shared knowledge with a global model is a common strategy for personalized FL. For example, when each client is considered as a single cluster, the objective function of CAM is similar to [ref 1] and [ref 2]. In addition, the state-of-the-art pFL baselines based on feature sharing can be included to validate the performance of CAM.\n\n(4) The convergence analysis of CAM is provided based on the assumption that clustering remains stable during training. But the results in subsection 5.2 might not validate the clustering stability. This is because Figures 1&3 provides the clustering results using only cluster size. It is not guarantee that clients in each cluster remains the same during training. More quantitively evaluations could be provided to verify the quality and stability of clustering in CAM. \n\n[ref 1] Deng, Yuyang, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. \"Adaptive personalized federated learning.\" arXiv preprint arXiv:2003.13461 (2020).\n[ref 2] Mansour, Yishay, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. \"Three approaches for personalization with applications to federated learning.\" arXiv preprint arXiv:2002.10619 (2020).\n\nQUESTIONS:\n(1) It shows that CAM tends to learn balanced clusters and determine the number of clusters automatically. Thus, CAM might perform well when the ground-truth of client clustering is balanced. But when the ground-truth of client clustering is highly skewed, would the balanced clusters learned by CAM result in sub-optimal solutions?\n\n(2) Why is there a constant 1/m in Eq. (9) for updating global model parameters?\n\n(3) In Eq. (13), what is the notation “\\eta”? If it denotes the learning rate, why is it only applied to the first term of Eq. (12) with classification loss? \n\n(4) In line 176, should it be “Eq. (13)” instead of “Eq. (3)” regarding the K-means regularization term?\n\n(5) Why do IFCA-CAM and FeSEM-CAM use different warmup strategies? Is the selection of warmup strategy related to the clustering methods (e.g., minimum loss values in IFCA, K-means over parameters in FeSEM) in FL?\n\n(6) What is FedAlt in line 217?\n\n(7) In FL system setting, would the hyper-parameter \\lambda be selected based on the best performance over the validation set? If so, how are the train/validation/test sets splits in the experiments?\n\nLIMITATIONS:\nNot discussed.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new federated learning with a specified structure for the prediction produced for each client, called \"clustered additive modelling\", which adds the prediction of a global model to the prediction of the model trained on a local cluster of clients. It is a modification to existing clustered federated learning methods and fixes several critical issues of these methods such as clustering collapse, fragility to outlier clients, and sensitivity to initialization. Some of the issues are fundamental and widely exist in many existing methods. The paper shows that the proposed CAM can effectively address these issues in different clustered FL methods under different non-IID settings, consistently leading to a promising improvement.\n\nSTRENGTHS:\n1. This tackles a fundamental problem in clustered FL, i.e., clustering collapse. This problem might be widely observed but has not been well investigated and effectively solved before. Hence, I believe this work can be helpful and have a broad impact on FL problems studying statistical heterogeneity and structures among clients. \n2. The motivation of this paper is clear. The solution provided by clustered additive modelling is principal, easy to understand, effective across different settings, and generalizable to various existing FL methods. \n3. The proposed algorithms in Section 3.1 and 3.2 are intuitive and convincing with detailed explanations per step and theoretical guarantee (convergence analysis). I appreciate the authors for providing two examples of CAM algorithms following the same high-level idea. \n4. The experiments conducted on 2 datasets x 4 non-IID settings show great advantages of the proposed method over existing FL and clustered FL methods. The analysis of clustering collapse also provides a nice explanation and empirical evidence for the improvement. \n5. The code is provided for high reproducibility.\n\nWEAKNESSES:\n1. Compared to existing clustered FL methods, the proposed method needs to train an additional model. How much extra computation does it require? It would be helpful to provide the computation complexity of CAM.\n2. While other parts of the paper are easy to understand, the convergence analysis part is limited in length and left many details in the appendix. Can you elaborate a little bit more on the proof idea?\n3. The comparison and experiments on each dataset are thorough to me but the empirical results would be stronger if experimenting on more datasets.\n\nQUESTIONS:\n1. Can you provide more details and some examples of the cluster-wise non-IID setting?\n2. What is the cost of the warmup stage? How is the performance affected by reducing the warmup rounds?\nOverall, I think this is a solid paper that successfully addressed a widely existing fundamental problem in FL that has not been investigated before. The motivation and proposed idea are principal. Most parts of the paper are clear. The analysis and experiments are convincing. Experiments on more datasets and more guidance on the theoretical analysis can further strengthen the draft.\n\nLIMITATIONS:\nYes. The authors have addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel clustered federated learning method using additive modeling to tackle the clustering collapse problem. The paper contributes to advance the research domain of clustered federated learning which is an important and practical solution to solve non-iid problem in federated settings.\n\nSTRENGTHS:\n1. The proposed method is simple yet effective. The overall solution is technique sounds. This work is very likely will become a new baseline of clustered FL.\n\n2. The targeting problem is an inherent challenge of clustered FL. The Introduction section provides insight analysis of current clustered FL. \n\n3. The overall flow description is sufficiently clear.\n\n4. The provided theoritical analhysis is well suited to the clustered FL settings.\n\n5. The claims are well supported by the experiments.\n\nWEAKNESSES:\n1. The readability of this paper could be improved. For example, there are many abbreviation names, e.g. CAM, FED-CAM, IFCA, FeSEM. The symbol system is a little bit complicated. \n \n2. The paper discussed several insightful challenge of clustered FL. However, it is how the proposed method can solve these challenges. \n\n3. Some description needs to be improved. For example, it is unclear how the Figure 1 and 2 link to the Introduction section. Moreover, in Introduction section, there are many symbols without a clear definition.\n\n4. This method may require storing multiple models locally (cluster-level model, global model, and local model), which could increase costs. It would be better to compare and discuss with other methods to explore potentially more favorable alternatives.\n\nQUESTIONS:\n1. In Eq.1, why you use two symbols h(.) and f(.) to represent global model and cluster-specific model. Are they using different model architecture?\n\n2. In Eq. 2, why is the additive results is Y rather than 2Y?\n\n3. Please discuss how the proposed method can solve the mentioned challenges in Introduction section.\n\n4. Is it necessary to have a comparison experiments with some personalized FL method, e.g. FedAvg + Finetuning?\n\n5. In Ensemble FL, whatis the different part of K times experiments? \n\n6. How about to add a parameter to adjust the importance weight between two models h and f? Say h(.) + a*f(.).\n\n7. In Notations (line 128), what does the symbol \"0\" represent in the context of the symbol $\\theta_{1:m}^0$ ?\n\n8. Is the performance of the global model or the cluster-level model compared during testing?\n\n9. How many clusters obtained in training phase? Is it necessary to keep the same number of clusters as the CAM method for methods that need to specify the number of clusters?\n\nLIMITATIONS:\nYes. The authors have addressed the limitations. There is no negative societal impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-12"}