Year: 2021

Abstract:
We propose a novel probabilistic model over hierarchies on graphs obtained by continuous relaxation of tree-based hierarchies. We draw connections to Markov chain theory, enabling us to perform hierarchical clustering by efficient end-to-end optimization of quality metrics such as Dasgupta cost or Tree-Sampling Divergence (TSD). We show that our model learns rich, high-quality hierarchies present in 12 real world graphs, including a large graph with 2.3M nodes. Our model consistently outperforms recent as well as strong traditional baselines such as average linkage. We demonstrate that our model also obtains competitive results on link prediction despite not being trained on this task, highlighting the quality of the hierarchies discovered by our model.

text: We propose a novel probabilistic model over hierarchies on graphs obtained by continuous relaxation of tree-based hierarchies. We draw connections to Markov chain theory, enabling us to perform hierarchical clustering by efficient end-to-end optimization of quality metrics such as Dasgupta cost or Tree-Sampling Divergence (TSD). We show that our model learns rich, high-quality hierarchies present in 12 real world graphs, including a large graph with 2.3M nodes. Our model consistently outperforms recent as well as strong traditional baselines such as average linkage. We demonstrate that our model also obtains competitive results on link prediction despite not being trained on this task, highlighting the quality of the hierarchies discovered by our model.

heading: 1 INTRODUCTION
text: Clustering is a fundamental problem in unsupervised learning, both in theory and practice. In contrast to traditional ‘flat’ clustering, hierarchical clustering brings several advantages. It enables us to naturally analyze (Zhao & Karypis, 2002) and visualize (Himberg et al., 2004) a given dataset on different scales. Further, different downstream tasks may require different granularities. Hierarchical clustering lets us pick the desired granularity after learning. It can also be used for personalized recommendation (Zhang et al., 2014) or to solve record linkage tasks (Wick et al., 2012). In an influential paper, Eisen et al. (1998) perform hierarchical clustering on gene expression data and argue that it could be used to discover the yet-unknown meaning of certain genes, which could lead to advances in medicine. Besides traditional vector data, real-world graphs are often scale-free and hierarchically organized (Ravasz & Barabási, 2003; Barabási & Pósfai, 2016). This makes them interesting candidates for uncovering the underlying hierarchy via hierarchical clustering. Thus, the hierarchical clustering method we introduce is particularly useful for analyzing scale-free real-world networks such as Web graphs, citation networks, flight networks, or biological networks.
Related work. We can group hierarchical clustering algorithms into discrete and continuous optimization methods. Discrete optimization methods can themselves be decomposed into agglomerative approaches and divisive approaches. Agglomerative methods – like the famous linkage algorithms (Gower & Ross, 1969; Jardine & Sibson, 1968; Bonald et al., 2018) – follow a bottom-up approach and iteratively aggregate the two most similar clusters w.r.t. a given similarity measure. Divisive methods follow a top-down approach and recursively split larger clusters into smaller ones using, e.g., k-means (Steinbach et al., 2000) or spectral clustering (Charikar & Chatziafratis, 2016). While Moseley & Wang (2017); Kamvar et al. (2002) have shown connections with explicit objective functions, those methods are mainly based on heuristics to avoid exhaustive combinatorial optimization.
On the other hand, continuous optimization methods circumvent this issue by (i) relaxing the discrete tree structure, (ii) defining an explicit quality metric for hierarchies, and (iii) optimizing it using gradient-based optimizers. Chierchia & Perret (2019); Monath et al. (2017; 2019); Chami et al. (2020) focus on minimizing the Dasgupta cost (Dasgupta, 2016) (see Eq. (1)) to fit an ultrametric, hyperbolic embeddings, or a probabilistic model over cluster assignments. Monath et al. (2017) learns a probabilistic assignment of the samples to leaves of a tree with a fixed binary structure. The method of Monath et al. (2019) requires approximation for the computation of lowest common ancestors (LCA) probabilities and regularization losses. Beyond the Dasgupta cost, Charpentier & Bonald (2019) propose the Tree Sampling Divergence (TSD), a quality metric for hierarchies on graphs. Both
metrics have the benefit to be internal metrics, i.e., they do not require ground truth information to evaluate the quality of a hierarchy. Most of the aforementioned methods focus on vector data.
Contributions. In this work, we focus on hierarchical clustering in homophilic graphs, where the similarity of two nodes is made explicit by the presence or absence of an edge. We (1) propose a probabilistic model over hierarchies via continuous relaxation of a tree’s parent assignment matrices. We (2) theoretically analyze the model by drawing connections to absorbing Markov chains, which allows (3) efficient and exact computation of relevant quantities (e.g., LCA probabilities). This enables us to (4) learn hierarchies on graphs by efficient, end-to-end optimization of quality metrics such as Dasgupta cost and the TSD score. Our extensive experimental evaluation on 12 real-world graphs, including a massive graph with more than 2M nodes and 60M edges, highlights the effectiveness of our Flexible Probabilistic Hierarchy model (FPH). It outperforms all baselines, traditional and recent, on the quality of the learned hierarchies measured both by TSD and Dasgupta cost. Remarkably, FPH also performs competitively on link prediction despite not being trained on this task.
2 HIERARCHICAL GRAPH CLUSTERING PRELIMINARIES
We define a graph G = (E, V ) with n nodes and m undirected edges. We let w(vi, vj) be equal to the weight of edge (vi, vj) if they are connected and to 0 otherwise, and w(vi) = ∑ j w(vi, vj) to be equal to the weight of node vi. This allows us to define the edge distribution P (vi, vj) ∝ w(vi, vj) normalized over all edges s.t. ∑ ij P (vi, vj) = 1 and the node distribution
P (vi) ∝ w(vi) s.t. ∑
i P (vi) = 1, which are later needed to compute the quality metrics of learned hierarchies.
Discrete hierarchical clustering. We define a hierarchical clustering T of a graph G as a tree structured partitioning of the nodes V . We denote the nodes of the tree T with n′ ∈ N internal nodes by {x1, ..., xn+n′}. The tree nodes can be decomposed into the leaf nodes
{v1, ..., vn} corresponding to the nodes of the input graph, and the internal nodes {z1, ..., zn′} representing clusters of leaf nodes; overall {x1, ..., xn+n′} = {v1, ..., vn}∪ {z1, ..., zn′}. Since any tree is a directed acyclic graph (DAG), we impose w.l.o.g. a topological ordering on the internal nodes of the tree in the remainder of this work. That is, zi cannot be a parent of zj if i < j; zn′ denotes the root node. We can uniquely define a tree as T = (Â, B̂) with two adjacency matrices, Â ∈ {0, 1}n×n′ and B̂ ∈ {0, 1}n′×n′ . Â describes the connections from the leaves to the internal nodes, i.e. Âi,k = 1 if internal node zk is the parent of leaf node vi. B̂ describes the connections between the internal nodes, i.e. Bk,l = 1 if internal node zl is the parent of internal node zk. Since in a tree any node has exactly one parent, Â is row-stochastic and B̂ is row-stochastic; except its last row B̂n′ corresponding to the root, whose sum is 0. We denote by vi ∧ vj the internal node which is the lowest common ancestor (LCA) of vi, vj in the tree. The maximum possible number of internal nodes is n′ = n− 1. Probabilistic hierarchical clustering. The central idea of FPH is to define a probabilistic tree via continuous relaxation of the binary-valued adjacency matrices Â, B̂. I.e., we define A, B such that A ∈ [0, 1]n×n′ and B ∈ [0, 1]n′×n′ while keeping the row-stochasticity constraints, i.e., ∑n′ j Ai,j = 1 for 1 ≤ i ≤ n, ∑n′ l Bk,l = 1, 1 ≤ k < n′, and for the root node we have∑n′
l Bn′,l = 0. This describes a probabilistic hierarchy; the probability of an internal node zj to be a parent of a leaf or an internal node is Aij := p(zj |vi) or Bij := p(zj |zi), respectively. See Fig. 1 for a visualization. Our key technical contribution are efficient closed-form expressions for the lowest common ancestor probabilities that arise when sampling discrete trees from the probabilistic one.

heading: 2.1 INTERNAL METRICS FOR HIERARCHICAL CLUSTERING
text: Internal metrics are commonly used to evaluate clustering methods, and they have the important advantage of not requiring external labels. For hierarchical clustering on graphs, two internal metrics have been proposed to assess the quality of a discrete hierarchy T given an input graph G: Dasgupta cost (Das.) (Dasgupta, 2016), and Tree-Sampling Divergence (TSD) (Charpentier & Bonald, 2019).
While the Dasgupta cost can be used for graphs and vector data, TSD is tailored specifically to graphs.
Das(T ) = ∑
vi,vj∈V P (vi, vj) · c(vi ∧ vj), TSD(T ) = KL(p(z)||q(z)) = ∑ z p(z) log p(z) q(z) (1)
For Das., c(z) = |z| denotes the number of leaves for which the internal node, i.e., z = vi ∧ vj , is an ancestor. Intuitively, the cost is large if vi and vj are connected by an edge and their lowest common ancestor contains many leaves, i.e., is close to the root. Thus, Dasgupta favors leaves connected by an edge to have LCAs low in the hierarchy. For Das., lower scores indicate better hierarchies.
TSD defines two distributions over the internal nodes, p(z) = ∑
vi,vj :z=vi∧vj P (vi, vj), induced by the edge distribution, and q(z) = ∑ vi,vj :z=vi∧vj P (vi)P (vj), induced by the independent node distribution. Intuitively, we expect the graph distribution p(z) (induced by edge sampling) and the null model distribution q(z) (induced by independent node sampling) to differ significantly if the tree T indeed represents the hierarchical structure of the graph, leading to high TSD scores. Continuous relaxation. To enable end-to-end gradient-based optimization of Dasgupta or TSD, we replace the discrete hierarchy T = (Â, B̂) with its continuous relaxation T = (A,B). This leads to probabilistic parent assignments of leaves to internal nodes. By itself, this does not help us in hierarchical clustering as it is not clear how to compute the lowest common ancestor probabilities when sampling discrete hierarchies from the relaxed one, i.e., the probability that an internal node is the lowest common ancestor of two leaves. Our main theoretical contribution is to derive closed-form expressions for the lowest common ancestor probabilities in Sec. 3. This gives rise to continuous versions of TSD and Dasgupta (see equations in Appendix A.1), which we will refer to as Soft-TSD and Soft-Das., respectively. For the first time, this allows us to directly optimize for hierarchical clustering quality metrics in an end-to-end fashion instead of proxy losses or heuristic algorithms.
For optimization over hierarchies, we either minimize the Dasgupta cost or maximize TSD:
min A,B Soft-Das(T = (A,B)) or max A,B Soft-TSD(T = (A,B)), (2)
where A ∈ [0, 1]n×n′ ,B ∈ [0, 1]n′×n′ are the continuous relaxations of the parent assignment matrices as described in Sec. 2. Note that for both metrics we recover the same score as their discrete versions in the case of a deterministic probabilistic model, i.e. when A, B are binary-valued.

heading: 2.2 SAMPLING DISCRETE HIERARCHIES
text: Given the probability matrices A and B we can easily recover a discrete hierarchical clustering by sampling discrete matrices Â and B̂, which in turn describe a (discrete) tree T . For each leaf and internal node we independently sample its parent from the categorical distribution described by the respective row in A or B. As we show in Appendix A.2, this tree-sampling procedure (denoted T = (Â, B̂) ∼ PA,B(T )) leads to valid tree hierarchies, and we can directly compute the probability of any discrete hierarchy given the probabilistic one. We denote probabilities associated with the tree-sampling perspective as p(T )(·). Note that we can easily obtain a discrete tree given continuous A,B in a deterministic way by selecting for each leaf and internal node its most likely parent.

heading: 3 EFFICIENT, DIFFERENTIABLE HIERARCHIES VIA MARKOV CHAINS
text: Outline. This section is organized as follows. The goal is to derive efficient, closed-form equations to compute (lowest common) ancestor probabilities which are consistent with the tree-sampling procedure explained in the previous section. For this, we draw connections to absorbing Markov chains in Sec. 3 and show in Secs. 3.1, 3.2 how to compute the desired quantities efficiently under this simplifying Markov Chain perspective. Then, we show that these equations under the Markov chain and the tree-sampling perspectives are equivalent, i.e., we do not introduce any error by the Markov chain assumptions. Finally, we show in Sec. 3.3 how to exploit the independence assumptions of the Markov chain to compute the (lowest common) ancestor probabilities efficiently and in a vectorized way.
We start by showing that the probabilistic model described in the previous section indeed defines an absorbing Markov chain. Intuitively, A and B can be interpreted as transition matrices from leaves to internal nodes and among internal nodes, respectively.
Definition 1 (Tree Markov Chain). Let A ∈ [0, 1]n×n′ , ∑n′
j=1Aij = 1 ∀1 ≤ i ≤ n and B ∈ [0, 1]n ′×n′ , ∑n′ j=1Bij = 1 ∀1 ≤ i < n′, ∑n′
j=1Bn′j = 0. We define a Markov chain M = (S ∪ {ω},T ) with state set S ∪ {ω} and transition matrix T , where S = {v1, v2, . . . , vn} ∪ {z1, z2, . . . , zn′}. Further,
T ∈ R|S|+1×|S|+1 = [ Q w 0 1 ] ,Q ∈ R|S|×|S| = [ 0 A 0 B ] is the transition matrix of the Markov chain M in canonical form, where w ∈ R|S| = (0, 0, . . . , 0, 1)T is the vector of transition probabilities to state ω. Note that we add here an auxiliary state ω, which acts as final absorbing state from the root. Theorem 1. LetM be a tree Markov chain as defined in Definition 1.M is an acyclic absorbing Markov chain with ω being its only absorbing state. (See proof in App. A.3)
Thus, our probabilistic hierarchy defined by A and B describes an absorbing and acyclic Markov chain. In our hierarchical clustering interpretation, a random walk starts at one of the leaves and first randomly transitions to the internal nodes based on A, followed by transitions among the internal nodes via B. Once the root is reached, the random walk is absorbed after one last step to ω.
In the remainder, we denote the probabilities arising from the tree Markov-chain as p(M)(·) to distinguish them from their tree-sampling counterparts p(T )(·). In general, the two perspectives are not equal, e.g., when considering pairs of leaf nodes. Here, the Markov chain introduces an independence assumption of pairs of random walks. Under tree-sampling, two leaves’ paths to the root are dependent in general, i.e., two paths traversing some internal node zk implies that all subsequent transitions are identical. While Markov-chain probabilities p(M)(·) often have efficient and analytical solutions grounded in Markov chain theory, it is unclear a-priori how to leverage them for hierarchical clustering. In the following, we relate central tree-sampling probabilities p(T )(·), i.e., (lowest common) ancestor probabilities, to their Markov chain counterparts p(M)(·).

heading: 3.1 ANCESTOR PROBABILITIES
text: First, we want to compute the probability p(T )anc (zk|vj) of an internal node zk being an ancestor of leaf vj under the row-wise tree-sampling procedure (see Sec. 2.2). We call this an ancestor probability.
Theorem 2. Let T = (Â, B̂) ∼ PA,B(T ) be a discrete hierarchy obtained by tree-sampling. Further, letW(vi) be a random walk onM rooted in vi resulting in path r̂i.We have that
p(T )anc (zk|vj) = p(zk ∈ r̂i) =: p(M)anc (zk|vj). (3) (See proof in App. A.4)
Note that any random walk onM must result in a path since, becauseM is acyclic, no transient state can be visited more than one time in a random walk. Theorem 2 means that we can use ancestor probabilities arising from the tree-sampling and Markov chain views interchangeably.

heading: 3.2 LOWEST COMMON ANCESTOR PROBABILITIES
text: For the hierarchical clustering problem we are interested in the lowest common ancestor (LCA) of two leaf nodes vi and vj , which we denote by vi ∧ vj . The LCAs are required to compute both the Dasgupta cost as well as the Tree-Sampling Divergence (TSD). Therefore, when optimizing for good Dasgupta or TSD scores, it is crucial to exactly and efficiently compute the LCA probabilities that arise. We denote these LCA probabilities by p(T )(zk = vi∧vj). Previous work resorts to heuristics in approximating the LCA probabilities (Monath et al., 2019). In contrast, our connection to absorbing Markov chains admits efficient closed-form computation of LCA probabilities, as we now show. Theorem 3. Let A and B be transition matrices describing a (soft or discrete) hierarchy as introduced in Section 2. Let ri = (r (1) i , . . . , r (T−1) i , zn′) denote a path from leaf vi = r (1) i to the root zn′ , where |ri| = T . Then, the probability of internal node zk being the lowest common ancestor vi ∧ vj of leaf nodes vi, vj under the tree-sampling procedure from Section 2.2 is
p(T )(zk = vi ∧ vj) = ∑
(ri,rj): zk=vi∧vj
p(ri) · p(rj), (4)
where ri is the part of a path from vi up to (and including) internal node zk, i.e., ri = (r (1) i , . . . , zk). (See proof in App. A.5)
From Eq. (4) it seems that a straightforward way to compute the LCA probabilities is to enumerate the set {(ri, rj) : zk = vi ∧ vj} of all pairs of potential paths from vi and vj to the internal node zk and sum their probabilities. However, this is intractable in practice because the size of the set grows exponentially in the ‘depth’ of the internal node k (see App. A.8, Theorem 7).
In Theorem 4 we show how to compute the LCA probabilities that arise when sampling pairs of random walks from the Markov chainM. Theorem 4. Let M be a Markov chain as defined in Definition 1, zk be an internal node, and vi, vj be leaf nodes. Let W(vi),W(vj) be independent random walks on M rooted in vi and vj , respectively. Then, the probability that internal node zk is the lowest internal node (as by the topological order) traversed by both random walks, i.e., the lowest common ancestor of vi and vj , is
p(M)(zk = vi ∧ vj) = panc(zk|vi)panc(zk|vj)− k−1∑ k′=1 p(M)(zk′ = vi ∧ vj)panc(zk|zk′)2 (5)
(See proof in App. A.6)
Thus, due to the independence of pairs of random walks in the Markov chainM, we can compute LCA probabilities efficiently in closed form on the Markov chain. However, the joint ancestor probability p (M) anc (zk|vi, vj) does not reflect the underlying process of sampling T = (Â, B̂) ∼ PA,B(T ) any more, i.e., in general p(M)anc (zk|vi, vj) = panc(zk|vi) · panc(zk|vj) 6= p(T )anc (zk|vi, vj). As our main theoretical result, we show next that, remarkably, the LCA probabilities obtained from the Markov chain are indeed equivalent to the tree-sampling LCA probabilities: Theorem 5. Let A and B be transition matrices describing a (soft or discrete) hierarchy as described in Sec. 2. The LCA probabilities arising from tree-sampling and the Markov chain are equal, i.e.,
p(T )(zk = vi ∧ vj) = p(M)(zk = vi ∧ vj). (6) (See proof in App. A.7)
This result is surprising at first, since we made the assumption that pairs of random walks are independent in the Markov chain sampling process. However, recall that two paths in the treesampling process are disjoint and independent until they meet at their LCA zk. Thus, for this subset of the pairs of paths leading to zk, the independence assumption inM does not lead to an error. The result is also very useful in practice, since it means that we can use efficient computations from the Markov chain view to exactly compute the tree-sampling LCA probabilities. As we show in the following, we can jointly compute LCA probabilities for all pairs of leaves in a vectorized way.

heading: 3.3 EFFICIENT VECTORIZED COMPUTATION
text: Definition 2 (Fundamental Matrix). LetM be an absorbing Markov chain and Q as in Definition 1. The fundamental matrix N of Markov chainM is
N ∈ R|S|×|S| = (I −Q)−1 = [ I A(I −B)−1 0 (I −B)−1 ] . (7)
Nij equals the expected number of visits of state j when starting a random walk in state i. SinceM is acyclic, each transient state can be visited at most once, i.e., Nij is the probability of state j being on a random path to the root starting from i. Observing the block structure of N in Eq. (7), we define
P anc ∈ Rn×n ′ := A(I −B)−1, P̃ anc ∈ Rn ′×n′ := (I −B)−1 − I (8)
P ancij = panc(zj |vi) is the probability of internal node zj ending up being an ancestor of vi under the tree-sampling procedure. Analogously, P̃ anc provides the ancestor probabilities among the internal nodes, i.e. P̃ ancij = panc(zj |zi). Note that we subtract the identity matrix in Eq. (8) since the diagonal entries of the fundamental matrix block (I −B)−1 are always trivially 1, as they correspond to the probability of traversing an internal node zk when also starting from zk. By subtracting I we obtain the ancestor probabilities by enforcing that an internal node is not its own ancestor. Since B is strictly upper triangular, the inverse (I −B)−1 always exists and is efficient to compute in O(n′2).
Theorem 6. The vector of LCA probabilities of all internal nodes w.r.t. leaf nodes vi and vj can be computed in a vectorized way via
P LCAvi,vj ∈ R n′ = (P ancvi P anc vj ) T · (I + P̃ anc P̃ anc)−1, (9) where denotes the element-wise (Hadamard) product. (See proof in App. A.9)
Complexity analysis. As we show in App. A.10, we can exploit sparsity in real-world graphs and thus do not have to construct the full Rn×n×n′ LCA tensor P LCA. This leads to a complexity of O(m× n′2) for both Soft-Das. and Soft-TSD, which is efficient since typically we have n′ n. All equations are vectorized and thus benefit from GPU acceleration. More details in App. A.12.

heading: 3.4 INTEGRAL SOLUTIONS
text: In App. A.13 we analyze the properties of the relaxed problem Soft-TSD which our method FPH optimizes. We prove in Theorem 10 that the optimization problem is integral, i.e., the global maximum is discrete. This is remarkable, since we are actually optimizing over continuous hierarchies parameterized by A and B. This implies that the global maximum of the relaxed problem is the same as for the combinatorial problem of optimizing over discrete hierarchies. Soft-Das. is not obviously convex or concave, thus not obviously integral.

heading: 3.5 FURTHER CONSIDERATIONS
text: Choice of n′. The number of internal nodes n′ is an important hyperparameter of our method (as well as most baselines). Similar to, e.g., the number of clusters k in k-means, large numbers of internal nodes n′ lead to more expressive hierarchies, which on the other hand are less interpretable by a human and require more memory and computation. In Fig. 2, we show how the expressiveness (as measured by TSD/Dasgupta) improves for increasing values of n′, and in App. B.6 we visualize learned hierarchies with different n′. Since FPH typically trains within a few minutes, our general recommendation is to use the elbow method (Thorndike, 1953) to determine n′.
Constrained vs. unconstrained optimization. Since our probabilistic hierarchy model leads to fully differentiable metrics (i.e., Soft-Das. and Soft-TSD), we can optimize the metrics in an end-to-end fashion via gradient descent. Note that the matrices A and B are constrained to be row-stochastic; we therefore experiment with two optimization schemes: unconstrained optimization of A and B (e.g. using Adam optimizer and softmax to obtain row-stochastic matrices), or constrained optimization via projected gradient descent (PGD). In our ablation study we found that the PGD optimization consistently leads to better results (see Fig. 4). We attribute this performance difference to the severe gradient re-scaling in the softmax operation when the parameters become very large or small, leading to very small step sizes (Niculae, 2020); thus, unless otherwise stated, FPH uses PGD optimization.
Scaling to large graphs. For fast training on large graphs on commodity GPUs, we propose a simple yet effective batching scheme. We uniformly pick K random leaves from the graph at each iteration. Then we select the induced subgraph of these K nodes and their neighbors while capping the total number of nodes by some constant C. We then compute the loss and perform the update based on the selected subgraph. By using our batching procedure we do not need to have all parameters on the GPU, which enables scaling to very large graphs.
Node embeddings. In our default setting, we directly parameterize the matrices A and B as learnable parameters via gradient descent. That is, the direct parent probabilities p(zk|vi) and p(zk|zk′) are the only trainable parameters. As an alternative, we also experiment with learning a node embedding for each leaf and internal node. We compute the parent probabilities A and B via softmax on the negative Euclidean distances of the embeddings. Given a fixed embedding size, this leads to a parameterization that scales linearly in the number of leaf nodes and internal nodes.
Initialization. We have found that our model (as well as most of the baselines) can greatly benefit from a “smart” initialization scheme. For the direct parameterization, we have found initializing from the solution obtained from the average linkage algorithm (Jardine & Sibson, 1968) to work well. Unless stated otherwise, FPH uses initialization from average linkage. In contrast to vector data, linkage algorithms are fast on graph data and can be performed in O(m) (Benzécri, 1982; Murtagh & Contreras, 2012), thus not affecting the complexity of FPH. In the embedding parameterization, we experiment with initializing the embeddings using DeepWalk (Perozzi et al., 2014).

heading: 4 EXPERIMENTS
text: Datasets. We use 11 real world datasets (McCallum et al., 2000; Sen et al., 2008; Aspert et al., 2019; Amunts et al., 2013; Cho et al., 2014; Adamic & Glance, 2005; Patokallio; Wang et al., 2020; Yang & Leskovec, 2015), including the very large ogbn-products dataset with around 2.3M nodes and 62M edges (Hu et al., 2020). We always select the largest connected component (LCC) as a preprocessing step and convert each graph to an undirected one. We provide further information about the datasets in Table 7, including number of nodes, number of edges and mutual information (MI) between nodes which is an upper-bound on the TSD score (Charpentier & Bonald, 2019). In our experiments we report normalized TSD in percent.
Baselines. We compare our method with the following deep learning baselines: Routing Gradientbased clustering (RGHC) (Monath et al., 2017), Hyperbolic Gradient-based Clustering (HGHC) (Monath et al., 2019), Ultrametric Fitting (UF) (Chierchia & Perret, 2019), and HypHC (Chami et al., 2020). Importantly RGHC, HGHC and HypHC also optimize a relaxed version of the Dasgupta cost. Moreover, we compare to the average linkage algorithm (Avg. link.). as well as the Ward linkage algorithm (Ward) (Ward Jr., 1963). Finally, we compare to the Louvain method (Blondel et al., 2008) (Louvain). For RGHC, HGHC, and HypHC, which require vector data as input (as opposed to graphs), we use DeepWalk embeddings as node features. Note that linkage algorithms return a full hierarchy with n′ = n − 1. Therefore we use the compression scheme introduced in (Charpentier & Bonald, 2019) to reduce the size of the hierarchies to the desired number of internal nodes. See App. B.7 for our hyperparameter choices.
For the results on Dasgupta and TSD, we train each baseline 5 times and report results for the best run.1 For all models, we apply a post-processing treatment consisting in pruning unused internal nodes. For FPH, we select for each node its most likely parent to obtain a discrete tree.

heading: 4.1 RESULTS
text: Hierarchical clustering. We report results on hierarchical clustering with n′ = 512 internal nodes in Table 1. Results for Louvain are in Table 9. FPH outperforms all baselines on all datasets on the Dasgupta cost function as well as TSD. Remarkably, the average linkage algorithm outperforms most of the deep-learning based baselines. We partially attribute this to the fact that the baselines use DeepWalk embeddings as node attributes.
1Since we are not using any label information, this does not lead to data leakage or overfitting on test data; we show standard deviations in Table 13. Since FPH is deterministic, we only run it once.
n = 100 n = 1000 Model Level 1 Level 2 Level 3 Level 1 Level 2 Level 3
Avg. 0.985 0.952 0.791 0.990 0.978 0.965 RGHC 1.0 0.881 0.766 0.567 0.897 0.745 HGHC 0.874 0.872 0.705 0.840 0.942 0.786 Ward 1.0 0.975 0.765 1.0 0.991 0.841 UF 1.0 0.921 0.707 1.0 0.987 0.791 HypHC 1.0 0.892 0.768 0.933 0.846 0.702 Louvain 0.693 0.955 0.795 0.678 1.0 0.993 FPH 1.0 0.994 0.829 1.0 1.0 0.994
Moreover, we noticed that HypHC requires an excessive amount of triplet samples to obtain good results. Even n2, as suggested by the authors, performs poorly. We use 50M triplets for all datasets, which is more than 4.5 times the recommended number of n2 on Wiki-Physics. This took almost 24 hours to compute, which is why we could not go higher. We clearly see a dependence of HypHC’s performance on the number of triplets, as it performs competitively on PolBlogs, our smallest dataset, and poorly on Wiki-Physics, the largest dataset on which we were able to run HypHC. In Tbls. 10, 11 we show baseline results with DW d = 128 embeddings, which did not improve results consistently.
In Figure 2 we show how the models perform for different numbers of internal nodes n′. As expected, the models generally obtain better scores for higher capacity. Again, we observe that FPH performs best across datasets and values of n′, highlighting the effectiveness of our approach. Specifically on the Dasgupta cost, FPH substantially outperforms the baselines. This is remarkable since RGHC, HGHC and HypHC also optimize a relaxed version of the Dasgupta cost.
External evaluation. As we typically have no knowledge about ground-truth hierarchies in realworld data, it is difficult to perform external evaluation. To address this, we propose the following:
(i) We compute the normalized mutual information (NMI) between the ground-truth class labels and the learned hierarchies when cutting them appropriately to divide the leaves into the same number of (flat) clusters. We set n′ = 256 for all models (except RGHC, since it ran OOM; we use n′ = 128 instead) and set FPH to optimize TSD. Note, though, that the node labels are not necessarily solely based on connectivity of nodes, but e.g., also on node attributes, to which the models do not have access. Thus, we cannot expect perfect correlation of “good” hierarchies in the sense of explaining edges in the graph and NMI w.r.t. ground-truth node labels. In Table 2 (more datasets in Tab. 12) we show the results; the hierarchies learned by FPH achieve strong scores, highlighting its effectiveness.
(ii) Similarly, we can use the ground-truth clustering from synthetic hierarchical stochastic blockmodel (HSBM) (Lyzinski et al., 2016) graphs to compute the NMI scores at each level of the hierarchy. We use five HSBM graphs with 100 and 1000 nodes, respectively. The graphs have three levels (see Figs. 6,7 for example graphs). On the small graphs we set n′ = 64; on the large graphs, we use the same settings as in (i). In Table 3 we show the results. FPH outperforms all baselines and is able to recover the ground-truth hierarchies to a very high accuracy.
Ablation study. In our ablation study, we compare (i) FPH (i.e., with constrained optimization using PGD) vs. unconstrained optimization with Adam (FPH-U) and using softmax; (ii) average linkage initialization vs. random initialization (FPH-R, FPH-UR); (iii) direct optimization of A,B vs. learning node embeddings (FPH Emb.). Here, we experiment with random initialization (FPH-R Emb.) as well as DeepWalk initialization (FPH Emb. DW). The results, obtained on the WikiPhysics dataset, are displayed in Figure 4 (App.). We observe that FPH, i.e. constrained PGD optimization of A and B with initialization from average linkage performs best. In a similar way, “smart” initialization from DeepWalk tends to improve the results on FPH Emb., too, but the effect is less pronounced. In general, FPH variants learning embeddings perform worse than our full version; this may be due to the gradient-rescaling due to softmax, which we also suspect is a reason why unconstrained optimization achieves weaker results than FPH.
Scalability. We run experiments on the very large graph dataset ogbn-products, which has about 2.3M nodes and 62M edges. We train FPH both with n′=512 and 1024 internal nodes performing batching as explained in Sec. 3.5. We show the results in Table 4. FPH outperforms all baselines both on TSD as well as Dasgupta, and with both n′=512 and 1024. Remarkably, FPH achieves more than 50% better Dasgupta results compared to the best baseline Ward, and cuts the average linkage
results by a factor of more than 3 for n′=1024, highlighting that FPH outperforms all baselines by a large margin. Note that we could not compare with HypHC, since (a) the implementation constructs a dense n× n matrix, which would require more than 14TB memory assuming single precision; (b) sampling n2 triplets, as recommended by the authors, would take weeks. Further, we do not report results on UF since it did not converge. RGHC ran out of memory for n′=1024.
Runtime. Due to the efficient, vectorized computations, a full epoch of a complete dataset to compute Soft-TSD or Soft-Dasgupta using our model is very fast. A complete forward pass on Wiki-Physics with 512 internal nodes takes only about 130ms on GPU and about 2s on CPU. On the other hand, a full evaluation on ogbn-products takes only about 13 minutes (758s) on CPU, and training the model on a GPU takes approximately 3.5 hours for 2K epochs (with batching).
Link prediction. As the TSD score can be interpreted in terms of a reconstruction loss, we can use the scores of the reconstruction scheme in Eq. (25) (App. B.1) for predicting links. We compare with DeepWalk (DW) and the variational graph autoencoder (VGAE) (Kipf & Welling, 2016), where the link prediction score of two nodes is the cosine of their respective embeddings. Further, we compare with degree-corrected stochastic blockmodels (DCSBM) (Karrer & Newman, 2011) and Adamic Adar (Ad./Ad.) (Adamic & Adar, 2001), which are established strong baselines for link prediction in non-attributed graphs. See App. B.2 for details. Importantly, none of the models use node features, including FPH. As shown in Table 5, FPH achieves best or second-best performance on all datasets. This is remarkable, since FPH is not trained specifically for this task, which highlights the generality and usefulness of the hierarchies discovered by FPH.
learning (L2, row 2), symbolic AI (L2, row 4), or
variational inference (L2, row 6). Again, FPH did not see any text information (node attributes) and performed the clustering based on citations (edges) alone. Besides offering qualitative evidence that the learned hierarchies are reasonable, this also showcases a potential real-world application of FPH: a scholar can explore the hierarchy of topics in an academic field to discover relevant papers.

heading: 5 CONCLUSION
text: We propose a new probabilistic model over hierarchies on graphs which can be learned using end-to-end gradient-based optimization. By drawing connections to absorbing Markov chains we can compute relevant quantities such as lowest common ancestor probabilities exactly and efficiently in a vectorized way. For the first time, this allows to directly optimize for quality metrics for hierarchical clustering such as Dasgupta cost or Tree-Sampling Divergence (TSD) in an end-to-end fashion. Our Flexible Probabilistic Hierarchy model outperforms strong traditional as well as recent deep-learning-based baselines on nearly all datasets and tasks considered and easily scales to massive graphs with millions of nodes.

heading: ETHICS STATEMENT
text: Since our method does not directly focus on a specific real-world application, good or bad societal outcomes depend on how practitioners and researchers use it in practice. This means that it could potentially be abused, e.g., by corporations or governments, to identify groups and hierarchies of dissidents via recorded metadata. For instance, phone call metadata is routinely collected at scale by service providers and governments. By using our method on the large-scale call graph it could be possible to identify groups of political dissidents and to repress them. On the other hand, we argue that our method can also have positive impact, e.g., by making corpora of literature more accessible to users or by enabling scientists, e.g., to discover cliques and hierarchies in biological networks. Moreover, our contribution is algorithmic in nature and does not consider the effects of potential biases or discrimination in the underlying network data.

heading: REPRODUCIBILITY STATEMENT
text: For reproducibility and verifiability of our theoretical results, we provide complete proofs of all ten theorems of our work in Appendices A.3-A.13. We further make explicit all assumptions and definitions we use to derive our results. For reproducibility of our experimental results, we first highlight that our model’s core implementation is a straightforward PyTorch implementation of the matrix equations in this work. Further, we detail our experimental approach in Sec. 4 and provide hyperparameter choices for our method as well as the baselines in Table 8 in the appendix. For the baselines, we use the authors’ official implementations and use the suggested hyperparameters. To compute the Dasgupta and TSD metrics (as well as to obtain the results for the Louvain algorithm), we use the sknetwork Python library.2 We share our codebase with the reviewers in a private comment on OpenReview; upon acceptance, we will make our implementation available to the general public.

heading: A PROOFS AND THEORETICAL ANALYSIS

heading: A.1 CONTINUOUS VERSIONS OF TSD AND DASGUPTA
text: Our probabilistic hierarchy model enables us to replace the discrete parent assignments in the Dasgupta cost and TSD with the parent probabilities from the relaxed adjacency matrices A and B. This, in turn, leads to LCA probabilities which are consistent under the tree-sampling procedure. For the first time, this allows us to directly and efficiently optimize for hierarchical clustering quality metrics in an end-to-end fashion instead of proxy losses or heuristic algorithms.
Hence, we propose the soft Dasgupta cost (Soft-Das) by replacing the cardinality of the lowest common ancestor |vi ∧ vj | by the expected cardinality under the probabilistic model, i.e. c(z) =∑
v panc(z|v).
Soft-Das(PA,B(T )) = ∑
vi,vj∈V P (vi, vj) ∑ z p(z = vi ∧ vj)c(z)
In the same way, we propose a differentiable version of the soft Tree Sampling Divergence (SoftTSD). To this end, we replace the discrete assignments in the distributions p(z) and q(z) with probabilistic assignments, i.e.
Soft-TSD(PA,B(T )) = KL(p(z)||q(z)) = ∑
z
p(z) log p(z) q(z)
where p(z) = ∑ vi,vj p(z = vi ∧ vj)P (vi, vj)
q(z) = ∑ vi,vj p(z = vi ∧ vj)P (vi)P (vj)
Note that for both metrics we recover the same score as their discrete formulations in the case of a deterministic probabilistic model, i.e. when A and B are binary-valued.

heading: A.2 TREE-SAMPLING PROCEDURE
text: Recall our assumption that the internal nodes are ordered, and that Bij = 0 if j ≤ i. This implies that there are no possible cycles, or equivalently that B is a strictly upper-triangular matrix, i.e., it describes a directed acyclic graph (DAG). Combined with the fact that each node in a tree (except the root) has exactly one parent, we see that the sampled discrete hierarchy is indeed a tree. We denote this tree-sampling process by T = (Â, B̂) ∼ PA,B(T ). We can also compute the probability of any tree T under the sampling procedure described above:
PA,B(T = (Â, B̂)) = ∏ i,j A Âi,j i,j ∏ i′,j′ B B̂i′,j′ i′,j′ (10)
Note that Ai,j and Bi′,j′ are the probabilities of internal nodes zj and zj′ to be a parent of leaf and internal nodes vi and zi′ respectively while Âi,j and B̂i′,j′ are equal to 1 if these connections exist in the tree T , else 0.

heading: A.3 PROOF OF THEOREM 1
text: Proof. For M to be absorbing (i) it must have at least one absorbing state, and (ii) at least one absorbing state must be reachable from any state in a finite number of steps. For (i), ω is an absorbing state since its self-transition probability Tk,k = 1, where k = |S|+ 1 is its corresponding index in the transition matrix T . Thus, once reached, a random walk cannot leave the state ω. To show (ii), we note that since the transient state transition matrix Q is a strictly upper-triangular matrix, any random walk onM must lead to state zn′ . From state zn′ , the random walk transits to ω with probability wn′ = 1. Thus, since n′ is finite, state ω can be reached from any state in a finite number of steps. Further, ω is the only absorbing state since all self-transition probabilities of states in S are zero as diag(Q) = 0. Since Q is strictly upper-triangular, none of the transient states can be visited more than once on a random walk, and thereforeM is acyclic.

heading: A.4 PROOF OF THEOREM 2
text: Proof. We can arbitrarily define the order in which we sample from the categorical distributions in A and B because of the independence of the sampling steps. We choose to start by sampling first from Ai, i.e., the row corresponding to the leaf node vi under consideration: w(1) ∼ Cat(Ai). Next, we sample from the row corresponding to w(1), and repeat until we reach zk, i.e.
w(t) ∼ Cat(Bw(t−1)) for 1 < t ≤ T, where w(T ) = zk. For the remaining entries, we continue in arbitrary order. Observe that (w(1), . . . , w(T )) are the ancestors of leaf vi in T . Further observe that this sampling procedure is identical to how the path r̂i is generated in the random walkW(vi), completing the proof.

heading: A.5 PROOF OF THEOREM 3
text: Proof. First, recall that all paths end in the root node, such that ri necessarily ends in zn′ = r (T ) i . When reasoning about the lowest common ancestors, it is no longer sufficient to consider the ancestors (or, equivalently, path to the root) of a single leaf node in isolation. Instead, we need to consider pairs of dependent paths (ri, rj), i 6= j rooted in vi and vj , respectively. Note that ri and rj necessarily converge at some internal node zk = vi ∧ vj — the latest at the root node zn′ .
Thus, we denote with ri = (r (1) i , . . . , vi ∧ vj) the first part of the path ri until (and including) its lowest common ancestor with rj , i.e., vi ∧ vj . Analogously, ri = (r (|ri|+1) i , . . . , zn′), such that ri = (ri, ri). Further, note that the paths ri and rj are on the same underlying hierarchy T , thus ri = rj , as both paths have the same trajectory to the root once they have reached their lowest common ancestor, i.e., they are dependent.
The probability of observing the pair of paths (ri, rj) under the tree-sampling perspective is
p(T )((ri, rj)) = p(r (1) i |vi) · |ri|∏ t=2 p(r (t) i |r (t−1) i ) · p(r (1) j |vj) · |rj |∏ t=2 p(r (t) j |r (t−1) j ) · |ri|∏ t=|ri|+1 p(r (t) i |r (t−1) i ) (11) More compactly,
p(T )((ri, rj)) = p(ri) · p(rj) · p(ri) = p((ri, rj)) · p(ri) = p((ri, rj)) · p(rj) (12)
Importantly, we can see from Eq. (12) that, in general, p(T )((ri, rj)) 6= p(ri) · p(rj) i.e., the paths ri and rj are not independent. We denote p (T ) anc (zk|vi, vj) the probability of the internal node zk to be the ancestor of leaf nodes vi and vj under the tree-sampling perspective. Hence, the probability of the internal node zk to be the ancestor of leaf nodes vi and vj under dependent and independent random walks are different i.e. p(T )anc (zk|vi, vj) 6= panc(zk|vi) ·panc(zk|vj). This makes intuitive sense because knowing that zk′ is an ancestor of vi and vj in a tree T , additional knowledge that zk, k > k′ is an ancestor of vi implies that zk is also an ancestor of vj .
However, for the parts of ri and rj before they converge, Eq. (12) shows that p((ri, rj)) = p(ri)·p(rj). This is because all transitions in ri and rj are disjoint thus independent. This is an important insight because it means that the probability of observing two paths both converging at an internal node zk factorizes.

heading: A.6 PROOF OF THEOREM 4
text: Proof. We start by reorganizing Eq. (5):
(i)︷ ︸︸ ︷ panc(zk|vi)panc(zk|vj) = (ii)︷ ︸︸ ︷ p(M)(zk = vi ∧ vj) + (iii)︷ ︸︸ ︷ k−1∑ k′=1 p(M)(zk′ = vi ∧ vj)panc(zk|zk′)2 (13)
In words, we can split the event “zk is an ancestor of vi and vj” (i) into two mutually exclusive events: (ii) zk is the lowest common ancestor of vi and vj ; or (iii) some internal node lower in the topological
order is the LCA of vi and vj , and further, both random walks also traverse through zk. (ii) and (iii) are mutually exclusive since exactly one internal node is the LCA for vi and vj on any two random walks.
Since the two random walks are independent, the probability of zk being traversed on both walks factorizes. Thus, p(M)anc (zk|vi, vj) = panc(zk|vi)panc(zk|vj) and therefore (i) is the probability of zk being an ancestor of vi and vj in our Markov chainM. The events in (iii) can indeed be expressed:
p(M) (zk′ = vi ∧ vj , zk ∈ anc (vi, vj)) = p(M) (zk′ = vi ∧ vj) · p(M) (zk|zk′ ∈ anc (vi, vj)) (14) where in the last step we exploit that zk′ = vi ∧ vj implies zk′ ∈ anc (vi, vj) as well as the Markov property of the random walks. Further, note that
p(M)anc (zk|zk′ ∈ anc(vi, vj)) = panc(zk|zk′ ∈ anc(vi)) · panc(zk|zk′ ∈ anc(vj)) = panc(zk|zk′)2,
(15)
where we first exploit factorization due to independence and in the last step again the Markov property of the random walks.

heading: A.7 PROOF OF THEOREM 5
text: Proof. Let r̂i ∈ P(vi), r̂j ∈ P(vj) be two independent random walks onM rooted in vi and vj , respectively. Then,
p(M)(zk = vi ∧ vj) = ∑
(r̂i,r̂j):zk=vi∧vj
p((r̂i, r̂j)), (16)
where r̂i = (r̂ (1) i , . . . , zk) is the first part of r̂i until it reaches zk. Note that the second part of the paths, r̂i and r̂j , which are theoretically independent under our Markov chain model, are marginalized out in the LCA formula.
However, due to the independence of the first part of the paths r̂i and r̂j under both models (see Eq. (4)), we can write:
p(M)(zk = vi ∧ vj) = ∑
(r̂i,r̂j):zk=vi∧vj
p((r̂i, r̂j))
= ∑
(r̂i,r̂j):zk=vi∧vj
p(r̂i) · p(r̂j)
= p(T )(zk = vi ∧ vj).
(17)

heading: A.8 NUMBER OF PAIRS OF LCA PATHS.
text: Theorem 7. LetM be a Markov chain as defined in Definition 1. The number of pairs of paths from two leaves vi, vj for which an internal node zk is the lowest common ancestor is 3k−1.
Proof. Proof by induction over k. For the base case k = 1 we have one pair of paths for which zk is the LCA, i.e. directly from vi to zk and vj to zk. Assume that for internal node zk there are 3k−1 unique pairs of paths for which zk is the LCA. For each of these paths we can generate three unique paths for which zk+1 is the LCA. (1) rewire the last transition of vi’s path to go to zk+1 instead of zk. (2) do the same but for vj . (3) rewire both vi’s and vj’s last transition to go to zk+1 instead of zk. Thus, the number of pairs of paths from vi and vj to zk+1 is 3 · 3k−1 = 3k.

heading: A.9 PROOF OF THEOREM 6
text: We provide here the proof for the fast vectorized computation of P LCAvi,vj in Theorem. 6.
Proof. We start by reorganizing Eq. (9):
P ancvi P anc vj = P LCA vi,vj + P LCA,T vi,vj · P̃ anc P̃ anc (18)
Note that the inverse (I + P̃ anc P̃ anc)−1 is guaranteed to exist and is efficient to compute because P̃ anc P̃ anc is a strictly upper triangular and therefore nilpotent matrix. The k-th entry is thus:[
P ancvi P anc vj ] k = P LCAvi,vj ,vk + n′∑ k′=1 P LCAvi,vj ,vk′ · [ P̃ anc P̃ anc ] k′,k .
Plugging in the definitions of Eq. (8), using Theorem 2, and observing that due to the upper triangular structure P̃ anck′,k = 0 for k ′ > k we obtain
panc(zk|vi) · panc(zk|vj) =p(zk = vi ∧ vj) + k−1∑ k′=1 p(zk′ = vi ∧ vj) · panc(zk|zk′)2.
A.10 VECTORIZED COMPUTATIONS
All quantities involved in Soft-Das and Soft-TSD can be computed in closed-form based on the Markov chainM and its fundamental matrix. However, their computation should not be done naively, as this involves unnecessary computations. Constructing the full tensor of LCA probabilities is expensive since P LCA ∈ Rn×n×n′ . Note, however, that to compute the Soft Dasgupta loss or the distribution p(z) in TSD we only require the LCA probabilities p(z = vi ∧ vj) for pairs of leaves connected by an edge (i.e., P (vi, vj) > 0). That is, we only need to construct an LCA probability matrix of shape Rm×n′ . Thus, we can exploit the sparsity of real world graphs, as typically m n2. In a similar way, the computation of the distribution q(z) does also not require the expensive explicit computation of p(zk = vi ∧ vj) for all pairs of leaf nodes. Instead, we can again exploit insights from the Markov chainM. First, observe that the equation of q(z) in Soft-TSD describes an expectation:
q(z) = Evi,vj∼P (v) [p(z = vi ∧ vj)] . (19)
Defining p̂anc = pT · P anc, the computation of this expectation can be vectorized similarly to Theorem 6 (see derivation in App. A.11): q = (p̂anc p̂anc)T · (I + P̃ anc P̃ anc)−1.
A.11 VECTORIZED q COMPUTATION.
We provide here the proof for the fast vectorized computation of q in Eq. 19.
Proof. We first rewrite the expectation Eq. 19 in vectorized form: q = ∑ i,j pviP LCA vi,vjpvj
where we denote P (vi) = pvi . Subsequently, we can plug the P LCA formula Eq.. (9) and pull pvi into the Hadamard product which is done over the internal node dimension:
q = ∑ i,j pvipvj (P anc vi P anc vj ) T · (I + P̃ anc P̃ anc)−1
= ∑ i pvi P ancvi ∑ j pvjP anc vj T · (I + P̃ anc P̃ anc)−1 =
∑ i pviP anc vi ∑ j pvjP anc vj T · (I + P̃ anc P̃ anc)−1 = (p̂anc p̂anc)T · (I + P̃ anc P̃ anc)−1.
where p̂anc = pT · P anc.

heading: A.12 COMPLEXITY ANALYSIS.
text: Both Dasgupta and TSD computations require to compute the ancestor probabilities (Eq. 8) which can be done in O(n×n′2) (i.e. inverse of triangular matrix (I−B) ∈ Rn′×n′ , plus matrix multiplication with A ∈ Rn×n′). Then, Soft-Das. or the distribution p(z) for the Soft-TSD loss require the LCA probabilities (Eq. 9) for all leaves connected by an edge only, amounting to O(m× n′2) operations, wherem is the number of edges in the graph. Note that similarly to Eq. (8), the inverse computation in Eq. (9) can be done in O(n′2). Additionally, Soft-TSD requires the computation of q(z) ( complexity O(n× n′2)). Both Soft-Das. and Soft-TSD computations are dominated by the O(m× n′2) term. This leads to an efficient time complexity as long as we assume a small number of internal nodes n′ n, which is reasonable in practice.

heading: A.13 CONVEXITY OF SOFT-TSD
text: Theorem 8. Let H ∈ [0, 1]n′×n×n be a tensor whose elements Hkij = p(zk = vi ∧ vj) are the LCA probabilities of internal nodes w.r.t. pairs of leaf nodes. Soft-TSD(H) is convex in H .
Proof. Let H(1), H(2) be two LCA probability tensors defined as above, and 0 ≤ α ≤ 1. We first compute the distribution p induced by the edge distribution. We have that
p(αH (1) k + (1− α)H (2) k ) = ∑ vi,vj P (vi, vj)(αH (1) kij + (1− α)H (2) kij)
= α ∑ vi,vj P (vi, vj)H (1) kij + (1− α) ∑ vi,vj P (vi, vj)H (2) kij
= α · p(H(1)k ) + (1− α) · p(H (2) k ).
(20)
Analogously we compute the distribution q induced by the independent ndoe distribution. We have that
q(αH (1) k + (1− α)H (2) k ) = α · q(H (1) k ) + (1− α) · q(H (2) k ). (21)
We combine this with the well-known fact that KL-divergence is convex w.r.t. pairs of distributions, i.e.,
KL(αp1(z)+(1−α)p2(z), αq1(z)+(1−α)q2(z)) ≤ αKL(p1(z), q1(z))+(1−α)KL(p2(y), q2(z)), to obtain the desired result:
Soft-TSD(αH(1) + (1− α)H(2)) ≤ α · Soft-TSD(H(1)) + (1− α) · Soft-TSD(H(2)). (22)
Theorem 9. LetHFPH = {H : ∃A,B ∈ Φ(n, n′) : H = FPH(A,B)} denote the set of probabilistic hierarchies which can be represented by FPH. Here, Φ(n, n′) are the constraints FPH places on A,B (see Sec. 2), and FPH (A,B) is shorthand the mapping from transition matrices to lowest common ancestor probability tensors defined in Theorems 4 and 6. Here, Hkij = p(zk = vi ∧ vj) are the LCA probabilities of internal nodes w.r.t. pairs of leaf nodes. This setHFPH is convex.
Proof. We start by recalling form Theorems 4 and 5 that FPH computes lowest common ancestor probabilities p(zk = vi ∧ vj) from the continuous parent probability matrices A and B such that the LCA probabilities are consistent with the expected result from the tree-sampling procedure described in Sec. A.2. More formally,
Hkij := p(zk = vi ∧ vj) = EĤ∼(A,B) [I[zk = vi ∧ vj ]] = EĤ∼(A,B) [ Ĥkij ] = E [ Ĥ ] kij , (23)
where Ĥ ∈ {0, 1}n′×n×n is a discrete hierarchy obtained via tree-sampling from A and B. By definition of the expectation we write
H = EĤ∼(A,B) [ Ĥ ] = ∑
Ĥ∈H(n,n′)
p(Ĥ|A,B) · Ĥ, (24)
where H(n, n′) is the set of all valid discrete hierarchies with n leafs and n′ internal nodes. Thus, any continuous hierarchy H learned by FPH is a convex combination of discrete hierarchies Ĥ . This completes the proof.
Theorem 10. The Soft-TSD optimization problem solved by FPH is integral. That is, the global maximum of the Soft-TSD optimization problem solved by FPH is the same as the global optimum of the discrete optimization problem of optimizing TSD over discrete hierarchies.
Proof. This follows from Theorems 8 and 9. Theorem 8 establishes that the Soft-TSD objective function is convex in the hierarchy tensors H; Theorem 9 proves that the set of hierarchies FPH optimizes over is convex. When maximizing a convex function over a convex set, we are guaranteed to find the global optimum at a vertex of the constraint set, which are discrete hierarchies in the case of FPH. Thus, the global maximizer is a discrete hierarchy; this discrete hierarchy must also be the maximizer of the discrete TSD optimization problem, since our relaxation optimizes over a superset of all discrete hierarchies.
Implications of Theorems 8, 9, and 10. In the previous theorems, we have shown that we are maximizing a convex function over a convex set. In general, maximizing a convex function over a convex set is NP-hard (Benson, 1995). Thus, we cannot hope to efficiently recover the global optimum. However, our continuous relaxation brings several practical benefits for the optimization.
First, observe that directly optimizing over the convex set of continuous hierarchies described in Theorem 9 is not practical. This is because there are exponentially many corners of the set, and encoding the constraints of the set is very difficult. Our parameterization of (continuous) hierarchies via A, B and being able to efficiently compute the expected lowest common ancestor probabilty tensor enables us to optimize over a fairly low-dimensional and convex set. The constraints on A,B, i.e., entries in [0, 1], unit row sums and upper-triangular structure of B, are easy to encode and enforce during optimization. This comes at the cost that mapping from A and B to the LCA tensor H is nonconvex (yet describes, as per Theorem 9, a convex set over hierarchies). Thus, we can solve the optimization problem with off-the-shelf methods such as projected gradient descent and benefit from the elaborate techniques from nonconvex optimization. While it is possible that FPH gets stuck in a non-discrete local optima during optimization, we can easily obtain a discrete and valid hierarchy given the non-discrete local optimizer via tree-sampling or selecting the most likely parent for all leaves and internal nodes under A and B, as described in Sec. 2.2.

heading: B EXPERIMENT INFORMATION

heading: B.1 LINK PREDICTION WITH SOFT-TSD
text: The TSD can be interpreted in terms of retrieved information when reconstructing the original graph from the tree representation (Charpentier & Bonald, 2019). In this case, the reconstruction scheme for the edge weights of the reconstructed graph Ĝ is:
ŵ(vi, vj) = w(vi)w(vj) p(vi ∧ vj) q(vi ∧ vj)
(25)

heading: B.2 LINK PREDICTION SETUP
text: For all datasets, we randomly select 10% of edges to hold out for testing while making sure that the graph remains connected. Further, we set n′ = 256 and minimize Soft-TSD via FPH. For DC-SBM, we use the Python package ‘graph-tool’ and follow the documentation3 with default parameters to learn the model. For VGAE, we use the default hyperparameters by the authors (one hidden layer, latent dimensions [32, 16], learning rate 0.01, training for 200 epochs). We use the variant described in the paper which replaces the node attributes by the n× n identity matrix. For DeepWalk, we set the embedding dimension to 10.
3https://graph-tool.skewed.de/static/doc/demos/inference/inference.html

heading: B.3 DATASET SUMMARY
text: See Table 7 for an overview of the datasets we used.

heading: B.4 ADDITIONAL RESULT FIGURES
text: In Fig. 3 we show results for four more datasets.

heading: B.5 ABLATION STUDY
text: See Fig 4 for the comparison of different FPH model variants, and our full discussion in Sec. 4.1.

heading: B.6 HIERARCHY VISUALIZATION
text: We conduct a qualitative study of the structure discovered by FPH. In Figure 5 we compare the hierarchies learned by FPH when optimizing for TSD or Dasgupta, respectively. We cut at different levels of the dendrogram to obtain a coarse hierarchy (10 clusters) and fine-grained structure (50 clusters). Comparing Figure 5 (a) and (d), we notice that the coarse structure learned by optimizing Soft-Dasgupta looks more appealing, as TSD essentially splits only into the Americas and the rest
of the world. At 50 clusters, however, we observe the opposite: TSD splits the airports across the world into meaningful, coherent geographical regions, whereas Dasgupta looks mostly unchanged from the coarse version, highlighting the complementarity of both quality metrics. In addition, the dendrogram learned by TSD in (f) appears to be of higher quality and more balanced than the Dasgupta dendrogram in (c).

heading: B.7 HYPERPARAMETERS
text: We use the hyperparameters for models and baselines described in Tab. 8. Note that we train FPH for 1,000 epochs and restore the best hierarchy after training. For ogbn-products we train for 2,000 epochs. Further, to ensure convergence we reduce the learning rate to 0.02 and introduce weight decay of 1e-5 for FPH (Das.) for n′ ≥ 128 on Polblogs, Cora-ML, and Brain, and train for 2,000 epochs. Similarly, we use different learning rates for A and B for FPH (Das.) on ogbn-arxiv, ogbl-collab, and DBLP (lrA = 1e− 2, lrB = 1e− 5).

heading: B.8 COMPUTING INFRASTRUCTURE
text: We train all models on a single GPU (NVIDIA GTX 1080 Ti or NVIDIA GTX 2080 Ti, 11 GB memory) in our own in-house compute cluster. The machines have 10-core Intel CPUs. We use Python 3 and PyTorch for all our experiments.

heading: B.9 HSBM GRAPHS
text: We generated HSBMs with n = 100 leaf nodes and n = 1000 leaf nodes for our external evaluation. The small HSBMs have 3 levels with edge probabilities in [.01, .1, .3, .6], a branching factor of 2 and core community sizes in [10, 15]. The large HSBMs have 3 levels with edge probabilities in [.001, .01, .1, .4], branching factor in [2, 3, 4] and core community sizes in [30, 35]. In Fig. 6 and Fig. 7, we plot one of the five synthetic HSBM graphs used in our experiments (right), and their corresponding dendrograms (left). The graphs have n = 100 and n = 1000 leaf nodes and three levels of hierarchy.
C ADDITIONAL RESULTS

Review summaries:
The paper proposes a probabilistic clustering method based on two losses: Dasgupta loss and Tree-sampling divergence. A scalable sampling-based optimization is proposed. The modes is optimized with the projected gradient descent.

The paper proposes a probabilistic model of tree-based hierarchical clustering by considering a continuous relaxation. The authors establishes the theorical connections between tree-sampling and the Markov chain that is specifically constructed. This allows an end-to-end gradient-based learning to optimize the internal metrics of graph clustering, namely the Dasgupta cost and Tree Sampling Divergence. This method allows efficient and effective modeling of the graph clustering. Extensive experiments results on different types of graph, including some of the large scale ones, suggests the proposed method behave favorably. 

This paper presents a hierarchical clustering over graphs using a probabilistic model via continuous relaxation of the discrete tree-based hierarchies. It uses two metrics the Dasgupta cost and the Tree-Sampling Divergence. The paper presents an efficient vectorized computation and discusses several details on the choice of hyper parameters. The article presents empirical evaluation on 8 real datasets and compares the technique with several baselines on hierarchical clustering (using NMI), comparing also results on controlled clusters (ground truth class labels, synthetic stochastic block-model) and link prediction (best or second based performance), as well as on meta-analysis which includes ablation, scalability, runtime, and a qualitative assessment.