Wasserstein $k$-means with sparse simplex projection

We denote scalars with lower-case letters $(a,b,\ldots)$, vectors as bold lower-case letters $(\mbox{\boldmath$a$},\mbox{\boldmath$b$},\ldots)$, and matrices as bold-face capitals $(\mbox{\bf A},\mbox{\bf B},\ldots)$. The $i$-th element of $a$ and the element at the $(i,j)$ position of A are represented as $\mbox{\boldmath$a$}_{i}$ and $\mbox{\bf A}_{ij}$, respectively. $\mbox{\boldmath$1$}_{d}$ is used for the $d$-dimensional vector of ones. Operator $(\cdot)^{T}$ stands for the matrix transpose. Operator $[\mbox{\boldmath$a$}_{i}]_{+}$ represents $\mathrm{max}(a,0)$ that outputs $a$ when $a\!\geq\!0$, and $0$ otherwise. Given a set $\mathcal{S}\subseteq\mathcal{N}=\{1,\ldots,n\}$, the complement $\mathcal{S}^{c}$ is defined with respect to $\mathcal{N}$, and the cardinality is $|\mathcal{S}|$. The support set of $a$ is ${\rm supp}(\mbox{\boldmath$a$})=\{i:\mbox{\boldmath$a$}_{i}\neq 0\}$. $\mbox{\boldmath$a$}_{|\mathcal{S}}$ extracts the elements of $\mathcal{S}$ in $a$, of which size is $|\mathcal{S}|$. $\mathbb{R}_{+}^{n\times m}$ represents a nonnegative matrix of size ${n\times m}$. The unit-simplex, simply called simplex, is denoted by $\Delta_{n}$, which is the subset of $\mathbb{R}^{n}$ comprising all nonnegative vectors whose sums is 1. $\lfloor a\rfloor$ represents the floor function, which outputs the greatest integer less than or equal to $a$.

The goal of the standard Euclidean $k$-means algorithm is to partition given $q$ data $\mbox{\bf X}=\{\mbox{\boldmath$x$}_{1},\ldots,\mbox{\boldmath$x$}_{q}\}\subset\mathbb{R}^{d}$ into $k$ separated groups, i.e., $k$ clusters, as $\mathcal{C}=\{C_{1},\ldots,C_{k}\}$ such that the following $k$-means cost function is minimized:

where $\mbox{\bf C}=\{\mbox{\boldmath$c$}_{1},\ldots,\mbox{\boldmath$c$}_{k}\}\subset\mathbb{R}^{d}$ is a set of $k$ cluster centers, designated as centroids hereinafter. Also, $\|\cdot\|_{2}^{2}$ represents the squared Euclidean distance. Consequently, our aim is to find $k$ cluster centroids C. It is noteworthy that the problem defined in (1) is non-convex. It is known to be an NP-hard problem [20, 21]. Therefore, many efforts have been undertaken to seek a locally optimized solution of (1). Among them, the most popular but simple and powerful algorithm is the one proposed by Lloyd [2], called Lloyd’s $k$-means. This algorithm consists of two iterative steps where the first step, called the assignment step, finds the closest cluster centroid $\mbox{\boldmath$c$}_{j}(\in[k])$ for each data sample $\mbox{\boldmath$x$}_{i}(\in[q])$. It then updates the cluster centroids C at the second step, called the update step. These two steps continue until the clusters do not change anymore. The algorithm is presented in Algorithm 1.

Although Lloyd’s $k$-means is simple and powerful, the obtained result is not guaranteed to reach a global optimum because the problem is not convex. For that reason, several efforts have been launched to find good initializations of Lloyd’s algorithm [22, 23, 24]. Apart from the lack of the global guarantee, Lloyd’s algorithm is adversely affected by high computation complexity. The highest complexity comes from the assignment step, where the Euclidean distances among all the data samples and all cluster centroids must be calculated. This cost of complexity is $\mathcal{O}(qk)$, which is infeasible for numerous data samples or clusters. It is also problematic for higher dimension $d$. For this particular problem, many authors have proposed efficient approaches. One category of approaches is to generate a hierarchical tree structure of the data samples and to make the assignment step more efficient [3, 4, 5]. However, this also entails high costs to maintain this structure. It becomes prohibitively expensive for larger dimension data samples. Another category is to reduce redundant calculations of the distance at the assignment step. For example, if the assigned centroid changes closer to a data sample $x$ in successive iterations, then it is readily apparent that all other centroids that did not change cannot move to $x$. Consequently, the distance calculations to the corresponding centroids can be omitted. In addition, the triangle inequality is useful to skip the calculation. This category includes, for example, [6, 7, 8, 9, 10]. However, it remains unclear whether these approaches are applicable to Wasserstein $k$-means described later in Section 2.4.

Let $\Omega$ be an arbitrary space, $d(\cdot,\cdot)$ a metric on that space, and $P(\Omega)$ the set of Borel probability measures on $\Omega$. For any point $\nu\in\Omega$ is the Dirac unit mass on $\nu$. Here, we consider two families $\mbox{\boldmath$\nu$}=(\nu_{1},\ldots,\nu_{m})$ and $\mbox{\boldmath$\mu$}=(\mu_{1},\ldots,\mu_{n})$ of point in $\Omega$, where $a$ and $b$ respectively satisfy $\mbox{\boldmath$a$}\!=\!(a_{1}\ldots,a_{m})^{T}\in\mathbb{R}_{+}^{m}$ and $\mbox{\boldmath$b$}\!=\!(b_{1}\ldots,b_{n})^{T}\in\mathbb{R}_{+}^{n}$ with $\sum_{i}^{m}a_{i}=\sum_{j}^{n}b_{j}=1$. The ground cost matrix $\mbox{\bf C}_{\scriptsize\mbox{\boldmath$\nu\mu$}}$, or simply C, represents the distances between elements of $\nu$ and $\mu$ raised to the power $p$ as $\mbox{\bf C}_{{\scriptsize{\mbox{\boldmath$\nu\mu$}}}}\ =\ \mbox{\bf C}:=[d(\nu_{u},\mu_{v})^{p}]_{uv}\ \ \in\mathbb{R}_{+}^{m\times n}$. The cost transportation polytope $\mathcal{U}_{mn}$ of $a\in\Delta_{m}$ and $b\in\Delta_{n}$ is a feasible set defined as the set of $m\times n$ nonnegative matrices such that their row and column marginals are respectively equal to $a$ and $b$. It is formally defined as

Consequently, the optimal transport matrix $\mbox{\bf T}^{*}$ is defined as a solution of the following optimal transport problem [17]

The Wasserstein $k$-means algorithm replaces the distance metric $\|\mbox{\boldmath$x$}_{i}-\mbox{\boldmath$c$}_{j}\|_{2}^{2}$ in (1) with the following the $p$-Wasserstein distance $W_{p}(\mbox{\boldmath$\mu$},\mbox{\boldmath$\nu$})$, which is formally defined as shown below.

In addition, the calculation in Step 2 in Algorithm 1, i.e., the update step, is replaced with the Wasserstein barycenter calculation, which is defined as presented below.

Addressing that $m=n$ holds in this case, we assume $\{\mbox{\bf T},\mbox{\bf C}\}\subset\mathbb{R}^{n\times n}_{+}$ hereinafter. Then, the overall algorithm is summarized in Algorithm 2.

The computation of the Wasserstein barycenters has recently gained increasing attention in the literature because it has a wide range of applications. Many fast algorithms have been proposed, and they include, for example, [26, 27, 28, 16, 29, 30]. Although they are effective at the update step in the Wasserstein $k$-means algorithm, they do not directly tackle the computational issue at the assignment step that this paper particularly addresses. [31] proposes a mini-batch algorithm for the Wasserstein $k$-means. However, it is not clear whether it gives better performances than the original Wasserstein $k$-means algorithm due to the lack of comprehensive experiments. Regarding another direction, the direct calculation of a Wasserstein barycenter of sparse support may give better performances on the problem of interest in this paper. However, as suggested and proved in [32], the finding such a sparse barycenter is a hard problem, even in dimension $2$ and for only $3$ measures.

The optimal transport problem in (2) described in Section 2.3 is reformulated as

Consequently, this problem is reduced to the linear programming (LP) problem with $(m+n)$ linear equality constraints. The LP problem is an optimization problem with a linear objective function and linear equality and inequality constraints. The problem is solvable using general LP problem solvers such as the simplex method, the interior-point method and those variants. The computational complexity of those methods, however, scales at best cubically in the size of the input data. Here, if the measures have $n(=m)$ bits, then the number of unknowns $[\mbox{\bf T}]_{uv}$ is $n^{2}$, and the computational complexities of the solvers are at best $\mathcal{O}(n^{3}\log n)$ [33, 34]. To alleviate this difficulty, the entropy-regularized algorithm, known as Sinkhorn algorithm, is proposed under the entropy-regularized Wasserstein distance. This algorithm regularizes the Wasserstein metric by the entropy of the transport plan and enables much faster numerical schemes with complexity $\mathcal{O}(n^{2})$ or $\mathcal{O}(n\log n)$. It is nevertheless difficult to obtain a higher-accuracy solution. In addition, it can be unstable if the regularization parameter is reduced to a smaller value to improve the accuracy.

Subsequently, this paper continues to describe the original LP problem in (3.1). We attempt to reduce the size of the measures. More specifically, we adopt the following approaches: (i) we make the input data samples $\nu$ and the centroids $c$ sparser than the original one by projecting them onto sparser simplex, and (ii) we shrink the projected data sample $\nu$ and centroid $c$ and the corresponding ground cost matrix $\mbox{\bf C}_{\mbox{\boldmath$\nu c$}}$ by removing their zero elements. Noteworthy points are that the first approach enables to speed-up the computation while maintaining degradation of the clustering quality as small as possible. The second approach yields no degradation.

We first consider the sparse simplex projection $\text{Proj}^{\gamma(t)}(\cdot)$, where $\gamma(t)\in(0,1]$ is the control parameter of the sparsity, which is described in Section 3.4. $\text{Proj}^{\gamma(t)}(\cdot)$ projects the $i$-th data sample $\mbox{\boldmath$\nu$}_{i}$ and the $j$-th centroid $\mbox{\boldmath$c$}_{j}$, respectively, into sparse $\hat{\mbox{\boldmath$\nu$}}_{i}$ and $\hat{\mbox{\boldmath$c$}}_{j}$ on $\Delta_{n}$. For this purpose, we exploit the projection method proposed in [35], which is called GSHP.

More concretely, denoting the original $\mbox{\boldmath$\nu$}_{i}$ or $\mbox{\boldmath$c$}_{j}$ as $\mbox{\boldmath$\beta$}\in\Delta_{n}$, and the projected $\hat{\mbox{\boldmath$\nu$}}_{i}$ or $\hat{\mbox{\boldmath$c$}}_{j}$ as $\hat{\mbox{\boldmath$\beta$}}\in\Delta_{n}$, then $\hat{\mbox{\boldmath$\beta$}}$ is generated as

where $\kappa=\lfloor n\cdot\gamma(t)\rfloor$, and where $\mathcal{S}^{\star}$ is defined as $\mathcal{S}^{\star}={\rm supp}(\mathcal{P}_{\kappa}(\mbox{\boldmath$\beta$}))$, where $\mathcal{P}_{\kappa}(\mbox{\boldmath$\beta$})$ is the operator that keeps the $\kappa$-largest elements of $\beta$ and sets the rest to zero. Additionally, the $v(\in[n])$-th element of $\mathcal{P}_{\Delta_{\kappa}}(\mbox{\boldmath$\beta$}_{|\mathcal{S}^{\star}})$ is defined as

where $\tau$ is defined as

This computational complexity is $\mathcal{O}(n\min(\kappa,\log(n)))$.

The sparse simplex projection in the previous section increases the number of zeros in the centroid and data samples. We reduce the computational complexity by harnessing this sparse structure. The vector shrinking operator, denoted as $\text{shrink}(\cdot)$, removes the zero elements from the projected sample $\hat{\mbox{\boldmath$\nu$}}_{i}$ and centroid $\hat{\mbox{\boldmath$c$}}_{i}$, and generates $\tilde{\mbox{\boldmath$\nu$}}_{i}$ and $\tilde{\mbox{\boldmath$c$}}_{j}$, respectively. It should be noted that this operation does not produce any degradations in terms of the solution of the LP optimization problem. More specifically, $\hat{\mbox{\boldmath$\nu$}}_{i}$ can be calculated as

where $\mathcal{S}_{\rm samp}={\rm supp}(\tilde{\mbox{\boldmath$\nu$}}_{i})$. Similarly to this, we calculate

where $\mathcal{S}_{\rm cent}={\rm supp}(\tilde{\mbox{\boldmath$c$}}_{i})$.

Accordingly, we must also shrink the ground cost matrix $\mbox{\bf C}_{\mbox{\boldmath$\nu c$}}$ based on the change of the size of $\tilde{\mbox{\boldmath$c$}}_{i}$ and $\tilde{\mbox{\boldmath$\nu$}}_{i}$. For this purpose, we newly introduce the matrix shrinking operator $\text{Shrink}(\cdot)$, which removes the row and the column vectors from $\mbox{\bf C}_{\mbox{\boldmath$\nu c$}}$. The removal is performed against $(\mathcal{S}_{\rm cent})^{c}$ and $(\mathcal{S}_{\rm samp})^{c}$. Consequently, $\mbox{\bf C}_{\mbox{\boldmath$\nu c$}}$ is compressed into $\tilde{\mbox{\bf C}}$ using $\text{Shrink}(\cdot)$ as

It should also be noted that the size of $\tilde{\mbox{\bf C}}$ is $|\mathcal{S}_{\rm samp}|\times|\mathcal{S}_{\rm cent}|$. Those sizes change according to the control parameter of the sparse ratio $\gamma(t)$, which is described in the next subsection.

As described in Section 3.2, the control parameter of the sparsity is denoted as $\gamma(t)$, where $t$ represents the iteration number $t$. We also propose three $\gamma(t)$ control algorithms, which are denoted respectively as ‘FIX’ (fixed), ‘DEC’ (decrease), and ‘INC’(increase). They are mathematically formulated as

where $\gamma_{\rm min}\in\mathbb{R}$ is the minimum value, and where $T_{\rm max}\in\mathbb{N}$ is the maximum number of the iterations. The overall algorithm of the proposed SSPW $k$-means is presented in Algorithm 3.

We first conducted a preliminary experiment using the COIL-100 object image dataset⁴⁴4http://www1.cs.columbia.edu/CAVE/software/softlib/coil-100.php [37], which includes $7,200$ images of $100$ objects. From this dataset, we randomly select $10$ classes and $10$ images per class. We first convert the pixel information into intensity with the range of $0$ to $255$. Then we generate a one-dimensional histogram of which the bin size is $255$ by removing the intensity of zero. We set $\gamma_{\rm min}=\{0.7,0.8\}$.

The averaged clustering performances on Purity, NMI and Accuracy over $5$ runs are summarized in TABLE 1. For ease of comparison, the difference values against the baseline method are presented in Fig. 1, of which panels (a), (b) and (c) respectively show results of Purity, NMI, and Accuracy. Here, positive values represent improvements against the baseline method. Additionally, Fig. 1 (d) presents the speed-up ratio of the computation time against that of the baseline method, where the value more than than $1.0$ indicate faster computations than that of the baseline. The figures specifically summarize the results in terms of the combinations of different $\gamma(t)$ algorithms and different $\gamma_{\rm min}$. They also show the differences among the different projection types, i.e., the projection of the centroid, data sample, and both the centroid and data sample. From Fig. 1, one can find that the case with the centroid projection stably outperforming other cases. Especially, the INC algorithm with the centroid projection stably outperforms others. However, it requires a greater number of computations. Moreover, the effectiveness in terms of the computation complexity reduction is lower. On the other hand, the FIX algorithm with $\gamma_{\rm min}=0.7$ engenders comparable improvements while keeping the computation time much lower.

In this experiment, comprehensive evaluations have been conducted. We investigate a wider range of $\gamma_{\rm min}$ than that of the earlier experiment, i.e., $\gamma_{\rm min}=\{0.5,0.6,0.7,0.8\}$. For this purpose, we used the USPS handwritten dataset⁵⁵5http://www.gaussianprocess.org/gpml/data/, which includes completely $9298$ handwritten single digits between $0$ and $9$, each of which consists of $16\times 16$ pixel image. Pixel values are normalized to be in the range of $[-1,1]$. From this dataset, we randomly select $10$ images per class. We first convert the pixel information into intensity with the range of $0$ to $255$. Then we generate a two-dimensional histogram for which the sum of intensities of all pixels is equal to one.

The averaged clustering qualities over $10$ runs are presented in TABLE 2. As shown there, the FIX algorithm outperforms others in terms of both the clustering quality and the computation time. Similarly to the earlier experiment, for ease of comparison, the difference values against the baseline method are also shown in Fig. 2. It is surprising that the proposed algorithm maintains the clustering quality across all the metrics as higher than the baseline method does, even while reducing the computation time. Addressing individual algorithms, although the INC algorithm exhibits degradations in some cases, the DEC algorithm outperforms the baseline method in all settings. Furthermore, as unexpected, the performances of the algorithms with lower $\gamma_{\rm min}$, i.e., $\gamma_{\rm min}=\{0.5,0.6\}$, are comparable with or better than those with $\gamma_{\rm min}=\{0.7,0.8\}$.

Second, the computation time under the same settings is presented in Fig. 2(d). From this result, it is apparent that the computation time is approximately proportional to $\gamma_{\rm min}$ in each $\gamma(t)$ control algorithm. In other words, the case with $\gamma_{\rm min}=0.5$ requires the lowest computation time among all. Moreover, with regard to the different projection types, the case in which the sparse simplex projection is performed onto both the centroid and data sample requires the shortest computation time, as expected.

This subsection investigates the convergence performances of the proposed algorithm. The same settings as those of earlier experiments are used. We first address the convergence speed in terms of the computing time with respect to the different projection types as well as different $\gamma_{\rm min}$. For this purpose, we examine the DEC algorithm. However, the two remaining algorithms behave very similar. Fig. 3 shows convergence of Accuracy in terms of computation time at the first instance of $10$ runs in the earlier experiment. Fig. 3(a)–Fig. 3(c) respectively show the results obtained when the projection is performed onto the data sample, centroid, and both the centroid and data sample. Regarding the different $\gamma_{\rm min}$, the case with the smaller $\gamma_{\rm min}$ indicates faster convergence. Apart from that, the convergence behaviors among three projection types are similar. However, the case with the which projection of both the centroid and data sample yields the fastest results among all. This is also demonstrated explicitly in Fig. 5 under $\gamma_{\rm min}=0.5$.

We also address performance differences in terms of the different algorithms of the $\gamma(t)$ control parameters. Fig. 5 presents the convergence speeds obtained when projecting both the centroid and data sample with $\gamma_{\rm min}=0.5$, where the plots explicitly show individual iteration. As might be readily apparent, the FIX control algorithm and the INC control algorithm require much less time than the others because these two algorithms can process much smaller centroids and data samples than others at the beginning of the iterations. However, the computation time of the INC control algorithm increases gradually as the iterations proceed, whereas that of the DEC control algorithm decreases gradually. Overall, the FIX control algorithm outperforms the other algorithms.

Finally, we compared performances on different sparsity ratios $\{0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9\}$ of $\gamma_{\rm min}$ under the same configurations of the earlier experiment. We specifically address the FIX algorithm of $\gamma(t)$ and the case where the projection is performed onto both the centroid and sample. Regarding the time comparison, we use the same speed-up metric against the baseline method as the earlier experiment. Therefore, the range of the values is more than $0.0$, and the value more than $1.0$ means a speed-up against the baseline method. For this experiment, we have re-conducted the same simulation, thereby the obtained value are slightly different from those in the earlier experiment due to its randomness.

The results are summarized in TABLE 3. Fig. 6 also demonstrates the results for ease of comparison. From both of the results, the cases with $\gamma_{\rm min}=\{0.3,0.4\}$ yield higher performances than other cases with respect to the clustering quality as well as the computation efficiency. Surprisingly, we obtained around $4$ points higher values in the clustering quality metrics with more than $20$-fold speed-up. On the other hand, $\gamma_{\rm min}=0.2$ results in the worst performance. Therefore, the optimization of $\gamma_{\rm min}$ to give best performances is one of the important topics in the future research.

From the exhaustive experiments above, it can be seen that the proposed proposed SSPW $k$-means algorithm dynamically reduce the computational complexity by removing lower-valued histogram elements and harnessing the sparse simplex projection while maintaining degradation of the clustering quality lower. Although the performance of each setting of the proposed algorithm depends on the dataset, we conclude that the FIX algorithm with the projections of both the centroid and sample under lower sparsity ratios $\gamma_{min}$ is the best option because it gives stably better classification qualities with the fastest computing times among all the settings. In fact, this algorithm gives comparable performances to the best one in TABLE 1, and outperforms the others in TABLE 2. The algorithm also yiedls around $4$ points higher values in the clustering quality metrics with more than $20$-fold speed-up as seen in TABLE 3.