Unsupervised Embedding of Hierarchical Structure in Euclidean Space

Hierarchical clustering aims to find an iterative grouping of increasingly large subsets of data, terminating when one cluster contains all of the data. This results in a tree structure, where each level determines a partition of the data. Organizing the data in such a way has many applications, including automated taxonomy construction and phylogeny reconstruction balaban2019treecluster; friedman2001elements; manning2008introduction; shang2020nettaxo; zhang2018taxogen. Motivated by these applications, we address the simultaneous goals of recovering an underlying clustering and building a hierarchical tree on top of the clusters. We imagine that we have access to many samples from each individual cluster (e.g., images of plants and animals), while we lack labels or any direct information about the hierarchical relationships. In this scenario, our objective is to correctly cluster the samples in each group and also build a dendrogram that matches the natural supergroups.

Despite significant effort, hierarchical clustering remains a challenging task, theoretically and empirically. Many objectives are NP-Hard (charikar2017approximate; charikar2018hierarchical; dasgupta2002performance; dasgupta2016cost; heller2005bayesian; MoseleyWang), and many theoretical algorithms may not be practical because they require solving a hard subproblem (e.g., sparsest cut) at each step (ackermann2014analysis; ahmadian2019bisect; charikar2017approximate; charikar2018hierarchical; dasgupta2016cost; lin2010general). The most popular hierarchical clustering algorithms are linkage-based methods king1967step; lance1967general; sneath1973numerical; ward1963hierarchical. These algorithms first partition the dataset into singleton clusters. Then, at each step, the pair of most similar clusters is merged. There are many variants depending on how the cluster similarity is computed. While these heuristics are widely used, they do not always work well in practice. At present, there is only a superficial understanding of theoretical guarantees, and the central guiding principle is that linkage-based methods recover the underlying hierarchy as long as the dataset satisfies certain strong separability assumptions (abboud2019subquadratic; balcan2019learning; Cohen-Addad; emamjomeh2018adaptive; ghoshdastidar2019foundations; grosswendt2019analysis; kpotufepruning; reynolds2006clustering; roux2018comparative; sharma2019comparative; szekely2005hierarchical; tokuda2020revisiting).

We consider embedding the data in Euclidean space and then clustering with a linkage-based method. Prior work has shown that a variational autoencoder (VAE) with Gaussian mixture model (GMM) prior leads to separated clusters via the latent space embedding dilokthanakul2016deep; nalisnick2016approximate; jiang2017variational; uugur2020variational; xie2016unsupervised. We focus on one of the best models, Variational Deep Embedding (VaDE) jiang2017variational, in the context of hierarchical clustering. Surprisingly, the VaDE embedding seems to capture hierarchical structure. Clusters that contain semantically similar data end up closer in Euclidean distance, and this pairwise distance structure occurs at multiple scales. This phenomenon motivates our empirical and theoretical investigation of enhancing hierarchical clustering quality with unsupervised embeddings into Euclidean space.

Our experiments demonstrate that applying Ward’s method after embedding with VaDE produces state-of-the-art results for two hierarchical clustering objectives: (i) dendrogram purity for classification datasets heller2005bayesian and (ii) Moseley-Wang’s cost function MoseleyWang, which is a maximization variant of Dasgupta’s cost dasgupta2016cost. At a high level, both measures reward clusterings for merging more similar data points before merging different ones. For brevity, we restrict our attention to these two measures, while our qualitative analysis suggests that the VaDE embedding would also improve other objectives.

As another contribution, we propose an extension of the VaDE embedding that rescales the means of the learned clusters (attempting to increase cluster separation). On many real datasets, this rescaling improves both clustering objectives, which is consistent with the general principle that linkage-based methods perform better with more separated clusters. The baselines for our embedding evaluation involve both a standard VAE with isotropic Gaussian prior kingma2014auto; rezende2014stochastic and principle component analysis (PCA). In general, PCA provides a worse embedding than VAE/VaDE, which is expected because it cannot learn a non-linear transformation. The VaDE embedding leads to better hierarchical clustering results than VAE, indicating that the GMM prior offers more flexibility in the latent space.

Our focus on Euclidean geometry lends itself to a synthetic distribution for hierarchical data that serves as a model of the VaDE embedding. We emphasize that we implicitly represent the hierarchy with pairwise Euclidean distances, bypassing the assumption that the algorithm has access to actual similarities. Extending our model further, we demonstrate a shifted version, where using a non-linear embedding of the data improves the clustering quality. Figure 1 depicts the original and shifted synthetic data with 8 clusters forming a 3-level hierarchy. The 3D plot of the original data in Figure 1(a) has ground truth pairwise distances in Figure 1(b). In Figure 1(c), we see the effect of shifting four of the means by a cyclic rotation, e.g., $(4,2,1,0,0,0)\mapsto(0,0,0,4,2,1)$. This non-linear transformation increases the distances between pairs of clusters while preserving hierarchical structure. PCA in Figure 1(d) can only distinguish between two levels of the hierarchy, whereas VaDE in Figure 1(e) leads to more concentrated clusters and better identifies multiple levels.

To improve the theoretical understanding of Ward’s method, we prove that it exactly recovers both the underlying clustering and the planted hierarchy when the cluster means are hierarchically separated. The proof strategy involves bounding the distance between points sampled from spherical Gaussians and then extending recent results on Ward’s method for separated data grosswendt2019analysis; grosswendt_2020. Intuitively, we posit that the rescaled VaDE embedding produces vectors that resemble our planted hierarchical model, and this enables Ward’s method to recover both the clusters and the hierarchy after the embedding.

Related Work.   For flat clustering objectives (e.g., $k$-means or clustering accuracy), there is a lot of research on using embeddings to obtain better results by increasing cluster separation aljalbout2018clustering; dilokthanakul2016deep; fard2018deep; figueroa2017simple; guo2017deep; jiang2017variational; li2018discriminatively; li2018learning; min2018survey; tzoreff2018deep; xie2016unsupervised; yang2017towards. However, we are unaware of prior work on learning unsupervised embeddings to improve hierarchical clustering. Studies on hierarchical representation learning instead focus on metric learning or maximizing data likelihood, without evaluating downstream unsupervised hierarchical clustering algorithms ganea2018hyperbolic; goldberger2005hierarchical; goyal2017nonparametric; heller2005bayesian; klushyn2019learning; nalisnick2017stick; shin2019hierarchically; sonthalia2020tree; teh2008bayesian; tomczak2018vae; vasconcelos1999learning. Our aim is to embed the data into Euclidean space so that linkage-based algorithms produce high-quality clusterings.

Many researchers have recently identified the benefits of hyperbolic space for representing hierarchical structure (de2018representation; gu2018learning; linial1995geometry; mathieu2019continuous; monath2019gradient; nickel2017poincare; tifrea2018poincar). Their stated motivation is that hyperbolic space can better represent tree-like metric spaces due to its negative curvature (sarkar2011low). We offer an alternative perspective. If we commit to using a linkage-based method to recover the hierarchy, then it suffices to study Euclidean geometry. Indeed, we do not need to approximate all pairwise distances — we only need to ensure that the clustering method produces a good approximation to the true hierarchical structure. Our simple approach facilitates the use of off-the-shelf algorithms designed for Euclidean geometry, whereas other methods, such as hyperbolic embeddings, require new tools for downstream tasks.

Regarding our synthetic model, prior work has studied analogous models that sample pairwise similarities from separated Gaussians balakrishnan2011noise; Cohen-Addad; ghoshdastidar2019foundations. Due to this key difference, previous theoretical techniques for similarity-based planted hierarchies do not apply to our data model. We extend the prior results to analyze Ward’s method directly in Euclidean space (consistent with the embeddings).

Planted Hierarchies with Noise.   We introduce a simple way of constructing a mixture of Gaussians that also encodes hierarchical structure via Euclidean distance. The purpose of this distribution is twofold: (i) it approximately models the latent space of VaDE, and (ii) it allows us to test the effectiveness of VaDE to increase the separation between pairs of clusters.

We called it a Binary Tree Gaussian Mixture (BTGM), and it is a uniform mixture of spherical Gaussians with shared variance and different means (that encode the hierarchy). There are three parameters: the height $h$ and the margin $m$ and the dimension $d$. For a BTGM with height $h$ and margin $m$, we construct $2^{h}$ spherical Gaussian distributions with mean $\mu_{i}$ and unit variance. The $j^{th}$ coordinate of $\mu_{i}$ is given by

Intuitively, the $2^{h}$ means are leaves of a complete binary tree with height $h$. The crucial property is that the distance between any two means $\mu_{i},\mu_{j}$ is related to their distance in the binary tree. Precisely, $\left\lVert\mu_{i}-\mu_{j}\right\rVert_{2}\geq m\cdot 2^{\mathsf{ht}(\mathsf{LCA}(\mu_{i},\mu_{j}))}$, where $\mathsf{ht}(v)$ denotes the height of node $v$ in the tree, and $\mathsf{LCA}$ is the least common ancestor. Figure 1 shows an example of the BTGM with $m=4$, $h=3$ and $d=100$. Figure 1(a) depicts the 3D tSNE visualization of the data and Figure 1(b) shows the pairwise distance matrix (normalized to [0,1]) of each cluster as well as the ground-truth hierarchy. We also define a non-linear shifting operation to the BTGM. In experiments, we shift the nonzero entries of the first $\frac{k}{2}$ components in BTGM to an arbitrary location, e.g., for a BTGM with $k=8,h=3$, and $m=2$, the first mean $\mu_{1}$ will shift $(4,2,1,0,0,0)\mapsto(0,0,0,4,2,1)$. After the shifting, the hierarchical structure of BTGM is less clear from the pairwise distances in the $h$-dimensional Euclidean subspace.

Objectives.   Index the dataset as $\mathcal{X}=\{x_{1},x_{2},\ldots,x_{n}\}$. Let $\mathsf{leaves}(\mathsf{T}[i\vee j])\subseteq\mathcal{X}$ denote the set of leaves of the subtree in $\mathsf{T}$ rooted at $\mathsf{LCA}(x_{i},x_{j})$, the least common ancestor of $x_{i}$ and $x_{j}$.

Given pairwise similarities $w_{ij}$ between $x_{i}$ and $x_{j}$, Dasgupta’s cost dasgupta2016cost aims to minimize

Moseley-Wang’s objective MoseleyWang is a dual formulation, based on maximizing

Dendrogram purity heller2005bayesian is a measure using ground truth clusters $\mathcal{C}=\{C^{*}_{1},\ldots,C^{*}_{k}\}$. Let $\mathcal{P}^{*}$ be the set of pairs in $\mathcal{X}$ residing in the same cluster: $\mathcal{P}^{*}=\{(x^{i},x^{j})\mid\mathcal{C}(x^{i})=\mathcal{C}(x^{j})\}$. Then,

Ward’s Method.   Let $C_{1},\ldots,C_{k}$ be a $k$-clustering, and let $\mu(C)=\frac{1}{|C|}\sum_{x\in C}x$ denote the cluster mean. Ward’s method ward1963hierarchical; grosswendt2019analysis merges two clusters to minimize the increase in sum of squared distances: $\min_{i,j}\sum_{x\in C_{i}\cup C_{j}}\|x-\mu(C_{i}\cup C_{j})\|_{2}^{2}-\sum_{x\in C_{i}}\|x-\mu(C_{i})\|_{2}^{2}-\sum_{x\in C_{j}}\|x-\mu(C_{j})\|_{2}^{2}$.

VaDE jiang2017variational considers the following generative model $p(x,z,c)=p(x\mid z)p(z\mid c)p(c)$. We use $\mathsf{Cat}(.)$ to denote a categorical distribution, $\mathcal{N}(.)$ is a normal distribution, and $k$ is the predefined number of clusters. In the following, $\mu_{c}$ and $\sigma^{2}_{c}$ are parameters for the latent distribution, and $\hat{\mu}_{z}$ and $\hat{\sigma}^{2}_{z}$ are outputs of $f(z;\theta)$, parametrized by a neural network. The generative process for the observed data is

The VaDE maximizes the likelihood of the data, which is bounded by the evidence lower bound,

where $q(z,c\mid x)$ is a variational posterior to approximate the true posterior $p(z,c\mid x)$. We use a neural network $g(x;\phi)$ to parametrize the model $q(z\mid x)$. We refer to the original paper for the derivation of Eq. (5) and for the connection of $L_{\text{ELBO}}$ to the gradient formulation of the objective jiang2017variational.

Improving separation by rescaling means.   The VaDE landscape is a mixture of multiple Gaussian distributions. As we will see in Section 4, the separation of Gaussian distributions is crucial for Ward’s method to recover the underlying clusters. Since Ward’s method is a local heuristic, data points that lie in the middle area of two Gaussians are likely to form a larger cluster before merging into one of the true clusters. This "problematic" cluster will further degrade the purity of the dendrogram. However, the objective that VaDE optimizes does not encourage a larger separation. As a result, we apply a transformation to the VaDE latent space to enlarge the separation between each Gaussians. Let $x^{\prime}_{i}$ be the embedded value of the data point $x_{i}$, and define

to be the cluster label assigned by the learned GMM model. Then, we compute the transformation

where $s$ is a positive rescaling factor. The transformation has the following properties: (i) for points assigned to the same Gaussian by the VaDE method, the transformation preserves their pairwise distances; (ii) if $x_{i}$’s assigned to the same $\mu$ by $\mathcal{C}$ are i.i.d. random variables with expectation $\mu$ (which coincides with the intuition behind VaDE), then the transformation preserves the ratio between the distances of the expected point locations associated with different mean values.

Appendix C illustrates how to choose $s$ in detail. To improve both DP and MW, the rule of thumb is to choose a scaling factor $s$ large enough. We set $s=3$ for all datasets. In summary, we propose the following hierarchical clustering procedure:

1. 1.

   Train VaDE on the unlabeled dataset $\{x_{1},\ldots,x_{n}\}$.

2. 2.

   Embed $x_{i}$ using the value $x_{i}^{\prime}=\hat{\mu}_{\hat{z}}$ where $(\hat{\mu}_{\hat{z}},\log\hat{\sigma}_{\hat{z}}^{2})=f(\hat{z};\ \theta)$ and $\hat{z}=g(x_{i}\ ;\ \phi)$.

3. 3.

   Choose a rescaling factor $s$ and apply the rescaling transformation on each $x^{\prime}_{i}$ to get $x^{\prime\prime}_{i}$.

4. 4.

   Run Ward’s method on $\{x^{\prime\prime}_{i}\}_{i=1}^{n}$ to produce a hierarchical clustering.

We provide theoretical guarantees for Ward’s method on separated data, providing justification for using VaDE. Here, we focus on the mixture model of spherical Gaussians. We are given a sample of size $n$, say $X_{1},\ldots,X_{n}$, where each sample point $X_{i}\in\mathbb{R}^{d}$ is independently drawn from one of $k$ Gaussian components according to mixing weights $w_{1},\ldots,w_{k}$, with the $j$-th component having mean $\mu_{j}\in\mathbb{R}^{d}$ and covariance matrix $\Sigma_{j}\in\mathbb{R}^{d\times d}$. Specifically, for any component index $j$, we assume that the covariance $\Sigma_{j}$ satisfies $\Sigma_{j}=\sigma_{j}^{2}\cdot I_{d}$, where $\sigma_{j}>0$ and $I_{d}\in\mathbb{R}^{d\times d}$ is the identity matrix. Under this assumption, the model is fully determined by the set of parameters $\{(w_{i},\mu_{i},\sigma_{i}):\forall i\in[k]\}$. The goal is to cluster the sample points by their associated components.

The theorem below presents sufficient conditions for Ward’s method to recover the underlying $k$-component clustering when points are drawn from a $k$-mixture of spherical Gaussians. To recover a $k$-clustering, we terminate when there are $k$ clusters. Proofs appear in Appendix A.

As a recap, Ward’s method starts with singleton clusters, and repeatedly selects a pair of clusters and merges them into a new cluster, where the selection is based on the optimal value of the error sum of squares. In particular, recovering the underlying $k$-component clustering means that every cluster contains only points from the same Gaussian component when the algorithm produces $k$-clusters.

Intuition and Optimality.   On a high level, the mean separation conditions in Theorem 1 guarantee that with high probability, the radiuses of the clusters will be larger than the inter-cluster distances, and the sample lower bound ensures that we observe at least constant many sample points from every Gaussian, with high probability. In Appendix B, we illustrate the hardness of improving the theorem without pre-processing the data, e.g., applying an embedding method to the sample points, such as VaDE or PCA. In particular, we show that i) without further complicating the algorithm, the separation conditions in Theorem 1 have the right dependence on $d$ and $n$; ii) for exact recovery, the sample-size lower bound in Theorem 1 is also tight up to logarithmic factors of $k$.

These results and arguments indeed justify our motivation for combining VaDE with Ward’s method. To further demonstrate our approach’s power, we experimentally evaluate the clustering performance of three methods on high-dimensional equally separated Gaussian mixtures – the vanilla Ward’s method, the PCA-Ward combination, and the VaDE-Ward combination. As shown in Figure 2(a), in terms of the dendrogram purity, Ward’s method with VaDE pre-processing tops the other two, for nearly every mean-separation level. In terms of the final clustering result in the (projected) space, Figures 2(b) to 2(d) clearly shows the advantage of the VaDE-Ward method as it significantly enlarges the separation among different clusters while making much fewer mistakes in the clustering accuracy.

Next, we show that given the exact recovery in Theorem 1 and some mild cluster separation conditions, Ward’s method recovers the target hierarchy.

For any index set $I\subseteq[k]$ and the corresponding Gaussian components $\mathcal{G}_{i},i\in I$, denote by $w_{I}$ the total weight $\sum_{i\in I}w_{i}$, and by $C_{I}$ the union of sample points from these components. We say that a collection $\mathcal{H}$ of nonempty subsets of $[k]$ form a hierarchy of order $\tau$ over $[k]$ if there exists a index set sequence $T_{\ell}$ for $\ell=1,\ldots,\tau$ such that i) for every $\ell$, the set $\mathcal{H}_{\ell}$ partitions $[k]$; ii) $\mathcal{H}_{k}$ is the union of $\mathcal{H}_{\ell}:=\{I_{\ell t}:t\in T_{\ell}\}$; iii) for every $\ell<\tau$ and $t\in T_{\ell}$, we have $I_{\ell t}\subseteq I_{(\ell+1)t^{\prime}}$ for some $t^{\prime}\in T_{\ell+1}$.

For any sample size $n$ and Gaussian component $\mathcal{G}_{i}$, denote by $S_{i}:=\sigma_{i}(\sqrt{d}+2\sqrt{\log n})$, which upper bounds the radius of the corresponding sample cluster, with high probability. The triangle inequality then implies that for any $i\not=j\in[k]$, the distance between a point in the sample cluster of $\mathcal{G}_{i}$ and that of $\mathcal{G}_{j}$ is at most $D_{ij}^{+}:=\left\lVert\mu_{i}-\mu_{j}\right\rVert_{2}+S_{i}+S_{j}$, and at least $D_{ij}^{-}:=\left\lVert\mu_{i}-\mu_{j}\right\rVert_{2}-S_{i}-S_{j}$. Naturally, we also define $D_{I,J}^{+}:=\max_{i\in I,j\in J}D_{ij}^{+}$ and $D_{I,J}^{-}:=\min_{i\in I,j\in J}D_{ij}^{-}$ for any $I,J\subseteq[k]$.

The above theorem resembles Theorem 1, with separation conditions comparing the cluster radiuses with the inter-cluster distances at different hierarchy levels. Informally, the theorem states that if the clusters in each level of the hierarchy are well-separated, Ward’s method will be able to recover the correct clustering of that level, and hence also the entire hierarchy.

We empirically determine which embedding method produces the best results for dendrogram purity and Moseley-Wang’s objective. We use Ward’s method for clustering after embedding (Appendix C has other algorithms). We provide results for our synthetic model (BTGM) and three real datasets.

Set-up and methods.   We compare VaDE against PCA (linear baseline) and a standard VAE (non-linear baseline). All the embedding methods are trained on the whole dataset. Then we run Ward’s method on the sub-sampled embedded data and the original data to get the dendrograms. In the experiments, we only use 25 classes from CIFAR-100 that fall into one of the five superclasses “aquatic animals”, “flowers”, “fruits and vegetables”, “large natural outdoor scenes” and “vehicle1”. We call it CIFAR-25, and we also use CIFAR-25-s to refer to the images labeled with the five superclasses. We average the results over 100 runs, each with 10k random samples from each dataset (except 20newsgroup where we use all 18k samples). Details in Appendix C.

Synthetic Data (BTGM).   We vary the depth and margin of the BTGM evaluate different regimes of hierarchy depth and cluster separation. After generating $k$ different means for the BTGM, we shift all the non-zero entries of the first $\frac{k}{2}$ means to an arbitrary location. The goal is to show that this simple non-linear transformation is challenging for linear embedding methods (i.e., PCA), while preserving the hierarchical structure. To generate the dataset of $k$ clusters with shifted means, we sample 2000 points from each Gaussian distributions of the BTGM of height h such that $2^{h}=k$. Then, we set the embedding dimension of PCA, VAE and VaDE to be $h$.

Dendrogram Purity (DP).   We compute the DP of various methods using (3). We construct ground truth clusters either from class labels of real datasets or from the leaf clusters of the BTGM model.

Moseley-Wang’s Objective (MW).   To evaluate the MW objective in (2), we need pairwise similarities $w_{ij}$. For classification datasets, we let $w_{ij}=\mathbbm{1}_{\mathcal{C}(x^{i})=\mathcal{C}(x^{j})}$ to be the indicator function for point $i$ and $j$ being in the same class. When there are hierarchical labels, we instead define $w_{ij}=2^{l-1}*\mathbbm{1}_{\mathcal{C}^{l}(x^{i})=\mathcal{C}^{l}(x^{j})}$, where $\mathcal{C}^{l}(x^{i})$ denotes the label for point $x^{i}$ at the $l^{th}$ level of the hierarchical tree from top to bottom. We measure the ratio achieved by various methods to the optimal MW value opt that we compute as follows. For $k$ flat clusters, the optimal MW value is obtained by first merging points in the same clusters, then building a binary tree on top of the $k$ clusters. When there are hierarchical labels, we compare against the optimal tree with ground truth levels.

We investigated the effectiveness of embedding data in Euclidean space to improve unsupervised hierarchical clustering methods. We saw that rescaling the VaDE latent space embedding, and then applying Ward’s linkage-based method, leads to the best results for dendrogram purity and the Moseley-Wang objective on many datasets. To better understand the VaDE embedding, we proposed a planted hierarchical model using Gaussians with separated means. Theoretically, we proved that Ward’s method recovers both the clustering and the hierarchy with high probability in this model.

While we focus on embedding using variational autoencoders, an open direction for further research could involve embedding hierarchical structure using other representation learning methods agarwal2007generalized; borg2005modern; caron2018deep; chierchia2019ultrametric; mishne2019diffusion; nina2019decoder; salakhutdinov2012one; tsai2017learning; yadav2019supervised. Another direction is to better understand the differences between learned embeddings, comparison-based methods, and ordinal relations emamjomeh2018adaptive; ghoshdastidar2019foundations; kazemi2018comparison; jamieson2011low. Finally, it would be worth exploring the use of approximate clustering methods to improve the running time abboud2019subquadratic; kobren2017hierarchical; rashtchian2017clustering.