Title: Variational Pólya Tree
Venue: NeurIPS 2025
Accepted: yes
Keywords: Nonparametric Bayesian

Variational Inference

Tree Models

Deep Density Estimation
Abstract: Density estimation is essential for generative modeling, particularly with the rise of modern neural networks. While existing methods capture complex data distributions, they often lack interpretability and uncertainty quantification. Bayesian nonparametric methods, especially the Pólya tree, offer a robust framework that addresses these issues by accurately capturing function behavior over small intervals. Traditional techniques like Markov chain Monte Carlo (MCMC) face high computational complexity and scalability limitations, hindering the use of Bayesian nonparametric methods in deep learning. To tackle this, we introduce the variational Pólya tree (VPT) model, which employs stochastic variational inference to compute posterior distributions. This model provides a flexible, nonparametric Bayesian prior that captures latent densities and works well with stochastic gradient optimization. We also leverage the joint distribution likelihood for a more precise variational posterior approximation than traditional mean-field methods. We evaluate the model performance on both real data and images, and demonstrate its competitiveness with other state-of-the-art deep density estimation methods. We also explore its ability in enhancing interpretability and uncertainty quantification. Code is available at https://github.com/howardchanth/var-polya-tree.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new prior for density estimators called the Variational Polya Tree (VPT). The idea is based on the Polya tree model, where a probability density over an interval is modeled as a tree that recursively partitions the interval. That is, the probability of a subinterval is proportional to the probability of traversing the tree from the root to a leaf node that represents the subinterval.\n\nThe Polya tree model is non-parametric and can be very expensive to work with. The main idea of this paper is to variationally infer the posterior of the Polya tree prior. That is, the authors parametrize the Polya tree density by learnable neural networks. Posterior inference is then cast as an optimization problem.\n\nThis prior can then be useful as a prior for density estimation problems. Specifically, the authors proposed to refine the VPT prior with flow-based models. Experiments in various settings (synthetic data, UCI datasets, and image datasets) show that VPT can yield good performance and is flexible since the tree can be refined further by adding more levels.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1. VPT is a novel idea that makes the Polya tree model tractable. Its usage for density estimation with flow-based models and VAEs is particularly interesting.\n2. Unlike other priors, VPTs are interpretable due to their tree structure.\n3. VPTs are efficient in terms of the number of parameters and runtime costs.\n4. Comprehensive experiments are presented to evaluate VPTs.\n\n\n**Weaknesses**\n\n1. The main weakness is the presentation of the VPT itself in Sec. 3.2. More detailed, step-by-step explanation needs to be added. \n    1. For example, it is unclear to me what it means by \"... parametrize tree partitions $\\Pi$ by learnable neural networks. ...\". Are the partitions' endpoints being predicted by NNs? Or something else?\n    2. Another example is in the discussion about the loss function itself. Why not incorporate the KL term in the ELBO maximization? Why is the loss description different than the equation in Alg. 1? Etc.\n    3. I would suggest the authors rewrite this section.\n\nQUESTIONS:\n1. See Weaknesses above.\n2. Are VPTs compatible (not just theoretically, but in terms of costs and practicality) with large-scale generative models like diffusion models?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose to utilize Bayesian nonparametrics, especially the Polya tree, to enhance interpretability and offer uncertainty quantification for deep neural density estimators. They develop a stochastic variational inference method that leverages the tree structure and internal conjugacy to obtain an efficient and precise posterior approximation. The authors then demonstrate that this nonparametric prior works well as a \"prior\" for the base distribution of deep neural density estimators as normalizing flows on UCI tabular datasets as well as MNIST and CIFAR-10. They show that if paired with good NF transformation layers, it performs competitively (or better) compared to state-of-the-art deep neural density estimation methods.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- Integrating Bayesian Nonparametrics methods for uncertainty estimation in deep neural density estimators is novel to me. Furthermore, the approach taken by the authors seems promising and more efficient than related approaches based on Bayesian Neural networks/Ensembles.\n- Results for density estimation are promising, especially on MNIST and CIFAR-10.\n\nWeaknesses:\n- As the authors also discuss in *end*, the work solely focuses on univariate Polya trees. Title, the intro, and methods sections strongly focus on the \"general\" variational Polya tree. **But all the experiments** use a factorized Polya tree as a \"base\" distribution for an existing normalizing flow architecture to achieve the results presented in the paper. Not a single experiment tests the VPT and its properties itself. This point should be made more clearly in the abstract/intro and raises the question if the VPT is responsible for the improved performance or simply a more complex base distribution combined with the same normalizing flow architecture (see comments).\n- The largest parts of the experiments focus on plain \"density estimation\" performance and only slightly on the main *claimed* improvements by using BNP, i.e., interpretability and uncertainty quantification. A main strong point of this approach and using Bayesian Nonparametrics in the first place is getting \"free\" interpretability + uncertainty, but the evaluation of these points is very limited.\n- While most of the methods are well described, some parts are not fully clear to me (see questions).\n\nComments: \n\nAs mentioned in the weaknesses, the experiments are more like benchmarking the Block-NAF normalizing flow block with a VPT base distribution instead of testing the VPT itself. Since a truncated factorized VPT is essentially a piecewise constant PDF (i.e., a Histogram distribution with learned binning), this basically benchmarks Block-NAF with a learnable, somewhat complex base distribution against Block-NAF with a simple base distribution (only on UCI, it is missing MNIST/CIFAR-10). This raises a major concern that is not addressed in the manuscript: Are the performance gains simply due to using an arbitrarily more complex base distribution, such as a learnable Histogram or a MoG, rather than the VPT itself? \n\nThis question could have been addressed by providing a non-BNP but equally flexible baseline (e.g. histogram with learnable bins equal to the number of bins a l-level VPT can have e.g. $2^{d*l}$) with the exact same Block-NAF blocks. Although I am not sure why a VPT should be better in this setting (see questions).\n\nSpecifically, the quite efficient uncertainty estimation seems like a great feature of this methodology (and is relevant for certain domains, i.e., [1],[2], which achieves this using BNN within the NF [which also should be discussed in the paper]). Unfortunately, the evaluation of quality falls short in the manuscript (basically, Fig 6 \"looks meaningful\") and no empirical results are presented that the estimated uncertainties are \"good\" or at least comparable or better than from alternative approaches (BNN, Ensembles ...).\n\n\n[1] Lemos, Pablo, et al. \"Robust simulation-based inference in cosmology with Bayesian neural networks.\" Machine Learning: Science and Technology 4.1 (2023): 01LT01.\n\n[2] Delaunoy, Arnaud, et al. \"Low-Budget Simulation-Based Inference with Bayesian Neural Networks.\" arXiv preprint arXiv:2408.15136 (2024).\n\nQUESTIONS:\n- Can the authors go into more details on how the VPT variational posterior distribution is parameterized? (Sec. 3.2 mentions that the tree partitions are parameterized by neural networks, but it doesn't go into detail, i.e., in 3.3, one samples the partitions from its \"variational\" posterior.\n- Would the author expect a learnable histogram base-distirbution for the exact same Block-NAF block to perform worse that a VPT? Especially if it's truncated, the MAP of the posterior should be a histogram with a fixed number of bins (but with variable bin size, which can be learned).\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces the Variational Pólya Tree (VPT), which integrates Pólya tree priors with deep neural networks for density estimation. The authors develop a variational inference framework that avoids mean-field independence assumptions by exploiting the hierarchical structure and conjugacy properties of Pólya trees. They demonstrate their approach on synthetic data, UCI datasets, and MNIST in a generative setting, showing competitive performance while claims on uncertainty quantification and interpretability.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: The technical execution is solid - they've carefully worked out how to make Pólya trees tractable for deep learning by leveraging conjugacy to avoid mean-field approximations, which is genuinely clever. The experimental coverage spans multiple domains (synthetic, tabular, images) and shows improvements over standard Gaussian/logistic priors.  The theoretical foundation is sound, building properly on established Pólya tree literature while making the key insight about preserving hierarchical dependencies during variational inference. The interpretability angle is compelling too - Figure 5 shows how their latent space interpolations follow more semantically meaningful paths than standard priors; however, it is limited.\n\nWeaknesses: The novelty feels somewhat incremental - this is essentially taking a well-established Bayesian nonparametric prior and plugging it into existing architectures like normalizing flows. The dimensional independence assumption is quite limiting for many real applications where cross-variable dependencies matter, and they don't adequately explore when this breaks down. I'm skeptical about computational scalability - while they provide some timing comparisons, the O(2^L × d) scaling with tree depth and dimensions could become prohibitive for high-dimensional problems. The experimental setup often uses relatively simple architectures (like Block-NAF instead of more modern flows), which makes me wonder how much of the improvement comes from the prior versus just having more parameters. The uncertainty quantification claims feel oversold - they show posterior variance visualizations but don't rigorously validate calibration or compare against other uncertainty methods.\n\nQUESTIONS:\n1) Computational Scalability and Practical Limits: Your method scales as O(2^L × d) with tree depth and dimensions. For high-dimensional applications like modern image generation (where d could be thousands), how does this constraint limit practical applicability? Can you provide concrete memory and time complexity comparisons beyond the simple timing tables?\n\n2) Independence Assumption Violations: You assume dimensional independence, which seems problematic for many real-world scenarios. How much performance degrades when strong cross-dimensional dependencies exist? Have you tested on datasets where this assumption clearly fails?\n\n3) Uncertainty Calibration: You claim superior uncertainty quantification, but only show variance visualizations. Are your prediction intervals actually well-calibrated? How do your uncertainty estimates compare to other Bayesian approaches like ensemble methods or Monte Carlo dropout in terms of reliability?\n\n4) Architecture Dependence: Your experiments often use relatively simple base architectures. How sensitive are the improvements to the choice of backbone network? Would we see similar gains when using state-of-the-art flows like coupling layers with more sophisticated transformations?\n\n5) How is $\\beta$ in line 153 defined?\n\n6) In line 167, $P(B_\\epsilon)$ should be dependent on the index $L$; however, it seems indepedent of $L$, why so?\n\n7) Are $B_\\epsilon$ probabilities or random variables?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-29"}

comments: SUMMARY:
This paper combines BNP techniques Polya trees, which are often under high computational costs, and VI to do density estimation, named VPT. The paper also provides some illustrations on this new method.

STRENGTHS AND WEAKNESSES:
1. I really like the BNP methods used. But it is too slow. Although VI techniques are used, have you tried scaling the questions up? For larger models, it may be not good enough. Can you compare the running time with other methods?
2. Ferguson 1974 pointed out that DP are special cases of Polya trees (See Ma 2017). But DP are the only tail-free process in which the choice of partitions does not affect inference, which shows on big advantage of BNO to solve the model misspecification problem. See, for example, Hanson and Johnson 2002. Why do you think VPT is better, or why do you think using BNP worth?
3. VIDEO is biased. 
4. Line 40: MFVI not valid, but have you tried other VIs, like FFVI?
5. Auto-regressive flows are a little bit outdated. 
6. Datasets used are not large. That may cause a good performance for BNP methods. 
7. The numerical experimental results are not enough. And some comparison with traditional methods are needed.

QUESTIONS:
Please see Strengths And Weaknesses.

LIMITATIONS:
N/A.
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-03