Research map / Probabilistic and scientific modeling / Variational and generative sampling
TreeVI: Reparameterizable Tree-structured Variational Inference for Instance-level Correlation Capturing
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes TreeVI, a variational inference method whose posterior is a tree-structured Gaussian (and mixtures of trees, MTreeVI) that captures instance-level correlations with reparameterization costing only about 2-3x mean-field.
Problem
Mean-field variational inference assumes independence among data instances, which fails when instances are correlated through graph structure or constraints; existing structured VI methods target dimension correlations and do not scale to instance-level correlation.
Approach
Approximate the instance-level posterior correlation with a tree (or mixture of trees) in which only tree-edge correlation parameters are learned and non-adjacent correlations are products along paths; ancestral sampling gives a matrix-form reparameterization needing only O(N) network calls, and tree structure is learned via a NOTEARS-style acyclicity constraint.
Key results
- On constrained clustering, MTreeVI reaches 97.5% ACC on MNIST versus 96.5% for DC-GMM, and TreeVI reaches 97.4%.
- On synthetic user matching (MovieLens 20M), MTreeVI attains test RR 0.7521 versus 0.7129 for CVAE with correlation.
- Training cost is about 5 s/epoch (TreeVI) and 9 s/epoch (MTreeVI) on MNIST versus 2-3 s/epoch for DC-GMM.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the mean-field posterior with TreeVI (ablation) | VAE · Synthetic VAE | ELBO: -11.1535 -> -10.8998 | Helped |
| Use tree T2 instead of tree T1 (ablation) | TreeVI · Synthetic VAE | ELBO: -10.8998 -> -10.6137 | Helped |
| Add a second tree component to form MTreeVI (ablation) | TreeVI · Synthetic VAE | ELBO: -10.6137 -> -10.3586 | Helped |
| Increase mixture from one tree to two trees (ablation) | TreeVI · Synthetic VAE (further experiment) | ELBO: -8.8838 -> -8.8729 | Helped |
| Increase mixture from two trees to three trees (ablation) | MTreeVI · Synthetic VAE (further experiment) | ELBO: -8.8729 -> -8.8690 | Helped |
| Use greedy-searched rather than random tree initialization, without constrained optimization (ablation) | TreeVI · MNIST | ACC: 96.58 -> 96.70 | Helped |
| Apply constrained optimization with random tree initialization (ablation) | TreeVI · MNIST | ACC: 96.58 -> 97.31 | Helped |
| Apply constrained optimization with greedy-searched initialization (ablation) | TreeVI · MNIST | ACC: 96.70 -> 97.45 | Helped |
| Increase the number of tree components (ablation) | MTreeVI · Constrained clustering | ACC | Helped |
| Use TreeVI instead of DC-GMM's mean-field posterior | Constrained clustering · MNIST | ACC: 96.5 ±0.2 -> 97.4 ±0.3 | Helped |
| Use MTreeVI instead of CVAEcorr | User matching · MovieLens 20M | Test RR: 0.7129 ± 0.0096 -> 0.7521 ± 0.0101 | Helped |
| Use MTreeVI instead of CVAEcorr | Link prediction · Epinions | Test NCRR: 0.0171 ± 0.0009 -> 0.0203 ± 0.0014 | Helped |
Surprising findings
- Constrained optimization of the tree structure improves MNIST clustering accuracy only modestly, and random versus greedy tree initializations perform similarly after it (97.31 vs 97.45).
Limitations
- Correlation is restricted to tree structures, which may not capture the full graph-structured posterior correlation.
- Correlation is modeled only between instances; dimension-level correlation within each latent is omitted.
Experimental setup
Datasets: MNIST Fashion MNIST Reuters STL-10 MovieLens 20M Epinions synthetic graph-structured latent data
Benchmarks: MNIST Fashion MNIST Reuters STL-10 MovieLens 20M user matching Epinions link prediction synthetic ELBO test
Metrics: ACC NMI ARI reciprocal rank (RR) normalized cumulative reciprocal rank (NCRR) ELBO
Baselines: VAE CVAE GraphSAGE VaDE SDEC C-IDEC DGG DC-GMM PCKmeans DEC DCN
Builds on: VaDE DC-GMM CVAE NOTEARS variational autoencoder
Closest papers by meaning
- Least squares variational inference (NeurIPS 2025 · similarity 0.87)
- EigenVI: score-based variational inference with orthogonal function expansions (NeurIPS 2024 · similarity 0.87)
- SoftCVI: Contrastive variational inference with self-generated soft labels (ICLR 2025 · similarity 0.87)
- Provably Scalable Black-Box Variational Inference with Structured Variational Families (ICML 2024 · similarity 0.87)
- Nearly Dimension-Independent Convergence of Mean-Field Black-Box Variational Inference (NeurIPS 2025 · similarity 0.86)
- Variational Inference with Mixtures of Isotropic Gaussians (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.