Clustering with UMAP: Why and How Connectivity Matters

UMAP first constructs a high dimensional graph representation of the data (graph construction phase), which is used to represent the “fuzzy simplicial complex”. Next it optimizes low-dimensional vectors to be as similar as possible to the graph representation in higher dimensions, by minimizing the cross-entropy of two fuzzy sets with the same underlying elements (datapoints).

The default choice of $\mathcal{M}_{i}$ are the nearest neighbors of $x_{i}$. However, previous research has shown that a weighted k-NN graph captures noisy relationships, and may not be an accurate representation of the underlying local structure for a high dimensional dataset due to the “curse of dimensionality” (Liu and Zhang 2012; Radovanović, Nanopoulos, and Ivanovic 2010). Instead we utilize a weighted mutual $k$-nearest neighbor graph (weighted mutual k-NN graph) to represent the underlying structure of our data. A mutual k-NN graph is defined as a graph that only has an undirected edge between two vertices $x_{i}$ and $x_{j}$ if $x_{i}$ is in $k$-nearest neighbors of $x_{j}$ and $x_{j}$ is in the $k$-nearest neighbors of $x_{i}$. Therefore our local neighborhood is $\mathcal{M}_{i}=\{x_{j}|x_{j}\in knn(x_{i}),x_{i}\in knn(x_{j})\}$.

A mutual k-NN graph can hence be interpreted as a subgraph of the original k-NN graph. Mutual k-NN graphs have been shown to contain many desirable properties in comparison to its k-NN counterpart when combating the “curse of dimensionality”. One such property is that edges in mutual k-NN graphs have been shown to be a stronger indicator of similarity than the unidirectional nearest neighborhood relationship, alleviating the distance concentration effect (Jégou et al. 2010). Mutual k-NN graphs have also been shown to be less prone to the hub effect than the regular k-NN graph, since each vertex in a mutual k-NN graph is guaranteed to be at most k-degrees (Ozaki et al. 2011). They have also shown strong performance when used in both unsupervised clustering (Maier, Hein, and Luxburg 2009; Sardana and Bhatnagar 2014; Abbas and Shoukry 2012) and semi-supervised graph based classifications algorithms (Ozaki et al. 2011; de Sousa, Rezende, and Batista 2013). In Section 3, we describe how we utilize the mutual k-NN graph to create our high dimensional graph representation.

Since a mutual k-NN graph only retains a subset of the edges from the original k-NN graph(Section 2.2), the resulting mutual k-NN graph often contains disconnected components and potential isolated vertices. Isolated vertices violate one of UMAP’s primary assumptions that the underlying manifold is locally connected, and disconnected components negatively impact the spectral initialization (McInnes, Healy, and Melville 2018). Spectral initialization is crucial for preserving the global structure of the high dimensional data in the low dimensional vectors (Kobak and Linderman 2021).

Naively, we could increase $k$, the number of nearest neighbors used for calculating the mutual k-NN graph, to obtain a more connected graphical representation. However increasing $k$ does not guarantee that the resulting graph has no isolated vertices or fewer disconnected components. We therefore consider the following methods for increasing the Connectivity of the mutual k-NN graph, presented in order of increasing connectivity.

1. 1.

   NN: Add an undirected edge between each isolated vertex and its original nearest neighbor (de Sousa, Rezende, and Batista 2013).

2. 2.

   MST-min: To achieve a connected graph, add the minimum number of edges from a maximum spanning tree to the mutual-kNN graph that has been weighted with similarity-based metrics(Ozaki et al. 2011). We adapt this by calculating the minimum spanning tree for distances (Algorithm 1).

3. 3.

   MST-all: Adding all the edges of the MST.

We call this graph $G^{\prime}$, the locally connected mutual k-NN graph, and run experiments with $G^{\prime}$ in Section 4.

Next we wish to obtain a “fuzzy simplicial complex” representation as described in Section 2.1, from the connected mutual k-NN graph, $G^{\prime}$ (Section 2.2). To achieve this, we first need to obtain the new local neighborhood, $\mathcal{M}_{i}$ for each $x_{i}\in X$ using $G^{\prime}$. The most straightforward way is to use adjacent vertices in $G^{\prime}$. However, this excludes all points that are more than 1 hop away, even if they are relatively close.

Instead, we compute $G$ from $G^{\prime}$ by refining the local neighborhood $\mathcal{M}_{i}$ using the Shortest Path Distance between any two nodes using Dijkstra’s Algorithm (See Fig. 2 for example) to other nodes on $G^{\prime}$ (Algorithm 2). This shortest path distance can be considered a new distance metric as it directly aligns with UMAP’s definition of an extended-pseudo-metric space. That is, we replace $d(x_{i},x_{j}’)$ with $d_{\mathrm{path}}(x_{i},x_{j})$ in the graph construction phase (Section 2.1).

The complexity of constructing the mutual k-NN graph from the original k-NN graph is $O(kn)$. The additional complexity of connecting disconnected components varies depending on the method chosen. Adding edges between isolated vertices to their nearest neighbors (NN) is $O(n)$ while the MST variants have a known time complexity of $O(E\log(V))$ = $O(kn\log(n))$. Since we construct the MST from the original k-NN graph, we know that $E=nk$ and $V=n$. With the MST, we add either all the edges (MST-all), which has an upper bound of $O(n)$ since there are $n-1$ edges in the MST, or the minimum number of edges (MST-min) until we have one component, which has an upper bound of $O(n\log(n))$ since we need to sort the edges first.

Finally, using Path Neighbors on mutual k-NN graph has an additional cost since we have to perform a graph search. We use a variant of Djikstra’s algorithm which has a known time complexity of $O(V+E\log(V))$. Running the search for all $V=n$ nodes and $E\leq nk$, we get a final complexity of $O(n(n+kn\log(n)))$. It is important to note that the worst case for Djikstra’s algorithm assumes one may need to explore the whole graph as it is typically to find the shortest distance between two points. However, since we are using Djikstra’s to find the new closest points, we in practice are more likely to terminate much earlier. This search can be further optimized by parallelizing the individual searches needed to find new nearest neighbors for each point.

We used two standard image classification datasets, MNIST (LeCun and Cortes 2010) and Fashion MNIST (FMNIST; Xiao, Rasul, and Vollgraf (2017)), and two standard text classification datasets, 20 newsgroup (20NG),¹¹1http://qwone.com/ jason/20Newsgroups/ and AG’s News Topic (AG;Zhang, Zhao, and LeCun (2015)) for evaluation. For the image classification datasets, no preprocessing was applied and each image was flattened to create a 1D vector. For the text classification datasets, we lowercase tokens, remove stopwords, punctuation and digits, and exclude words that appear in less than 5 documents. After preprocessing, we converted each dataset into both Count vectors and TF-IDF vectors. We opted to use the Count based representation of text as an illustration of datasets that have very high-dimensional and sparse representations.²²2TF-IDF is a weighting on the sparse representation. While our method is applicable to word or document vectors, these are heavily dependent on the text encoder and it is more difficult to interpret findings as we do not know what the dense high-dimensional vector space represents in the first place. ³³3Code is available at https://github.com/adalmia96/umap-mnn.

Quantative Evaluation: We evaluate the methods by comparing the clustering performance of the resulting low dimensional vectors generated. We performed KMeans clustering, assigning the number of clusters to be the ground truth number of topics for each dataset. We use the Normalised Mutual Information (NMI), a widely used clustering metric that ranges from 0 to 1 (1 being a perfect score), to evaluate the clusters.⁴⁴4NMI was primarily used to evaluate clustering quality for UMAP and subsequent works(McInnes, Healy, and Melville 2018) As UMAP is a stochastic algorithm, the NMI scores presented in Table 1 are averaged across low-dimensional vectors from 5 random seeds.

Qualitative Evaluation: In addition, we consider several desirable properties for dimensionality reduction in qualitative visualisations, local structure and global structure, and separation of classes. A dimensionality reduction method which preserves the local structure means that intra-class distances are preserved; i.e. the datapoints belonging to the same class are close to each other. A method which preserves the global structure means that inter-class distances are preserved; i.e., clusters of similar classes appear closer together and clusters of dissimilar classes appear farther from each other. Finally, the classes should be well separated (not clumped together).

Our starting point for all connectivity methods is the mutual k-NN, which we compare against UMAP which uses the default k-NN. As described in Section 3.1, we tested NN, which connects the nearest neighbor, MST-min which uses the minimum edges from the minimum spanning tree (MST), and MST-all which uses all the edges from the minimum spanning tree. This gives us a connected graph $G^{\prime}$. Next, we consider two methods for obtaining the local neighborhood for each point, $\mathcal{M}_{i}$ which is used for the final graph $G$ (Section 3.2). For Adjacent Neighbors, $G^{\prime}=G$ as there is no change to the local neighborhood. For Path Neighbors, we additionally consider nodes from $G^{\prime}$ which are more than one hop away until we get $k$ neighbors, to get the ‘new’ local neighborhood $\mathcal{M}_{i}$ which is used for $G$.

Using a mutual k-NN representation results in improved separation between similiar classes. In general we observe that for most of the mutual k-NN graph based vectors (Fig. 3), there is a better separation between similar classes than the original UMAP vectors regardless of Connectivity (NN, MST variants). We observed the desired separation between similar classes such as with the 4, 7, 9 in MNIST and the footwear classes for FMNIST. mutual k-NN graphs have previously been shown as a useful method for removing edges between points that are only loosely similar, directly reducing distance concentration and hub effects.

We consider three methods for connecting the mutal k-NN. In terms of connectivity, NN $<$ MST-min $<$ MST-all.

NN performs worse than MST variants, with vectors that are randomly scattered in 2d space. For both MNIST and FMNIST, we see that NN, which connects isolated vertices to their nearest neighbor, had multiple small clusters of points scattered throughout the vector space. Given that KMeans is sensitive to outliers, these randomly projected points negatively affect clustering performance as seen in Table 1. Since a mutual k-NN graph only retains a subset of the edges from the original k-NN graph, it can result in a very sparse representation. When constructing the mutual k-NN graph, we observed that $\approx$ 1000($3\%$) points were separated into small components for MNIST and $\approx$ 7000($10\%$) points for FMNIST. Our results show stronger notions of connectivity than NN are required.

We would expect that having higher connectivity that reduces random scattering of points would be better for clustering. However, we observe that too much connectivity from using all the edges from the MST (MST-all) with Path Neighbors can hurt performance on FMNIST (Section 5.3).

We consider two methods, Adjacent Neighbors and Path Neighbors, for finding the local neighborhood $\mathcal{M}_{i}$ of each point, after obtaining a connected mutual k-NN graph $G^{\prime}$.

Path Neighbors achieves the best clustering performance together with MST-min. We observe a similar clustering performance when we use the minimum number of edges from the MST (MST-min), vs all the edges from the MST (MST-all). In the case of FMNIST, using MST-all results in a worse clustering performance of 0.698 vs 0.649 for 2 dims, and 0.698 vs 0.679 for 64 dims. Although using Adjacent vs Path Neighbors resulted in similar clustering performance for image based dataset, Path Neighbors produced better results for text based datasets. From Fig. 3, both MST-min and MST-all produced better results when using Path vs Adjacent Neighbors for text based datasets ($p<0.05$).

Adjacent Neighbors produce a poorer local structure than Path Neighbors. Visually, the vector generated using the Adjacent Neighbors and MST-min result in disperse dense clusters of points e.g, the footwear classes in FMNIST and the recreation topics in 20 NG. However when we apply Path Neighbors, the groups of points belonging to the same class are less dispersed (Fig. 3). This indicates that Adjacent Neighbors have a poorer local structure than Path Neighbors; i.e. same class datapoints are far from each other.

To investigate further why Adjacent Neighbors produce worse results for text datasets, we compute the variance along the dimensions of the 2D vectors for each class from the 20NG Count vectors and plot them in Fig. 4. We find that Adjacent Neighbors has greater variance for each class than the vectors for both Path Neighbors and original UMAP, indicating greater dispersion of points within-class. Having high variance is bad for clustering and indicates a poor lower dimensional representation, as there is no distinct range of values associated with the class label.

Path Neighbors increases the connectivity for the final $G$, and therefore can rely on a more refined notion of $G^{\prime}$ with MST-min instead of MST-all. We can interpret Path Neighbors as a method which strictly increases the general connectivity of $G$. Consider the local neighborhood $\mathcal{M}_{i}$ for Adjacent Neighbors, which is not guaranteed to have $k$ connected neighbors for $\mathcal{M}_{i}$. Points with smaller $\mathcal{M}_{i}$ will be close to primarily few adjacent neighbors and repelled further away from the other points. This creates small groups of points that belong to the same class being spread across the vector space. On the other hand, for original UMAP and Path Neighbor vectors, $|\mathcal{M}_{i}|=k$,and local groups of points are more likely to be connected to other groups within the same class. This increase in connectivity explains why visually Path Neighbors methods generate vectors which are less dispersed (within class) than the Adjacent Neighbors method, while still preserving the underlying structure of the data. Across Table 1, we see that using Path Neighbors performs consistently well with MST-min across both image and text datasets, and allows more flexibility in building a locally and globally connected fuzzy simplicial set.

We also found that clustering performance was consistent across different dimensions for image datasets but was better at higher dimensions for text dataset with TF-IDF representations. Table 1 show results for 2 and 64 dimensions. While it is not surprising that having 64 dimensions allows the high-dimensional text datasets ($>$26000 features for AG and $>$34000 features for 20NG) to preserve more information, it is interesting that MST variants with Path Neighbors does not produce worse results at 2 vs 64 dimensions.