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

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