SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$k$ -Median Clustering via Metric Embedding: Towards Better Initialization with Differential Privacy
Chenglin Fan, Ping Li, Xiaoyun Li
10900 NE 8th St. Bellevue, WA 98004, USA
{fanchenglin, pingli98, lixiaoyun996}@gmail.com
In clustering, the choice of initial centers is crucial for the convergence speed of the algorithms. We propose a new initialization scheme for the  $k$ -median problem in the general metric space (e.g., discrete space induced by graphs), based on the construction of metric embedding tree structure of the data. We propose a novel and efficient search algorithm which finds initial centers that can be used subsequently for the local search algorithm. The so-called HST initialization method can produce initial centers achieving lower error than those from another popular method  $k$ -median++, also with higher efficiency when  $k$  is not too small. Our HST initialization are then extended to the setting of differential privacy (DP) to generate private initial centers. We show that the error of applying DP local search followed by our private HST initialization improves prior results on the approximation error, and approaches the lower bound within a small factor. Experiments demonstrate the effectiveness of our proposed methods.
Clustering is an important classic problem in unsupervised learning that has been widely studied in statistics, data mining, machine learning, network analysis, etc. (Punj and Stewart, 1983; Dhillon and Modha, 2001; Banerjee et al., 2005; Berkhin, 2006; Abbasi and Younis, 2007). The objective of clustering is to divide a set of data points into clusters, such that items within the same cluster exhibit similarities, while those in different clusters distinctly differ. This is concretely measured by the sum of distances (or squared distances) between each point to its nearest cluster center. One conventional notion to evaluate a clustering algorithms is: with high probability,  $cost(C, D) \leq \gamma OPT_k(D) + \xi$ , where  $C$  is the centers output by the algorithm and  $cost(C, D)$  is a cost function defined for  $C$  on dataset  $D$ .  $OPT_k(D)$  is the cost of optimal clustering solution on  $D$ . When everything is clear from context, we will use  $OPT$  for short. Here,  $\gamma$  is called multiplicative error and  $\xi$  is called additive error. Alternatively, we may also use the notion of expected cost.
Two popularly studied clustering problems are 1) the  $k$ -median problem, and 2) the  $k$ -means problem. The origin of  $k$ -median dates back to the 1970's (e.g., Kaufman et al. (1977)), where one tries to find the best location of facilities that minimizes the cost measured by the distance between clients and facilities. Formally, given a set of points  $D$  and a distance measure, the goal is to find  $k$  center points minimizing the sum of absolute distances of each sample point to its nearest center. In  $k$ -means, the objective is to minimize the sum of squared distances instead. There are two general frameworks for clustering. One heuristic is the Lloyd's algorithm (Lloyd, 1982), which is built upon an iterative distortion minimization approach. In most cases, this method can only be applied to numerical data,
This work was initially submitted in 2020 and was made public in 2021.
typically in the (continuous) Euclidean space. Clustering in general metric spaces (discrete spaces) is also important and useful when dealing with, for example, the graph data, where Lloyd's method is no longer applicable. A more generally applicable approach, the local search method (Kanungo et al., 2002; Arya et al., 2004), has also been widely studied. It iteratively finds the optimal swap between the center set and non-center data points to keep lowering the cost. Local search can achieve a constant approximation (i.e.,  $\gamma = O(1)$ ) to the optimal solution (Arya et al., 2004). For general metric spaces, the state of the art approximation ratio is 2.675 for  $k$ -median Byrka et al. (2015) and 6.357 for  $k$ -means Ahmadian et al. (2017).
Initialization of cluster centers. It is well-known that the performance of clustering can be highly sensitive to initialization. If clustering starts with good initial centers with small approximation error, the algorithm may use fewer iterations to find a better solution. The  $k$ -median++ algorithm (Arthur and Vassilvitskii, 2007) iteratively selects  $k$  data points as initial centers, favoring distant points in a probabilistic way, such that the initial centers tend to be well spread over the data points (i.e., over different clusters). The produced initial centers are proved to have  $O(\log k)$  multiplicative error. Follow-up works further improved its efficiency and scalability, e.g., Bahmani et al. (2012); Bachem et al. (2016); Lattanzi and Sohler (2019); Choo et al. (2020); Cohen-Addad et al. (2021); Grunau et al. (2023); Fan et al. (2023). In this work, we propose a new initialization framework, called HST initialization, which is built upon a novel search algorithm on metric embedding trees constructed from the data. Our method achieves improved approximation error compared with  $k$ -median++. Moreover, importantly, our initialization scheme can be conveniently combined with the notion of differential privacy (DP) to protect the data privacy.
Clustering with Differential Privacy. The concept of differential privacy (Dwork, 2006; McSherry and Talwar, 2007) has been popular to rigorously define and resolve the problem of keeping useful information for machine learning models, while protecting privacy for each individual. DP has been adopted to a variety of algorithms and tasks, such as regression, classification, principle component analysis, graph distance release, matrix completion, optimization, and deep learning (Chaudhuri and Monteleoni, 2008; Chaudhuri et al., 2011; Abadi et al., 2016; Ge et al., 2018; Wei et al., 2020; Dong et al., 2022; Fan and Li, 2022; Fan et al., 2022; Fang et al., 2023; Li and Li, 2023a,b). Private  $k$ -means clustering has also been widely studied, e.g., Feldman et al. (2009); Nock et al. (2016); Feldman et al. (2017), mostly in the continuous Euclidean space. Balcan et al. (2017) considered identifying a good candidate set (in a private manner) of centers before applying private local search, which yields  $O(\log^3 n)$  multiplicative error and  $O((k^2 + d)\log^5 n)$  additive error. Later on, the private Euclidean  $k$ -means error axis further improved by Stemmer and Kaplan (2018), with more advanced candidate set selection. Huang and Liu (2018) gave an optimal algorithm in terms of minimizing Wasserstein distance under some data separability condition.
For private  $k$ -median clustering, Feldman et al. (2009); Ghazi et al. (2020) considered the problem in high dimensional Euclidean space. However, it is rather difficult to extend their analysis to more general metrics in discrete spaces (e.g., on graphs). The strategy of Balcan et al. (2017) to form a candidate center set could as well be adopted to  $k$ -median, which leads to  $O(\log^{3/2} n)$  multiplicative error and  $O((k^2 + d) \log^3 n)$  additive error in the Euclidean space where  $n$  is the sample size. In discrete space, Gupta et al. (2010) proposed a private method for the classical local search heuristic, which applies to both  $k$ -medians and  $k$ -means. To cast privacy on each swapping step, the authors applied the exponential mechanism of McSherry and Talwar (2007). Their method produced an  $\epsilon$ -differentially private solution with cost  $6OPT + O(\triangle k^2 \log^2 n / \epsilon)$ , where  $\triangle$  is the diameter of the point set. In this work, we will show that our proposed HST initialization can improve the DP local search for  $k$ -median of Gupta et al. (2010) in terms of both approximation error and efficiency. Stemmer and Kaplan (2018); Jones et al. (2021) proposed  $(\epsilon, \delta)$ -differentially private solution also with constant multiplicative error but smaller additive error.
The main contributions of this work include the following:
- We introduce the Hierarchically Well-Separated Tree (HST) as an initialization tool for the  $k$ -median clustering problem. We design an efficient sampling strategy to select the initial center set from the tree, with an approximation factor  $O(\log \min \{k,\triangle \})$  in the non-private setting, which is  $O(\log \min \{k,d\})$  when  $\log \triangle = O(\log d)$ . This improves the  $O(\log k)$  error of  $k$ -median++. Moreover, the complexity of our HST based method can be smaller than that of  $k$ -median++. When the number of clusters  $k$  is not too small ( $k\geq \log n$ ), which is a common scenario in practical applications.
- We propose a differentially private version of HST initialization under the setting of Gupta et al. (2010) in discrete metric space. The so-called DP-HST algorithm finds initial centers with  $O(\log n)$  multiplicative error and  $O(\epsilon^{-1}\triangle k^2\log^2 n)$  additive error. Moreover, running DP-local search starting from this initialization gives  $O(1)$  multiplicative error and  $O(\epsilon^{-1}\triangle k^2 (\log \log n)\log n)$  additive error, which improves previous results towards the well-known lower bound  $O(\epsilon^{-1}\triangle k\log (n / k))$  on the additive error of DP  $k$ -median (Gupta et al., 2010) within a small  $O(k\log \log n)$  factor. This is the first clustering initialization method with  $\epsilon$ -differential privacy guarantee and improved error rate in general metric space.
- We conduct experiments on simulated and real-world datasets to demonstrate the effectiveness of our methods. In both non-private and private settings, our proposed HST-based approach achieves smaller cost at initialization than  $k$ -median++, which may also lead to improvements in the final clustering quality.
2 Background and Setup
The definition of differential privacy (DP) is as follows.
Definition 2.1 (Differential Privacy (DP) (Dwork, 2006)). If for any two adjacent datasets  $D$  and  $D'$  with symmetric difference of size one and any  $O \subset \operatorname{Range}(\mathbb{A})$ , an algorithm  $\mathbb{A}$  with map  $f$  satisfies
$$
P r [ \mathbb {A} (D) \in O ] \leq e ^ {\epsilon} P r [ \mathbb {A} (D ^ {\prime}) \in O ],
$$
then algorithm  $\mathbb{A}$  is said to be  $\epsilon$ -differentially private ( $\epsilon$ -DP).
Intuitively, DP requires that after removing any data point from  $D$  (e.g., a node in a graph), the output of  $D'$  should not be too different from that of the original dataset  $D$ . The Laplace mechanism adds Laplace  $(\eta(f)/\epsilon)$  noise to the output where  $\eta(f) = \sup_{|D - D'| = 1} |f(D) - f(D')|$  is the sensitivity of  $f$ , which is known to achieve  $\epsilon$ -DP. The exponential mechanism is also a tool for many DP algorithms with discrete outputs. Let  $O$  be the output domain. The utility function  $q: D \times O \to \mathbb{R}$  is what we aim to maximize. The exponential mechanism outputs an element  $o \in O$  with probability  $P[\mathbb{A}(D) = o] \propto \exp(\frac{\epsilon q(D, o)}{2\eta(q)})$ . Both mechanisms will be used in our paper.
2.1  $k$ -Median Clustering and Local Search
In this paper, we follow the classic problem setting in the metric clustering literature, e.g. Arya et al. (2004); Gupta et al. (2010). Specifically, the definitions of metric  $k$ -median clustering problem (DP and non-DP) are stated as follow.
Definition 2.2 ( $k$ -median). Given a universe point set  $U$  and a metric  $\rho: U \times U \to \mathbb{R}$ , the goal of  $k$ -median to pick  $F \subseteq U$  with  $|F| = k$  to minimize
$$
k \text {m e d i a n :} \quad \operatorname {c o s t} _ {k} (F, U) = \sum_ {v \in U} \min  _ {f \in F} \rho (v, f). \tag {1}
$$
Let  $D \subseteq U$  be a set of "demand points". The goal of DP  $k$ -median is to minimize
$$
D P k \text {m e d i a n}: \quad \operatorname {c o s t} _ {k} (F, D) = \sum_ {v \in D} \min  _ {f \in F} \rho (v, f), \tag {2}
$$
and the output  $F$  is required to be  $\epsilon$ -differentially private with respect to  $D$ . We may drop "F" and use "cost $_k(U)$ " or "cost $_k(D)$ " if there is no risk of ambiguity.
Note that in Definition 2.2, our aim is to protect the privacy of a subset  $D \subset U$ . To better understand the motivation and application scenario, we provide a real-world example as below.
Example 2.3. Consider  $U$  to be the universe of all users in a social network (e.g., Facebook, LinkedIn, etc.). Each user (account) has some public information (e.g., name, gender, interests, etc.), but also has some private personal data that can only be seen by the data server. Let  $D$  be a set of users grouped by some feature that might be set as private. Suppose a third party plans to collaborate with the most influential users in  $D$  for e.g., commercial purposes, thus requesting the cluster centers of  $D$ . In this case, we need a differentially private algorithm to safely release the centers, while protecting the individuals in  $D$  from being identified (since the membership of  $D$  is private).
Algorithm 1: Local search for  $k$ -median clustering (Arya et al., 2004)
Input: Data points  $U$  , parameter  $k$  , constant  $\alpha$    
Initialization: Randomly select  $k$  points from  $U$  as initial center set  $F$    
while  $\exists x\in F,y\in U$  s.t. cost  $(F - \{x\} +\{y\})\leq (1 - \alpha /k)$  cost  $(F)$  do Select  $(x,y)\in F_i\times (D\setminus F_i)$  with arg  $\min_{x,y}\{\mathrm{cost}(F - \{x\} +\{y\})\}$  Swap operation:  $F\gets F - \{x\} +\{y\}$    
Output: Center set  $F$
The (non-private) local search procedure for  $k$ -median proposed by Arya et al. (2004) is summarized in Algorithm 1. First, we randomly pick  $k$  points in  $U$  as the initial centers. In each iteration, we search over all  $x \in F$  and  $y \in U$ , and do the swap  $F \gets F - \{x\} + \{y\}$  such that the new centers improve the cost the most, and if the improvement is more than  $(1 - \alpha / k)$  for some  $\alpha > 0^2$ . We repeat the procedure until no such swap exists. Arya et al. (2004) showed that the output center set  $F$  achieves 5 approximation error to the optimal solution, i.e.,  $cost(F) \leq 5OPT$ .
2.2  $k$ -median++ Initialization
Although local search is able to find a solution with constant error, it takes  $O(n^{2})$  per iteration (Re-sende and Werneck, 2007) in expected  $O(k\log n)$  steps (which gives total complexity  $O(kn^{2}\log n)$ ) when started from a random center set, which would be slow for large datasets. Indeed, we do not need such complicated/meticulous algorithm to reduce the cost at the beginning, i.e., when the cost is large. To accelerate the process, efficient initialization methods find a "roughly" good center set as the starting point for local search. In the paper, we compare our new initialization scheme mainly with a popular (and perhaps the most well-known) initialization method, the  $k$ -median++ (Arthur and Vassilitskii, 2007)<sup>3</sup> as presented in Algorithm 2. The output centers  $C$  by  $k$ -median++ achieve  $O(\log k)$  approximation error with time complexity  $O(nk)$ . Starting from the initialization, we only need to run  $O(k\log \log k)$  steps of the computationally heavy local search to reach a constant error solution. Thus, initialization may greatly improve the clustering efficiency.
Algorithm 2:  $k$ -median++ initialization (Arthur and Vassilvitskii, 2007)
Input: Data points  $U$  , number of centers  $k$  Randomly pick a point  $c_{1}\in U$  and set  $F = \{c_1\}$  for  $i = 2,\dots,k$  do Select  $c_{i} = u\in U$  with probability  $\frac{\rho(u,F)}{\sum_{u^{\prime}\in U}\rho(u^{\prime},F)}$ $F = F\cup \{c_i\}$  Output:  $k$  -median++ initial center set  $F$
3 Initialization via Hierarchical Well-Separated Tree (HST)
In this section, we propose our new initialization scheme for  $k$ -median clustering, and provide our analysis in the non-private case solving (1). The idea is based on the metric embedding theory. We will start with an introduction to the main tool used in our approach.
3.1 Hierarchically Well-Separated Tree (HST)
In this paper, for an  $L$ -level tree, we will count levels in a descending order down the tree. We use  $h_v$  to denote the level of  $v$ , and let  $n_i$  be the number of nodes at level  $i$ . The Hierarchically Well-Separated Tree (HST) is based on the padded decompositions of a general metric space in a hierarchical manner (Fakcharoenphol et al., 2004). Let  $(U,\rho)$  be a metric space with  $|U| = n$ , and we will refer to this metric space without specific clarification. A  $\beta$ -padded decomposition of  $U$
$^2$ Arya et al. (2004) chooses the first swap that improves the cost by  $(1 - \alpha /k)$ , instead of picking the smallest cost among all such swaps as in Algorithm 1. Both methods are valid and yield the same approximation error.
<sup>3</sup>In the original paper of Arthur and Vassilvitskii (2007), the  $k$ -median++ algorithm was called  $D^1$ -sampling.
Figure 1: An example of a 3-level padded decomposition and the corresponding 2-HST. Left: The thickness of the ball represents the level. The colors correspond to different levels in the HST in the right panel. “ $\triangle$ ”s are the center nodes of partitions (balls), and “ $\times$ ”s are the non-center data points. Right: The 2-HST generated from the padded decomposition. Bold indices represent the centers.
is a probabilistic partition of  $U$  such that the diameter of each cluster  $U_{i} \in U$  is at most  $\beta$ , i.e.,  $\rho(u, v) \leq \beta, \forall u, v \in U_{i}, i = 1, \dots, k$ . The formal definition of HST is given as below.
Definition 3.1. Assume  $\min_{u,v\in U}\rho (u,v) = 1$  and denote the diameter  $\triangle = \max_{u,v\in U}\rho (u,v)$ . An  $\alpha$ -Hierarchically Well-Separated Tree ( $\alpha$ -HST) with depth  $L$  is an edge-weighted rooted tree  $T$ , such that an edge between any pair of two nodes of level  $i - 1$  and level  $i$  has length at most  $\triangle/\alpha^{L - i}$ .
Our analysis will consider  $\alpha = 2$ -HST for conciseness, since  $\alpha$  only affects the constants in our theoretical analysis. Figure 1 is an example 2-HST (right panel) with  $L = 3$  levels, along with its underlying padded decompositions (left panel). Using Algorithm 3, a 2-HST can be built as follows: we first find a padded decomposition  $P_{L} = \{P_{L,1},\dots,P_{L,n_{L}}\}$  of  $U$  with parameter  $\beta = \triangle /2$ . The center of each partition in  $P_{L,j}$  serves as a root node in level  $L$ . Then, we re-do a padded decomposition for each partition  $P_{L,j}$ , to find sub-partitions with diameter  $\beta = \triangle /4$ , and set the corresponding centers as the nodes in level  $L - 1$ , and so on. Each partition at level  $i$  is obtained with  $\beta = \triangle /2^{L - i}$ . This process proceeds until a node has a single point (leaf), or a pre-specified tree depth is reached. It is worth mentioning that, Blelloch et al. (2017) proposed an efficient HST construction in  $O(m\log n)$  time, where  $n$  and  $m$  are the number of nodes and edges in a graph, respectively. Therefore, the construction of HST can be very efficient in practice.
The first step of our method is to embed the data points into an HST (see Algorithm 4). Next, we will describe our new strategy to search for the initial centers on the tree (w.r.t. the tree metric). Before moving on, it is worth mentioning that, there are polynomial time algorithms for computing an exact  $k$ -median solution in the tree metric (Tamir (1996); Shah (2003)). However, the dynamic programming algorithms have high complexity (e.g.,  $O(kn^2)$ ), making them unsuitable for the purpose of fast initialization. Moreover, it is unknown how to apply them effectively to the private case. The three key merits of our new algorithm are: (1) It is more efficient than  $k$ -median++ when  $k$  is not too small, which is a very common scenario in practice; (2) It achieves  $O(1)$  approximation error in the tree metric; (3) It can be easily extended to incorporating differential privacy (DP).
Algorithm 3: Build 2-HST(U,L)
Input: Data points  $U$  with diameter  $\triangle, L$   
Randomly pick a point in  $U$  as the root node of  $T$   
Let  $r = \triangle / 2$   
Apply a permutation  $\pi$  on  $U$  // so points will be chosen in a random sequence for each  $v \in U$  do  
Set  $C_v = [v]$   
for each  $u \in U$  do  
| Add  $u \in U$  to  $C_v$  if  $d(v, u) \leq r$  and  $u \notin \bigcup_{v' \neq v} C_{v'}$   
Set the non-empty clusters  $C_v$  as the children nodes of  $T$   
for each non-empty cluster  $C_v$  do  
| Run 2-HST $(C_v, L - 1)$  to extend the tree  $T$ ; stop until  $L$  levels or reaching a leaf node  
Output: 2-HST  $T$
3.2 HST Initialization Algorithm
Let  $L = \log \Delta$  and suppose  $T$  is a level- $L$  2-HST in  $(U, \rho)$ , where we assume  $L$  is an integer. For a node  $v$  at level  $i$ , we use  $T(v)$  to denote the subtree rooted at  $v$ . Let  $N_v = |T(v)|$  be the number of data points in  $T(v)$ . The search strategy for the initial centers, NDP-HST initialization ("NDP" stands for "Non-Differentially Private"), is presented in Algorithm 4 with two phases.
Subtree search. The first step is to identify the subtrees that contain the  $k$  centers. To begin with,  $k$  initial centers  $C_1$  are picked from  $T$  who have the largest score  $(v) = N(v)\cdot 2^{h_v}$ . This is intuitive, since to get a good clustering, we typically want the ball surrounding each center to include more data points. Next, we do a screening over  $C_1$ : if there is any ancestor-descendant pair of nodes, we remove the ancestor from  $C_1$ . If the current size of  $C_1$  is smaller than  $k$ , we repeat the process until  $k$  centers are chosen (we do not re-select nodes in  $C_1$  and their ancestors). This way,  $C_1$  contains  $k$  root nodes of  $k$  disjoint subtrees.
Algorithm 4: NDP-HST initialization
Input:  $U,\triangle ,k$    
Initialization:  $L = \log \triangle ,C_0 = \emptyset ,C_1 = \emptyset$    
Call Algorithm 3 to build a level-  $L2\mathrm{-HST}T$  using  $U$    
for each node  $v$  in  $T$  do  $N_{v}\gets |U\cap T(v)|$  score  $(v)\leftarrow N_v\cdot 2^{h_v}$    
while  $|C_1| <   k$  do Add top  $(k - |C_1|)$  nodes with highest score to  $C_1$  for each  $v\in C_1$  do  $C_1 = C_1\backslash \{v\}$  if  $\exists v^{\prime}\in C_{1}$  such that  $v^{\prime}$  is a descendant of  $v$ $C_0 =$  FIND-LEAF(T,C1)   
Output: Initial center set  $C_0\subseteq U$
Leaf search. After we find  $C_1$  the set of  $k$  subtrees, the next step is to find the center in each subtree using Algorithm 5 ("FIND-LEAF"). We employ a greedy search strategy, by finding the child node with largest score level by level, until a leaf is found. This approach is intuitive since the diameter of the partition ball exponentially decays with the level. Therefore, we are in a sense focusing more and more on the region with higher density (i.e., with more data points).
The complexity of our search algorithm is given as follows. All proofs are placed in Appendix B.
Proposition 3.2 (Complexity). Algorithm 4 takes  $O(dn\log n)$  time in the Euclidean space.
Remark 3.3 (Comparison with  $k$ -median++.) The complexity of  $k$ -median++ is  $O(dnk)$  in the Euclidean space (Arthur and Vassilvitskii, 2007). Our algorithm would be faster when  $k > \log n$ , which is a common scenario. Similar comparison also holds for general metrics.
3.3 Approximation Error of HST Initialization
We provide the error analysis of our algorithm. Firstly, we show that the initial center set produced by NDP-HST is already a good approximation to the optimal  $k$ -median solution. Let  $\rho^T(x,y) = d_T(x,y)$  denote the "2-HST metric" between  $x$  and  $y$  in the 2-HST  $T$ , where  $d_T(x,y)$  is the tree distance between nodes  $x$  and  $y$  in  $T$ . By Definition 3.1 and since  $\triangle = 2^L$ , in the analysis we assume