TreeVI: Reparameterizable Tree-structured Variational Inference for Instance-level Correlation Capturing

We thank the reviewer for the comments and suggestions. Below we respond to each of the raised questions.

Q1: What do the numbers in the nodes represent in Figure 1?

A1: Sorry for the confusing notations used in the Figure 1. In Figure 1 (b) and (c), the numbers 1, 2, ... 5 in the trees should also be written as $\mathbf{z}_1, \mathbf{z}_2, \cdots, \mathbf{z}_5$. In our model, we assume for each data instance $\mathbf{x}_i$, it corresponds to a latent variable $\mathbf{z}_i$, and Fig. 1 (a), (b), (c) are used to discribe the correlation structures among differnt latent variables $\mathbf{z}_i$. The notations in this paper are different from those in TreeVAE because TreeVAE focuses on how to model the hierarchical structure of different dimensions within one instance, while ours devotes to modeling the correlation among different instances.

Q2: I also believe that the paper could benefit from a more thorough analysis of the computational complexity of the proposed appraoch.

A2: Thanks for the suggestions. For amortized variational inference, the computational complexity mainly comes from running the re-parameterization neural network, thus the complexity can be measured by how many runs of the re-parameterization neural network are required. For mean-field VI, each instance only need to pass the neural network once, thus for a dataset with $N$ instances, the number runs of the neural network is $O(N)$.

In our method, by restricting the correlation matrix $\mathbf{R}$ to a form constructed from a tree, we prove that all of the non-zero elements in $\mathbf{L}$ can be explicitly computed from the $|\mathcal{E}|$ parameters $\gamma_{ij}$ for $(i,j)\in \mathcal{E}$. Thus, we only need to re-parameterize the $|\mathcal{E}|$ parameters $\gamma_{ij}$, and run re-parameterization neural network $|\mathcal{E}|$ times, in addition to re-parameterizing the mean and variance, which requires a complexity of $O(N)$. For a tree, it is known that $|\mathcal{E}|\le N$, Thus, the total number of runs of the neural network is approximately $O(2N)$. This is also consistent with our observation in the experiments that the training time of our method is generally 2~3 times of mean-field methods.

Q3: Standard deviations are not reported in the table, and few experimental settings are explained. E.g. At what value is M set in the experiments?

A3: We thank the reviewer for pointing out the missing details. The results of constrained clustering with standard deviations are supplemented in our attachment to the global rebuttal, and will be included into our final version.

In our experiments, the value of $M$ is set to 3, and more detailed experimental settings are elaboratred in Appendix E.2. We will relocate some of the settings into the main context of the paper in our final version.

Q4: In 4.2 it would be more convincing to provide evidence that the learnt tree are consistent with the instance-level constraints.

A4: Thanks for the suggestion. Our method is mainly used to capture the inherent correlations among different instances. Thus, the correlations captured by the trees are not necessary to be consistent with the constrains, but instead are more inclined to align with the similarities among instances. To validate this, we plot the tree learned over the constrained clustering experiment on MNIST dataset, as shown in the Figure 1 in the attachment of global rebuttal. From the figure, we can see that instances from the same category are connected more tightly than those from different categories in the learned tree. This demonstrates that our method has the abilities to capture the underlying inherent correlations among different instances.

Q5: Could the authors provide a proof of Equation 4?

A5: Equation 4 is a result of the Markov tree dependent distributions, for a complete proof we can refer to [1] which is mainly done by induction. By the way, this equation can also be regarded as a special case of the Hammersley-Clifford theorem. We will add these citations in our final version.

Reference

1. Meeuwissen A.M.H. and Cooke R.M., Tree dependent random variables Report 94-28, Delft University of Technology, 1994.

Q6: 4.1 Why the tree structures are defined a-priori and not learnt from the data? Validate the correctness of the the trees in the synthetic scenario.

A6: The experiment in Sec. 4.1 is a toy experiment to illustrate the contribution of each component of our method in capturing the accurate correlation. For better controlability, we choose to use pre-defined trees, instead of trees learned from the data. That is because if we learn the trees from the data, the tree $\mathbb{T}_1$ in MTreeVI($\mathbb{T}_1$, $\mathbb{T}_2$) and the tree $\mathbb{T}_1$ in TreeVI($\mathbb{T}_1$) could be different. Thus, we choose to fix the two trees $\mathbb{T}_1$ and $\mathbb{T}_2$. As seen from the experiment, when using a single tree, either $\mathbb{T}_1$ or $\mathbb{T}_2$, better performance is obained comparing with the mean-field method. When the two trees are used simultaneously, we can see that the performance is further boosted.

Since the 4 dimensions in the toy experiment are all closely correlated, any tree would be good for capturing the correlation among the 4 dimensions. Thus, it is hard to clearly show which tree is better than the other. But we have demonstrated the consistency between the learned tree and the underlying correlation in the constrained clustering experiments, which is stated in the answer to Q4.

Q7: How to compute the latent embedding distance

A7: We utilize the same measurement of latent embedding distance as CVAE for fair comparison, which is computed based on the expected quadratic L2 distance $\mathbb{E}_{q(\mathbf{z}_i,\mathbf{z}_j)}[\| \mathbf{z}_i-\mathbf{z}_j \|^2]$ of the latent distribution for VAE-based methods. Additionally for GraphSAGE, we simply use quadratic L2 distance $\| \mathbf{z}_i-\mathbf{z}_j \|^2$.

We thank the reviewer for the comments and suggestions. Below we respond to each point raised by the reviewer:

Q1: It is unclear how important the initialization is and if there always exists a straightforward initialization procedure.

A1: Below we show the constrained clustering accuracies (%) on the MNIST dataset, with different initializations (random tree or greedy-searched tree), with or without our constrained optimization. It can be seen that without constrained optimization, both initializations underperform. And with constrained optimization, both initializations converge to substantially better performances, with the greedy-searched initialization slightly better. Overall, the constrained optimization procedure makes our proposed TreeVI less sensitive to initializations.

Q2: The maximum number of edges in a DAG with $n$ nodes is $\mathcal{O}(n^2)$. So in the worst-case instances the computational cost could still be quite high

A2: In our method, we require to use a tree, instead of a DAG, to capture the correlation among instances. For a tree, the number of edges $|{\mathcal{E}}|$ is at most $N-1$, where $N$ denotes the number of vertices. Thus, the maximum number of edges is $O(N)$, instead of $O(N^2)$. As proved in our paper, by using our method, we only need to re-parameterize the edge parameters $\gamma_{ij}$ for $(i,j)\in {\mathcal{E}}$, and thus only need to run the neural network by $|{\mathcal{E}}|$ times for each epoch, in addition to the runs required by a mean-field method, which is $O(N)$. Thus, the total number of runs of the re-parameterization neural networks is approximately $2N$, roughtly two times of that required by mean-field variational inference. That's why we observed that the training time of our method is generally 2~3 times of the mean-field methods in our experiments.

Q3: Depending on the task it might be more or less clear how to compute meaningful neighborhood probabilities.

A3: The probability with respect to each pair of data points is calculated according to the cosine similarities between them, which can be leveraged to generate a meaningful neighborhood for each data instance. For more details we refer to Appendix D.2.

Q4: Something small that I noticed. The acronyms ARI and NMI have not been introduced and might not be familiar to readers not familiar with clustering.

A4: Sorry for the missing explanations for the acronyms of clustering metrics. These details will be included into our final version.

We appreciate your detailed comments, but we believe that you misunderstood our paper deeply. We hope our clarifications could help you recognize our contributions correctly.

You mistakenly think that our main contribution is simply to Cholesky decompose the correlation matrix $\mathbf{R}$ as $\mathbf{R}=\mathbf{L}\mathbf{L}^\top$ and then re-parameterize the elements $\ell_{ij}$ in the lower-triangular matrix $\mathbf{L}$ by a neural network $f_\phi(\cdot)$, that is, $\ell_{ij}=f _ \phi(\mathbf{x} _ i, \mathbf{x} _ j)$, with $\mathbf{x} _ i$ denoting the $i$-th data instance. We want to emphasize that this is not our contribution. If this approach is adopted, we have to re-parameterize as many as $N(N+1)/2$ elements, i.e., all elements from the lower-triangular positions of $\mathbf{L}$ need to be re-parameterized, where $N$ is the number of instances in training dataset. That means we need to run the neural network $f_\phi(\cdot, \cdot)$ by $O(N^2)$ times for every epoch, which is computationally unacceptable, especially considering that $N$ could be as large as millions in practice.

Our main contribution lies at finding a way to reduce the required times of running the neural network $f_\phi(\cdot, \cdot)$ from $O(N^2)$ to $O(N)$, making it applicable to scenarios even with millions of instances. To achieve this goal, we propose to restrict the correlation matrix $\mathbf{R}$ to a special form constructed from a tree $\mathbb{T}(\mathcal{V}, \mathcal{E})$. It should be pointed out that although $\mathbf{R}$ is constructed from a tree, it doesn't represent a tree. Actually, it is a dense matrix, with its $(i,j)$-th element $[\mathbf{R}] _ {ij}$ for $(i,j)\notin \mathcal{E}$ set as $\prod_{(s,t)\in \mathbb{P} _ {i\to j}}\gamma_{st}\triangleq \tilde \gamma_{ij}$.

We prove that under the restricted correlation matrix $\mathbf{R}$, the lower-triangular matrix $\mathbf{L}$ possesses a very elegant form, as shown in Eq. (11) in the paper. The elegance lies at that although $\mathbf{L}$ still has $O(N^2)$ non-zero elements, all of these non-zero elements can be explicitly computed from the $|\mathcal{E}|$ parameters $\\{\gamma_{ij}\\} _ {(i,j)\in \mathcal{E}}$. For instance, the $(1, 4)$-th element in Eq.(11) can be computed as $\prod_{(s,t)\in \mathbb{P} _ {1\to 4}}\gamma_{st}$, where $\mathbb{P} _ {1\to 4}$ denotes the path conneting node 1 and 4 in the tree ${\mathbb{T}}$. The importance of this conclusion is that now we only need to re-parameterize the parameters $\\{\gamma_{ij}\\} _ {(i,j)\in \mathcal{E}}$, while computing the other non-zero elements in $\mathbf{L}$ with $\\{\gamma_{ij}\\} _ {(i,j)\in \mathcal{E}}$. In this way, we only need to re-parameterize $|\mathcal{E}|$ parameters, and thus only need to run $|\mathcal{E}|$ times of neural network $f_\phi(\cdot, \cdot)$, instead of $O(N^2)$ times in the vanilla method. For a tree, the number of edges $|\mathcal{E}|\le N -1$, thus we only need to additionally run $O(N)$ times of neural networks $f_\phi(\cdot, \cdot)$, in addition to the runs required by the mean-field method.

For a mean-field variational inference that assumes independence among instances, for each epoch, it also need to run $O(N)$ times of neural network. Thus, the complexity of our proposed method is roughly only 2 times of the mean-field method. This is consistent with the elapsed time observed in the experiments.

Below we respond to each point raised by the reviewer:

Q1: A lot of the derivation in Section 2.1 is redundant. The authors could have simply parameterized the Cholesky factors.

A1: As discussed above, directly parameterizing the Cholesky factors is computationally consuming, which requires to run the neural network $f_\phi(\cdot, \cdot)$ by $O(N^2)$ times. Section 2.1 is mainly devoted to how to achieve a correlation-capturing variational inference method that only requires running the neural network by $O(N)$ times.

Q2: Why to take a tree structure among various different correlation graph structures.

A2: As elaborated above, by restricting the correlation matrix ${\mathbf{R}}$ to a form that is constructed from a tree, we can prove that all of the non-zero elements in ${\mathbf{L}}$ can be explicitly computed from the $|{\mathcal{E}}|$ parameters $\gamma_{ij}$ for $(i,j)\in {\mathcal{E}}$. Then, we only need to re-parameterize the $|{\mathcal{E}}|$ parameters $\gamma_{ij}$, instead of all the non-zero elements in ${\mathbf{L}}$. That's why we have to restrict the correlation matrix ${\mathbf{R}}$ to be constructed from a tree. Again, we want to point out that ${\mathbf{R}}$ is only constructed from a tree, but it is a dense matrix.

Q3: What the constrained optimization problem in Eq 17 is supposed to do.

A3: As discussed above, we need to construct a tree to learn the instance-level correlations. Besides trivial methods like heuristic methods and greedy search algorithm, it can also be learnt from data by our constrained optimization. Specifically, we use a real matrix to represent the correlation parameters, which satisfies the acyclicity constraint if and only if the indicator function $h(\mathbf{A}) = 0$ as shown in Eq. 17. This allows us to formulate a constrained optimization problem with a smooth equality constraint, which can be solved with Lagrangian multipliers by incorporating a penalty term for the indicator function.

Q4: Imprecise about amortized variational inference not just variational inference.

A4: We apologize for the impreciseness brought by our terminologies, and we will reflect this in our title of the final version.

Q5: Sticking-the-landing gradient and MCMC-based augmentation methods.

A5: We thank the reviewer for pointing out some feasible alternatives. However, the MCMC-based augmentation methods may cause inefficiency by introducing time-consuming sampling procedure.

We thank the reviewer for the comments and suggestions. Below we respond to each point raised by the reviewer:

Q1: The study could be enhanced by providing some insights derived from the tree structures obtained through training.

A1: Thanks for the suggestions. Our method is mainly used to capture the inherent correlations among different instances. Thus, the learned trees are inclined to reflect the correlation among different instances. To see this, we plot the tree learned over the constrained clustering experiment on MNIST dataset, which is shown in the Figure 1 in the attachment of global rebuttal. From the figure, we can see that instances from the same category are connected more tightly than those from different categories in the learned tree. This demonstrates that our method has the abilities to capture the underlying inherent correlations among different instances. The figure and the explanation to it will be included in our final paper.

Q2: Discrepancy of notations used in the synthetic VAE experiment

A2: The discrepancy arises becuase in this toy experiment, to demonstrate our method is able to catpure correlaitons, we use it to capture correlation among different dimensions within one instance, ranther than capturing correlations among different instances. That is why we see the number of instances $N=6000$, but we only set the vertex set as ${\mathcal{V}}=\{1, 2, 3, 4\}$. That is because the number of dimensions of the latent space is 4 in this toy experiment. We will make this point clearer in our final version.

Yes, your understanding is correct. But the subscript $(i, j)\in \\{1, \cdots, D \\}^2$ need to be written as $(i, j)\in {\mathcal{E}}$ since only a fraction of instances are connected in the tree.

I believe the original equation (4) was the approximate joint posterior distribution of $\boldsymbol{z}\_1, \boldsymbol{z}\_2, \dots , \boldsymbol{z}\_N$ given $\boldsymbol{x}\_1, \boldsymbol{x}\_2, \dots , \boldsymbol{x}\_N$. The equation above seems to be for a single instance. Is my understanding correct that the original equation (4) can be written out as follows?

$q\_{\boldsymbol{\phi}}^{\mathbb{T}}(\boldsymbol{Z}|\boldsymbol{X}) = \prod\_{n=1}^N \left( \prod_{i \in \\{ 1, \dots , D \\}} q\_{\boldsymbol{\phi}} (z\_{n,i}|\boldsymbol{x}\_n) \prod\_{(i,j) \in \\{ 1, \dots , D \\}^2} \frac{q\_{\boldsymbol{\phi}}(z\_{n,i}, z\_{n,j}|\boldsymbol{x}\_n)}{q\_{\boldsymbol{\phi}}(z\_{n,i}|\boldsymbol{x}\_n)q\_{\boldsymbol{\phi}}(z\_{n,j}|\boldsymbol{x}\_n)} \right)$,
where $z_{n,i}$ denotes the $i$-th element of the $n$-th latent embedding $\boldsymbol{z}_n$.