Title: TreeVI: Reparameterizable Tree-structured Variational Inference for Instance-level Correlation Capturing
Venue: NeurIPS 2024
Accepted: yes
Keywords: Probabilistic Models

Variational Inference

Instance-level Correlation

Reparameterization

Reparameterized Variational Inference
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.

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

Abstract

Mean-field variational inference (VI) is computationally scalable, but its highlydemanding 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 involving graphs or explicit constraints among instances. In this paper, we developed the Tree-structured Variational Inference (TreeVI) 2 , which uses a tree to capture the correlation among latent variables in the posterior. 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 the efficacy of TreeVI in capturing instance-level correlation in posteriors and enhancing the performance of downstream applications.

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

2 Tree-structured Variational Inference

Consider a latent-variable model p θ ( X , Z ) = p θ ( X | Z ) p ( Z ) , where X = [ 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)

̸

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

̸

2.1 TreeVI: Variational Inference under a Single Tree

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 ND is the mean and Σ z = diag( σ z ) R diag( σ z ) is the ND × ND covariance matrix; σ z = [ σ ⊤ 1 , · · · , σ ⊤ N ] ⊤ ∈ R ND is the standard deviation; and R means the correlation matrix, with its ( i, j ) -th block taking the form

γ 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

̸

𝐳

!

𝐳

"

𝐳

#

𝐳

$

𝐳

%

(a) Fully Connected

1

12

𝕋

2

3

5

4

23

24

35

(b) TreeVI

1

12

𝕋

!

2

3

5

4

23

24

35

1

2

3

4

5

𝕋

"

⋯

12

13

24

35

(c) MTreeVI: Extension to Multiple Trees

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 is

where 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 as

where the marginal distribution q ϕ ( z i | x i ) = N ( z i ; µ i , diag( σ 2 i )) , and pairwise marginal distribution for adjacent variables ( i, j ) ∈ E is q ( z i , z j | x i , x j ) = N ( z i , z j ; µ ij , Σ ij ) ; and the mean µ ij = [ µ ⊤ i , µ ⊤ j ] ⊤ and covariance matrix

γ 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

where 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 is

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

Theorem 1. Suppose that N latent variables Z = [ z 1 , · · · , z N ] ⊤ follow a tree-structured posterior distribution q T ϕ ( Z | X ) with the tree structure T = ( V , E ) , the joint sample ( z 1 , z 2 , · · · , z N ) ∼ q T ϕ ( Z | X ) can be expressed as

̸

where ϵ 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 sample Z = [ z ⊤ 1 , · · · , z ⊤ N ] ⊤ can be re-parameterized in a matrix-form as

where ϵ i ∼ N ( 0 D , I D ) for i ∈ V ; and L ( T ) z is a ND × ND matrix, with its ( i, j ) -th block [ L ( T ) z ] ij = diag( ℓ ij ) , and

̸

Actually, 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 that L ( 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

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 ϕ ( · , · ) as

according 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 ϕ ( z i , z j | x i , x j ) q ϕ ( z i | x i ) q ϕ ( z j | x j ) 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 TreeVI

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

2.2 Extension to Multiple Trees

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 distribution

where 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 T m ϕ ( Z | X ) with respect to the m -th tree component is expressed in the form of Eq. (4). Each tree component 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 Γ T m = { γ m ij : ( i, j ) ∈ E m } to be learned, and a weight controlled by a tree coefficient π m , where m = 1 , · · · , M .

With the mixture-of-trees posterior, the data log-likelihood has the evidence lower bound log p ( X ) ≥ ∑ M m =1 π m E Z ∼ q T m ϕ ( Z | X ) [log p θ ( X , Z )] + H [ q MT ϕ ( Z | X )] , where the first term can be estimated by tree-structured reparameterization Z ( T m ) ∼ q T m ϕ ( 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 T m ϕ ( 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 MTreeVI

where Z ( T m ) 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.

2.3 Learning the Tree Structure from Data

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 if

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 structure T ( 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 problem

̸

̸

where ℓ ( θ , ϕ , 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.

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

4 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 V AE [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.

4.1 Synthetic VAE

We design a synthetic dataset with a graph-structured latent variable model. The dataset contains N = 6000 data points x 1 , · · · , x N ∈ R D with D = 4 , each independently generated from the conditional distribution p ( x | z ) = N ( x ; z , σ 2 I 4 ) given latent embeddings z 1 , · · · , z N ∈ R D where σ 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 E 1 = { (1 , 2) , (1 , 3) , (2 , 4) } and E 2 = { (1 , 2) , (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.

4.2 Constrained Clustering

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 ; µ c i , diag( σ 2 c i )) , 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 W

̸

where π = { π 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 δ c i c j ) 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 , and W 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

9

10

8

10

0

10

10

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.

4.3 User Matching

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 u i ( 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 nodes V = { 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 2 N ; ( I 2 N + λ 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 2 N eval -1 synthetic user candidates in terms of latent embedding distances.

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

4.4 Link Prediction

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 u i 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 = ∑ ( u i ,u j ) ∈E test |{ 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.

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

Acknowledgment This work is supported by the National Natural Science Foundation of China (No. 62276280, U1811264), Guangzhou Science and Technology Planning Project (No. 2024A04J9967).

References

Agrawal, A. and Domke, J. (2021). Amortized variational inference for simple hierarchical models. In Neural Information Processing Systems .

Basu, S., Banerjee, A., and Mooney, R. J. (2004). Active semi-supervision for pairwise constrained clustering. In Proceedings of the 2004 SIAM international conference on data mining .

Bedford, T. and Cooke, R. M. (2001). Probability density decomposition for conditionally dependent random variables modeled by vines. Annals of Mathematics and Artificial Intelligence , 32:245-268.

Blei, D. M., Kucukelbir, A., and McAuliffe, J. D. (2016). Variational inference: A review for statisticians. Journal of the American Statistical Association , 112:859 - 877.

Byrd, R. H., Lu, P., Nocedal, J., and Zhu, C. (1995). A limited memory algorithm for bound constrained optimization. SIAM Journal on scientific computing , pages 1190-1208.

Caterini, A., Cornish, R., Sejdinovic, D., and Doucet, A. (2021). Variational inference with continuously-indexed normalizing flows. In Uncertainty in Artificial Intelligence .

Coates, A., Ng, A., and Lee, H. (2011). An analysis of single-layer networks in unsupervised feature learning. In International Conference on Artificial Intelligence and Statistics .

Hamilton, W. L., Ying, Z., and Leskovec, J. (2017). Inductive representation learning on large graphs. In Neural Information Processing Systems .

Harper, F. M., Konstan, J. A., and A., J. (2016). The movielens datasets: History and context. ACM Trans. Interact. Intell. Syst.

He, K., Zhang, X., Ren, S., and Sun, J. (2015). Deep residual learning for image recognition. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pages 770-778.

Hoffman, M. D. and Blei, D. M. (2014). Stochastic structured variational inference. In International Conference on Artificial Intelligence and Statistics .

Huszár, F. (2017). Variational inference using implicit distributions. ArXiv , abs/1702.08235.

Janke, T., Ghanmi, M., and Steinke, F. (2021). Implicit generative copulas. In Neural Information Processing Systems .

Jiang, Z., Zheng, Y., Tan, H., Tang, B., and Zhou, H. (2016a). Variational deep embedding: An unsupervised and generative approach to clustering. In International Joint Conference on Artificial Intelligence .

Jiang, Z., Zheng, Y., Tan, H., Tang, B., and Zhou, H. (2016b). Variational deep embedding: An unsupervised and generative approach to clustering. In International Joint Conference on Artificial Intelligence .

Kim, M. and Pavlovic, V. (2020). Recursive inference for variational autoencoders. Advances in Neural Information Processing Systems .

Kingma, D. P. and Ba, J. (2014). Adam: A method for stochastic optimization. CoRR , abs/1412.6980.

Kingma, D. P. and Welling, M. (2013). Auto-encoding variational bayes. CoRR .

Kipf, T. and Welling, M. (2016). Variational graph auto-encoders. ArXiv , abs/1611.07308.

Lambert, M., Chewi, S., Bach, F., Bonnabel, S., and Rigollet, P. (2022). Variational inference via wasserstein gradient flows. Advances in Neural Information Processing Systems .

LeCun, Y. and Cortes, C. (2010). MNIST handwritten digit database.

Lewis, D. D., Yang, Y., Russell-Rose, T., and Li, F. (2004). Rcv1: A new benchmark collection for text categorization research. Journal of machine learning research , pages 361-397.

Lin, W., Hubacher, N., and Khan, M. E. (2018). Variational message passing with structured inference networks. arXiv: Machine Learning .

Manduchi, L., Chin-Cheong, K., Michel, H., Wellmann, S., and Vogt, J. (2021). Deep conditional gaussian mixture model for constrained clustering. Advances in Neural Information Processing Systems .

Manduchi, L., Vandenhirtz, M., Ryser, A., and Vogt, J. (2023). Tree variational autoencoders. Advances in Neural Information Processing Systems .

Moens, V., Ren, H., Maraval, A., Tutunov, R., Wang, J., and Ammar, H. (2021). Efficient semi-implicit variational inference. CoRR , abs/2101.06070.

Morningstar, W. R., Vikram, S., Ham, C., Gallagher, A., and Dillon, J. V. (2020). Automatic differentiation variational inference with mixtures. In International Conference on Artificial Intelligence and Statistics .

Nemirovski, A. (1999). Optimization II: Standard Numerical Methods for Nonlinear Continuous Optimization .

Ou, Z., Su, Q., Yu, J., Liu, B., Wang, J., Zhao, R., Chen, C., and Zheng, Y. (2021). Integrating semantics and neighborhood information with graph-driven generative models for document retrieval. In ACL/IJCNLP (1) , pages 2238-2249.

Papamakarios, G., Pavlakou, T., and Murray, I. (2017). Masked autoregressive flow for density estimation. Advances in neural information processing systems , 30.

Pearce, M., Chiappa, S., and Paquet, U. (2018). Comparing interpretable inference models for videos of physical motion. In 1st symposium on advances in approximate bayesian inference .

Ren, Y., Hu, K., Dai, X., Pan, L., Hoi, S. C. H., and Xu, Z. (2019). Semi-supervised deep embedded clustering. Neurocomputing , pages 121-130.

Smith, M. S. and Loaiza-Maya, R. (2021). Implicit copula variational inference. Journal of Computational and Graphical Statistics .

Smith, M. S., Loaiza-Maya, R., and Nott, D. J. (2019). High-dimensional copula variational approximation through transformation. Journal of Computational and Graphical Statistics .

Song, J., Meng, C., and Ermon, S. (2021). Denoising diffusion implicit models. In 9th International Conference on Learning Representations .

Tang, D., Liang, D., Jebara, T., and Ruozzi, N. (2019). Correlated variational auto-encoders. In International Conference on Machine Learning .

Tölle, M. and Schlaefer, A. (2021). A mean-field variational inference approach to deep image prior for inverse problems in medical imaging. In International Conference on Medical Imaging with Deep Learning .

Wang, Q. and Van Hoof, H. (2020). Doubly stochastic variational inference for neural processes with hierarchical latent variables. In International Conference on Machine Learning .

Wang, Y., Liu, F., and Schiavazzi, D. E. (2022). Variational inference with nofas: Normalizing flow with adaptive surrogate for computationally expensive models. Journal of Computational Physics .

Xiao, H., Rasul, K., and Vollgraf, R. (2017). Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. ArXiv , abs/1708.07747.

Xie, J., Girshick, R. B., and Farhadi, A. (2016). Unsupervised deep embedding for clustering analysis. In ICML .

Yang, L., Cheung, N.-M., Li, J., and Fang, J. (2019). Deep clustering by gaussian mixture variational autoencoders with graph embedding. In Proceedings of the IEEE/CVF international conference on computer vision .

Yin, P., Zhou, C., He, J., and Neubig, G. (2018). Structvae: Tree-structured latent variable models for semi-supervised semantic parsing. ArXiv , abs/1806.07832.

Zhang, H., Basu, S., and Davidson, I. (2019). A framework for deep constrained clustering algorithms and advances. In ECML/PKDD .

Zheng, X., Aragam, B., Ravikumar, P., and Xing, E. P. (2018). Dags with no tears: Continuous optimization for structure learning. In Neural Information Processing Systems .

A Precision Matrix of Tree-structured Posterior

When latent variables Z = [ z 1 , · · · , z N ] ⊤ are all 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 as

with the inverse of tree-structured correlation matrix given by

demonstrating that only correlation parameters along the edge set E = { (1 , 2) , (2 , 3) , (2 , 4) , (3 , 5) } 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 } :

s

35

=

-

γ

35

⊙

(

1

D

-

γ

2

35

)

,

(28)

B Ancestral Sampling for Tree-structured Posterior

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

̸

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

-

1

Proof. For each latent variable z j , there exist a unique path P 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 by

The ancestral sampling starts by sampling from q ϕ ( z v 1 ) : z v 1 = µ v 1 + σ v 1 ⊙ ϵ v 1 , where ϵ v 1 ∼ N ( 0 D , I D ) . Given the observation of z v 1 , then the latent variable z v 2 can be sampled using the conditional distribution Eq. (7):

Assume that at the k -th step of the ancestral sampling, we have

where ˜ γ v i v k = γ v i v i +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 sampling

Therefore by induction, the ancestral sampling for the latent variable z j can be expressed as

which is consistent with the expression of Eq. (31).

B.2 Example of Ancestral Sampling Procedure

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.

Start by sampling from q ϕ ( z 1 ) :

where ϵ ∼ N ( 0 , I )

Then sample from q ϕ ( z 2 | z 1 ) given the observation of z 1 :

Then sample from q ϕ ( z 3 | z 2 ) given the observation of z 2 :

where ϵ 3 ∼ N ( 0 D , I D ) ;

Then sample from q ϕ ( z 4 | z 2 ) given the observation of z 2 :

where ϵ 4 ∼ N ( 0 D , I D ) ;

Finally, sample from q ϕ ( z 5 | z 3 ) given the observation of z 3 :

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

C Evidence Lower Bound

C.1 Evidence Lower Bound for TreeVI

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 by

where Z ( T ) denotes the tree-structured reparameterization for latent variables. The first term can be directly computed by

where z ( T ) i is 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 )

where the entropy of singleton posterior q ϕ ( z i | x i ) = N ( z i ; µ i , diag( σ 2 i )) is given by

where Σ i = diag( σ 2 i ) , and the step ⋆ relies on several properties of the trace operator:

And the entropy of pairwise posterior q ϕ ( z i , z j | x i , x j ) = N ( z i , z j ; µ ij , Σ ij ) with mean µ ij = [ µ i , µ j ] ⊤ and covariance matrix

is similarly given by

where we denote z ij = [ z i , z j ] ⊤ .

C.2 Evidence Lower Bound for MTreeVI

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 by

where Z ( T m ) 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 by

where z ( T m ) i is 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 ) .

D Constrained Optimization Details

D.1 Proof of Proposition 1

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 :

Suppose that A ∈ { 0 , 1 } N × N and r ( A ) < 1 , then A is acyclic if and only if

Suppose that A ∈ { 0 , 1 } N × N , then A is acyclic if and only if

Suppose that A ∈ R N × N , then A is acyclic if and only if

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 lengthk closed walks in a directed graph.

Clearly the directed graph induced by the symmetric matrix A will only have self-loops, and hence tr A 2 k +1 = 0 for all k = 1 , · · · , ∞ . In other words, A has no cycles if and only if f ( A ) = ∑ ∞ k =1 ∑ N i =1 ( A 2 k +1 ) ii = 0 , then

The desired result follows.

Similar to the proof of Proposition 1.1 by noting that A has no cycles if and only if ( A 2 k +1 ) ii = 0 for all k ≥ 1 and all i ∈ { 1 , · · · , N } , which is true if and only if

The proof is similar to Proposition 1.2 by replacing A with A ⊙ A , which counts weighted closed walks.

D.2 Greedy Searched Initialization

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 [ · ] means randomly choosing one index according to the indicator function; ID [ · ] represents the set of node indexes satisfying the indicator condition and N ( i ) denotes the neighbors of node i .

Algorithm 1 DFS Algorithm for Tree Generation

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

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.

D.3 Solving Constrained Optimization

In section 2.3, we establish a continuous characterization of acyclicity, leading to the following equality-constrained program (ECP)

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

Algorithm 2 Algorithm for Constrained Optimization

Input: Initial guess ( A 0 , α 0 ) ; progress rate c ∈ (0 , 1) ; tolerance ϵ > 0 ; threshold ω > 0 Output: Threshold matrix 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:

with 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 α

is the augmented Lagrangian. And the goal is to find a local solution to the dual problem

Let 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):

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

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

E Experimental Details

E.1 Datasets

The datasets used in the experiments are the followings:

MovieLens 20M: A dataset describing ratings and free-text tagging activities from MovieLens, 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: Adataset 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].

E.2 Implementation Details

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

is shared for all data points to capture dimension-level correlation. Our synthetic V AE is trained for 200 epochs for all methods. We apply stochastic gradient optimizations with a step size of 0.005, and use the Adam algorithm [17] to adjust the learning rates.

Constrained Clustering. To implement our methods, we are careful in maintaining a fair comparison with the baseline methods. In particular, we adopt the same encoder and decoder feed-forward architecture used by the baseline method: four layers of 500, 500, 2000, D units respectively, where D = 10 unless stated otherwise. VAE-based baselines and the VAEs equipped with our posterior approximation methods are pretrained for 10 epochs while the DEC-based baselines involve 50 epochs of pretraining for each layer and 100 epochs of pretraining as finetuning. Each dataset is divided into training and test sets, where the former one is used for training and our reported results are operated on the latter one. The pairwise constraints used in our experiments are randomly chosen within the training set, by randomly sampling any two data points and assigning a must-link if they have the same label and a cannot-link otherwise. The absolute values of the elements | W ij | in the pairwise prior information matrix are set to 10 4 for all datasets for convenience, and 6000 pairwise constraints are sampled for training both our methods and other baselines. Following DC-GMM [24], the hyper-parameters are universally set for four different datasets, as shown in Table 6. The learning rate is set to 0.001 and it decreases every 20 epochs with a decay rate of 0.9. The number of tree components adopted in our MTreeVI is set to M = 3 .

User Matching & Link Prediction. For both tasks of user matching and link prediction, we set the dimensionality of latent embeddings as D = 100 for all methods. For V AE-based methods (including TreeVI and MTreeVI), 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 posterior approximation q ϕ ( z i | x i ) . The model likelihood function p θ ( X | Z ) is a multinomial distribution, and the singleton posterior distributions q ϕ ( z i | x i ) are all diagonal normal distributions. For CVAE corr and our methods, we also learn a two-layer feed-forward neural network that takes the concatenation [ x i ; x j ] as input and output the correlation parameter between z i and z j on each of the D dimensions. To determine the tree structure for TreeVI and MTreeVI, we take average of these correlation parameters across all D dimensions and obtain a correlation matrix A ∈ R N × N or a set of correlation matrices that is used for constrained optimization. For GraphSAGE, we choose to use K = 2 aggregation steps and use the mean aggregator function. We use Q = 20 negative samples to optimize the loss function. For all methods, we apply stochastic gradient optimizations with a step size of 0.001, and use the Adam optimizer to adjust the learning rates. The number of tree components adopted in our MTreeVI is set to M = 3 for both tasks.

E.3 Further Experiments

Synthetic Data. We design a synthetic dataset with a graph-structured latent variable model. The dataset contains N = 6000 data points x 1 , · · · , x N ∈ R D with D = 4 , each independently generated from the conditional distribution p ( x | z ) = N ( x ; θ z , σ 2 I 4 ) given latent embeddings z 1 , · · · , z N ∈ R D where σ 2 is a fixed value and set to 0.5. The latent embeddings z 1 , · · · , z N are drawn from a standard normal distribution p ( z ) = N ( z ; 0 4 , I 4 ) , and the graph-structured correlation

is incorporated into the sample generation via a coefficient matrix θ

And 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 kernel

where 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 as

̸

The clustering performances of our proposed methods against baseline methods are shown in Table 8.

E.4 Resource Usage

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.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The main claims in the abstract and introduction have clearly reflected the paper's main contributions in the following context.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Limitations of our work has been discussed in Section 5 in the paper.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

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.

4. Experimental Result Reproducibility

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.

Guidelines:

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

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., 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.

5. Open access to data and code

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.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

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

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

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

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

Answer: [Yes]

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

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

11. 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 authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: No new assets are introduced or released in our paper.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Our work does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our work does not involve crowdsourcing nor research with human subjects.

Guidelines:

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nWhile using tree structures in VAEs have been already explored [1], the setting used by this paper deviates from [1] as it considers correlations across samples and not across hierarchical latent variables. I believe that the idea is novel and the method sound. The experiments show improvements over related work and show a broad applicability of the proposed inference.\n\n[1] Manduchi et al. Tree Variational Autoencoders, NeurIPS 2023, https://proceedings.neurips.cc/paper_files/paper/2023/file/ac58b418745b3e5f10c80110c963969f-Paper-Conference.pdf.\n\nWEAKNESSES:\n1. 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.\n2. 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).\n3. 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?\n4. 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?). \n5. I think the paper is not always clearly written and it contains several typos (see line 350 or 316, 317 as examples).\n\nQUESTIONS:\n1. Could the authors provide a proof of Equation 4 (in the appendix)?\n2. 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.\n3. 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?\n\nLIMITATIONS:\nI believe the authors have addressed the limitations of the proposed approach.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nTheir methods are methodologically interesting, well described. The experiments compare the proposed method to relevant prior work and demonstrate that the method is performing well empirically.\n\nWEAKNESSES:\nIt 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.\n\nQUESTIONS:\nThe 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?\n\n- 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?\n\n- 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?\n\n- 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.\n\n--- Edit ---\n\nThe 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).\n\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* The approach successfully incorporates inter-instance correlations into the approximate posterior distribution while maintaining efficient training by utilizing a tree structure.\n* The proposed methods show improved performance over existing techniques across various tasks.\n\nWEAKNESSES:\nIn 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.\n\nQUESTIONS:\nOn 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?\n\nLIMITATIONS:\nThe limitations are appropriately addressed within the paper, and I have no additional suggestions in this regard.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* In variational inference with structured variational families, learning the appropriate correlation structure from data is an important problem that hasn't been fully addressed.\n\nWEAKNESSES:\n* 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.\n* 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.\n* Given the comment above, it is unclear *why* we have to take a tree structure among various different correlation graph structures.\n* 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.\n    * 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.\n    * 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\"\n\nOverall, 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.\n\n## Additional Comments\n* Section 2.2: The additional Jensen gap can be avoided by using the sticking-the-landing gradient [13].\n* 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.\n* Line 55: Cite something for the ELBO. It is typical to cite [9,10].\n* Line 28: The mean-field approximation was first proposed by [11] and popularized by [12].\n* 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.\n\n## References\n(Disclaimer: I am not the author of any of the papers below.)\n1. Katzfuss, Matthias, and Joseph Guinness. \"A general framework for Vecchia approximations of Gaussian processes.\" (2021): 124-141.\n2. Hoffman, Matthew, and David Blei. \"Stochastic structured variational inference.\" Artificial Intelligence and Statistics. PMLR, 2015.\n3. Titsias, Michalis, and Miguel Lázaro-Gredilla. \"Doubly stochastic variational Bayes for non-conjugate inference.\" International conference on machine learning. PMLR, 2014.\n4. Kucukelbir, Alp, et al. \"Automatic differentiation variational inference.\" Journal of machine learning research 18.14 (2017): 1-45.\n5. Caterini, Anthony L., Arnaud Doucet, and Dino Sejdinovic. \"Hamiltonian variational auto-encoder.\" Advances in Neural Information Processing Systems 31 (2018).\n6. Geffner, Tomas, and Justin Domke. \"MCMC variational inference via uncorrected Hamiltonian annealing.\" Advances in Neural Information Processing Systems 34 (2021): 639-651.\n7. Doucet, Arnaud, et al. \"Score-based diffusion meets annealed importance sampling.\" Advances in Neural Information Processing Systems 35 (2022): 21482-21494.\n8. Thin, Achille, et al. \"Monte Carlo variational auto-encoders.\" International Conference on Machine Learning. PMLR, 2021.\n9. Jordan, Michael I., et al. \"An introduction to variational methods for graphical models.\" Machine learning 37 (1999): 183-233.\n10. 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.\n11. Peterson, Carsten, and Eric Hartman. \"Explorations of the mean field theory learning algorithm.\" Neural Networks 2.6 (1989): 475-494.\n12. 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.\n13. 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).\n\nQUESTIONS:\nn/a\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}