Federated Geometric Monte Carlo Clustering to Counter Non-IID Datasets

Federated learning (FL) frameworks [1, 2] are commonly used when regulations (such as the GDPR²²2https://gdpr.eu/) or mere data volume prevent data exchanges. Datasets are distributed across the members (clients) of the federation and the global machine learning optimization problem can approximately be split up into smaller sub-problems that are distributed across independently acting clients. Client solutions are aggregated either by a central server or in a distributed manner. However, this approximation only holds in an idealised settings where client datasets are sampled from the same distribution and data samples form independent events. Deviating from independent and identically distributed datasets (IID) poses a challenge to most FL approaches and results in global models that fail to accurately represent the entire aggregated dataset [3, 4, 5, 6]. Sources of non-IID-ness can be found in both the label space and the feature space. In the former, for a given feature, different labels might be attributed due to regional differences (e.g., different sentiments in fashion). This is also termed concept shift and leads ultimately to clients with incongruent labelling participating to the training of a common FL model. In addition, classes may be imbalanced due to one class of labels being over-represented with respect to others. This paper focuses on the latter, i.e., feature space IID violations. They appear for example when a particular feature is over-represented in a member’s dataset compared to other datasets. A biobank might sample genes that bias towards a particular local population [7] with the consequence of possibly overlooking regional differences, such as the well known preponderance of a specific gene mutation coding for sickle cell disease in the Sub-Saharan Africa population [8]. For a given label, features might also vary across datasets, such as different handwriting styles for the same alphanumeric character (e.g., 7 with and without bar). Non-IID-ness may lead to feature and label skew [3] and accuracy degradation [4, 6, 9] by leaving the choice between starting from different initialization weights, causing models to diverge, or accepting diversity suppression when starting from the same weights [5].

Besides statistical heterogeneity, FL algorithms need to cope with system heterogeneity, leading to asynchronous model updates or incomplete local training. The first proposed FL attempt, FedAvg [1], aggregates all client models by averaging their weights, leading to the above inaccuracies in the presence of non-IID-ness. Variations [6, 10] of FedAvg have been proposed, with satisfactory results when the non-IID-ness is less pronounced and exclusively in the label space (see Sec. II). Federated Clustering [11, 12, 13, 14] tackles this problem by assigning clients to separate clusters and hence limiting the deleterious transfers of negative knowledge between member models that have been trained on distinct datasets. These approaches result in the construction of trained models with reasonable accuracy despite IID violations, however, at the cost of significant communication and/or computational overheads.

In this paper, focusing on non-IID-ness in the feature space, we propose a novel Monte Carlo Clustering approach, called FedGMCC, to counter non-IID-ness without compromising final model accuracy and while maintaining low communication costs and data privacy. We show that in order to train a global FL model in a non-IID setting, the objective function to be minimized has to be conceptually split up into two components: the first encompasses the IID component for which the usual FedAvg solution holds; the second captures the non-IID contribution, which we show can be solved by introducing interaction between client models. We derive this interaction by leveraging an observation about the geometry of training loss manifolds  [15, 16, 17, 18], namely that seemingly different, stationary solutions to the training loss minimization problem, obtained at the individual member sites, are often connected via simple parametric curves where the training loss is approximately flat. The existence of these curves reveals the pair-wise interaction between models needed to solve the non-IID problem, which leads us to establish a criterion for clustering seemingly different solutions and suppressing negative knowledge transfer between non-connected ones. More importantly, we show that drawing from [19], averaging models along the curve parametrization, can produce global models with enhanced accuracy and generalization. As a summary, we make the following contributions.

1. 1.

   We separate the FL objective in an IID and a non-IID components and, based on the geometry of curved training loss manifolds, put forward an Ansatz solution that simultaneously minimizes both components.

2. 2.

   We demonstrate how to construct this solution using a Monte Carlo sampling of the received client model output spaces, and present novel model weight clustering and aggregation rules.

3. 3.

   We evaluate FedGMCC using the EMNIST62 dataset and a genomic dataset with different non-IID-ness. FedGMCC always outperforms FedAvg [1] and FedProx [6]. It outperforms CFL [13] in case of high non-IID-ness.

Among the myriad of published FL algorithms, most of them rely on clients to upload model weight updates onto a central server that proceeds to aggregate them and to redistribute the result back to clients. The main difference often relies in the way the aggregation rule is applied. FedAvg [1] was first to allow clients to train a global model locally, instead of transferring data. FedAvg returns to a central server only for aggregating all local models by averaging their weights. FedProx [6] inhibits local updates, by adding a proximal term to the loss function, which restrains divergence between local and global model weights. FedBn [10] achieves accurate results on non-IID datasets by batch normalizing the local neural networks’ input layer before aggregation. FedFV [20] detects inconsistent gradient updates and corrects them before the aggregation step. Our approach (FedGMCC) aggregates models by averaging weights along their geometric connection.

Federated clustering groups certain client models into separate clusters to prevent negative knowledge transfer such as the iterative federated clustering algorithm (IFCA) [11] which solves an incongruency that occurs when non-IID datasets exhibit concept shifts by training multiple models on local data and returning the one with the lowest aggregation loss. However, these works primarily focus on label-space non-IID-ness, leading to inaccurate results in the presence of feature-space non-IID-ness. Moreover, training multiple models induce high communication and computation costs.

Merging local models with approximately the same weights (according to cosine similarity, L1 or L2 norm) reduces the number of models in the ensemble [12, 13]. For example, clustered FL (CFL) [13] merges models if their weights (or gradients) compare well enough based on a specified metric. Again CFL focuses on label non-IID-ness. In contrast, FedGMCC explicitly considers feature-space non-IID-ness. Among the aforementioned FL approaches, CFL is the most reminiscent to our approach, with two differences. First in CFL, the central server applies clustering after multiple FedAvg iterations whereas FedGMCC clusters after each client-server communicating round. Second, CFL’s clustering rule relies on measuring whether the gradients of individual client-weight models are coherent (parallel) via the cosine similarity measure. Our approach verifies that a parallel transport of one client model weights to the other is feasible. Hence, FedGMCC extends CFL’s gradient coherency measure to include intermediary models along the affine connection.

We consider $K$ clients that aim to locally minimize a local objective $\mathcal{L}_{k}{=}\mathcal{L}(\operatorname{\mathbf{X}}_{k},\operatorname{\mathbf{w}})$. Clients send their respective solutions (client model weights) $\operatorname{\mathbf{w}}_{k}^{*}{=}\arg\min_{\operatorname{\mathbf{w}}}\mathcal{L}_{k}$ to a central server to approximate with aggregate $\operatorname{\mathbf{w}}_{t}^{f}$ a global solution for the loss minimization problem $\operatorname{\mathbf{w}}^{*}{=}\arg\min_{\operatorname{\mathbf{w}}}\mathcal{L}(\operatorname{\mathbf{x}},\operatorname{\mathbf{w}})$, which maps the $C$ classes of the compact input space $\mathcal{X}$ into the label space $\mathcal{Y}=[\mathbf{C}]$, where $[\mathbf{C}]={1,...,C}$. We consider the neural network as a function $f(\operatorname{\mathbf{w}}):\mathcal{X}\mapsto\mathcal{S}$, parametrized by weights $\operatorname{\mathbf{w}}$, that maps $\operatorname{\mathbf{X}}$ to the probability simplex $\mathcal{S}{=}\left\{\mathbf{z}|\sum_{i=1}^{C}z_{i}=1,z_{i}\geq 0,\forall i\in[\mathbf{C}]\right\}$. We write $f_{i}$ for the probability of the $i$-th class and define the population loss $\mathcal{L}(\operatorname{\mathbf{X}}_{k},\operatorname{\mathbf{w}}_{k})$ at federation member $k$ as the cross-entropy loss $\mathcal{L}(\operatorname{\mathbf{X}}_{k},\operatorname{\mathbf{w}}){=}\frac{1}{N}\sum_{x_{i}\in\operatorname{\mathbf{X}}}\sum_{j=1}^{C}p(y{=}j)\operatorname{\mathbb{E}_{\operatorname{\mathbf{x}}|y=i}}\left[\log f_{j}(x_{i},\operatorname{\mathbf{w}}_{k})\right]$.

In an IID setting, the optimization problem minimizing $\mathcal{L}_{IID}$ can be expressed as $K$ sub-problems that train local models via stochastic gradient descent (SGD), which are then communicated to a server for aggregation after each communication round (indexed by $t$):

FedAvg [1] makes use of the weighted averaging aggregation function $\operatorname{\mathbf{w}}_{t}^{f}=\sum_{k}^{K}\frac{n_{k}}{n}\operatorname{\mathbf{w}}_{t,k}^{*}$, where $\frac{n_{k}}{n}$ is the fraction of the $k$-th member’s local dataset with repect to the overall dataset size. Unfortunately, in a non-IID setting, this aggregation is not a valid approximation.

We follow [5] in quantifying non-IID-ness with the help of the Earth Mover Distance (EMD) [21] and write $d(\mathcal{D}_{1},...,\mathcal{D}_{K})$ for the EMD of $K$ distributions $\mathcal{D}_{1},...,\mathcal{D}_{K}$ obtained by averaging the set of pairwise EMDs (cf. Apx. A).

Weighted averaging fails as an aggregation function under non-IID datasets [5] because: (1) if weights in client models are all initialized identically, the label-space non-IID-ness between client dataset distributions will be the dominant factor and weights will diverge; and (2) if weights are initialized differently, client models risk converging to different solutions ³³3Without a formal proof,a third contribution term to weight divergence could be added in terms of feature-space non-iid. The steepness of local loss function pockets in relation to the local flatness of the loss manifold exacerbates the divergence after aggregation and hence leads to sub-optimal global model performance [22], which we address with Geometric Monte Carlo Clustering.

In a non-IID setting, individual client models are set to converge during training towards distinct local minima, i.e., different coordinates in weight space on the training loss manifold (see Fig. 1). In particular, when the non-IID-ness lies exclusively in the feature space between two client datasets $\{\operatorname{\mathbf{X}}_{1},\operatorname{\mathbf{y}}_{1}\}$ and $\{\operatorname{\mathbf{X}}_{2},\operatorname{\mathbf{y}}_{2}\}$, we hypothesize the existence of a continuous transformation $\mathcal{T}$ (e.g. pixel rotation, color inversion, shape distortion, etc.) that maps, on average, one subset of features $\operatorname{\mathbf{X}}_{1}^{\prime}{\in}\operatorname{\mathbf{X}}_{1}$ detained by one client to a subset of features $\operatorname{\mathbf{X}}_{2}^{\prime}{\in}\operatorname{\mathbf{X}}_{2}$ of a different client $\mathcal{T}{:}\operatorname{\mathbf{X}}_{1}{\rightarrow}\operatorname{\mathbf{X}}_{2}$. As a consequence, we suppose that the model weights of one client are trained to encode a set of transformed features with respect to another client. Hence we hypothesize the existence of a continuous transformation $\Gamma$ that maps one subset of weights to another one $\Gamma:\operatorname{\mathbf{w}}^{\prime}_{1}\subseteq\operatorname{\mathbf{w}}_{1}\rightarrow\operatorname{\mathbf{w}}^{\prime}_{1}\subseteq\operatorname{\mathbf{w}}_{1}$ and that can alternatively be modeled by a continuous affine connection via a curve parametrization $\gamma_{\theta}(u)$. This curve connects $\operatorname{\mathbf{w}}_{1}$ and $\operatorname{\mathbf{w}}_{2}$ (with $\gamma_{\theta}(0){=}\operatorname{\mathbf{w}}_{1}$ and $\gamma_{\theta}(1){=}\operatorname{\mathbf{w}}_{2}$) ideally on a loss surface where the loss value does not vary. That is, two neural nets parametrized by $\operatorname{\mathbf{w}}_{1}$ and $\operatorname{\mathbf{w}}_{2}$ agree on the output space given the same input. Along this curve of invariant loss reside a family of intermediary model weights that can be aggregated to produce a richer global model, incorporating the knowledge of both client models. In contrast, under label-space IID violations, particularly concept shifts, we expect no continuous curve to be found (e.g. between $\operatorname{\mathbf{w}}_{3}$ and $\operatorname{\mathbf{w}}_{4}$ on Fig. 1) since they disagree on the classification of a same input feature. In this case, a well constructed clustering rule should separate both models to avoid negative knowledge transfer. This curve finding is at the heart of our novel Federated Geometric Monte Carlo Clustering algorithm (FedGMCC).

To evaluate our approach, we show that real-life data is in fact non-IID in the feature space and compare the accuracy of our approach — Federated Geometric Monte-Carlo Clustering (FedGMCC) — against state-of-the art federated learning approaches (FedAVG [1], FedProx [6], CFL [13] and standard SGD). We leverage the well known image dataset EMNIST62 [23] and two sequential genomic datasets (SENSG-A and SENSG-R) we created (see Appx. B). EMNIST62 contains 814255 handwritten characters, labeled as 62 unbalanced classes in 28x28 pixel format. The genomic dataset contains reads (i.e., sequences of the bases A, T, C, G) of size 150 and every position is labelled according to whether or not it is part of a genomic variation. Such labelling is important, for example, to filter out and protect private information in the genomic information processing pipeline [24]. We introduce a concept shift in SENSG-R by randomly flipping the label of the most recurrent genomic features. EMNSIT62 and SENSG-A were distributed among $K$ clients following the procedure detailed in Appx. B. This partitioning introduces non-IID-ness in the feature space but also in the label space because certain classes might be over-represented in some client datasets compared to others. The genomic dataset SENSG-R is naturally partitioned using the populations of individuals (Asian, European, African). In addition, we partition both sets artificially into 10 client datasets to consider a wider range of non-IID-ness. Models were trained on an AMD Ryzen7 3700x system with 8 3.6 GHz cores, a NVIDIA Geforce RTX 3090 GPU with 10496 CUDA cores and 24 GB of GDDR6X memory. We used the binary-cross entropy loss function for the training of all classifiers and Tensorflow 2.6 [25].

We presented FedGMCC, a federated learning framework that consists of novel clustering rules and a new aggregation procedure. Substituting real datasets by a surrogate Monte Carlo dataset, we show how the curve finding procure can reveal the geometric connection between congruent models serving as a clustering rule of client model weights. Furthermore, FedGMCC aggregation rule averages all the intermediary model weights on the curve parametrization, leading to better generalization and extending the classical FedAvg aggregation rule to weight spaces of curves training loss manifold.

FedGMCC outperformed other FL training algorithms on the EMNSIT62 and genomic sequence sensitivity classification tasks where we controlled the non-IID-ness using an artificial partitioning technique and the EMD measure. In high non-IID setting FedGMCC yielded convergence rates respectively 36% and 8% faster than those of FedAvg and FedProx on EMNSIT62. FedGMCC outperformed FedAvg, FedProx and CFL on the genomic datasets by 63$\%$, 19$\%$ and 4 $\%$.