Federated $K$-Means Clustering via Dual Decomposition-based Distributed Optimization

While computational performance remains an issue for many machine learning problems, the increase in computing power and in the efficiency of optimization algorithms can render many challenging problems tractable. However, the inability to share data due to confidentiality reasons still necessitates the use of distributed algorithms. Fig. 1(a) shows a setting in which training data is stored across two different nodes. Each node can use its local data to train an individual machine learning model. By including a coordination layer the two training processes can be guided in a way that a global model is trained, without the need to share confidential data. Distributed training of a global model without sharing individual training data is often referred to as federated optimization or federated learning [8]. If the underlying optimization problems are still hard to solve, or if the training data is heterogeneous, the training process can be further divided into subproblems. Fig. 1(b) depicts the situation in which models of different node clusters are trained in a distributed manner which in turn are again coordinated to obtain a global model. This approach is often referred to as clustered federated learning [9].

Most algorithms for federated learning involve an averaging step of the model parameters of the individual nodes [10]. Federated learning methods have been applied in the context of manufacturing [11], healthcare [12], mobile devices [13], and smart city sensing [14]. Li et al. [15] and Liu et al. [16] provide surveys on federated learning while Chamikara et al. examine the privacy aspects related to external attacks [17]. Applying federated learning to heterogeneous data sets can lead to the deterioration of the model quality of individual nodes in regard to their training data, which might hinder their willingness to participate in such a setting. This issue is addressed through personalized federated learning [18, 19].

$K$-means clustering describes an unsupervised machine learning problem in which a set of observations/data is divided into $K$ disjoint clusters according to a similarity measure [20]. Clustering problems can be found in many practical applications such as image segmentation [21], customer market segmentation [22], or the identification of similar operating points in a production plant [23]. While $K$-means clustering is a well-studied problem, federated clustering, i.e., clustering of data across multiple nodes, has not been studied extensively in the literature yet. Dennis et al. present a one-shot federated clustering algorithm with heterogeneous data, where the clustering problem of each node is solved by Lloyd’s heuristic algorithm [24, 25]. Kumar et al. apply federated averaging to pre-trained models on separate devices and present an update strategy when new data points are added [26]. Li et al. address the security aspect of federated clustering by encoding the data of each node and applying Lloyd’s algorithm to the encoded data [27]. Stallmann and Wilbik extend fuzzy $c$-means clustering, i.e., a clustering problem where each data point can be assigned to multiple clusters, to a federated setting [28]. A similar approach to federated $c$-means clustering was previously presented by Pedrycz [29]. Wang et al. use model averaging and gradient sharing for federated clustering of data in a smart grid [30].

A common feature of the federated clustering approaches described in the literature is the use of a heuristic algorithm to solve the individual clustering problems. In this paper mixed-integer programming is used to solve the individual clustering problems and a dual decomposition-based distributed optimization approach is employed to coordinate the solutions of the different nodes. While an averaging step is still performed to obtain feasible primal solutions, the use of duality enables the computation of valid lower bounds on the objective of the global clustering problem. It should be noted that the mixed-integer programming problems resulting from the $K$-means clustering problem are very hard to solve due to their weak integer relaxations. Thus, the approach presented in this paper can be regarded as preliminary work that can become a viable option with the continuous improvement of mixed-integer programming solvers [31, 32].

We denote vectors with bold lower-case letters ($\mathbf{x}$) and matrices with bold upper-case letters ($\mathbf{X}$). The $j$th element of a vector $\mathbf{x}$ is denoted by $[\mathbf{x}]_{j}$. The Euclidean norm is denoted by $\|\cdot\|_{2}$. The value of a variable $\mathbf{x}$ in iteration $t$ is denoted by $\mathbf{x}^{(t)}$.

This section presents the mixed-integer programming-based formulation of the $K$-means clustering training problem. The formulation is subsequently extended to the case of distributedly stored data, which gives rise to a federated learning problem. Consensus constraints are used to couple the training problems of different nodes. These constraints can be dualized such that the federated learning problem can be solved via dual decomposition-based distributed optimization. Since the underlying optimization problem contains integrality constraints it is not convex and thus strong duality does not hold. However, a feasible primal solution can be computed in each iteration through an averaging heuristic.

By defining the vector of stacked cluster centroids of each node $i$,

the consensus constraints can be rewritten as

Constraints (8) can subsequently be rewritten in matrix form

or in a more compact way

with $\mathbf{A}_{i}\in\mathbb{R}^{K\cdot n_{\mathbf{y}}\cdot(N_{s}-1)\times K\cdot n_{\mathbf{y}}}$. By introducing dual variables $\bm{\lambda}\in{\mathbb{R}^{n_{K\cdot n_{\mathbf{y}}\cdot(N_{s}-1)}}}$ for the consensus constraints (10) the Lagrange function of problem (5) can be defined,

The minimization of the Lagrange function for a fixed value of the dual variables $\bm{\lambda}$ gives the corresponding value of the dual function.

The dual function has two important properties. First, the value of the dual function is always a lower bound on the solution of its corresponding primal problem, in this case, problem (5) [37]. The problem of finding the dual variables that result in the best lower bound is referred to as the dual optimization problem,

The resulting dual problem can be solved in a distributed manner by solving the individual clustering problems for the current value of the dual variables,

Second, the dual function (12) is always concave, regardless of whether the primal problem is convex or not [37]. Therefore the dual problem (13) is a convex optimization problem. However, the dual function is usually nondifferentiable due to a changing set of active individual constraints, which means that problem (13) is a nonsmooth optimization problem [38]. The following subsections present some algorithms for the solution of the dual problem, namely the subgradient method, the bundle trust method, and the quasi-Newton dual ascent algorithm.

The dual decomposition-based distributed clustering approach was evaluated on a set of benchmark problems of varying sizes. The data for each benchmark problem was generated randomly. First, initial cluster centroids $\mathbf{m}_{k}^{0}$ were generated, with $[\mathbf{m}_{k}^{0}]_{l}\in\mathcal{U}_{c}(-1,1),\;l=1,\dots,n_{\mathbf{y}}$. Then, for each cluster $k$ five random data points were added within a radius of 0.5 from the generated centroid. The parameters of the benchmark problems were varied as follows:

Five benchmark problems were generated for each combination of nodes, dimensions, and clusters, resulting in a total of 90 benchmark problems. A benchmark problem is characterized by its number of nodes, dimension of the data, and number of clusters. For instance, problem $\text{3N2D4K}_{\text{5}}$ is the 5th benchmark problem comprised of 3 nodes with 2-dimensional data sorted into 4 clusters.
The benchmark problems were solved using the subgradient method, the bundle trust method, and the quasi-Newton dual ascent algorithm.

The initial step size (subgradient method)/ trust region (BTM, QNDA) parameter was set to $\alpha^{(0)}=0.5$ and varied according to

The size of the bundle for BTM and QNDA was set to $\tau=50$ points. All algorithms were initialized with $\bm{\lambda}^{(0)}=\mathbf{0}$ and the initial approximated Hessian of the QNDA algorithm was set to the negative identity matrix. The algorithms were terminated either when the Euclidean norm of the primal residual

i.e., the violation of the consensus constraints, lied below a threshold of $\epsilon_{p}=10^{-2}$ or when the relative duality gap (35) reached a value of $\epsilon_{DG}=0.25\;\%$. The used parameters for the different algorithms are summarized in Tab. I. The MIQCP clustering problems of all nodes were solved using the commercial solver Gurobi [45] and the total computation time was computed as

where $N_{\text{iter}}$ is the number of required iterations, $T_{\text{comm}}=800\;ms$ is the required communication time between the coordinator and the subproblems, which is assumed to be constant, $T_{\text{update}}^{(t)}$ is the time required by the coordinator to update the dual variables in iteration $t$ and $T_{\text{sub},i}^{(t)}$ is the solution time for the clustering problem of node $i$ in itereation $t$.

Fig. 6 shows the evolution of the relative duality gap for benchmark problem $\text{2N2D4K}_{\text{3}}$. The subgradient method converges rather slowly. In comparison, the BTM and QNDA algorithms exhibit a faster rate of convergence. Between these two algorithms, BTM exhibits an oscillatory behavior before converging. In contrast, the QNDA algorithm does not exhibit oscillations and therefore converges earlier. Additionally, it should be noted that the QNDA algorithm achieves a relative duality gap of $0\;\%$, i.e., it converges to a proven global optimum.

Fig. 7 further illustrates the results. Fig. 7(a) and 7(b) show the results of the clustering in the first iteration, i.e., the individual global optima. Fig. 7(c) and 7(d) depict the solutions upon convergence of the QNDA algorithm. Each node computes the same cluster centroids corresponding to the globally optimal solution with respect to the entire data set, but not to the individual data sets. It is therefore possible to compute a global model locally in each node while only accessing local data.

As shown in the previous section, solving the MIQCP clustering problems is computationally expensive. This is due to the weak integer relaxation of problem (2), which means that the solution of the relaxed problem within the branch-and-bound algorithm is far away from the integer solution. This results in slow-moving relative integrality gaps and slow convergence of the solution algorithm. While the main motivation of the distributed clustering approach is the training of a global model without the exchange of local data, it can also be used to efficiently solve larger clustering problems. Fig. 8 depicts the evolution of the relative duality gap of the QNDA algorithm as well as the evolution of the relative integrality gap of Gurobi for the complete data set of benchmark problem $\text{4N4D4K}_{3}$. The clustering problems of the individual nodes were solved sequentially in the case of QNDA, also using Gurobi. While the relative gap of the central solution improves very slowly, the QNDA algorithm quickly converges to a solution close to the global optimum. Note, that both relative gaps prove a worst-case distance to the global optimum. Hence, decomposing a large clustering problem into smaller subproblems and coordinating the solutions via a distributed optimization algorithm can offer significant performance improvements compared to a central solution.