Variational Polya Tree

TL;DR

The paper tackles scalable, interpretable density estimation with calibrated uncertainty by integrating a continuous Polya tree prior into deep generative models through variational inference. The Variational Pólya Tree (VPT) preserves the hierarchical, Beta-distributed structure of the PT while parameterizing splits with neural networks, enabling exact posterior updates and end-to-end training with flows and VAEs. Empirical results across synthetic, real, and image data show improved density estimation, competitive or superior log-likelihoods, and meaningful uncertainty quantification, with the added benefit of interpretable latent structure. This approach offers a principled Bayesian nonparametric prior that complements existing deep density estimators and scales to large datasets with minimal overhead.

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 \polya 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 \polya 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.

Figures (11)

• Figure 1: Graphical illustration of the Pólya tree construction. The PT prior randomly splits $B_{\epsilon_{1:j}}$ into two subintervals $B_{\epsilon_{1:j} 0}$ and $B_{\epsilon_{1:j} 1}$ with Beta-distributed probabilities.
• Figure 2: Density estimation results of an isotropic Gaussian prior, a 2-level VPT, and a 3-level VPT with 2D synthetic datasets.
• Figure 3: Negative log-likelihood in bits-per-dimension on the MNIST and CIFAR-10 test sets. Our method, VPT (shown in 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.
• 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.
• Figure 5: Posterior variance of the trained 4-level VPT model on the MNIST dataset. On each dimension/pixel ($d=1,\dots,D$, where $D=28\times28$), we visualize the mean over the variances of terminal node Beta distribution.