Variational Polya Tree

2025-10-26

AI Summary

Existing generative models for density estimation suffer from limited interpretability and difficulty in quantifying uncertainty, while traditional Bayesian nonparametric methods—such as Markov chain Monte Carlo (MCMC)—exhibit high computational cost and poor scalability to deep learning frameworks. To address these challenges, we propose the Variational Polya Tree (VPT) model: the first framework that tightly integrates the Polya tree Bayesian nonparametric prior with stochastic variational inference, yielding a differentiable, end-to-end trainable density estimator. By optimizing the joint distribution likelihood, VPT achieves more accurate posterior approximation and overcomes the efficiency bottlenecks inherent in MCMC. Empirically, VPT scales effectively to high-dimensional data and attains competitive performance with state-of-the-art deep density estimators on both real-world tabular and image data. Crucially, it significantly enhances model interpretability and enables principled uncertainty quantification. The implementation is publicly available.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Relational Probabilistic Models
• Computer Vision: Large Vision Models

Application Category

• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods, algorithms and applications for the development of semantic models, knowledge graphs and other forms of structured data models with machine-interpretable semantics

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.

Problem

• Developing interpretable density estimation with uncertainty quantification
• Overcoming computational limitations of Bayesian nonparametric methods
• Providing flexible nonparametric priors compatible with gradient optimization

Innovation

• Variational Polya Tree model with stochastic variational inference
• Flexible nonparametric Bayesian prior for latent densities
• Joint distribution likelihood for precise posterior approximation