# END-TO-END LEARNING OF PROBABILISTIC HIERARCHIES ON GRAPHS

**Daniel Zügner, Bertrand Charpentier, Sascha Geringer, Morgane Ayle,**
**Stephan Günnemann**
Technical University of Munich
{zuegnerd, charpent, geringer, ayle, guennemann}@in.tum.de

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 relaxed versions of quality metrics such as Dasgupta cost or TreeSampling Divergence (TSD). We show that our model learns rich, high-quality
hierarchies present in 11 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. Our model also obtains strong results on link
prediction despite not being trained on this task, highlighting the quality of the
hierarchies discovered by our model.

1 INTRODUCTION

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 opti-**
mization 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; Charpentier, 2019) – 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. 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 relaxed versions of quality metrics such as Dasgupta cost and the TSD score. Our
extensive experimental evaluation on 11 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 = (E, V ) with n nodes and m undi- 𝑝(𝑧𝑗[|𝑧]𝑖[)] 𝑧1
_G_ 𝑧2
rected edges. We let w(vi, vj) be equal to the weight 𝐵
of edge (vi, vj) if they are connected and to 0 other- 𝑧3
wise, and w(vi) = _j_ _[w][(][v][i][,][ v][j][)][ 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_ 𝑣1 𝑣2 𝑣3 𝑣4 𝑣5 𝑣6 𝑣7 𝑣8 𝐴
_∝[P]_ Learned Hierarchy

𝑝(𝑧𝑗[|𝑧]𝑖[)] 𝑧1

𝑧2

𝑧3

[|𝑣]𝑖[)]

𝑣1 𝑣2 𝑣3 𝑣4 𝑣5 𝑣6 𝑣7 𝑣8

edges s.t. _ij_ _[P]_ [(][v][i][,][ v][j][) = 1][ and the][ node distribution]

_P_ (vi) _w(vi) s.t._ _i_ _[P]_ [(][v][i][) = 1][, which are later needed] 𝑣1 𝑣2 𝑣5 𝑣6 Matrix Form
_∝_

to compute the quality metrics of learned hierarchies.[P] 𝑣3 𝑣4 𝑣7 𝑣8

𝑣1 𝑣2 𝑣5 𝑣6

𝑣3 𝑣4 𝑣7 𝑣8

Observed Graph

**Discrete hierarchical clustering.[P]** We define a hierarchi_cal clustering_ _T[ˆ] of a graph G as a tree structured partition-_ Figure 1: Model overview.
ing 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[′] repre_{_ _}_ _{_ _}_
senting 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′ correspond-_**
ing to the root, whose sum is 0. We denote byancestor (LCA) of vi, vj in the tree. The maximum possible number of internal nodes is vi ∧ vj the internal node which is the lowest common n[′] = n − 1.

**Probabilistic hierarchical clustering. The central idea of FPH is to define a probabilistic tree**
_T 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_ **_[A][i,j][ = 1][ for][ 1][ ≤]_** _[i][ ≤]_ _[n][,][ P]l[n][′]_ **_[B][k,l][ = 1][,][ 1][ ≤]_** _[k < n][′][, and for the root node we have]_

_l_ **_[B][n][′][,l][ = 0][. This describes a probabilistic hierarchy; the probability of an internal node][ z][j][ 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

Pa visualization. Our key technical contribution are efficient closed-form expressions for the lowest[P][n][′] _|_ _|_

common ancestor probabilities that arise when sampling discrete trees from the probabilistic one.


2.1 INTERNAL METRICS FOR HIERARCHICAL CLUSTERING

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( [ˆ] ) = _P_ (vi, vj) I[z=vi vj ]c(z), TSD( [ˆ] ) = KL(p(z) _q(z))_ (1)
_T_ _∧_ _T_ _||_

vi,Xvj _∈V_ Xz

For Das., c(z) = vi∈V [I][[][z][∈][anc][(][v]i[)]][ denotes the number of leaves for which the internal node][ z][ 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

[P]

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 [I][[][z][=][v]i[∧][v]j []][P] [(][v][i][,][ v][j][)][, induced]

by the edge distribution, and q(z) = vi,vj [I][[][z][=][v]i[∧][v]j []][P] [(][v][i][)][P] [(][v][j][)][, induced by the independent node]

distribution. Intuitively, we expect the graph distribution p(z) (induced by edge sampling) and the

[P]

null model distribution q(z) (induced by independent node sampling) to differ significantly if the tree

[P]

_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 (details 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 or max
**_A,B_** [Soft-Das][(][T][ = (][A][,][ B][))] **_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.

2.2 SAMPLING DISCRETE HIERARCHIES

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
ˆ = (A[ˆ], **_B[ˆ])_** _PA,B(_ )) leads to valid tree hierarchies, and we can directly compute the probability
_T_ _∼_ _T_
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._**

3 EFFICIENT, DIFFERENTIABLE HIERARCHIES VIA MARKOV CHAINS

**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 pro_cedure 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][′], _j=1_ _[A][ij][ = 1][ ∀][1][ ≤]_ _[i][ ≤]_ _[n][ and][ B][ ∈]_

[0, 1][n][′][×][n][′], _j=1_ _[B][ij][ = 1][ ∀][1][ ≤]_ _[i < n][′][,][ P]j[n]=1[′]_ _[B][n][′][j][ = 0][. We define a Markov chain][ M][ =]_

( _ω_ _, T ) with state set_ _ω_ _and transition matrix[P][n] T[′]_ _, where_ = v1, v2, . . ., vn
_S ∪{_ _}_ _S ∪{_ _}_ _S_ _{_ _} ∪_
z1, z2, . . ., z[P]n′[n][′]. Further,
_{_ _}_

**_Q_** **_w_** **0** **_A_**
**_T_** R[|S|][+1][×|S|][+1] = _, Q_ R[|S|×|S|] =
_∈_ **0** 1 _∈_ **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. Let M 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][)](·).

3.1 ANCESTOR PROBABILITIES

First, we want to compute the probability p[(]anc[T][ )][(][z]k[|][v]j[)][ of an internal node][ z]k [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, let W(vi) be a random walk on M rooted in vi resulting in path ˆri.We have that_

_p[(]anc[T][ )][(][z][k][|][v][j][) =][ p][(][z][k]_ _ri) =: p[(]anc[M][)][(][z][k][|][v][j][)][.]_ (3)

_[∈]_ [ˆ]
_(See proof in App. A.4)_

Note that any random walk on M must result in a path since, because M 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.

3.2 LOWEST COMMON ANCESTOR PROBABILITIES

For the hierarchical clustering problem we are interested in the lowest common ancestor (LCA) of
two leaf nodesDasgupta cost as well as the Tree-Sampling Divergence (TSD). Therefore, when optimizing for good vi and vj, which we denote by vi ∧ vj. The LCAs are required to compute both the
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 intro-**
_duced in Section 2. Let ri = (ri[(1)], . . ., ri[(][T][ −][1)], zn′_ ) denote a path from leaf vi = ri[(1)] _to the root zn′_ _,_
_whereof leaf nodes |ri| = T vi. Then, the probability of internal node ̸= vj under the tree-sampling procedure from Section 2.2 is zk being the lowest common ancestor vi ∧_ vj

_p[(][T][ )](zk = vi ∧_ vj) = _p(ri) · p(rj),_ (4)

(ri,rj ):X zk=vi∧vj


-----

_where ri is the part of a path from vi up to (and including) internal node zk, i.e., ri = (ri[(1)], . . ., zk)._
_(See proof in App. A.5)_

For vi=vj, an LCA probability is trivially the parent probability, i.e., p(zk=vi vi) = p(zk vi). From
_∧_ _|_
Eq. (4) it seems that a straightforward way to compute the LCA probabilities is to enumerate the
setand sum their probabilities. However, this is intractable in practice because the size of the set grows {(ri, rj) : zk = vi ∧ vj} of all pairs of potential paths from vi and vj to the internal node zk
exponentially in the ‘depth’ of the internal node k (see App. A.8, Thm. 7). In Thm. 4 we show how
to compute the LCA probabilities that arise when sampling pairs of random walks from M.
**Theorem 4. Let M be a Markov chain as defined in Definition 1, zk be an internal node, and**
vvij ̸, respectively. Then, the probability that internal node= vj be leaf nodes. Let W(vi), W(vj) be independent random walks on zk is the lowest internal node (as by the M rooted in vi and
_topological order) traversed by both random walks, i.e., the lowest common ancestor of vi and vj, is_


_k−1_

_p[(][M][)](zk′ = vi_ vj)panc(zk zk′ )[2] (5)
_k[′]=1_ _∧_ _|_

X


_p[(][M][)](zk = vi ∧_ vj) = panc(zk|vi)panc(zk|vj) −

_(See proof in App. A.6)_


Thus, due to the independence of pairs of random walks in the Markov chain M, we can compute LCA
probabilities efficiently in closed form on the Markov chain. However, the joint ancestor probability
_p[(]anc[M][)][(][z]k[|][v]i[,][ v]j[)][ does not reflect the underlying process of sampling][ T][ = (][ ˆ]A,_ **_B[ˆ]) ∼_** _PA,B(T ) any_
more, i.e., in general p[(]anc[M][)][(][z]k[|][v]i[,][ v]j[) =][ p]anc[(][z]k[|][v]i[)][ ·][ p]anc[(][z]k[|][v]j[)][ ̸][=][ p][(]anc[T][ )][(][z]k[|][v]i[,][ v]j[)][.]

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 in does not lead to an error.
_M_

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.

3.3 EFFICIENT VECTORIZED COMPUTATION

**Definition 2 (Fundamental Matrix). Let M be an absorbing Markov chain and Q as in Definition 1.**
_The fundamental matrix N of Markov chain M is_

**_I_** **_A(I_** **_B)[−][1]_**
**_N_** R[|S|×|S|] = (I **_Q)[−][1]_** = _−_ _._ (7)
_∈_ _−_ **0** (I **_B)[−][1]_**
 _−_ 

**_Nij equals the expected number of visits of state j when starting a random walk in state i. Since M_**
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] R[n][×][n][′] := A(I **_B)[−][1],_** **_P˜_** [anc] R[n][′][×][n][′] := (I **_B)[−][1]_** **_I_** (8)
_∈_ _−_ _∈_ _−_ _−_
_Pij[anc]_ = 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[˜]ij[anc]_ = 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_

**_Pv[LCA]i=vj_** vi **_Pv[anc]j_** [)][T][ ·][ (][I][ + ˜]P _[anc]_ **_P_** _[anc])[−][1]_ **_Pv[LCA]i,vi_** [=][ A][v][i][,] (9)
_̸_ _[∈]_ [R][n][′][ = (][P][ anc] _⊙_ _⊙_ [˜]

_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 R[n][×][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.

3.4 INTEGRAL SOLUTIONS

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.

3.5 FURTHER CONSIDERATIONS

**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).


-----

Dasgupta cost (lower is better) Normalized TSD (higher is better)

Alg. Ward Louv. UF HypHC HGHC RGHC Avg. lk. FPH Ward Louv. UF HypHC HGHC RGHC Avg. lk. FPH

Brain 618.81 777.14 712.33 571.64 749.40 556.57 556.68 **503.67** 31.72 29.28 28.61 17.48 24.18 22.05 28.91 **32.34**
OpenFlight 382.45 633.66 393.58 463.43 487.96 488.90 363.40 **355.61** 55.48 51.51 53.89 39.08 49.50 39.56 52.02 **57.72**
Genes 202.17 247.26 251.01 495.26 366.53 247.07 196.50 **183.63** 66.80 67.47 62.95 20.66 53.33 51.81 66.72 **67.69**
Citeseer 92.27 178.23 98.61 215.62 150.26 131.89 83.69 **77.16** 69.43 68.45 67.40 37.22 57.61 50.61 67.80 **69.37**
Cora-ML 281.82 336.86 342.86 442.09 411.49 350.00 292.77 **254.78** 56.47 57.51 53.06 30.73 46.76 42.68 55.30 **58.02**
PolBlogs 377.63 443.48 350.74 330.58 354.86 433.77 355.61 **262.48** 27.54 25.93 25.23 22.21 23.94 19.41 25.25 **31.41**
WikiPhysics 736.11 986.32 753.81 759.07 840.15 740.87 658.04 **537.95** 45.28 46.03 43.40 32.02 39.70 38.39 43.15 **49.97**
ogbn-arxiv 22,870 31,655 52,666 OOM 22,076 24,077 20,671 **14,354** 36.77 37.75 24.75 OOM 26.05 25.21 33.64 **39.66**
ogbl-collab 13,835 20,664 91,807 OOM 34,934 21,057 15,716 **13,493** 45.33 46.12 27.90 OOM 24.80 34.07 45.44 **48.36**
DBLP **31,138 40,744** 148,439 OOM 94,384 44,424 36,463 31,686 38.26 40.92 20.21 OOM 15.96 27.82 38.99 **41.66**


Table 1: Hierarchical clustering results (n[′] = 512). Bold/underline indicate best/second best scores.

FPH Avg. link. RGHC HGHC HypHC UF Ward

500 60

400

50 25

400

40

Dasgupta 300 20

Norm. TSD

300 30

15

3264 128 256 512 3264 128 256 512 3264 128 256 512 3264 128 256 512

Num. internal nodes Num. internal nodes Num. internal nodes Num. internal nodes


(a) Dasgupta - Cora-ML


(b) Dasgupta - PolBlogs


(c) TSD - Cora-ML


(d) TSD - PolBlogs


Figure 2: Hierarchical clustering results measured by Dasgupta cost (lower is better) and TSD (higher
is better). More results in Fig. 3 in the appendix.

4 EXPERIMENTS

**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 6, 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 Gradient-**
based 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.

4.1 RESULTS

**Hierarchical clustering. We report results on hierarchical clustering with n[′]** = 512 internal nodes
in Table 1. FPH outperforms all baselines on all datasets on TSD and on all except one dataset on the
Dasgupta cost. 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 11. Since FPH is deterministic, we only run it once.


-----

Model Citeseer Cora Polblogs ogbn-arxiv DBLP

Avg. 0.367 0.420 0.507 0.216 0.526
RGHC 0.281 0.400 **0.730** 0.286 0.510
HGHC 0.365 0.379 0.177 0.290 0.408
Ward 0.368 0.504 0.702 **0.411** 0.591
UF 0.347 0.428 0.676 0.254 0.598

HypHC 0.270 0.121 0.691 OOM OOM
Louvain 0.329 0.500 0.640 0.395 0.558
FPH **0.398** 0.462 0.680 0.251 0.560
FPH (Louv.) 0.380 0.507 0.614 0.399 0.564
FPH (Ward) 0.393 **0.516** 0.708 0.401 **0.604**

Table 2: NMI results on real-world datasets. FPH (Louv.) and FPH (Ward) refer to FPH initialized
from the solutions of Louvain and Ward, respectively.

Moreover, we noticed that HypHC requires an excessive amount of triplet samples to obtain good
results. Even n[2], 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 n[2] 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. 8, 9
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 real-**
world 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 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 block-**
model (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 10 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 3. FPH outperforms all baselines
TSD and is competitive on Dasgupta, with both n[′]=512 and 1024. Note that we could not compare


-----

**Dataset** **FPH** **DCSBM** **DW** **Ad./Ad.** **VGAE**

Cora-ML 95.7 95.5 94.3 86.5 **95.9**
Citeseer **96.2** 93.6 96.0 76.8 94.8
PolBlogs 94.3 **94.9** 84.8 92.6 92.8
WikiPhysics **97.2** 96.9 92.9 96.6 97.0

Brain 94.1 **95.2** 83.8 90.7 93.2
OpenFlight **99.3** 99.0 94.3 98.4 99.0

Genes 69.8 66.9 **70.3** 53.0 66.6


Dasgupta Norm. TSD
_n[′]_ 512 1024 512 1024

Ward **144,157** **127,968** 37.05 40.59
HGHC 219,959 168,851 35.90 39.53
RGHC 184,688 -  35.35 - 
Avg. link. 175,571 168,753 40.33 43.21

FPH 147,169 142,404 **43.79** **45.57**

Table 3: Results on ogbn-products.


Table 4: AUC-PR score (%) for link prediction.


with HypHC, since (a) the implementation constructs a dense n × n matrix, which would require
more than 14TB memory assuming single precision; (b) sampling n[2] 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 4, 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.

Root Level 1 Level 2

**Qualitative analysis. In Table 5 we provide insight** asocs, generalization, dynamic

neural,

into our model trained on the Cora-ML dataset. We paper, genetic, reinforcement, search

data, models, training

show the first three levels of internal nodes of the networks reasoning, knowledge, planning
learned hierarchy. Each cell corresponds to an inter- chain, markov, series, time, bayesian
nal node and contains the three most frequent words learning, sampler gibbs, distribution, mcmc

algorithm, dimacs, sum, binary, comparison

of the abstracts of all papers (i.e., leaf nodes) as- problem evolutionary, perfect, graphs, parameterized
signed to the respective internal nodes. Words from species technical, report, polynomial
a node are excluded from its children to avoid dupli- stability, control, gain, output
cates. Terms tend from more general at the root (e.g., systems,linear proved, trajectory, boundedlyapnov, state, online
‘problem’) towards more specific at the lower levels wavelet, minimax,
(e.g., ‘mcmc’). Further, we can identify categories estimation
of machine learning approaches, e.g., reinforcement

Table 5: Cora-ML hierarchy visualization.

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.

5 CONCLUSION

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 relaxed versions of 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.


-----

ETHICS STATEMENT

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.

REPRODUCIBILITY STATEMENT

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 7 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] Our implementation is available at
[https://www.daml.in.tum.de/fph](https://www.daml.in.tum.de/fph)

REFERENCES

Lada A. Adamic and Eytan Adar. Friends and neighbors on the web. SOCIAL NETWORKS, 2001.

Lada A. Adamic and Natalie Glance. The political blogosphere and the 2004 u.s. election: Divided
they blog. In LinkKDD, 2005.

Katrin Amunts, Claude Lepage, Louis Borgeat, Hartmut Mohlberg, Timo Dickscheid, Marc-Étienne
Rousseau, Sebastian Bludau, Pierre-Louis Bazin, Lindsay Lewis, Ana-Maria Oros-Peusquens,
N. Shah, Thomas Lippert, Karl Zilles, and Alan Evans. Bigbrain: An ultrahigh-resolution 3d
human brain model. Science, 2013.

Nicolas Aspert, Volodymyr Miz, Benjamin Ricaud, and Pierre Vandergheynst. A graph-structured
dataset for wikipedia research. In WWW, 2019.

Albert-László Barabási and Márton Pósfai. Network Science. Cambridge University Press, Cambridge,
1 edition, 2016. ISBN 978-1-107-07626-6.

Harold P. Benson. Concave Minimization: Theory, Applications and Algorithms, pp. 43–148. Springer
US, Boston, MA, 1995. ISBN 978-1-4615-2025-2. doi: 10.1007/978-1-4615-2025-2_3. URL
[https://doi.org/10.1007/978-1-4615-2025-2_3.](https://doi.org/10.1007/978-1-4615-2025-2_3)

Jean-Paul Benzécri. Fast hierarchical clustering using reciprocal nearest-neighbor chain algorithm.
_Notebook of the data analysis, 1982._

Vincent D Blondel, Jean-Loup Guillaume, Renaud Lambiotte, and Etienne Lefebvre. Fast unfolding
of communities in large networks. Journal of statistical mechanics: theory and experiment, 2008
(10):P10008, 2008.

Aleksandar Bojchevski and Stephan Günnemann. Bayesian robust attributed graph clustering: Joint
learning of partial anomalies and group structure. In AAAI, 2018.

[2https://scikit-network.readthedocs.io/en/latest/](https://scikit-network.readthedocs.io/en/latest/)


-----

Thomas Bonald, Bertrand Charpentier, Alexis Galland, and Alexandre Hollocou. Hierarchical graph
clustering using node pair sampling. KDD Workshop MLG, 2018.

Ines Chami, Albert Gu, Vaggos Chatziafratis, and Christopher Ré. From trees to continuous embeddings and back: Hyperbolic hierarchical clustering. In NeurIPS, 2020.

Moses Charikar and Vaggos Chatziafratis. Approximate Hierarchical Clustering via Sparsest Cut and
Spreading Metrics. arXiv e-prints, 2016.

Bertrand Charpentier. Multi-scale clustering in graphs using modularity. DiVA, 2019.

Bertrand Charpentier and Thomas Bonald. Tree sampling divergence: An information-theoretic
metric for hierarchical graph clustering. In IJCAI, 2019.

Giovanni Chierchia and Benjamin Perret. Ultrametric fitting by gradient descent. In NeurIPS. 2019.

Ara Cho, Junha Shin, Sohyun Hwang, Chanyoung Kim, Hongseok Shim, Hyojin Kim, Hanhae Kim,
and Insuk Lee. Wormnet v3: A network-assisted hypothesis-generating server for caenorhabditis
elegans. Nucleic acids research, 2014.

Sanjoy Dasgupta. A cost function for similarity-based hierarchical clustering. In ACM Symposium
_on Theory of Computing, pp. 118–127, 2016._

Michael B. Eisen, Paul T. Spellman, Patrick O. Brown, and David Botstein. Cluster analysis and
display of genome-wide expression patterns. Proceedings of the National Academy of Sciences,
[95(25):14863–14868, 1998. ISSN 0027-8424. URL https://www.pnas.org/content/](https://www.pnas.org/content/95/25/14863)
[95/25/14863.](https://www.pnas.org/content/95/25/14863)

J. C. Gower and G. J. S. Ross. Minimum spanning trees and single linkage cluster analysis. Journal
_of the Royal Statistical Society: Series C, 1969._

Johan Himberg, Aapo Hyvärinen, and Fabrizio Esposito. Validating the independent components
of neuroimaging time series via clustering and visualization. NeuroImage, 22(3):1214–1222,
[July 2004. doi: 10.1016/j.neuroimage.2004.03.027. URL https://doi.org/10.1016/j.](https://doi.org/10.1016/j.neuroimage.2004.03.027)
[neuroimage.2004.03.027.](https://doi.org/10.1016/j.neuroimage.2004.03.027)

Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta,
and Jure Leskovec. Open Graph Benchmark: Datasets for Machine Learning on Graphs. In
_NeurIPS, 2020._

N. Jardine and R. Sibson. The Construction of Hierarchic and Non-Hierarchic Classifications. The
_Computer Journal, 1968._

Sepandar D. Kamvar, Dan Klein, and Christopher D. Manning. Interpreting and Extending Classical
Agglomerative Clustering Algorithms using a Model-Based approach. In ICML, 2002.

Brian Karrer and M.E.J. Newman. Stochastic blockmodels and community structure in networks.
_Physical review. E, 2011._

Thomas N Kipf and Max Welling. Variational graph auto-encoders. NIPS Workshop on Bayesian
_Deep Learning, 2016._

Vince Lyzinski, Minh Tang, Avanti Athreya, Youngser Park, and Carey E Priebe. Community
detection and classification in hierarchical stochastic blockmodels. IEEE Transactions on Network
_Science and Engineering, 4(1):13–26, 2016._

Andrew Kachites McCallum, Kamal Nigam, Jason Rennie, and Kristie Seymore. Automating the
construction of internet portals with machine learning. Inf. Retr., 2000.

Nicholas Monath, Manzil Zaheer, Daniel Silva, Andrew McCallum, and Amr Ahmed. Gradient-based
hierarchical clustering. NeurIPS Workshop DISCML, 2017.

Nicholas Monath, Manzil Zaheer, Daniel Silva, Andrew McCallum, and Amr Ahmed. Gradient-based
hierarchical clustering using continuous representations of trees in hyperbolic space. In KDD,
2019.


-----

Benjamin Moseley and Joshua Wang. Approximation bounds for hierarchical clustering: Average
linkage, bisecting k-means, and local search. In NeurIPS. 2017.

Fionn Murtagh and Pedro Contreras. Algorithms for hierarchical clustering: an overview. WIREs
_Data Mining and Knowledge Discovery, 2012._

V. Niculae. Optimizing with constraints: reparametrization and geometry. 2020.

Jani Patokallio. Openflight. online https://openflights.org.

Bryan Perozzi, Rami Al-Rfou, and Steven Skiena. Deepwalk: Online learning of social representations. In KDD, 2014.

Erzsébet Ravasz and Albert-László Barabási. Hierarchical organization in complex networks. Physical
_review E, 67(2):026112, 2003._

Prithviraj Sen, Galileo Namata, Mustafa Bilgic, Lise Getoor, Brian Galligher, and Tina Eliassi-Rad.
Collective classification in network data. 2008.

Michael Steinbach, George Karypis, and Vipin Kumar. A comparison of document clustering
techniques. KDD Workshop TM, 2000.

Robert L. Thorndike. Who belongs in the family? Psychometrika, 18(4):267–276, Dec 1953. ISSN
[1860-0980. doi: 10.1007/BF02289263. URL https://doi.org/10.1007/BF02289263.](https://doi.org/10.1007/BF02289263)

Kuansan Wang, Zhihong Shen, Chiyuan Huang, Chieh-Han Wu, Yuxiao Dong, and Anshul Kanakia.
Microsoft academic graph: When experts are not enough. Quantitative Science Studies, 1(1):
396–413, 2020.

Joe H. Ward Jr. Hierarchical grouping to optimize an objective function. Journal of the American
_Statistical Association, 1963._

Michael Wick, Sameer Singh, and Andrew McCallum. A discriminative hierarchical model for
fast coreference at large scale. In Proceedings of the 50th Annual Meeting of the Association
_for Computational Linguistics (Volume 1: Long Papers), pp. 379–388, Jeju Island, Korea, July_
[2012. Association for Computational Linguistics. URL https://aclanthology.org/](https://aclanthology.org/P12-1040)
[P12-1040.](https://aclanthology.org/P12-1040)

Jaewon Yang and Jure Leskovec. Defining and evaluating network communities based on ground-truth.
_Knowledge and Information Systems, 42(1):181–213, 2015._

Yuchen Zhang, Amr Ahmed, Vanja Josifovski, and Alexander J Smola. Taxonomy discovery for
personalized recommendation. In ACM International Conference on Web Search And Data Mining
_(WSDM), 2014._

Ying Zhao and George Karypis. Evaluation of hierarchical clustering algorithms for document
datasets. In Proceedings of the Eleventh International Conference on Information and Knowledge
_Management, CIKM ’02, pp. 515–524, New York, NY, USA, 2002. Association for Computing_
[Machinery. ISBN 1581134924. doi: 10.1145/584792.584877. URL https://doi.org/10.](https://doi.org/10.1145/584792.584877)
[1145/584792.584877.](https://doi.org/10.1145/584792.584877)


-----

A PROOFS AND THEORETICAL ANALYSIS

A.1 CONTINUOUS VERSIONS OF TSD AND DASGUPTA

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 relaxed versions of hierarchical
clustering quality metrics in an end-to-end fashion instead of proxy losses or heuristic algorithms.

Recall the equation of the discrete Dasgupta cost, Das( T[ˆ] ) = vi,vj _∈V_ _[P]_ [(][v][i][,][ v][j][)][P]z [I][[][z][=][v]i[∧][v]j []][c][(][z][)][.]

_c(z) is the number of leaves for which internal node z is an ancestor, i.e., c(z) =_ v∈V [I][[][z][∈][anc][(][v][)]][.]

Thus, we get [P]

Das( [ˆ] ) = _P_ (vi, vj) I[z=vi vj ] I[z anc(v)]. [P]
_T_ _∧_ _∈_

vi,Xvj _∈V_ Xz vX∈V

We propose the soft Dasgupta cost (Soft-Das) by replacing the indicators I[ ] with their expectations,

_·_
i.e., with their associated probabilities:


Soft-Das(PA,B(T )) = vi,Xvj _∈V_ _P_ (vi, vj) Xz _p(z = vi ∧_ vj) Xv _panc(z|v)_

Intuitively, Soft-Das is the Dasgupta cost of the expected hierarchy obtained via sampling from
**_A, B, i.e., Soft-Das(PA,B(_** )) Das(E ˆ **_PA,B_** ( )[[ ˆ] ]). Ideally, we would like to optimize
_T_ _≡_ _T ∼_ _T_ _T_
the expected Dasgupta cost when sampling discrete hierarchies, i.e., Exp-Das(PA,B( )) =
_T_

E ˆ **_PA,B_** ( ) Das( [ˆ] ), i.e., to put the expectation outside the Das function, which is nontrivial. We
_T ∼_ _T_ _T_
leave this extension for future work.h i

In the same way, we propose a differentiable version of the soft Tree Sampling Divergence (Soft**TSD). To this end, we replace the discrete assignments in the distributions p(z) and q(z) with**
probabilistic assignments, i.e.


_p(z) log_ _[p][(][z][)]_

_q(z)_


Soft-TSD(PA,B( )) = KL(p(z) _q(z)) =_
_T_ _||_


where p(z) =

_q(z) =_


_p(z = vi_ vj)P (vi, vj)
vi,vj _∧_

X

_p(z = vi_ vj)P (vi)P (vj)
vi,vj _∧_

X


Analogously to above, Soft-TSD effectively computes the TSD score of the expected hierarchy
obtained via sampling from A, B. Extending our model to compute Exp-TSD(PA,B( )) =
_T_

E ˆ **_PA,B_** ( ) TSD( [ˆ] ) is left for future work. Note that for both Soft-TSD and Soft-Das we
_T ∼_ _T_ _T_
recover the same score as their discrete formulations in the case of a deterministic probabilistic model,h i
i.e. when A and B are binary-valued.

A.2 TREE-SAMPLING PROCEDURE

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 = (A[ˆ], **_B[ˆ])_** _PA,B(_ ). We can also compute the probability of any
_T_ _∼_ _T_
tree T under the sampling procedure described above:


_AAi,jˆi,j_
_i,j_

Y


_BiBˆ[′],ji′[′],j′_ (10)
_i[′],j[′]_

Y


_PA,B(_ = (A[ˆ], **_B[ˆ])) =_**
_T_


-----

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.

A.3 PROOF OF THEOREM 1

_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 on must lead to state zn′. From state zn′, the random walk transits to ω with probability
_M_
**_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 therefore M is acyclic.

A.4 PROOF OF THEOREM 2

_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 ˆri is generated in the random walk W(vi), completing the proof.

A.5 PROOF OF THEOREM 3

_Proof. First, recall that all paths end in the root node, such that ri necessarily ends in zn′ = ri[(][T][ )]._
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 ̸= 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 = (ri[(1)], . . ., vi vj) the first part of the path ri until (and including) its

lowest common ancestor with rj, i.e., v ∧i vj. Analogously, ri = (ri(|ri|+1), . . ., 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


_|rj_ _|_

_p(rj[(][t][)]_ _j_ )
_t=2_ _[|][r][(][t][−][1)]_ _·_

Y


_|ri|_

_p(ri[(][t][)]_ _i_ ) _p(rj[(1)]_
_t=2_ _[|][r][(][t][−][1)]_ _·_ _[|][v][j][)][ ·]_

Y


_|ri|_ _|rj_ _|_ _|ri|_

_p[(][T][ )]((ri, rj)) = p(ri[(1)]_ vi) _p(ri[(][t][)]_ _i_ ) _p(rj[(1)]_ _p(rj[(][t][)]_ _j_ ) _p(ri[(][t][)]_ _i_
_|_ _·_ _tY=2_ _[|][r][(][t][−][1)]_ _·_ _[|][v][j][)][ ·]_ _tY=2_ _[|][r][(][t][−][1)]_ _·_ _t=Y|ri|+1_ _[|][r][(][t][−][1)]_

(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)) = p(ri) _p(rj) i.e., the paths_
_̸_ _·_
_ri and rj are not independent. We denote p[(]anc[T][ )][(][z]k[|][v]i[,][ v]j[)][ the probability of the internal node][ z]k_ [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[(]anc[T][ )][(][z]k[|][v]i[,][ v]j[)][ ̸][=][ p]anc[(][z]k[|][v]i[)] _[·]_ _[p]anc[(][z]k[|][v]j[)][. 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.

A.6 PROOF OF THEOREM 4

_Proof. We start by reorganizing Eq. (5):_

(iii)

(i) (ii)

_k−1_ (13)

_panc(zk_ vi)panc(zk vj) = _p[(][M][)](zk = vi_ vj) + _p[(][M][)](zk′ = vi_ vj)panc(zk zk′ )[2]
_|_ _|_ _∧_ zk[′]=1 }| ∧ _|_ {
z }| { z }| { X

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[(]anc[M][)][(][z]k[|][v]i[,][ v]j[) =][ p]anc[(][z]k[|][v]i[)][p]anc[(][z]k[|][v]j[)][ and therefore (i) is the probability of][ z]k
being an ancestor of vi and vj in our Markov chain M.

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 thatproperty of the random walks. Further, note that zk′ = vi ∧ vj implies zk′ ∈ anc (vi, vj) as well as the Markov

_p[(]anc[M][)][(][z][k][|][z][k][′][ ∈]_ [anc][(][v][i][,][ v][j][)) =][ p][anc][(][z][k][|][z][k][′][ ∈] [anc][(][v][i][))][ ·][ p][anc][(][z][k][|][z][k][′][ ∈] [anc][(][v][j][))]

(15)
= panc(zk zk′ )[2],
_|_

where we first exploit factorization due to independence and in the last step again the Markov property
of the random walks.

A.7 PROOF OF THEOREM 5

_Proof.respectively. Then, Let ˆri ∈P(vi), ˆrj ∈P(vj) be two independent random walks on M rooted in vi and vj,_
_p[(][M][)](zk = vi_ vj) = _p((ˆri, ˆrj)),_

(16)

_∧_

(ˆri,rˆj ):Xzk=vi∧vj

where ˆri = (ˆri[(1)], . . ., zk) is the first part of ˆri until it reaches zk. Note that the second part of the
paths, ˆri and ˆrj, 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 ˆri and ˆrj under both models (see
Eq. (4)), we can write:


_p[(][M][)](zk = vi ∧_ vj) = _p((ˆri, ˆrj))_

(ˆri,rˆj ):Xzk=vi∧vj

= _p(ˆri)_ _p(ˆrj)_

_·_
(ˆri,rˆj ):Xzk=vi∧vj

= p[(][T][ )](zk = vi ∧ vj).


(17)


-----

A.8 NUMBER OF PAIRS OF LCA PATHS.

**Theorem 7. Let M 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 3[k][−][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 3[k][−][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 · 3[k][−][1] = 3[k].

A.9 PROOF OF THEOREM 6

We provide here the proof for the fast vectorized computation of Pv[LCA]i,vj [in Theorem. 6.]

_Proof. For the case vi_ = vj, we start by reorganizing Eq. (9):
_̸_

**_Pv[anc]i_** _[⊙]_ **_[P][ anc]vj_** = Pv[LCA]i,vj [+][ P][ LCA]vi,vj[,T] _·_ **_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:


_n[′]_

_k[′]=1_ **_Pv[LCA]i,vj_** _,vk′_ _[·]_ **_P˜_** [anc] _⊙_ **_P[˜]_** [anc][i]k[′],k _[.]_

X h


**_Pv[anc]i_** vj

_[⊙]_ **_[P][ anc]_**


_k_ [=][ P][ LCA]vi,vj _,vk_ [+]


Plugging in the definitions of Eq. (8), using Theorem 2, and observing that due to the upper triangular
structure **_P[˜]k[anc][′],k_** [= 0][ for][ k][′][ > k][ we obtain]


_k−1_

_p(zk′ = vi_ vj) _panc(zk_ zk′ )[2].
_k[′]=1_ _∧_ _·_ _|_

X


_panc(zk|vi) · panc(zk|vj) =p(zk = vi ∧_ vj) +


For vi = vj, we have Pv[LCA]i,vi [=][ A][v]i [, i.e., again simply the parent probabilities of v][i][.]

A.10 VECTORIZED COMPUTATIONS

All quantities involved in Soft-Das and Soft-TSD can be computed in closed-form based on the
Markov chain M 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] _∈_ R[n][×][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 R[m][×][n][′] . Thus, we can exploit the sparsity of real world graphs, as typically m ≪ _n[2]._

In a similar way, the computation of the distribution q(z) does also not require the expensive explicit
computation ofthe Markov chain p( Mzk =. First, observe that the equation of vi ∧ vj) for all pairs of leaf nodes. Instead, we can again exploit insights from q(z) in Soft-TSD describes an expectation:

_q(z) = Evi,vj_ _P (v) [p(z = vi_ vj)] . (19)
_∼_ _∧_

Defining ˆp[anc] = p[T] _· 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]_** _−_ (p ⊙ **_p)[T]_** _· P_ [anc] _⊙_ **_P_** [anc]) · (I + P[˜] [anc] _⊙_ **_P[˜]_** [anc])[−][1] + (p ⊙ **_p)[T]_** **_A._**


-----

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:_


**_pvi_** **_Pv[LCA]i,vj_** **_[p][v]j_**
_i,j_

X


**_q =_**


where we denoteinto the Hadamard product which is done over the internal node dimension. P (vi) = pvi. Subsequently, we can plug the P [LCA] formula Eq. (9) and pull pvi


_i,j_ **_pvi_** **_pvj_** (Pv[anc]i _[⊙]_ **_[P][ anc]vj_** [)][T][ ·][ (][I][ + ˜]P [anc] _⊙_ **_P[˜]_** [anc])[−][1]

X


**_qˆ =_**



_· (I + P[˜]_ [anc] _⊙_ **_P[˜]_** [anc])[−][1]

_· (I + P[˜]_ [anc] _⊙_ **_P[˜]_** [anc])[−][1]


**_Pv[anc]i_**



 _[⊙]_


**_pvj_** **_Pv[anc]j_**

**_pvj_** **_Pv[anc]j_**


**_pvi_**


= **_pvi_** **_Pv[anc]i_** **_pvj_** **_Pv[anc]j_** (I + P[˜] [anc] **_P_** [anc])[−][1]

 _i_ _[⊙]_ _j_  _·_ _⊙_ [˜]

X

[X] 

= (ˆp[anc] _⊙_ **_pˆ[anc])[T]_** _· (I + P[˜]_ [anc] _⊙_ **_P[˜]_** [anc])[−][1].

where ˆp[anc] = p[T] _· P_ [anc]. However, recall that for vi = vj the LCA probabilities are Avi; thus, we
need to correct for this difference to obtain q:


**_qˆ =_** **_pvi_** **_Pv[LCA]i,vj_** **_[p][v]j_** [+]
Xi≠ _j_ Xi

= q − **_pvi_** **_Pv[LCA]i,vi_** **_[p][v]i_** [+]

_i_

X

_⇔_ **_q = ˆq +_** _i_ **_pvi_** **_Pv[LCA]i,vi_** **_[p][v]i_** _[−]_
X


**_pvi_** **_P[ˆ]v[LCA]i,vi_** **_[p][v]i_**

**_pvi_** **_P[ˆ]v[LCA]i,vi_** **_[p][v]i_**
_i_

X

**_pvi_** **_P[ˆ]v[LCA]i,vi_** **_[p][v]i_** [= ˆ]q +

X


**_pvi_** **_P[ˆ]v[LCA]i,vi_** **_[p][v]i_**


**_pvi_** **_Avi_** **_pvi_**
_−_


= ˆq + (p **_p)[T]_** **_A_** (p **_p)[T][ ˆ]Pv[LCA]i,vi_**
_⊙_ _·_ _−_ _⊙_
= ((ˆp[anc] _⊙_ **_pˆ[anc])[T]_** _−_ (p ⊙ **_p)[T]_** _· P_ [anc] _⊙_ **_P_** [anc]) · (I + P[˜] [anc] _⊙_ **_P[˜]_** [anc])[−][1] + (p ⊙ **_p)[T]_** **_A,_**

where **_P[ˆ]v[LCA]i,vi_** [=][ P][ anc][ ⊙] **_[P][ anc][ ·][ (][I][ + ˜]P_** [anc] _⊙_ **_P[˜]_** [anc])[−][1] (obtained by using Eq. equation 9).

Note that, unless stated otherwise, we use ˆq as a proxy for q in our experiments for simplicity.

A.12 COMPLEXITY ANALYSIS.

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) ∈_** R[n][′][×][n][′], plus matrix multiplication
with A ∈ R[n][×][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,
where m 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._

A.13 CONVEXITY OF SOFT-TSD

**Theorem 8.LCA probabilities of internal nodes w.r.t. pairs of leaf nodes. Let H ∈** [0, 1][n][′][×][n][×][n] _be a tensor whose elements Soft-TSD Hkij =(H p) is convex in(zk = vi ∧_ **_Hvj). are the_**


-----

_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(αHk[(1)]_ + (1 − _α)Hk[(2)][) =]_


_P_ (vi, vj)(αHkij[(1)] [+ (1][ −] _[α][)][H]kij[(2)][)]_
vi,vj

X


_P_ (vi, vj)Hkij[(1)] [+ (1][ −] _[α][)]_
vi,vj

X


_P_ (vi, vj)Hkij[(2)]
vi,vj

X


(20)


= α


= α · p(Hk[(1)][) + (1][ −] _[α][)][ ·][ p][(][H]k[(2)][)][.]_

Analogously we compute the distribution q induced by the independent node distribution. We have
that
_q(αHk[(1)]_ + (1 − _α)Hk[(2)][) =][ α][ ·][ q][(][H]k[(1)][) + (1][ −]_ _[α][)][ ·][ q][(][H]k[(2)][)][.]_ (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. Let HFPH = {H : ∃A, B ∈** Φ(n, n[′]) : H = FPH(A, B)} denote the set of probabilis_tic 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,are the LCA probabilities of internal nodes w.r.t. pairs of leaf nodes. This set H HkijFPH = is convex. p(zk = vi ∧_ vj)

_Proof. We start by recalling form Theorems 4 and 5 that FPH computes lowest common ancestor_
probabilitiesLCA probabilities are consistent with the expected result from the tree-sampling procedure described p(zk = vi ∧ vj) from the continuous parent probability matrices A and B such that the
in Sec. A.2. More formally,


**_Hkij := p(zk = vi_** vj) = E ˆH (A,B) [[][I][[][z][k][ =][ v][i][ ∧] [v][j][]]]
_∧_ _∼_

= E ˆH∼(A,B) **_Hˆ_** _kij_ = E **_Hˆ_**
h i h


(23)
_kij_ _[,]_


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 = E ˆH_** (A,B) **_Hˆ_** = _p( H[ˆ]_ **_A, B)_** **_H[ˆ]_** _,_ (24)
_∼_ _|_ _·_
h i **_Hˆ_** _∈HX(n,n[′])_

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.


-----

**Dataset** **Nodes (LCC)** **Edges (LCC)** **MI (LCC)** License

_PolBlogs, (Adamic & Glance, 2005)_ 1,222 16,715 2.39 n/a
_Brain, Amunts et al. (2013)_ 1,770 8,957 3.37 n/a
_Citeseer, Sen et al. (2008)_ 2,110 3,694 5.69 n/a
_Genes, Cho et al. (2014)_ 2,194 2,688 6.12 n/a
_Cora-ML, McCallum et al. (2000); Bojchevski & Günnemann (2018)_ 2,810 7,981 5.23 n/a
_WikiPhysics, Aspert et al. (2019)_ 3,309 31,251 3.44 n/a
_OpenFlight, Patokallio,_ 3,097 18,193 3.44 ODbL
_Ogbn-products, Hu et al. (2020)_ 2,385,902 61,806,367 9.29 Amazon license
_Ogbn-arxiv, Hu et al. (2020); Wang et al. (2020)_ 169,343 1,157,799 7,40 ODC-BY
_Ogbl-collab, Hu et al. (2020); Wang et al. (2020)_ 232,865 961,883 9.02 ODC-BY
_DBLP, Yang & Leskovec (2015)_ 317,080 1,049,866 9.64 n/a

Table 6: Dataset summary; we convert directed datasets to undirected and select the largest connected
component (LCC).

**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.

B EXPERIMENT INFORMATION

B.1 LINK PREDICTION WITH SOFT-TSD

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][(][v][i][ ∧]_ [v][j][)] (25)

_q(vi_ vj)
_∧_

B.2 LINK PREDICTION SETUP

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 documentation[3] 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.

B.3 DATASET SUMMARY

See Table 6 for an overview of the datasets we used.

[3https://graph-tool.skewed.de/static/doc/demos/inference/inference.html](https://graph-tool.skewed.de/static/doc/demos/inference/inference.html)


-----

600

500

400

|FPH Avg. link. RGHC HGHC HypHC UF Ward 800 30 TSD 700 50 600 Norm. 40 20 500|Col2|Col3|Col4|Col5|Col6|Col7|Col8|Col9|Col10|Col11|
|---|---|---|---|---|---|---|---|---|---|---|
||||||||||||
||||||||||||
||||||||||||
||||||||||||
||||||||||||


Num. internal nodes

(a) Dasgupta - OpenFlight


Num. internal nodes

(b) Dasgupta - Brain


Num. internal nodes

(c) TSD - OpenFlight


Num. internal nodes

(d) TSD - Brain


Figure 3: Results on hierarchical clustering measured by Dasgupta cost (lower is better) and TSD
(higher is better).

50


40

30

|Col1|Col2|Col3|Col4|Col5|Col6|
|---|---|---|---|---|---|
|||||||
|||||||
|||||||
|32 64 128 256 512 Num. internal nodes||||||


FPH FPH-R FPH-R Emb.

FPH Emb. DW FPH-U FPH-UR

800

Dasgupta

600

3264 128 256 512 3264 128 256 512

Num. internal nodes Num. internal nodes


(a) TSD - WikiPhysics


(b) Dasgupta - WikiPhysics


Figure 4: Ablation study results.

B.4 ADDITIONAL RESULT FIGURES


In Fig. 3 we show results for four more datasets.

B.5 ABLATION STUDY


See Fig 4 for the comparison of different FPH model variants, and our full discussion in Sec. 4.1.

B.6 HIERARCHY VISUALIZATION


(a) Dasgupta - 10 clusters (b) Dasgupta - 50 clusters (c) Das. - Dendrogram


(d) TSD - 10 clusters (e) TSD - 50 clusters (f) TSD - Dendrogram

Figure 5: Visual comparison of trees obtained after Soft-TSD and Soft-Dasgupta optimization on
OpenFlight.


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).

B.7 HYPERPARAMETERS

We use the hyperparameters for models and baselines described in Tab. 7. For smaller datasets, we
train FPH for 1,000 epochs and restore the best hierarchy after training. For ogbn-products, ogbnarxiv, ogbl-collab, and DBLP, we train for 2,000 epochs. Similarly, we use different learning rates

Model Hyperparameter Value

Learning rate 150
FPH (TSD) Batch size K** 10,000
Batch cutoff C** 200,000

FPH-R (TSD) Learning rate 200

Learning rate 0.05
FPH (Das.) Batch size K** 10,000
Batch cutoff C** 200,000

FPH Emb. Learning rate 0.1

Routing NN dim 128
RGHC Iterations 5000
Learning rate 0.0001

Init. method K-means + agglom. linkage
HGHC Iterations 10
Learning rate 0.1

Loss Closest + cluster size
UF Epochs 500
Learning rate 0.1


Num. triples 50M
Epochs 50
Learning rate 0.001
Temperature 0.1


HypHC


Embedding dim 10
DeepWalk
Embedding dim* 32

-  Used for ogbn-products, ogbn-arxiv, ogbl-collab, DBLP.
** Used for ogbn-products

Table 7: Hyperparameter settings.

for A and B for FPH (Das.) on ogbn-arxiv, ogbl-collab, ogbn-products, and DBLP (lrA = 1e − 2,
lrB = 1e − 9).

B.8 COMPUTING INFRASTRUCTURE

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.

B.9 HSBM GRAPHS

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


-----

Figure 6: Example HSBM graph with n = 100, n[′] = 7, and three levels in the hierarchy.

Figure 7: Example HSBM graph with n = 1000, n[′] = 53, and three levels in the hierarchy.

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

Dasgupta cost (lower is better) Normalized TSD (higher is better)
Alg. Ward UF HypHC HGHC RGHC Ward UF HypHC HGHC RGHC

Brain 596.73 938.49 568.18 894.87 650.53 32.43 26.28 17.68 17.31 16.54
OpenFlight 416.05 643.45 423.80 477.51 469.98 55.59 49.88 40.06 47.52 45.84
Genes 221.76 258.21 467.12 482.11 444.82 66.87 63.73 23.94 50.59 40.98
Citeseer 105.12 280.66 271.80 224.63 200.91 69.28 62.95 31.74 52.53 47.66
Cora-ML 301.47 673.27 441.21 516.87 499.49 57.22 47.96 29.10 42.10 35.19
PolBlogs 383.51 726.34 334.69 428.52 376.94 27.01 10.73 21.60 20.30 20.78
WikiPhysics 808.87 958.20 701.14 919.27 790.39 45.54 41.55 33.85 34.51 36.51
ogbn-arxiv 22,046 64,950 OOM 37,177 26,286 37.43 26.22 OOM 17.55 25.20
ogbl-collab 14,834 101,562 OOM 112,048 17,964 45.20 30.50 OOM 11.11 37.73
DBLP 33,349 160,742 OOM 171,975 41,796 38.87 22.62 OOM 5.61 29.9

Table 8: Hierarchical clustering results (n[′] = 512, d = 128).


-----

Model Citeseer Cora Polblogs DBLP ogbn-arxiv

RGHC 0.218 0.394 0.756 0.510 0.358
HGHC 0.304 0.362 0.604 0.655 0.385
Ward 0.363 0.445 0.436 0.587 0.402
UF 0.180 0.242 0.102 0.395 0.143
HypHC 0.285 0.390 0.740 -  - 

Table 9: NMI results for d = 128 DeepWalk embeddings.

_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**

Table 10: NMI results on synthetic HSBM graphs.

TSD standard deviation Dasgupta standard deviation
Alg. HypHC HGHC RGHC HypHC HGHC RGHC

Brain 0.53 (3%) 0.06 (<0.5%) 0.6 (3%) 19.96 (3%) 38.56 (5%) 14.18 (3%)
OpenFlight 1.15 (3%) 0.4 (1%) 0.7 (2%) 26.61 (6%) 31.14 (6%) 23.9 (5%)
Genes 0.61 (3%) 0.12 (<0.5%) 1.19 (2%) 13.3 (3%) 2.88 (1%) 17.31 (7%)
Citeseer 0.19 (1%) 0.09 (<0.5%) 1.29 (3%) 6.93 (3%) 5.56 (4%) 8.02 (6%)
Cora-ML 0.96 (3%) 0.13 (<0.5%) 0.61 (1%) 17.39 (4%) 21.78 (5%) 12.98 (4%)
PolBlogs 0.13 (1%) 0.09 (<0.5%) 0.18 (1%) 4.04 (1%) 1.92 (1%) 3.4 (1%)
WikiPhysics 0.47 (1%) 0.14 (<0.5%) 0.56 (1%) 5.77 (1%) 6.83 (1%) 30.88 (4%)
ogbn-arxiv -  0.15 (1%) 1.31 (5%) -  405 (2%) 765 (3%)
ogbl-collab -  0.16 (1%) 0.61 (2%) -  1,592 (5%) 625 (3%)
DBLP -  0.47 (3%) 0.81 (3%) -  2,482 (3%) 1,005 (2%)

Table 11: Standard deviations of non-deterministic baselines. In parentheses we report the standard
deviation in relation to the best value reported in Table 1 in percent.


-----

