Tree Variational Autoencoders
Authors: Laura Manduchi
Moritz Vandenhirtz
Alain Ryser
Julia E Vogt
Venue: NeurIPS 2023 spotlight (2023)
Keywords: hierarchical clustering
hierarchical VAE
representation learning
VAE
deep clustering
TL;DR: We propose a new generative hierarchical clustering model that learns a flexible tree-based posterior distribution over latent variables.

Abstract: We propose Tree Variational Autoencoder (TreeVAE), a new generative hierarchical clustering model
  that learns a flexible tree-based posterior distribution over latent variables. TreeVAE hierarchically divides samples according to their intrinsic characteristics, shedding light on hidden structures in the data. It adapts its architecture to discover the optimal tree for encoding dependencies between latent variables. The proposed tree-based generative architecture enables lightweight conditional inference and improves generative performance by utilizing specialized leaf decoders. 
  We show that TreeVAE uncovers underlying clusters in the data and finds meaningful hierarchical relations between the different groups on a variety of datasets, including real-world imaging data. 
  We present empirically that TreeVAE provides a more competitive log-likelihood lower bound than the sequential counterparts. 
  Finally, due to its generative nature, TreeVAE is able to generate new samples from the discovered clusters via conditional sampling.

Tree Variational Autoencoders
Laura Manduchi; Moritz Vandenhirtz, Alain Ryser, Julia E. Vogt Department of Computer Science ETH Zurich Switzerland
Abstract
We propose Tree Variational Autoencoder (TreeVAE), new generative hierarchi- cal clustering model that learns flexible tree-based posterior distribution over latent variables. TAE hierarchically divides samples according their intrin- sic characteristics, shedding light on hidden structures in the data. It adapts its architecture to discover the optimal tree for encoding dependencies between latent variables. The proposed tree-based generative architecture enables lightweight conditional inference and improves generative performance by utilizing specialized leaf decoders. We show that Tree VAE uncovers underlying clusters in the data and finds meaningful hierarchical relations between the different groups on variety of datasets, including real-world imaging data. We present empirically that Tree- VAE provides a more competitive log-likelibhe -igliikll lower bound than the sequential counterparts. Finally, due to its generative nature, Tree is able to generate new samples from the discovered clusters via conditional sampling.
1 Introduction
Discovering structure and hierarchies the data has been a long-standing goal in machine learning (Bishop) 2006 Bengio et al. 2012 Jordan & Mitchell 2015). Interpretable supervised methods, such as decision trees (Zhou & Feng 201 Tanno et al.) 2019), have proven to be successful unveiling hierarchical relationships within data. However, the expense of annotating large quantities of data has resulted in a surge of interest in unsupervised approaches (LeCun et al. 2015). Hierarchical clustering (Ward 1963) offers unsupervised path to find hidden groups in the data and their hierarchical relationship (R. G. B. Campello et al.) 2013). Due to its versatility, interpretability, and ability to uncover meaningful patterns in complex data, hierarchical clustering has been widely used in variety of applications, including phylogenetics (Sneath & Sokal 1962), astrophysics (McConnachie et al.) 2018), and federated learning (Briggs et al. 2020). Similar to how the human brain automatically categorizes and connects objects based on shared attributes, hierarchical clustering algorithms construct dendrogram
tree-like structure of clusters that organizes data into nested groups based on their similarity. Despite its potential, hierarchical clustering has taken a step back in light of recent advances in self-supervised deep learning (Chen al. 20220) and only few deep learning based methods have been proposed in recent years (Goyal et al. 2017 Mautz et al. 2020).
Deep latent variable models (Kingma & Welling 2019), class of generative models, have emerged as powerful frameworks for unsupervised learning and they have been extensively used to uncover hidden structures in the data (Dilokthanakul et al. 2016)] Manduchi et al.) 2021). They leverage the flexibility of neural networks to capture complex patterns and generate meaningful representations of high-dimensioonl data. By incorporating latent variables, these models can uncover the underlying factors of variation of the data, making them a valuable tool for understanding and modeling complex data distributions. In recent years, variety of deep generative methods have been proposed to incorporate more complex posterior distributions by modeling structural sequential dependencies
'Equal contribution. Correspondence (laura manduchi. moritz. vandenhirtz)@inff.cmm ethz.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

animals
vehicles
horses
dogs
trucks
cars
ships
planes
birds
cats
frogs
deers
Figure 1: The hierarchical structure discovered by VAE on the CIFAR-10 10 dataset. We display random subsets of images that are probabilistically assigned to each leaf of the tree.
between latent variables (Sonderby et al. 2016 He et al. 2018 Maalóe et al., 2019 Vahdat Kautz 2020a), thus offering different levels of abstraction for encoding the data distribution.
Our work advances the state-of-the-art in structured VABs by combining the complementary strengths of hierarchical clustering algorithms and deep generative models. We propose TreeVAE novel tree-based generative model that encodes hierarchical dependencies between latent variables. We introduce a training procedure to learn the optimal tree structure to model the posterior distribution of latent variables. An example of a tree learned by Tree VAE is depicted and split are encoded by neural networks, while the circles depict latent variables. Fach sample associated with probability distribution over paths. The resulting tree thus organizes the data into an interpretable hierarchical structure in an unsupervised fashion, optimizing the amount of shared information between samples. In CIFAR-10, 10, for example, the method divides the vehicles and animals into two different subtrees and similar groups (such planes and ships) share common ancestors.
Fig. [ Each edge
Our main contributions are as follows: (i) We propose novel, deep probabilistic approach to hierarchical clustering that learns the optimal generative binary tree mimic the hierarchies present in the data. (ii) We provide thorough empirical assessment of the proposed approach on MNIST, Mashion-MNIST; 20Newsgroups, and Omniglot. partieular, we show that VAe (a) outperforms related work on deep hierarchical clustering, (b) discovers meaningful patterns in the data and their hierarchical relationships, and (c) achieves a more competitive log- likelihood lower bound compared to VAE and LadderVAE, its sequential counterpart. (iii) We propose an extension of Tree VAE that integrates contrastive learning into its tree structure. Relevant prior knowledge, expertise, or specific constraints can be incorporated into the generative model via augmentations, allowing for more accurate and contextually meaningful clustering. We test the contrastive version of Tre on CIFAR-10, 10, CIFAR- 100, and CelebA, and we show that the proposed approach achieves competitive hicrarchical clustering performance compared to the basclines.
2 Tree VAE
We propose Vre,VAE, a novel deep generative model that learns flexible tree-based posterior distribution over latent variables. Each sample travels through the tree from root to leaf in probabilistic manner as VAE learns sample-specific probability distributions of paths. As result, the data is divided in a hierarchical fashion, with more refined concepts for deeper nodes in the tree. The proposed graphical model is depicted in Fig. 2 The inference and generative models share the same top-down tree structure, enabling interaction between the bottom-up and top-down architecture, similarly to Sønderby et al. (2016).
2.1 Model Formulation
Given H. the maximum depth of the tree, and dataset X, the model is defined by three components that are learned during training:
'Tee code is available htt.ppatttttttaae/.......ooo//an

---

the global structure of the binary tree T. which specifies the set of nodes V = {O..... .......V, VJ. for different
the set of leaves where L C V, and the set of edges See Fig. examples of tree structures learned by the model.
the sample-specific latent embeddings z {zo. (....... Zv which are random variables assigned to each node in V. Each embedding characterized by a Gaussian distribution function of the realization of the parent node. The dimensions of the latent embeddings are defined by their depth, with 2; E RReeepece where depth(i) is the depth of the node i, and he)pth(()) the embedding dimension for that depth.
whose parameters
the sample-specific decisions с = Cy Cy- which are Bernoulli random variables defined by the probability of going to the right (or left) child of the underlying node. They take values E {0, 1 for i E V L, with c, if the left child is selected. A decision path, Pt. indicates the path from root leaf given the tree and is defined by the nodes in the path, e.g., in Fig. 2 P. {0, 1, 1,4, 5J. The probability of Pt is the product of the probabilities of the decisions in the path.
The tree structure is shared across the entire dataset and is learned iteratively by growing the tree wise. The latent embeddings and the decision paths, on the other hand, are learned using variational inference by conditioning the model the current tree structure. The generative/infernnce
model and the learning objective conditioned on T are explained in Sec. 2.2 2.3 3 2.4 respectively, we elaborate on the efficient growing procedure of the tree.
while in 2.5
2.2
Generative Model
The generative process of VAE for given T is depicted in Fig. 2 (right). The generation of a new sample starts from
O
Zu
Zu
the root. First, the latent embedding of the root node Zo is sampled from standard Gaussian pe (zo) = N (20 0, I). Then, given the sampled zu, the decision of go- ing to the left or the right node is sampled from Bernoulli distribution p(co 20) (er(r,o(zz))) where {rp.: E V \ L} are functions parametrized by neural net- works defined routers, and cause the splits in Fig. [ The subscript p is used to indicate the parameters of the generative model. The latent embedding of the selected child, let us assume it is ø1, is then sam- pled from a Gaussian distribution pe (z1 20) =N)=((2) (zi Pp.1 (2o) of,1 (zo)). where
Zi
Z2
shared
7-
Z3
Z4
Z3
Figure 2: The proposed inference (left) and generative (right) models for Tree Circles stochastic vari- ables while diamonds are deterministic. The global topology of the tree is learned during training.
(Mp.i. i E \ {0}) are functions parametrized by neural networks defined as transformations. They are indicated by the top-down arrows in Fig. 2 This process continues until leaf is reached.
Let us define the set of latent variables selected by the path Pt, which goes from the root to the leaf l, as Zp, {z, E Pi). the parent node of the node i as pa(i), and p(Cpa(t)=> Epa(s)) the probability of going from pa(i) to i. Note that the path Pi defines the sequence of decisions. The prior probability of the latent embeddings and the path given the tree T can be summarized as
po(2p, Pt) p(2))
II
P(Cpatt)-xs Epa(t) =pa(s))p(2; 2po(t))-
(1)
tEPi\(O)
Finally, is sampled from distribution that is conditioned the selected leaf. If we assume that is real-valued, then
Po( kp, Pt) = =V Mr. (z)) (zi)),
(2)
where (12,t, Zz,1 I. 1 E T.} are functions parametrized by leaf-specific neural networks defined as decoders.
3

---

2.3 Inference Model
The inference model is described by the variational posterior distribution of both the latent embeddings and the paths. It follows a similar structure as in the prior probability defined in ( with the difference that the probability of the root and of the decisions are now conditioned on the sample æ:
(zp, Pe x) = q(z0 x)
II
((Cpati)-<i x ())(2; xpa(s))
(3)
(EPM(()
To compute the variational probability distribution of the latent embeddings g(zo x) and {(z; ppa(t)) where
((2 m) = N (zo (1,,((2),
(z: Xpa(s)) = -N (N(z Mai (Epate)) o p()(t)) Vi E Pt,
(5)
follow similar approach to the one proposed by /Sondebyy et al. (2016). Note that the subscript 9 to indicate the parameters of the inference model. First, deterministic bottom-up up pass computes the node-specific approximate likelihood contributions
d MLP (d++1)
(6)
pg.. Linear (dleept((s) i
(7)
a²= Softplus (Linear (ddepth(e)) .
(8)
9.4
where d is parametrized by domain domain-specific neural network defined as encoder, and MLP(dd,) for h E (1,...: H). indicated by the bottom-up arrows in Fig. are neural networks, shared among the parameter predictors, ug.i o², of the same depth. They are characterized by the same architecture
as the transformations defined in Sec2..2
stochastic downward pass then recursively computes the approximate posteriors defined as
pp.0pp
9
Migi
9.a
g p.
where all operations are performed elementwise. Finally, the variational distributions of the decisions q(e; æ) are defined as
q(; x)) = g(c; ddepsh(c)) ) = Ber(rgg a (cdeepth(i)))
(10)
where (Te, e L) are functions parametrized by neural networks and are characterized by the same architecture as the routers of the generative model defined in Sec. 2.2
2.4 Evidence Lower Bound
The parameters of both the generative model (defined as p) and inference model (defined as q).
consisting of the encoder (J4,,0, a..)). the transformations (f (
(Ha.cs )g4)) E V (0)}), the decoders ({Ma,, nx E IL.)) and the routers (frm,i Tq.. i E W.L.), are learned by maximizing the Evidence Lower Bound (ELBO) (Kingma & Welling, 2014 Rezende et al. (014).. Each leaf / associated with only one path Pi, hence we can write the data likelihood conditioned on T as
p( () =
p(z, Zp. Pi)
Σ
Po(z,, Pi)po(< 2p,, Pi).
(11)
LEL zPL
JeL 2P:
We use variational inference to derive the ELBO of the log-
L(x T) (,((,,, [log p( æp, Pi) KL (g (zp, Pi x))lp(Z), Pi)).
(12)
The first term of the ELBO is the reconstruction term:
free Eytxp, Py m)) [log p(x Zp, Pi)
(13)
Σ
((20 x)) II
d(patt))>i x))(Z; Zpa(s) log zp, Pt)
(14)
TET.
zr
sEP((())
M
Σ
Σ
P(l: log (x Med (z)m)) oz. (z(m)))
(15)
M
m= IET.
P(i:c)
II
a(Cpoti)-3 x)) for E
(16)
JEPA(0)

---

where P. for i E V is the path from root to node i, P():e) is the probability of reaching node i, which is the product over the probabilities of the decisions in the path until i, 2((T) are the Monte Carlo (MC) samples, and M the number of the MC samples. Intuitively, the reconstruction loss is the sum of the leaf-wise reconstruction losses weighted by the probabilities of reaching the respective leaf. Note that here we sum over all possible paths in the tree, which is equal to the number of leaves.
The second term of 12 is the Kullback-Leibhe divergence (KI.) between the prior and the variational posterior of the tree. It can be written as sum of the KL of the root, the nodes, and the decisions:
KL (q (zp,, Pi a)Il (ZP, Pi)) = KLroot + KLnodde KLddccsioos
(17)
Koooot [Z |>x)
=))p((z)))
(18)
M
KLnodes R
Σ
Σ
P();c) (L(((Z
ppl(z (0(2(()))) M[p(=
ppa(z)=))
(19)
M
m=1 kEV
M
q(ci x)) p(cs x ((m))
KLddecsioos
Σ c(::)) c.E(U,1)
Σ g(c, | x) log
(20)
M
where M is the number of MC samples. We refer to Appendix A for the full derivation. The KLooot is the KL between the standard Gaussian prior p(zo) and the variational posterior of the root g(50 æ). thus enforcing the root to be compact. The KLnodea is the sum of the node- specific KLs weighted by the probability of reaching their node i: c(t; The node-specific KL of node is the KL between the two Gaussians (zz pa(zi)), p(z; pa;)))). Finally, the last term, KLdecisionn: is the weighted sum of all the KLs of the decisions, which are Bernoulli random variables,
l))oog KL LL(g((, x> p(c: x))) = Eo,5(0.1; q(c= x>) log ([og) ). The hierarchical specification of the binary tree allows encoding highly expressive models while retaining the computational efficiency of fully factorized models. The computational complexity is described in Appendix A.2
2.5 Growing The Tree
In the previous sections, we discussed the variational objective to learn the parame- ters of both the generative and the inference model given defined tree structure T. Here we discuss how learn the structure of the binary tree T. Tree starts by training a tree composed of a root and two leaves, see Fig. (left,, for N₂ epochs by optimizing the ELBO. Once the model converged, a leaf selected, e.g.. 21 in Fig- 3 and two children are attached to it. The leaf selection criteria can vary depending on the application and can be determined by, c.g., the reconstruc- tion loss or the ELBO. In our experiments, we chose to select the nodes with the max- imum number of samples to retain balanced leaves. The sub-tree composed of the new leaves and the parent node is then trained for N, epochs by freezing the weights of the rest of the model, see Fig. 3 (right), resulting in computing the ELBO of the nodes of the subtree. For efficiency, the subtree is trained using only the subset of data that have a high probability (higher than threshold f) of being assigned the parent node. The process is repeated until the tree reaches its maximum capacity (defined by the maximum depth) or until condition (such predefined maximum number of leaves) is met. The entire model is then fine-tuned for N2 epochs by unfreezing all weights. During fine-tuning, the tree is pruned by removing empty branches (with the expected number of assigned samples lower than threshold).
Zo
Zo
d
Zi
x
X
Z3
ds
X
Step
Step 2
The first two steps of the growing process to learn the global structure of the tree during training. Highlighted in red are the trainable weights.
Figure 3:
2.6 Integrating Prior Knowledge
Retrieving semantically meaningful clustering structures of real-world images is extremely challeng- ing, there are several underlying factors according to which the data can be clustered. Therefore,
5

---

is often crucial to integrate domain knowledge that guides the model toward desirable cluster assign- ments. Thus, we propose an extension of VAE where we integrate recent advances in contrastive learning (van den Oord et al. 2018 Chen et al. 2020 Li et al. (02T), whereby prior knowledge on data invariances can be encoded through augmentations. For batch X with N samples, we randomly augment every sample twice to obtain the augmented batch X with 2N samples. For all positive pairs (1, where Ø; and ă; stem from the same original sample, we utilize the NT-Xeni (Chen et
al. 2020), which introduces losses L = - log E
=== 3 [j exp exp (8,,//)) where Si,j denotes the cosine similanty between the representations of ã; and x and is a temperature parameter. We integrate f in both the bottom-up and the routers Tree VAE. In the bottom-up, similar to Chen et al. (2020), we compute C.,; on the projections gh (d,). For the routers, we directly compute the loss on the predicted probabilities (g.;(d,) Finally, we average the terms over all positive pairs and add them to the negative ELBO (12) in real-worlh image experiments. Implementation details can be found in
exp (81.//)
Appendix E while a loss ablation is shown in Appendix C.3
3 Related Work
Deep latent variable models automatically learn structure from data by combining the flexibility of deep neural networks and the statistical foundations of generative models (Mattei & Frellsen, 2018).
Variational autoencoders (VAEs) (Rezende et al. 2014 Kingma & Welling (0144) are among the
most used frameworks (Nasiri & Bepler] [022] Bae et al. 2023] Bredell et al.) 2023). A variety of works has been proposed to integrate more complex empirical prior distributions. thus reducing the gap between approximate and true posterior distributions (Ranganath et al.) 2015 Webb et al.
2017: Klushyn et al., 2019). Among these, the most related our work is the VAE-nCRP (Goyal it al. 2017 Shin et al. 2219 and the TMC-VAE (Vikram et al. 2018). Both works use Bayesian nonparametrie hierarchical clustering based on the nested Chinese restaurant process (nCRP) prior (Blei et al. 2003), and on the time- tmme-maginallied coalescent (TMC). However, even if they allow more flexible prior distributions these models suffer from restrictive posterior distributions (Kingma et al.
2016).To overcome the above issue, deep hierarchical VAEs (Gregor et al. [015]] Kingma et al. (0166) have been proposed to employ structured approximate posteriors, which are composed of hierarchies
of conditional stochastic variables that are connected sequentially. Among a variety of proposed
methods (Vahdat & Kautz] [220b]] Falck et al. 2022 T. Xiao Bamler 2023), Ladder VAE
(Sønderby et (016) is most related to Tree VAF. The authors propose to model the approximate posterior by combining a "bottom-up" recogmition distribution with the "top-down" prior. Further extensions include BIVA (Maaløe et al.) 2019), which introduces a bidirectional inference network, and Gra (He et al. 2019), that introduces gated dependencies over a fixed number of latent variables. Contrary to the previous approaches, Tree VAE models tree-based posterior distribution of latent variable, thus allowing hierarchical clustering of samples. For further work on hierarchical clustering and its supervised counterpart, decision trees, we refer to Appendix B
4 Experimental Setup
Datasets and Metrics: We evaluate the clustering and generative performance of Tree VAE on
MNIST (LeCun et al. 1998). Fashion-MNIST (H. Xiao et al.] (017)). 20Newsgroups (Lang) 1995). Omniglot (Lake al. 2015). and Omniglot-5, where only 5 vocabularies (Braille, Glagolitic, Cyrillic, Odia, and Bengali) are selected and used as true labels. We assess the hierarchical clustering performance by computing dendrogram purity (DP) and leaf purity (LP), as defined by (Kobren et al. 2017a) using the datasets labels, where we assume the number of true clusters is unknown. We also report standard clustering metrics, accuracy (ACC) and normalized mutual information (NMD), by setting the number of leaves for Tree VAE and for the baselines to the true number of clusters. In terms of generative performance, we compute the approximated true log-likelihood calculated using 1000 importance-weigheee
samples, together with the ELBO (2) and the reconstruction loss (6). We also perform hierarchical clustering experiments on real-world imaging data, namely CIFAR- 10. CIFAR-101 100 (Krizhevsky & Hinton 2009) with 20 superclasses as labels, and CelebA (Z. Liu et al. 2013) using the contrastive extension (Sec. 2.6). We refer to Appendix D for more dataset details.
Baselines: We compare the generative performance of VAE to the VAE (Rezende et al. 2014
Kingma & Welling 2014), its non-hierarchical counterpart, and the LadderVAE. (Sonderby et al.
2016), its sequential counterpart. For fair comparison, all methods share the same architecture and

---

Table 1: Test set hierarchical clustering performances (%) of TreeVAE compared with baselines.
Means and standard deviations are computed across 10 runs with different random model initialization. The star "*" indicates real-world image datasets on which contrastive approaches were applied.
Dataset
Method
DP
LP
ACC
NMI
MNIST
Agg VAE + Agg
63.7 0.0
78.6
69.5 0.0
0.0
79.9 2.2
90.8
8..6
81.6 2.0
LadderVAE Agg
81.6
90.
80.3
8.1
DeepECT
74.6
90.7
74.9
76.7
Tre (ours)
87.9
96.0 1.9
90.2 7.5
90.0 ::
Fashion
Agg
4..0
6..6
51.3
52.6
VAE + Agg
4..3
65.9
4.4
56.1 3.2
LadderVAF. Agg
49.5
67.6
55.9
60.7 1.4
DeepECT
44.9 +..
67.8
51.8 5.7
57.7 #.7
VAE (ours)
54.4
7..4
63.6
64.7
20Newsgroups
Agg VAE Agg
1..1
30.
2..1
27.5
0.3
18.1
15.2
11.6
LadderVAE Agg
9.0
20.0
17.4
17.8
DeepECT
9.3
17.2
15.6
18.1
VAE (ours)
17.5
38.4
32.8 2.3
34.4
Omniglot-5
Agg VAE Agg LadderVAE Agg
0.0
63.7 0.0
53.2 1..200000
33.3 0.0
46.3
68.1
4.29
34.4
49.8
71.3
59.6
41.2
DeepECT
33.3
55.1
41.1 #1.2222
23.5
VAE (ours)
58.8 1:
77.7 3.9
63.9 14.977.00 7.0
50.0 6.9
CIFAR-I0*
VAE + Agg
10.54 0.112
16.33 + 0.16
14.43 0.19
1.86 1.G6
LadderVAE Agg
12.81 0.020
25.37 0.66
19.29 0.60
7.41 0.442
DeepECT
10.01
10,30
10.31 0.39
0.18
Tre (ours)
35.30
53.85 1:23
52.98 1.34
41.44
CIFAR-100*
VAE Agg LadderVAE Agg
5.27 0.02
9.86 0.11
0.12
2.16
6.36 0.07
16.08
14.01 0.441
8.99
DeepECT
5.28 4: 0.18
6.97 4: 0.69
6.97 +:
1.71 0.86
Tre (ours)
10.44 + 0.38
24.16 0.0.6
21.82 0.7.7
17.80+0042
hyperparameters whenever possible. We compare TreeVAE to non-generative hicrarchical clustering basclines for which the code was publicly available: Ward's s minimum variance agglomerative clustering (Agg) (Ward) 1963; Murtagh Legendre 2014), and the DeepECT (Mautz et al. 2020). We propose two additional basclines, where we perform Ward's agglomerative clustering on the latent space of the VAE (VAE Agg) and of the last layer of the LadderVAE (LadderVAE + Agg). For the contrastive clustering experiments, we apply a contrastive loss similar to Tre to the VAE and the LadderVAE, while for DeepECT we use the contrastive loss proposed by the authors.
Implementation Details: While we believe that more complex architectures could have substantial choose employ rather simple settings validate
impact on the performance of VAE, the proposed approach. We set the dimension of all latent embeddings z = (zo..... Zy } to 8 for MNIST, Fashion, and Omniglot, CelebA. The maximum depth of the tree is set to 6 for all datasets, except 20Newsgroups where we increased the depth to maximum of 30 leaves for 20Newsgroups and CIFAR-101 and 20 for the rest, while for ACC and NMI fix the number of leaves the number of true classes. The transformations consist of one-layer MLPs of size 128 and the routers of two-layers of size 128 for all datasets except for real-world imaging data where we slightly increase the ML. complexity to 512. Finally,
for 20Newsgroups, and to 64 for CIFAR-IO. CIFAR-100, and
to capture more clusters. To compute DP and LP. we allow the tree to grow
encoder and decoders consist of simple CNNs and MI.ss. The trees are trained for N, = 150 epochs each growth step, and the final tree is finetuned for N, 200 epochs. For the real-world imaging experiments, set the weight of the contrastive loss to 100. See Appendix E for additional details.

---

Table 2: Test set generative performances of Tree with 10 leaves compared with baselines. Means and standard deviations are computed across 10 runs with different random model initialization.
Dataset
Method
LL
RL
ELBO
MNIST
VAE
-00.9 0.2
0.3
101.6
LadderVAE
0.5
8.7
103.2
VAE (ours)
-92.9 0.2
80.3 0.2
0.2
Fashion
VAE
-242.2 a: 0.2
231.7 0.5
-24..4 0.5
LadderVAR
239.. 0.5
231.5 0.6
244.00
VAE (Ours)
23..7=0..1
226.5 2
-..
20Newsgroups
VAE
44.26 =: 0.01
45.52 +: 0.03
44.61 0.01
LadderVAE
44.30 0.00 0.03
43.52 0.00
44.62 0.00
VAE (Ours)
511.67
45.83 0.33
5...79
Omniglot
VAE
-115.3
101.6
-118.2
LadderVAE
113.1
100.7
-17.55 0.6
TreeVAE (Ours)
-100.4
96.9 + 0.5
-114.6
Results
Hierarchical Clustering Results Table shows the quantitative hierarchical clustering results averaged across 10 seeds. First, we assume the true number of clusters is unknown and report DP and LP. Second, we assume we have access the true number of clusters K and compute ACC and NMI. As can be seen, TreeVAB outperforms the baselines in both experiments. This suggests that the proposed approach successfully builds an optimal tree based on the data's intrinsic characteristics. Among the different baselines, agglomerative clustering using Ward's 's method (Agg) trained on the last layer of LadderVAE shows competitive performances. To the best of our knowledge, we are the first to report these results. It is noteworthy to observe that it consistently improves over VAE + Agg, indicating that the last layer LadderVAE. captures more cluster information than the VAE.
Generative Results In Table 2 evaluate the generative performance of the proposed approach, Vre, compared to the VAE, its non-hierarchical counterpart, and LadderVAE, its sequential counterpart. Tree VAE outperforms the baselines the majority of datasets, indicating that the proposed ELBO (12 can achieve a tighter lower bound of the log-likelihood. The most notable improvement appears be reflected in the reconstruction loss, showing the advantage of using cluster-specializzd decoders. However, this improvement comes at the expense of a larger neural network architecture and an increase the number of parameters (as VAE has decoders). While this requires more computational resources at training time, during deployment the tree structure of VAE permits lightweight inference through conditional sampling, thus matching the inference time of LadderVAE. It is also worth mentioning that results differ from (Sønderby et al. 2016) as we adapt their architecture to match our experimental setting and consequently use smaller latent dimensionality. Finally, we notice that more complex methods are prone to overfitting on the 20Newsgroups dataset, hence the best performances are achieved by the VAE.
Real-world Imaging Data & Contrastive Learning Clustering real-world imaging data is ex- tremely difficult as there are endless possibilities of how the data can be partitioned (such as the colors, the landscape, etc). We therefore inject prior information through augmentations guide Tree VAE and the baselines to semantically meaningful splits. Table (bottoom) shows the hierarchical clustering performance of TreeVAF: and its baselines, all employing contrastive learning, CIFAR- 10 and CIFAR-100. We observe that DeepECT struggles in separating the data as their contrastive approach leads to all samples falling into the same leaf. In Table face attributes using the tree learned by Tree For all datasets, TreeVAE is able to group the data into contextually meaningful hierarchies and groups, evident from its superior performance compared to the baselines and from the distinct attribute frequencies in the leaves and subtrees.
present the naf-frequencies of various
Discovery of Hierarchies addition to solely clustering data, Tre able to discover meaningful hierarchical relations between the clusters, thus allowing for more insights into the
8

---

(a)
(b)
Figure Hierarchical structures learned by VAE on Fashion. Subtree (a) encodes tops, while (b)
encodes shoes, purses, and pants.
++ comp. windows.x
++ comp-graphics
comp.os. mammp.osssswwiiii. www.m-siddws..iis
(a)
(b)
++ comp.sys. stppp.ymssbm.pp.......mm..... httpp.s..bb....aadwaarrr pa.hard
4
7
8
A
T
p
comp.sys.mas wwt.s/ma.haarwww... hardware
misc.forsale
1
N
C
4
%
rec.autos / motorcycles
&
N
8
rec.motorcycles
misc.forsale
++ sci.med
Figure 6: Hierarchical structures learned by Tree Omniglot-5. Subtree (a) learns a hierarchy over Braille and the Indian alphabets, while (b) groups Slavic alphabets.
+ talk. religion.misc
altatheism
soc.religion.crristan
poc.religion.chrittaa
Attribute
2
3
4
7
8
talk, politicss.ideeeee
Female
97.2
55.0
97.7
86.6
23.1
30.7
46.6
43.7
polk.politics.gum misc
Bangs
1.6
1.2
24.1
61.7
3.4
11.1
9.1
11.4
rec.sport.hocke recssport
Blonde
1.1
3.7
66.7
2.2
5.8
2.6
26.1
7.1
Makeup
75.7
43.4
76.6
59.7
15.0
12.4
16.3
12.8
sci.crypt
Smiling
54.3
66.6
66.4
51.2
54.7
42.4
37.3
22.4
++ sci.space
Hair Loss
3.6
17.8
3.0
0.2
18.9
6.9
21.2
10.6
Beard
20.6
0.4
3.7
39.5
36.5
21.3
21.4
Figure 5:
Hierarchical structure
learned by Vre on 20News- groups.
Table 3: We present the frequency (in %) of selected at- tributes for each leaf of Tree VAE with eight leaves in CelebA.
Figure 7: Hierarchical structure learned by TrE with eight leaves on the CelebA dataset with generated images through conditional sampling. Generally, most females are in the left subtree, while most males are in the right subtree. We observe that leaf is associated with dark-haired females, leaf 2 with smiling, dark-haired individuals, leaf 3 with blonde females, leaf 4 with bangs, leaf 7 with a receding hairline, and leaf 8 with non-sminnng non-smiling people. See Table for further details.
dataset. In the introductory Fig. and we present the hierarchical structures learned by
VAE while Fig. and we additionally display conditional cluster generations from the leaf-
specific decoders. In Fig. Tre separates the fashion items into two subtrees, one containing shoes and bags, and the other containing the tops, which are further refined into long and short sleeves. In Fig. 5 we depict the most prevalent ground- -trtt topic label in each leaf. VAE learns to

---

separate technological and societal subjects and discovers semantically meaningful subtrees. In Fig. TreeVAE learns split alphabets into Indian (Odia and Bengali) and Slavic (Glagolitic and Cyrillic) subtrees, while Braille is grouped with the Indian languages due to similar circle-like structures. For CelebA, Fig. and 3 the resulting tree separates genders in the root split. Females (left) are further divided by hair color and hairstyle (bangs). Males (right) are further divided by smile intensity, beard, hair loss, and age. In Fig. [ and Appendix [[we show how Tree can additionally be used to sample unconditional generations for all clusters simultaneously, by sampling from the rool and propagating through the entire tree. The generations differ across the leaves by their cluster-specific features, whereas cuuserrinndepnndee properties retained across all generations.
leaf
leaf
leaf 3
leaf
leaf 5
leaf
leaf 7
leaf 8
Figure 8: Selected unconditional generations of CelebA. One row corresponds one sample from the root, for which we depict the visualizations obtained from the 8 leaf-decoders, The overall face shape, skin color, and face orientation are retained among leaves from the same row, while several properties (such as make-up, beard, mustache, glasses, and hair) vary across the different leaves.
Conclusion
In this paper, we introduced VAE, a new generative method that leverages tree-based posterior distribution of latent variables to capture the hierarchical structures present the data. TreeVAE optimizes the balance between shared and specialized architecture, enhancing the learning and adaptation capabilities of generative models. Empirically, we showed that our model offers substantial improvement in hierarchical clustering performance compared to the related work, while also providing a tighter lower bound to the log-likelihood of the data. We presented qualitatively comprehensive understanding of
how the hierarchical structures learned by Vre enable the data, thereby facilitating enhanced analysis, interpretation, and decision-making Our findings highlight the versatility of the proposed approach, which we believe to hold significant potential for unsupervised representation learning, paving the way for exciting advancements in the field.
Limitations & Future Work: Currently, Tree uses simple heuristic on which node to split that might not work on datasets with unbalanced clusters. Additionally, the contrastive losses on the routers encourage balanced clusters. Thus, more research is necessary to convert the heuristics to data-driven approaches. While deep latent variable models, such as VAEs, provide framework for modeling explicit relationships through graphical structures, they often exhibit poor performance on synthetic image generation. However, more complex architectural design (Vahdat & Kautz) (020a) or recent advancement in diffusion latent models (Rombach et al. (021) present potential solutions to enhance image quality generation, thus striking an optimal balance between generating high-quality images and capturing meaningful representations.
Acknowledgments and Disclosure of Funding
We thank Thomas M. Sutter for the insightful discussions throughout the project, Jorge da Silva Gonçalves for providing interpretable visualizations of the Tre model, and Gabriele Manduchi for the valuable feedback on the notation of the ELBO. LM is supported by the SDSC PhD Fellowship #1-001568 037. MV supported by the Swiss State Secretariat for Education, Research and Innovation (SERI) under contract number MB22.00047. AR is supported by the StimuLoop grant #1-007811-00 1-002 and the Vontobel Foundation.
10

---

References
Arenas, M., Barceló, P., Orth, M. R., & Subercaseaux, B. (2022). On computing probabilistic expla- nations for decision trees. In Neurips. Retrieved from httpt::ppppppes...a..tt http://pppppppppp...pppppp http:p///ppppprrsppppccpaappppp Daaer/2022/hasie httpa///888ee1100000000333333110
Abstracct--ooffrrnnce html
Bae, J., Zhang, M. R., Ruan, M., Wang. E., Hasegawa, S., Ba, J., & Grosse, R. B. (2023).
Multi-rate VAE: Train once, get the full rate-distortion curve. In The eleventh international conference on learning representations. Retrieved from http:://wpeeeeew.neteeeee (JMNS)
Basak, J., & Krishnapuram, R. (2005). Interpretable hierarchical clustering by constructing an unsupervised decision tree. IEEE Trans. Knowl. Data Eng., 17(1), 121-132.
Bengio, Y., Courville, A. C., & Vincent, P. (2012). Representation learning: A review and new perspectives. IEEE Transactions on Pattern Analysis and Machine Intelligence, 35, 1798-1828.
Bishop, C. M. (2006). Pattern recognition and machine learning (information science and statistics). Berlin, Heidelberg: Springer-Verlag.
Blei, D. M., Jordan, M. I., Griffiths, T. L.., & Tenenbaum, B. (2003). Hierarchical topic models and the nested chinese restaurant process. In Proceedings of the 16th international conference on neural information processing systems (p. 17-24). Cambridge, MA, USA: MIT Press.
Blockeel, H., & De Raedt, L. (1998). Top-down induction clustering trees. In Proceedings of the fifteenth international conference machine learning (pp. 55-63).
Bredell, G., Flouris, K., Chaitanya, K., Erdil, E., & Konukoglu, E. (2023). Explicitly minimizing the blur error of variational autoencoders. In The eleventh international conference on learning representations. Retrieved from https: http:///oppeeeveeeeeee
Breiman, L.,, Friedman, J. H., Olshen, R. A., & Stone, C. (1984). Classification and regression frees. Wadsworth.
Briggs, C., Fan, Z., & András, (2020). Federated learning with hierarchical clustering of local updates to improve training on non-iid data. 2020 International Joint Conference Neural Networks (IJCNN), 1-9.
Campello, R. J., Moulavi, D., Zimek, A., & Sander, J. (2015). Hierarchical density estimates for data clustering, visualization, and outlier detection. ACM Transactions on Knowledge Discovery from Data (TKDD), 10(1), 1-51.
Campello, J. G. B., Moulavi, D., Zimek, A., & Sander, J. (2015). Hierarchical density estimates for data clustering, visualization, and outlier detection. ACM Transactions on Knowledge Discovery from Data (TKDD), 10, 1 51.
Chen, T., Kornblith, S., Norouzi, M., & Hinton, G. (2020). A simple framework for contrastive learning of visual representations. In Proceedings of the 37th international conference on machine learning. JMLR.org.
Dilokthanakul, N., Mediano, P. A. M., Garnelo, M., Lee, M. C. H., Salimbeni, H., Arulkumaran, K., & Shanahan, M. (2016). Deep unsupervised clustering with Gaussian mixture variational autoencoders. 11I.iv:611..02442
Ester, M., Kriegel, H., Sander. J., & Xu, X. (1996). A density-based algorithm for discovering clusters in large spatial databases with noise. E. Simoudis, J. Han, & U. M. Fayyad (Eds.). Proceedings of the second international conference on knowledge discovery and data mining (kdd-96), portland, oregon, USA (pp. 226 221-231). AAAI Press. Retrieved from http::// //wp:////wwwww. aaai www./orrrrrrrr.. phppppp)
Falck, F., Williams, C., Danks, D.,, Deligiannidis, G., Yau, C., Holmes, C. C.,... Willetts, M. (2022). A multi-resolutionn framework for u-nets with applications to hierarchical VAEs. In A. H. Oh, A. Agarwal, D. Belgrave, & K. Cho (Eds.), Advances neural information processing systems. Retrieved from https::/ http//pprrrreee.eeeffooourmmm///
Fraiman, R., Ghattas, B., & Svare, M. (2013). Interpretable clustering using unsupervised binary trees. Adv. Data Anal. Classif., 7(2), 125-145. Retrieved from httpps://doi..... http::///ooo.or//10000077
s 811634-00 013---0111-1--99 doi: 10. 11.10007.00 10000000//1160 1634-0131134-

---

Frosst, N., & Hinton. G. E. (2017). Distilling a neural network into a soft decision tree. In T. R. Besold & O. Kutz (Eds.), Proceedings of the first international workshop on comprehensibilliy and explanation in Al and MI. 2017 co-located with 16th international conference of the italian association for artificial intelligence (ai*ia 2017), bari, italy, november 16th and 17th, 2017 (Vol. Retrieved from https: / ctte///eruus//// argg/Vol-2771/ExAA. 2011
2071). CEUR-WS.on
_paper_3.pdf
Ghojogh, B., Ghodsi, A., Karray, F., & Crowley, M. (2021). Uniform manifold approximation and
projection (UMAP) and its variants: Tutorial and survey. CoRR, abs/2109.02508. Retrieved from
https://arxiv. wttg/ot/////aeee abw. 02508
Goyal, P., Hu, Z., Liang. X., Wang, C., Xing. E. P., Mellon, C. (2017). Nonparametric variational auto-encoders for hierarchical representation learning, 2017 IEEE International Conference Computer Vision (ICCV), 5104-5112.
Gregor, K., Danihelka, I., Graves, A., Rezende, D., & Wierstra, D. (2015). Draw: A recurrent neural network for image generation. In F. Bach & D. Blei (Eds.), Proceedings of the 32nd international conference on machine learning (Vol. 37, pp. 1462-1471). Lille, France: PMLR. Retrieved from htttps://proceedingg mlr wwwpppesss//.. wttpa////v///rgrr555/5 httll
He, J., Gong, Y., Marino, J., Mori, G., & Lehrmann, A. M. (2018). Variational autoencoders with jointly optimized latent dependency structure. In International conference on learning representations.
He, J., Gong, Y., Marino, J., Mori, G., & Lehrmann, A. M. (2019). Variational autoencoders with jointly optimized latent dependency structure. In Iclr.
Heller, K. A., Ghahramani, Z. (2005). Bayesian hierarchical clustering. Proceedings of the 22nd international conference machine learning (pp. 297-304).
Jordan, M. I., & Mitchell, T. M. (2015). Machine learning: Trends, perspectives, and prospects. Science, 349(6245), 255-260, Retrieved from https::/www.scieettt. http://wwwwccceee.e.ooo////////a wwwp////g//g/////////ooo/e/oetst 11111.11.. 1126/ Schence.aaa8444 doi: 10. 10.1126/cieecce.aaa8//
Kingma, D. P., Salimans, T., Jozefowicz, R., Chen, X., Sutskever, I., & Welling, M. (2016). Improved variational inference with inverse autoregressive flow. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, & R. Garnett (Eds.), Advances in neural information processing systems (Vol. 29).
Curran Associates, Inc, Retrieved from https: httpm/ttttss/oeeeeeeeeeee httpa/////coccceoi..e..e nwwrips...pp... cctppppr_fiileesseeeeeeeeeeeeeee påf
Kingma, D. P., & Welling, M. (2014). Auto-encoding variational bayes. arXiv preprint
arXiv: 1312.61 14.
Kingma, D. P.. Welling, M. (2019). An introduction to variational autoencoders. Found. Trends Mach. Learn., 12, 307-392.
Klushyn, A., Chen, N., Kurle, R., Cseke, B., & Smagt, P. v. d. (2019). Learning hierarchical priors in vaes. Proceedings of the 33rd international conference on neural information processing systems. Red Hook, NY, USA: Curran Associates Inc.
Kobren, A., Monath, N., Krishnamurthy, A., & McCallum, A. (2017a). A hierarchical algorithm for extreme clustering. In Proceedings of the 23rd ACM SIGKDD international conference knowledge discovery and data mining, halifax, ns, canada, august 13 17, 2017 (pp. 255- 264). ACM. Retrieved from hattss://do/ org 114.0.11/309918833 3098079 3097983.3098079
doi: 10.1
Kobren, A., Monath, N., Krishnamurthy, A., & McCallum, A. (2017b). A hierarchical algorithm for extreme clustering. In Proceedings of the 23rd acm sigkdd international conference on knowledge discovery and data mining (pp. 255-264).
Krizhevsky, A., & Hinton, G. (2009). Learning multiple layers of features from tiny images (Tech. Rep. No. 0). Toronto, Ontario: University of Toronto.
Lake, B. M., Salakhutdinov, R., & Tenenbaum, B. (2015). Human-level concept learning through probabilistic program induction. Science, 350(6266), 1332-1338. Retrieved from https: // www. www.sscience.cm wwtp/////gg/////////
2cc2////cinnc...e. aab3050 doi: 10.1126/scicccc.uab333
Lang, K. (1995). Newsweeder: Learning to filter netnews. In Proceedings of the twelfth international
conference on machine learning (p. 331-339).
12

---

Laptev, D., & Buhmann, J. M. (2014). Convolutional decision trees for feature learning and segmentation. In Pattern recognition: 36th german conference, gepr 2014, münster, germany, september 2-5, 2014, proceedings 36 (pp. 95-106).
LeCun, Y., Bengio, Y., & Hinton, G. (2015). Deep learning. Nature, 521(3))3), 436-444.
LeCun, Y., Bottou, L., Bengio, Y., & Haffner, P. (1998). Gradient-based learning applied to document recognition. Proc. IEEE, 86, 2278-2324
Li, Y., Hu, P.. Liu, J. Z., Peng, D.. Zhou, J. T., & Peng, X. (2021). Contrastive clustering. In Thirty- fifth AAAI conference on artificial intelligence, AAAI 2021. thirty-third conference on innovative applications of artificial intelligence, IAAI 2021, the eleventh symposium on educational advances
in artificial intelligence, EAAI 2021, virtual event, february 2-9, 202/ (pp. 8547-8555) AAAI
Press. Retrieved from btttps://of aaai org/ wewgx/idex.../ php/ AAA.AAAAA//aaaaa.. AAAT/article/view///TO00
Li, Y., Yang, M., Peng, D., Li, T., Huang, J., & Peng. X. (2022). Twin contrastive learning for online clustering. International Journal of Computer Vision, 130(9), 2205-2221.
Liu, B., Xia, Y., & Yu, P. S. (2000). Clustering through decision tree construction. In International conference on information and knowledge management.
Liu. Z., Luo, P., Wang. X., & Tang, X. (2015). Deep learning face attributes in the wild. In Proceedings of international conference computer vision (icev).
Maaløe, I.., Fraccaro, M., Liévin, V., & Winther, O. (2019). Biva: A very deep hierarchy of latent variables for generative modeling. In Neurips.
Manduchi, I.., Chin-Cheong. K., Michel, H., Wellmann, S., & Vogt, J. E. (2021). Deep conditional gaussian mixture model for constrained clustering. In Neural information processing systems.
Mathieu, E., Le Lan, C., Maddison, J., Tomioka, R., Teh, Y. W. (2019). Continuous hierarchical representations with poincaré variational auto-encoders. Advances in neural information processing systems, 32.
Mattei, P.-A., & Frellsen, J. (2018). Leveraging the exact likelihood of deep latent variable models. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, & R. Gar- nett (Eds.), Advances in neural information processing systems (Vol, 31). Curran Associates, Inc. Retrieved from http::t:t:::/cttetttttitt wew.eeeedddinnnnnnnieeeee.e.eeee newrips. httppp..cpppppo cc/ wwtpaper_fiileeeecee fearrrfile/pppee /ttpa//////pppe//////e////// Paper pdf
Mautz, D., Plant, C., & Böhm, C. (2020). Deepect: The deep embedded cluster tree. Data Science
and Engineering, 5, 419 432.
McConnachie, A. W., Ibata, R. A., Martin, N., Ferguson, A. N., Collins, M. L. M., Gwyn, S. J., Widrow, L. M. (2018) 8). The large-scale structure of the halo of the andromeda galaxy. il hierarchical structure in the pan-andromeda archaeological survey. The Astrophysical Journal, 868.
Monath, N., Zaheer, M., Silva, D., McCallum, A., & Ahmed, A. (2019). Gradient- bbased hierarchical clustering using continuous representations of trees in hyperbolic space. In Proceedings of the 25th aem sigkdd international conference on knowledge discovery & data mining (pp. 714-722).
Moshkovitz, M., Yang, Y., & Chaudhuri, K. (2021). Connecting interpretability and robust- ness in decision trees through separation. In M. Meila & T. Zhang (Eds.), Proceedings of the 38th international conference on machine learning, ICML 2021, 18-24 july 2021, virtual event (Vol. 139, pp. 7839-7849). PMLR. Retrieved from http: http:://proceeedngg..... mlr press/vi39) moshkovitz2 ittooviti2aaaa html
Murtagh, F., & Contreras, P. (2012). Algorithms for hierarchical clustering: an overview. Wiley Interdisciplinary Reviews: Data Mining and Knowledge Discovery, 2(1), 86-97.
Murtagh, F., & Legendre, P. (2014), Ward's ' hierarchical agglomerative clustering method: Which algorithms implement ward's criterion? Journal of Classification, 31, 274-295.
Nasiri. A., & Bepler, T. (2022). Unsupervised object representation learning using translation and rotation group equivariant VAE. In A. H. Oh, A. Agarwal, D. Belgrave, & K. Cho (Eds.), Advances in neural information processing systems. Retrieved from httpss://peereeiieeee net/
forum? 1d 1d=qmm__... dd= JMIMLLL
Neal, R. M. (2003). Density Modeling and Clustering using Dirichlet Diffusion Trees. In
J. M. Bernardo et al. (Eds.), Bayesian Statisties 7 (p. 619-629). Oxford University Press.
13

---

Nistér, D., & Stewénius, H. (2006). Scalable recognition with a vocabulary tree. 2006 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR'O0) (06), 2, 2161-2168.
Pace, A., Chan, A. J., van der Schaar, M. (2022). POETREE: interpretable policy learning with adaptive decision trees. In The tenth international conference on learning representations, ICLR 2022, virtual event, april 25-29, 2022. OpenReview.net. Retrieved from https::/
net
Ram, P., & Gray, A. G. (2011). Density estimation trees. In Knowledge discovery and data mining. Ranganath, R., Tran, D., & Blei, D. M. (2015). Hierarchical variational models. ArXiv,
abs/1511.02386.
Rezende, D. J., Mohamed, S., & Wierstra, D. (2014). Stochastic backpropagation and approximate inference in deep generative models. In Proceedings of the 31st international conference on machine learning (Vol. 32, pp. 1278-1286) PMLR.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., & Ommer, B. (2021). High-resolution image synthesis with latent diffusion models. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 10674-10685.
Rota Bulo, S., & Kontschieder, P. (2014). Neural decision forests for semantic image labelling. In Proceedings of the ieee conference on computer vision and pattern recognition (pp. 81-88).
Shin, S.-J., Song, K., & Moon, I.-C. (2019). Hierarchically clustered representation learning. In Aaai conference on artificial intelligence.
Sneath, P. H. (1957). The application of computers to taxonomy. Microbiology, 17(1), 201-226.
Sneath, P. H., & Sokal, R. R. (1962). Numerical taxonomy. Nature, 193, 855-860.
Sohn, K. (2016). Improved deep metric learning with multi-class n-pair loss objective. In D. D. Lee, M. Sugiyama, U. von Luxburg, I. Guyon, & R. Garnett (Eds.), Advances in neural information pro- cessing systems 29: Annual conference on neural information processing systems 2016, december 5-10, 2016, barceloma, spain (pp. 1849-1857). Retrieved from http::t:/tppceeednnnnoee .eurips ccc/ httpppppe//////aaa.l /tt0/////aa////10000aaaaaaa22210 2as-t2200aaaa---aaaaaatt.t. html
Sønderby, C. K., Raiko, T., Maalóe, L., Sonderby, S. K., & Winther, O. (2016). Ladder variational autoencoders. Advances in neural information processing systems, 29.
Souza, V. F., Cicalese, F., Laber, E. S., Molinaro, M. (2022). Decision trees with short explainable rules. In Neurips. Retrieved from http::/ nips cc/ pape/p Latees/paper/2000 witaafff555 5fbaba Ceanstraact-ooffeeeeee html
Steinbach, M. M.S Karypis, G., & Kumar, V. (2000). A comparison of document clustering techniques..
Suárez, A., & Lutsko, J. F. (1999). Globally optimal fuzzy decision trees for classification and regression. IEEE Trans. Pattern Anal. Mach. Intell., 21(12), 1297-1311. Retrieved from https::/ // doi wtt.:o../...00..101 1109/110//4.... /4.14999
doi: 10.10/944.817400
Tanno, R., Arulkumaran, K., Alexander, D., Criminisi, A., & Nori, A. (2019). Adaptive neural trees. In K. Chaudhuri & R. Salakhutdinov (Eds.), Proceedings of the 36th international conference on machine learning (Vol. 97, pp. (166-.175). PMLR. Retrieved from https: mlr. wwwpppeeeeessaeeeee http:////ttannnnnaaaaa. html
Vahdat, A., & Kautz, J. (2020a). Nvac: A deep hierarchical variational autoencoder. In Proceedings of the 34th international conference neural information processing systems. Red Hook, NY, USA: Curran Associates Inc.
Vahdat, A., & Kautz, J. (2020b). Nvae: A deep hierarchical variational autoen- coder. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, & H. Lin (Eds.), Ad- vances in neural information processing systems (Vol. 33, pp. 19667-19679). Curran Associates, Inc. Retrieved from ltttp::// httpps://proceedingss new. htt.a/ips...c//ppp/////////oo//e httpppoeeddings..eiirs......e pddf
van den Oord, A., Li, Y., & Vinyals, O. (2018). Representation learning with contrastive predictive coding. CoRR, abs/1807.0374l Retrieved from hattp:///rxxi.......//// org/ atg/abs/1807. 03748
Vikram, S., Hoffman, M. D., & Johnson, M. J. (2018). The loracs prior for vaes: Letting the trees speak for the data. ArXiv, abs/1810.06891.
14

---

Wan, A., Dunlap. L., Ho, D., Yin, J., Lee, S., Petryk, S.,.. Gonzalez, J. E. (2021). NBDT: neural- backed decision tree. In 9th international conference on learning representations, ICLR 2021, virtual event, austria, may 3-7, 2021. OpenReview.net. Retrieved from hommtttss://ppereviiee hett/forumrccmmm....ipppp..
Ward, J. H. (1963). Hierarchical grouping optimize objective function. Journal of the American Statistical Association, 58, 236-244.
Webb, S., Golinski, A., Zinkov, R., Narayanaswamy, S., Rainforth, T., Teh, Y. W., Wood, F. (2017).
Faithful inversion of generative models for effective amortized inference. In Neural information processing systems.
Williams, C. (1999). meme approach hierarchical mixture modelling. Advances in Neural Information Processing Systems, 12.
Wu, Z., Xiong, Y., Yu, X., & Lin, D. (2018). Unsupervised feature learning via non con-parametric instance discrimination. In 2018 IEEE conference on computer vision and pattern recognition, CVPR 2018, salt lake city, ut, usa, june 18-22, 2018 (pp. 3733-3742). Computer Vision Foun- dation IEEE Computer Society. Retrieved from hotpppttpp/paacccsssst thecvf com/ content html doi:
_cvpr_20 2018/0011tttt Wu_
10.1109/CVPR.2013 8.0093
Xiao, H., Rasul, K., & Vollgraf, (2017). Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. CoRR, abs/1708.07744 Retrieved from http///xii..oog//ba/////// arb/
1708.07747
Xiao, T. Z.,& & Bamler, R. (2023). Trading information between latents in hierarchical variational autoencoders. The eleventh international conference on learning representations. Retrieved
You, C., Robinson, D. P., & Vidal, R. (2015). Scalable sparse subspace clustering by orthogonal matching pursuit. 2016 IEEE Conference on Computer Vision and Puttern Recognition (CVPR), 3918-3927.
Zharmagambetov, A., & CurreiraPPerpinaa M. Á. learning with decision trees: Graph laplacian tree alternating optimization. In Neurips. Retrieved from http://papers: wwwp..ccpapp.... filfffills/ffllllf
(2022). Semi-supervised
wwwes//ppee/eee// htlpppppp///00 papprr/022/haah) coteee----bsttaatt---ereeeeeee html
Zhou, 7.-H., & Feng, J. (2017). Deep forest: Towards alternative deep neural networks. In
Proceedings of the twenty-sixth international joint conference on artificial intelligence, IJCAI- 17 (pp. 3553-3559) Retrieved from https://doi. org/ 10 0tt.22493//ijce hwtoe//ijcaa.. 11496////cc1.201744557
doi:
10.24963/jcaii21114447
15

---

A Evidence Lower Bound
In this section, provide closer look at the loss function of VAE. We focus the derivations of the Kullback-Leibler divergence term of the Evidence Lower Bound and provide an interpretable factorization. Furthermore, we address the computational complexity, thus offering an in-depth understanding of the loss function, its practical implications, and the trade-offs involved in its computation.
A.1 ELBO Derivations
[2). which the Kullback-Leible divergence (KL) between the prior and the variational posterior of VAE. Additionally, we give details about the underlying distributional assumptions for computing the reconstruction loss.
In this section, we derive the KL loss (17) of the ELBO
Let us define P/ the decision path from root O to leaf 1, L is the number of leaves, which is equal to
the number of paths in T, zp; {zz | PP}} the set of latent variables selected by the path P;, the
parent node of the node as pa(i). p Cp (1) the probability of going from pa(i) to i. For
example, if consider the path in Fig. 2 (right) we will observe co ci 1, and c4 where
= O means the model selects the left child of node i. The KL loss can be expanded using Eq.
KL (Zp:: Pt x)) lpp (ZP,, Pt))
(21)
KL g((0 x)
II ((Cpa(t))== x)q(z; Zpa(t))
kEPAMO)
(22)
(|((z))
II iEPr((0)
((Cpa(1)-1 Zpo()))((z =pa(t))
= _
g(z
x)
II (Cpa(())+s Φ)q(Zz ppa(s))
IEL ZP
tEP:((0)
(23)
<(20 x) [ii (((pat))=> x)) x))(Zz)
Epa(s))
x log
p(zo)
IIKEP (() #po(k)) (p(())))((z Zz
*pa(k))
g(z0 x))
Σ
x)
(Cpa((i)+s x)q(z; Zpp(s)) log
(24)
p(zo)
IEL ET
(EP:(())
II (0)E [INEP: ((Cpa(k)--x =pa(k))
x))
()) II (Cpa(t))>s ()q(zz zpa(s)) log
(25)
REPi\(U)
II {(z, >xp(s))
Σ/
x:)
II ((Cpa())= x)g(2; Z)a(t)) log
(26)
[IKEP) ((0 p(z* xpp(k))
Tel. sp,
REPi\/0)
In the following, we will simplify each of the three terms 24
25 and 26 separately.
A.1.1
KL. Root
The term (2 corresponds to the KL of the root node. We can integrate out all the latent variables z; for o 0 and all decisions c. The first term can be then written as follows:
(2 xx)
KLroot
Σ
d(z0 x)g(co-4s zu) log
(27)
p(zu)
sE(1,)
/
"y(0 xx)
9(0 x))
Σ a(co-ri | 20) log
(28)
p(zo)
20
(E(1,2)
/ g(20
((0 x))
()) [g(Co = 0 2n) + g(ooo = 1 20)] log
(29)
p(20)
20
[((0 xx)
g(z0
x)) log
= KL (g(20 .))pp(z))),
(30)
p(zn)
16

---

where g(cu 0 zo] + g(cu 20) and KL (g(z) ())[p(z))) is the KL between two
Gaussians, which can be computed analytically.
A.1.2 KI. Decisions
The second term log, yiclding
25
corresponds to the KL of the decisions. We can pull out the product from the
KT decissions
Σ IEL #PL
g(20
x)) II iEP:M()
((Cpa(t)++s x))(2z Zpa(t))
(31)
((pa(t))+} x))
log
II jEPiMO)
P(Cpa(i)->) Zpals))
- -E/ TEL (EPi((0)
E
((0 ()) II (Cpaa())+ x)q(Zz ppa(s) log ((Cpali))>3 | xx) (p(epa()) ppa(1)) (32)
(EP()(0)
Let us define as Piej all paths that go through node j. as PEy (denoted as P, in the main text for brevity) the unique path that ends in the node and as Pa5 all the possible paths that start from the node j and continue to a leaf E Similarly, let us define as ž5y all the latent embeddings that are contained in the path from the root node and ø>j all the latent embeddings of the nodes i that can be reached from node j. To factorize the above equation, summing over all possible leaves in the tree (Zeee) and then over each contained node (jer ((0))
first change from pathwise view nodewise view. Instead of
we sum over all nodes (o))v((o)) and then over each path that leads through the selected node
EPrES
q(Cpa(5)->> | xx)
Σ
Σ
g x)) II ((poa(t)-a x)g(2; Zpa(t)) log
(p(ppl()) -s *xp((i))
IEL 2P1 (EPA((0)
tEPi\(0)
(33)
1(20 x:) II (Cpa(())) a))(2; zpa(i)) log g(Cpa(5)=+) ())
Σ
Σ
((Cpa Zpali))
(EV Pres
tEP)(0)
The above can be proved with the following Lemma, where we rewrite Exies EvEL. 1/ E P.).
Lemma A.1. Given binary tree T defined in Section 2.7
composed of of nodes V {0, ...,, V} and leaves LL C V, where Pi is the decision path from root 0 to leaf I, and zp, {z; i Pi) the set of latent variables selected by the path Pt. Then it holds
Σ/ Σ 1011,2)= IEL EPI (OP:)(0)
Σ Σ9
1 E 1,)
(34)
(0V LEL. 2P,
17

---

Proof. The proof is as follows:
Σ JEV(O) TEL. km
e P)) E
E/
f(11(Ep))
E
1i= j]
(35)
JEV((O) TEL.
(EP!(()
Σ Σ
E
E
2p,
(36)
IEL HEV
SPI iEP()())
= Σ
Σ
E
f(1,1,2,)) 5
(37)
LEL (0)
(EPI(()
Σ
f(L,,22)
Σ
1|i= 5
(38)
lEL
API iEP:YO)
JEV
Σ
Σ
f(i,4zp))
(39)
TEL IPI (EP:
E/
E
(40)
#PI JEPi
L
Having proven the equality, we can continue with the KL of the decisions as follows:
KLdecisions
Σ
g(20
x)) II (((ppa())= [))((2) 3pa(s)) (aa(s)))oggg log
((pra(t)) -2 | x))
(41)
p(Cpo(i) Zpa(s))
JEV((0) Pies EP,
iEPiMO)
Σ Σ
[a(20 x=) II ((p x))(zz Zpa(t)) log
((Cpa(5))>> x))
p(Cpa(()=+) ppa(t)) (42)
(0) Pies
2PL
(EPs,\(0)
x
II
(Cpp((k))>k
x)){(z 2pm(k))
kEPs>
g(Cpa(i)+g) (x)
= Σ Σ
19(=0 x)) II (Cpp(())-> x)g(z; xpa(s)) log (o))
((ppatj) 2pa(s))
JEV } Pas
(43)
x
II
((Cpa(k))-s
x x){((z xpa(k))]
kEPs5
g(z ()) II ((Cpa(t)++s x)q(2; I Zpa(t)) log ((Cpat))>>s 1 x))
Σ
p((pa(t)=3 Rpa(s))
[EV
iEPs)\(0)
(44)
Σ
II kEPs,
( (Cpo(t)) x))(Zz =po(ks)
P>1 E>3
From Eq. 41 to Eq.
42 we split the inner product into the nodes of the paths Piej that are before and
after the node j.
From Eq. 42 to Eq. 43 we observe that the sum over all paths going through can be reduced to
the sum over paths starting from j. because there is only one path to j. which is specified in the
product that comes after.
From Eq. 43 to Eq. [44] we observe that the sum over paths starting from j and integral over z> concern only the terms of the second line. Observe that the term on the second line of Eq. 44
18

---

integrates out to and we get
KT decisions
Σ
(20
a)) Il ((po(()) x)g(2; 3pa(i)) log (Pg))(0)
((Cpa(t) m2)
(45)
p((po(s)={ ?pa(t))
JEV([O)
#5J
[(Cpo(j)>>3 2
Σ
((20 x3) II ((Cpa(i)) (0)((z 2pa(t)) log
(46)
((Cpa(()->j ppa(si)
1EV(10)
iEPc)
/
[(Cpa(j)+t) 1 x
E
((20 a II (Cppa(i)->i x)g(2z xpa(t)) log
p(Cpa(j))>) 3pa(s))
[EV 2Cj
iEP (00)
(47)
x ((Cpa())>>) x))((z} Epa(j))
Σ [EV =cs
/ ((20 (+)
II iEP)s((0)
1 Spa(t) Spo(t))((Cpa((-1) x))
(48)
{(Cpa(i))>> x))
x log
((Cp(())+5 ppa(s))
From Eq. to Eq.
47
we single out the term in the product that corresponds to j = i, which is the
only term that depends on Ja.-
From Eq. 47 to Eq. 48 we observe that in the singled-out term, Sa, {(zz xpa(s)) 1, which leaves
only ((ppat)))>) x)
This equation can be rewritten in more interpretable way. Let define the probability of reaching node j and observing 2j as
P(j: 2, = 9(20 x))
II
(Cpo(())+ x))((2; =pa(i))-
(49)
iEP=;Mot
Then the KL term of the decisions can be simplified as
log g(Cpa(t))t) x
KT ddcissions
Σ
2,
(50)
jEV((0) =EJ
(((ppa()) Zpa(i)
E
Σ LEV)LEE(O.1) Fes
P(izz,c)qle = x)) log
q(c k
(51)
p(c; = 4: zi
This term requires Monte Carlo sampling for the expectations over the latent variables while we can analytically compute the sum over all decisions Pt.
KT diccssions
Σ
P(i:z,
LEVIL
Zer
(52)
q(c; 0 2x)
[((; = 1 ②)) p(c; 2;)
X a(c: = x)) log
plc; = 2;) + g(c, 1| [1x) log
Σ
P(i:z(m) c)
M
iEViL
(53)
a(c; 0 0
g(c: = 1 x3)
x g(c. 0) log
q((; =1 | x)
p(c; = 0 2, (m))
p(c; = 1 (m)
where P(i; (m) IJEP3, ((Cpo(5)-33 x) is the probability of reaching node i, defined as
P(ū; c) in Eq. 16 1919 20 for simplicity.
19

---

A.1.3 KL Nodes
Finally, we can analyze the last term of the KL term, which corresponds the KL of the nodes (26).
The reasoning is similar to the equations above and we will use the same notation. The KL of the nodes can be written as
Klaoodes
Σ
8(50 x) II ((Cpa(E)~i x)q(Zz ppa(i))
iel. 17)
iEPA/O)
(54)
[I g(2)| spa(s))
x log
[IKEP\(0)) p(zz >pa(k))
<=0 x) II ((Cpa(t)=+s x)q(Z= Ppa(t)) log [((;} Epa(j))
(55)
p(zz Zpa(j))
JEL
EP (EP))())
(EP)M()
We now change from a pathwise view a nodewise view.
Σ E/ JEV Piej
x) II ((Cpa(i)~> x))(2z xpa(s)) log ((2) ppa(1))
((=0
(56)
((z)) Apa(1))
iEPA)))
II (Cppa())>> x)g(2; xpa(t)) (oe(()))oggg log [(2; 2pa(i))
Σ
g(z0
x
p[zz ppa(j))
JEV
(57)
E/
(Cpo(k))+* x))((z Epa(k))
Pzj 233
kEP>
Σ
g(z0 x) II ((cpa(t)<r x)q(zz =pa(t)) log
g(zz
ppolj))
(58)
p(z;
2pa(j))
JEV\(O}
iEPs>\(0)
Σ
2) x){(zz =pa(t)) log
{(z;
z)a(;i)
(59)
Res J2,
p(zz
2pa(t)
JEV
Σ JEV)(O)
P(pa(j); clop(p))= (x) q(zz =pp(1)) log
{(z}
pp)
(60)
zes
p(2)
ppa(1))
Σ
P(pa(j); 2))((ppoc))) (x) KL (a(z) =pals)) p(z; xpo(s)) ))
(61)
*cs
M
Σ
Σ P(pa(i); ())))))((pp(())))) x)
M
(62)
m=1
10)
<L(g(> (m)
ppa(z (p)pp())))lpp( (m)) pa(z!"))))
1 M
Σ
Σ
(m) c) (L(( (T)) (m)()("))) ([p(( (mm) pa(2 (m)) ))),
(63)
M
m=
where P(pa(j); 2, e) is defined in Eq. 49 and where P(i; (1m) .)) = P(i;e) = IIEPSS A((pa())=3) x) is the probability of reaching node
20

---

A.1.4 KL terms
Using the above factorization, the KL term of the ELBO can be written as
KL (q (z, P.
x) (2, KL (9(20
2: llp(z)))
+
1
Σ
E
(m) c)
Σ
g x)) log
q(cs ())
M
p(c; 2 2((m)))
m=LiEVL
c
(64)
Σ
Σ
(m)),
(L(((z) (m) pa(z!"))) Ip(z) pa(z) (())(m))))
M
m= iEV (0)
where P(i; (m))
c)) IlsePs (CCoott)~3 x)).
A.1.5 Reconstruction Loss
Finally, to compute the full ELBO, the KL terms are added to the reconstruction loss defined in (16). Here, assumptions about the distribution of the inputs are required. For the grayscale datasets such as MNIST, MINSS, and Omniglot, as well as the one-hot-encoded 20Newsgroup, we assume that the inputs are Bernoulli distributed, such that the resulting reconstruction loss is the binary cross entropy. On the other hand, for the colored datasets CIFAR-10, 10, CIFAR-100, and CelebA, we assume that their pixel values are normally distributed, which leads to the mean squared error as loss function, where we assume that 1.
A.2 Computational Complexity
All terms of the Evidence Lower Bound, Equation 12 can be computed efficiently and the compu- tational complexity of single joint update of the parameters is O(M BV DCp). where M is the number of MC samples, B is the batch size, V is the number of nodes in the tree, D is the maximum depth of the tree, and C₂ the cost compute the KL between Gaussians. It should be noted that the computational complexity is, in practice, reduced to O(LBVC,), as the term P(i; c) can be computed dynamically from parent nodes.
B Related Work
In addition to the review presented in Section 3 which encompasses a broad range of related work in the field of deep latent variable models, we provide review of relevant work in the domains of hierarchical clustering and decision trees. By doing so, we hope to shed further light on the current state-of-te saiee-offthhee ttt-hhe-rtt approaches and contribute a deeper understanding of the challenges and opportunities that lie in the intersection of hierarchical clustering, decision trees, and latent variable models.
B.1 Hierarchical Clustering
Hierarchical clustering algorithms have long been employed the field of data mining and machine learning extract hierarchical structures from data (Sneath) 1957; Ward 1963) Murtagh & Contreras 2012). Agglomerative clustering is among the earliest and most well-known hierarchical clustering algorithms. These methods start with each data point as individual cluster and then iteratively merge the closest pairs of clusters, according to a predefined distance metric, until a stopping criterion met. While single-linkage and complete- baselines, we observe better performance when using the bottom-up strategy proposed by Ward (1963). Ward's s minimum variance criterion minimizes the total within-cluster variance (Murtagh & Legendre 2014). thus providing balanced and compact clusters. In contrast, the Bayesian Hierarchical Clustering (BHC) proposed by Heller & Ghahraman] (2005) takes different approach by employing hypothesis testing to determine when to merge the clusters. The divisive clustering algorithms, on the other hand, provide a different strategy to hierarchical clustering. Unlike agglomerative methods, divisive clustering starts with all data points in single cluster and recursively splits clusters into smaller ones. The proposed VAE an example of a divisive clustering method. Among variety of proposed methods, the Bisecting- Bisecting-K-means algorithm (Steinbach et al. 2000 Nistér & Stewénius
agglomeration clustering are widely used as
21

---

2006) is widely used for its simplicity; it applies k-means with two clusters recursively. Similarly, Williams (1999) learn a hierarchical probabilistic Gaussian mixture model. Neal (2003) introduce a hierarchical generalization of Dirichlet process mixture models.
More recent approaches include PERCH (Kobren et al. 2017b), which is non-greedy, incremental algorithm that scales to both the number of data points and the number of clusters, GHC (Monath et al. 2019), which leverages continuous representations of trees in hyperbolic space and optimizes a differentiable cost function, and Mathieu et al. (2019), which similarly use a Poincaré ball model of hyperbolic geometry as latent space to learn continuous hierarchical representations with a VAE.Fuurherrore, RSSCOMP (You et al. 2015), which explores a subspace clustering method based on an orthogonal matching pursuit. Finally, Deep ECT (Mautz et al.) (020)) proposes divisive hierarchical embedded clustering method, which jointly optimizes an autoencoder that compresses the data into an embedded space and a hierarchical clustering layer on top of it. Density-based clustering algorithms, such as DBSCAN, belong to distinct category of clustering techniques, They aim to identify regions in dataset where points are densely concentrated and classify outliers as noise (Ester et al. 1996). R. J. Campello et al. 2015)) build on this idea to learn a hierarchy based on the distances between datapoints where distance is roughly determined by the density. However, one limitation of density-based methods lie in their performance when confronted with complex datasets requiring high-dimensional representations. In such cases, estimating density requires an exponentially growing number of data points, which leads to scalability issues that TreeVAE does not have.
B.2 Decision Trees
Decision trees (Breiman et al. 1984) are interpretable, non-parametric supervised learning techniques commonly employed in classification and regression tasks. They rely on the data itself to build hierarchical structure. This is done by recursively partitioning the data into subsets, each of which corresponds to specific node in the tree. At each node, decision rule is generated based on one of the input features that best discriminates the data in that subset. One of the key advantages of decision trees is their interpretabiliiy. The learned tree structure can be easily visualized to get insights into the data and the model's brtision-making process. [Suárez & Lutsko (1999) argue that deterministic splits lead to overfitting and introduce fuzzy decision trees with probabilistic decisions, implicitly allowing for backpropagation. With the advancement of neural networks, many works (Rota Bulo & Kontschieder 2014 Laptev & Buhmann 2014 Frosst & Hinton (017)
leverage MLPs or CNNs for learning a more complex decision rule. However, the input at every node remains the original features, which limits their performance, as they are unable learn meaningful representations. Thus, Tanno et al. (2019) introduces Adaptive Neural Trees (ANT), a method that learns flexible, hicrarchical representations through NNs, hereby facilitating hicrarchical separation of task-relevant features. Additionally, ANTs architectures grow dynamically such that they can adapt to the complexity of the training dataset. At inference time. ANTs allow for lightweight conditional computation via the most likely path, leading to inbuilt interpretable decision- decision-making.
While decision trees were initially designed with the goal of achieving high predictive performance, they are also used for many auxiliary goals such as semi-supervised learning (Zharmagambetoy
& Carreira-Perpihann 2022), robustness (Moshkovitz et al. 2021) or interpretability (Wan et al. Souza et al.) 2022] Arenas et al. 2022 Pace |l.] 2022). In the context of interpretabilliy, various approaches have been proposed to enhance accuracy and interpretability. |Wan et al.) (2021) introduce Neural-Bace Backed Decision Trees (NBDTss, which improve accuracy by replacing the final layer of a neural network with a differentiable sequence of decisions to increase its interpretability while retaining predictive performance. [Souza et al. (2022) focus on optimizing the structural parameters of decision trees, introducing the concept of "explanation size" as being the expected number of attributes required for prediction to measure interpretabilityy |ace et al. (2022) develop the POETREE framework for interpretable policy learning via decision trees in time-varyin mireingyingg clinical decision environments. AAreas et al. (2022) investigate explanations in decision trees, considering both deterministic and probabilistic approaches and showing the limitations thereof. These works collectively contribute to advancing the use of decision trees in interpretable machine learning, providing insights into trade-offss, criteria, and frameworks for improving accuracy and interpretability.
22

---

While decision trees are most commonly known for their application in supervised tasks, they can be repurposed to partition the data space into clusters in an unsupervised way (Blockeel & De Raedt 1998
B. Liu et al. 2000 Basak & Krishnapuram 2005 Ram & Gray 2011 Fraiman et al. (013) Most methods, however, have the downside not learning meaningful representations for their splits, such that they are unfit for modelling complex interactions. Therefore, similar to Tanno et al. (2019) in the supervised setting, in this work, we combine the simplicity and interpretability of clustering decision trees with the flexibility of NNs.
C Further Experiments
Here, we provide additional experimental results, including further hierarchical structures, uncondi- tional generation samples, and ablation experiments.
C.1 Discovery of Hierarchies
We show further hierarchical structures learned by Tree VAE MNIST in Fig. 211 Fashion-MNIST Fig. [1 Omniglot-50 in Fig. 12 and CIFAR-100 Fig. 13 Additionally, the figures that include conditional generations generated by sampling from the root and following the most probable path to generate every leaf-specific image.
9 9
9
6
6
6
2
6
6
6
2
9 9
9
7
6
6
6
2
5
5
8 8
8
3
3
3 3
5
5
5
8
8
8
3
0
5
5
8
8
8
3
Figure Hicrarchical structure learned by Tree VAE with ten leaves on MNIST dataset with generated images through conditional sampling. The right subtree contains rounded digits, while the left subtree contains digits with a straight line.
3 6 6
Figure 10: Hierarchical structure learned by Tree VAE with 20 leaves on MNIST dataset with generated images through conditional sampling. The digits are further divided according to style.
C.2
Generations
We further show how Tre can be used to sample unconditional generations for all clusters simultaneously for MNIST, in Fig. [4] and CelebA, Fig. T5] Differently from the conditional generation, where we follow the path of the tree with the highest probability, here we follow all
23

---

Figure 11: Hierarchical structure learned by Tree VAE with 20 leaves on Fishion-MNIST dataset with generated images.
b0
6 8
0
C h
6
8
P
*
1 #
x 7
M
40
4
3/
A 8
31
2
M
a
Hierarchical structure learned by Tre on the full Omniglot dataset. We display random subsets of images that are probabilistically assigned to each leaf of the tree. Similar to the results in the main text, we can again find regional hierarchies in some individual subtrees. For instance, the leftmost subtree seems to find structures in different Indian alphabets. We can also observe that this subtree seems cluster smaller characters its left child, whereas the right child contains bigger shapes. Another example is the rightmost tree, which seems to encode the more straight shapes of Japanese writing styles. There, left child encodes more complex shapes that contain many different strokes, and the right subtree groups simpler shapes with only few lines.
Figure 12:
Hierarchical structure learned by TreeVAE with 20 leaves on CIFAR- dataset. We display random subsets of images that are probabilistically assigned to each leaf of the tree.
Figure 13:
paths of the tree regardless of their probabilities. a result, we generate L images, where is the number of leaves, corresponding to the number of decoders. These generated samples exhibit distinct characteristics based on their respective cluster-specific features while maintaining cluster- independent properties across all generated instances.

---

2
6
4
7
9
3
8
5
2
6
7
4
7
9
3
8
6
O
2
6
4
7
9
3
8
5
0
2
6
7
9
3
8
5
0
leaf 1
leaf
leaf 3
leaf
leaf 5
leaf
leaf
leaf 8
leaf 9
leaf 10
Figure 14: Selected unconditional generations of MNIST. One row corresponds one sample from the root, for which we depict the visualizations obtained from the 10 leaf-decoders. Each row retains similar properties across the different leaves. In the first row, all digits are rotated; the second are bold; the third are straight; and the last are squeezed.
leaf
leaf 2
leaf
leaf 4
leaf 5
leaf
leaf 7
leaf 8
Figure 15: Selected unconditional generations of CelebA. One row corresponds one sample from the root, for which we depict the visualizations obtained from the 8 leaf-decoders, The overall face shape, skin color, and face orientation are retained among leaves from the same row, while several properties (such as beard, mustache, hair) vary across the different leaves.
C.3 Ablation Study
In this section, we provide additional studies on the effect of selected hyperparameters and assump- tions on the behavior of VAE In Table we provide ablation selected hyperparameters of Tre on Fashion MNIST. In the following subsection we will provide further in-depth explo- rations of Tree
25

---

Table 4: Ablation Study on Fashion MNIST: (A) dimensionality of latent embeddings ordered by depth, (B) dimensionality of MLP and Router hidden layer ordered by depth, (C) Increasing dimensionality, (D) number of hidden layers of MI., (E) number of hidden layers of router, (F) pre-determined maximum depth of tree.
Adayth
Paye
L Routeryye depth
Acc
NMI
ELBO
base
8,8, M. 8, 8,
128, 125, 128, 128, 128
2
63.0
63.9 2.2
289.2 + 0.5
(A)
2.
2.2
49.8 + 4.0
2.7
2299.5
0.4
(A]
4,4,4,4
60.0
61.03.7
-242.6
16,16,
16,1
630+24
-
62.7 63.5+28
-238.7
-241.
(8)
32. 32.
32,32
60.6 5.6 62.7=2.6
-240.4 0.4
(B)
64. 64,
64,66 64
62.6+388 63.5+14
-209.0 0.3
(B)
256,
256,
256,
256, 256
62 63.1
-239.2
(C)
64,
32,
8,
128, 64,
16,
56.7 59.9
-239.2 2.0.210..4
64,
-
(C)
128
32,
8,
128, 62,
58 613
-239. 139.11004
(C)
128,
64,
52,
16, 8,
256,
125,
64, 32,
62.6 63.6
-238.8 2.3
(C)
256,
128,
64.
18.8
256.
64,
64, 32,
611+ 61.7=29
128,
64,
16,8
256
64
(D)
61.4 02.9333
-339.0 0.3
(8)
39.6 623+20
-23922
(E)
62.1 + 4.4
63.6 6 3.1
-239.1 + 0.4
(
8,8,
128, 128. 128, 128
64.0 64.4
-239.2
(P)
8.
8,8,8,
128, 1288, 128, 128, 128 128. 1288, 128, 128, 128.
61.0 6294 2.9
-239.2
62.1+ 6.76+2.7 -239.1
C.3.1 Training time parameter matching
Tree has the advantage of lightweight inference, each datum ends up in one leaf. Still, have to learn the full tree with all paths during training. Thus, we have conducted an ablation study on Fashion MNIST that increases the effective number of parameters of LadderVAE to match the parameter count Tree VAE at training instead inference. The results are presented Table 5
Table 5: Test set performances (%) of TreeVAE and LadderVAE with matched training parameter count. Means and standard deviations are computed across 10 runs with different random model initializations.
Dataset
Method
ACC
NMI
LL
ELBO
Fashion
LadderVAE Agg
1..4
62.3 2.0
-233.7+
221.8 = 0.3
-238.9 0.3
Tre
63.6
64.7
-231.7 +0.1
226.5
-239.2+
C.3.2
Tree Structure
First, we analyze the necessity of the iterative growing procedure as opposed fixing an (informed) tree structure priori. Thus, we present ablation, where we fix reasonable tree structure and directly train the full tree as opposed to iteratively grow. The results Table 6 suggest that the iterative learning is vital to the performance of TAE. Intuitively, by fixing a tree structure, the model loses the ability to learn one out of multiple equivalent tree structures (c.g. left and right subtrees can be exchanged) and instead needs to fit one specific structure. This leads to less flexibility and more influence of the random weight initialization, as well as local optima.
Second, we analyze whether the learnt tree structure is informed through the learnt embeddings or rather an artifact of induced biases. For this, we show the UMAP (Ghojogh et al. (02T) representation of the learnt root embedding in Fig. [6] As can be seen, tops well as shoes are clustered together, while bags are somewhat between. Trousers on the other hand are completely separate. Our interpretation is that the root split is made between shoes and tops, where it is unclear in which subtree bags and trousers should fall as they are both not clearly assigned to one of the two groups. Therefore, depending initialization, they might end up either side. To analyze the effect of
26

---

Table Test set performances (%) of Tree VAE with a fixed structure standard Tre. Means and standard deviations are computed across 10 runs with different random model initializations.
Dataset
Method
ACC
NMI
LL
RL
ELBO
Fashion
fixed VAE
4.8
54.1
-237.4 20.
228.9
-211.0 111.0
Tre
6.3
61.7
-231.7+0.1
226.5
-239.2
random weight initialization on the learned dendrogram, we have made further efforts to find dendrogram that best summarizes the learned dendrogram of multiple runs. We show it in Fig. [7 For this, we first align the leaves of the different trees by maximizing the overlapping samples, then we store the number of edges between any clusters for every tree and then average this number over all trees. Thus, we have computed a distance matrix of all clusters, which is averaged over all trees. We then used average and complete linkage cluster according the average distance matrix. The recovered dendrograms of the two algorithms are identical, apart from the Bag and the Trouser cluster, which switch places. That is, all tops are in one subtree and all shoes in the other, while bags and trousers are assigned to either the shoes tops subtree, depending the clustering algorithm used. What this indicates is that the groups shoes and tops are consistently recovered, no matter the random initialization, while the assignment of bags and trousers varies, which is aligned with the interpretation of the root embedding in Fig. 16. as there is no clear assignment thereof.
Anbll bost
12
Beg
10
D3666
Pul Pulover
1.0000 250 2002
troueer
Figure 17: "Average" dendrogram learnt on Fashion MNIST by averaging cluster distances across 10 runs and applying average linkage clustering.
Tisiit
Figure 16: UMAP visualization of root embed- ding for Fashion MNIST.
C.3.3 Contrastive Losses
In this section, we present an ablation on the effect of the contrastive losses that are applied in real world datasets. The experiments are performed on slightly modified version of CIFAR-II and shown in Table 6 As can be seen, the experiments showed that the proposed combination of (A) and (B) leads to the best results the methods complement each other.
C.3.4 Unknown number of clusters
In an unsupervised setting, is unclear a priori what the optimal number clusters is. Thus, we are interested in the sensitivity of Tre with respect to misspecifications of the ground truth number of clusters. The CIFAR- 100 dataset is an optimal candidate for this ablation, as the ground truth number of clusters can be defined differently depending on whether classes. As a metric, we will only use NMI, because DP, LP. as well Accuracy values are not fairly comparable across different number of predicted clusters. To determine the sensitivity of Tre with respect to the misspecification of number of clusters, we evaluate its performance with 10, 20, and 32 predicted classes, compared to ground truth 20 superelasses. The results are depicted in Table 8 Clearly, Vre is not very sensitive to the correct specification of the number of clusters. This also supported by Fig. TO IT which show that even as we grow the tree further than the ground truth number of classes, VAE VAE still finds meaningful subsets of inside the classes.
look at coarse or fine-grained
27

---

Table 7: Ablation study of contrastive losses. (A) corresponds to the NT-Xent regularization on the routers, (B) to the NT-Xent regularization on the bottom- embeddings, (C) on transposed
version of the NT-XXen regularization on the routers, following Li al. (2021), (D) to the NT-Xent
regularization of only the output of the encoder, and (E) to minimizing the Jensen-Shannon divergence between the router probabilities the two augmented inputs.
Dataset
Contrastive Method
NMI
CIFAR-I0*
Ours (=A+B)
44.1 0.7
(A)
21.7 4.0
(B)
6.0 6.0
(C)
23.0
(D)
0.0
(E)
1.0
Table 8: Test hierarchical clustering performances (%) of Tree VAE for different number of predicted clusters. Means and standard deviations are computed across 10 runs with different random model initializations. The star "*" indicates that contrastive learning was applied.
Dataset
Predicted clusters
NMI
CIFAR-100
15.99 0.86
20
17.80 0.42
32
19.33 0.41
As a next ablation, we are interested in the case of the number of ground truth clusters being misspecified unclear. Note that the difference to the previous ablation is whether the ground truth or the predicted number of classes is misspecified. Again, CIFAR-100 is a natural choice for this ablation, the ground truth number of classes either be 20 100. Thus, in Table how the performance of TreeVAB would change if we chose a different granularity of ground truth classes while keeping the number of predicted clusters fixed 20. We can observe that the NMI for finer granularity is even higher than for the superclasses used, indicating that Tree learned to differentiate classes of the same superclass.
present
D Datasets
This section provides supplementary background information and preprocessing steps of the different datasets used in our experiments.
D.I MNIST
The MNIST (LeCun al. 1998) dataset is widely used benchmark dataset in machine learning. is composed of grayscale images representing handwritten digits from zero to nine. Each image of size 28 28 pixels that rescale [0, 1]. The dataset is balanced, meaning that there is an equal number of examples for each digit class, with total of 60'000 training images and 10'000 test images. The digits in the dataset are manually labeled, providing ground truth information for clustering evaluation.
D.2 Fashion-MNIST
The FasHion-MNIST (H. Xiao al., 2017) dataset consists of a collection of grayscale images depicting various clothing items and accessories. The dataset contains balanced distribution of ten different classes, representing different fashion categories such as T-shirts, trousers, dresses, footwear, and more. Each image in the dataset has resolution of 28 x 28 pixels that we again rescale to [0, 1). Similar to the original MNIST dataset, Fashion MNIST also includes 60, 000 training images and 10'000 test images.
28

---

Table 9: Test set hierarchical clustering performances (%) of Tree VAE for different number ground truth clusters. Means and standard deviations are computed across 10 runs with different random model initializations. The star "*" indicates that contrastive learning was applied.
Dataset
Ground truth clusters
NMI
CIFAR-100
20
17. 17.80 0.42
100
22. 0.36
D.3 20Newsgroups
The 20Newsgroups (Lagg 1995) dataset is widely-used benchmark dataset in the field of natural language processing and text classification, known for its inherent hierarchical structure. It consists of approximately 20'000 newsgroup documents organized into 20 different topics or categories. These categories include sports. politics, technology, science, religion, and more. One of the most important characteristics of this dataset is the presence of hierarchical relationships within the topies. We employ TF-IDF vectorizer for the 2'000 most frequent words and use this text embedding as input to our model. We retain 60% of the original dataset for training and separate the other 40% for the test set.
D.4
Omniglot
The Omniglot Lake al. (2015) dataset contains handwritten characters from 50 different alphabets. Every alphabet contains various individual characters, amounting to total of 1623 different characters the dataset. Given the heterogeneous nature of the dataset, Omniglot contains relatively few samples. The dataset consists of 32' 460 samples, which is about 650 samples per alphabet or 20 samples per character. our experiments, we are interested in whether Tree can find alphahet-specific patterns that stay invariant across different characters. Given that many alphabets developed either from or parallel to each other, many alphabets are very hard distinguish, both for humans and machines. While we use the full Omniglot dataset for our generative experiments to have more data, we only use a subset of the alphabets for our clustering experiments. We call this subset Omniglot-5 and it consists of the five alphabets Braille, Glagolitic. Cyrillic, Odia, and Bengali. We selected them by ensuring that we have languages from different origins (artificial, Slavic, and Indian) and that there learnable hierarchies. Old Church Slavonic (Cyrillic) developed from Glagolitic and Odia and Bengali are languages with their own alphabet from geographically similar regions in India. The resulting Omniglot-5 results in 4'160 samples. For both dataset versions and all experiments, we split the dataset into train/test splits with 80%/20% of the samples, respectively. and stratifying across the characters of the dataset. We resized the images to 28 28 align with images in MNIST and Mashion-MNISTT transformed the images to grayscale, and rescaled pixel values to [0, 1]. For description of the augmentations we use, we refer Appendix E.1.2
D.5 CIFAR-10
The CIFAR-10 (Krizhevsky & Hinton 2009) dataset contains colored images of ten balanced classes. These classes are (Airplane, Automobile, Bird, Cat, Deer, Dog, Frog, orsee Ship, Truck), Therefore, some classes are more similar (for example Automobile and Truck) than others (for example Frog and Airplane), indicating that there a structure across clusters that can be captured. The dataset consists of 50 '000 training and 10'000 test images of resolution 32 x 32 pixels and 3 color channels, which we rescaled [0, 1].
D.6 CIFAR-100
The CIFAR- 100 (Krizhevsky & Hinton 2009) dataset originally contains colored images of 100 classes. order to not have to grow until Tree has 100 clusters, we apply the common grouping of the 100 classes into 20 balanced superclasses. In Table 9 we also present ablation if instead we chose the full 100 clusters. The dataset consist 500000 training and 10'000 test images of resolution 32 x 32 pixels and 3 color channels, which rescaled to [0, 1].
29

---

D.7 CelebA
The CelebA (Z. Liu et al., 2015) dataset contains images celebrities with 40 attribute annotations per image. This dataset is well-suited for exploratory analysis, as there are no agreed-upon ground truth clusters. As a result, we can visualize the clusters that are learned by VAE and validate whether they align with our intuition well as with the given attributes. Thus, in Table 3 additionally define "Hair Loss" as balding or having a receding hairline, and "Beard" as having a beard or the so-called 5 o' clock shadow. For training, we select subset of 100'000 random images from the training set and evaluate on the given test set 19'962 images. Lastly, we crop the images to 64 x 64 pixels with 3 color channels and rescale them [0, 1].
E
Implementation Details
The following section provides a comprehensive overview of the various aspects and components regarding the practical implementation of our proposed framework. We provide a detailed description of our proposed model's training process and architectural design. Furthermore, we outline the specific techniques and methodologies used for the contrastive learning extension. We also provide the implementation details of the Variational Autoencoder and the Ladder Variational Autoencoder that are used as baselines. Finally, present the computational resources required for executing the proposed framework. Together, these subsections aim to provide a comprehensive view of the implementation details, enabling readers to replicate and build upon our work effectively.
E.1 Training Details
Hierarchical variational models, such as LadderVAE, are prone to local minima and posterior collapse, and Tree shows similar behavior. Therefore, few technical choices, such as batch normalization and KL annealing, needed to converge desirable optimum. For all datasets, train TreeVAE only on the training set and evaluate the trained model on the separate test set.
E.1.1 Training the Tree
For MNIST, Fashion-MNIST, Omniglot, and 20Newsgroup, the trees are trained for M = 150 epochs at each growth step, and the final tree is finetuned for M, 200 epochs. To reduce the risk of posterior collapse during training, we anneal the KL terms of the ELBO. Starting from 0, we increase the weight of the KL terms every epoch by 0.001 except for the final finetuning of the full tree. Here, we linearly increase the weight by 0.01 until we reach 1, such that the final model is trained on the complete ELBO. Additionally, for the aforementioned datasets, we finetune the full third growing step for 80 epochs with KL annealing of 0.001, so the otherwise frozen parameters can adapt the new structure. For CIFAR-10, CIFAR- 100, and CelebA, the splits are mostly determined by the contrastive losses, and therefore, we avoid finetuning the full tree to reduce the computational time. The KI. term is annealed every epoch by 0.01 such that the weight reaches after 100 of the 150 epochs every growth step. While the annealing procedure is required to achieve good performances. the exact annealing decay plays smaller role the final performances, as well as the in between finetuning of the trees, and the number of epochs at each training step.
after every
E.1.2 Augmentations
Due to the scarcity and heterogeneity of Omniglot, we trained our models using data augmentation by randomly rotating images by up to 10 degrees, horizontally//erticallyl shifting images by up to +2 pixels, applying up to 1% shearing, and zooming by up to +10% Note that for Omniglot, no contrastive learning was applied; the augmentations were solely used for overcoming the scarcity of data points.
On the other hand, in CIFAR-10 10, CIFAR- and CelebA, augmentations were applied to use contrastive learning as introduced in Section 2.6 and Appendix E.3 For CIFAR-10 10 and CIFAR-100, first, each image is randomly cropped and resized with scale (0.2, 1) and aspect ratio E (3/4, 4/3). Second, each image is randomly flipped in the horizontal direction with probability 0.5. Third, with probability 0.8, color changes are performed, which consist of changes in brightness, contrast, and saturation, all with their respective factors E (0.6, 1.4), as well as change in hue with its parameter
30

---

(-1. Lastly, with probability 0.2, we transform the image to grayscale. The chosen parameters are largely adopted from Li et.js (022) weak augmentation scheme.
While these augmentations lead to superior group separation, in CelebA, we are also interested in generative results. Therefore, obtain generations without non-sensical artifacts, reduce the augmentation strength by removing the hue and grayscale augmentations, reducing the other color parameters by factor 2, and reducing the cropping parameters scale E (3/4, 1) and aspect ratio (4/5, 5/4).
E.2 Model Architecture
In general, model architectures for all datasets are similar and consist of an encoder, followed by the tree structure with the transformations, routers, and the bottom-up, as depicted in Figures and and identical decoders for each Icaf of the tree. We always apply batch normalization followed by Leaky ReLU non-lineanne linearities after all convolutional and dense layers. For further details on the specific implementation of the sections below, we direct the reader to consult the accompanying code.
E.2.1 Traasformaation/BottomUU
The architecture of the transformations and the bottom-up is identical across all experiments. They take as input the previous latent vector of dimension 4 (20Newsgroup), 8 (MNIST, Fashion-MNIST, Omniglot), 64 (CIFAR-10. CIFAR-100. CelebA) and map them to some intermediate representation of dimension 128 (20New MNIST, Rashion-MNIST. Omniglot), or 512 (CIFAR-10. CIFAR- 100, CelebA) using dense layer, followed by dense layer without activation to compute and one with Softplus activation to compute a of the approximate posterior/conditional prior of the respective tree level.
E.2.2 Routers
Similar transformations, the architectures routers are identical across all experiments. The generative routers take as input 2;, the latent representation of a node i, while the inference routers use ddeeph(s)- Then, both sequentially apply two dense layers that both map them 128 (20Newsgroup, MNIST, Fashion-MNIST, Omniglot), or 512 (CIFAR-10, CIFAR-100, CelebA) dimensions, after which final dense layer with Sigmoid activation function computes the probability of the binary decision that infers the next child node.
F.2.3 MNIST/Fashion-MMNIST
1x288
32@3x3
1607x7
1@28x28 80 8014x14
Figure 18: MNIST and Fashion-MNIST encoder architecture.
For MNIST and Fashion-MNIST we use identical encoder/decoder architectures. Both datasets have inputs of dimension 28 x 28 and 8 dimensional latent spaces. We use small encoder as depicted in Figure 18 and symmetric decoder. The encoder uses 3 convolutions with stride 2, whereas the decoder uses transposed convolutions.
31

---

1012
1e25d
Figure 19; 20Newsgroups encoder architecture.
E.2.4 20Newsgroups
20Newsgroups is text dataset, so the encoder/decoder architecture differs from the other experiments. The samples of 20Newsgroups are preprocessed in such a way, that the input becomes 2000D vector (see Section D.3). Hence, the encoder for 20Newsgroups a simple MLP with 5 dense layers, as depicted in Figure 19 and the decoder again mirrors the encoder architecture.
E.2.5 Omniglot
128@5x5
128@ 1x1
32@28x28
32@13x13
64@ 13x13
64@5x
1@28x28
1x8
Figure 20: Omniglot encoder architecture.
For Omniglot, to Section E.2.3 we increase the receptive field the encoder by using x convolutions and make it 4 times wider and times deeper. We refer Figure 20 for schematic of the encoder. As in the other architectures, the decoder is symmetric to the encoder.
architecture similar the for MNIST and Fashion-MNIST. In contrast
E.2.6 CIFAR-1O/CIEAR-100/GGGA
For real-world image datasets, complex architectures required. We first define a ResnetBlock as two consecutive convolutions with kernel size 3 and stride 1, whose output gets weighted by 0. and combined with residual connection the input of the first convolution. In case input and output dimensions do not match, the residual connection consists of a convolution with kernel size and stride 1 to match dimensionality. The encoder consists of consecutive ResnetBlocks with 2D average-powe peraageppooling of kernel size and stride 2. Here, number of filters per ResnetBlock are increasing by factor of starting from 32 and going up to 256 until the representation is which point we flatten it to get dy. We refer to Figure 21 for schematic of the encoder. The decoder structure similar the reversed encoder, where the pooling layers are replaced by 2D upsampling with bilinear interpolation and before the ResnetBlocks, we pass the input through dense layer and consecutive reshaping in order to match the expected input dimensions. Slight differences are required for CelebA due to the higher resolution. To arrive at the same internal representation, add one more ResnetBlock with 256 filters the encoder, as well one more ResnetBlock the start of the decoder, starting from 16 instead of 32 filters.
256,
32

---

256@4x4
1286 8x3
64@16116
32@32x32
S@32x32
1x84
Figure 21: CIFAR-10, CIFAR-100, and CelebA encoder architecture. The initial convolutional layer comprises a ResnetBlock, while the remaining convolutional layers additionally incorporate average pooling prior to the ResnetBlock.
E.3 Contrastive Learning
In real-world datasets, data splits, as determined by the ELBO and especially by the reconstruction loss, might not coincide with the human-aligned notion of clusters. For example, we observed that plain VAE utilizes specialized decoders to focus on good reconstructions of the various background colors. These data splits are rather intuitive from a reconstruction point of view, however, these might not encode contextually meaningful splits. Thus, we introduce an additional contrastive loss to guide Tree VAE toward meaningful clusters. The idea of contrastive learning (Sohn) 2016) van den Oord et al. 2018 Wu et al. 2018 Chen et al.) 2020) is simple; constrain the space of potential clustering rules by defining certain data augmentations with respect which the model's cluster assignment should become invariant. The model should assign a sample to the same cluster independent of whether and which subset of the chosen augmentations augmentations that were used for contrastive learning, we refer to Appendix [E.1.2]
applied to it. For the
To achieve the desired behavior, we utilize the normalized temperature-scaled cross-entropy loss (NT- Xent) (Sohn 2016 van den Oord et al. 2018 Wu et al. 2018 Chen et al. 2020), whose definition we repeat here for ease of reading: L,, = - log S2N (euss) exp (s,, where St,j denotes similarity
exp (x..71
measure, which we define as the cosine similarity between the representations of ã; and my, and is temperature parameter. This is computed for all 2NN pairs of augmented samples ă; and 2; originating from the same initial sample and then averaged. We compute this depth-wise for every embedding d, by first passing it through a separate MIP with one hidden layer of dimension 512, followed by batch normalization and LeakyReLU and the final linear layer with output dimension 64, on which the NT-Xeet with 7 = 0.5 is applied. Then, we average the pairwise losses and sum them up over all d;, where we, to be invariant to the depth the bottom-up. divide by the maximum depth. With this, augmented pair of samples close to each other, thus constraining the learning of information that the model should be invariant to.
regularize the bottom-up embeddings learned by the model to encode every
Even with the previous regularization, di might still contain some of information that relevant for clustering (such as background color), and therefore the routers might base their split on such characteristics. Therefore, we build upon ideas from Li et al. [002 2022), which propose not only to regularize the embeddings, but also the cluster probabilities themselves. While [Li et al. (202 2022) change from a sample-wise view to a cluster- wise view for calculating the NT-Xent on the 2K clusters, we instead keep the sample-wise view and apply the NT-Xent with 1 directly on the
33

---

outputs of every bottom-up router g(c; d:). We would like to emphasize here that we do not project them to another space, because, in contrast to previous methods (Chen et al. 2020), we wish to remove all irrelevant information with respect to the augmentations from the routers. The computed pairwise loss then weighted by the probability of falling into the given node. Lastly, we sum up the weighted pairwise losses and divide them by the sum of the weights to ensure equal regularization in every node. This has three advantages: (?) We can use the proposed loss during the training of the subtrees, thus gaining computational efficiency. (ii) We implicitly encourage the splitting of the data into balanced classes, following the overall design choice for splitting the data, outlined in Section 2.5 (fii) We can omit the computationally intensive finetuning of the full tree, the splits, as enforced by the contrastive losses, are local, which alleviates the need of finetuning the full tree for global optimization.
Lastly, we add these two contrastive parts together and weigh them by 100 match the gradients of the ELBO. Thus, we naturally introduce tradeoff between reconstruction quality and clustering quality, as we constrain the model to a clustering solution that does not optimally make use of specialized decoders.
E.4 Baselines
For a fair comparison, we retain the same architectural choices made for VAE for both the Variational Autoencoder and Ladder Variational Autoencoder, used as baselines. In particular, we retain the same encoder and decoder structure as explained in Section E.2 For datasets on which VAE applies contrastive learning, similarly apply the regularization the bottom-up embeddings of the baselines. Below we describe specific training details that deviate from Tree VAE, due to the differences in the graphical models.
E.4.1 Variational Autoencoder
The minhheerarchicll counterpart of VAE is classical Variational Autoencoder (VAE). As it does not contain hierarchical structure, the VAE consists solely an encoder and decoder architecture, which follows the architectural choices of VAF.. We train the VAE for 500 epochs and use the KI. annealing procedure proposed by LadderVAE, which linearly increases the weight of the KL terms from 0 to 1 in 200 epochs.
E.4.2 Ladder Variational Autoencoder
The sequential counterpart of VAE is the LadderVAE (Sonderby et al. 2016). We retain the same encoder, decoder, and transformations used by VAE, and we set the depth to for all datasets. We train the LadderVAE for 1000 epochs and use the KL annealing procedure proposed by the authors, which linearly increases the weight of the KL terms from
200 epochs.
E.4.3 Agglomerative Clustering
We train Ward's minimum variance agglomerative clustering (Ward 1963 Murtagh & Legendre 2014) on the plain input space (Agg), the latent space of the VAH (VAE + Agg) and the last layer of LadderVAE (LadderVAB. + Agg). We observed that clustering performed substantially better than when using the root (or the top layer). We chose Ward's method as it performed better than other classical hierarchical clustering methods (such single-linkage agglomeration clustering or bisecting K-means). However, Ward's method cannot be tested on held-out dataset, it inherently does not allow for prediction new data. For this reason, we train and test Ward's method on the test set embeddings given by the baseline models, which were trained on the training data, for fairer comparison with the other methods.
the last layer of LadderVAB
E.5 Resource Usage
In the following, we describe the resource usage of our experiments. All our experiments were run on RTX3080 GPUs, except for CelebA, where we increased the memory requirement and use a RTX3090. Training TreeVAE with 10 leaves on MNIST, Fashion-MNIST, and Omniglot-50 takes between 1h and 2h, Omniglot-5 30 minutes, CIFAR-10 5h. Training TreeVAE with leaves on 20Newsgroup takes approximately 30 minutes, and on CIFAR-100 9h. Training TreeVAR. CelebA
34

---

takes approx 8h. Please note that we only report the numbers to generate the final results but not the development time.
35


Reviews:
[{"review_id": "dfQiQPPiBq", "reviewer": "Reviewer_UPHe", "summary": "This paper introduces a new architecture for variational auto encoders with a binary tree structured generative model. The approach takes the architecture of the Ladder VAE, but introduces a binary routing variable at each stochastic layer that allows generation to continue down one of two possible paths. The authors derive an evidence lower bound objective for their approach and a simple heuristic for learning the structure of the tree at training time. In their experiments the authors compare the clustering and generative modeling performance on a variety of benchmark image datasets.", "questions": "- How does the model perform if you start with a pre-specified tree structure?\n- Equations 3 and 5 suggest that $q(z_i)$ is only dependent on the parent of $z_i$ and not $x$, which is inconsistent with equations 6-9 and the diagrams.\n- Line 134, I'm unclear exactly what is meant by the routers having the same architecture as the generative model. Don't they depend on $x$?\n- Equations 19 and 20 don't use $l$ anywhere within the summation, so it's unclear exactly how multiple samples are used.\n\n", "limitations": "The authors do discuss the limitations appropriately.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This seems to be promising work and the model introduced in this paper does appear to have a number of advantages. The model is clearly introduced and the objective function derived by the authors seems sound to me, though I haven't gone through the details in depth. Unsupervised clustering of complex, high-dimensional data is a challenging and relevant problem. Both the quantitative and qualitative evaluations of the clustering performance seem to show that the method does find interesting structures in the data that correspond to known class labels, without having access to these labels at training time. The generative performance is also promising, though the gains are unsurprising given that the model can have substantially more parameters than its competitors. The method is also straightforward to implement and apply in practice.", "weaknesses": "While I think this work is interesting, I do have a number of concerns:\n - The computational complexity of training seems to be an issue, as computing the training objective requires evaluating every path in the tree. This likely means that this model is limited to very small trees and is not applicable to discovering large numbers of clusters. Similarly, the number of parameters is large for the depth. If clustering was not a concern it seems likely that the computational complexity would likely be better spent simply creating a larger hierarchical VAE model, though this matched parameter count comparison is not explored in the work. I do appreciate that for generation, the complexity only depends on the depth.\n- I feel the work would benefit from further discussion and analysis of the tree-building routine. The proposed heuristic seems to work reasonably well, but is not well justified apart from the reasoning that the authors want to keep the tree balanced. This reasoning may fall apart if the true clusters are unbalanced. I would like to see discussion of how this choice interacts with the likelihood bound and whether the fact that some leaves can have significantly more generative capacity (as they are deeper in the hierarchy), has issues.\n- The VAE/network architectures used for the experiments are somewhat out-of-date. The ladder VAE is substantially outperformed in generative performance by hierarchical VAEs like the NVAE or \"very-deep\" VAE of Child. State-of-the-art generative performance isn't necessarily the goal of this work, but on the chosen datasets the model presented is not competitive. The VAE + agg clustering baseline could similarly be updated to these newer (pertained) models or applied to something like a vector quantized VAE.\n- The experiments are all on benchmark image datasets, which is fine, but I would be more interested in an application to a real-world dataset where the approach could provide new insights, rather than just rediscovering already known image labels.\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful comments and feedback!\n\nW1: We agree that a naive training algorithm would have a heightened computational complexity, however, due to the structure of the model we can alleviate this issue: Firstly, we can store the values of every visited node such that we need to compute each only once, making it linear in the number of nodes. Secondly, the subtrees of all nodes with the same depth are conditionally independent, which implies that their computation can be parallelized. This reduces computational complexity to linear in the depth of the tree. Furthermore, during the greedy growing procedure, most of the tree is frozen, shown in Fig. 3 (right) of the submission. Lastly, even the greedy growing procedure can be parallelized, as subtrees are again conditionally independent. Thus, with enough resources available, we are able to recover the computational complexity of an iteratively grown version of LadderVAE, while simultaneously being able to cluster the data. We do appreciate that if not enough resources are available, one has to omit a few of the aforementioned speed-ups, however, we argue that this is the case for any parallelizable system. Nonetheless, we agree that the effective number of parameters \\textit{during training} is higher for TreeVAE than for LadderVAE, with the notable difference that TreeVAE offers lightweight inference. Therefore, we have additionally conducted an experiment with matched the parameter count of LadderVAE. The results for LadderVAE are \n|Method|Acc.|NMI |LL |RL|ELBO|\n|-|-|-|-|-|-|\n|LadderVAE|$56.4 \\pm 1.9$|$62.3\\pm2.0$|$-233.7 \\pm 0.1$|$224.8\\pm0.3$|$-238.9\\pm 0.3$|\n|TreeVAE|$63.6 \\pm 3.3$|$64.7 \\pm 1.4$|$-234.7 \\pm 0.1$|$226.5 \\pm 0.3$|$-239.2 \\pm 0.4$|. As expected, this big LadderVAE is performing slightly better generatively, however at inference time, LadderVAE does not offer lightweight inference that reduces the effective parameter count to what is denoted in the main paper. We are quite satisfied that the differences, especially in the ELBO, are only marginal, while the significant difference in clustering performance stays similar. Note that FMNIST was chosen as the worst-case dataset for clustering performance, so we would expect the differences to be bigger in the other datasets.\n\nW2: Regarding the proposed tree-building routine, we took inspiration from decision trees for the iterative, greedy growing procedure. Our empirical experiments suggest that the number of samples as simple splitting criterion supports a consistent improvement of the ELBO, as can be seen in Fig. 2 of the additional PDF, where we monitored the ELBO after the convergence of every subtree. \nAn alternative splitting criterion would be reconstruction loss, for which we conducted experiments that result in a slightly worse test ELBO of $-239.3 \\pm 0.2$. Thus, we believe our criterion is a good starting point on which future work, e.g. combining our work with [1], could measure their proposed improvements and we will make sure to include this in our Limitations section of the final version.\nLastly, the fact that some leaves can have more generative capacity is an intended design choice because it gives the model more freedom to adapt to the data. Having the flexibility to find the right depth for each cluster gives the model more rather than less capabilities.\n\nW3: With respect to the generative capabilities of our work, we agree that it is not state-of-the-art, which is intentional, as this is not the focus of this work. Rather, we want to introduce a new model class that is able to perform generative hierarchical clustering. This is also reflected in our choice of datasets. We want to emphasize that the performance difference to NVAE [3] is not due to the theoretical model formulation but rather due to the chosen architecture and designs. \nWe argue that NVAE is in some way a LadderVAE with smart design choices, where we see TreeVAE on the design-choice-level similar to LadderVAE and leave it up to future work, to focus on developing smarter and more involved design choices for our model, which would be an entire work in itself. Thus, we expect that similarly to how TreeVAE outperforms LadderVAE, a version of TreeVAE with highly optimized architecture would outperform NVAE. \n\nW4: Similar arguments to Weakness 3 can be made for the choice of datasets, however, we are currently working on a medical application of our model, as well as are in contact with practitioners in the field of robotics. We decided against the pursuit of including such applications in the paper in order to, given the space constraints, put more focus on the theoretical side of the method. Here, we would also like to note that Newsgroup20 is not an image but a text dataset and that the presented cluster enrichment of CelebA is intended as an exploratory data analysis, rather than for simply rediscovering categorical labels.\n\nQ: In response to your questions, we have conducted an additional experiment, where we a priori fixed the tree structure depicted in Fig. 4 (right) of the submission and trained the full tree without growing. The results are as follows: Acc. $35.1 \\pm 4.8$, NMI $54.1 \\pm 5.8$, LL $-237.4 \\pm 0.4$, RL $228.9 \\pm 0.5$, ELBO $241.0 \\pm 0.5$. Expectedly, the results are much worse, as a fixed tree is more prone to local optima, providing support of our growing scheme. Regarding Eq. 3 \\& Eq. 5, we adopted the notation of LadderVAE (see their Eq. 16) in order to be consistent. For Line 134, the routers of the inference model depend on the intermediate representation $\\mathbf{d}_{depth(i)}$ and thus take as input a latent embedding and not $x$. Regarding the type in Eq. 19 \\& Eq. 20, thank you for bringing it up, we have changed $z_i$ to $z_i^{(l)}$.\n\nWe hope our general response and the response to your review address your concerns. We are happy to address any additional feedback or questions you might have to alleviate your concerns."}, {"review_id": "pOs0CfXEgn", "reviewer": "Reviewer_hJ8G", "summary": "This paper introduces a new deep generative model, the tree variational autoencoder, designed to discover latent hierarchical clusters in data. The generative model makes a number of latent binary choices over a pre-learned tree structure, sampling a continuous representation at each node, then finally decoding the observed data point at the particular leaf of the binary tree that is reached. A scheme for performing variational inference on the model and a heuristic for learning the tree structure are given, and results presented for hierarchical clustering on common small image and text datasets.\n", "questions": "Perhaps I’m missing something simple here but how is learning performed over the Bernoulli routers after the tree structure has been learned and all weights unfrozen? During the tree formation each learning problem has a single Bernoulli variable that is summed out but if the weights are unfrozen then aren’t there a combinatorial number of values to the Bernoulli variables?", "limitations": "Acknowledgments of a few limitations on Pg 9. I suspect there are specific implementation details where an alternative choice would make learning fail. It would strengthen the paper to have an ablation study where the sensitivity to model and training design choices are tested (i.e. going beyond the experiments in C.3).", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Simple but novel technical idea for model and training, breaking it down into a number of simpler learning problems. Interesting qualitative results in terms of what hierarchy of clusters are discovered and strong qualitative results compared to baselines. Related work includes all relevant that I am aware of.\n", "weaknesses": "The datasets used seem quite small/simple in this day and age. I think it would be interesting to perform hierarchical clustering on the embeddings from a foundation model like Stable Diffusion on a much larger dataset. I think this may be what is being suggested in Line 318 in “Limitations & Future Work”. I would upgrade my score to Strong Accept if these results were present and favorable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and your thorough reading of the paper.\n\nW1: We agree with this statement, as the purpose of this work is to introduce a new model class that jointly learns generation as well as hierarchical clustering. We are currently working on improving the model architecture, similar to how NVAE [3] improves upon LadderVAE, in order to reach SOTA performance on SOTA benchmarks, however, these results will not be available by the end of the rebuttal period and might justify a separate paper by themselves.\n\nQ1: During the training of the full tree, we sum over all $K$ paths to the leaves. Each term in the sum only depends on the realization within the path, that is, we can ignore all Bernoulli variables that do not lie in the path. Furthermore, for paths that share certain edges and nodes (think of two adjacent leaves), we only have to compute the intermediate values once, as they are shared between the paths. This reduces the computational complexity to being linear in the number of nodes. \n\nL1: We agree and refer to Table 1 of the additional PDF.\n\nIf you have any additional questions, comments, or feedback, we will be happy to address them."}, {"review_id": "thaJdWtZqX", "reviewer": "Reviewer_bo43", "summary": "introduces a new generative hierarchical clustering model called Tree Variational Autoencoders (TreeVAE) that uncovers hidden structure in data by adapting its architecture to discover the optimal tree for encoding dependencies between latent variables. The authors compare TreeVAE to other generative models and demonstrate its effectiveness in uncovering underlying clusters and hierarchical relations in real-world datasets. They also discuss the advantages and disadvantages of using TreeVAE, such as its ability to provide a tighter lower bound of the log-likelihood at the expense of a larger neural network architecture and an increase in the number of parameters. Overall, this paper presents an innovative approach to generative modeling that has the potential to improve our understanding of complex data structures.", "questions": "Please refer to the above sections for detailed discussion.", "limitations": "Yes, there is a section at the end of the paper discussing limitations and future research opportunities.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The quality of exposition is commendable. The paper presents a clear narrative, beginning with the intuitive premise, followed by a comprehensive description of the model, including a highly useful algorithmic representation. The logical sequencing of these sections facilitates reader comprehension. The authors have succeeded in presenting a complex topic in an accessible and cogent manner.\n\n2. The innovation presented through the proposition of a tree structure in the latent space of a Variational Autoencoder (VAE) constitutes a significant contribution to the field. This unique concept not only manifests novelty but also offers functional value as it endows the model with enhanced performance capabilities. The authors have adeptly demonstrated how this structural adjustment can add a new dimension to the potential applications of VAEs and its progenies.\n\n3. robust set of qualitative and quantitative evaluations are delivered in the paper. The authors have not shied away from providing extensive empirical evidence to assert the superiority of their model. The data presented is both extensive and convincing, demonstrating the improved performance of their model compared to common benchmark models. The evaluative components of the paper are meticulously presented, reflecting the comprehensive and rigorous experimental protocols the authors have employed. The authors' due diligence in this aspect lends credibility to their claims and effectively underlines the practical implications of their research.\n\n", "weaknesses": "1. The authors have stated that the clusters learned by their model possess semantic meaning. However, it is prudent to query if this semantic characterization emerges from explicit feedback or training, or if it is simply a by-product of the model's intrinsic clustering tendencies. For instance, in Figure (4), items such as pants are grouped alongside shoes and bags, while being distinct from tops. This calls into question the semantic integrity of the learned representation. Hence, it would be advantageous for the authors to conduct an additional layer of embedding analysis, in order to ascertain the proximity of these items beyond mere qualitative observation. Such an endeavor would bolster the readers' understanding and lend more weight to the claim that the cluster learned is indeed semantic.\n\n2. he paper could benefit from a more extensive ablation study that investigates the effect of different components and hyperparameters on the model's performance. For instance, an intriguing area of exploration could be the impact of modifying the number of layers in the tree. It would be worthwhile to examine if augmenting this parameter might encourage a more detailed clustering. A comprehensive study of this nature could elucidate the relative contributions of each component, making clear which are indispensable to the model's performance, and which are more auxiliary in nature. By adding more depth to the discussion of how these elements influence the final result, the authors would be equipping their audience with a more nuanced understanding of their model's operation.\n\nThese enhancements would provide further depth and rigor to the paper, effectively strengthening the authors' claims and helping readers to more fully grasp the novelty and impact of the presented research.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer bo43,\n\nThank you for your comments and your thorough reading of the paper.\n\nW1: The idea behind the claim of semantic meaningfulness is that samples of the same cluster should have a similar latent representation and also that clusters that are close to each other in the hierarchical structure should have more similar latent representations than clusters far apart. This is integrated in the model design, as the learning of the decision boundary is mainly guided through the weighted reconstruction loss. That is, the split in every node is optimized such that the two decoders can specialize themselves as well as possible on the subgroup of data, and therefore the subspace of representations, they observe. Thus, the routers try to find splits that partition the data into two subgroups that are as distinct as possible while being as similar as possible within, which coincides with partitioning the latent space at every node into two semantically different subspaces. With this objective, our model should be able to recover semantically meaningful clusters through the training procedure. A direct way of checking this is by exploring the learned tree of the 20Newsgroup dataset in Fig. 5 of the main work. For each leaf, we denoted the majority class that falls into it. Note that the name of each class entails multiple hierarchical layers. As can be seen, the learned tree recovers the subgroups of different hierarchical levels very well, for example, all computer-related topics are grouped in one subtree. Thus, the human-assigned semantic structure corresponds to the unsupervisedly learned semantic structure. Furthermore, we have conducted an additional layer of analysis for the embeddings of FMNIST. Firstly, we have applied UMAP to reduce the root embedding to 2 dimensions and visualize them in Fig. 1 of the additional PDF. As can be seen, tops as well as shoes are clustered together, while bags are somewhat in between. Trousers on the other hand are completely separate. Our interpretation is that the root split is made between shoes and tops, where it is unclear in which subtree bags and trousers should fall as they are both not clearly assigned to one of the two groups. Therefore, depending on initialization, they might end up on either side. To analyze the effect of random weight initialization on the learned dendrogram, we have made further efforts to find a dendrogram that best summarizes the learned dendrogram of multiple runs. For this, we first align the leaves of the different trees by maximizing the overlapping samples, then we store the number of edges between any two clusters for every tree and then average this number over all trees. Thus, we have computed a distance matrix of all clusters, which is averaged over all trees. We then used average and complete linkage to cluster according to the average distance matrix. The recovered dendrograms of the two algorithms are identical, apart from the Bag and the Trouser cluster, which switch places. That is, all tops are in one subtree and all shoes in the other, while bags and trousers are assigned to either the shoes or tops subtree, depending on the clustering algorithm used. What this indicates is that the two clear groups of shoes and tops are consistently recovered, no matter the random initialization, while the assignment of bags and trousers varies, which is to be expected, as there is no clear right choice.\n\nW2: We agree that analyzing the effects of different components is important and have thus conducted an extensive ablation study presented in Table 1 of the additional PDF. Additionally, we would like to mention the further experiments performed in App. C, where for example we let the tree grow deeper, which lead to the discovery of different subgroups within the same class, for example, the differentiation of straight and tilted $9$'s or dark and light T-shirts.\n\nWe hope our general response and the response to your review address the established weaknesses. We are happy to address any additional feedback or questions that might arise."}, {"review_id": "tQ7NJZA2zH", "reviewer": "Reviewer_NuUb", "summary": "The paper presents a novel method of unsupervised hierarchical clustering by encoding structural sequential dependencies between hidden variables within the framework of variational auto-encoders. The authors adopt similar designs of top-down and bottom-up dependency structure to Ladder VAE, but imposes a binary tree or decision tree based structure to prior and posterior distribution of latent variables instead of fixed chain structure adopted by Ladder VAE, to discover hierarchical and well-organized dependencies among encoded representations. In the proposed decision tree structure, all hidden variables are organized in a top-down architecture with a nested set of binary gating/decision variables controlling the direction of every move, and altogether optimized with a unified objective function derived through variational inference. To make the binary tree structure flexible and learnable, a sequential tree growth procedure is also developed. The paper compares TreeVAE with both non-hierarchical and hierarchical baselines, and it shows superior clustering performance and achieves competitive log-likelihood lower bound.", "questions": "1) The pre-defined maximum depth of the binary tree or number of total leaves should be a major factor for hierarchical methods. So it would be more convincing to take different choices of this factor into account and make comparisons accordingly.\n2) Could the authors explain how the labeled hierarchical structure shown in Figure 5 is obtained? If it is an empirical summary, the results will be not convincing enough since there might exist mismatch between ground-truth labels and the clustering hierarchies learned by the model, with clustering accuracies less than 50 percent as shown in Table 1.\n3) The results of Figure 6 and Table 3 are really confusing and vaguely explained. What does the term ‘leaf-frequency’ mean and what can be concluded from the results? Could the authors explain it in a clearer way?\n4) Since the binary tree structure is learned in a sequential way, would it be better to term it a ‘suboptimal’ or ‘locally optimal’ tree instead of ‘optimal’ tree adopted in the paper?\n", "limitations": "see above", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) The paper is well written, and is easy to follow. Informative figures and detailed formulae are given to explain how the model is built, which is very helpful.\n2) Although the perspective of modeling hierarchical clustering as a binary decision problem is not novel, the idea of extending Ladder VAE to the hierarchical clustering task is interesting and  looks natural.\n3) Quite a few illustrative examples are introduced to validate the meaningful hierarchical structures learned by TreeVAE (e.g., Figure 1 and Figure 4), which makes the paper easier to follow and the results more convincing.\n4) Extensive experiments are conducted on multiple datasets (e.g., MNIST, Fashion, 20Newsgroups, Omniglot, CIFAR-10/100, CelebA) and multiple tasks (e.g., hierarchical clustering, sample generation and hierarchy discovery) and with quantitative results compared with state-of-the-art baseline methods.\n", "weaknesses": "1) The paper provide comprehensive experimental results, which is good. But The qualitative results are not explained very clearly, prompting people to doubt if the proposed method can really generate proper hierarchies that is consistent with the dataset’s intrinsic structure.\n2) Some important ablation studies are missing, such as data augmentation and the pre-determined maximum depth of the binary tree.\n3) For a generative model, sample generation should be an important part of the evaluation, but the implementation and experimental results are unclear in the context and mostly postponed to the appendix, making it quite confusing.\n4) Since the tree growth procedure is proceeded in an iterative and sequential way, the corresponding training efficiency is doubtful. And even so, before the procedure we still have to set up the maximum depth of the tree or the number of leaves depending on specific dataset.\n5) Data augmentation is a relatively weak form of prior knowledge especially for the hierarchical clustering task, it would be appreciated to incorporate other easily accessible forms if possible.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer NuUb,\nthank you for your feedback and constructive criticism.\n\nW1: Due to the space limitations, we have refrained from explaining the qualitative results in greater detail, however, we will make sure to include this in the camera-ready version. There are also additional figures in Appendix C, which we described in more detail. We also want to point toward Fig. 1 of the additional PDF combined with our response to reviewer bo43, where we perform an analysis of the root embedding of FMNIST. We show that the in the learned embeddings, bags are embedded in between shoes and tops, while trousers are an entirely separate category for themselves. Additional evidence is presented in the 20Newsgroup dataset, which we talk about in your Q2 below.\n\nW2: Additionally to the ablation studies in Appendix C, We have added an extensive ablation study in Table 1 of the additional PDF, as well as presented a comparison of different contrastive approaches in our response to reviewer 822K. We want to emphasize that without data augmentation, the clustering performance for the CIFAR datasets would lower greatly, as the model clusters according to background information, which usually influences reconstruction loss more than the foreground information. \n\nW3: Thank you for this comment. As we are solving two tasks in one (those being generation and hierarchical clustering), we decided to emphasize the discovery of hierarchies in the main text. We will make sure to move selected generative figures from the Appendix into the main text for the camera-ready version.\n\nW4: We have experimented with multiple growing procedures and the iterative way has proven to be the most stable one, see Fig. 2 of the additional PDF for a visualization of the training ELBO behavior on FMNIST. By iteratively growing, we divide the complex problem into simpler subproblems, which is inspired by agglomerative clustering. Additionally, with enough computational resources, one could grow all nodes of the same depth simultaneously, as the resulting subtrees are conditionally independent, which would further speed up training. Note also that during training of the subtree, the weights of the rest of the tree are frozen (see Fig. 3 (right) of the main paper), which improves training efficiency.\nRegarding the stopping criterion of the clustering procedure, we want to mention that it is common to set the number of clusters $K$ a priori. Nevertheless, we agree that this is unsatisfactory and will focus future efforts on finding smarter ways in going about this issue, as well as mention it in the Limitations section of the final version. \n\nW5: A goal of our method is to perform hierarchical clustering with as little inductive biases as possible, in order to be truly unsupervised and also stay true to maximizing the evidence lower bound. Therefore, we only add additional supervision in the weak form of data augmentation to guide the model towards desirable splits, when it is required, which is the case for the CIFAR datasets and CelebA. Naturally, future extensions can be made where stronger inductive biases can be leveraged in order to guide and support the clustering procedure, however, this is not the focus of our current work, as we want to put focus on the fact that the model itself already works very well. This is, we propose a theoretically founded model class, where the contrastive extension shows that combining it with forms of supervision is possible.\n\nQ1: We refer to Table 1 of the additional PDF, where we have conducted an ablation analysis for different pre-defined maximum depths of the tree while keeping the maximum number of leaves $K$ constant. Additionally, the metrics Dendrogram Purity (DP) and Leaf Purity (LP) in the main paper allow unbiased evaluations for dendrograms where the number of learned clusters different from the ground-truth number, which is why we compute these two metrics with trees that contain $20$ leaves. Lastly, we also want to mention the further experiments in Appendix C, where we depict trees that are grown to more than $10$ clusters, as well as quantitative results for varying the number of ground-truth clusters, as well as the number of learned clusters.\n\nQ2: 20Newsgroup consists of $20$ classes that also have a hierarchical structure. The labels depicted in the figure are the $20$ ground-truth labels of the dataset, where we visualize the majority class of the samples falling into each leaf. The words in each class label indicate the group assignments on each hierarchical level, where each level is separated by a dot. For example, \"comp.graphics\" indicates a first hierarchical layer of \"computer\" and a second layer of \"graphics\". To be clear, we did not choose the hierarchical level of each class in Fig. 5; it is predetermined by the creators of the dataset. Our experiments show that TreeVAE discovers the correct hierarchies, which supports the claim that we uncover meaningful hierarchies. \n\nQ3: The CelebA experiments are performed to give an alternative experiment to simply recovering already known labels. Here, we let the model grow and try to interpret the results in a more explorative fashion. We will make sure to put more information in the camera-ready version. Every number $v_{a,k}$ in Table 3 indicates for a given attribute $a$ and a given leaf $k$ the percentage of all samples that fall in $k$, in which attribute $a$ is present. Thus, we observe that people with blonde hair mostly fall into cluster $8$ (given the knowledge that the number of samples is evenly distributed).\n\nQ4: Thank you for the suggestion, we agree with the pointed-out issue and will adopt \"locally optimal\" in the camera-ready version.\n\nWe hope our general response and the response to your review address the established weaknesses and questions. We are happy to address any additional feedback or questions that might arise."}, {"review_id": "931W3MtNvs", "reviewer": "Reviewer_822K", "summary": "This paper introduces a new class of variational autoencoders that define a binary tree structure. Given a fixed tree, the generative model is a top-down hierarchical model with binary routing components and the inference network follows a top-down approach in the style of the Ladder VAE. The authors propose to grow the tree iteratively every N training steps. They also suggest using `NT-Xent` to force aligning the learned representation with human perception.\n\nBased on binary image, natural image and a text dataset, the authors demonstrate the effectiveness of the method based on modelling (likelihood) and clustering performances. Examples of learned trees are presented in the paper and are convincing: coupled with `NT-Xent`, Tree VAE learns convincing data hierarchies without explicit supervision.\n", "questions": "1. What is the modelling performance (likelihood) of Tree VAE without `NT-Xent`?", "limitations": "Limitations of the method are not sufficiently discussed. The authors might want to discuss some of the following points more thoroughly\n1. the multi-stage training of TreeVAE adds complexity, and result might vary depending on the choice of expansion criterion\n2. the results depend of the choice of \"prior knowledge\" model", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "A strong, creative and well-written paper!\n\n1. This paper makes a significant contribution towards learning meaningful hierarchical models, which is an important problem in the field\n2. The solution described in this paper is creative and well presented (structured presentation: generative model vs. inference network, detailed derivations, intuitive notation)\n3. The experiments are extensive (5 datasets, 4 baselines, 10 seeds for each run)\n4. The visual results are very convincing\n5. In-depth appendix", "weaknesses": "1. Section 2.6 is unclear as such. I suggest adding the equation corresponding to the auxiliary loss term that comes with the use of `NT-Xent`\n2. The high quality of the results might be tightly tied to the use of `NT-Xent` .  According to appendix E.3, a lot of work is required to get this right. I think this work would greatly benefit from an ablation study targeted at the use of choice of methods for integrating prior knowledge (e.g., without `NT-Xent`, with another method)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and feedback! We appreciate your support for this paper.\n\nW: We will make sure to extend Section 2.6 for the camera-ready version. We intentionally kept it short, as we only use `NT-Xent` for real-world image datasets, that is CIFAR-10, CIFAR-100, and CelebA. Here, the reconstruction loss is not sufficient for guiding the tree toward a meaningful clustering, which is why for those datasets, we inject inductive biases through contrastive learning. As such, the results of MNIST, FMNIST, 20Newsgroup, and Omniglot are solely tied to the maximization of the ELBO and the presented generative modeling performance of TreeVAE is without the `NT-Xent`. For the choice of method for integrating prior knowledge, we have experimented with multiple options. They are all inspired by the fact that two augmented versions of the same input should behave similarly. Please note that the following results were obtained on a curated subset of CIFAR-10 as this ablation was performed during the model development stage. Thus, the absolute values are not directly comparable to the values in our paper, however, the relative differences in methods remained comparable over different experiments, so we believe, in order to save computational resources, that these results are still meaningful. \n|Method|NMI|\n|-|-|\n|Ours (=A+B)|$44.1 \\pm 0.7$|\n|(A)|$21.7 \\pm 4.0$|\n|(B)|$6.0 \\pm 6.0$|\n|(C)|$23.0 \\pm 1.2$|\n|(D)|$0.2 \\pm 0.0$|\n|(E)|$1.0 \\pm 1.3$|\nHere, (A) corresponds to the `NT-Xent` regularization on the routers, (B) to the `NT-Xent` regularization on the bottom-up embeddings, (C) on a transposed version of the `NT-Xent` regularization on the routers, following Li et al. [2], (D) to the `NT-Xent` regularization of only the output of the encoder, and (E) to minimizing the Jensen-Shannon divergence between the router probabilities of the two augmented inputs. Further experiments showed that the combination of (A) and (B) leads to the best results. \n\nQ: We hope that we were able to answer your question in the previous section noting that the modelling performance in the submission is already without `NT-Xent`.\n\nL1: Lastly, regarding the multi-stage training, we agree with the reviewer and leave the investigation of a more principled approach regarding the expansion criterion to future work. We will make sure to include this in the limitations section of the camera-ready version. Additionally, we would also like to refer to our answer to Reviewer UPHe, where we show that given enough computational resources, the computational complexity is the same as for an iteratively grown LadderVAE.\n\nL2: Discussed in the weaknesses section.\n\nIf you have any additional questions, comments, or feedback, we will be happy to address them."}]

Author rebuttal:
Dear reviewers,

We deeply appreciate your insightful questions, constructive comments and helpful feedback! Your reviews suggest that you invested a considerable amount of time and effort into understanding our work, for which we are very grateful and thankful. We provide individualized responses to your Weaknesses (W), Questions (Q), and Limitations (L) in separate sections below and hope that our responses address the raised concerns.

Nevertheless, we would like to mention a few points here that might be interesting to everybody:
 
In the additional PDF, we present an extensive ablation study in Table 1. We have decided a priori to run the rebuttal experiments on Fashion MNIST, as our clustering performance on this dataset is closest to the baselines, therefore, this provides a worst-case analysis for all comparisons. The detail-focused person will observe that the presented values of the default TreeVAE configuration, denoted by "base", are marginally lower(<1%p) than in the paper, which is due to the fact that we trained for slightly shorter, in order to finish the rebuttal in time. All ablation experiments are run over $10$ seeds. Additionally, in Figure 1, we analyze the root embedding of our tree on FMNIST to support why Figure 4 (right) of the main paper clusters bags and trousers together with shoes (for a more detailed discussion, have a look at our response to reviewer bo43). Lastly, Figure 2 shows that the proposed splitting heuristic leads to a stable improvement in the ELBO. 

Please find the full bibliography utilized in all responses below.

We are happy to address any further questions or comments from the reviewers to improve the paper and look forward to a fruitful discussion!

[1] "Decision Jungles: Compact and Rich Models for Classification", Shotton et al., 2013

[2] "Twin Contrastive Learning for Online Clustering
", Li et al., 2022

[3] "NVAE: A Deep Hierarchical Variational Autoencoder
", Arash Vahdat and Jan Kautz, 2022

Meta-review: All the reviewers liked the paper. As the paper represents a novelty (tree-based hierarchies) in the well-studied variational autoencoder (VAE) field, this AE is of the opinion that the paper deservers a spotlight.
Decision: Accept (spotlight)