Tree Bandits for Generative Bayes

TL;DR

This work reframes likelihood-free ABC as a bandit-style learning problem by partitioning the parameter space into boxes and treating each box as an arm. An inner loop uses Thompson Sampling to learn efficient ABC proposals within a fixed partition, while an outer loop adaptively refines the partition itself, yielding ABC-Tree for posterior sampling and ABC-MAP for likelihood-free MAP estimation. The approach comes with theoretical regret guarantees and practical demonstrations on tasks like masked MNIST image classification, showing substantial reductions in simulator calls while maintaining accurate posterior approximations. The combination of recursive partitioning with bandit-based proposal learning offers a scalable path for high-dimensional, simulator-based Bayesian inference. The methods leverage adaptive discretization (CART, BART, or dyadic partitions) and regularized exploitation to balance exploration and sampling efficiency in complex, likelihood-intractable settings.

Abstract

In generative models with obscured likelihood, Approximate Bayesian Computation (ABC) is often the tool of last resort for inference. However, ABC demands many prior parameter trials to keep only a small fraction that passes an acceptance test. To accelerate ABC rejection sampling, this paper develops a self-aware framework that learns from past trials and errors. We apply recursive partitioning classifiers on the ABC lookup table to sequentially refine high-likelihood regions into boxes. Each box is regarded as an arm in a binary bandit problem treating ABC acceptance as a reward. Each arm has a proclivity for being chosen for the next ABC evaluation, depending on the prior distribution and past rejections. The method places more splits in those areas where the likelihood resides, shying away from low-probability regions destined for ABC rejections. We provide two versions: (1) ABC-Tree for posterior sampling, and (2) ABC-MAP for maximum a posteriori estimation. We demonstrate accurate ABC approximability at much lower simulation cost. We justify the use of our tree-based bandit algorithms with nearly optimal regret bounds. Finally, we successfully apply our approach to the problem of masked image classification using deep generative models.

Figures (20)

• Figure 1: The nested structure of ABC-Tree. The outer loop actively learns the optimal partition and the inner loop actively learns the optimal proposal distribution on the current partition. Red dots are accepted proposals.
• Figure 2: Evolution of the proposal distributions in Algorithm \ref{['alg:seqtreeabc']} (partitioning with CART). The algorithm decreases $\varepsilon$ and re-partitions $\Theta$. Within outer loops (each row), the partition and $\varepsilon$ stay fixed, while the weight for each box is updated after every inner loop iteration.
• Figure 3: Empirical regret comparison. We compare various methods for sequentially updating proposals in a finite parameter space example. See Example \ref{['ex:1']} for details.
• Figure 4: Fixed grid. Performance of Algorithm \ref{['alg:qe_abc']} on a univariate Gaussian mixture example using a regular grid of various sizes for partitioning $\Theta$. Left: kernel density estimates from accepted samples. Right: Effect of bin size on accuracy of posterior sampling.
• Figure 5: Sequential refinements of the proposal distribution using various tree-based partitioning methods on the Gaussian mixture example discussed in Example \ref{['ex:2']}. Top: CART. Middle: BART. Bottom: Dyadic partitioning.

Theorems & Definitions (22)

• Lemma 1
• Remark 1
• Theorem 3.1
• Example 1
• Example 2
• Example 3
• Theorem 4.1
• Theorem A.1
• proof
• Lemma 2