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

Authors: junxi  xiao
qinliang  su
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 15:58:38

Abstract: Mean-field variational inference (VI) is computationally scalable, but its highly-demanding independence requirement hinders it from being applied to wider scenarios. Although many VI methods that take correlation into account have been proposed, these methods generally are not scalable enough to capture the correlation among data instances, which often arises in applications with graph-structured data or explicit constraints. In this paper, we developed the Tree-structured Variational Inference (TreeVI), which uses a tree structure to capture the correlation of latent variables in the posterior distribution. We show that  samples from the tree-structured posterior can be reparameterized efficiently and parallelly, making its training cost just 2 or 3 times that of VI under the mean-field assumption. To capture correlation with more complicated structure, the TreeVI is further extended to the multiple-tree case. Furthermore, we show that the underlying tree structure can be automatically learned from training data. With experiments on synthetic datasets, constrained clustering, user matching and link prediction, we demonstrate that the TreeVI is superior in capturing  instance-level correlation in posteriors and enhancing the performance of downstream applications.

Main paper: 
 introduction: 
 Variational inference is a probabilistic method that is widely used for approximating the exact posterior in latent-variable models p(X, Z) = p(X|Z)p(Z), with X and Z being the observed data and latent variable, respectively. When not considering the existence of any relation among instances, we can write the model into a factorized form as p(X, Z) = i p(x i |z i )p(z i ), where x i and z i denote the i-th data instance and latent variable. Thanks to the factorized form, we can reasonably use the mean-field posterior q(Z|X) = i q(z i |x i ) for model inference and training. However, there are also many circumstances, under which complex relations among instances exist. For instances, in graph-structured data, VGAE [19] uses two instances' latent representations to model the probability of existing an edge e ij between them via p(e ij |z i , z j ), which, obviously, makes all instances coupled together. Also, in the constrained clustering, DC-GMM [24] designs a prior p(c; W) to specify the assignment constraints and then uses it to construct the prior distribution p(Z, c) = p(c; W) i p(z i |c i ), which will also cause all z i to be correlated with each other. Obviously, under these circumstances, due to the existence of correlation among instances, it is unreasonable to assume a factorized form for the posterior anymore.There has been many work on incorporating correlation structures into posterior approximation in variational inference recently [23, 31, 43] , and they are mainly focused on modeling correlation structure among different dimensions within each latent variable. However, the amount of latent dimensions is usually no more than a few hundred, at a completely different level from the size of dataset which might be as large as millions. Therefore, these variational inference methods modeling correlations among dimensions cannot be leveraged to capture instance-level correlation especially for large datasets. To investigate correlations among different instances, Tang et al. [36] recently used a tree structure over latent variables to take instance-level correlations into consideration, but only capable of capturing correlations between immediately connected nodes.In this work we propose a novel posterior approximation for variational inference, the Tree-structured Variational Inference (TreeVI) , that models the latent correlation structure with a tree-structureinduced distribution and enables modeling high-order correlations between non-adjacent latent variables. Our method provides a matrix-form reparametreization for tree-structured correlated latents based on ancestral sampling, that is scalable to large datasets with computational complexity comparable to mean-field variational inference methods. To better capture structural correlation, we also generalize TreeVI to multiple trees, and propose a continuous optimization algorithm to stochastically learn a theoretically correlation-rich tree or mixture-of-trees structure. 
tree-structured variational inference: 
 Consider a latent-variable model p θ (X, Z) = p θ (X|Z)p(Z), whereX = [x ⊤ 1 , x ⊤ 2 , • • • , x ⊤ N ] and Z = [z ⊤ 1 , z ⊤ 2 , • • • , z ⊤ N ]; and x i and z i ∈ R D denote the i-th data instance and its corresponding latent variable, respectively. Given the latent-variable model p θ (X, Z), variational inference aims to find a distribution that is closet to the true posterior p θ (Z|X) from a distribution family Q, which is achieved by maximizing the evidence lower bound (ELBO)EQUATIONwhere q ϕ (Z|X) ∈ Q denotes the approximate posterior parameterized by ϕ. In most of current works, the latent-variable model is assumed to take a factorized form as p θ (X, Z) = N i=1 p θ (x i |z i )p(z i ) due to the observed independence among data instances. With the factorized form for the model, it can be seen that the true posterior p θ (Z|X) also takes a factorized form, hence it is reasonable to choose a fully factorized form q ϕ (Z|X) = N i=1 q ϕ (z i |x i ) for the approximate posterior. However, for applications like constrained clustering with generative model and generative modeling on graph data, we often need to impose some constraints among the latent variables {z i } N i=1 by using a correlated prior distribution p(Z) with p(Z) ̸ = N i=1 p(z i ), or use several latent variables z j to be responsible for generating a data instance x i (e.g., x i representing an edge), which will result in p θ (X|Z) ̸ = N i=1 p θ (x i |z i ). Obviously, under these scenarios, latent variables {z i } N i=1 from the true posterior p θ (Z|X) are not independent anymore. Thus, if a factorized posterior q ϕ (Z|X) = N i=1 q ϕ (z i |x i ) is still employed to approximate the true posterior, the model will lose the ability to capture the correlations among data instances. To alleviate the deviation caused by mean-field approximation, we propose to approximate the correlation structure with a tree or a mixture-of-tree structure as shown in Fig. 1 , respectively.To capture the correlations among latent variables {z i } N i=1 in the posterior distribution, a multivariate normal distribution q(Z|X) = N (Z; µ z , Σ z ) can be used, whereµ z = [µ ⊤ 1 , • • • , µ ⊤ N ] ⊤ ∈ R N D is the mean and Σ z = diag(σ z )R diag(σ z ) is the N D × N D covariance matrix; σ z = [σ ⊤ 1 , • • • , σ ⊤ N ] ⊤ ∈ R N Dis the standard deviation; and R means the correlation matrix, with its (i, j)-th block taking the formEQUATIONγ ij ∈ (-1, 1) D for i ̸ = j denotes the correlation strength between latent variable z i and z j andγ ii ≜ 1 D for i ∈ V ≜ {1, 2, • • • , N }.Here, the diagonal structure assumed for the (i, j)-th block [R] ij implies that we are only interested in capturing the correlation among different variables, without considering the correlation of different dimensions within one variable. Actually, due to the (c) MTreeVI: Extension to Multiple Trees Figure 1 : Comparison between the fully-connected correlation structure and approximation by single and multiple tree-structured variational inference relatively low dimensions of a latent variable z i , the correlation within a variable can be captured by combining with existing VI methods, but we here omit the modeling of correlation among dimensions for simplicity. In the following context, the diagonal operators are omitted for conciseness.A crucial part of the variational inference is to determine appropriate values for the correlation parameters Γ 0 = {γ ij : i ̸ = j ∈ V}. For a general multivariate normal distribution, all correlation parameters in Γ 0 are learnable, which admits a highly expressive correlation structure, but also resulting in high computational complexity. The widely-used mean-field approximation reduces the cost by assuming an independence structure within any pair of latent variables, i.e., γ ij = 0 D for any i ̸ = j, completely ignoring the correlation among latent variables. In order to achieve a balance between expressiveness and complexity, we propose to approximate the fully-connected correlation structure with a tree structure, introducing our TreeVI, namely tree-structured variational inference.The idea of tree-structured variational inference is to impose a tree correlation structure T = (V, E) among latent variables, as illustrated in the quin-variate example in Fig. 1b . In our TreeVI, given a tree structure T = (V, E), we only specify the correlation parameters γ ij for latent variables (z i , z j ) adjacent on the tree, and for any non-adjacent latent variables (z i , z j ), the correlation between then are directly computed as the multiplication of correlation parameters along P i→j , that isEQUATIONwhere P i→j denotes the path connecting the latent variable z i to z j on the tree T, and the vector multiplication is assigned as Hadamard product. Therefore, only the correlation strengths w.r.t. the edges of the tree T are learnable, which can be denoted as Γ T = {γ ij : (i, j) ∈ E} ⊂ Γ 0 . The parameterization leads to a correlation matrix [R (T) ] ij = γij that is totally determined by parameters in Γ T , with γii = 1 D for i ∈ V. For example, for the quin-variate tree structure in Fig. 1b , the correlation parameter between z 1 and z 5 is fixed as γ15 = γ 12 ⊙ γ 23 ⊙ γ 35 .With the tree-structured correlation defined above, it can be easily shown that the distribution q T ϕ (Z|X) forms a Markov random field defined by the tree T [3] , and its joint probability can be expressed asEQUATIONwhere the marginal distribution q ϕ (zi |x i ) = N (z i ; µ i , diag(σ 2 i )), and pairwise marginal distribution for adjacent variablesEQUATIONγ ij is the correlation parameter between two adjacent variables. Furthermore, we can also see that the distribution q T ϕ (Z|X) can be represented by an acyclic Bayesian network of the form q T ϕ (Z|X) = q ϕ (z 1 )EQUATIONwhere the conditioning on data x i is omitted for conciseness; z 1 is assumed to be the root node; the conditional distribution of latent variables with respect to an edge (i, j) ∈ E isEQUATIONWith the conditional distribution q ϕ (z j |z i ), joint samples (z 1 , z 2 , • • • , z N ) can be drawn from q T ϕ (Z|X) with ancestral sampling. By sampling z i one by one with ancestral sampling, we can show that the joint sample (z 1 , z 2 , • • • , z N ) can be equivalently represented by a set of N independent Gaussian noises {ϵ i } N i=1 with ϵ i ∼ N (0 D , I D ), as stated in the theorem below.EQUATIONwhere ϵ i ∼ N (0 D , I D ) is a Gaussian random noise and pa(i) denotes the parent node of z i with respect to i ∈ V.For the proof and a concrete example, we refer to Appendix B. With the reparameterization of sample z j in Eq. ( 8), the joint sampleZ = [z ⊤ 1 , • • • , z ⊤ N ]⊤ can be re-parameterized in a matrix-form asEQUATIONwhereEQUATIONEQUATIONActually, if the variables are indexed according to their positions on the tree T from left to right and then top to bottom, it can be easily shown thatL (T)z is a lower-triangular block matrix. For example, the matrix L (T) z corresponding to the quin-variate example of Fig. 1b can be written as follows 11) where diag(•) has been omitted for conciseness. By parameterizing the mean µ i , standard deviation σ i and the correlation parameters γ ij for (i, j) ∈ E with a neural network f ϕ (•, •) asEQUATIONEQUATIONaccording to Eq. ( 9), samples drawn from q T ϕ (Z|X) can be reparameterized in the form of neural network parameters ϕ and random Gaussian noise ϵ, which can facilitate the training of variational inference significantly. Importantly, except the mean and standard deviation, we only need to reparameterize the parameters Γ T = {γ ij : (i, j) ∈ E}, while computing the other non-zero elements in L (T) z with Γ T . In this way, we only need to re-parameterize |E| parameters, and thus only need to run |E| times of neural network f ϕ (•, •), instead of O(N 2 ) times in the vanilla method. For a tree, the number of edges |E| ≤ N -1, thus we only need to additionally run O(N ) times of neural networks f ϕ (•, •), 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 needs 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. As seen in the experiments on constrained clustering, the consumed time of our TreeVI is only 2 to 3 times that of mean-field variational inference methods.It should be pointed that if we simply apply Cholesky decomposition to the correlation matrix R to produce R = LL ⊤ and then directly re-parameterize the elements ℓ ij in the lower-triangular matrix L by a neural network f ϕ (•, •), that is, ℓ ij = f ϕ (x i , x j ), then we have to re-parameterize as many as N (N + 1)/2 elements, i.e., all elements from the lower-triangular positions of 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 ϕ (•, •) by O(N 2 ) times for every epoch, which is computationally unacceptable, especially considering that N could be as large as millions in practice. And our main contribution lies at finding a way to reduce the required times of running the neural network f ϕ (•, •) from O(N 2 ) to O(N ) by restricting the correlation matrix R to a special form R (T) constructed from a tree T = (V, E), which is actually a dense matrix with its (i, j)-th element [R (T) ] ij = γij defined as Eq. ( 3) for (i, j) / ∈ E. Under the restricted correlation matrix R (T) , the lower-triangular matrix L (T) z possesses a very elegant form, as shown in Eq. (10) . The elegance lies at that although L (T) z still has O(N 2 ) non-zero elements, all of these non-zero elements can be explicitly computed from the |E| parameters Γ T = {γ ij : (i, j) ∈ E}.With the tree-structured posterior, the data log-likelihood has the evidence lower bound log p(X) ≥E Z∼q T ϕ (Z|X) [log p θ (X, Z)] + H[q T ϕ (Z|X)],where the first term can be estimated by reparameterization Z (T) ∼ q T ϕ (Z|X), and the entropy H of the tree-structured posterior distribution q T ϕ (Z|X) = i∈V q ϕ (z i |x i ) (i,j)∈E q ϕ (zi,zj |xi,xj ) q ϕ (zi|xi)q ϕ (zj |xj ) can be decomposed as entropy terms with respect to singleton posterior q ϕ (z i |x i ) and pairwise posterior q ϕ (z i , z j |x i , x j ) which can be both directly computed (for detailed expressions we refer to Appendix). Therefore, the optimization of variational inference with tree-structured posterior approximation can be simplified as maximizing the following sampling-based evidence lower bound of our proposed TreeVIEQUATIONwhere Z (T) denotes the tree-structured reparameterization for latent variables with respect to the tree structure T. For detailed computation for the evidence lower bound, we refer to Appendix C.1.The expressiveness of a single tree-structured posterior is still restrictive. To alleviate this issue, we propose to approximate the true posterior with a mixture-of-trees posterior distributionEQUATIONwhere we use a weighted mixture of tree structures MT = {T 1 , • • • , T M }, as shown in Fig. 1c , to approximate the underlying correlation structure, and the posterior q Tm ϕ (Z|X) with respect to the m-th tree component is expressed in the form of Eq. ( 4). Each tree component With the mixture-of-trees posterior, the data log-likelihood has the evidence lower bound log p(X)T m = (V, E m ) corresponds to a set of latent indices V = {1, • • • , N }, a set of pairwise connections E m , a set of correlation parameters Γ Tm = {γ m ij : (i, j) ∈ E m } to be learned,≥ M m=1 π m E Z∼q Tm ϕ (Z|X) [log p θ (X, Z)] + H[q MT ϕ (Z|X)],where the first term can be estimated by tree-structured reparameterization Z (Tm) ∼ q Tm ϕ (Z|X) within each tree component T m for m = 1, • • • , M , but the entropy term has no explicit expression. Since the entropy H[p] is concave in the probability distribution p, the entropy of the mixture-of-trees posterior distribution is lower bounded by the weighted sum of entropies with respect to each tree component M m=1 π m H[q Tm ϕ (Z|X)], which can be similarly computed by the singular and pairwise marginal posterior distributions with respect to each tree component T m , m = 1, • • • , M . Therefore, the optimization of variational inference with mixture-of-trees posterior approximation can be simplified as maximizing the following sampling-based evidence lower bound of our proposed MTreeVIEQUATIONwhere Z (Tm) denotes the tree-structured reparameterization for latent variables with respect to the tree structure T m . For detailed computation for the evidence lower bound, we refer to Appendix C.2.Our proposed tree-structured and mixture-of-trees structured posterior approximation make variational inference with pairwise latent correlations feasible, but a specific tree structure over latent variables need to be determined in advance. In this section, our goal is to develop an efficient algorithm for learning correlation-rich tree structures to approximate the underlying posterior correlation structure.To learn a meaningful tree structure from the latent embeddings, we adopt a symmetric binary matrix A ∈ {0, 1} N ×N with entries 0 along the diagonal, where each element a ij ∈ {0, 1} represents an edge between latents z i and z j , i ̸ = j ∈ {1, • • • , N } (here we assume D = 1 for convenience).Proposition 1. Suppose that the symmetric adjacency matrix A ∈ {0, 1} N ×N , then the undirected structure induced by adjacency matrix A is acyclic if and only ifEQUATION)where tr(•) and exp(•) represent the trace and exponential of a matrix respectively, and ⊙ is the Hadamard product.For the proof we refer to the Appendix D.1. Inherited from the idea of Zheng et al. [45] , we use a similar indicator function h(A) in Eq. ( 16) to check the acyclicity of the structure induced by the symmetric binary matrix A. If and only if h(A) = 0, the adjacency matrix A determines a unique acyclic undirected graph, which can be leveraged to build a tree structureT(A) = (V, E) with nodes V ⊂ {1, • • • , N } and edges E = {(i, j) : a ij ̸ = 0, i ̸ = j ∈ V}.Under the constraints Eq. ( 16), we seek to establish the following continuous optimization problemEQUATIONwhere ℓ(θ, ϕ, X, A) = -L T(A) (θ, ϕ, X) is the negative of the evidence lower bound given by Eq. ( 13), and the correlation parameters Γ T(A) = {γ ij : a ij ̸ = 0, i ̸ = j ∈ V} of the learned tree is determined by the binary adjacency matrix A with γ ij = f ϕ (x i , x j ) calculated by the neural network. A similar optimization can be implemented to stochastically learn a meaningful mixture of trees by using a set of symmetric adjacency matrices A = {A 1 , • • • , A M } to represent multiple tree components, and further maximizing the evidence lower bound Eq. ( 15) under acyclic constraints Eq. ( 16) of each matrix in A. The constrained optimization problem above can be further converted into unconstrained subproblems with Lagrangian multiplies and efficiently solved by numerical algorithms or stochastic gradient methods. For implementation details, we refer to Appendix D.3.To initialize the tree structure before constrained optimization, the easiest way is to randomly build from the fully connected graph over data instances by using depth-first search (DFS) algorithm. To enrich the initialized tree structure with neighboring correlation information, the uniform sampling process in the DFS algorithm can be further modified to generate a meaningful neighborhood for each instance, by assigning the probability of sampling a neighbor of each instance according to their similarity. For detailed implementation of the tree initialization, we refer to the Appendix D.2. 
related work: 
 Variational inference is broadly used for approximating intractable posterior in latent variable models, but notorious for its restricted variational families, especially mean-field variational family which is still widely used by modern generative models [4, 37] . So far there has been a wide variety of variational inference methods that attempt to improve on traditional mean-field variational inference by modeling dependence structures within latent posterior distribution, such as constructing the variational distribution with a normalizing flow [30, 6, 39] , which deterministically transforms a simple probability distribution over latent variable to a complex one through a sequence of invertible and differentiable functions with tractable Jacobians. Besides deterministic transformations, the form of structured variational inference can be diverse, such as constructing with implicit models [12, 35, 26] , modeling dependencies between local and global parameters [11, 38] , constructing with a mixture of variational distributions [16, 27, 20] , determining pairwise dependencies between univariate marginals with copula functions [34, 13, 33] , and designing variational distribution with hierarchies [38, 1, 25] . However, these work are mostly focused on distributional assumption over latent dimensions, that fail to be directly extended to capture instance-level correlation structure.Recently, there has been some work on incorporating instance-level correlation structure in variational inference [24, 36, 29] . Manduchi et al. [24] designs a prior information matrix to express must-link and cannot-link constraints between data in an explicit way, and integrates the instance-level prior information into the framework of variational inference by conditioning on the prior clustering preferences. But the instance-level correlation is only considered in the generative process regardless of the correlation structure induced from latent posterior by still adopting an amortized mean-field variational distribution. The work of Tang et al. [36] is the most related to ours, which takes instancelevel correlation structure into consideration when learning latent representations, but restricts both the prior and posterior distributions to be identically tree-structured for tractable optimization. However, this assumption is unrealistic in most variational inference scenarios, and high-order correlations are unable to be modeled within the tree structure. 
experiments: 
 Tasks & Datasets. We evaluate our methods with four different tasks: synthetic evidence lower bound test, constrained clustering, user matching and link prediction, on synthetic dataset, standard datasets (MNIST, Fashion MNIST, Reuters and STL-10), public movie rating and product rating datasets, respectively. We refer to Appendix E.1 for more dataset details.Baselines & Implementation Details. For baselines on the user matching and link prediction tasks, we include the standard variational auto-encoders [18] and a recent modification to VAE [36] taking instance-level correlation structure into consideration. And we also compare our method to the state-of-the-art method learning latent embeddings with graph convolutional networks, GraphSAGE [8] . With regard to the constrained clustering task, we take a variety of constrained clustering methods, e.g., the traditional pairwise constrained K-means (PCKmeans, [2] ), deterministic deep constrained clustering method based on DEC (SDEC, [32] ) and constrained IDEC (C-IDEC, [44] ), as stateof-the-art constrained clustering methods for comparison. We also include the generative models, the unsupervised VaDE [14] , the graph augmented VaDE (DGG, [42] ), and the weakly-supervised DC-GMM [24] . For comparison, we experiment on our methods TreeVI and MTreeVI, where the number of tree components for the mixture-of-trees posterior is fixed as M = 3. Aadditional information related to experimental implementation details are in Appendix E.2. We design a synthetic dataset with a graph-structured latent variable model. The dataset contains N = 6000 data pointsx 1 , • • • , x N ∈ R D with D = 4, each independently gener- ated from the conditional distribution p(x|z) = N (x; z, σ 2 I 4 ) given latent embeddings z 1 , • • • , z N ∈ R Dwhere σ 2 is a fixed value and set to 0.5. The latent embeddings z 1 , • • • , z N are drawn from a zero-mean normal distribution p(z) = N (z; 0 4 , Σ z ), and the graph-structured correlation is incorporated into the latent covariance matrix Σ z = I 4 + λA via an affinity matrix A ∈ R 4×4 assigned with a loopy graph structure where λ is leveraged to control the overall correlation strength and set to 0.5. Two tree structures over latent dimensions T 1 = (V, E 1 ) and T 2 = (V, E 2 ) are designed with vertex set V = {1, 2, 3, 4} and edge sets (1, 4) , (2, 3)}, respectively. The posterior correlation structure is modeled by T 1 , T 2 and their mixture model MTreeVI(T 1 ,T 2 ), with the estimated evidence lower bounds and ground truth log-likelihood shown in Table 1 . It can be seen that tree-structured posteriors can learn more correlation information than traditional mean-field approximations. Moreover, different choices of the tree structures influence the amount of correlation information, and mixture of tree components can obtain more correlations than each individual.EQUATIONE 1 = {(1, 2), (1, 3), (2, 4)} and E 2 = {(1, 2),Constrained clustering tasks differ from the classic clustering scenario with access to instance-level constraints, consisting of must-links if two samples are believed to belong to the same cluster, and cannot-links, otherwise. Based on the variational deep embedding (VaDE) framework [14] , constrained clustering can be formulated as a probabilistic clustering problem with joint probability p θ (X, Z, c) = p θ (X|Z)p(Z|c)p(c), where the sample x i is generated from a normal distribution conditioned on z i , z i is sampled from p(z i |c i ) = N (z i ; µ ci , diag(σ 2 ci )), and the cluster assignments c = {c i } N i=1 are sampled from a categorical distribution. Following previous work [24] , we incorporate the clustering preference through a conditional probability p(c|W) with a pairwise prior information matrix WEQUATIONwhere π = {π k } K k=1 are the weights associated to each cluster, Ω(π) is the normalization factor and h i (c, W) = j̸ =i exp(W ij δ cicj ) is a weighting function. The pairwise prior information matrix W is defined as a symmetric matrix containing the pairwise constraints: W ij > 0 if there is a must-link constraint between x i and x j , W ij < 0 if there is a cannot-link constraint between x i and x j , andW ij = 0 otherwise. The values |W ij | ∈ [0, ∞)reflect the degree of certainty in the constraints, and are set to 10 4 for all datasets. And 6000 pairwise constraints are used for experiments on both our methods and other constrained clustering baselines.The variational posterior distribution is defined as q ϕ (Z, c|X) = q ϕ (Z|X)q(c|Z), where the probability of cluster assignments is factorized as q(c|Z) = i q(c i |z i ) which can be easily computed by Bayes theorem. In the work of DC-GMM [24] , the posterior distribution q ϕ (Z|X) for latent variables is assumed to be mean-field which fails to capture the posterior correlation structure, while in our methods TreeVI and MTreeVI, the latent posterior is approximated by tree-structured and mixture-of-trees distribution, respectively. In Table 2 we report the averaged clustering performances across 10 turns of both our proposed methods against baseline methods. Accuracy (ACC), Normalized Mutual Information (NMI), and Adjusted Rand Index (ARI) are used as evaluation metrics. It can be easily observed that our models reach the state-of-the-art clustering performance in all metrics and datasets, benefiting from the correlation structures captured by our designed correlated posteriors. Moreover, due to high efficiency of tree-structured reparameterization, our methods have computational complexity comparable to DC-GMM that adopts the fully factorized posterior distribution. For example, the training of DC-GMM on MNIST dataset each epoch takes 2 to 3 seconds on GeForce RTX 3090, while our TreeVI and MTreeVI take around 5 and 9 seconds, respectively. Further, we plot the tree learned over the constrained clustering experiment on MNIST dataset, as shown in Figure 2 . 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.We evaluate our methods against the baselines on a public movie rating dataset MovieLens 20M. The ratings of each user u i are binarized as a bag-of-word vector x ui (i = 1, 2, • • • , N ), and only ratings for movies that have been rated over 1000 times are preserved for simplicity. In our experiments, the watch history of each user u i is randomly split into two halves, leading to two synthetic users u A i and u B i with the most similar movie preference, and a correlation graph G = (V, E) with nodesV = {u A i , u B i : i = 1, • • • , N } and edges E = {(u A i , u B i ) : i = 1, • • • , N }.Suppose that the joint distribution of rating data X and the corresponding latent embeddings Z is modeled by p θ (X, Z) = p θ (X|Z)p(Z), and the variational posterior distribution is modeled by q ϕ (Z|X). In the work of CVAE [36] , both the prior distribution p(Z) and posterior distribution q ϕ (Z|X) are designed as weighted sums of tree-structured distributions with respect to each maximal acyclic subgraph of the correlation graph G. While in our work, the latent representations are drawn from a Gaussian distribution of the form p(Z) = N (Z; 0 2N ; (I 2N + λA) ⊗ I D ), where the hyper-parameter λ ∈ (0, 1) is used to control the overall correlation strength, and the affinity matrix A = [a ij ] i,j∈V is derived from the correlation graph with a ij = 1 if (i, j) ∈ E and a ij = 0 otherwise. And our latent posterior distribution q ϕ (Z|X) is modeled by tree-structured and mixture-of-trees distributions, respectively. And our goal is to identify the dual user u B i given a synthetic user u A i from a held-out set in terms of the latent embedding distance. The user data are implemented a random train/test split with a 90/10 ratio, and the synthetic user pairs from the training set are used to train all the methods. To evaluate the user matching accuracies, a fixed number N eval = 1000 of synthetic user pairs are selected from the test set and for each synthetic user u A i (or u B i ), we summarize the ranking of its dual user u B i (or u A i ) among all other 2N eval -1 synthetic user candidates in terms of latent embedding distances. 0.0188 ± 0.0007 MTreeVI (Ours) 0.0203 ± 0.0014In Table 3 we show the average Reciprocal Ranking (RR) for all the methods, which demonstrates superiority of our model over both implementations for the baseline method CVAE.We perform link prediction task on a constructed undirected correlation graph G = (V, E) within the public product rating dataset Epinions. The rating data are binarized into bag-of-words feature vectors x ui for each user u i , and only products that have been rated at least 100 times are kept and users who have rated these products at least once are considered. To construct the undirected graph G from the single-directional "trust" statements between users u i , u j ∈ V provided by the dataset, we only build an edge (u i , u j ) ∈ E if both u i trusts u j and u j trusts u i . The experiment setting of our TreeVI and MTreeVI are similar to the user matching task, by defining a correlation graph incorporated prior distribution and approximating latent posterior distribution with tree-structured and mixture-of-trees distributions, respectively.The product rating dataset is split for each user u i ∈ V into training and test sets, with max 1, 1 20 • degree(u i ) edges held out for test edge set E test . The remaining edges for the training edge set E train are used to train all methods on the product rating data. To evaluate the link prediction accuracies, we calculate for each user u i the Normalized Cumulative Reciprocal Rank NCRR i of the ratings of u i 's test edges among all possible connections except for the training edges, in terms of latent embedding distance metrics. Formally, the NCRR value is the [0, 1]-normalization of the Cumulative Reciprocal Rank (CRR) formulated as CRR i = (ui,uj )∈Etest |{k :(u i , u k ) / ∈ E train , d ik ≤ d ij }| -1, where d ij represents the latent embedding distance between users u i and u j for 1 ≤ i ̸ = j ≤ N . Larger NCRR i indicates better ability to predict held-out test links with respect to each user u i , and the averaged results for all methods are reported in Table 4 .Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: The datasets, baseline methods and implementation details needed to reproduce the experimental results are included in the Section 4, and we refer to Appendix E for more experimental details. 
conclusion: 
 In this work, we present a novel variational inference method called TreeVI, that approximates the intractable latent posterior distribution with a tree-structured distribution. This induces a Bayesian network whose ancestral sampling gives a matrix-form reparameterization for the correlated latents and enables efficient optimization. To enrich the correlation structure, TreeVI is further extended to MTreeVI with a mixture of trees to better approximate the underlying posterior. Furthermore, the tree and mixture-of-trees structures are allowed stochastically learned from data by solving a constrained optimization problem under our proposed acyclicity constraints. With correlated posteriors, we show that our methods can capture more correlation information and achieve superior performances in real-world tasks.Limitations & Future Work The proposed method requires a tree structure to approximate the graph-structured posterior correlation structure. This limitation is mitigated by adopting a weighted mixture of trees and stochastically learning a correlation-rich tree or mixture-of-trees structure with our proposed constrained optimization. For future work, we will further investigate on correlation structures with higher expressivity for approximating the latent posterior.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [No] Justification: The datasets and experimental settings are provided in Section 4 and Appendix E, and our code will be submitted to the GitHub repository soon.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: Training details of all our experiments have been specified in the Appendix E.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: The statistical significances of our experiments that support the main claims of the paper are shown in Table 1 , Table 2 , Table 3 and Table 4 , respectively. And the related metrics have been depicted in the context.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: The computer resources needed to reproduce our experiments have been specified in Appendix E.4.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Our work mainly focuses on basic theory about combination of variational inference and instance-level correlation structure, meaning that there is no societal impact to be addressed.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our experiments are conducted on standard and public datasets available for everyone, and our proposed methods focus on basic theory with regards to variational inference without safety risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: Datasets used in our experiments are all public, and their related papers have been cited in our paper.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: When latent variables Z = [z 1 , • • • , z N ] ⊤ areall Gaussian, it is a well-known consequence of the Hammersley-Clifford theorem that the entries of the precision matrix P z = Σ -1 z correspond to rescaled conditional correlations. Denote P z = [diag(p ij )] i,j=1,••• ,N with p ij ∈ R D , then the magnitude of p ij expresses the correlation of z i and z j conditioned on other latent variables. In particular, the sparsity pattern of p ij reflects the edge structure of the correlation graph:p ij = 0 D if and only if z i ⊥ ⊥ z j |z {1,••• ,N }\{i,j} .The tree structure T = (V, E) with non-adjacent correlation parameters defined as Eq. ( 3) forms a Markov random field with joint probability given by Eq. ( 4), which follows the principle of conditional independence. Therefore, the precision vector p ij ̸ = 0 D if and only if latent variables z i and z j are adjacent, that is, (i, j) ∈ E. Example. In the example of tree-structured correlation approximation with 5 latent variables as shown in Fig. 1b , the precision matrix corresponding to the tree-structured latent covariance matrix can be expressed asEQUATIONwith the inverse of tree-structured correlation matrix given by [R (T) ] -1 =     demonstrating that only correlation parameters along the edge setEQUATION} of tree structure T has been essentially captured in the underlying correlation structure, where all of the matrix elements are simply combinations of correlation parameters from Γ T = {γ 12 , γ 23 , γ 24 , γ 35 }:EQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONB Ancestral Sampling for Tree-structured PosteriorB.1 Proof of Theorem 1 Suppose that N latent variables Z = [z 1 , • • • , z N ]⊤ follow a tree-structured posterior distribution given the tree structure T = (V, E) with V = {1, • • • , N }, the ancestral sampling for each latent variable z j , j ∈ V can be expressed byEQUATIONwhere ϵ i ∼ N (0 D , I D ) is a randomly sampled noise from the standard normal distribution and pa(i) denotes the parent node of z i with respect to i ∈ V.Proof. For each latent variable z j , there exist a unique pathP 1→j = (v 1 , v 2 , • • • , v n )along the tree-structured Bayesian network that starts from the root node v 1 = 1 and ends at the goal node v n = j. The joint distribution along the path P 1→j is given byEQUATIONThe ancestral sampling starts by sampling from q ϕ (z v1 ): z v1 = µ v1 + σ v1 ⊙ ϵ v1 , where ϵ v1 ∼ N (0 D , I D ). Given the observation of z v1 , then the latent variable z v2 can be sampled using the conditional distribution Eq. ( 7):EQUATIONEQUATION)Assume that at the k-th step of the ancestral sampling, we haveEQUATIONwhereγviv k = γ vivi+1 ⊙ • • • ⊙ γ v k-1 v k for i < k and γv k v k = 1 D .Then at the (k + 1)-th step of the ancestral sampling, the latent variable z v k+1 can be similarly sampled by conditional samplingEQUATIONEQUATIONTherefore by induction, the ancestral sampling for the latent variable z j can be expressed asEQUATIONwhich is consistent with the expression of Eq. (31).For the quin-variate example in Fig. 1b , we assume z 1 as the root node, and the detailed ancestral sampling procedure for the tree-structured Bayesian network Eq. ( 6) is shown as follows.1. Start by sampling from q ϕ (z 1 ):EQUATIONwhere ϵ 1 ∼ N (0 D , I D );2. Then sample from q ϕ (z 2 |z 1 ) given the observation of z 1 :EQUATIONwhere ϵ 2 ∼ N (0 D , I D );3. Then sample from q ϕ (z 3 |z 2 ) given the observation of z 2 :EQUATION) where ϵ 3 ∼ N (0 D , I D ); 4. Then sample from q ϕ (z 4 |z 2 ) given the observation of z 2 :EQUATIONwhere ϵ 4 ∼ N (0 D , I D );5. Finally, sample from q ϕ (z 5 |z 3 ) given the observation of z 3 :EQUATIONwhere ϵ 5 ∼ N (0 D , I D );It is easily noted that, the expression of each latent sample in this ancestral sampling procedure is consistent to Eq. ( 8).For a tree-structured posterior approximation with respect to the tree structure T = (V, E), the evidence lower bound of our proposed TreeVI is given byEQUATIONwhere Z (T) denotes the tree-structured reparameterization for latent variables. The first term can be directly computed byEQUATIONwhere z (T) iis the reparameterization for latent variable z i , i = 1, • • • , N . And the entropy of the treestructured posterior q T ϕ (Z|X) can be factorized as entropy terms with respect to singleton posterior q ϕ (z i |x i ) and pairwise posterior q ϕ (z i , z j |x i , x j )EQUATION)where the entropy of singleton posterior q ϕ (zEQUATIONwhere Σ i = diag(σ 2 i ), and the step ⋆ relies on several properties of the trace operator:EQUATION)And the entropy of pairwise posterior q ϕ (z i ,EQUATIONis similarly given byEQUATIONEQUATION)where we denotez ij = [z i , z j ] ⊤ .For a mixture-of-trees structured posterior approximation with respect to the mixture of tree components MT = {T 1 , • • • , T M }, the evidence lower bound of our proposed MTreeVI is given byEQUATIONwhere Z (Tm) denotes the tree-structured reparameterization for latent variables with respect to the m-th tree component T m , m = 1, • • • , M . The first weighted summations can be directly computed byEQUATIONwhere z (Tm) iis the reparameterization for latent variable z i , i = 1, • • • , N with respect to the m-th tree component, m = 1, • • • , M . And the weighted summations of the entropy terms can be similarly factorized and calculated as Appendix C.1 by entropy terms with respect to singleton posteriors q ϕ (z i |x i ) and pairwise posteriors q ϕ (z i , z j |x i , x j ).Recall that the spectral radius r(•) is the largest absolute eigenvalue of a matrix. The following show characterizations for acyclicity of the undirected structure induced by the symmetric matrix A:1. Suppose that A ∈ {0, 1} N ×N and r(A) < 1, then A is acyclic if and only ifEQUATION)2. Suppose that A ∈ {0, 1} N ×N , then A is acyclic if and only ifEQUATION)3. Suppose that A ∈ R N ×N , then A is acyclic if and only ifEQUATION)where tr(•) and exp(•) represent the trace and exponential of a matrix respectively, and ⊙ is the Hadamard product.Proof. The proof relies on the fact that tr A k counts the number of length-k closed walks in a directed graph.1. Clearly the directed graph induced by the symmetric matrix A will only have self-loops, and hence tr A 2k+1 = 0 for all k = 1, • • • , ∞. In other words, A has no cycles if and only ifEQUATION)The desired result follows.2. Similar to the proof of Proposition 1.1 by noting that A has no cycles if and only if (A 2k+1 ) ii = 0 for all k ≥ 1 and all i ∈ {1, • • • , N }, which is true if and only ifEQUATION)3. The proof is similar to Proposition 1.2 by replacing A with A ⊙ A, which counts weighted closed walks.The constrained optimization of tree structure requires an initialization by construction from data, and the easiest way of tree construction is to randomly build from the fully-connected graph by using depth-first-search (DFS) algorithm. Algorithm 1 shows the DFS algorithm for our tree initialization.In the algorithm, RC for k = 0, • • • , M -1 do 4: V = [F alse] |V| ▷ Visited node list 5:while F alse in V do6: i ← RC [V ==F alse] ▷ Choose node 7: Q = [i] ▷ Initial queue 8:while len(Q) > 0 do9: i ← Q[0]10:V [i] ← T rue 11: N = ID [V [N (i)==F alse]]12:if len(N ) == 0 then 13:P OP (Q, -1) end for 23: end procedure To enrich the constructed spanning tree with neighboring correlation information, the uniform sampling process (line 16 in Algorithm 1) in the DFS algorithm can be further modified to generate a meaningful neighborhood for each data instance, by assigning the probability of sampling neighbor j of instance i as exp(cos(x ⊤ j x i )/α)EQUATION)where α is the temperature parameter controlling the trade-off between the precision and diversity of edges, and we find the best configuration of α on the validation set with the values in{0.1, 0.2, • • • , 1.0}.In Table 5 , we show the constrained clustering accuracies on the MNIST dataset, with different initializations, and 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. In section 2.3, we establish a continuous characterization of acyclicity, leading to the following equality-constrained program (ECP)EQUATIONSolving the constrained optimization requires classical techniques from the mathematical optimization literature. However, this is a nonconvex program since {A : h(A) = 0} is a non-convex constraint, and hence we are aiming to find its stationary points with non-convex optimization. The algorithm for solving Eq. (68) consists of three steps: (i) converting the constrained problem into a series of unconstrained subproblems, (ii) optimizing the unconstrained subproblems, and (iii) thresholding. The full algorithm is outlined in Algorithm 2.Input: Initial guess (A 0 , α 0 ); progress rate c ∈ (0, 1); tolerance ϵ > 0; threshold ω > 0 Output:Threshold matrix A 1: for t = 0, 1, 2, • • • do 2: Solve A t+1 ← arg min A L ρ (A, α t ) with ρ such that h(A t+1 ) < ch(A t ) 3: α t+1 ← α t + ρh(A t+1 ) ▷ Dual ascent 4:if h(A t+1 ) < ϵ thenA ← A t+1end if 8: end for9: A ← A ⊙ 1(| A| > ω)Converting constrained optimization into unconstrained subproblems. The augmented Lagrangian method [28] can be leveraged to solve the equality-constrained program Eq. (68) by first augmenting the original problem with a quadratic penalty:EQUATIONwith a penalty parameter ρ > 0, which approximates well the solution of the original constrained problem by the solution of unconstrained problems without increasing ρ to infinity. Then the algorithm implements dual ascent for Eq. ( 69) by defining a dual function with Lagrange multiplier α D(α) = minEQUATIONwhereEQUATIONis the augmented Lagrangian. And the goal is to find a local solution to the dual problemEQUATIONLet A * α be the local minimizer of the Lagrangian (70) at α, i.e., D(α) = L ρ (A * α , α). Since the dual objective D(α) is linear in α with the derivative simply given by ∇D(α) = h(A * α ), one can perform dual gradient ascent to optimize the dual problem (72):EQUATIONSolving the unconstrained subproblem. The augmented Lagrangian converts the constrained problem (69) into a series of subproblems (70), and our goal is to solve these subproblems efficiently. Let a = vec(A) ∈ R p , with p = N 2 . The unconstrained subproblem (70) can be considered as a typical minimization problem over real vectors:EQUATIONwhereEQUATIONis a smooth objective, for which a number of efficient numerical algorithms are available, such as L-BFGS [5] . In our experiments, the values of a can be learned by a neural network that encodes pairwise correlation.Thresholding. Motivated by post-processing estimates of coefficients via hard thresholding, we threshold the edge weights as follows: after obtaining a stationary point A of (69) given a fixed threshold ω > 0, set any weights smaller than ω in absolute value to zero. This strategy also has the effect of "rounding" the numerical solution of the augmented Lagrangian (69), since the solution satisfies h( A) ≤ ϵ for some all tolerance ϵ (set to ϵ = 10 -8 in our experiments) instead of h( A) = 0 strictly due to numerical precisions. However, a small threshold ω suffices to rule out cycle-inducing edges since h( A) explicity quantifies the acyclicity of A. Following [45] , a fixed value of threshold ω = 0.3 is set in all our experiments.The datasets used in the experiments are the followings:• MovieLens 20M: A dataset describing ratings and free-text tagging activities from Movie-Lens, a movie recommendation service. It contains 20,000,263 ratings and 465,564 tag applications across 27,278 movies created by 138,493 users [9] .• Epinions: A dataset that records ratings and trust statements issued by users from Epinions, a consumers opinion site where users can review items and assign them numeric ratings in the range 1 to 5, and also express their Web of Trust by issuing trust statements. It consists of 49,290 users who rated a total of 139,738 different items at least once. The total number of review is 664,824. The total number of issued trust statements is 487,181.• MNIST: It consists of 70,000 handwritten digits. The images are centered and of size 28 by 28 pixels, each reshaped to a 784-dimensional vector [21] .• Fashion MNIST: A dataset of Zalando's article images consisting of a training set of 60,000 examples and a test set of 10,000 examples [40] .• Reuters: It contains 810,000 English news stories [22] . Following the work of [41] , we used 4 root categories: corporate/industrial, government/social, markets, and economics as labels and discarded all documents with multiple labels, which results in a 685,071-article dataset. We computed tf-idf features on the 2000 most frequent words to represent all articles. A random subset of 10,000 documents is then sampled.• STL-10: It contains color images of 96-by-96 pixel size. There are 10 classes with 13,000 examples each [7] . As pre-processing, we extracted features from the STL-10 image dataset using a ResNet-50 [10] , as in previous works [15] .Synthetic Data. For synthetic VAEs, we apply a two-layer feed-forward neural network for the generative model p θ (x i |z i ) and a two-layer feed-forward neural network for the variational posterior approximation q ϕ (z i |x i ), with each z i ∈ R D where D = 4. The traditional meanfield approximation assumes the posterior distribution for each latent to be fully factorized: q ϕ (z i |x i ) = D d=1 q ϕ (z id |x id ), while in our methods the posterior is assumed to be tree-structured and mixture-of-trees structured, which can be factorized into singleton posterior q ϕ (z id |x id ) and pairwise posterior q ϕ (z id , z ie |x id , x ie ) both assumed to be Gaussian. Each edge in the tree structure corresponds to a correlation parameter γ de to be learned in the bi-variate normal distribution, which EQUATIONAnd we adopt different number of tree components to model the posterior correlation, the comparison of their estimated evidence lower bounds to the ground truth log-likelihood is presented in Table 7 . It can be observed that tree-structured poterior benefits from capturing more pairwise correlations, and the mixture of more tree components better approximates the ground truth log-likelihood.Clustering with Constructed Affinity Matrix. Besides the weak supervision of pairwise constraints, we also analyze the unsupervised scenario of constructed affinity matrix in this work, where we construct an affinity matrix under a specific similarity measure and each entry is a value between [0, 1] reflecting the similarity between a sample pair. To construct the affinity matrix W, we find a set of nearest neighbors for a given data point and compute their similarity using a predefined kernel function, such as Gaussian kernelEQUATIONwhere s i is a predefined scalar, N (x i ) denotes the set consisting of the nearest N s neighbors of x i .To enhance robustness to different datasets, we train a Siamese network to measure the similarity between data points. With the constructed affinity matrix W, the conditional distribution of the cluster assignment can be similarly defined as Eq. ( 19), with the weighting function modified asEQUATIONThe clustering performances of our proposed methods against baseline methods are shown in Table 8 .Experiments were conducted on an internal computing cluster. Each experiment configuration used one NVIDIA GPU (either a 2080TI or 3090TI), 16 CPUs and a total of 24GB of memory.Justification: Assumptions for the theoretical results are given in the descriptions of theorem or proposition, and the proofs are included in the Appendix 31 and D.1, respectively.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.Justification: We have reviewed the NeurIPS Code of Ethics, and make sure that our research follows the Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 


Decision: The paper presents a method to perform amortized VI with a distribution family in between mean-field and fully connected Gaussians. Instead, the authors proposed to use a tree-structured Gaussian graphical model over data instances. Reviewers found the paper to be interesting, novel, and sound and with experiments that supported the main claims. There was considerable back-and-forth during the rebuttal and author-reviewer discussion phase. Much of it surrounding technical clarification of the methods, with the result that two reviewers raised their ratings and one lowered their rating. In the end, all reviewers appreciated the contributions and unanimously recommended accept. The most significant issue that can be addressed to improve the paper is to resolve potential ambiguity surrounding the problem setting and motivation. For example, the authors are encouraged to address the post-rebuttal comments of Reviewer ptS8 about highlighting more clearly the main motivation of using a tree structure, as well as the post-rebuttal comments by Reviewer BYBL about the discrepancy of using an instance-level tree in one of the experiments.

Reviews: review: paper_summary: paper_summary: The paper proposes to perform amortized variational inference, where the variational approximation has a tree-dependence structure across the instance-level latent variables. For this, the scale is decomposed into variance and correlation, where both are amortized through a neural network. A non-convex constrained optimization procedure is proposed to also learn the correlation structure, which ensures that the correlation structure forms an acyclic graph.
strength_weakness: strength_weakness: * The proposed methodology is unclear at key places, making it hard to grasp exactly what is going on. For instance, it isn't clear to me what the constrained optimization problem in Eq 17 is supposed to do. Isn't an acyclic graph that spans *all* latent variables a solution to this problem? In that case, the resulting variational approximation is full-rank, which seems to defeat the whole point. Therefore, how exactly the authors obtain a sparse correlation structure needs to be clarified.
* A lot of the derivation in Section 2.1 is redundant. For instance, it isn't clear why the Cholesky factor has to be decomposed into variance and correlations involving a bunch of square roots. That is, the authors could have simply parameterized the $(i,j)$th element of the Cholesky factor as $L_{ij} = f_{\phi}\left(x_i, x_j\right)$. This would cut down the text of Section 2.1 by more than half since it is now obvious how to reparameterize. In fact, it is well known that manipulating the elements of the Cholesky factor is equivalent to manipulating the graph structure. (See Proposition 3 in [1] for example. I also suspect there should be a textbook result somewhere, but couldn't find such result.) This would also avoid all the square roots.
* Given the comment above, it is unclear *why* we have to take a tree structure among various different correlation graph structures.
* The paper is imprecise about the fact that it is solving a problem specific to *amortized variational inference* not just variational inference. In fact, a full-rank/dense scale matrix is strictly more expressive than any sparser graph approximation. Therefore, the proposed method does not make sense outside of the scalable amortized VI context. In that sense, the title and abstract are not specific enough about the setting and should be changed.
    * Line 31-36 "... struggling to alleviate the independence structure adopted by standard variational inference ...": We knew how to alleviate independence structure since the work of [2,3]. The statement is only true for instance level correlation in amortized VI.
    * Line 41: In "the study of variational inference with instance-level correlation remains under-explored," "the study of variational inference" should have been "the study of scalable amortized variational inference"

Overall, the crucial contribution of the paper seems to be about automatically learning the correlation structure, but it is unclear how this is done and how sparsity in the correlation structure is enforced. Furthermore, the motivation for tree structure (out of all structures) is unclear.

## Additional Comments
* Section 2.2: The additional Jensen gap can be avoided by using the sticking-the-landing gradient [13].
* While it is true that Gaussian mean-field families commonly used in deep latent variable models are restrictive, the recent development of MCMC-based augmentation methods [5,6,7,8] is able to overcome this problem.
* Line 55: Cite something for the ELBO. It is typical to cite [9,10].
* Line 28: The mean-field approximation was first proposed by [11] and popularized by [12].
* Line 38 "... share the same tree structure which is unrealistic in real applications": All variational inference is performing approximations. There is nothing realistic or realistic about any of the approximations. There are *accurate* or *less accurate* approximations.

## References
(Disclaimer: I am not the author of any of the papers below.)
1. Katzfuss, Matthias, and Joseph Guinness. "A general framework for Vecchia approximations of Gaussian processes." (2021): 124-141.
2. Hoffman, Matthew, and David Blei. "Stochastic structured variational inference." Artificial Intelligence and Statistics. PMLR, 2015.
3. Titsias, Michalis, and Miguel Lázaro-Gredilla. "Doubly stochastic variational Bayes for non-conjugate inference." International conference on machine learning. PMLR, 2014.
4. Kucukelbir, Alp, et al. "Automatic differentiation variational inference." Journal of machine learning research 18.14 (2017): 1-45.
5. Caterini, Anthony L., Arnaud Doucet, and Dino Sejdinovic. "Hamiltonian variational auto-encoder." Advances in Neural Information Processing Systems 31 (2018).
6. Geffner, Tomas, and Justin Domke. "MCMC variational inference via uncorrected Hamiltonian annealing." Advances in Neural Information Processing Systems 34 (2021): 639-651.
7. Doucet, Arnaud, et al. "Score-based diffusion meets annealed importance sampling." Advances in Neural Information Processing Systems 35 (2022): 21482-21494.
8. Thin, Achille, et al. "Monte Carlo variational auto-encoders." International Conference on Machine Learning. PMLR, 2021.
9. Jordan, Michael I., et al. "An introduction to variational methods for graphical models." Machine learning 37 (1999): 183-233.
10. Blei, David M., Alp Kucukelbir, and Jon D. McAuliffe. "Variational inference: A review for statisticians." Journal of the American statistical Association 112.518 (2017): 859-877.
11. Peterson, Carsten, and Eric Hartman. "Explorations of the mean field theory learning algorithm." Neural Networks 2.6 (1989): 475-494.
12. Hinton, Geoffrey E., and Drew Van Camp. "Keeping the neural networks simple by minimizing the description length of the weights." Proceedings of the sixth annual conference on Computational learning theory. 1993.
13. Roeder, Geoffrey, Yuhuai Wu, and David K. Duvenaud. "Sticking the landing: Simple, lower-variance gradient estimators for variational inference." Advances in Neural Information Processing Systems 30 (2017).
questions: questions: n/a
limitations: limitations: n/a
paper_summary: paper_summary: This paper proposes Tree-structured Variational Inference (TreeVI), a novel method for capturing instance-level correlations in the posterior distribution of latent variables. TreeVI represents correlations between latent variables using a tree structure. This enables reparameterization of latent variables using ancestral sampling, allowing for efficient training. Furthermore, the authors propose Multiple Tree-structured Variational Inference (MTreeVI) to represent more complex correlation structures using multiple tree structures. They also propose a method to learn the tree structure from data. The proposed methods demonstrate superior performance compared to existing approaches on various tasks using both synthetic and real-world datasets.
strength_weakness: strength_weakness: In the experiments on real-world data, there is insufficient explanation of what the tree structure represents. The tree structure is used merely as a means to relax the independence assumption of the approximate posterior distribution while maintaining training efficiency. The study could have been enhanced by providing some insights derived from the tree structures obtained through training.
questions: questions: On line 78, $\mathcal{V} = \{ 1, 2, \cdots , N \}$ is defined. However, in the Synthetic VAE experiment in Section 4.1, despite $N=6000$, $\mathcal{V}= \{ 1,2,3,4 \}$ is used. Could you explain this discrepancy?
limitations: limitations: The limitations are appropriately addressed within the paper, and I have no additional suggestions in this regard.
paper_summary: paper_summary: The authors propose to use a variational approximation that captures tree-structured correlations as a middle ground between scalable fully-factorized approximations and more expressive approximations which aim to capture the full correlation structure. They show how these tree structured correlations give rise to variational distributions that factorize into tree-structured Markov random fields, whose ancestral sampling procedure informs a matrix-form reparameterization. The authors further extend their methods to mixtures of tree-structured variational distributions and show how to learn tree-structure by jointly optimizing over possible tree-constraint adjacency matrices.
strength_weakness: strength_weakness: It is unclear how important the initialization is and if there always exists a straightforward initialization procedure (see related questions). It would be nice to see some result on how the methods perform without initialization.
questions: questions: The authors report that their method is only about 2-3 times more costly than mean-field VI, however, the maximum number of edges (i.e. additional correlations that have to be learned) 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. Are these instances just not learned in practice or do the authors explicitly regularize to encourage a degree of sparsity in the learned trees? Or is there another reason why this cost is not seen in practice?

- It seems that in order to learn a rich correlation structure initialization is quite important. How well does the proposed method perform without instance specific initialization?

- The initialization itself requires to compute pairwise *neighborhood probabilities* on the data, however, depending on the task it might be more or less clear how to compute meaningful neighborhood probabilities. Is the assumption that we know how to compute distances (or some form of meaningful similarity) between data points explicitly?

- Something small that I noticed. The acronyms ARI (adjusted rand index) and NMI (normalized mutual information) have not been introduced and might not be familiar to readers not familiar with clustering.

--- Edit ---

The authors addressed my concerns regarding the computational complexity (the graph structure is a tree not a DAG) and the importance of the initialization. Consequently I am raising my score to an Accept (7).
limitations: limitations: -
paper_summary: paper_summary: The paper introduces a new type of variational inference that allows correlations across training samples in the form of a tree, making it suitable for applications with graph-structured data or explicit constraints. The method extends to multiple trees for capturing more complex correlations and includes a mechanism for automatically learning the underlying tree structures from data. Experimental results demonstrate that the proposed approach provide advantages in settings such as constrained clustering, user rating and link prediction.
strength_weakness: strength_weakness: 1. I find the setting of the paper very confusing, I worked on hierarchical VAEs and therefore had some difficulty in understanding the notation of the paper. Figure 1 did not help as notations are misleading: what do the numbers in the nodes represent? Why in GraphVAE the nodes are labeled z_i while in the trees they are labeled with just a number? I believe the paper would be stronger if (a) notations could be better explained and (b) if comparison with TreeVAE [1] could be explained to differentiate the two settings.
2. I also believe that the paper could benefit from a more thorough analysis of the computational complexity  of the proposed appraoch(there is only an empirical comparison with DC-GMM).
3. Experiments are limited to simple datasets, standard deviations are not reported in the table, which I believe are of crucial importance especially in clustering (given the high instability of the model with different initialisation), and few experimental settings are explained. E.g. At what value is M set in the experiments?
4. The construction of the trees, while interesting, have not been properly validated. For example, in 4.2 it would be more convincing to provide evidence that the learnt tree are consistent with the instance-level constrains (are similar pair of samples close in the tree?). 
5. I think the paper is not always clearly written and it contains several typos (see line 350 or 316, 317 as examples).
questions: questions: 1. Could the authors provide a proof of Equation 4 (in the appendix)?
2. 4.1 Why the tree structures are defined a-priori and not learnt from the data? That would be an interesting experiment to validate the correctness in constructing the tree in the synthetic scenario.
3. Line 317: how do you compute the latent embedding distance? Is it the distance in the tree or the e.g. L2 distance in the embedding space?
limitations: limitations: I believe the authors have addressed the limitations of the proposed approach.

Comments: title: Re: Rebuttal by Authors
Response
Additional Questions
comment: Thank you for the clarification. I raised my score to weak acceptance.
Thank you for providing multiple responses to my questions.  
I believe I now have an accurate understanding of the experiments in Section 4.1.

Regrettably, I have to sαy the current manuscript lacks sufficient explanation. Moreover, providing a detailed explanation of the formulation for capturing inter-dimensional correlations in an appendix or elsewhere may not fully solve this issue. Capturing inter-dimensional correlations is not the main focus of this paper, so doing so might be misleading. In other words, it seems that the purpose and design of the experiments in Section 4.1 were not sufficiently aligned with the main theme of this paper.

While I regret that the significant time you spent answering my questions has led to this outcome, I will have to lower the score by one point. I deeply appreciate your patience in responding to my inquiries.
Yes, your understanding is fully correct.
So, if we explicitly show the instance subscript $n$, can we say that equations (6), (7), and (12) can be written as follows?

$q_{\boldsymbol{\phi}}^\mathbb{T}(\boldsymbol{Z}|\boldsymbol{X}) = \prod_{n=1}^N \left( q_{\boldsymbol{\phi}}(z_{n,1}) \prod_{(i,j) \in \mathcal{E}} q_{\boldsymbol{\phi}} (z_{n,j} | z_{n,i}) \right)$

$q_{\boldsymbol{\phi}} (z_{n,j} | z_{n,i}) = \mathcal{N} \left(z_{n,j}; \mu_{n,j}+\gamma_{n,ij}\odot\sigma_{n,j}\odot\sigma_{n,i}^{-1}\odot(z_{n,i}-\mu_{n,i}), \sigma_{n,j} \odot \sqrt{1-\gamma_{n,ij}^2} \right)$

Are there any quantities in the right-hand side of the above equation that do not depend on $n$?

$\\{\gamma_{n,ij}\\}\_{(i,j) \in \mathcal{E}} = f\_{\boldsymbol{\phi}}(\boldsymbol{x}\_n)$

Is it correct to consider $f_{\boldsymbol{\phi}}(\cdot)$ as a vector-valued function of dimension $|\mathcal{E}|$?

Finally, in the traditional mean-field approximation in the experiments of Section 4.1, is it correct that independence both between instances and between dimensions are assumed?
Thanks a lot for saying that you are inclined to raise your score to Accept (7). This really means a lot to us, and we hope our response below can convince you to take the final action of raising the score .

Yes, we cannot guarantee the cosine similarity will always find similar instances, but as evidenced by its successes in a wide range of scenarios, cosine similarity is often not a bad choice when measuring the similarity of instances. 

Moreover, in addition to the greedy-search method, in our paper, we also proposed an optimization method to learn the tree from data, where the greedy search method is only used to provide an initialization for the optimization method. Our experimental results show that even if a random tree is used as the initialization, we can still get a good performance. Thus, with the proposed optimization-based tree learning method, the initialization does not play a very important role.

Thanks for spotting the missing parentheses for us in Eq. 67, we will correct it 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$.
Thank you for addressing my questions. After reading all of the reviews and corresponding rebuttals I'm inclined to raise my score to an Accept (7). 

Below are some comments and follow-up questions regarding the addressed points.

**Q2:** This makes perfect sense—my apologies for my blatant confusion of graph structures.

**Q3:** I had another look at Appendix D.2. and believe I understand how the procedure assigns the probability of sampling neighbors based on pairwise cosine similarities. However, while I understand that the cosine similarity can be a useful measure of similarity in certain settings, it is a priori unclear to me why we would expect it to be a good measure of similarity between arbitrary instance embeddings in Euclidean space. In any case, it seems that cosine-similarities are a reasonable first choice and that the procedure can easily be adapted to incorporate other measures of similarity. 

One more small thing. There is a small typo in Appendix D.2. The neighborhood operator in Equation 67 is missing a closing parentheses.
Thanks for your reply.

> If so, the index $i$ of $\boldsymbol{x}_i$ in Eq. (4) moves in $\{ 1, 2, \dots , D\}$. Is it correct? Could you write down the modified version of Eq. (4) here?

Since we are capturing correlations among dimensions within a latent embedding $\mathbf{z}$, we don't need to feed two instances $\mathbf{x} _ i$ and $\mathbf{x} _ j$ as input in Eq. (4), but only one instance $\mathbf{x}$ is needed instead. So the modified version of Eq. (4) should be
$$
q _ \phi^{\mathbb{T}} (\mathbf{z}|\mathbf{x}) = \prod _ {i \in \mathcal{V}} q_\phi (z _ i|\mathbf{x}) \prod _ {(i,j) \in \mathcal{E}} \frac{q _ \phi (z _ i,z _ j|\mathbf{x})}{q _ \phi (z _ i|\mathbf{x}) q _ \phi (z _ j|\mathbf{x})},
$$
where $\mathcal{V} = \\{1,\cdots,D\\}$ denotes the set of dimension indices and $\mathcal{E}$ denotes the set of edges that connect different dimensions. Here the index $i$ represents different dimensions of the latent embedding. These modified expressions and subtle details regarding to dimension-level correlations will be included in appendix to make this point clearer in our final version.
Thank you for your explanation.

> Eq. (4),(6) and (7) will be unchanged, only changing the meaning of the index $i$ from $i$-th instance to $i$-th dimension.

If so, the index $i$ of $\boldsymbol{x}_i$ in Eq. (4) moves in $\\{ 1, 2, \dots , D\\}$. Is it correct?  
Could you write down the modified version of Eq. (4) here?
Thanks for your reply to our rebuttal. The reason why we cannot expect that the links imposed by constraints are largely aligned with the edges in the tree is that for two non-adjacent nodes in a tree, if they are close, they still could be strongly correlated. For example, for a tree with 3 nodes (node 1 connected to node 2, and node 2 connected to node 3), although node 1 is not directly connected to node 3, node 1 and node 3 still could be strongly correlated. 

Thus, for two nodes that are constrained to be linked, it is fine to see they are not directly connected in the tree, as long as they are close to each other in the tree.

To investigate whether the learned tree has reflected the constraints correctly, we calculate the averaged correlation coefficient between nodes with must-link constraints as well as the averaged coefficient between nodes with cannot-link constraints (still conducted on the MNIST dataset). We observe that the averaged correlation coefficient of nodes with must-link constraints is 0.6481, showing strong positive correlation. On the other hand, the averaged correlation coefficient of nodes with cannot-link constraints is -0.5757, showing strong negative correlation. The results clearly demonstrate that the learned tree reflects the information imposed by the must-link and cannot-link constraints correctly.
We thank the reviewer for the comments and suggestions. Capturing correlations among dimensions within a specific latent embedding $\mathbf{z}$ can be done by simply replacing the latent variables from $\mathbf{z} _ i$ to $z_d$, $d \in \mathcal{V} = \\{1,\cdots,D\\}$, where $\mathbf{z} _ i \in \mathbb{R}^D$ denotes the latent variable corresponding to instance ${\mathbf{x}} _ i$; and $z_d$ just denotes the $d$-th dimension of latent variable ${\mathbf{z}} \in {\mathbb{R}}^D$. So the expressions of Eq. (4),(6) and (7) will be unchanged, only changing the  meaning of the index $i$ from $i$-th instance to $i$-th dimension. Eq. (12) will be slightly modified. When capturing the correlation among dimensions within one instance, we don't need to feed two instances ${\mathbf{x}} _ i$ and ${\mathbf{x}} _ j$ into the re-parameterization network $f_\phi(\cdot, \cdot)$, but instead only need to feed one instance ${\mathbf{x}}$. Thus, Eq. (12) becomes $\\{\gamma _ {ij}\\} _ {(i,j)\in {\mathcal{E}}} = f _ \phi({\mathbf{x}})$, where here $i$ and $j$ mean the $i$-th and $j$-th dimension; and ${\mathcal{E}}$ denotes the set of edges that connect different dimensions.
We appreciate the reviewer for carefully reading our explanations and recognizing our original contributions in this paper. As for the clarity issue, we agree that due to the limited space, we haven't provided sufficient  explanation to the computational complexity as well as the elegance of the proposed method. But as an additional page is allowed in the final version, we believe more detailed explanations on the complexity and the distinguished properties of the proposed method, as stated in our response to you above, could be easily incorporated into the paper. Thus, we think this issue could be easily addressed in the final version.

As for the question on computing the correlation coefficients from a tree (Eq. (3)), yes, the proposed correlation matrix ${\mathbf{R}}$ constructed from a tree can only represent a subset of all legitimate correlation matrices. We will make this point more clear and cite appropriate literatures here.
Thank you for the response.

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

Now, I believe I better understand the contribution. Thank you for the explanation. However, there are some clarity issues, and I believe the writing could be substantially improved. Let me be more precise. In Section 2.1, which is crucial to understanding the contribution of the paper, there are zero citations. Therefore, it is somewhat hard to distinguish which part of the claims are mathematical facts that are original to this paper, approximations, or well-known results. For instance, if I understood correctly, Eq 3 is saying we can retrieve the full correlation structure just from the tree right? What is unclear here is whether all fully-correlated Guassians can be retrieved from this formula or a subset of correlated Gaussians. If this is a known fact in graphical models, a citation would be needed here.

Furthermore, the main motivation for using trees here is that one can model $N^2$ interactions from $N$ queries to a neural network, which is nowhere mentioned in the introduction! Both the abstract and the introduction only mention that the method proposed here is more "scalable" and can be "reparameterized efficiently." However, the derivation of the reparameterization trick itself appears fairly straightforward and, in my opinion, is not the most critical contribution here. Basically, if one doesn't carefully read Lines 133-135, it is easy to miss the fundamental motivation for the tree structure.

With that said the role of the optimization problem is now clear. I do agree that the paper has some original contributions, so I am happy to raise my score. However, as other reviewers pointed out and pointed out above, the paper has some clarity issues and I feel that a thorough revision and an additional round of review would be very beneficial to the paper.
A4: I understand that the learnt correlation should capture the similarity of instances but I would assume it should also reflect the given constraints. Could the authors provide further explanation of why the correlations captured by the tree should not be guided by the pairwise constraints (and as a consequence why the tree should not reflect the given constraints)?
Thank you for your response.

I think the figure in the answer to Q1 is interesting. Including it in the final paper would enhance the value of the paper.

However, I still have questions regarding the answer to Q2. Since many equations in the paper are based on the assumption that $\mathcal{V}=\\{1, 2, \dots , N\\}$, it is not clear to me whether it can be directly applied to capturing correlations between dimensions within one instance. For example, if $\mathcal{V}=\\{1, 2, \dots , D\\}$, how would equations (4), (6), (7), and (12) be modified?