SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 a hierarchical structure. This is done by recursively partitioning the data into subsets, each of which corresponds to a specific node in the tree. At each node, a 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 interpretability. The learned tree structure can be easily visualized to get insights into the data and the model's decision-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, 2017) 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 to learn meaningful representations. Thus, Tanno et al. (2019) introduces Adaptive Neural Trees (ANT), a method that learns flexible, hierarchical representations through NNs, hereby facilitating hierarchical 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-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 (Zharmagambetov & Carreira-Perpinañán, 2022), robustness (Moshkovitz et al., 2021) or interpretability (Wan et al., 2021; Souza et al., 2022; Arenas et al., 2022; Pace et al., 2022). In the context of interpretability, various approaches have been proposed to enhance accuracy and interpretability. Wan et al. (2021) introduce Neural-Backed Decision Trees (NBDTs), 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 interpretability. Pace et al. (2022) develop the POETREE framework for interpretable policy learning via decision trees in time-varying clinical decision environments. Arenas 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-offs, criteria, and frameworks for improving accuracy and interpretability.
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., 2013). Most methods, however, have the downside of 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.
Here, we provide additional experimental results, including further hierarchical structures, unconditional generation of samples, and ablation experiments.
C.1 Discovery of Hierarchies
We show further hierarchical structures learned by TreeVAE on MNIST in Fig. 9/10, Fashion-MNIST in Fig. 11, Omniglot-50 in Fig. 12, and CIFAR-100 in Fig. 13. Additionally, the figures that include conditional generations are generated by sampling from the root and following the most probable path to generate every leaf-specific image.
Figure 9: Hierarchical structure learned by TreeVAE 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.
Figure 10: Hierarchical structure learned by TreeVAE 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 TreeVAE can be used to sample unconditional generations for all clusters simultaneously for MNIST, in Fig. 14, and CelebA, in Fig. 15. Differently from the conditional generation, where we follow the path of the tree with the highest probability, here we follow all
Figure 11: Hierarchical structure learned by TreeVAE with 20 leaves on Fashion-MNIST dataset with generated images.
Figure 12: Hierarchical structure learned by TreeVAE 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 to cluster smaller characters in 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, the left child encodes more complex shapes that contain many different strokes, and the right subtree groups simpler shapes with only a few lines.
Figure 13: Hierarchical structure learned by TreeVAE with 20 leaves on CIFAR-100 dataset. We display random subsets of images that are probabilistically assigned to each leaf of the tree.
paths of the tree regardless of their probabilities. As a result, we generate  $L$  images, where  $L$  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.
Figure 14: Selected unconditional generations of MNIST. One row corresponds to 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.
Figure 15: Selected unconditional generations of CelebA. One row corresponds to 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 assumptions on the behavior of TreeVAE. In Table 4, we provide an ablation on selected hyperparameters of TreeVAE on Fashion MNIST. In the following subsection we will provide further in-depth explorations of TreeVAE.
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 MLP, (E) number of hidden layers of router, (F) pre-determined maximum depth of tree.
zdepth	MLPdepth	MLPlayers	Routerlayers	depth	Acc	NMI	ELBO
base	8, 8, 8, 8, 8, 8	128, 128, 128, 128	1	2	6	63.5 ± 4.1	63.9 ± 2.2	-239.2 ± 0.3
(A)	2, 2, 2, 2, 2, 2	49.8 ± 4.0	51.8 ± 2.7	-249.5 ± 0.4
(A)	4, 4, 4, 4, 4, 4	60.0 ± 4.7	61.0 ± 3.7	-242.6 ± 0.4
(A)	16, 16, 16, 16, 16, 16	61.0 ± 4.5	63.0 ± 2.4	-238.8 ± 0.4
(A)	32, 32, 32, 32, 32, 32	62.7 ± 4.4	63.5 ± 2.8	-238.7 ± 0.4
(B)	16, 16, 16, 16, 16	58.3 ± 3.5	60.6 ± 2.1	-241.5 ± 0.3
(B)	32, 32, 32, 32, 32	60.6 ± 3.6	62.7 ± 2.5	-240.4 ± 0.4
(B)	64, 64, 64, 64, 64	62.6 ± 3.8	63.5 ± 1.4	-239.6 ± 0.3
(B)	256, 256, 256, 256, 256	62.3 ± 4.9	63.4 ± 3.9	-239.2 ± 0.4
(C)	64, 32, 16, 8, 4, 2	128, 64, 32, 16, 8	56.7 ± 4.0	59.9 ± 3.1	-239.2 ± 0.4
(C)	64, 32, 16, 8, 4, 2	256, 128, 64, 32, 16	60.9 ± 5.3	62.6 ± 2.4	-238.6 ± 0.3
(C)	128, 64, 32, 16, 8, 4	128, 64, 32, 16, 8	58.5 ± 4.4	61.3 ± 2.5	-239.1 ± 0.4
(C)	128, 64, 32, 16, 8, 4	256, 128, 64, 32, 16	62.6 ± 3.4	63.6 ± 1.5	-238.8 ± 0.3
(C)	256, 128, 64, 32, 16, 8	256, 128, 64, 32, 16	61.1 ± 3.7	61.7 ± 2.9	-238.9 ± 0.3
(C)	256, 128, 64, 32, 16, 8	512, 256, 128, 64, 32	60.5 ± 4.7	61.1 ± 3.1	-239.1 ± 0.3
(D)	2	60.3 ± 5.3	61.6 ± 3.5	-239.0 ± 0.3
(D)	3	61.4 ± 4.2	62.9 ± 3.2	-239.0 ± 0.3
(E)	1	59.6 ± 4.0	62.3 ± 2.0	-239.2 ± 0.4
(E)	3	62.1 ± 4.4	63.6 ± 3.1	-239.1 ± 0.4
(F)	8, 8, 8, 8, 8	128, 128, 128, 128	5	64.0 ± 4.6	64.4 ± 2.7	-239.2 ± 0.3
(F)	8, 8, 8, 8, 8, 8, 8	128, 128, 128, 128, 128, 128	7	61.0 ± 4.9	62.9 ± 2.9	-239.2 ± 0.3
(F)	8, 8, 8, 8, 8, 8, 8, 8	128, 128, 128, 128, 128, 128	8	62.1 ± 5.2	63.6 ± 2.7	-239.1 ± 0.4
C.3.1 Training time parameter matching
TreeVAE has the advantage of lightweight inference, as each datum ends up in one leaf. Still, we 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 of TreeVAE at training instead of inference. The results are presented in 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	RL	ELBO
Fashion	LadderVAE + Agg	56.4 ± 1.9	62.3 ± 2.0	-233.7 ± 0.1	224.8 ± 0.3	-238.9 ± 0.3
TreeVAE	63.6 ± 3.3	64.7 ± 1.4	-234.7 ± 0.1	226.5 ± 0.3	-239.2 ± 0.4
C.3.2 Tree Structure
First, we analyze the necessity of the iterative growing procedure as opposed to fixing an (informed) tree structure a priori. Thus, we present an ablation, where we fix a reasonable tree structure and directly train the full tree as opposed to iteratively grow. The results in Table 6 suggest that the iterative learning is vital to the performance of TreeVAE. Intuitively, by fixing a tree structure, the model loses the ability to learn one out of multiple equivalent tree structures (e.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., 2021) representation of the learnt root embedding in Fig. 16. 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
Table 6: Test set performances (\%) of TreeVAE with a fixed structure standard TreeVAE. Means and standard deviations are computed across 10 runs with different random model initializations.
Dataset	Method	ACC	NMI	LL	RL	ELBO
Fashion	fixed TreeVAE	35.1 ± 4.8	54.1 ± 5.8	-237.4 ± 0.4	228.9 ± 0.5	-241.0 ± 0.5
TreeVAE	63.6 ± 3.3	64.7 ± 1.4	-234.7 ± 0.1	226.5 ± 0.3	-239.2 ± 0.4
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. We show it in Fig. 17. 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 groups of 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.
Figure 16: UMAP visualization of root embedding for Fashion MNIST.
Figure 17: "Average" dendrogram learnt on Fashion MNIST by averaging cluster distances across 10 runs and applying average linkage clustering.
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 a slightly modified version of CIFAR-10 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 as the two methods complement each other.
C.3.4 Unknown number of clusters
In an unsupervised setting, it is unclear a priori what the optimal number of clusters is. Thus, we are interested in the sensitivity of TreeVAE 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 we look at coarse or fine-grained classes. As a metric, we will only use NMI, because DP, LP, as well as Accuracy values are not fairly comparable across different numbers of predicted clusters. To determine the sensitivity of TreeVAE with respect to the misspecification of number of clusters, we evaluate its performance with 10, 20, and 32 predicted classes, compared to the ground truth 20 superclasses. The results are depicted in Table 8. Clearly, TreeVAE is not very sensitive to the correct specification of the number of clusters. This is also supported by Fig. 10/11, which show that even as we grow the tree further than the ground truth number of classes, TreeVAE still finds meaningful subsets of inside the classes.
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-up embeddings, (C) on a transposed version of the NT-Xent regularization on the routers, following Li et 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 of the two augmented inputs.
Dataset	Contrastive Method	NMI
CIFAR-10*	Ours (=A+B)	44.1 ± 0.7
(A)	21.7 ± 4.0
(B)	6.0 ± 6.0
(C)	23.0 ± 1.2
(D)	0.2 ± 0.0
(E)	1.0 ± 1.3
Table 8: Test set hierarchical clustering performances (%) of TreeVAE 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*	10	15.99 ± 0.85
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 or 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, as the ground truth number of classes can either be 20 or 100. Thus, in Table 9, we present how the performance of TreeVAE would change if we chose a different granularity of ground truth classes while keeping the number of predicted clusters fixed to 20. We can observe that the NMI for finer granularity is even higher than for the superclasses used, indicating that TreeVAE learned to differentiate classes of the same superclass.
This section provides supplementary background information and preprocessing steps of the different datasets used in our experiments.
The MNIST (LeCun et al., 1998) dataset is a widely used benchmark dataset in machine learning. It is composed of grayscale images representing handwritten digits from zero to nine. Each image is of size  $28 \times 28$  pixels that we rescale to  $[0,1]$ . The dataset is balanced, meaning that there is an equal number of examples for each digit class, with a 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 et al., 2017) dataset consists of a collection of grayscale images depicting various clothing items and accessories. The dataset contains a 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 a resolution of  $28 \times 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^{\prime}000$  test images.
Table 9: Test set hierarchical clustering performances (\%) of TreeVAE for different number of 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.80 ± 0.42
100	22.74 ± 0.36
D.3 20Newsgroups
The 20Newsgroups (Lang, 1995) dataset is a 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^{\prime}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 topics. We employ a TF-IDF vectorizer for the  $2^{\prime}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 et al. (2015) dataset contains handwritten characters from 50 different alphabets. Every alphabet contains various individual characters, amounting to a total of 1623 different characters in 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. In our experiments, we are interested in whether TreeVAE can find alphabet-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 to 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 are 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 \times 28$  to align with images in MNIST and Fashion-MNIST, transformed the images to grayscale, and rescaled pixel values to [0, 1]. For a description of the augmentations we use, we refer to 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, Horse, Ship, Truck}. Therefore, some classes are more similar (for example Automobile and Truck) than others (for example Frog and Airplane), indicating that there is a structure across clusters that can be captured. The dataset consists of  $50^{\prime}000$  training and  $10^{\prime}000$  test images of resolution  $32 \times 32$  pixels and 3 color channels, which we rescaled to [0, 1].
D.6 CIFAR-100
The CIFAR-100 (Krizhevsky & Hinton, 2009) dataset originally contains colored images of 100 classes. In order to not have to grow until TreeVAE has 100 clusters, we apply the common grouping of the 100 classes into 20 balanced superclasses. In Table 9, we also present an ablation if instead we chose the full 100 clusters. The dataset consists of  $50'000$  training and  $10'000$  test images of resolution  $32 \times 32$  pixels and 3 color channels, which we rescaled to  $[0,1]$ .
D.7 CelebA
The CelebA (Z. Liu et al., 2015) dataset contains images of 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 TreeVAE and validate whether they align with our intuition as well as with the given attributes. Thus, in Table 3, we 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 a subset of  $100'000$  random images from the training set and evaluate on the given test set of  $19'962$  images. Lastly, we crop the images to  $64 \times 64$  pixels with 3 color channels and rescale them to  $[0,1]$ .
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, we 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.
Hierarchical variational models, such as LadderVAE, are prone to local minima and posterior collapse, and TreeVAE shows similar behavior. Therefore, few technical choices, such as batch normalization and KL annealing, are needed to converge to a desirable optimum. For all datasets, we train TreeVAE only on the training set and evaluate the trained model on the separate test set.
E.1.1 Training the Tree