SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Laura Manduchi\*, Moritz Vandenhirtz\*, Alain Ryser, Julia E. Vogt
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.
Discovering structure and hierarchies in 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, 2017; Tanno et al., 2019), have proven to be successful in 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 an unsupervised path to find hidden groups in the data and their hierarchical relationship (R. J. G. B. Campello et al., 2015). Due to its versatility, interpretability, and ability to uncover meaningful patterns in complex data, hierarchical clustering has been widely used in a 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 a dendrogram - a 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 et al., 2020), and only a 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), a 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-dimensional 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, a variety of deep generative methods have been proposed to incorporate more complex posterior distributions by modeling structural sequential dependencies
*Equal contribution. Correspondence to {laura.manduchi, moritz.vandenhirtz}@inf.ethz.ch
Figure 1: The hierarchical structure discovered by TreeVAE on the CIFAR-10 dataset. We display random subsets of images that are probabilistically assigned to each leaf of the tree.
between latent variables (Sønderby 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 VAEs by combining the complementary strengths of hierarchical clustering algorithms and deep generative models. We propose TreeVAE<sup>1</sup>, a 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 TreeVAE is depicted in Fig. 1. Each edge and split are encoded by neural networks, while the circles depict latent variables. Each sample is associated with a 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, for example, the method divides the vehicles and animals into two different subtrees and similar groups (such as planes and ships) share common ancestors.
Our main contributions are as follows: (i) We propose a novel, deep probabilistic approach to hierarchical clustering that learns the optimal generative binary tree to mimic the hierarchies present in the data. (ii) We provide a thorough empirical assessment of the proposed approach on MNIST, Fashion-MNIST, 20Newsgroups, and Omniglot. In particular, we show that TreeVAE (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 TreeVAE 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 TreeVAE on CIFAR-10, CIFAR-100, and CelebA, and we show that the proposed approach achieves competitive hierarchical clustering performance compared to the baselines.
2 TreeVAE
We propose TreeVAE, a novel deep generative model that learns a flexible tree-based posterior distribution over latent variables. Each sample travels through the tree from root to leaf in a probabilistic manner as TreeVAE learns sample-specific probability distributions of paths. As a 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 a dataset  $\mathbf{X}$ , the model is defined by three components that are learned during training:
<sup>1</sup>The code is publicly available at https://github.com/lauramanduchi/treevae-pytorch.
- the global structure of the binary tree  $\mathcal{T}$ , which specifies the set of nodes  $\mathbb{V} = \{0, \dots, V\}$ , the set of leaves  $\mathbb{L}$ , where  $\mathbb{L} \subset \mathbb{V}$ , and the set of edges  $\mathcal{E}$ . See Fig. 1/4/5/6/7 for different examples of tree structures learned by the model.
- the sample-specific latent embeddings  $\mathbf{z} = \{\mathbf{z}_0, \dots, \mathbf{z}_V\}$ , which are random variables assigned to each node in  $\mathbb{V}$ . Each embedding is characterized by a Gaussian distribution whose parameters are a function of the realization of the parent node. The dimensions of the latent embeddings are defined by their depth, with  $z_i \in \mathbb{R}^{h_{\mathrm{depth}(i)}}$  where  $\mathrm{depth}(i)$  is the depth of the node  $i$ , and  $h_{\mathrm{depth}(i)}$  is the embedding dimension for that depth.
- the sample-specific decisions  $\mathsf{c} = \{\mathsf{c}_0, \dots, \mathsf{c}_{V - |\mathbb{L}|}\}$ , which are Bernoulli random variables defined by the probability of going to the right (or left) child of the underlying node. They take values  $c_i \in \{0,1\}$  for  $i \in \mathbb{V} \setminus \mathbb{L}$ , with  $c_i = 0$  if the left child is selected. A decision path,  $\mathcal{P}_l$ , indicates the path from root to leaf given the tree  $\mathcal{T}$  and is defined by the nodes in the path, e.g., in Fig. 2,  $\mathcal{P}_l = \{0,1,4,5\}$ . The probability of  $\mathcal{P}_l$  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 node-wise. The latent embeddings and the decision paths, on the other hand, are learned using variational inference by conditioning the model on the current tree structure. The generative/inference model and the learning objective conditioned on  $\mathcal{T}$  are explained in Sec. 2.2/2.3/2.4 respectively, while in 2.5, we elaborate on the efficient growing procedure of the tree.
The generative process of TreeVAE for a given  $\mathcal{T}$  is depicted in Fig. 2 (right). The generation of a new sample  $\pmb{x}$  starts from the root. First, the latent embedding of the root node  $\mathbf{z}_0$  is sampled from a standard Gaussian  $p_{\theta}(z_0) = \mathcal{N}(z_0 \mid \mathbf{0}, \mathbf{I})$ . Then, given the sampled  $z_0$ , the decision of going to the left or the right node is sampled from a Bernoulli distribution  $p(\mathrm{c}_0 \mid z_0) = \text{Ber}(r_{p,0}(z_0))$ , where  $\{r_{p,i} \mid i \in \mathbb{V} \setminus \mathbb{L}\}$  are functions parametrized by neural networks defined as routers, and cause the splits in Fig. 2. 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  $z_1$ , is then sampled from a Gaussian distribution  $p_{\theta}(z_1 \mid z_0) = \mathcal{N}(\mathbf{z}_1 \mid \mu_{p,1}(z_0), \sigma_{p,1}^2(z_0))$ , where
Figure 2: The proposed inference (left) and generative (right) models for TreeVAE. Circles are stochastic variables while diamonds are deterministic. The global topology of the tree is learned during training.
$\{\mu_{p,i}, \sigma_{p,i} \mid i \in \mathbb{V} \setminus \{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 a leaf is reached.
Let us define the set of latent variables selected by the path  $\mathcal{P}_l$ , which goes from the root to the leaf  $l$ , as  $\mathbf{z}_{\mathcal{P}_l} = \{\mathbf{z}_i \mid i \in \mathcal{P}_l\}$ , the parent node of the node  $i$  as  $pa(i)$ , and  $p(c_{pa(i) \to i} \mid z_{pa(i)})$  the probability of going from  $pa(i)$  to  $i$ . Note that the path  $\mathcal{P}_l$  defines the sequence of decisions. The prior probability of the latent embeddings and the path given the tree  $\mathcal{T}$  can be summarized as
$$
p _ {\theta} \left(\boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l}\right) = p \left(\boldsymbol {z} _ {0}\right) \prod_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} p \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {z} _ {p a (i)}\right) p \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right). \tag {1}
$$
Finally,  $\pmb{x}$  is sampled from a distribution that is conditioned on the selected leaf. If we assume that  $\pmb{x}$  is real-valued, then
$$
p _ {\theta} \left(\boldsymbol {x} \mid \boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l}\right) = \mathcal {N} \left(\mathbf {x} \mid \mu_ {x, l} \left(\boldsymbol {z} _ {l}\right), \sigma_ {x, l} ^ {2} \left(\boldsymbol {z} _ {l}\right)\right), \tag {2}
$$
where  $\{\mu_{x,l},\sigma_{x,l}\mid l\in \mathbb{L}\}$  are functions parametrized by leaf-specific neural networks defined as decoders.
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 (1), with the difference that the probability of the root and of the decisions are now conditioned on the sample  $x$ :
$$
q \left(\boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l} \mid \boldsymbol {x}\right) = q \left(\boldsymbol {z} _ {0} \mid \boldsymbol {x}\right) \prod_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right). \tag {3}
$$
To compute the variational probability distribution of the latent embeddings  $q(\mathbf{z}_0 \mid \mathbf{x})$  and  $q(\mathbf{z}_i \mid \mathbf{z}_{pa(i)})$ , where
$$
q \left(\boldsymbol {z} _ {0} \mid \boldsymbol {x}\right) = \mathcal {N} \left(\boldsymbol {z} _ {0} \mid \mu_ {q, 0} (\boldsymbol {x}), \sigma_ {q, 0} ^ {2} (\boldsymbol {x})\right) \tag {4}
$$
$$
q _ {\phi} \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right) = \mathcal {N} \left(\boldsymbol {z} _ {i} \mid \mu_ {q, i} \left(\boldsymbol {z} _ {p a (i)}\right), \sigma_ {q, i} ^ {2} \left(\boldsymbol {z} _ {p a (i)}\right)\right), \forall i \in \mathcal {P} _ {l}, \tag {5}
$$
we follow a similar approach to the one proposed by Sønderby et al. (2016). Note that we use the subscript  $q$  to indicate the parameters of the inference model.
First, a deterministic bottom-up pass computes the node-specific approximate likelihood contributions
$$
\mathbf {d} _ {h} = \operatorname {M L P} \left(\mathbf {d} _ {h + 1}\right) \tag {6}
$$
$$
\hat {\boldsymbol {\mu}} _ {q, i} = \operatorname {L i n e a r} \left(\mathbf {d} _ {\text {d e p t h} (i)}\right), i \in \mathbb {V} \tag {7}
$$
$$
\hat {\boldsymbol {\sigma}} _ {q, i} ^ {2} = \operatorname {S o f t p l u s} \left(\operatorname {L i n e a r} \left(\mathbf {d} _ {\text {d e p t h} (i)}\right)\right), i \in \mathbb {V}, \tag {8}
$$
where  $\mathbf{d}_H$  is parametrized by a domain-specific neural network defined as encoder, and  $\mathrm{MLP}(\mathbf{d}_h)$  for  $h\in \{1,\dots ,H\}$ , indicated by the bottom-up arrows in Fig. 2, are neural networks, shared among the parameter predictors,  $\hat{\mu}_{q,i},\hat{\sigma}_{q,i}^2$ , of the same depth. They are characterized by the same architecture as the transformations defined in Sec.2.2.
A stochastic downward pass then recursively computes the approximate posteriors defined as
$$
\boldsymbol {\sigma} _ {q, i} ^ {2} = \frac {1}{\hat {\boldsymbol {\sigma}} _ {q , i} ^ {- 2} + \boldsymbol {\sigma} _ {p , i} ^ {- 2}}, \quad \boldsymbol {\mu} _ {q, i} = \frac {\hat {\boldsymbol {\mu}} _ {q , i} \hat {\boldsymbol {\sigma}} _ {q , i} ^ {- 2} + \boldsymbol {\mu} _ {p , i} \boldsymbol {\sigma} _ {p , i} ^ {- 2}}{\hat {\boldsymbol {\sigma}} _ {q , i} ^ {- 2} + \boldsymbol {\sigma} _ {p , i} ^ {- 2}}, \tag {9}
$$
where all operations are performed elementwise. Finally, the variational distributions of the decisions  $q(\mathbf{c}_i \mid \mathbf{x})$  are defined as
$$
q \left(\mathbf {c} _ {i} \mid \boldsymbol {x}\right) = q \left(\mathbf {c} _ {i} \mid \mathbf {d} _ {\text {d e p t h (i)}}\right) = B e r \left(r _ {q, i} \left(\mathbf {d} _ {\text {d e p t h (i)}}\right)\right), \tag {10}
$$
where  $\{r_{q,i} \mid i \in \mathbb{V} \setminus \mathbb{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  $(\mu_{q,0},\sigma_{q,0})$ , the transformations  $\{(\mu_{p,i},\sigma_{p,i}),(\mu_{q,i},\sigma_{q,i})\mid i\in \mathbb{V}\setminus \{0\} \}$ , the decoders  $(\{\mu_{x,l},\sigma_{x,l}\mid l\in \mathbb{L}\})$  and the routers  $(\{r_{p,i},r_{q,i}\mid i\in \mathbb{V}\backslash \mathbb{L}\})$ , are learned by maximizing the Evidence Lower Bound (ELBO) (Kingma & Welling, 2014; Rezende et al., 2014). Each leaf  $l$  is associated with only one path  $\mathcal{P}_l$ , hence we can write the data likelihood conditioned on  $\mathcal{T}$  as
$$
p (\boldsymbol {x} \mid \mathcal {T}) = \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} p (\boldsymbol {x}, \boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l}) = \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} p _ {\theta} \left(\boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l}\right) p _ {\theta} \left(\boldsymbol {x} \mid \boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l}\right). \tag {11}
$$
We use variational inference to derive the ELBO of the log-likelihood:
$$
\mathcal {L} (\boldsymbol {x} \mid \mathcal {T}) := \mathbb {E} _ {q \left(\boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l} \mid \boldsymbol {x}\right)} \left[ \log p (\boldsymbol {x} \mid \boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l}) \right] - \mathrm {K L} \left(q \left(\boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l} \mid \boldsymbol {x}\right) \| p \left(\boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l}\right)\right). \tag {12}
$$
The first term of the ELBO is the reconstruction term:
$$
\mathcal {L} _ {r e c} = \mathbb {E} _ {q \left(\boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l} \mid \boldsymbol {x}\right)} [ \log p (\boldsymbol {x} \mid \boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l}) ] \tag {13}
$$
$$
= \sum_ {l \in \mathbb {L}} \int_ {\boldsymbol {z} _ {\mathcal {P} _ {l}}} q \left(\boldsymbol {z} _ {0} \mid \boldsymbol {x}\right) \prod_ {i \in \mathcal {P} _ {l} \backslash \{0 \}} q \left(c _ {p a (i) \rightarrow i} \mid \boldsymbol {x}\right) q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {p a (i)}\right) \log p (\boldsymbol {x} \mid \boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l}) \tag {14}
$$
$$
\approx \frac {1}{M} \sum_ {m = 1} ^ {M} \sum_ {l \in \mathbb {L}} P (l; \boldsymbol {c}) \log \mathcal {N} \left(\mathbf {x} \mid \mu_ {x, l} \left(\boldsymbol {z} _ {l} ^ {(m)}\right), \sigma_ {x, l} ^ {2} \left(\boldsymbol {z} _ {l} ^ {(m)}\right)\right), \tag {15}
$$
$$
P (i; \boldsymbol {c}) = \prod_ {j \in \mathcal {P} _ {i} \backslash \{0 \}} q \left(c _ {p a (j) \rightarrow j} \mid \boldsymbol {x}\right) \text {f o r} i \in \mathbb {V}, \tag {16}
$$
where  $\mathcal{P}_i$  for  $i\in \mathbb{V}$  is the path from root to node  $i$ ,  $P(i;c)$  is the probability of reaching node  $i$ , which is the product over the probabilities of the decisions in the path until  $i$ ,  $z_{l}^{(m)}$  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-Leibler divergence (KL) between the prior and the variational posterior of the tree. It can be written as a sum of the KL of the root, the nodes, and the decisions:
$$
\operatorname {K L} \left(q \left(\boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l} \mid \boldsymbol {x}\right) \| p \left(\boldsymbol {z} _ {\mathcal {P} _ {l}}, \mathcal {P} _ {l}\right)\right) = \operatorname {K L} _ {\text {r o o t}} + \operatorname {K L} _ {\text {n o d e s}} + \operatorname {K L} _ {\text {d e c i s i o n s}} \tag {17}
$$
$$
\mathrm {K L} _ {\text {r o o t}} = \mathrm {K L} (q (\boldsymbol {z} _ {0} \mid \boldsymbol {x}) \| p (\boldsymbol {z} _ {0})) \tag {18}
$$
$$
\mathrm {K L} _ {\text {n o d e s}} \approx \frac {1}{M} \sum_ {m = 1} ^ {M} \sum_ {i \in \mathbb {V} \backslash \{0 \}} P (i; c) \operatorname {K L} \left(q \left(\boldsymbol {z} _ {i} ^ {(m)} \mid p a \left(\boldsymbol {z} _ {i} ^ {(m)}\right)\right) \| p \left(\boldsymbol {z} _ {i} ^ {(m)} \mid p a \left(\boldsymbol {z} _ {i} ^ {(m)}\right)\right)\right) \tag {19}
$$
$$
\mathrm {K L} _ {\text {d e c i s i o n s}} \approx \frac {1}{M} \sum_ {m = 1} ^ {M} \sum_ {i \in \mathbb {V} \backslash \mathbb {L}} P (i; \boldsymbol {c}) \sum_ {c _ {i} \in \{0, 1 \}} q \left(c _ {i} \mid \boldsymbol {x}\right) \log \left(\frac {q \left(c _ {i} \mid \boldsymbol {x}\right)}{p \left(c _ {i} \mid \boldsymbol {z} _ {i} ^ {(m)}\right)}\right), \tag {20}
$$
where  $M$  is the number of MC samples. We refer to Appendix A for the full derivation. The  $\mathrm{KL}_{root}$  is the KL between the standard Gaussian prior  $p(z_0)$  and the variational posterior of the root  $q(z_0\mid x)$ , thus enforcing the root to be compact. The  $\mathrm{KL}_{nodes}$  is the sum of the node-specific KLs weighted by the probability of reaching their node  $i$ :  $P(i;c)$ . The node-specific KL of node  $i$  is the KL between the two Gaussians  $q(z_{i}\mid pa(z_{i}))$ ,  $p(z_{i}\mid pa(z_{i}))$ . Finally, the last term,  $\mathrm{KL}_{decisions}$ , is the weighted sum of all the KLs of the decisions, which are Bernoulli random variables,  $KL(q(\mathsf{c}_i\mid \boldsymbol {x})\mid p(\mathsf{c}_i\mid \boldsymbol {z}_i))) = \sum_{c_i\in \{0,1\}}q(c_i\mid \boldsymbol {x})\log \left(\frac{q(c_i\mid\boldsymbol{x})}{p(c_i\mid\boldsymbol{z}_i)}\right)$ . 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 parameters of both the generative and the inference model given a defined tree structure  $\mathcal{T}$ . Here we discuss how to learn the structure of the binary tree  $\mathcal{T}$ . TreeVAE starts by training a tree composed of a root and two leaves, see Fig. 3 (left), for  $N_{t}$  epochs by optimizing the ELBO. Once the model converged, a leaf is selected, e.g.,  $z_{1}$  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, e.g., the reconstruction loss or the ELBO. In our experiments, we chose to select the nodes with the maximum number of samples to retain balanced
Figure 3: 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.
leaves. The sub-tree composed of the new leaves and the parent node is then trained for  $N_{t}$  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 a threshold  $t$ ) of being assigned to the parent node. The process is repeated until the tree reaches its maximum capacity (defined by the maximum depth) or until a condition (such as a predefined maximum number of leaves) is met. The entire model is then fine-tuned for  $N_{f}$  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 a threshold).
2.6 Integrating Prior Knowledge
Retrieving semantically meaningful clustering structures of real-world images is extremely challenging, as there are several underlying factors according to which the data can be clustered. Therefore, it