Tree Variational Autoencoders

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

Strengths

A strong, creative and well-written paper! 1. This paper makes a significant contribution towards learning meaningful hierarchical models, which is an important problem in the field 2. The solution described in this paper is creative and well presented (structured presentation: generative model vs. inference network, detailed derivations, intuitive notation) 3. The experiments are extensive (5 datasets, 4 baselines, 10 seeds for each run) 4. The visual results are very convincing 5. 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` 2. 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)

Questions

1. What is the modelling performance (likelihood) of Tree VAE without `NT-Xent`?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

Limitations of the method are not sufficiently discussed. The authors might want to discuss some of the following points more thoroughly 1. the multi-stage training of TreeVAE adds complexity, and result might vary depending on the choice of expansion criterion 2. the results depend of the choice of "prior knowledge" model

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.

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. 2) 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. 3) 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. 4) 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.

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. 2) Some important ablation studies are missing, such as data augmentation and the pre-determined maximum depth of the binary tree. 3) 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. 4) 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. 5) 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.

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. 2) 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. 3) 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? 4) 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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

see above

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.

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. 2. 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. 3. 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.

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. 2. 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. These 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.

Questions

Please refer to the above sections for detailed discussion.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes, there is a section at the end of the paper discussing limitations and future research opportunities.

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.

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.

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.

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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

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

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.

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

Questions

- How does the model perform if you start with a pre-specified tree structure? - 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. - 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$? - Equations 19 and 20 don't use $l$ anywhere within the summation, so it's unclear exactly how multiple samples are used.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors do discuss the limitations appropriately.

Thank you for the detailed response!

Having read through the rebuttal and global response, I am willing to change score to recommend acceptance. The authors have address several of my concerns. I still feel that some of the shortcomings in terms of the architectures, datasets and tree-building routine are worth further exploration, but I'm inclined to agree they could be left for future follow-up work. I will be interested to see such extensions!