TreeVI: Reparameterizable Tree-structured Variational Inference for Instance-level Correlation Capturing
Authors: Junxi Xiao
Qinliang Su
Venue: NeurIPS 2024 poster (2024)
Keywords: Probabilistic Models
Variational Inference
Instance-level Correlation
Reparameterization
Reparameterized Variational Inference
Primary area: probabilistic_methods

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.

Tree VI: Reparameterizable Tree-structured Variational Inference for Instance-level Correlation Capturing
Junxi Xiao
Qinliang Su
School of Computer Science and Engineering, Sun Yat-sen University, Guangzhou, China "Gunngdong Key Laboratory of Big Data Analysis and Processing, Guangzhou, China xiaojx7 (rjix7mmaa222 12. sysu. edu. cn suqliangdmail. sysu. wyw..syu..uu..... edu.cn
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 involving graphs or explicit constraints among instances, In this paper, we developed the Tree-structured Variational Inference (TreeVV)² which uses tree 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 or times that of VI under the mean mean-field assumption. To capture correlation with more complicated structure, the TreeVi is further extended to the multiple-tre 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 probabilistic method that is widely used for approximating the exact posterior latent-variable models p(X, Z) = p(X(Z))((Z) Z)) with X and being the observed data and latent variable, respectively. When not considering the existence of any relation among instances, we can write the model into factorized form as p(X, Z) = II, p(x:222 (x.)))))((.). where x; and z; denote the i-th data instance and latent variable. Thanks to the factorized form, reasonably use the mean-field posterior q( X) = II, q(z, x,) 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 edge ei between them via p(eis [2;, 2,),, which, obviously, makes all instances coupled together. Also, in the constrained clustering, DC-GMM [24] designs prior p(c; W) specify the assignment constraints and then uses it to construct the prior distri- bution pZZ, c) = p(c; w) I1 p(z.((())) which will also cause all z; to be correlated with each other. Obviously, under these circumstances, due the existence of correlation among instances, it is unreasonable to assume factorized form for the posterior anymore.
"Corresponding author:
"Code available at: http:///gihuuu.commmtpteessshpet
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 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 novel posterior approximation for variational inference, the Tree-structured Variational Inference (TreeVI), that models the latent correlation structure with tree-structure- induced distribution and enables modeling high-order correlations between non-adjacent latent variables. Our method provides 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 VI to multiple trees, and propose a continuous optimization algorithm to stochastically learn theoretically correlation-rich tree or mixture-of-tree structure.
2 Tree-structured Variational Inference
Consider a latent-variable model po (X, Z) = (o( where X = [xx] x x and ]]: and x; and 2; E TR. denote the i-th data instance and its corresponding
Z = [2 2 latent variable, respectively. Given the latent-variable model po (X, Z), variational inference aims to find a distribution that is closet to the true posterior pe (Z]X) from distribution family Q, which is achieved by maximizing the evidence lower bound (ELBO)
[(0, Φ) = (Z)) [log X Z) log L (Z[)),,
where qs ( E Q denotes the approximate posterior parameterized by Φ. In most of current works, the latent-variable model is assumed to take a factorized form as px Z) -I pe (xx z, (((,) due to the observed independence among data instances. With the factorized form for the model, it can be seen that the true posterior po Po (Z) also takes a factorized form, hence it is reasonable to choose a fully factorized form 4 (Z = I 9o (2, (x,) 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 by using a correlated prior distribution p(Z) with p(Z) * [[;=1 p(z)), or use several latent variables zj to be responsible for generating a data instance x, (e.g., X, representing an edge), which will result in pe Po(X (X + II (x;((zz)) |z.) Obviously, under these scenarios, latent variables } from the true posterior po [Z X) are not independent anymore. Thus, if factorized posterior 9s (Z() = II Io (2,1x,)
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 VI: Variational Inference under Single Tree
To capture the correlations among latent variables {z: J in the posterior distribution, a multi- variate normal distribution (Z|x) = NZ M2, Ex) can be used, where μ4, = -IM [MT, , e RND is the mean and Σ, = diag(or,) R. diag((z, is the ND x ND covariance matrix; o2 = ........ . IJ" RND is the standard deviation; and R means the correlation matrix, with its
(i, block taking the form
Rly = diag(r.);
(2)
Yij E (-D, for i + j denotes the correlation strength between latent variable z, and zj and
7u A In for i E V A 21, ,N. Here, the diagonal structure assumed for the ji, block [R]y implies that we are only interested capturing the correlation among different variables, without considering the conrelation of different dimensions within one variable. Actually, due to the

---

Z
T₁
M
12
13
Z-2
Z
2
3
35
35
4
5
Z
5
5
(a) Fully Connected
(b) TreeVI
(c) MTreeVI: Extension to Multiple Trees
Figure I: Comparison between the colnecconeected correlation structure and approximation by single and multiple tree-structured variational inference
relatively low dimensions a latent variable z;, 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 p (7ig i # E V}. For a general multivariate normal distribution, all correlation
parameters in r are learnable, which admits highly expressive correlation structure, but also resulting in high computational complexity. The widely-used mean-field approximation reduces the cost by assuming independence structure within any pair of latent variables, i.e. Vij On for any i # j. completely ignoring the correlation among latent variables. In order to achieve a balance between expressiveness and complexity, we propose approximate the fully-connected correlation structure with a tree structure, introducing our Tree namely tree-structured variational inference.
The idea of tree- structurdd variational inference is to impose a tree correlation structure T (V, among latent variables, as illustrated in the quin-variate example in Fig. lb. In our TIe given a tree structure T (V, E), we only specify the correlation parameters Yig for latent variables (Z, zj) adjacent the tree, and for any non adjjeent latent variables (z,, 2ÿ). the correlation between then are directly computed as the multiplication of correlation parameters along Pi-ej, that is
Yis
II Y (E.))
(a.t)ER(-)
where Pr-.y denotes the path connecting the latent variable z, z, the tree T, and the vector multiplication assigned as Hadamard product. Therefore, only the correlation strengths w.r.t the edges of the tree are learnable, which can be denoted as r = (T (1.j) E E} c T The parameterization leads to correlation matrix [R [T) Yu that totally determined by parameters r", with 7, = 1p for i E V. For example, for the quin-variate tree structure in Fig. 1b, the correlation parameter between Z and Z; is fixed as 715 V12 723 Y35
With the tree-structured correlation defined above, it can be easily shown that the distribution o0= (Z.X) forms Markov random field defined by the tree T [3], and its joint probability can be expressed as
Jos 2y X;)
[Z
II VEV
a (z((x)
II
(x: xx))
(K,)EE
where the marginal distribution (z, (x,) Nazz u4, diag o?)), and pairwise marginal distribution
for adjacent variables (i,j) E E is q(z,, 2j xx, = N (2,, Zj; Mijz Eiy); and the mean Mij
[MTT. 14, J and covariance matrix
0:00;
Yoj 10,00,0
Σg
(5)
Yoj 11,00,00
Yay is the correlation parameter between two adjacent variables. Furthermore, also that the distribution o X) can be represented by an acyclic Bayesian network of the form
Te (Z(x) = = (z()))
II
(0 (z,,
(6)
(KJ)EE
3

---

where the conditioning on data x; is omitted for conciseness; 21 is assumed to be the root node; the conditional distribution of latent variables with respect to an edge (i.j) E Es
(0<2))2;)) = N (3jle + Vig о o. @ (z: - p,),o 2 VID 1D - 23
(7)
With the conditional distribution go (z, joint samples (z), 22. 2N) can be drawn from Ug (Z(X) with ancestral sampling. By sampling 2, one by one with ancestral sampling, we can show that the joint sample (Z1, Z2, ... , can be equivalently represented by a set of N independent Gaussian noises (c; ( with c; (OD, ID), stated in the theorem below.
Theorem 1. Suppose that N latent variables Z !Z1.-.. ZN] follow tree-structured posterior distribution 4 (Z) with the tree structure T EV.E), the joint sample (Z1, 22, .2N)
Lo (Z()) can be expressed as
2 H, 71 ① 61
00
Σ 9i
II
ID
P pa(i). OG for j jEV;
(8)
VEP1-j
(OD, ID) is a Gaussian random noise and pa(i) denotes the parent node of with respect to E V.
where E
For the proof and concrete example, refer to Appendix B. With the reparameterization of sample
2j in Eq. (8), the joint sample Z = [2] ... .2N] can be re-parameterized a matrix-form as
Z Pz L I. with E 6= [e] ,T
where 6 ~ NO,, In) for i V; and L (T)
is a ND x ND matrix, with its (i,j)-th block
[L 11 1ig diag(l.g). ). and
71 1 j.
= 1, 7 pa(i), Ooj, f Prjjje
Yis 3 In
(10)
On,
otherwise.
Actually, if the variables are indexed according to their positions the tree T from left to right and then top to bottom, can be easily shown that L is a lower-triangular block matrix. For example, the matrix Lg corresponding to the quin-variate example of Fig. lb can be written as follows
DD
712
VD - 782
L = diag(oz) 2
713
723
@ / In -712
Ip Y23
714
724
① VIn 7712
0p
VID -2
715
725 @ /1n
rī2
Y35 @ /1n 723
On
V1n
73._
(11)
where diag(-) has been omitted for conciseness.
By parameterizing the mean Me standard deviation o, and the correlation parameters Yi for (i.j) with neural network fo(
Yij xx),,, (i.5) E E,
(12)
according to Eq. (9), samples drawn from o (Z can be reparameterized in the form of neural network parameters & and random Gaussian noise e, which can facilitate the training variational inference significantly. Importantly, except the mean and standard deviation, we only need to re- parameterize the parameters r (ry (i.j) E E}. while computing the other non-zero elements
in L with I". In this way, we only need to re perammeteriize El parameters, and thus only need to run |8| times of neural network (+:)), instead of ((22) times the vanilla method. For a tree, the number of edges El < N 1, thus we only need to additionally run ( times of neural networks (.. in addition to the runs required by the mean-field method. For mean-fiele 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 the experiments on constrained clustering, the consumed time of our Tree 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 lij in the lower-triangular matrix L by a neural network (-.)), that is, Ejj (x x,), then we have re-parameterize as many as (N + 1)/2 elements, i.e., all elements from the lower-triangular positions of L need to be re-pann pa-aameterrzed, where N the number of instances in training dataset. That means we need to run the neural network fa(:-) -) by ((N2) times for every epoch, which is computationally unacceptable, especially considering that could be as large millions in practice. And our main contribution lies at finding a way to reduce the required times of rumning the neural network (-:-) from (NN²) ((N) by restricting the correlation matrix R to a special form R constructed from tree T (V,E), which is actually dense matrix with its (i. element [R.]] Fg defined as Eq. (3) for (i.j) & E. Under the restricted correlation matrix R (T), the lower-triangular matrix L2 possesses very elegant form, as shown in Eq. (10). The elegance lies at that although L still has ((N²) non-zero elements, all of these non-zero elements can be explicitly computed from the |8] parameters r" (Yiy" (i.j) E E}.
With the tree-structured posterior, the data log- likelihood has the evidence lower bound log p(X) > Ez [lo Poogoo((X (X, 2 + H (2()), where the first term can be estimated by reparam- eterization (TT)
u (Z X), and the entropy H of the tree-structured posterior distribution 90(2,,,1xx,xx11 1 (2))) Tley (2, (x;) II (ES)EE ga(2, (9)0,,)x)))081883. can be decomposed as entropy terms with respect to singleton posterior 4s (z: (x;) and pairwise posterior (z(2,, Zz |x,, xj) 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-bast -asedd evidence lower bound of our proposed TreeVI
(" Φ. X) = log pe X, Z (T) + H 4 (Z|X)],
(13)
where (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 Appendix C.I.
2.2 Extension to Multiple Trees
The expressiveness a single tree-structured posterior is still restrictive, To alleviate this issue, we propose to approximate the true posterior with mixture-of-trees posterior distribution
M
(Z
=
Σ m==1
m L (Z(X),
(14)
where weighted mixture of tree structures MIT (Ti,..., Tw.. as shown in Fig. Ic, to approximate the underlying correlation structure, and the posterior 9 " (Z(X) with respect to the m-th tree component is expressed in the form Eq. (4), Each tree component Tm = (V. E corresponds to set of latent indices V = {l, .N. set of pairwise connections Em set of correlation parameters r coefficient km, where m = 1,.................. ,..
{>]) (i,j) E } to be learned, and weight controlled by tree
With the mixture-of-trees posterior, the data log- ligllikel has the evidence lower bound log (X) > M
EmEn Thp E-zzgg (Z[x) Dog (X, Z H (Z X)), where the first term can be estimated by
tree-structured reparameterization (T... ~ 0 T. (Z within each tree component Tm 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 lower bounded by the weighted sum of entropies with respect to each tree component ΣM= T. 7(q " (Z] which can be similarly computed by the singular and pairwise marginal posterior distributions with respect to each tree component T m 1, M.. Therefore, the optimization of variational
5

---

inference with mixture- treet posterior approximation can be simplified as maximizing the following sampling-base based evidence lower bound of our proposed VT
C (0 X)
Ti log pe (X, (X, Z (F)) + Hlgg (2/x)]
(15)
where (Tm) denotes the tree-structured reparameterization for latent variables with respect to the tree structure Tm- 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-tree structured posterior approximation make variational inference with pairwise latent correlations feasible, but specific tree structure over latent variables need to be determined in advance. In this section, our goal is develop an efficient algorithm for learning correlation-rich tree structures to approximate the underlying posterior correlation structure.
To learn meaningful tree structure from the latent embeddings, adopt symmetric binary matrix A E (0. 1) JNNN with entries O along the diagonal, where each element dgg E (0,1) 1} represents an edge between latents z, and zz. i + j (N, (here we assume
1 for convenience).
Proposition Suppose that the symmetric adjacency matrix A E {0, 1] NxN, then the undirected structure induced by adjacency matrix A is acyclic if and only if
h(A) = tr exp(A2) =0
(16)
where tr(.) and exp(.) represent the truce and exponential of matrix respectively, and @ is the Hadamard product.
For the proof we refer to the Appendix D.I. Inherited from the idea of Zheng et al. [45], we use similar indicator function 1i( (A) in Eq. (16) to check the acyclicity of the structure induced by the symmetric binary matrix A. If and only if (A) O, 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.C. {1, _N3 and edges & {(i,j) ai # 0, i + E V}. Under the constraints Eq. (16), we seek to establish the following continuous optimization problem
min { Φ, A),, subject to h(A) 0. 0.0.A
(17)
where - Φ, X) is the negative of the evidence lower bound given by Eq. (13), and the correlation parameters ( {Yij ai # #ii j V} of the learned tree determined by the binary adjacency matrix A with Yi network. A similar optimization can be implemented to stochastically learn a meaningful mixture of trees by using set of symmetric adjacency matrices 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.
calculated by the neural
(A1, A&r) } to represent multiple tree
To initialize the tree structure before constrained optimization, the casiest 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 neighbor of each instance according to their similarity. For detailed implementation of the tree initialization, we refer to the Appendix D.2.
Related Work
Variational inference is broadly used for approximating intractable posterior in latent variable models, but notorious for its restricted variational families, especially mean-field variational family which is still widely used by modern generative models [4, 37). So far there has been a wide variety of variational inference methods that attempt to improve on traditional mean-field variational inference by modeling dependence structures within latent posterior distribution, such constructing the

---

variational distribution with a normalizing flow [30, 6, 39], which deterministically transforms simple probability distribution over latent variable 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- levell correlation structure.
Recently, there has been some work on incorporating instance-level correlation structure in variational inference [24, 36, 29J. Manduchi et al. [24] designs 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 instance- level 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 unable to be modeled within the tree structure.
Experiments
Tasks & Datasets. We evaluate our methods with four different tasks: synthetic evidence lower bound test, constrained clustering, user matching and link prediction, on synthetic dataset, standard datasets (MNIST, Fashion MNIST, Reuters and STL-110,, public movie rating and product rating datasets, respectively. We refer Appendix E. for more dataset details.
Baselines & Implementation Details. For baselines on the user matching and link prediction tasks, we include the standard variational auto-encoders [18] and a recent modification to VAE [36] taking instance-level correlation structure into consideration. And we also compare our method to the state-of-the-arr 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 state- of-the-at constrained clustering methods for comparison. We also include the generative models, the unsupervised VaDE [14], the graph augmented VaDE (DGG, [42]), and the seaerv-superrvsed experiment on methods Tree and VT, where the number of tree components for the mixture-of posterior is fixed as M = 3. Aadditional information related experimental implementation details are in Appendix E.2.
DC-GMM [24]. For comparison,
4.1 Synthetic VAE
Table 1: Estimated lower bounds (ELBO) of VAE with posterior dis- tributions approximation by mean- field distributions, tree structures T1 and T2, as well their mix- ture model VI(T,, com- pared ground truth log- log-likelihood log p(X).
We design a synthetic dataset with a graph-structured latent variable model. The dataset contains N 6000 data points X1. ..y E Rº' with D = 4, each independently gener- ated from the conditional distribution p(x zz) N(x: 2, I given latent embeddings Z1, Zx E T D where o² is fixed value and set 0.5. The latent embeddings 21, ... 2, are drawn from zero-mean normal distribution p(z) (2 04, Σz). and the graph-structured correlation is incor- porated into the latent covariance matrix Σ, 14 + XA via 4 assigned with a loopy graph
Methods
Lower Bound
an affinity matrix structure
Mean-field
-11.1535
Tree (T1)
-10.8998
0
1
0
0.3
Tree (T2)
-10.6137
0
1
0.3
VT
-10.3586
A
(18)
0
1
0
0.4
log p(X)
-10.3417
0.3
0.3
0.4
7

---

Table 2: Clustering performances (%) of our proposed methods TreeVI and MT compared with baselines. Means and standard deviations are computed across 10 runs with different random initializations. Results taken from DC-GMM [24]
Datasct
Metric
VaDEt
SDEC
C-IDECt
DGG
DC-GMM
TrecVI
MTTre VI
MNIST
ACC
89.0 +5.0
86.2 #0.1
96.3 #0.2
95.8 #0.1
96.5 #0.2
97.4 #0.3
97.5 #0.4
NMI
82.8 #3.0
84.2 10.1
91.8 +1.0
91.2 #0.2
91.4 40.3
93.1 #0.6
93.1 #0.6
ARI
80.9 =5.0
80.1 #0.1
92.1 #0.4
91.4.0.3
92.5 40.5
93.7 +0.7
94.0 #0.5
IMNIST
ACC
55.1 =2.2
54.0 #0.2
68.1 43.0
79.9 #0.4
80.5 #0.8
81.4 +0.6
82.1 #0.7
NMI
57.9 =2.7
57.3 #0.1
66.7 +2.0
70.1 #0.3
72.0 20.4
73.9 20.6
74.1 =0.6
ARI
43.6
40.2 #0.1
52.3
64.9 #0.3
66.4 40.5
67.9 +0.9
68.1 =0.6
Reuters
ACC
=0.7
82.1 #0.1
94.7 #0.6
93.5 #0.6
95,4.0.2
95.9 20.6
96.1 =0.7
NMI
=1.3
=0.
81.4
81.2 =0.8
82.7 #0.7
83.4 +0.5
83,9 =0.5
ART
58.0
66.7 #0.1
87.7 40.9
87.8 #0.5
89.0 40.6
90.2 40.4
90.5 #0.4
STL-10
ACC
79.2 #0.5 =0.5
81.6 #3.8
89.9 =0.3
89.5 #0.5
90.4
90.7 =0.9
NMI
70.6 #0.4
78.6 #0.1
77.3 +1.7
80.9 #0.5
80.2 +0.7
81.3 #0.8
81.6 =0.7
ARI
62.7 #0.4
71.0 10.1
71.8 43.4
79.0 #0.4
78.4 40.9
79.5 +0.7
79.7 #0.9
where À is leveraged to control the overall correlation strength and set to 0.5. Two tree structures over latent dimensions 1 (V.&,) and T2 (V, E2) are designed with vertex set V 2, 2, 3, 4] and edge sets & = {(1, 2). (1, 3), (2, })} and & = {11, 2), 4), (2,3)}, respectively. The posterior correlation structure modeled by T1, Tz and their mixture model MTrre VI(T..T with the estimated evidence lower bounds and ground truth log- log-likelihood shown in Table 1. It can be seen that tree-structured posteriors can learn more correlation information than traditional mean-fiell 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 samples are believed to belong to the same cluster, and cannot-links, otherwise. Based on the variational deep embedding (VaDE) framework [14], con- strained clustering can be formulated as a probabilistic clustering problem with joint probability Pe( X. Z. c) = pe X Z) p(c))(()) (c), where the sample x, is generated from normal distribution conditioned on z,, zł is sampled from p(z,|c; N(z Me, diag(o?,) )). and the cluster assignments
с = { are sampled from a categorical distribution. Following previous work [24], we in- corporate the clustering preference through a conditional probability P(c[W) with a pairwise prior information matrix W
11, hicc W) W) E. (yc W)
(((TT) II F.. hulc, W)
(19)
where (m k are the weights associated to each cluster, ((T)) is the normalization factor and (4,,, W) = II ززز exp(W< (x,,,,)),,, is a weighting function. The pairwise prior information matrix W defined as a symmetric matrix containing the pairwise constraints: Wij > 0 if there is a must-link constraint between x; and x,. W., O if there is cannot-link constraint between x; and xj, and Wiy = otherwise. The values [Wayl E [0, oo) reflect the degree of certainty in the constraints, and are set to 1 for all datasets. And 6000 pairwise constraints are used for experiments on both our methods and other constrained clustering baselines.
The variational posterior distribution defined as qΦ (Z. c qo (Z) (c Z), where the prob- ability of cluster assignments is factorized as q(c Z) = II, ((c,/z,) which can be easily computed by Bayes theorem. In the work of DC-GMM [24], the posterior distribution go (Z(X) for latent variables is assumed to be mean-field which fails to capture the posterior correlation structure, while our methods TreeVI and MIreeVI, the latent posterior 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 evaluation metrics. It can be

---

-32
Figure 2: T-SNE visualization of MNIST samples embedded space and the learnt tree structure of our proposed VIe 100 samples are randomly selected to plot their instance-level tree structure (colored in black).
easily observed that our models reach the state-of-thh-ar clustering performance in all metrics and datasets, benefiting from the correlation structures captured by our designed correlated posteriors. Moreover, due high efficiency of tree-structured reparameterization, our methods have computa- tional 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 seconds on GeForce RTX 3090, while our TreeVI I and MTree VI take around 5 and 9 seconds, respectively. Further, we plot the tree learned over the constrained clustering experiment 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.
MNIST dataset, as shown in Figure
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; are binarized bag-of-word vector Xx, 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, randomly split into two halves, leading to synthetic users us and B with the most similar movie preference, and a correlation graph G = (V,E) with nodes V= (uA u ii= ,,..., and edges & {(u : 1.....N
Suppose that the joint distribution of rating data X and the corresponding latent embeddings Z is modeled by (X, Z) = pe X (Z), and the variational posterior distribution modeled by 1s (Z) In the work of CVAE [36], both the prior distribution p(Z) and posterior distribution as () are designed as weighted sums of tree-structured distributions with respect to each maximal acyclic subgraph of the correlation graph G. While in work, the latent representations drawn from a Gaussian distribution of the form (ZZ) = A (Z; O2v; (I2N + XA) & Ip). where the hyper-parameter
10,1) is used to control the overall correlation strength, and the affinity matrix A [a:g] j derived from the correlation graph with a; = if (i.j) E & and aij 0 otherwise. And our latent posterior distribution 4 ( is modeled by tree-structured and mixture-of-tree distributions, respectively.
And our goal is to identify the dual user uf given synthetic user u from held-out set in terms of the latent embedding distance. The user data are implemented a random train/test split with 90/10 ratio, and the synthetic user pairs from the training set are used train all the methods. To evaluate the user matching accuracies, fixed number Aeeaa = 1000 of synthetic user pairs are selected from the test set and for each synthetic user ut (or ),,,, summarize the ranking of its dual user u, (or with) among all other Neval 1 synthetic user candidates in terms of latent embedding distances.

---

Table 3: Synthetic user matching test RR
Table 4: Link prediction test Normalized CRR
Methods
Test RR
Methods
Test NCRR
VAE
0.3498 0.0167
VAE
0.0052 0.0007
CV
0.6608 + 0.0066
GraphSAGE
0.0115 0.0025
CVAEcoorr
0.7129 + 0.0096
CVABind
0.0160 + 0.0004
TreeVI (Ours)
0.7408 = 0.0124
CVAEcorr
0.0171 = 0.0009
MI (Ours)
0.7521 + 0.0101
TreeVI (Ours)
0.0188 + 0.0007
VT (Ours)
0.0203 + 0.0014
In Table we show the average Reciprocal Ranking (RR) for the methods, which demonstrates superiority of our model over both implementations for the baseline method CVAB.
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-word, feature vectors Xz, for each user u, and only products that have been rated least 100 times are kept and users who have rated these products at least once considered. To construct the undirected graph from the single-directional "trust" statements between users W, Uj provided by the dataset, we only build an edge (u,, ug) E & if both u, trusts új and uj trusts of our TreeVI and MTree VI 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-erees distributions, respectively.
The experiment setting
The product rating dataset is split for each user max (1, yu degree(u:)) edges held out for test edge set Eeest. The remaining edges for the training edge set Etrain are used to train all methods on the product rating data. To evaluate the link prediction accuracies, we calculate for each user u, the Normalized Cumulative Reciprocal Rank NCRR, of the ratings of u,'s 's test edges among all possible connections except for the training edges, in terms of latent embedding distance metries. Formally, the NCRR value is the [o, 1] n)-oomaaiization of the Cumulative Reciprocal Rank (CRR) formulated as CRR, Σ ((,,,, K (u, ux) &
V into training and test sets, with
Etrain dik I diyH1'' where d, represents the latent embedding distance between users u, and uj for 1 jj S N. Larger NCRR, indicates better ability predict held-out test links with respect to each user u4, and the averaged results for all methods are reported Table 4.
Conclusion
In this work, we present novel variational inference method called TreeVI, that approximates the intractable latent posterior distribution with 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 MTree VI with a mixture trees better approximate the underlying posterior. Furthermore, the tree and mixture-of of-rrress--- structures are allowed stochastically learned from data by solving 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 tree structure to approximate the graph-structure posterior correlation structure. This limitation mitigated by adopting a weighted mixture of trees and stochastically learning correlation-rich tree or mixture-of-treed 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 supported by the National Natural Science Foundation of China (No. 62276280, U1811264), Guangzhou Science and Technology Planning Project (No. 2024A0419967).
10

---

References
[1] Agrawal, A. and Domke, (2021). Amortized variational inference for simple hicrarchical models. In Neural Information Processing Systems.
[2] Basu, S., Banerjee, A.. and Mooney, R. J. (2004). Active semi-supervision for pairwise con- strained clustering. In Proceedings of the 2004 SIAM international conference on data mining.
[3] 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.
[4] Blei, D. M., Kucukelbir, A., and McAuliffe, D. (2016). Variational inference: A review for statisticians. Journal of the American Statistical Association, 112:859 -777
[5] 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.
[6] Caterini, A., Cornish, R., Sejdinovic, D., and Doucet, A. (2021). Variational inference with continuously-idexxel normalizing flows. In Uncertainty in Artificial Intelligence.
[7] Coates, A., Ng, A., and Lee, (2011). An analysis of single-layer networks in unsupervised feature learning. In International Conference on Artificial Intelligence and Statistics.
[8] Hamilton, W. L., Ying, Z., and Leskovec, J. (2017). Inductive representation learning on large graphs. In Neural Information Processing Systems.
[9] Harper, M., Konstan, A., and A., J. (2016). The movielens datasets: History and context. ACM Trans. Interact. Intell. Syst.
[10] He, K., Zhang, X., Ren, S., and Sun, J. (2015). Deep residual learning for image recognition. 2016 IEEE Conference on Computer Vision and Puttern Recognition (CVPR), pages 770-778.
[11] Hoffman, M. D. and Blei, D. M. (2014). Stochastic structured variational inference. In International Conference on Artificial Intelligence and Statistics.
[12] Huszár, F. (2017), Variational inference using implicit distributions. ArXiv, abs/ 702.08825.
[13] Janke, T., Ghanmi, M., and Steinke, F. (2021). Implicit generative copulas. In Neural Informa- tion Processing Systems.
[14] 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 Artificial Intelligence.
[15] Jiang, Z.,, Zheng, Y., Tan, H., Tang, B., and Zhou, H. (2016b). Variational deep embedding: An unsupervised and generative approach clustering. In International Joint Conference on Artificial Intelligence.
[16] Kim, M. and Pavlovic, V. (2020). Recursive inference for variational autoencoders. Advances Neural Information Processing Systems.
[17] Kingma, D. P. and Ba, J. (2014). Adam: method for stochastic optimization. CoRR. abs/1412.6980.
[18] Kingma, D. P. and Welling, M. (2013). Auto-encoding variational bayes. CoRR.
[19] Kipf, T. and Welling, M. (2016). Variational graph auto-encoders. ArXiv, abs/1611.07300 1.07308.
[20] Lambert, M.. Chewi, S., Bach, F., Bonnabel, S., and Rigollet, P. (2022). Variational inference
via wasserstein gradient flows. Advances in Neural Information Processing Systems.
[21] LeCun, Y. and Cortes, C. (2010). MNIST handwritten digit database.
[22] Lewis, D. D., Yang, Y., Russell-Rose, T., and Li, F. (2004). Rev1: A new benchmark collection for text categorization research. Journal of machine learning research, pages 361-397.

---

[23] Lin, W., Hubacher, N., and Khan. M. E. (2018). Variational message passing with structured inference networks. arXiv: Machine Learning.
[24] Manduchi, L., Chin- ChinnCheong, K., Michel, H., Wellmann, S., and Vogt, J. (2021). Deep con- ditional gaussian mixture model for constrained clustering. Advances in Neural Information Processing Systems.
[25] Manduchi, L., Vandenhirtz, M., Ryser, A., and Vogt, J. (2023). Tree variational autoencoders. Advances in Neural Information Processing Systems.
[26] Moens, V., Ren, H., Maraval, A., Tutunov, R., Wang. J., and Ammar, H. (2021). Efficient semi-implicit variational inference. CoRR, abs/2101.06070.
[27] 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.
[28] Nemirovski, A. (1999). Optimization II: Standard Numerical Methods for Nonlinear Continuous Optimization.
[29] Ou, Z., Su, Q., 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/IJNNPP (I), pages 2238-2249.
[30] Papamakarios, G., Pavlakou, T., and Murray, I. (2017). Masked autoregressive flow for density estimation. Advances in neural information processing systems, 30.
[31] Pearce, M., Chiappa, S., and Paquet, U. (2018). Comparing interpretable inference models for videos of physical motion. In 1st symposium on advances approximate bayesian inference.
[32] Ren, Y., Hu, K., Dai, X., Pan, L., Hoi, S. C. H., and Xu, (2019). Semi-supervised deep embedded clustering. Neurocomputinng, pages 121-130.
[33] Smith, M. S. and Loaiza-Maya, R. (2021). Implicit copula variational inference. Journal of Computational and Graphical Statistics.
[34] Smith, M. S., Loaiza-Maya, R., and Nott, D. (2019). High-dimensional copula variational approximation through transformation. Journal of Computational and Graphical Statistics.
[35] Song, J., Meng, C., and Ermon, S. (2021). Denoising diffusion implicit models. In 9th International Conference on Learning Representations.
[36] Tang, D., Liang, D., Jebara, T., and Ruozzi, N. (2019). Correlated variational auto-encoders. In International Conference on Machine Learning.
[37] Tölle, M. and Schlaefer, A. (2021). A mean-field variational inference approach deep image prior for inverse problems in medical imaging. In International Conference on Medical Imaging with Deep Learning.
[38] Wang, Q and Van Hoof, H. (2020). Doubly stochastic variational inference for neural processes with hierarchical latent variables. In International Conference on Machine Learning.
[39] 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.
[40] Xiao, H., Rasul, K., and Vollgraf, R. (2017). Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. ArXiv, abs/1708.07747.
[41] Xie, J., Girshick, B., and Farhadi, A. (2016). Unsupervised deep embedding for clustering analysis. In ICML.
[42] Yang, L.,, Cheung, N.-M., Li, J., and Fang, J. (2019). Deep clustering by gaussian mixture variational autoencoders with graph embedding. In Proceedings the IEEE/CV international conference computer vision.
12

---

[43] Yin. P., Zhou, C., He, J., and Neubig, G. (2018). Structvae: Tree-structured latent variable models for supe-supervised semantic parsing. ArXiv, 806.080...
[44] Zhang, H., Basu, S., and Davidson, (2019). A framework for deep constrained clustering algorithms and advances. In ECML/PKDD.
[45] Zheng, X., Aragmm, B., Ravikumar, P., and Xing, E. P. (2018). Dags with tears: Continuous optimization for structure learning. In Neural Information Processing Systems.
13

---

A Precision Matrix of Tree-structured Posterior
When latent variables Z = 11.-.. .2N] are all Gaussian, it is a well-known consequence of the Hammerrley--Clffiiee theorem that the entries of the precision matrix Pz Σx correspond to rescaled conditional correlations. Denote Pz
with Pij E. then the magnitude of Pij expresses the correlation of z; and 2j conditioned on other latent variables. In particular, the sparsity pattern Pij reflects the edge structure the correlation graph: Pij Op and only if 2; II. 3jj "A,,,)..,,,
[diag(Pr))l..=....
The tree structure T = (V,E) with non-adjacent correlation parameters defined as Eq. (3) forms Markov random field with joint probability given by Eq- (4), which follows the principle of conditional independence. Therefore, the precision vector Pij f Op if and only if latent variables z; and 2j are adjacent, that is, (i,j) E
Example. the example 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
E diag Oxx
R (
diag((oz))
(20)
with the inverse of tree-structured correlation matrix given by
S11
$12
812
822
823
824
R
111111
$23
$33
835
(21)
S24
841
$35
$65
demonstrating that only correlation parameters along the edge set = (11, (2,3), (2, (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 r {712. Y23, 724 3s
$11 (1D 712)-
(22)
$12 Y12 (1D 712)
2
(23)
$22 = (1D 2v 0 723
24
2
712
① 723
2
712 2 Y
2
224
Y23
O 72,
(24)
@ (1D - (1D 22
-1 (10 724) 7-2)-1
$23 = -2 0 (1p
Y23 72))
(25)
$24 = Y24 0 (1p 224
(26)
833 = (1D 2 0
355
2
② (1p
2 -1 0 (1, -335) -1
(27)
= Yas 0 (1p 335
(28)
(1n Y21
2
(29)
$55 (1) 7/3 -1
(30)
B Ancestral Sampling for Tree-structured Posterior
B.1 Proof of Theorem 1
Suppose that N latent variables Z = [2.... ZN] follow a srre-structured posterior distribution
given the tree structure T (V, with V = {1, ... N), the ancestral sampling for each latent variable zy. j V can be expressed by
2, = M4 71 1600 +
Σ IEPI-j.A
9u 2 VID p OG00, j jEV,
(31)
where E; O0 In) a randomly sampled noise from the standard normal distribution and pa(i) denotes the parent node of with respect to i V.
14

---

Proof. For each latent variable 2,, there exist a unique path P1-35 = (V1, U2, ..) along the tree-structured Bayesian network that starts from the root node 21 1 and ends the goal node Vn = j. The joint distribution along the path Pi is given by
q (Zer Zuzs Zu..)
(z II 9s (Zectz IZv.).
(32)
The ancestral sampling starts by sampling from do (Z., ): Ze, Mii + es O Cer where Ev, ~ (OD, Ib). Given the observation of Zuy then the latent variable Zuz can be sampled using the conditional distribution Eq. (7):
(Zon -Mo) + VID Yo,
Zuz 14g Yores
Cuy C deg
(33)
M2 Yor
Co o + 1D Po Cuz az.
(34)
Assume that at the k-th step of the ancestral sampling, we have
Zuz = Mo + Feve Eor O ou + Σ Jov VIL
1D - 11--2 Ev. o
(35)
where Yorn 7 70
o Yo the for and Youm 1p. Then the (k + 1)-th step of the ancestral sampling, the latent variable Zez+a can be similarly sampled by conditional sampling
(Zer Mug + VID - 7 Yogue+
Zozt1 Murai + Voguett
uk+1
Evp+1
(36)
8++
Muui Yornis 2 Cp
① 0x-1
Σ
Yornn
1n 0
22 in
2 Cu.
① 0x+1
(37)
1=2
Therefore by induction, the ancestral sampling for the latent variable zj can be expressed as
Zu, Mv, Verv, Ev O B.
Juny 2 11 P2. 3 Evz C Bays
(38)
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 2 as the root node, and the detailed ancestral sampling procedure for the tree-structured Bayesian network Eq. (6) is shown follows.
1. Start by sampling from 9 (z)):
21 = 14 + 0 O 61,
(39)
where E ~ (O,,,
2. Then sample from G ( (2)) given the observation of Z
Z2 H2 + 712 2
0
1D
722
02 0 E2.
(40)
where E2 NVO0, Ip);
15

---

3. Then sample from (0(zz) z2) given the observation of Zg:
23 13 + 723 0
712 61 + VID -2 062 + 1D 2
03
0.
= =13 + 712 2 Y23 ① 0g 0 72 O VID -712 0
03 62
+ +V1D 1D 723 3
(41)
where 63 ~ (O)
4. Then sample from 90 (Z) given the observation of Zz:
24 = u4 + 712 2 724 2 04 61
Y 2 v 1n 722 0.
04
(42)
+VID 1-0-221-0 00 0444
where E4 ~ Ip):
5. Finally, sample from g (2z 23) given the observation of Zz:
25 = 143 + 712 ! Y23 @ Y35
0g
①
61
+ 723 2 Ys5 2 V
1D
712
2
6
(43)
+ Ya5
ID 123
+ 1 33 0 66
VI
where 65 ~ Io)
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,8), the evidence lower bound of our proposed TreeVI is given by
CT (0. X)
log pe X, (T) (Z|X),
(44)
where (T) denotes the tree-structured reparameterization for latent variables. The first term can be directly computed by
N
log Pa X, (T))
Σ
log Pa (x;(((()) + log p(z.") (T).
(45)
where x; (T) is the reparameterization for latent variable z;, i = 1,... N. And the entropy of the tree- structured posterior q3 (Z can be factorized as entropy terms with respect to singleton posterior q (x) and pairwise posterior qs (z,, 2j|X;, Xi, x,)
M[g T (Z[))
Σ ivv
(2, + Σ
Hlga (z, (2,, 2,1x,, - H Higs (z, xx)) - Hlgs (z. [x]]] (K))E
(46)
16

---

where the entropy of singleton posterior 96 (z/(x;) = (zz M4, diag(o?)) is given by
H- (2: (2][x]]] = - Ez mo (z, (x.)
log 1s Z [x)x]
(47)
)]
--2a Ez no (m.
log (2x))
exp
(z; 14) E (zi m4)
(48)
D 21 + [oog(2))] log |
(49)
D
[1 + log((T)]
+ 2
Σ
log
(50)
2
d=1
where E, diag(o?) and the step relies several properties of the trace operator:
Ezz [(=: m4)) E (z; m.]
Fr [tr (z: m1;) Σ (z+ [s))
(51)
1lz, tr [>; (z. 14) T (z. [s)) tr (E Ezz [zz m4) (z: 14)))
(52)
(53)
tr tr(E, tr(In) D.
(54)
And the entropy of pairwise posterior [Po My]" and covariance matrix
Xj)
NN Mij Eij) with mean Ptj
0:00: Yij 00,0,
Yij 3
Σij =
(55)
similarly given by H l (zi, uz² xs, x,)] = - En [log ue (ni, 2ij [x/1x; x;]]
(56)
1 (2 Hij) Enj (zi - mk3))
= -E2, log (2x)) -1- | exp
(57)
D + Dlog(2T)) + log Eiyl
(58)
+ Dlog(2m) + [lo((1 Faya) + log oid + log oja]
(59)
where we denote Zij = [z; 2,1T.
C.2 Evidence Lower Bound for MTreeVI
For mixture-of -otterrr------ess trees structured posterior approximation with respect the mixture of tree compo- nents MT = (Ti, ... TTs) }. the evidence lower bound of our proposed MTTreeVI is given by
M
MT x. =
Σ Tm log XX, Z (Tm) + + [Z|x)
(60)
where ( (T.) denotes the tree-structured reparameterization for latent variables with respect to the m-th tree component Tm, m = 1,... M. The first weighted summations can be directly computed by
Tmm log Po (X 2 (Tm))
Σ Tm Σ log po
(x,|2, (Tm)) + log p(z(TT (Tm)).
(61)
where 2 (Tm) the reparameterization for latent variable 2;, 1, .N with respect to the m-th tree component, m 1,... .M. And the weighted summations of the entropy terms can be similarly
17

---

factorized and calculated as Appendix C.I by entropy terms with respect to singleton posteriors 9 xx) and pairwise posteriors qΦ (z Z,) X,, Xy).).
D Constrained Optimization Details
D.1 Proof of Proposition
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 E 10,11NxN and r rAA) < 1, then A acyclic if and only if
tr A(Ix - A 0.
(62)
2. Suppose that A E (0, 1]NNN, then A is acyclic if and only if
tr A exxp(A²) = 0.
(63)
3. Suppose that A E NNN then is acyclic if and only if
(r A) e O A) = =0
(64)
where tr(:) and exp(-) represent the trace and exponential matrix respectively, and @ is the Hadamard product.
Proof. The proof relies on the fact that tr A counts the number of length-k closed walks in a directed graph.
1. Clearly the directed graph induced by the symmetric matrix A will only have self-loops, and hence tr A 2k++1 0 for all k 1, co. In other words, A has cycles if and only
if f(A) = 1 Di 2b++1 )as 0, then
N
tr A A3)-1 = tr
2k+1
A 2k+ 0.
(65)
k=1
The desired result follows.
2. Similar to the proof of Proposition 1.1 by noting that A has cycles if and only if (A 26+1) = 0 for all 2 and all i (1,--. NN. which is true if and only
N
tr A exp(A2) tr
26+1
2k+ ): O.
(66)
k=
3. The proof is similar to Proposition 1.2 by replacing A with c A, which counts weighted closed walks.
D.2 Greedy Searched Initialization
The constrained optimization of tree structure requires 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 first-search (DFS) algorithm. Algorithm shows the DFS algorithm for 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 A(i) denotes the neighbors of node i.
18

---

Algorithm 1 DFS Algorithm for Tree Generation
Input: Fully-connected graph G; number of trees M
Output: Edges list of generated trees E
1: procedure L'REEGGENERATION(M.
Input: fttree M
2: E
Initial edges list
3:
for k = 0,.... M - 1 do
4:
V = [[aasss [M]
Visited node list
5:
while False in V do
6:
[V[V=== Palaee]
Choose node
7:
Initial queue
8:
while len( Q) > do
9:
Q00
10:
i True
N ID Faalel
12:
if
(en(NN == 0 then
13:
POP(Q,-1)
14:
break
15:
end if
16:
j f [N
D Choose neighbor
17:
Vij + True APPEND(Q.)) APPEND(E,
18:
19:
20:
end while end while
21:
22: end for
23:
end procedure
To enrich the constructed spanning tree with neighboring correlation information, the uniform sampling process (line 16 in Algorithm 1) in the DFS algorithm can further modified generate meaningful neighborhood for each data instance, by assigning the probability of sampling neighbor of instance i
exp(cos(x] x))aa
(67)
EKENG exp(cos(x x))aa)
where a is the temperature parameter controlling the trade-off between the precision and diver- sity of edges, and we find the best configuration of the validation set with the values in
In Table 5, we show the constrained clustering accuracies the MNIST dataset, with different initializations, and with or without our constrained optimization. It can be seen that without con- strained optimization, both initializations underperform. And with constrained optimization, both initializations converge to substantially better performances, with the searched initialization slightly better. Overall, the constrained optimization procedure makes our proposed TreeVI less sensitive to initializations.
Table Constrained clustering accuracies (%) on the MNIST dataset, with random tree or greedy- searched tree initializations, with or without our constrained optimization
Random Tree
Greedy Search
w/o CO
96.58
96.70
w/ CO
97.31
97.45
19

---

D.3 Solving Constrained Optimization
In section 2.3, we establish continuous characterization of acyclicity, leading the following equality-constrained program (ECP)
min
E(A)
A
(68)
subject to h(A) O,
Solving the constrained optimization requires classical techniques from the mathematical optimization literature. However, this nonconvex program since (A ((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: (t) converting the constrained problem into a series of unconstrained subproblems, (ii) optimizing the unconstrained subproblems, and (iii) thresholding. The full algorithm is outlined Algorithm
Algorithm/ 2 Algorithm for Constrained Optimization
Input: Initial guess (Ao, ao); progress rate E (0, tolerance > 0; threshold >> 0 Output: Threshold matrix
I: for = 0, 2, do
2: Solve A4t1 arg minA L"(A, a4) with p such that h(A4+1) ch(At)
3: 0++1 t- + ph (A+1)
Dual ascent
4:
h(A+1] < c then
5:
A f- A
6:
break
7: end
8: end for
9: A - A 1
Converting constrained optimization into unconstrained subproblems. The augmented La- grangian method [28] can be leveraged to solve the equality-constrained program Eq. (68) by first augmenting the original problem with a quadratic penalty:
min
{(A) + h(A)
AERNXN
(69)
subject to h(A)= = 0,
with a penalty parameter / > O, 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 a
D(a)= min AERN
L'(A,
(70)
where TP(A, a) ((A) P [(A)))2 ah(A)
(71)
the augmented Lagrangian. And the goal is find local solution the dual problem
max D(a).
(72)
aER
Let A; be the local minimizer of the Lagrangian (70) at a, i.e., D(a) = LA o a). Since the dual objective D(a) is linear in with the derivative simply given by VD(a) hh(A,, ), one perform dual gradient ascent to optimize the dual problem (72):
a a ph(A
(73)
Solving the unconstrained subproblem. The augmented Lagrangian converts the constrained problem (69) into series of subproblems (70), and our goal is to solve these subproblems efficiently, Let (ec(A) E R², with 2 The unconstrained subproblem (70) can be considered
20

---

typical minimization problem over real vectors:
min f(a);
(74)
aER
where f(a) = ((A)+ P [b()) ah(A)
(75)
a smooth objective, for which a number of efficient numerical algorithms are available, such as L-BFGS [5]. In our experiments, the values of can be learned by a neural network that encodes pairwise correlation.
Thresholding. Motivated by post-process prospesccssing estimates of coefficients via hard thresholding, threshold the edge weights follows: after obtaining a stationary point A of (69) given fixed threshold w > 0, set any weights smaller than w 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 ((() A) S for some all tolerance (set to = 10-8 in our experiments) instead of ((A) = strictly due to numerical precisions. However, small threshold w suffices to rule out cycle- sycle-inducing edges since h( A) explicity quantifies the acyclicity of A. Following [45], fixed value of threshold w 0.3 is set 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 Movic- Lens, a movic recommendation service. contains 20,000,263 ratings and 465,564 tag applications across 27, 7,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 to 5, and also express their Web of Trust by issuing trust statements. It consists 49,290 users who rated a total of 1 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 dineemnnnio vector [21].
Fashion MNIST: dataset of Zalando's article images consisting a training set of 60,000
examples and test set of 10,000 examples [40].
Reuters: It contains 8 810,000 English news stories [22]. Following the work of [41], we used root categories: corporate/industrial,, government/social,, markets, and economics as labels and discarded all documents with multiple labels, which results in 685,071 -article dataset. We computed if-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 10 classes with 13,000 examples each [7]. As pre-processing, we extracted features from the STL-10 image dataset using ResNet-50 [10], as in previous works [15].
E.2 Implementation Details
Synthetic Data. For synthetic VAFs, we apply a two-layer feed-forward neural network for the generative model (x; 24) and two-layer feed-forward neural network for the variational posterior approximation g (o(2((x) with each 2; E TR. where D = 4. The traditional mean- field approximation assumes the posterior distribution for each latent be fully factorized: g (z. xx) = I yo (zaa(a))))) while our methods the posterior assumed to be tree-structured and miixtuee-of-tree structured, which can be factorized into singleton posterior q (zsal (raa) and pairwise posterior 4 (Znd, Ziel Fed, Nie) both assumed to be Gaussian. Each edge in the tree structure corresponds to correlation parameter "de to be learned in the bi-variat normal distribution, which
g
21

---

Table Hyperparameters setting of constrained clustering task.
MNIST
IMNIST
Reuters
STL-10
Batch size
256
256
256
256
Epchs
1000
500
500
500
Learning rate
0.001
0.001
0.001
0.001
Decay
0.9
0.9
0.9
0.9
Epochs decay
20
20
20
20
shared for all data points to capture dimension-leved correlation. Our synthetic VAE is trained for 200 epochs for all methods. We apply stochastic gradient optimizations with step size of 0.005, and use the Adam algorithm [17] adjust the learning rates.
Constrained Clustering. To implement our methods, we careful in maintaining fair compari- son with the baseline methods. In particular, we adopt the same encoder and decoder ferwaforward 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 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 must-link if they have the same label and a cannot-link otherwise. The absolute values the elements [Way in the pairwise prior information matrix are set to 10 for all datasets for convenience, and 6000 pairwise constraints sampled for training both methods and other baselines. Following DC-GMM [24], the hyper-pate hyper-parameters 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 decay rate of 0.9. The number of tree components adopted in our MTTreeV VT is set to 3.
is used for training and reported results
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 VAE-based methods (including TreeVI and MTTree VT), we apply a two-layer feed-forward neural network for the generative model (x./)/z; and a two-layer feed-forward neural network for the posterior approximation (s(2, (x;), The model likelihood function (X Z) is a multinomial distribution, and the singleton posterior distributions 10 (z; xx) are all diagonal normal distributions. For CVAEcorr and our methods, we also learn a two-layer feed-forward neural network that takes the concatenation [x;; x,] as input and output the correlation parameter between 2; and zj on each of the D dimensions. To determine the tree structure for VI and MIr we take average of these correlation parameters across all D dimensions and obtain correlation matrix A NXN or set of correlation matrices that is used for constrained optimization. For AGE, we choose to use K 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 adjust the learning rates. The number of tree components adopted in MTrre VI is set to for both tasks.
E.3 Further Experiments
Synthetic Data. We design synthetic dataset with a graph-structured latent variable model. The dataset contains N' 6000 data points "I,... XN E DRD with D = 4, each independently generated from the conditional distribution p(x[z) N((x Oz, o',) given latent embeddings 21,- nn RD where a2 is fixed value and set to 0.5. The latent embeddings .., 2N are drawn from standard normal distribution p(z) = (2; O v L)), and the graph-structured correlation
22

---

Table 7: Estimated lower bounds (ELBO) of VAE with posterior distributions approximation by mean- field, tree-structured, and mixture-of-treet distributions, compared to ground truth log- ool-iiklihood log (X)..
Methods
Lower Bound
Mean-field
-9.022
TreeVI Tree)
-8.8838
VT (2 Trees)
-8.8729
MTree VT (3 Trees)
-8.8690
log p(X)
-8.8142
Table 8: Clustering with constructed affinity matrices performances (%) of our proposed methods TT and VI compared with baselines
MNIST
IMNIST
Reuters
STL-10
ACC
NMI
ARI
ACC
NMI
ARI
ACC
NMI
ARI
ACC
NMI
ARI
VADE
92.00
90.80
82.90
63.20
62.90
48.30
83.35
63.10
67.09
84.45
79.10
73.98
DEC
87.93
80.71
78.35
60.61
63.70
47.41
76.33
51.91
57.31
78.56
72.47
63.91
DCN
83.10
82.50
75.60
53.10
57.10
38.60
76.60
61.10
65.40
78.40
65.50
60.70
DGG
97.58
92.50
93.40
81.33
73.20
67.90
93.50
81.23
87.85
90.59
81.58
79.80
DC-GMM
95.28
90.01
90).5
77.24
70.05
65.73
92.38
80.68
88.05
88.70
78.90
77.50
TrevV
96.30
92.21
92.82
78.90
72.08
66.39
95.12
82.40
89.02
90.22
80.97
79.30
MTrevVI
97.66
92.96
93.52
81.57
73.61
68.20
95.31
82.64
89.20
90.50
81.22
79.62
incorporated into the sample generation via a coefficient matrix
2
0
07
2
(76)
0
1
0.3
0.3
0.4
2
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- likelikeehood is presented in Table 7. It can be observed that tree-structured poterior benefits from capturing more pairwise correlations, and the mixture of more components better approximates the ground truth log-likelihood.
Clustering with Constructed Affinity Matrix. Besides the weak supervision of pairwise con- straints, 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 value between [0, reflecting the similarity between sample pair. construct the affinity matrix W, find set of nearest neighbors for given data point and compute their similarity using predefined kernel function, such Gaussian kernel
exp
|xx:
1xx-xx1
x
if E (Nx,)
W ij
25
(77)
0,
otherwise
where s; a predefined scalar, Nxx;) denotes the set consisting of the nearest N, neighbors of x;- 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
(, W) = Il exp Wejdere + (1 Wsy) (1
(78)
jzi
The clustering performances of our proposed methods against baseline methods are shown in Table 8.
23

---

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

---

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 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.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
Limitations our work has been discussed in Section 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 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 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
25

---

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.I, 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 cross- referenced.
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 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 ex- perimental results of the paper to the extent that 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 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 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.
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 a specific model and empirical evaluation, 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 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 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 submis- sions 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 reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which 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 have some path reproducing verifying the results.
5. Open access to data and code
26

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: The datasets and experimental settings 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 NeurlPS code and data submission guidelines (https://nips.co
hetrbbmmsssionoll1111oem oolllycy) for more details.
While encourage the release of code and data, 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 new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / /ttpss.ccpppliiiiim wwtppccppcc////ccccccccccccccclt http//uudeeee/eeee https//ooSSSuuiiseeeeeeee
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 proposed method and baselines. only 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 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, hyper- parameters, 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 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, as supplemental material.
7. Experiment Statistical Significance
Does the report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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 emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
27

---

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 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 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 figures symmetric error bars that would yield results that are out 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 com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
The computer resources needed to reproduce our experiments have been specified in Appendix E.4.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https :ttp::://neuuse hwtpp///cc/cccccccciieeeeeeeeeee
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 reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal 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 no societal impact to be addressed.
28

---

Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., denloymemt 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 being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have 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 variational inference without safety risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.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 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 dataset.
29

---

The authors should state which version of the asset is used and, if possible, include 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 souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi whwode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 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 release new assets.
Researchers should communicate the details the datasct/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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create 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 sereenshots, 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 contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved 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 crowdsoureing research with human subjects
Guidelines:
30

---

The answer NA means that the paper does not involve crowdsoureing 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 expect authors adhere 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.
31


Reviews:
[{"review_id": "DnAHp2Kdn9", "reviewer": "Reviewer_ptS8", "summary": "The paper proposes to perform amortized variational inference, where the variational approximation has a tree-dependence structure across the instance-level latent variables. For this, the scale is decomposed into variance and correlation, where both are amortized through a neural network. A non-convex constrained optimization procedure is proposed to also learn the correlation structure, which ensures that the correlation structure forms an acyclic graph.", "questions": "n/a", "limitations": "n/a", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* In variational inference with structured variational families, learning the appropriate correlation structure from data is an important problem that hasn't been fully addressed.", "weaknesses": "* 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your detailed comments, but we believe that you misunderstood our paper deeply. We hope our clarifications could help you recognize our contributions correctly.\n\nYou mistakenly think that our main contribution is simply to Cholesky decompose the correlation matrix $\\mathbf{R}$ as $\\mathbf{R}=\\mathbf{L}\\mathbf{L}^\\top$ and then re-parameterize the elements $\\ell_{ij}$ in the lower-triangular matrix $\\mathbf{L}$ by a neural network $f_\\phi(\\cdot)$, that is, $\\ell_{ij}=f _ \\phi(\\mathbf{x} _ i, \\mathbf{x} _ j)$, with $\\mathbf{x} _ i$ denoting the $i$-th data instance. We want to emphasize that this is not our contribution. If this approach is adopted, we have to re-parameterize as many as $N(N+1)/2$ elements, i.e., all elements from the lower-triangular positions of $\\mathbf{L}$ need to be re-parameterized, where $N$ is the number of instances in training dataset. That means we need to run the neural network $f_\\phi(\\cdot, \\cdot)$ by $O(N^2)$ times for every epoch, which is computationally unacceptable, especially considering that $N$ could be as large as **millions** in practice.\n\n\n**Our main contribution lies at finding a way to reduce the required times of running the neural network $f_\\phi(\\cdot, \\cdot)$ from $O(N^2)$ to $O(N)$**, making it applicable to scenarios even with millions of instances. To achieve this goal, we propose to restrict the correlation matrix $\\mathbf{R}$ to a special form constructed from a tree $\\mathbb{T}(\\mathcal{V}, \\mathcal{E})$. It should be pointed out that although $\\mathbf{R}$ is constructed from a tree, it doesn't represent a tree. Actually, it is a dense matrix, with its $(i,j)$-th element $[\\mathbf{R}] _ {ij}$ for $(i,j)\\notin \\mathcal{E}$ set as $\\prod_{(s,t)\\in \\mathbb{P} _ {i\\to j}}\\gamma_{st}\\triangleq \\tilde \\gamma_{ij}$. \n\nWe prove that under the restricted correlation matrix $\\mathbf{R}$, the lower-triangular matrix $\\mathbf{L}$ possesses a very elegant form, as shown in Eq. (11) in the paper. The elegance lies at that although $\\mathbf{L}$ still has $O(N^2)$ non-zero elements, all of these non-zero elements can be explicitly computed from the $|\\mathcal{E}|$ parameters $\\\\{\\gamma_{ij}\\\\} _ {(i,j)\\in \\mathcal{E}}$. For instance, the $(1, 4)$-th element in Eq.(11) can be computed as $\\prod_{(s,t)\\in \\mathbb{P} _ {1\\to 4}}\\gamma_{st}$, where $\\mathbb{P} _ {1\\to 4}$ denotes the path conneting node 1 and 4 in the tree ${\\mathbb{T}}$. The importance of this conclusion is that now we only need to re-parameterize the parameters $\\\\{\\gamma_{ij}\\\\} _ {(i,j)\\in \\mathcal{E}}$, while computing the other non-zero elements in $\\mathbf{L}$ with $\\\\{\\gamma_{ij}\\\\} _ {(i,j)\\in \\mathcal{E}}$. In this way, we only need to re-parameterize $|\\mathcal{E}|$ parameters, and thus only need to run $|\\mathcal{E}|$ times of neural network $f_\\phi(\\cdot, \\cdot)$, instead of $O(N^2)$ times in the vanilla method. For a tree, the number of edges $|\\mathcal{E}|\\le N -1$, thus we only need to additionally run $O(N)$ times of neural networks $f_\\phi(\\cdot, \\cdot)$, in addition to the runs required by the mean-field method. \n\nFor a mean-field variational inference that assumes independence among instances, for each epoch, it also need to run $O(N)$ times of neural network. Thus, the complexity of our proposed method is roughly only 2 times of the mean-field method. This is consistent with the elapsed time observed in the experiments.\n\n\n\n\nBelow we respond to each point raised by the reviewer:\n\n**Q1: A lot of the derivation in Section 2.1 is redundant. The authors could have simply parameterized the Cholesky factors.**\n\nA1: As discussed above, directly parameterizing the Cholesky factors is computationally consuming, which requires to run the neural network $f_\\phi(\\cdot, \\cdot)$ by $O(N^2)$ times. Section 2.1 is mainly devoted to how to achieve a correlation-capturing variational inference method that only requires running the neural network by $O(N)$ times.\n\n\n**Q2: Why to take a tree structure among various different correlation graph structures.**\n\nA2: As elaborated above, by restricting the correlation matrix ${\\mathbf{R}}$ to a form that is constructed from a tree, we can prove that all of the non-zero elements in ${\\mathbf{L}}$ can be explicitly computed from the $|{\\mathcal{E}}|$ parameters $\\gamma_{ij}$ for $(i,j)\\in {\\mathcal{E}}$. Then, we only need to re-parameterize the $|{\\mathcal{E}}|$ parameters $\\gamma_{ij}$, instead of all the non-zero elements in ${\\mathbf{L}}$. That's why we have to restrict the correlation matrix ${\\mathbf{R}}$ to be constructed from a tree. Again, we want to point out that ${\\mathbf{R}}$ is only constructed from a tree, but it is a dense matrix.\n\n\n**Q3: What the constrained optimization problem in Eq 17 is supposed to do.**\n\nA3: As discussed above, we need to construct a tree to learn the instance-level correlations. Besides trivial methods like heuristic methods and greedy search algorithm, it can also be learnt from data by our constrained optimization. Specifically, we use a real matrix to represent the correlation parameters, which satisfies the acyclicity constraint if and only if the indicator function $h(\\mathbf{A}) = 0$ as shown in Eq. 17. This allows us to formulate a constrained optimization problem with a smooth equality constraint, which can be solved with Lagrangian multipliers by incorporating a penalty term for the indicator function.\n\n**Q4: Imprecise about amortized variational inference not just variational inference.**\n\nA4: We apologize for the impreciseness brought by our terminologies, and we will reflect this in our title of the final version.\n\n**Q5: Sticking-the-landing gradient and MCMC-based augmentation methods.**\n\nA5: We thank the reviewer for pointing out some feasible alternatives. However, the MCMC-based augmentation methods may cause inefficiency by introducing time-consuming sampling procedure.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4vbjOIHfou", "reviewer": "Reviewer_BYBL", "summary": "This paper proposes Tree-structured Variational Inference (TreeVI), a novel method for capturing instance-level correlations in the posterior distribution of latent variables. TreeVI represents correlations between latent variables using a tree structure. This enables reparameterization of latent variables using ancestral sampling, allowing for efficient training. Furthermore, the authors propose Multiple Tree-structured Variational Inference (MTreeVI) to represent more complex correlation structures using multiple tree structures. They also propose a method to learn the tree structure from data. The proposed methods demonstrate superior performance compared to existing approaches on various tasks using both synthetic and real-world datasets.", "questions": "On line 78, $\\mathcal{V} = \\{ 1, 2, \\cdots , N \\}$ is defined. However, in the Synthetic VAE experiment in Section 4.1, despite $N=6000$, $\\mathcal{V}= \\{ 1,2,3,4 \\}$ is used. Could you explain this discrepancy?", "limitations": "The limitations are appropriately addressed within the paper, and I have no additional suggestions in this regard.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* 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.", "weaknesses": "In the experiments on real-world data, there is insufficient explanation of what the tree structure represents. The tree structure is used merely as a means to relax the independence assumption of the approximate posterior distribution while maintaining training efficiency. The study could have been enhanced by providing some insights derived from the tree structures obtained through training.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and suggestions. Below we respond to each point raised by the reviewer:\n\n**Q1: The study could be enhanced by providing some insights derived from the tree structures obtained through training.**\n\nA1: Thanks for the suggestions. Our method is mainly used to capture the inherent correlations among different instances. Thus, the learned trees are inclined to reflect the correlation among different instances. To see this, we plot the tree learned over the constrained clustering experiment on MNIST dataset, which is shown in the Figure 1 in the attachment of global rebuttal. From the figure, we can see that instances from the same category are connected more tightly than those from different categories in the learned tree. This demonstrates that our method has the abilities to capture the underlying inherent correlations among different instances. The figure and the explanation to it will be included in our final paper.\n\n\n\n**Q2: Discrepancy of notations used in the synthetic VAE experiment**\n\nA2: The discrepancy arises becuase in this toy experiment, to demonstrate our method is able to catpure correlaitons, we use it to capture correlation among different dimensions within one instance, ranther than capturing correlations among different instances. That is why we see the number of instances $N=6000$, but we only set the vertex set as ${\\mathcal{V}}=\\{1, 2, 3, 4\\}$. That is because the number of dimensions of the latent space is 4 in this toy experiment. We will make this point clearer in our final version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HTAFeIGj7Z", "reviewer": "Reviewer_96LZ", "summary": "The authors propose to use a variational approximation that captures tree-structured correlations as a middle ground between scalable fully-factorized approximations and more expressive approximations which aim to capture the full correlation structure. They show how these tree structured correlations give rise to variational distributions that factorize into tree-structured Markov random fields, whose ancestral sampling procedure informs a matrix-form reparameterization. The authors further extend their methods to mixtures of tree-structured variational distributions and show how to learn tree-structure by jointly optimizing over possible tree-constraint adjacency matrices.", "questions": "The authors report that their method is only about 2-3 times more costly than mean-field VI, however, the maximum number of edges (i.e. additional correlations that have to be learned) in a DAG with $n$ nodes is $\\mathcal{O}(n^2)$. So in the worst-case instances the computational cost could still be quite high. Are these instances just not learned in practice or do the authors explicitly regularize to encourage a degree of sparsity in the learned trees? Or is there another reason why this cost is not seen in practice?\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).", "limitations": "-", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Their 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.", "weaknesses": "It is unclear how important the initialization is and if there always exists a straightforward initialization procedure (see related questions). It would be nice to see some result on how the methods perform without initialization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and suggestions. Below we respond to each point raised by the reviewer:\n\n**Q1: It is unclear how important the initialization is and if there always exists a straightforward initialization procedure.**\n\nA1: Below we show the constrained clustering accuracies (%) on the MNIST dataset, with different initializations (random tree or greedy-searched tree), with or without our constrained optimization. It can be seen that without constrained optimization, both initializations underperform. And with constrained optimization, both initializations converge to substantially better performances, with the greedy-searched initialization slightly better. Overall, the constrained optimization procedure makes our proposed TreeVI  less sensitive to initializations.\n\n|  | Random Tree | Greedy Search |\n| -------- | -------- | -------- |\n| w/o CO | 96.58 | 96.70 |\n| w/ CO  | 97.31 | 97.45 |\n\n**Q2: The maximum number of edges in a DAG with $n$ nodes is $\\mathcal{O}(n^2)$. So in the worst-case instances the computational cost could still be quite high**\n\nA2: In our method, we require to use a tree, instead of a DAG, to capture the correlation among instances. For a tree, the number of edges $|{\\mathcal{E}}|$ is at most $N-1$, where $N$ denotes the number of vertices. Thus, the maximum number of edges is $O(N)$, instead of $O(N^2)$. As proved in our paper, by using our method, we only need to re-parameterize the edge parameters $\\gamma_{ij}$ for $(i,j)\\in {\\mathcal{E}}$, and thus only need to run the neural network by $|{\\mathcal{E}}|$ times for each epoch, in addition to the runs required by a mean-field method, which is $O(N)$. Thus, the total number of runs of the re-parameterization neural networks is approximately $2N$, roughtly two times of that required by mean-field variational inference. That's why we observed that the training time of our method is generally 2~3 times of the mean-field methods in our experiments.\n\n**Q3: Depending on the task it might be more or less clear how to compute meaningful neighborhood probabilities.**\n\nA3: The probability with respect to each pair of data points is calculated according to the cosine similarities between them, which can be leveraged to generate a meaningful neighborhood for each data instance. For more details we refer to Appendix D.2.\n\n**Q4: Something small that I noticed. The acronyms ARI and NMI have not been introduced and might not be familiar to readers not familiar with clustering.**\n\nA4: Sorry for the missing explanations for the acronyms of clustering metrics. These details will be included into our final version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Y6t2v5NAIR", "reviewer": "Reviewer_smr5", "summary": "The paper introduces a new type of variational inference that allows correlations across training samples in the form of a tree, making it suitable for applications with graph-structured data or explicit constraints. The method extends to multiple trees for capturing more complex correlations and includes a mechanism for automatically learning the underlying tree structures from data. Experimental results demonstrate that the proposed approach provide advantages in settings such as constrained clustering, user rating and link prediction.", "questions": "1. 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?", "limitations": "I believe the authors have addressed the limitations of the proposed approach.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "While 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.", "weaknesses": "1. I find the setting of the paper very confusing, I worked on hierarchical VAEs and therefore had some difficulty in understanding the notation of the paper. Figure 1 did not help as notations are misleading: what do the numbers in the nodes represent? Why in GraphVAE the nodes are labeled z_i while in the trees they are labeled with just a number? I believe the paper would be stronger if (a) notations could be better explained and (b) if comparison with TreeVAE [1] could be explained to differentiate the two settings.\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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and suggestions. Below we respond to each of the raised questions. \n\n**Q1: What do the numbers in the nodes represent in Figure 1?**\n\nA1: Sorry for the confusing notations used in the Figure 1. In Figure 1 (b) and \\(c), the numbers 1, 2, ... 5 in the trees should also be written as $\\mathbf{z}_1, \\mathbf{z}_2, \\cdots, \\mathbf{z}_5$. In our model, we assume for each data instance $\\mathbf{x}_i$, it corresponds to a latent variable $\\mathbf{z}_i$, and Fig. 1 (a), (b), \\(c) are used to discribe the correlation structures among differnt latent variables $\\mathbf{z}_i$. The notations in this paper are different from those in TreeVAE because TreeVAE focuses on how to model the hierarchical structure of different dimensions within one instance, while ours devotes to modeling the correlation among different instances.\n\n**Q2: I also believe that the paper could benefit from a more thorough analysis of the computational complexity of the proposed appraoch.**\n\nA2: Thanks for the suggestions. For amortized variational inference, the computational complexity mainly comes from running the re-parameterization neural network, thus the complexity can be measured by how many runs of the re-parameterization neural network are required. For mean-field VI, each instance only need to pass the neural network once, thus for a dataset with $N$ instances, the number runs of the neural network is $O(N)$. \n\nIn our method, by restricting the correlation matrix $\\mathbf{R}$ to a form constructed from a tree, we prove that all of the non-zero elements in $\\mathbf{L}$ can be explicitly computed from the $|\\mathcal{E}|$ parameters $\\gamma_{ij}$ for $(i,j)\\in \\mathcal{E}$. Thus, we only need to re-parameterize the $|\\mathcal{E}|$ parameters $\\gamma_{ij}$, and run re-parameterization neural network $|\\mathcal{E}|$ times, in addition to re-parameterizing the mean and variance, which requires a complexity of $O(N)$. For a tree, it is known that $|\\mathcal{E}|\\le N$, Thus, the total number of runs of the neural network is approximately $O(2N)$. This is also consistent with our observation in the experiments that the training time of our method is generally 2~3 times of mean-field methods.\n\n**Q3: Standard deviations are not reported in the table, and few experimental settings are explained. E.g. At what value is M set in the experiments?**\n\nA3: We thank the reviewer for pointing out the missing details. The results of constrained clustering with standard deviations are supplemented in our attachment to the global rebuttal, and will be included into our final version.\n\nIn our experiments, the value of $M$ is set to 3, and more detailed experimental settings are elaboratred in Appendix E.2. We will relocate some of the settings into the main context of the paper in our final version.\n\n**Q4: In 4.2 it would be more convincing to provide evidence that the learnt tree are consistent with the instance-level constraints.**\n\nA4: Thanks for the suggestion. Our method is mainly used to capture the inherent correlations among different instances. Thus, the correlations captured by the trees are not necessary to be consistent with the constrains, but instead are more inclined to align with the similarities among instances. To validate this, we plot the tree learned over the constrained clustering experiment on MNIST dataset, as shown in the Figure 1 in the attachment of global rebuttal. From the figure, we can see that instances from the same category are connected more tightly than those from different categories in the learned tree. This demonstrates that our method has the abilities to capture the underlying inherent correlations among different instances.\n\n**Q5: Could the authors provide a proof of Equation 4?**\n\nA5: Equation 4 is a result of the Markov tree dependent distributions, for a complete proof we can refer to [1] which is mainly done by induction. By the way, this equation can also be regarded as a special case of the Hammersley-Clifford theorem. We will add these citations in our final version.\n\nReference\n\n1. Meeuwissen A.M.H. and Cooke R.M., Tree dependent random variables Report 94-28, Delft University of Technology, 1994.\n\n**Q6: 4.1 Why the tree structures are defined a-priori and not learnt from the data? Validate the correctness of the the trees in the synthetic scenario.**\n\nA6: The experiment in Sec. 4.1 is a toy experiment to illustrate the contribution of each component of our method in capturing the accurate correlation. For better controlability, we choose to use pre-defined trees, instead of trees learned from the data. That is because if we learn the trees from the data, the tree $\\mathbb{T}_1$ in MTreeVI($\\mathbb{T}_1$, $\\mathbb{T}_2$) and the tree $\\mathbb{T}_1$ in TreeVI($\\mathbb{T}_1$) could be different. Thus, we choose to fix the two trees $\\mathbb{T}_1$ and $\\mathbb{T}_2$. As seen from the experiment, when using a single tree, either $\\mathbb{T}_1$ or $\\mathbb{T}_2$, better performance is obained comparing with the mean-field method. When the two trees are used simultaneously, we can see that the performance is further boosted.\n\nSince the 4 dimensions in the toy experiment are all closely correlated, any tree would be good for capturing the correlation among  the 4 dimensions. Thus, it is hard to clearly show which tree is better than the other. But we have demonstrated the consistency between the learned tree and the underlying correlation in the constrained clustering experiments, which is stated in the answer to Q4.\n\n**Q7: How to compute the latent embedding distance**\n\nA7: We utilize the same measurement of latent embedding distance as CVAE for fair comparison, which is computed based on the expected quadratic L2 distance $\\mathbb{E}_{q(\\mathbf{z}_i,\\mathbf{z}_j)}[\\| \\mathbf{z}_i-\\mathbf{z}_j \\|^2]$ of the latent distribution for VAE-based methods. Additionally for GraphSAGE, we simply use quadratic L2 distance $\\| \\mathbf{z}_i-\\mathbf{z}_j \\|^2$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their careful reading and detailed feedback.

Please refer to the attached PDF for new results. We note that we have provided additional insights about our learnt tree structure obtained through training. In Figure 1, we plot the tree structure learnt over constrained clustering task on the MNIST dataset, showing tight connection between our learnt correlation structure to the category distribution. Further, we supplement the constrained clustering results with both means and standard deviations for stability analysis as shown in Table 1.

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