SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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_f = 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 tree after every third growing step for 80 epochs with KL annealing of 0.001, so the otherwise frozen parameters can adapt to 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 KL term is annealed every epoch by 0.01 such that the weight reaches 1 after 100 of the 150 epochs every growth step. While the annealing procedure is required to achieve good performances, the exact annealing decay plays a smaller role in the final performances, as well as the in-between finetuning of the trees, and the number of epochs at each training step.
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  $\pm 10$  degrees, horizontally/vertically shifting images by up to  $\pm 2$  pixels, applying up to  $1\%$  shearing, and zooming by up to  $\pm 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, CIFAR-100, and CelebA, augmentations were applied to use contrastive learning as introduced in Section 2.6 and Appendix E.3. For CIFAR-10 and CIFAR-100, first, each image is randomly cropped and resized with scale  $\in (0.2, 1)$  and aspect ratio  $\in (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  $\in (0.6, 1.4)$ , as well as change in hue with its parameter
$\in (-0.1, 0.1)$ . Lastly, with probability 0.2, we transform the image to grayscale. The chosen parameters are largely adopted from Li et al.'s (2022) weak augmentation scheme.
While these augmentations lead to superior group separation, in CelebA, we are also interested in generative results. Therefore, to obtain generations without non-sensical artifacts, we reduce the augmentation strength by removing the hue and grayscale augmentations, reducing the other color parameters by a factor 2, and reducing the cropping parameters to scale  $\in (3/4, 1)$  and aspect ratio  $\in (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 2 and 3, and identical decoders for each leaf of the tree. We always apply batch normalization followed by Leaky ReLU non-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 Transformations/Bottom-Up
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), or 64 (CIFAR-10, CIFAR-100, CelebA) and map them to some intermediate representation of dimension 128 (20Newsgroup, MNIST, Fashion-MNIST, Omniglot), or 512 (CIFAR-10, CIFAR-100, CelebA) using a dense layer, followed by a dense layer without activation to compute  $\mu$  and one with Softplus activation to compute  $\sigma$  of the approximate posterior/conditional prior of the respective tree level.
E.2.2 Routers
Similar to transformations, the architectures of routers are identical across all experiments. The generative routers take as input  $\mathbf{z}_i$ , the latent representation of a node  $i$ , while the inference routers use  $\mathbf{d}_{depth(i)}$ . Then, both sequentially apply two dense layers that both map them to 128 (20Newsgroup, MNIST, Fashion-MNIST, Omniglot), or 512 (CIFAR-10, CIFAR-100, CelebA) dimensions, after which a final dense layer with a Sigmoid activation function computes the probability of the binary decision that infers the next child node.
E.2.3 MNIST/Fashion-MNIST
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 \times 28$  and 8 dimensional latent spaces. We use a small encoder as depicted in Figure 18, and a symmetric decoder. The encoder uses  $3 \times 3$  convolutions with stride 2, whereas the decoder uses transposed convolutions.
Figure 19: 20Newsgroups encoder architecture.
E.2.4 20Newsgroups
20Newsgroups is a 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 a 2000D vector (see Section D.3). Hence, the encoder for 20Newsgroups is a simple MLP with 5 dense layers, as depicted in Figure 19, and the decoder again mirrors the encoder architecture.
E.2.5 Omniglot
Figure 20: Omniglot encoder architecture.
For Omniglot, we use an architecture similar to the one for MNIST and Fashion-MNIST. In contrast to Section E.2.3, we increase the receptive field of the encoder by using  $4 \times 4$  convolutions and make it 4 times wider and 2 times deeper. We refer to Figure 20 for a schematic of the encoder. As in the other architectures, the decoder is symmetric to the encoder.
E.2.6 CIFAR-10/CIFAR-100/CelebA
For real-world image datasets, more complex architectures are required. We first define a ResnetBlock as two consecutive convolutions with kernel size 3 and stride 1, whose output gets weighted by 0.1 and combined with a residual connection to 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 1 and stride 1 to match dimensionality. The encoder consists of consecutive ResnetBlocks with 2D average-pooling of kernel size 3 and stride 2. Here, number of filters per ResnetBlock are increasing by a factor of 2, starting from 32 and going up to 256 until the representation is  $4 \times 4 \times 256$ , at which point we flatten it to get  $\mathbf{d}_L$ . We refer to Figure 21 for a schematic of the encoder. The decoder structure is similar to the reversed encoder, where the pooling layers are replaced by 2D upsampling with bilinear interpolation and before the ResnetBlocks, we pass the input through a 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, we add one more ResnetBlock with 256 filters to the encoder, as well as one more ResnetBlock at the start of the decoder, starting from 16 instead of 32 filters.
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 TreeVAE 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 TreeVAE 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 to 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 are applied to it. For the augmentations that were used for contrastive learning, we refer to Appendix E.1.2.
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:  $\ell_{i,j} = -\log \frac{\exp(s_{i,j} / \tau)}{\sum_{k=1}^{2N}\mathbb{1}_{[k\neq i]}\exp(s_{i,k} / \tau)}$ , where  $s_{i,j}$  denotes a similarity measure, which we define as the cosine similarity between the representations of  $\tilde{x}_i$  and  $\tilde{x}_j$ , and  $\tau$  is a temperature parameter. This is computed for all  $2N$  pairs of augmented samples  $\tilde{x}_i$  and  $\tilde{x}_j$  originating from the same initial sample  $x$  and then averaged. We compute this depth-wise for every embedding  $\mathrm{d}_l$  by first passing it through a separate MLP 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-Xent with  $\tau = 0.5$  is applied. Then, we average the pairwise losses and sum them up over all  $\mathrm{d}_l$ , where we, to be invariant to the depth of the bottom-up, divide by the maximum depth. With this, we regularize the bottom-up embeddings learned by the model to encode every augmented pair of samples close to each other, thus constraining the learning of information that the model should be invariant to.
Even with the previous regularization,  $\mathbf{d}_l$  might still contain some of information that is not 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. (2021, 2022), which propose not only to regularize the embeddings, but also the cluster probabilities themselves. While Li et al. (2021, 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  $\tau = 1$  directly on the
outputs of every bottom-up router  $q(\mathbf{c}_i \mid \mathbf{d}_l)$ . 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 is 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: (i) 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. (iii) We can omit the computationally intensive finetuning of the full tree, as 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 to match the gradients of the ELBO. Thus, we naturally introduce a 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.
For a fair comparison, we retain the same architectural choices made for TreeVAE 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 TreeVAE applies contrastive learning, we similarly apply the regularization on the bottom-up embeddings of the baselines. Below we describe specific training details that deviate from TreeVAE, due to the differences in the graphical models.
E.4.1 Variational Autoencoder
The non-hierarchical counterpart of TreeVAE is a classical Variational Autoencoder (VAE). As it does not contain a hierarchical structure, the VAE consists solely of an encoder and decoder architecture, which follows the architectural choices of TreeVAE. We train the VAE for 500 epochs and use the KL 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 TreeVAE is the LadderVAE (Sønderby et al., 2016). We retain the same encoder, decoder, and transformations used by TreeVAE, and we set the depth to 5 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 0 to 1 in 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 VAE (VAE + Agg) and the last layer of LadderVAE (LadderVAE + Agg). We observed that clustering on the last layer of LadderVAE 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 as single-linkage agglomeration clustering or bisecting K-means). However, Ward's method cannot be tested on a held-out dataset, as it inherently does not allow for prediction on 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 a fairer comparison with the other methods.
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 20 leaves on 20Newsgroup takes approximately 30 minutes, and on CIFAR-100 9h. Training TreeVAE on CelebA
takes approx 8h. Please note that we only report the numbers to generate the final results but not the development time.