Variational Pólya Tree
Authors: Lu Xu
Tsai Hor Chan
Lequan Yu
Kwok Fai Lam
Guosheng Yin
Venue: NeurIPS 2025 poster ()
Keywords: Nonparametric Bayesian
Variational Inference
Tree Models
Deep Density Estimation
Primary area: probabilistic_methods

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.

Variational Pólya Tree
Lu Xu" The University of Hong Kong xulu09080hku.ht hk
Tsai Hor Chan'
University of Pennsylvania
Tsaihor Seashor..hhnnPeennkeee cine. upenn edu
Kwok Fai Lam Hong Kong Metropolitan University kf kflam@hmu. edu.hk
Lequan Yu The University of Hong Kong lqyuehku. .k
Guosheng Yin
The University of Hong Kong gyinChku.he hk
Abstract
Density estimation is essential for generative modeling, particularly with the rise of modern neural networks. While existing methods capture complex data distri- butions, they often lack interpretability and uncertainty quantification. Bayesian nonparametric methods, especially the Pólya tree, offer robust framework that addresses these issues by accurately capturing function behavior over small in- tervals. 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 Polya tree (VPT) model, which employs stochastic variational inference to compute posterior distributions. This model provides flexible, nonparametric Bayesian prior that captures latent densities and works well with stochastic gra- dient optimization. We also leverage the joint distribution likelihood for 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-ard deep density estimation methods. We also explore its ability in enhancing interpretability and uncertainty quantification. Code is available at htttttss///giieeee com hwwphhhhwwaahhhooooooopppa var wwtpt/aaaa-pooaaaaaa tree
Introduction
Density estimation has gained significant importance generative modeling, particularly with the advent of modern neural networks 3 28), such as normalizing flows and autoregressive networks. These techniques have found numerous applications, including image generation and large language modeling, highlighting their versatility and effectiveness in capturing complex data distributions. While many recent efforts focus on enhancing feature learning and transforming densities [46] 9 these deep density estimation networks often struggle with interpretability. raising concerns about their trustworthiness
30
Equal contributions Corresponding author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Bayesian nonparametric methods provide an appealing alternative, enabling flexible modeling that naturally adapts complexity according to data and rigorously quantifies uncertainty [32]. Incorporating Bayesian nonparametric priors such as Dirichlet processes (DP) into deep neural architectures has attracted substantial attention, yielding notable progress in regression [42] and clustering [26].
However, DP-based approaches inherently induce discreteness, making them suboptimal for directly modeling continuous distributions.
In contrast, the Pólya tree (PT) is specifically suited to continuous distributions, producing probability
measures absolutely continuous with respect to the Lebesgue measure [35] 4 29] Yet. despite their appealing theoretical properties, PT priors have rarely been combined with modern deep learning frameworks due to computational challenges. Traditional posterior inference via Markov chain Monte Carlo (MCMC) suffers scalability limitations, becoming prohibitively expensive as data scales [48 4
To bridge this gap, propose the variational Pólya tree (VPT), the first variational Bayesian integration of PT priors with deep neural networks. Our VPT leverages stochastic variational inference, enabling scalable training via stochastic gradient optimization. Importantly, our method does not rely on the simplifying independence assumptions typical of mean-field approximations. Instead, we exploit the PT prior's s intrinsic hierarchical structure and conjugacy properties, allowing tractable and exact joint posterior over Beta-distributed node weights. This approach maintains the rich dependency structure across tree nodes, facilitating efficient and precise backpropagation.
We validate VPT across diverse tasks, including high- nigh-dimensional tabular data and image density estimation using normalizing flow architectures. Unlike existing deep density estimation works that mainly focus on transformation and feature learning architectures [18 6. 39 15), we present a novel prior and demonstrate how VPT can enhance interpretability and provide a measure of uncertainty.
Our contributions can be summarized as follows: (I) We introduce the first variational inference framework that makes PT priors practically useful and can be used as plug-and-play components for continuous density estimation and deep generative modeling, integrating seamlessly with modern neural architectures, e.g., flows, variational autoencoders (VAEs). (2) We leverage the joint posterior likelihood as the variational objective to provide an exact approximation of the posterior distribution of the PT while minimizing computational complexity. (3) Empirical evaluations various datasets demonstrate that our VPT prior can achieve superior performance in density estimation compared to existing methods. Additionally, ablation analysis reveals that deeper tree generally leads more accurate density estimation. The framework incurs minimal overhead-lest than 0.05% additional memory and at most 1.3x runtime- while remaining end-to-end trainable with standard deep learning architectures. (4) We demonstrate that with the VPT prior, the model can effectively learn both the estimated density distributions and the associated uncertainties. Our findings also indicate that the latent features acquired using the VPT prior are significantly more meaningful than those derived from traditional distributions.
2
Related Works
2.1 Deep Density Estimation
Deep generative models have achieved significant success in density estimation, especially for high- dimensional data like images 33 28 These models typically transform latent variables, sampled from known distributions, into data points, enabling flexible and expressive density modeling.
Normalizing flows use sequences of invertible transformations, parameterized by neural networks, to model complex distributions. They enable efficient likelihood evaluation by maintaining tractable Jacobian computations. For example, NICE [8] uses coupling layers, while RealNVP [9] extends this with additive and multiplicative coupling layers for improved expressiveness. Autoregressive methods decompose D D-dimensional density into -imeeinnnsional conditional densities, functioning as a specialized form of normalizing flow. MAF [39] uses masked neural networks to model these conditional distributions, while IAF 23 enables more efficient sampling through inverse transformations. VAEs 20 introduce stochastic transformations to learn compact representations but often suffer from posterior collapse due to the evidence lower bound (ELBO). Diffusion models 17., a hierarchical extension of VAEs, use forward noise addition and reverse denoising processes to generate data.

---

Despite these advances. most deep density estimators still rely on parametric assumptions and lack a principled way encode interpretability and uncertainty about the distribution. Most recent works in this area focus on learning powerful and tractable transformations from input to the latent density. For example, Glow [22 introduces invertible x 1 convolutions to improve expressiveness. Neural autoregressive flow (NAF) [18] employs autoregressive models within normalizing flow framework, allowing for efficient and expressive density estimation. MuLAN [41] utilizes adaptive noise techniques within diffusion models to improve sample quality and diversity. We propose Bayesian nonparametrie prior that can be compatible with these transformation approaches, which is not only flexible train but also provides interpretability and uncertainty quantification naturally.
2.2 Bayesian Nonparametric Methods
Bayesian nonparametric (BNP) models provide a framework for adapting model complexity to data 47 40
16
19 45) Classic examples include Dirichlet processes (DPs), which enable unbounded
mixture components, and have been extended to deep learning, such as stick-breaking VAEs and Beta-Bernoulii processes for infinite latent features [33 However, DPs concentrate probability on discrete distributions, making them unsuitable for continuous densities without additional smoothing
335. Variational inference for DPs also requires truncating the infinite process, which can result in biased estimates.
Pólya trees, contrast, naturally model continuous densities by recursively partitioning the domain with random probabilities. With appropriate hyper-parameters, they place probability 1 on the space of continuous distributions 225. Unlike DPs, PT's avoid discreteness and external smoothing. Despite this advantage, their use in modern machine learning has been limited, as existing applications relied on computationally expensive methods like MCMC (2). Crucially, no previous work has incorporated PT priors into deep neural architectures with variational Bayes.
We address this gap by introducing a PTT-ased variational inference framework for deep generative models. Our proposed framework is not a plug-and-play prior placement but a principled, nontrivial integration that bridges classical Bayesian nonparametries with modern deep learning. leverages the distinctive theoretical properties of PTs, enabling benefits such as robustness, interpretability. generalization, and calibrated uncertainty estimation. PTs' conjugacy allows efficient updates of branch probabilities via Beta-Binomial computation, enabling a structured variational posterior that retains dependencies across the tree. Unlike mean-field methods, which assume independence 45 14), our approach preserves hierarchical coherence, avoiding oversimplified approximations.
By integrating the flexibility of BNP priors with the scalability of deep learning, our method enables efficient end-to-end training of continuous density estimators. This approach overcomes the limita- tions of DP-based priors for continuous data and advances adaptive, interpretable density modeling in deep generative frameworks.
2.3 Tree-based Models
The use of binary trees is prevalent in machine learning 7 42
In addition to well-known methods like decision trees, random forests, boosting, and XGBoost (5) Bayesian approaches tree models have gained popularity, particularly with advancements in computational capacity.
Further, Bayesian inference has gained enormous popularity for modeling the distribution of trees, as seen in Bayesian additive regression trees (BART) M. Salazar [42] introduces a variational tree model for regression that establishes a Bayesian nonparametric prior the tree space. However, like traditional methods, it partitions the sample space rather than the probability space.
In contrast, our approach uses PT prior to model continuous densities. A tree space that partitions the probability space would offer greater modeling flexibility, accommodating a wider range of probability measures and enabling various tasks beyond regression. Our variational algorithm preserves dependencies across tree levels, ensuring coherent uncertainty propagation and avoiding the limitations of the mean-field approximation.
3

---

3 Methodology
3.1 The Pólya Tree Prior
The PT prior is a Bayesian nonparametric ap- proach for constructing random probability mea- sure. The construction is based on a recursive partition of the domain, which can be viewed as a random histogram with the sizes of the bins made sequentially smaller (38
B
0
4
B
=
Ant
Ain
B1
Consider the bins of a histogram defined by a partition of the sample space into nonempty sub- sets (Be.. with subsets indexed by an L-digit binary number Ei:L El EL, where each el- ement Ej (j = I,.... L) takes values of 0 or At level j. we define random probabilities p( for each bin. Consider refining the par- tition by splitting each bin into left and right parts, B Bry B U Bey,1 and define random probabilities for the refined histogram by condi- tional probabilities Y Yo P(B P(B ol Bey,)) and
Fn
j =
L
Figure I: Graphical illustration of the Pólya tree construction. The PT prior randomly splits Bey, into two subintervals Be and B Benyl with Beta- distributed probabilities.
Y Ya.,1 = P(B Ber.)) For any integer L 2 O, this recursive refinement of bins defines a sequence
of nested partitions.
random probability measure follows PT distribution PT(A, {3,})=1) with parameters A = {ae] on the sequence of partitions {33 1 where refers to the collection of Beta distribution random variables (i.e., the branching probabilities), if there exist splitting probabilities such that for any (1) The variables Y., are mutually independent and the posterior follow Beta Cley de,) distribution; (2) Y.
Y. S 1
Y. (3) For any L 2 0, P(B...)
IIi Ye....y.
Figure 1 illustrates this recursive representa- tion. summary, the PT prior defines a random probability distribution on distributions by assigning any partitioning subset B, with
[0, 1]
Algorithm 1 Training Procedure of Variational Pólya Tree (VPT)
probability P(B.) = 11 Y Ye x 11 Ye,........... - Algorithm shows
1: Input: Data points {ar dimension D, tree level L, number epochs Nepochh learning rate
the overall workflow of our VPT.
2: Compute total number of nodes per dimension:
Pnnodss 2 1
3.2 Variational Polya Tree
3: Initialize Beta parameters for all nodes:
Gen--10 1, 0m.--11 =
Variational Objective. Our VPT approach
4: for epoch 1,... Mepoch do
leverages the intrinsic hierarchical structure of the Pólya tree to avoid traditional independence
5:
for each dimension d = 1,
D do for each node j = 1,..., podos
6:
assumptions common in mean-field variational inference. Specifically, we parameterize the tree partitions {/3, f by learnable neural networks. Due to the conjugacy of the PT prior, the poste-
7:
Sample split:
8:
y-(a)
c1 oo x Beta(a e1.j e ens
(d)
Compute partition intervals ((d))
9:
rior retains the PT form, allowing closed-form 10:
end for
updates for each node's Beta parameters. This
11:
end for
structural property preserves hierarchical depen-
12:
Compute joint posterior in Eq. (1.
dencies among latent variables without any sim-
13:
p((3,)}=11 x), and LVPT = log Pogg(( (())) x)
plifying assumptions, which distinguishes VPT from the classical mean-field approximation.
14:
Update 0 10: Oen.j 11.
16: end for
Let x N be the observed data, (Br be the L-level tree partition proba-
bilities corresponding to the i-th observation, and D be the set of all Yeu, in the tree. Following

---

example, Paddock et al. [38. the joint posterior distribution of an L-level tree given can be written as
N
L II ;=1
((18)) = II
Y
((L
II
y (ons 915-10
10 sa- D (ons 1111) 915-11
(Bes
v(EL5-1)
where v/ ( (B(),..2)) =1 Big 1- (1- Be.)) si and p(Bc,) is the prior distribution for Bt.j. Data points 2; influence densities through components, the probabilities Bei.; indicating the likelihood of 24 falling into a specific subset, and the conditional probabilities Ye., that govern the likelihood of I, in child subsets given its parent subset. The gradients can be computed and backpropagated to update the parameters 1:1---10 and dc1:--1 of the Beta distributions within the Pólya tree.
We optimize VPT parameters by maximizing the log- ongliikeihhod (i.e., the ELBO) LVPT log (8))),1))) Note that our method does not explicitly include Kullback-Leible (KI.) divergence term, and the corresponding entropy of the posterior acts implicit regularization. This entropy term, which appears in the ELBO, plays a similar role by controlling the spread of the variational distribution and balancing data fit versus uncertainty. Nevertheless, incorporating an explicit KL (as shown in Appendix A) or entropy term can further balance exploration of the prior.
Density Estimation with Flow-basee Models. The VPT prior alone defines flexible density over a sample space. For complex high-dimensional data, we can further enhance its flexibility by embedding the VPT with deep neural networks. Taking flow-based networks use the PT density as the base distribution and apply an invertible deep transformation f from the latent space to the observed data space. Formally, let represent samples drawn from an unknown D-dimensional distribution. The goal of density estimation is to model the underlying distribution function F(-). Instead of assuming fixed parametric form for F()) in the previous work (42), we follow a Bayesian nonparametric approach by placing a VPT prior on F(-). thereby treating the distribution itself as random. Consider bijective function f X - 2 with the inverse function f-1 and Jacobian (z) Under this transformation, the density in the latent space Z can be computed as
(z))) px (xx(((())) det (=))),
where px (x:) is modeled by VPT prior. The determinant of the Jacobian ensures the validity of the density transformation, which can be computed following Dinh et al. [8]. The model is optimized by maximizing the log-likelihood distributions, VPT base distribution is learned and updated via the posterior of the PT. In practice, alternate analytic posterior updates of the PT parameters with gradient-based optimization.
L = log (2)). Unlike the standard flows with fixed base
3.3 Computation of Intervals
We use tree data structure to compute the intervals of each leaf node [38] For the i-th sample Ty we first sample {3.s) from its variational posterior distribution. For dimension where
(d) r E R. the interval (0, 1] is split into two subintervals with probabilities 3(() and .6 so that the subinterval (3(4) is of length g(d) and B(d) is of length 1 =p(a) Next, each of ( and Bld) is
further split into two subintervals where the splitting probability depends on (a5() as follows,
B (0, ((d) g(d)
Did (B(4) ,(())
Bid = (8ft) + (1 - (10)) 15 ,1]. The intervals at level j are obtained iteratively by applying the above rule.
Bi0 = (0( , + (1- 3(44) )o!"))
3.4 Posterior Inference and Interpretability
With the learned from the Beta distribution to obtain the probability of splits. Due to the conjugacy of the PT prior, the posterior predictive distribution has an intuitive form,
( and {ي for the i-th sample x, with D dimensions, we can sample
D L
1- (x) (1 (
(
p( yh ((B((x))
(d
(1)0
(z)
d=1j=1
5

---

where Dوh =Vrle
as-100 of. yrd)
is the (left) split variable at node é13-1 in dimension d, and under
fzg-10
the variational posterior y-(a)
(d) Beta Reta(oe -10 (1g--1)) with learned shape parameters (x) = 1 (z) is the parent index used at level j. Br(x) I B(d 11 (x)
-1:--10
(d)
(d)
fis- 111-111 5)
is the depth-L leaf containing æ, and (z) = I [B (w)
is its volume.
This hierarchical organization of posterior parameters makes the model highly interpretable. One can traverse the tree to see how the model allocates probability at different scales and locations. Each level of the tree gives coarse-to- view of the density.
Uncertainty quantification. As Bayesian nonparametric method, one of the key advantages of our VPT is to learn a distribution over the density estimation. This allows us to obtain the variance of the learned model. Given trained I.-level VPT, the posterior variance can be modeled as the mean over the variances of terminal node distributions, where the posterior variance of each subinterval can be computed easily from the Beta distribution.
4 Experiments
We evaluate VPT on multiple density estimation tasks. We first demonstrate the qualitative benefits of VPT compared to traditional Gaussian and logistic priors using simple synthetic data. We then assess VPT quantitatively on UCI and image datasets. Our findings indicate that models with VPT priors not only achieve superior likelihood results but also provide robust uncertainty estimates.
Implementation details. Training VPT involves optimizing a of Beta distribution parameters. To ensure the Beta distribution parameters remain positive, we reparameterize them via a softplus(-) transformation. To integrate VPT prior seamlessly with standard generative networks, we use sig- moid layer to project latent variables into the interval [o, 1]. This sigmoid mapping is computationally efficient, adding minimal complexity since its Jacobian determinant can be computed trivially.
4.1 Density Estimation with 2D Synthetic Data
We first illustrate the capability of our VPT prior using three common synthetic datasets, ring of 8 Gaussians, two interwoven spirals, and a checkerboard pattern. For these experiments, we adopt simple block neural autoregressive flow (Block-NAF)) architecture as the feature learning backbone, specifically employing one flow layer consisting of two hidden layers, each with 50 units. This modest architecture choice intentional, as our primary goal is not to achieve totte--f-the-att performance on synthetic benchmarks, but rather to clearly demonstrate the benefit of the VPT prior compared to standard priors (Gaussian and logis- tic) under identical model complexity.
Gaussian
VPT (level=2)
VPT (level=3)
we compare the learned densities using three different priors, a 2-level VPT, a 3- level VPT, and an isotropic Gaussian. While the isotropic Gaussian prior allows the model to ap- proximate the general shape of multimodal distri- butions, it struggles to capture sharp boundaries and disconnected regions accurately. In contrast, both the 2-level and 3-veve VPT priors clearly yield more precise representations, effectively cap- turing regions of low density and the multimodal structure inherent the data. Moreover, increas- ing the level of the VPT prior from two to three levels further enhances its modeling capacity, re- sulting in notably improved density estimation and clearer separation among data clusters.
In Figure 2
Figure 2: Density estimation results of an isotropic Gaussian prior, 2-level VPT, and a 3-level VPT with 2D synthetic datasets.

---

Table 1: Log- Log-likelihoods (standard deviations) the test sets for density estimation using VPT and baseline methods on real datasets with 5 runs. Higher values indicate better estimation.
Methods
POWER
GAS
HEPMASS
MINIBOONE
BSDS300
Real NVP 9
(.17
(o.3().14)
-1.07
(13.55(0.49)
153.28 (1.78) (53.28(1.78)
Glow
(.10(0
(.15
-18.92
(11.35
155.07 (0.08)
MADE MoG 33
(o.0)
8.47
-15.15(0.02)
-12.27
153.7 153.71
FFJORD 15
(.46().01)
(..
-14.92(0.09)
(0..43(0.04)
157.40(1.78)
MAF MoG [39
(o.0)
(0.02
(0..39(0.02)
-11.68
156.36 156.36(0.28)
TAN [36
(.60))
(2.06(0.02)
-13.78(0.02)
-11.01 (o.48)
159.80 (59.80(o.07))
NAF-DDSF
0.62(0.01)
11.96(0.33)
-15.09(0.40)
-8.86(0.15)
157.43
Block-NAF 6
(..0)
(0.10
- -15.06(0.08)
(.90
(o..33
VPT (L- 4)
0.67
(1.92(0.10)
(0.001(0.04)
(..2
158.59
VPT (L 6)
(..0))
(2.20
-13.94(0.05)
8.51(
162.72(0.94)
4.2 Density Estimation with Real Data
We perform density estimation five tabular UCI datasets, POWER, GAS, HEPMASS, MINI- BOONE, and BSDS300. Detailed information about these datasets is summarized in Appendix B (Table AT We follow the preprocessing procedure outlined in [39]. For this experiment, employ a Block-NAF [6] as our feature learning architecture, which has fewer parameters compared to the original neural autoregressive flow [18]. Consistent with the Block-NAF methodology, train stacked flows, each with 2 layers and 20D hidden units, where represents the input dimension. The model is trained using the Adam optimizer, with a learning rate of 10 for the Block-NAF flow and 0.1 for the variational Pólya tree.
summarizes the reported liglliiel values. With 4-level VPT, model outperforms Block-NAF across all datasets. Additionally, with 6-level VPT, our model demonstrates an even greater margin of improvement over Block-NAF, except the POWER dataset.
Results. Table I
By comparing our results with deeper density esti- mation models, we ob- serve that our VPT models achieve the state-of-the-ar performance across various datasets, with the exception of the HEPMASS dataset, where we attain results com- parable to TAN 36
Table 2: Number of model parameters, relative to Block-NAF. Block- 6 use the NAF and NAF use Gaussian priors, while VPTs with
Block-NAF backbone with the VPT prior. The notation "x" " indicates multiplicative factor, while "+" indicates an additive increase.
Datasets
Block-NAF
NAF
VPT(4)
VPT(6)
POWER GAS HEPMASS MINIBOONE
414, 213
x4.57
+180
+756
401, 741
x2.60
+240
+1008
9,272. 743
x35.88
+630
+2646
Number of parameters. We compare model plexity against NAF [88
7,487,321
x87.91
+1290
+5418
BSD300
36, 757,59,50 591
X x16.48
+1890
+7938
and Block-NAF (6). Our
VPT adds only (2- 1) x 2 x D parameters to Block-NAF, a negligible overhead compared NAF, Despite this minimal increase, VPT significantly surpasses Block-NAF in performance. demonstrating superior modeling power and efficiency (see Table
4.3 Image Density Estimation and Image Generation
We employ classic flow-based network NICE 8 as our feature learning backbone, and we use the same settings as in the original paper. Additional details of implementations are presented the Appendix.
We further test our method on two image datasets: MNIST and CIFAR-10 (2].
Results. In Figure
compare our VPT method with the vanilla NICE method g using two priors, Gaussian and logistic distributions. important to note that the metric employed in the original paper is leelliklsssooss in the logit space, which is not directly comparable across different
7

---

MNIST Density Estimation
CIFAR-10 Density Estimation
4.54
NCCllcggeee
6.10
5-42
VPT level-4
2.08
VPT Jeuele))
P.
(wwelit)
1.49
Fool NVF
4.54
6.0
RNCDE
0.87
1.52
MADE NoG
1.41
VCM
2.43
Pasl MVP
2:00
5
Bis Per Ctm
Bits Per Dim
Figure 3: Negative log-likelihood in biis-per-dimension on the MNIST and CIFAR- 10 test sets. Our method, VPT (shown gray), shares the same architecture as NICE (shown in black). The results from other methods (shown in light gray) are sourced from their respective papers.
2
2
6 6
& a
6
6
6
6
6
6
6
8
8
4
3
3
3
0
5
3
0
0
O
f.
3
3 9
9
9 8
9 &
9
9 @
9 8 4 7 9
Q 2
I
J
9
9
6
Ø
6 6 A
A
6 6
6
J
2
8
8
8
9
1
3
9 9 9
1
4 4
9
3
3
3
3
8
8
7
7
7
8
8
4
7
8 8
3
7
7
7
7 7
7 7 7
7 7
7
3
7 7 7 7
7
7
7
7
3
3
7
7
7
7
7
7
3
3
3
3
9 9
9
4
4
9
7
9
7
9
7
7
7
(a) Gaussian Prior
(b) Logistic Prior
VPT (L=4) Prior
Figure 4: Interpolation on the MNIST images. While with the Gaussian prior and logistic prior the latent spaces are focused on the pixel-level features, the mixed latent features with our VPT prior reveal how the particular tree nodes are structured.
estimation methods. Therefore, re-implemented their experiments and report the normalized negative log-likelihood in bits-prr-immnnion [39] log, p(x). Our results indicate that the VPT prior significantly enhances the likelihood for both datasets.
We also compare our results with a variety of density estimation methods in Figure 3 These include
flow-based methods such as Real NVP 19. MADE MoG [39], and MAF MoG [39;;
;
RNN-based
transformation approaches like TAN 36 the Neural ODE method (RNODE) 12 and more recently developed diffusion- ffffusiio--bseedd methods, including ScoreFlow (44), VDM (21, and MuLAN [41). We observe that our VPT prior performs best among all the flow-based density estimation methods, while demonstrating comparable performance with the recently developed diffusion- based methods. With the 4-level VPT, our VPT achieves a negative log-likelihood of 0.94 on MNIST, which is similar to 0.97 by RNODE. On the CIFAR-10 data, results with the 2-level VPT are comparable those of MuI.AN, while the performance can be further enhanced with the 4-level VPT.
4.4 Uncertainty Analysis and Interpretation
Uncertainty analysis. We compare uncertainty calibration with other uncertainty estimation methods, MC-dropout [3] and variational Bayesian neural network (BNN) with mean- mean-field Gaussian weights using the standardized squared error (SSE). For every sample and every dimension d 1,..., .,. D,
34
(d) (a; (d) a(d) (dd If the predictive we compute SSE = Ei { (d) /²/ND, where z;
variance is perfectly calibrated then E(z) = 1, thus SSE = 1, otherwise SSE > indicates under- estimated variance, SSE < indicates over-estimated variance.
In Table 3 the SSE results show that VPT slightly over-estimates its variance, while keeping the best calibration. MC-dropout and BNN underestimate their variances. The hierarchical structure of VPT updates the shrunk leaf probabilities toward their parents when data are sparse. This automatically inflates predictive variance in low-count regions. MC-Dropout [3] and mean mean-field BNNs (3

---

variance often fails to fully propagate to the output, leading to the systematic under-estimation.
Table 3: Predictive variance calibration with SSE closer to the better.
Figure 5) visualizes posterior variance estimates from our trained 4-level VPT MNIST. Higher uncertainty is observed around the central digit region, aligning well with intuition-this area ex- hibits greater variability across digits. This con- firms our model's capability to provide meaningful uncertainty quantification alongside accurate density estimation.
Datasets
VPT
MC-Dropout
BNN
POWER
0.92
1.30
1.17
GAS
0.90
1.35
1.14
MNIST
0.92
1.27
1.08
We emphasize that VPT's variances Carlo sampling is required at test time, whereas both baselines need extra computation per input. Thus VPT offers better calibration without extra computational burden.
analytic, Monte-
Effectiveness of the tree structure. To evaluate the effec- tiveness of the hierarchical structure of VPT, we compare with learnable histogram (LH) prior. We set the number of bins to K = 2 to match the number of nodes in an L-level VPT. For each dimension
we parametrize a increasing se- quence (()) < b)(1) < ... < by by setting the lower boundary b/4) and ((d) k-- soft softplus(8 The parameters og learned jointly with the backbone network, and softplus(-) guarantees positive bin widths. Given one dimension _d44) can find the active bin index ki, and compute the log-likelihood
Figure 5: Posterior variance of the trained 4-level VPT model on the MNIST dataset. On each dimen- sion/pixel (d = :.....D. where D 28 28), visualize the mean over the variances of terminal node Beta distribution.
as log p() log Pk.
(d) log(bb. bi()). All the hyper-
parameters are kept the same during training as those in VPT.
Table 4 shows the results UCI datasets with the Block-NAF backbone. The LH prior is used as a plug-in. We observe a clear advantage of VPT over LH, although they share the same number of parameters. We compare LH with VPT on CIFAR-10 10 with the NICE network. With L = 4, VPT obtains a result of 2.06 bits- tils-prrdimensionn and LH yields 2.57 BPD.
Table log-liielihoods (standard deviations) on the test for density estimation using our VPT and the learnable histogram (LH) prior under the same setting (L 4 and K 24) three real datasets with runs. Higher values indicate better estimation.
Priors
POWER
GAS
BSD300
Learnable histogram (LH)
0.62 0.67 (0.01) 11.92
(o.00)
156.98
Variational Pólya (VPT)
158.59
While a truncated factorized VPT may resemble learnable histogram in terms of producing piecewise-constant density, VPTs are more robust. Unlike an LH, where bins are estimated inde- pendently, VPT leverages the hierarchical conjugacy, enabling information sharing across scales. Specifically, every Beta-disttrbuted node in the VPT is coupled to ancestors. This hierarchical shrinkage allows small bins with few data points to borrow statistical strength from parent nodes, automatically adjusting the degree of smoothing. The VPT's variational posterior inference maintains the joint posterior over tree splits. This enables more coherent and uncertainty-ave density esti- mate, especially in regions of sparse data. In contrast, LHs often suffer from overfitting in low-density regions or collapsing of underpopulated bins.
Interpretation of VPTs. Besides density estimation, VPT naturally provides probability mass at multiple scales, providing 'coarse-to-fine insights. To better understand what the model has learned, we perform interpolation by linearly mixing the latent representations of MNIST images. Figure illustrates the interpolation paths between pairs of digits under different prior distributions.
Compared to Gaussian and logistic priors, VPT yields meaningful traversals in latent space. For example, when interpolating from '0' to '3', the intermediate representations pass through recognizable '1'. Similarly, interpolating from to '8' reveals a transition through '3'. These patterns suggest that, under the VPT prior, the latent space is hicrarchically organized, node

---

17771 7
9 7
4
9 4
8 5
8
8 8
6
6
1721
2
9
5
9
4
9
6
3
5
8
3
6 6
2
! 2
9
4
9
7
7
3
5
8
2
8
9
Figure 6: From left to right: images sampled from leaf nodes (nodes 3, 4, 5, 6) of a 3-1evel VAE- VPT.
corresponding the digit 'I' likely resides between nodes representing "0" and '3', while node associated with '3' is positioned between those representing '9' and '8', This hierarchical structure offers an interpretable organization of latent representations compared to standard unimodal priors.
5 VAEs with the VPT Prior
Beyond the above experiments, we demonstrate the flexibility of VPT prior by integrating into a VAE. We use a vanilla VAE architecture with fully connected encoder and decoder networks, each consisting of two hidden layers of 400 neurons and latent space dimension of By simply replacing the Gaussian prior with our VPT prior, we define a modified ELBO objective as
CxPT = (= [o - Σ
vEnodes
where the first term is the reconstruction error in the standard VAE and each node has a KL term comparing the approximated posterior (>(2, (2) to the VPT prior p(2.).
Specifically, we implement 3-level VPT prior using tree with 7 nodes, where node 0 is the root, nodes 1-2 are intermediate nodes, and nodes 3-6 are leaf nodes. After training, we generate images by directly sampling latent variables conditioned
each leaf node (see Figure The leaf nodes naturally cluster visually similar digit images. For instance, node 3 generates digit images resembling "1', '2' and '7', while node produces "9' and '4', Leaf node tends to generate curved digits like '3','s '5' and '8', and node 6 generates looped digits such as 'O' and '6', This demonstrates that the hierarchical partitioning induced by the VPT prior effectively captures meaningful groupings within the latent space, which enhances the interpretability compared to VAE with a Gaussian prior.
Conclusion
We introduce the variational Pólya tree, the first deep generative framework to integrate continuous Pólya tree priors with neural architectures. VPT effectively captures complex, high-dimensional densities and scales to large datasets. While classical BNP models such as Mondrian forests [31] adaptively grow deeper with more data to guarantee universal consistency, we compensate for the fixed tree level L. using expressive neural transformations that allow shallow structures in latent space. Although a fixed L. may limit asymptotic consistency, our approach maintains key theoretical properties. As supported by Orbanz (37 truncated PT can still achieve local a-Hölder continuity at non-dyadic points, preserving approximation capability without requiring strong continuity assumptions. Furthermore, we learn the Beta parameters via variational inference rather than fixing them the priori, enhancing adaptivity while preventing overfitting.
capture complex
Extensive experiments show that VPT significantly improves density estimation performance across various datasets. Beyond enhanced predictive accuracy, our approach provides meaningful uncertainty quantification and improved interpretability. Our findings demonstrate that VPT delivers competitive performance at different levels of the tree. This research opens promising avenues for advancing density estimation models using Bayesian nonparametric methods.
Acknowledgements
This work was supported part by the Research Grants Council of Hong Kong (27206123, 17200125, C5055-24G, and T45-401/22-N), the Hong Kong Innovation and Technology Fund (GHP/318/222) the National Natural Science Foundation of China (No. 62201483), and the Guangdong Natural Science Fund (No. 2024A1515011875
10

---

References
[1] BART: Bayesian additive regression trees. The Annals of Applied Statistics, 4(1):266-298.
2010.
[2] Naoki Awaya and Li Ma. Hidden markov Pólya trees for high-dimensional distributions. Journal of the American Statistical Association, 11 (115445):8992011 2024.
[3] Siavash A. Bigdeli, Geng Lin, L. Andrea Dunbar, Tiziano Portenier, and Matthias Zwicker. Learning generative models using denoising density estimators. IEEE Transactions on Neural Networks and Learning Systems, 35(12):17730-17744 2024.
[4] Ismaël Castillo. Pólya tree posterior distributions densities. Annales de L' L'Institul Henri Poincare Section (B) Probability and Statistics, 53(4):2074-2102, 2017.
[5] Tianqi Chen and Carlos Guestrin. XGBoost: A scalable tree boosting system. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pages 785-794, 2016.
[6] Nicola De Cao, Wilker Aziz, and Ivan Titov. Block neural autoregressive flow. In Uncertainty in artificial intelligence, pages 1263-1273, 2020.
[7] Barry De Ville. Decision trees. Wiley Interdisciplinary Reviews: Computational Statistics, 5(6): 448-455, 2013.
[8] Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: Non-linear independent components estimation. Workshop of International Conference Learning Representations, 2014.
[9] Laurent Dinh, Jascha bohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In International Conference on Learning Representations, 2017.
[10] Finale Doshi, Kurt Miller, Jurgen Van Gael, and Yee Whye Teh. Variational inference for the indian buffet process. In Proceedings of the Twelfth International Conference on Artificial Intelligence and Statisties, volume 5, pages 137-144, 2009.
[11] Michael Dusenberry, Ghassen Jerfel, Yeming Wen, Yian Ma. Jasper Snoek, Katherine Heller, Balaji Lakshminarayanan, and Dustin Tran. Efficient and scalable Bayesian neural nets with rank-1 1 factors. International Conference on Machine Learming. pages 2782-2792, 2020.
[12] Chris Finlay, -eerik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ODE: the world of Jacobian and kinetic regularization. In The 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 3154-31 1144-11-44 3164, 2020.
[13] Yarin Gal and Zoubin Ghahramani. Dropout a Bayesian approximation: representing model uncertainty in deep learning. In International Conference on Machine Learning, pages 1050-1059, 2016.
[14] Philip N Garner and Sibo Tong. Bayesian approach to recurrence in neural networks. JEEE Transactions Pattern Analysis and Machine Intelligence, 43(8):2527-2537, 2020.
[15] Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. FFJORD: Free-form continuous dynamics for scalable reversible generative models, 2018.
[16] Manuel Haußmann, Fred A Hamprecht, and Melih Kandemir. Variational Bayesian multiple instance learning with Gaussian processes. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 6570-6579, 2017.
[17] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pages 6840-6851, 2020.
[18] Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autore- gressive flows. In Proceedings of the 35th International Conference on Machine Learning, volume 80, pages 2078-2087 2018

---

[19] Kirthevasan Kandasamy, Willie Neiswanger, Jeff Schneider, Barnabas Poczos, and Eric P Xing. Neural architecture search with Bayesian optimisation and optimal transport. Advances in Neural Information Processing Systems, 31, 2018.
[20] Diederik P.
Kingma and Max Welling, Auto-encoding variational Bayes. In 2nd International Conference on Leurning Representations. 2014.
[21] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. On density estimation with diffusion models. In Advances in Neural Information Processing Systems, 2021.
[22] Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1xl convolutions.
In Advances in Neural Information Processing Systems, volume 31, 2018.
[23] Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling, Improved variational inference with inverse autoregressiive flow. In Advances in Neural Infor- mation Processing Systems, volume 29, 2016.
[24] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
[25] Michael Lavine. Some aspects of Pólya tree distributions for statistical modelling. The Annals of Statistics, pages 1222-1235, 1992.
[26] Naiqi Li, Wenjie Li. Yong Jiang, and Shu- Tao Xia. Deep Dirichlet process mixture models. In Uncertainty in Artificial Intelligence, pages 1138-1147, 2022.
[27] Zhihang Li, Teng Xi, Jiankang Deng, Gang Zhang, Shengzhao Wen, and Ran He. GP-NAS: Gaussian process based neural architecture search. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1 11933-11942, 2020.
[28] Qiao Liu, Jiaze Xu, Rui Jiang, and Wing Hung Wong. Density estimation using deep generative neural networks. Proceedings of the National Academy of Sciences, 118(15)ee210 10133411111111111 18, 2021.
[29] R Daniel Mauldin, William D Sudderth, and Stanley C Williams. Pólya trees and random distributions. The Annals of Statistics, pages 1203-1221, 1992.
[30] Shaun McDonald and David Campbell. A review of uncertainty quantification for density estimation. Statistics Surveys, 15:1 - 71, 2021.
[31] Jaouad Mourtada, Stéphane Gaïffas, and Erwan Scornet. AMF: Aggregated Mondrian forests for online learning. Journal of the Royal Statistical Society Series B: Statistical Methodology, 83(3):505-533, 2021.
[32] Peter Müller, Fernando Andrés Quintana, Alejandro Jara, and Tim Hanson. Bayesian Nonpara- metric Data Analysis, volume 1. Springer, 2015.
[33] Eric Nalisnick and Padhraic Smyth. Stick-breaking variational autoencoders. In International Conference on Learning Representations, 2017.
[34] Radford M Neal. Bayesian Learning For Neural Networks, volume 118. Springer Science & Business Media, 2012.
[35] Luis E Nieto- Nieto-Barajas and Peter Müller. Rubbery Pólya tree. Scandinavian Journal of Statistics, 39(1):166-184, 2012.
[36] Junier Oliva, Avinava Dubey, Manzil Zaheer, Barnabas Poczos, Ruslan Salakhutdinov, Eric Xing, and Jeff Schneider. Transformation autoregressive networks. In Proceedings of the 35th International Conference on Machine Learning, volume 80, pages
2018.
[37] Peter Orbanz. Projective limit random probabilities on polish spaces. Electronic Journal of Statistics, 5:1354 1373, 2011.
[38] Susan M Paddock, Fabrizio Ruggeri, Michael Lavine, and Mike West. Randomized Pólya tree models for nonparametric Bayesian inference. Statistica Sinica, pages 443-460, 2003.
12

---

[39] George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. In Advances in Neural Information Processing Systems, volume 30. 2017.
[40] Carl Edward Rasmussen and Christopher KI Williams. Gaussian Processes for Machine Learning, volume 1. Springer, 2006.
[41] Subham Sekhar Sahoo, Aaron Gokaslan, Christopher De Sa, and Volodymyr Kuleshov. Diffu- sion models with learned adaptive noise. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[42] Sebastian Salazar. VaRT: Variational regression trees. Advances in Neural Information Process- ing Systems, 36, 2024.
[43] Jasper Snoek, Hugo Larochelle, and Ryan P Adams. Practical Bayesian optimization of machine learning algorithms. Advances in Neural Information Processing Systems, 25, 2012.
[44] Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. In Advances in Neural Information Processing Systems, volume 34, pages 1415-1428, 2021.
[45] Ba-Hien Tran, Simone Rossi, Dimitrios Milios, and Maurizio Filippone. All you need is good functional prior for Bayesian deep learning. Journal of Machine Learning Research, 23(74): 1-56, 2022.
[46] Benigno Uria, lain Murray, and Hugo Larochelle. A deep and tractable density estimator. In Proceedings of the 31st International Conference on Machine Learning. volume 32, pages 467-475, 2014.
[47] Chenyang Zhang and Guosheng Yin. Bayesian nonparametric analysis of restricted mean survival time. Biometrics, 193))1683-1336, 2023.
[48] Luping Zhao, Timothy E Hanson, and Bradley P Carlin. Mixtures pólya trees for flexible spatial frailty survival modelling. Biometrika, 96(2):263-276, 2009.
13

---

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 abstract and introduction clearly state and accurately reflect our novel contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
We include the discussion of the potential limitations of our work in the conclusions section.
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect 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 low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails 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 faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification: All theoretical results are accompanied by clearly stated assumptions and complete proofs in the appendix.
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 the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results 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)?
Question:
Answer: [Yes]
Justification: We provide detailed descriptions of data processing and bootstrap- based confidence interval computations to enable full reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is 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 novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all 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 new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or 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 or verifying the results.
5. Open access to data and code
15

---

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: [Yes]
Justification: All code and data are provided in the anomonymized repository, with clear instructions for reproduction.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://niis.cc.
aublic/guudees//ddeettbiiss www.eeeeuuuuissioooooooooooo 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
/ wwtps/cccppplgguuuaaaeoeae
Cttelee//ooSSSSSeSSS.ee.a www.ssssiicggyooo 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: We specify data splits, hyperparameters, optimizers, hardware details, in the Implementation Details section.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to 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, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification:
We report the standard deviations over random runs as described in the main
text.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoo bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
16

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог 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 text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We detail GPU types, memory for all experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https://neurips: http://wwrrrspcccpppppeueeeeeeet cc/ppbblcc/thicccuciieeeeeeeeeee
Answer: [Yes]
Justification: guidelines.
We confirm that our research adheres fully to the NeurIPS Code of Ethics
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS 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 sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discuss positive applications in interpretable deep learning models and uncertainty quantification.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
17

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is 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 intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
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)?
Question:
Answer: [NA]
Justification: We determine that our work poses no high-risk assets requiring special release safeguards.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging. and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We credit all reused code and datasets with explicit citations.
Guidelines:
The answer NA means that the paper does
existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
18

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/dataseta
datasets.papersstt
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 has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [Yes]
Justification:
The code of our model is documented with usage details and licensing infor- mation.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Not applicable as our work does not involve human subject research or crowdsourced data.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information 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 in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
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?
Question:
Answer: [NA]
Justification:
Not applicable since our paper does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
19

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
Not applicable as we did not use LLMs in the core method development.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should not be described.
20

---

A Distributions and KL-Divergences
Beta distribution. Let Beta(a, 3) and Beta( 8) with a, 3, 6 O, and let p and denote their densities. The Kl-diivegeecc between the two Beta distributions is
((,,8) (a ()-()((()))))))))) v(a B)] + (8 (){{(3) vaa + B)], B(a,,3) 3)
KL || q) log
where (...)) is the Beta function and (-) is the digamma function. For X ~ Beta(a, (), the variance is
3 (a 3) + 3+ 1)
Var[X]
Multivariate Gaussian Distribution
The density of multivariate Gaussian distribution is defined as
1 (x p1) - m1). 2
p(x: μ. Σ)
exp
(2
where Rº is p-dimensional mean vector and E E R is the covariance matrix. The KL-divergence between two multivariate normal distributions N( Σn) and N ( Σ2)
Σg Σ1
KL(N(H1, Ei) Ez)) = log
p tr
E + (#2 m1) T Σg (M2 - m1)
B Experiment Details
The VPT is implemented with PyTorch. All experiments are conducted on a single RTX-3090 GPU.
Density estimation with real data.
The details of the dataset used in this experiment are summarized Table
In this experiment, we use Block-NAF [6] as the backbone feature learning network for VPT. We train the backbone network using Adam with Polyak averaging and apply an exponentially de- caying learning rate schedule, starting at the learn- ing rate of 10-2 with decay rate of À 0.5 and patience of 20 epochs for no improvement. In contrast, our VPT is trained with a constant learning rate of 0.1. All models are trained until convergence, with a maximum of 1, 000 epochs, stopping if there is improvement the valida- tion set for 100 epochs.
Table A1: Summary of five real UCI datasets.
Datasets
Dimensions
No. Samples
POWER
6
2, 049, 280
GAS HEPMASS MINIBOONE
8
1, 052, 065
21
525, 123
43
36, 488
BSDS300
63
1,300, 000
Density estimation and image generation. We use NICE [8] as the backbone network for this
experiment, employing a dequantized version the data. The architecture consists of a stack of four coupling layers, with diagonal positive scaling for the last stage. Each coupling function follows the same architecture: five hidden layers of 1,000 units for MNIST, and four layers of 2,000 units for SVHN and CIFAR-10.
The NICE models are trained with Adam with learning rate 10-3, momentum 0.9, B2 0.01, A and c = 10-4, Our VPT models are trained with Adam with learning rate 0.5.
Computational complexity analysis. Our complexity analysis shows that the computational cost of VPT scales primarily with the number of nodes per dimension, which is O(2 for tree of level L and the feature dimension of D. In our implementation, we mitigate this cost by choosing a small value of L (e.g., 2-4) and leveraging the independence across dimensions to parallelize computation. The recursive structure does introduce overhead due to tree traversal and node splitting, while these operations can be optimized via parallel processing across dimensions and nodes.
21

---

Table A2: Empirical profiling: training time per epoch for different datasets and configurations.
Datasets
No. Samples
D
Model
Tree Level (L)
Train Time/Epoch
Block-NAF
1 min 30
GAS
1, 052,, 065
8
VPT
2 min 10 sec
VPT
6
2 min 20 sec
Block-NAF
min
POWER
2, 0499 280
6
VPT
3 min 50 sec
VPT
min
NICE
30
MNIST
60, 000
784
VPT
35 sec
VPT
50 sec
NICE
50 sec
CIFARIO
60, 000
1024
VPT
2
min
VPT
1 min 20
Empirical profiling indicates that, when implemented on modern hardware, the additional overhead from VPT is modest compared to that of the normalizing flow backbone. Further optimizations- as batch processing of tree nodes, caching intermediate results, and efficient tree traversal algo- rithms-cot reduce computational burden even more. Future work will include detailed benchmark over various D and L settings to quantify these improvements.
Furthermore, the additional memory footprint introduced by VPT is negligible, as shown in Table|A3 The memory overhead remains below 0.05%, demonstrating that memory usage stays essentially constant when using VPT.
Table A3: Peak GPU memory usage (in gigabytes) during training with Adam and batch size of 128 using Block-NAF as the backbone. Compared to vanilla Block-NAF with a Gaussian prior, VPTs do not introduce any additional memory overhead.
Datasets
Gaussian Prior
VPT (L = 4)
VPT (L 6)
POWER
0.22
0.22
0.22
GAS
0.22
0.22
0.22
HEPMASS
5.00
5.00
5.00
MINIBOONE
4.04
4.04
4.04
BSD300
19.83
19.83
19.84
important to note that the dimensionality 1) referenced in the complexity analysis O typically corresponds to the latent representation rather than raw pixel dimensions. For image datasets, standard dimensionality reduction techniques such as convolutional layers effectively reduce D, making it practically manageable even for high- high-dimenssonal inputs.
C Additional Experiment Results on Density Estimation
Density estimation and image generation on the SVHN dataset.
In addition the experiments in Section
also conduct experiments the SVHN dataset.
Figure AT presents the density estimation results evaluated on the test set, while randomly generated images are shown in Figure A2 Similar to our findings with MNIST and CIFAR-10, our VPT enhances density estimation on SVHN without sacrificing perceptual quality.
presents more results on the MNIST interpolations. Similar to Figure 4 we can observe that the interpolation images with VPT are more clear and meaningful than those with the traditional Gaussian and logistic prior distributions.
MNIST interpolation. Figure A3
22

---

SVHN Density Estimation
NICE(Loggiste)
2.446
NICE(Gaussian)
2.915
VPT (level:4)
0.517
VPT (level=2)
0.547
Bits Per Dim
Figure Al: Negative log-likelihooe legelikelihodd in bits sitteeer-iimnnionn on the SVHN test set.
(a) Gaussian Prior
(b) Logistic Prior
(c) VPT (L=4) Prior
Figure A2: Randomly generated SVHN images with the NICE backbone under different priors.
Perceptual quality. Although perceptual quality not our primary focus, Figure A4]shows that ran- domly sampling from VPT prior produces images with clarity comparable to standard Gaussian logistic priors. This indicates that VPT improves log-likelihood without sacrificing image generation quality, providing advantageous balance between interpretability. uncertainty quantification, and generation performance.
3
3
3
3 9
3
3 3 3 3
3 3 3
8
8
8
8
5
6
8
8
8
3 8 3 3
3
3 8 3 9
6 6 6
5
8
8
8
7
7
7
7
7
7
7
7
7
9
9
7 7 7 4
4
@
O A
O
3
O
4
4 4 J
4
1
1
1
4
A
4
R
4 4 4
A
6
6
8
6 6
0
(a) Gaussian Prior
(b) Logistic Prior
(c) VPT (L=4) Prior
Figure A3: Interpolation on the MNIST images with the NICE backbone under different priors.
23

---

7
4 3
4
9
7
9 2
3
7
8
3
8
3
4
3
3
7
4
2 0 9 9 3 3
9
5
4
4
4
7
2
Q
9
3
4
9
9
3
3
0
8
3
6
3
0
8
8
5
5
4
3
4
7
5
8
8
9
4
7
6
3
4
5
C Ø 0
6
5
4
8
9
3
3
3
8
0
4
V
8
9 9
3
T
$
4
8
7
3
$
3
3
2
8
8
4
8
9
4
2
8
(a) Gaussian Prior
(b) Logistic Prior
(c) VPT (L=4) Prior
Figure A4: Randomly generated MNIST images with the NICE backbone under different priors.
24


Reviews:
[{"review_id": "GZaiQ1FBBC", "reviewer": "Reviewer_vPbf", "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.", "questions": "Please see Strengths And Weaknesses.", "limitations": "N/A.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "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?\n2. 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?\n3. VIDEO is biased. \n4. Line 40: MFVI not valid, but have you tried other VIs, like FFVI?\n5. Auto-regressive flows are a little bit outdated. \n6. Datasets used are not large. That may cause a good performance for BNP methods. \n7. The numerical experimental results are not enough. And some comparison with traditional methods are needed.", "was_revised": true, "final_justification": "I would like to raise the score because of the exploration of BNP novelty and authors' other replies.", "rebuttal": "Thank you for acknowledging the importance of developing novel deep Bayesian nonparametric frameworks. Below, we carefully address each of your concerns.\n\n>**Q1 Scalability**\n\nWe greatly appreciate your positive feedback on the BNP methods used in our work. Thank you also for highlighting the question of scalability, which is a natural concern for any BNP approach.\n\nTable 4 in our paper shows that VPT’s runtime overhead is modest: on the largest task (CIFAR-10) a depth-4 VPT adds only ≈ 0.6× of the original training time on an RTX 3090; on UCI datasets the extra cost is just 0.3–0.4× of the original training time.. These numbers use no parallelizations. As each VPT dimension is independent, so simple tensor-sharding can further trim runtime dramatically.\n\nScalability is not our main focus; our core contribution is a principled bridge between classical BNP and modern deep learning. Nothing prevents VPT from plugging into larger backbones: Appendix D already shows it working as a latent prior in a convolutional VAE with no memory issues, and we plan to extend it to transformer-based flows next.\n\nThese results show that VPT scales in practice and is not limited to small datasets or simple models. We chose modest backbones deliberately to isolate VPT’s benefits—robustness, calibrated uncertainty, and interpretability—without the confounding power of very large networks. Exploring larger architectures is ongoing work.\n\n>**Q2 Worth of BNP methods**\n\nWe thank the reviewer for this insightful question. While the Dirichlet Process (DP) does enjoy the **tail-free property**—inference remains invariant to partitioning choices—this comes at the cost of generating **discrete measures almost surely** (even if centered at a continuous base) [1]. In contrast, the Pólya tree (PT) can flexibly **learn data-adaptive partitions** and under proper hyperparameter settings induce **absolutely continuous densities**—a critical advantage when modeling complex continuous data [2].\n\nMoreover, Bayesian nonparametric alternatives like mixtures of Polya trees (e.g. Hanson \\& Johnson, 2002 [3]) show that PT-based methods can outperform DP mixtures when modeling regression errors or continuous residual distributions, thanks to their hierarchical smoothing and local adaptivity [4]. Our VPT method builds on these insights by integrating a **nonparametric, hierarchical prior with a modern variational inference framework**. We derive a **joint posterior objective over Beta weights** and enable **gradient backpropagation through the tree hierarchy**, preserving dependencies and smoothing across scales. These design choices go significantly beyond the generic application of BNP priors to deep models, delivering robustness, interpretability, and better uncertainty quantification in high-dimensional density estimation tasks.\n\n[1] Jara, Alejandro, and Timothy E. Hanson. \"A class of mixtures of dependent tail-free processes.\" Biometrika 98.3 (2011): 553-566.\n\n[2] Nieto‐Barajas, Luis E., and Peter Mueller. \"Rubbery polya tree.\" Scandinavian Journal of Statistics 39.1 (2012): 166-184.\n\n[3] Hanson, Timothy, and Wesley O. Johnson. \"Modeling regression error with a mixture of Polya trees.\" Journal of the American Statistical Association 97.460 (2002): 1020-1033.\n\n[4] Canale, Antonio. \"msBP: An R Package to perform Bayesian nonparametric inference using multiscale Bernstein polynomials mixtures.\" Journal of Statistical Software 78 (2017): 1-19.\n\n>**Q3 VIDEO is biased**\n\nWe are sorry for potential misunderstanding, but we did not find any contents introducing VIDEO in our article. It would be appreciated if you could provide more explanation on this point so that we can better address your concerns. \n\n>**Q4 Other potential VIs**\n\nThank you for the suggestion. While flexible variational inference (e.g., FFVI: Fixed-Form Variational Inference is a specific framework within VI where the parametric form of the variational posterior is pre-defined) can improve expressiveness, it often requires tailored probabilistic assumptions or model-specific reparameterizations, which are non-trivial to integrate with hierarchical priors like the Pólya tree. In contrast, mean-field VI is widely used due to its tractability and general applicability across models.\n\nOur approach avoids these limitations by directly optimizing the joint posterior likelihood, which captures richer dependencies across tree levels and enables more accurate posterior approximation than typical VI schemes. We agree that exploring FFVI for VPT is an interesting future direction and will discuss this in the revision.\n\n>**Q5 Autoregressive methods**\n\nWe appreciate the reviewer’s point. Our goal in using autoregressive flows is not to promote them as state-of-the-art, but rather to isolate and test the applicability of the Pólya tree prior within a standard and well-understood generative framework. By controlling for the generative backbone, we can more clearly attribute performance gains to the Pólya tree itself, rather than confounding effects from newer architectures. That said, our approach is modular and can be extended to more advanced models, which we plan to explore in future work.\n\n>**Q6&7 More Comprehensive Evaluation**.\n\nThank you for your suggestions. We include two additional numerical experiments here.\n\n**Compare our VPT prior with learnable histogram (LH) prior**\n\nWe set the number of bins of a LH to $K=2^L$ to match the number of leaves in a $L$ level VPT. For each dimension $d$, we parametrize a strictly increasing sequence $b_0^{(d)} < b_1^{(d)} < \\cdots < b_K^{(d)}$, by setting the lower boundary $b_0^{(d)}$ and $b_k^{(d)} = b\\_{k-1}^{(d)} + softplus(\\delta^{d}\\_{k-1})$. The parameters $\\delta^{d}_{k-1}$ are learned jointly with the backbone network, and $\\rm softplus(\\cdot)$ guarantees positive bin widths.\n\nGiven one dimension $x^{(d)}$, we can find the active bin index $k_*$, then compute the log-density as \n\\begin{align*}\n    \\log p(x^{(d)}) = \\log p_{k_*}^{(d)} - \\log (b_{k_*+1}^{(d)} - b_{k_*}^{(d)}).\n\\end{align*}\nThe overall log density of learnable histogram (LH) can be computed as $\\log p_{\\rm LH}(x)= \\sum^D_{d=1}\\log p(x^{(d)})$.\n\nThe following table shows the results on UCI datasets with Block-NAF backbone (same as the our paper). The learning histogram prior is used as a plug-in here. All the hyper-parameters are kept the same as those in VPT. \n\n**Table— Comparison of learnable histogram (LH) prior and VPT prior**.  Results are averages over 5 runs (mean ± std).\n| Dataset | POWER | GAS | BSD300|\n|---------|-------|-----|--------|\n|LH|0.62_(0.01)|11.00_(0.20)| 156.98_(0.93)|\n|VPT| 0.67_(0.01) |11.92_(0.10)| 158.59_(1.02)|\n\nFrom the results we can observe a clear advantage of VPT over LH, although they share the same number of parameters. Besides the results on UCI data, we also compare LH with VPT on CIFAR-10 with NICE network (same in the paper). With $L=4$, our VPT obtains a result of $2.06$ bits per dim (BPD), and LH obtain $2.57$ BPD. \n\nVPT uses a Bayesian-nonparametric Pólya-tree prior, coupling every Beta node to its ancestors. Evidence at coarse levels shrinks finer splits, so sparsely populated bins “borrow” strength from their parents and avoid over- or under-fitting. Learnable histograms lack this hierarchy and often overfit low-density regions or let empty bins collapse.\n\nUnlike MAP-learned histograms, VPT’s variational inference keeps the joint posterior over all splits, yielding a coherent, uncertainty-aware density— even with a truncated tree. The hierarchy also provides multi-scale probability mass, offering intuitive coarse-to-fine insights alongside robust density estimates.\n\n**Compare the uncertainty calibration with other uncertainty estimation methods**\nThe following is our experimental settings including the evaluation metrics and descriptions of comparable methods.\n\n**Metric**: Standardized squared error (SSE).\nFor every test example and every latent  dimension $d=1,\\ldots,D$ we compute\n$$\n z_{d}^{(n)} =\n\\frac{\\bigl(x_{d}^{(n)}-\\mu_{d}^{(n)}\\bigr)^{2}}\n          {\\sigma_{d}^{2\\,(n)}},\n   \\text{SSE}=\\frac1{ND}\\sum_{n,d} z_{d}^{(n)}.\n$$\nIf the predictive variance is perfectly calibrated, $\\mathbb E[z]=1\\Rightarrow {SSE}=1$; SSE $>1$ indicates under‐estimated variance (model over-confidence); SSE $<1$ indicates over‐estimated variance (model under-confidence).\n\n**Baselines**: We currently benchmark against two standard uncertainty methods\n- MC-dropout with a  rate of $0.15$ \\[3\\].\n\n- Variational BNN with mean-field Gaussian weights \\[4\\].\n\n[3] Yarin Gal etc. Dropout as a Bayesian Approximation. ICML 2016\n[4] Radford M. Neal. Bayesian Learning for Neural Networks. Ph.D. Thesis, 1995.\n\n|Dataset|VPT|MC-Dropout|BNN|\n|---------|-----|-----------|-----|\n|POWER|0.92|1.30|1.17|\n|GAS|0.90|1.35|1.14|\n|MNIST|0.92|1.27|1.08|\n\nThe table shows that VPT slightly over-estimates variance yet achieves the best calibration. MC-Dropout and BNN both under-estimate variance (BNN fares slightly better than MC-Dropout). We think VPT’s hierarchical shrinkage inflates variance in sparse regions, whereas MC-Dropout and mean-field BNNs fail to propagate uncertainty fully, leading to systematic under-confidence.\n\nWe also emphasize that VPT’s variances are analytic, no Monte-Carlo sampling is required at test time—whereas both baselines need $30$ forward passes per input. Thus VPT offers better calibration at virtually zero extra compute.\n\nWe will incorporate this calibration study—including additional datasets and seeds—into the revised paper to substantiate our uncertainty claims even more thoroughly.\n\nWe appreciate the reviewer’s thoughtful feedback. Our new results show that VPT scales well to larger models and highlight the benefits of its closed-form VI over other approaches. We have also added two new numerical studies, which will appear in the revision. We hope these clarifications fully address your concerns and it is much appreciated if you could consider raising your score."}, {"review_id": "BqxMd1sAms", "reviewer": "Reviewer_1hXT", "summary": "This 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.", "questions": "1) 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?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths: 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.", "was_revised": true, "final_justification": "The rebuttal addressed my questions,\n\n- (+) The computational overhead introduced is negligible.\n- (+) Clarifications address the limitations in the presence of cross-dimensional dependencies.\n- (+) They provided an experiment for calibration, comparing with the baselines.\n- (-) No experiments provided for the case with strong cross-dimensional dependencies.\n\nAs the paper has some notation inconsistencies, I would recommend a careful check by the authors. As my score is already positive, I opt to keep it as it is.", "rebuttal": "Thank you for the insightful feedback. Reviewers highlighted our solid technical execution, sound theory, and the value of preserving hierarchical dependencies in Pólya-tree variational inference, as well as our comprehensive empirical results and interpretability gains. Below, we address the remaining concerns—scalability, contribution scope, and experimental coverage.\n\n> **Q1:** The novelty feels somewhat incremental-established Bayesian nonparametric prior and plugging it into existing architectures like normalizing flows.\n\nUnlike generic applications of variational inference, our framework derives and optimizes the full joint posterior of the Beta weights, leveraging Pólya tree conjugacy (Lavine, 1992; Mauldin et al., 1992) and projectivity to maintain a closed-form structure. This makes our inference more expressive than standard mean-field approaches and specifically addresses the nonparametric nature of Pólya trees.\n\nBy integrating this variational design into deep density modeling, we obtain well-calibrated uncertainty estimates and improved generalization. Our empirical results confirm that this specialized integration—rooted in theoretical Pólya tree properties—outperforms both crude nonparametric baselines and standard flow architectures.\n\nThus, **our work is not a simple “plug-in prior,” but a rigorous bridge between classical Bayesian non-parametrics and modern deep learning, delivering robustness, interpretability, and reliable uncertainty estimates**.\n\n> **Q2 Assumption Violations:** How much performance degrades when strong cross-dimensional dependencies exist? Have you tested on datasets where this assumption clearly fails?\n\nWe appreciate the reviewer’s concern about the dimensional independence assumption. Our VPT assumes independence across dimensions due to considerations on tractability and computational costs. On the other hand, this assumption is mostly acceptable in practice: many real datasets exhibit only weak cross-dimensional correlations, and our empirical evaluations show that performance remains robust even when this assumption is not strictly true.\n\nWhen strong correlations exist, performance may degrade modestly. For example, multivariate density methods like mixture-of–Pólya-trees have been applied for ROC data and spatial modeling where dependencies are explicitly modeled through multivariate Pólya tree mixtures [1]. More sophisticated PT variants like Hidden Markov Pólya Trees adapt partitions based on latent state dependencies and improve performance under strong inter-variable coupling [2].\n\nWhile our current VPT design prioritizes computational efficiency through dimensional independence, this choice rarely impacts performance dramatically for most applied settings. In cases with strong latent dependencies, more flexible PT extensions or multivariate modeling can be considered to better capture correlations. This work is the first to validate the feasibility of scalable PT, which will motivate more interesting applications of VPT in neural networks.\n\n[1]Multivariate mixtures of Polya trees for modeling ROC data - Timothy E Hanson, Adam J Branscum, Ian A Gardner, 2008\n\n[2]Hidden Markov Pólya trees for high-dimensional distributions. Awaya and Ma. JASA. 2014\n\n> **Q3: 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\nThank you for raising this insightful question regarding computational scalability and practical limits. We provide detailed empirical results and practical clarifications below.\n\nAs detailed in Table 2 of the paper, VPT introduces a minimal number of additional parameters compared to the baseline model—at most 7,938 parameters on BSD300, representing less than $0.02\\%$ overhead.\n\nFurthermore, the additional memory introduced by VPT is negligible, as shown in Table below. The memory (in GB) overhead remains below $0.05\\%$.\n\n| Dataset   | Block–NAF | +VPT (4) | +VPT (6) |\n|-----------|-----------|----------|----------|\n| POWER|0.22| 0.22| 0.22|\n| GAS| 0.22| 0.22| 0.22|\n| HEPMASS   | 5.00| 5.00| 5.00|\n| MINIBOONE | 4.04| 4.04| 4.04 |\n| BSD300| 19.83| 19.83| 19.84 |\n\nIt's important to note that the dimensionality $d$ referenced in the complexity analysis $\\mathcal{O}(2^L d)$ typically corresponds to the **latent representation rather than raw pixel dimensions**. For image datasets, standard dimensionality-reduction techniques reduce $d$, making it **practically manageable even for high-dimensional inputs**.\n\nAdditionally, we emphasize that the tree depth $L$ can generally be small, as the method utilizes truncated trees. Indeed, as demonstrated throughout our experiments, a depth of $L\\leq6$ is sufficient to achieve excellent empirical results.\n\nFinally, since each dimension within VPT is processed independently, the computation is naturally parallelizable. If latent dimensionality becomes large, straightforward dimension-sharding across multiple GPUs can further alleviate scalability concerns.\n\n> **Q4: Architecture Dependence:** 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\nWe chose simple backbones deliberately: powerful flows can warp even a Gaussian prior into a complex shape, obscuring the added value of VPT. Using smaller, capacity-constrained networks lets us show VPT’s core benefit—better density estimates via base-distribution regularizations—without extra model complexity (<0.02 % parameters; Table 2).\n\nThis setting matters in practice—for low-data domains such as medical imaging where compact models are essential. In these cases, replacing a Gaussian prior with VPT boosts density estimates and gives far better-calibrated uncertainty with virtually no added cost.\n\nWhile large-model tests are outside our current scope, VPT is plug-and-play for any likelihood-based model. We have already swapped the VAE’s Gaussian encoder prior with VPT (see supplement). Rather than emphasizing scalability to larger backbones, we highlight our key contributions as follows:\n\n- First VI framework that makes Pólya-tree priors practical for density estimation and generative models.\n\n- Improves robustness and uncertainty calibration on tabular, image, and VAE tasks with minimal extra compute.\n\n- Provides interpretable, multi-scale mass allocations and closed-form predictive variances—capabilities missing from large black-box architectures.\n\n>**Q5: Uncertainty Calibration:** 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\nThank you for suggesting a quantitative uncertainty test and additional baselines. Below we outline the evaluation metrics, settings, and competing methods we used.\n\n**Metric**: Standardized squared error (SSE).\nFor every test example and every latent  dimension $d=1,\\ldots,D$ we compute\n$$\n z_{d}^{(n)} =\n\\frac{\\bigl(x_{d}^{(n)}-\\mu_{d}^{(n)}\\bigr)^{2}}\n          {\\sigma_{d}^{2\\,(n)}},\n   \\text{SSE}=\\frac1{ND}\\sum_{n,d} z_{d}^{(n)}.\n$$\nIf the predictive variance is perfectly calibrated, $\\mathbb E[z]=1\\Rightarrow {SSE}=1$; SSE $>1$ indicates under‐estimated variance (model over-confidence); SSE $<1$ indicates over‐estimated variance (model under-confidence).\n\n**Baselines**: We currently benchmark against two standard uncertainty methods\n- MC-dropout with a  rate of $0.15$ \\[3\\].\n\n- Variational BNN with mean-field Gaussian weights \\[4\\].\n\n[3] Yarin Gal etc. Dropout as a Bayesian Approximation. ICML 2016\n[4] Radford M. Neal. Bayesian Learning for Neural Networks. Ph.D. Thesis, 1995.\n\n| Dataset | VPT | MC-Dropout | BNN |\n|---------|-----|-----------|-----|\n| POWER   | 0.92 | 1.30 | 1.17 |\n| GAS     | 0.90 | 1.35 | 1.14 |\n| MNIST   | 0.92 | 1.27 | 1.08 |\n\nThe table shows that VPT slightly over-estimates variance yet achieves the best calibration. MC-Dropout and BNN both under-estimate variance (BNN fares slightly better than MC-Dropout). We think VPT’s hierarchical shrinkage inflates variance in sparse regions, whereas MC-Dropout and mean-field BNNs fail to propagate uncertainty fully, leading to systematic under-confidence.\n\nWe also emphasize that VPT’s variances are analytic, no Monte-Carlo sampling is required at test time—whereas both baselines need $30$ forward passes per input. Thus VPT offers better calibration at virtually zero extra compute.\n\nWe will incorporate this calibration study—including additional datasets and seeds—into the revised paper to substantiate our uncertainty claims even more thoroughly.\n\n>**Q6:** How is $\\beta$ in line 153 defined?\n\n$\\beta_j$ refers to the collection of Beta-distributed random variables (i.e., the branching probabilities) at depth level $j$ of the tree.\n\n>**Q7:** In line 167, $P(B_\\epsilon)$should be dependent on the index $L$; however, it seems indepedent of $L$, why so?\n\nWe apologize for the typo. The probability should be $P(B_{\\epsilon_{1:L}})$ instead of $P(B_{\\epsilon})$.\n\n>**Q8:** Are $B_\\epsilon$ probabilities or random variables?\n\nIn the standard formulation of the PT, $B_\\epsilon$ denotes the interval (or measurable subset) of the sample space indexed by the binary string $\\epsilon$, and $P(B_\\epsilon)$ is the random probability assigned to that interval.\n\nThank you for the helpful feedback. The new uncertainty-calibration experiments directly address your concerns and quantitatively support our claims. We will refine the presentation and clarify all noted points in the revision, and hope these improvements merit a higher rating."}, {"review_id": "fjjbPaJTYm", "reviewer": "Reviewer_Z5q1", "summary": "The 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.", "questions": "- 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).", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\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).", "was_revised": true, "final_justification": "The authors have addressed my main concerns as follows:\n\n- The ablation using learnable histograms effectively addresses my concern that the claimed 'performance' gains were solely due to a more flexible 'base distribution.'\n- The additional analysis on uncertainty quantification strengthens the paper.\n\nMy remaining concerns are related to presentation and writing. Given the authors' efforts in their rebuttal, I believe they can effectively address these issues.", "rebuttal": "We sincerely thank the reviewer for recognizing the novelty of integrating Bayesian nonparametrics into deep density estimation and for appreciating the efficiency and strong empirical performance of our approach. We would address your comments below.\n\n>**Q1:** 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\" VPT. 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. \n\nThank you for this insightful comment!\nWe agree that this work uses univariate Polya trees and demonstrates the feasibility of VPT in nueral networks. This is the first attempt in the literature to expand Polya trees with variational inference that enhances the scalability of Polya trees in high dimensions. Related research on nueral networks using Dirichlet process (DP) or mixture DP priors has been conducted extensively, but none in Polya trees. \n\nPer your suggestion, we will emphasize the application of univariate Polya trees to variational Bayes in the abstract and introduction. We will add discussion on further expanding VPT to diffusion models in the discussion, while this is non-trivial in practice due to the computational cost of recursive tree sampling and the challenge of adapting VPTs to the continuous-time dynamics of diffusion models. While promising, such integration requires careful algorithmic design\nand remains an open direction for future research.\n\n>**Q2:** 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\nDue to time limit for the rebuttal, we focus on the experimental setting with a learnable histogram (LH), while we will also add the mixture of Gaussians in the revision.\n\nWe set the number of bins to $K=2^L$ to match the number of leaves in a $L$ level VPT. For each dimension $d$, we parametrize a strictly increasing sequence $b_0^{(d)} < b_1^{(d)} < \\cdots < b_K^{(d)}$, by setting the lower boundary $b_0^{(d)}$ and $b\\_k^{(d)} = b\\_{k-1}^{(d)} + softplus(\\delta^{d}\\_{k-1})$. The parameters $\\delta^{d}_{k-1}$ are learned jointly with the backbone network, and $softplus(\\cdot)$ guarantees positive bin widths.\n\nGiven one dimension $x^{(d)}$, we can find the active bin index $k_*$, then compute the log-density as \n\\begin{align*}\n    \\log p(x^{(d)}) = \\log p_{k_*}^{(d)} - \\log (b_{k_*+1}^{(d)} - b_{k_*}^{(d)}).\n\\end{align*}\nThe overall log density of learnable histogram (LH) can be computed as $$\\log p_{\\rm LH}(x)= \\sum^D_{d=1}\\log p(x^{(d)}).$$\n\nThe following table shows the results on UCI datasets with Block-NAF backbone (same as the our paper). The learning histogram prior is used as a plug-in step here. All the hyper-parameters are kept the same as those in VPT. \n\n**Table — Comparison of learnable histogram (LH) prior and VPT prior**.  Results are averages over 5 runs (mean ± std).\n\n|Dataset|POWER|GAS|BSD300|\n|---------|-------|-----|--------|\n|LH| 0.62_(0.01)|11.00_(0.20)|156.98_(0.93)|\n|VPT| 0.67_(0.01)| 11.92_(0.10)|158.59_(1.02)|\n\nFrom the results we can observe a clear advantage of VPT over LH, although they share the same number of parameters. Besides the results on UCI data, we also compare LH with VPT on CIFAR-10 with NICE network (same in the paper). With $L=4$, our VPT obtains a result of $2.06$ bits per dimension (BPD), and LH yields $2.57$ BPD. \n\nWhile a truncated factorized VPT may resemble a learnable histogram in terms of producing a piecewise-constant density, our VPT provides the following advantages that go beyond histogram-like behaviors.\n\nFirst, unlike a learnable histogram (LH), where each bin is estimated independently, VPT leverages the Bayesian nonparametric prior and hierarchical conjugacy, enabling information sharing across scales. Specifically, every Beta-distributed node in the VPT is coupled to its ancestors, such that evidence at coarser levels regularizes finer partitions. This hierarchical shrinkage allows small bins with few data points to borrow statistical strength from parent nodes, automatically adjusting the degree of smoothing. In contrast, learnable histograms often suffer from overfitting in low-density regions or collapsing of underpopulated bins, since they lack this structural regularization.\n\nSecond, the VPT's variational posterior inference maintains the joint posterior over tree splits, rather than treating bins independently (as MAP-learned histograms do). This enables a more coherent and uncertainty-aware density estimate, especially in regions of sparse data. Even when truncated, the VPT retains these hierarchical dependencies, which are absent in flat, non-Bayesian alternatives.\n\nBesides robust density estimation, our VPT can naturally provide probability mass at multiple scales, giving intuitive `coarse-to-fine' insights. \n\n>**Q3 Uncertainty Estimation:**  Specifically, 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\nWe sincerely thank the reviewer for this valuable comment and the recommended references. We will include citations of these relevant works in our revised manuscript.\n\nBelow, we provide a clear description of our experimental setting, evaluation metrics, and comparisons with alternative uncertainty estimation methods to address the reviewer's concerns thoroughly.\n\n**Metric**: Standardized squared error (SSE).\nFor every test example and every latent  dimension $d=1,\\ldots,D$ we compute\n$$\n z_{d}^{(n)} =\n\\frac{\\bigl(x_{d}^{(n)}-\\mu_{d}^{(n)}\\bigr)^{2}}\n          {\\sigma_{d}^{2\\,(n)}},\n   \\text{SSE}=\\frac1{ND}\\sum_{n,d} z_{d}^{(n)}.\n$$\nIf the predictive variance is perfectly calibrated, $\\mathbb E[z]=1\\Rightarrow {SSE}=1$; SSE $>1$ indicates under‐estimated variance (model over-confidence); SSE $<1$ indicates over‐estimated variance (model under-confidence).\n\n**Baselines**: We currently benchmark against two standard uncertainty methods\n- MC-dropout with a  rate of $0.15$ \\[3\\].\n\n- Variational BNN with mean-field Gaussian weights [4].\n\n[3] Yarin Gal etc. Dropout as a Bayesian Approximation. ICML 2016\n\n[4] Radford M. Neal. Bayesian Learning for Neural Networks. Ph.D. Thesis, 1995.\n\n| Dataset | VPT | MC-Dropout|BNN|\n|---------|-----|-----------|-----|\n|POWER|0.92|1.30|1.17|\n|GAS| 0.90|1.35|1.14|\n|MNIST| 0.92 | 1.27 | 1.08 |\n\nFrom the above table, we observe that our VPT slightly over-estimates its variance, while keeping the best calibration value. MC-dropout and BNN underestimate the variances, while BNN have better calibration. The hierarchical structure of VPT updates the shrink leaf probabilities toward their parents when data are sparse. This automatically inflates predictive variance in low-count regions. MC-Dropout and mean-field BNNs variances often fail to fully propagate to the output, hence leading to the systematic under-estimation.\n\nWe also emphasize that VPT’s variances are analytic, and thus no Monte-Carlo sampling is required at test time—whereas both baselines need 30 forward passes per input.\nTherefore, VPT offers better calibration at virtually zero extra computational cost.\n\nWe will incorporate this calibration study—including additional datasets and seeds—into the revised paper to substantiate our uncertainty claims more thoroughly.\n\n>**Q4:** Can the authors go into more details on how the VPT variational posterior distribution is parameterized?\n\nIn our VPT framework, we parameterize the variational posterior by assigning a Beta distribution to each node in the tree. Specifically, for each split $\\epsilon$, we define $q(B\\_\\epsilon) = \\text{Beta}(\\hat{\\alpha}\\_{\\epsilon_{1:j-1}0}, \\hat{\\alpha}\\_{\\epsilon_{1:j-1}1})$, where the shape parameters $\\hat{\\alpha}\\_{\\epsilon_{1:j-1}0}$ and $\\hat{\\alpha}\\_{\\epsilon_{1:j-1}1}$ are output by a neural network conditioned on the input. This network takes the data point $x$ as input and returns parameters for each node along the path in the tree. Sampling from these Beta distributions recursively yields a full sample from the variational posterior over the hierarchical partitions.\n\n>**Q5:** Paper formatting concerns. \n\nThank you for your comments on paper formatting. We will resolve the above issues in the revised version of our manuscript.\n\nWe will trim down Sec. 1 and 2 by providing a more concise literature review.\nWe will remove Fig. 1, and cut the redundancy in Sec. 2.2. We will also rearrange Algorithm 1 in a proper place to have a better flow. This will give us some room to add all the new experiments as suggested by you, which will greatly strengthen the paper.\n\nThank you again for the insightful feedback. The new experiments and clarifications strengthen the paper in two key ways:\n\n- Learnable-histogram comparison: With equal parameter counts, VPT captures distributional structure markedly better.\n\n- Uncertainty calibration: VPT yields more accurate predictive variances than MC-Dropout and BNN baselines, at no extra computational cost.\n\nWe will add these results and clarifications to the revised manuscript for greater clarity and completeness. We sincerely hope the above discussion would warrant a higher score."}, {"review_id": "9YxsFQ4B4N", "reviewer": "Reviewer_Fojs", "summary": "This 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.", "questions": "1. 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?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "All of my doubts (see Weaknesses) have been sufficiently addressed by the authors' rebuttal and by seeing other reviews/discussions. I urge the authors to incorporate all of the suggestions.", "rebuttal": "Thank you for your positive assessment of VPT’s novelty, tractability, interpretability, and efficiency, and for noting its seamless integration with flow models and VAEs. We also appreciate your praise for our comprehensive experiments. Below, we address your remaining comment.\n\n>**Q1:** For example, it is unclear to me what it means by \"... parametrize tree partitions \n by learnable neural networks. ...\". Are the partitions' endpoints being predicted by NNs? Or something else?\n\nThank you for the question. We do not directly predict the partition endpoints with neural networks. Instead, we parameterize the shape parameters of the Beta distributions at each node in the tree using neural networks. These Beta distributions determine the probability mass split at each partition, which in turn defines the tree-based variational posterior over the sample space. We will clarify this in the revision.\n\n>**Q2:** 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\nWe appreciate the reviewer's observation. In our formulation, we focus on optimizing the posterior likelihood, as reflected in Algorithm 1. While we do not explicitly include a KL term in the ELBO, the corresponding entropy of the variational posterior acts as an implicit regularization. This entropy term, which appears in the negative log-likelihood, serves a similar role by controlling the spread of the variational distribution and balancing data fit versus uncertainty.\n\nThus, the loss function in our implementation prioritizes posterior inference while still benefiting from the regularizing effect of entropy, effectively aligning with the ELBO framework. We agree that making this connection more explicit can improve clarity and will revise the description accordingly.\n\n>**Q3:** I would suggest the authors rewrite this section.\n\nThank you for your suggestion. Sorry for the confusion. We will rewrite this section by adding more detailed, step-by-step explanations on VPT. \n\n>**Q4:** Are VPTs compatible (not just theoretically, but in terms of costs and practicality) with large-scale generative models like diffusion models?\n\nYes, VPTs can in theory be integrated into diffusion models by replacing the Gaussian prior with a VPT-based prior, potentially improving uncertainty modeling. However, this is non-trivial in practice due to the computational cost of recursive tree sampling and the challenge of adapting VPTs to the continuous-time dynamics of diffusion models. While promising, such integration requires careful algorithmic design and remains an open direction for future research.\n\nWe sincerely thank you for your generous evaluation. Your recognition of the work’s novelty, tractability, and empirical rigor is greatly appreciated and encourages us to further refine and extend this research."}]

Author rebuttal:

Meta-review: The reviewers agree this paper has many interesting ideas: Polya trees are not as popular as other BNP methods in the NeurIPS community, the idea of fusing ideas between Bayesian nonparametrics and deep learning is interesting. However, there were also a number of concerns raised, listed below. Overall though, the reviewers believe that these can be addressed while preparing the camera ready version. Please look at these and well as you discussion with the reviewers when preparing the camera ready version:

1) Clarity: The presentation of Polya trees is very brief and does not clearly articulate their benefits (e.g. obscuring the usefulness of the proposed variational approximation). Additionally, it seems to have a number typos/unclear notation. E.g. what is the difference between subscripts and superscripts for the betas?  I also think the paper oversells Polya trees a little. It mentions that Dirichlet processes "induce discreteness, making them suboptimal for directly modeling continuous distributions". Of course DP mixture models fix this. By contrast, the Polya tree approach here seems truncated at level L, which introduces its own issues. I think this warrants more discussion.
3) Claims vs results about interpretability and uncertainty quantification: The first half of the paper talks about the interpretability of the proposed method about 5 times, but the presented results are underwhelming. Reviewers were also not very satisfied with claims about uncertainty quantification. With their rebuttal, the authors left the reviewers feeling much better about these points, so please incorporate the latter into the final version
3) Scalability and empirical evaluation: this should be discussed more thoroughly, especially since this paper focuses (without being clear about this) on univariate Polya trees
Decision: Accept (poster)