Relax and Merge: A Simple Yet Effective Framework for Solving Fair $k$-Means and $k$-sparse Wasserstein Barycenter Problems

Authors: Shihong Song, Guanlin Mo, Hu Ding
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: clustering, k-means, fairness, approxiamte algorithm, optimal transport

Abstract:
The fairness of clustering algorithms has gained widespread attention across various areas, including machine learning, In this paper, we study fair $k$-means clustering in Euclidean space. 
  Given a dataset comprising several groups, the fairness constraint requires that each cluster should contain a proportion of points from each group within specified lower and upper bounds. 
  Due to these fairness constraints, determining the optimal locations of $k$ centers is a quite challenging task. 
  We propose a novel ``Relax and Merge'' framework that returns a $(1+4\rho + O(\epsilon))$-approximate solution,  where $\rho$ is the approximate ratio of an off-the-shelf vanilla $k$-means algorithm and $O(\epsilon)$ can be an arbitrarily small positive number. If equipped with a PTAS of $k$-means, our solution can achieve an approximation ratio of $(5+O(\epsilon))$  with only a slight violation of the fairness constraints, which improves the current state-of-the-art approximation guarantee. Furthermore, using our framework, we can also obtain a $(1+4\rho +O(\epsilon))$-approximate solution for the $k$-sparse Wasserstein Barycenter problem, which is a fundamental optimization problem in the field of optimal transport, and a $(2+6\rho)$-approximate solution for the strictly fair $k$-means clustering with no violation, both of which are better than the current state-of-the-art methods. In addition, the empirical results demonstrate that our proposed algorithm can significantly outperform baseline approaches in terms of clustering  cost.

Decision: accept (poster)
Reviewer scores: 8, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730436931383

thread_type:
official_review

<thread_object>
creator: reviewer_mejb

<content>
summary:
The paper gives an approximation algorithm for the fair k-means algorithm, where, in addition to minimizing the sum of squared distances, the clusters should also satisfy certain minimum and maximum size constraints (stated as fractions of the data points). The paper gives (1+4c+O(\eps))-approximation algorithm with running time that has an exponential dependence on the dimension d and polynomial dependence on other parameters. The suggested algorithm first finds an \eps-centroid set T using Matousek's algorithm. This is followed by assigning points to centers in the set T by solving an LP that guarantees fairness. The size of T is exponential in dimension d. Since k centers need to be produced, the algorithm finds k centers by running a c-approximation algorithm for the vanilla k-means algorithm over points that are "moved" to their respective centers in T. Finally, integer assignments are found using ideas of Bera et al. (LP rounding), which causes minor violations in fairness conditions. This two-step process (clustering points using T followed by clustering T) is called Relax-and-Merge. Similar ideas have been used in the past (see, e.g., [A]), and the approximation analysis of the two-step process is similar to previous results (see Theorem 4 in [A]).

[A] https://theory.stanford.edu/~nmishra/Papers/clusteringDataStreamsTheoryPractice.pdf

Additional comment: I have reviewed an earlier version of the paper, and much of the review is from my earlier review. The current version improves the number of violations in the special case of mutually disjoint groups. This is shown by setting up an appropriate minimum-cost circulation problem.

strengths:
- Fair clustering is a relevant problem, and the paper gives an approximation algorithm better than the known results in the Euclidean setting.

weaknesses:
- The centroid set-based algorithm is specific to the Euclidean setting and cannot be extended to general metric spaces.
- The exponential dependence of the running time on d may be problematic for high-dimensional data.
- The techniques have mainly been borrowed from previous works. So, the theoretical contribution has limited novelty. 
- PTAS for vanilla k-means either have exponential dependence on k or d. So, the ~ 5-approximation is at the cost of increasing the running time.

questions:
Some questions can be found in the other part of review.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730648959272

thread_type:
official_review

<thread_object>
creator: reviewer_9lbe

<content>
summary:
This paper considers the Euclidean fair $k$-means problem. In fair $k$-means, the input data set $G$ is partitioned into $m$ groups $G_1,..,G_m$ and $2m$ numbers $\alpha_1,..,\alpha_m,\beta_1,...\beta_m$ such that $\forall i\in[m], 0\leq \alpha_i\leq \beta_i\leq 1$. The task is to find $k$ clustering centers and partition $G$ into $k$ clusters so that for every $i\in [k],j\in [m]$ cluster $P_i$ must contain at least $\alpha_j|P_i|$  points and at most $\beta_j|P_i|$ points from group $G_j$.

The main contribution of this paper is two improved approximation algorithms for fair $k$-means. The first algorithm shows one can compute a $1+4\rho+\epsilon$ approximation with constant violation in fairness constraints by running upon an existing $\rho$-approximation algorithm for vanilla Euclidean $k$-means. This algorithm also relies on a centroid of size $O(n\epsilon^{-d}\log \epsilon^{-1})$ thus the running time is exponential on $d$. The second algorithm is designed for the strictly fair $k$-means, namely $\forall i, \alpha_i=\beta_i=\frac{1}{m}$. A $ 2+6\rho$ approximation without violation is obtained. Both algorithms have better approximation ratios than previous results. The authors also observe that the first algorithm can be extended to a related problem called $k$-sparse Wasserstein Barycenter.

Technically, authors observe that one can first compute an optimal assignment from the data set to a centroid set and then reduce the problem to the vanilla $k$-means by assigning each data point to its center in the centroid set. The idea is simple but effective. Experiments also show improvement over previous results.

Overall, I think this paper reaches the acceptance bar.

strengths:
1. Improved approximation algorithms for an important problem.
2. The idea is simple but very useful.
3. Experiments show it is also promising for practice.

weaknesses:
1. The running time of the first algorithm is exponential on $d$. Thus the result does not scale in high dimensions.
2. Eventually, the main contribution of this paper is a reduction via a centroid set. The performance relies on the centroid set construction algorithm and the approximation algorithm for vanilla $k$-means.

questions:
1. Can one extend the results to the fair $k$-median?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730715189637

thread_type:
official_review

<thread_object>
creator: reviewer_gbws

<content>
summary:
This paper studies the fair $k$-means clustering problem in Euclidean space, where the goal is to partition the given dataset into several groups such that each group should satisfy specific lower and upper bounds on the proportions of the assigned points. While the fair $k$-means clustering problem has been extensively studied in recent years, the existing approximation algorithms often incur at least a constant factor loss on the clustering quality guarantee in order to find good clustering centers for fair clustering instances that do not satisfy the “Voronoi properties”. To reduce the approximation loss, this paper introduces the notion of $\epsilon$-approximation centroid set, which can well approximate the locations for the optimal clustering centers under fairness constraints. Building on notion of the $\epsilon$-approximation centroid set, this paper proposes a Relax-and-Merge framework. The Relax-and-Merge framework first decomposes the given clustering instance into smaller partitions satisfying the fairness constraints (with more than $k$ partitions obtained) to construct a weighted instance. Then, by executing a weighted standard $k$-means algorithm on the weighted instance constructed, a solution satisfying the fairness constraints can be obtained  using linear programming rounding method, with better approximation guarantee on the clustering quality. Additionally, this paper shows that the proposed framework can also be extended to solve the $k$-sparse Wasserstein barycenter problem and strictly fair $k$-means problem. The experiments show that the proposed Relax-and-Merge framework achieves better clustering cost compared with other fair clustering algorithms, which further demonstrates the effectiveness of the proposed methods.

strengths:
The strength of this paper can be summarized as follows.

1. Fair clustering is a topic that receives much attention in machine learning. If equipped with a PTAS for the $k$-means problem, the proposed algorithm can achieve the best approximation ratio of $5+O(\epsilon)$, which improves the current state-of-the-art approaches on clustering quality guarantees.

2. The proposed Relax-and-Merge framework is simple and effective. The technical analysis proposed to bound the loss when reducing the number of partitions to exactly $k$ is interesting.

3. This paper is well-written and the proofs are easy to follow. 

4. The empirical evaluations demonstrate that the proposed framework achieves much better performances on clustering quality compared to the previous works.

weaknesses:
The weakness of this paper can be summarized as follows.

1. The approximation guarantees of the proposed method rely on the approximation loss of the weighted algorithm used to reduce the number of partitions. Although many weighted algorithms (such as the local search algorithms) can achieve near-optimal clustering results in practice, the worst-case analysis rely on the use of a PTAS algorithm to achieve better approximation guarantees. 

2. The proposed $\epsilon$-approximation centroid set is designed specifically to the Euclidean space. It is unclear that whether such methods can be extended to the general metric spaces.

questions:
Q1: It seems that the proposed Relax-and-Merge framework is designed specifically for the Euclidean space. However, many fair algorithms can also be used to solve the case when data points are in a general metric space. Can the proposed method achieve better performances in general metric spaces compared with other fair clustering algorithms？

Q2: Since the $\epsilon$-approximation centroid set construction method does not rely on a specific constraint during the clustering process, whether the proposed framework can be extended to solve other fair clustering problems using $\epsilon$-approximation centroid set construction method? 

Q3: In the experimental parts, the authors only compare the clustering quality of the proposed algorithms with other fair clustering algorithms. Can the authors give a detailed comparison between the running time of the proposed framework and other existing fair clustering algorithms?

Q4: The clustering quality of the proposed algorithm rely heavily on the weighted algorithm used to reduce the number of partitions. However, in the experimental parts, there is no detailed description of which weighted algorithm is chosen and how the weighted algorithm is executed. Which weighted $k$-means algorithm is used in the experiments for Algorithm 1? Did the authors choose a PTAS for the experiments?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730799898263

thread_type:
official_review

<thread_object>
creator: reviewer_9f2z

<content>
summary:
$ (\alpha, \beta) $ -fair k-means clustering is a group-based notion of fairness where the number of points belonging to each protected 'group' (race, gender etc.) inside each cluster is bounded. The paper provides algorithms for $ (\alpha, \beta) $ -fair k-means clustering and strictly $ (\alpha, \beta) $ -fair k-means clustering that have better approximation ratios than existing works. This is achieved by considering the group-based fairness constraints before choosing the location of the centers rather than first choosing centers and then trying to assign points to centers according to fairness constraints.  The paper also gives approximation guarantees for $k$- Sparse Wasserstein Barycenter problem that tries to find an 'average' distribution with a support size of $k$, for given discrete $m$ distributions, by treating the problem as a special case of fair clustering.  The theoretical results are supported by empirical evaluations.

strengths:
1) Fairness in clustering has gained a lot of attention in the past decade or so. For the case of $(\alpha, \beta) $ - fair clustering, the paper provides better approximation ratios than the existing works. 

2) The paper for the most part is written very nicely and is not difficult to follow the central ideas. Sufficient Proof sketches/ details are provided in the main part of the paper. Although I could not check all the proofs in details, the paper appears technically sound.

3) The paper also shows the effectiveness of the proposed algorithms by empirical results on different datasets. Some of the empirical results can be better presented (see weaknesses) but overall, it is a very good mix of theory and practice and very suitable for the venue.

4) The high-level idea of considering fairness while choosing the location of centers and not after choosing centers might give insights to design better algorithms other fairness definitions also.

weaknesses:
1) There is time complexity analysis of algorithm 1. However, there can be a comparison with the Bera et al. and Bohm et al. papers to get an idea of the tradeoff between approximation factor and time complexity. Empirical results should also result the time parameter. There are some experiments in the appendix, but they are just for your algorithm. Same goes for the violation ratio. 

2) In terms of presentation, I think, it would be better to move some experiments to the main paper and one more proof to the appendices.

questions:
1) Is it possible to obtain fair $k$- clustering without any fairness violation without the requirement of equally sized groups? 

2) Can the method be extended to other clustering costs like $k$-median? It appears that proofs rely heavily on the cost being squared Euclidean distance.

3) Can you clarify how the solution to the k-sparse Wasserstein barycenter solution is guaranteed to be a distribution? Is the optimization slightly changed?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732286513104

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
# Global Response

## Revised paper

We sincerely thank all the anonymous reviewers for their valuable and insightful comments. We have added some content to our revised paper according to the comments and questions. The changes are summarized below and marked in **blue** in the paper.
- *Page 1, Line 50-54*: We added more references about fair $k$-median algorithms.

 - *Page 9, Line 488*: We added a sentence stating that we use the $k$-means++ algorithm as our $k$-means solver.

- *Page 18, Appendix D*: We added a brief discussion on the k-sparse Wasserstein Barycenter problem for the sake of completeness.

 - *Appendix E*: We added detailed analysis on how to extend Algorithm 1 to $k$-median and $k$-means in general metric space.

 - *Appendix F.6*: We added an experiment comparing the running times of our Algorithm 1, Algorithm 2, and the baselines.

## Comparison of running time

For the comparison of running time, we added an experiment (the more detailed discussion is uploaded to Appendix F.6 of the revised paper). To shed some light, below we show an illlustrative example with $k=10$. We compare the running time of Algorithm 1, Algorithm 2 (for only strictly fair datasets), NIPS19 (Bera et al.), and ORL21 (Bohm et al., for only strictly fair datasets). The brief running time comparison is shown in the following table. 

|                 | Bank | Creditcard | Census1990 | Adult | Breastcancer | Moons | Cluto | Complex | Hypercube |
| :-------------: | ---- | :--------: | ---------- | :---: | :----------: | :---: | ----- | :-----: | --------- |
| Algorithm 1 (s) | 3.78 |    6.51    | 14.95      | 6.12  |     1.33     | 0.83  | 2.36  |  1.71   | 10.27     |
|   NIPS19 (s)    | 1.11 |    2.39    | 4.42       | 1.59  |     0.45     | 0.78  | 1.66  |  0.78   | 2.87      |
| Algorithm 2 (s) | /    |     /      | /          |   /   |      /       | 0.59  | 0.56  |  0.58   | 0.39      |
|    ORL21 (s)    | /    |     /      | /          |   /   |      /       | 0.30  | 0.72  |  0.72   | 4.33      |

From the table, we conclude that for most strictly fair datasets, Algorithm 2 runs fastest. However, in general (non-strictly fair) datasets, our Algorithm 1 takes longer running time compared to NIPS1. 

Here, we also provide some deeper insights for this phenomenon. Algorithm 1 and NIPS19 both have two important subprocedures: linear programming and the $k$-means algorithm. These two steps are the efficiency bottlenecks. Specifically, NIPS19 first runs the $k$-means algorithm (i.e., $k$-means++), and then  calls the LP solver once to compute the fractional assignment. A different part of our Algorithm 1 is that it calls the LP solver **twice**, once to compute the weights of candidate set $T$ and once to compute the fractional assignment, and calls the $k$-means algorithm once. In Algorithm 1, we only need to run $k$-means on $T$, which should be much smaller than the whole dataset, leading to less running time for the $k$-means subprocedure compared to NIPS19. However, the first call to the LP solver to compute the weight of $T$ consumes more time than the second call because $|T| > k$ usually. Generally speaking, LP takes more time than $k$-means, which means our Algorithm 1 usually runs slower than NIPS19. In fact, the first LP takes more than a half of overall time in most of the datasets according to the following table (the detailed comparison is shown in Table 5 and Table 6 of our revised paper).

|               | Bank | Creditcard | Census1990 | Adult | Cluto | Complex | Hypercube |
| :-----------: | ---- | :--------: | ---------- | :---: | ----- | :-----: | --------- |
| LP on $T$ (s) | 2.4  |    4.06    | 7.51       | 4.14  | 1.01  |  1.08   | 5.71      |
|  Overall (s)  | 3.78 |    6.51    | 12.99      | 6.12  | 2.36  |  1.71   | 10.27     |
|     Ratio     | 63%  |    62%     | 58%        |  68%  | 43%   |   63%   | 56%       |


Besides, we would like to point **an implicit benefit** of the construction of T is that it is irrelevant to the parameters $k$, $\alpha$, and $\beta$. So if we consider a real scenario that we need to repeatedly try different choices for these parameters (e.g., we may want to tune the value k and select the most satisfying result), the step of constructing T and performing linear programming on T can be seen as preprocessing on the data  before the tuning. Namely, we just need to run this preprocessing one time, and consequently the amortized cost over the whole tuning procedure can be reduced significantly.

As for strictly fair datasets, Algorithm 2 has an advantage in efficiency for most datasets.  The primary reason is that Algorithm 2 only calls the k-means algorithm once and does not need to solve the LP. As for ORL21, it needs to run k-means for each group and then choose the best one. As a result, ORL21 takes longer time than Algorithm 2, especially on the datasets with large number of groups.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The notion of $(\alpha, \beta)$ fair clustering was introduced to consider $k$-means and other clustering algorithms with fairness constraints. The paper provides improved approximation algorithms for this problem and also the problem of $k$-sparse Wasserstein Barycenter. The results are interest and the paper has some algorithmic novelty; the results are also supported by empirical evidence.