Maximum a Posteriori Estimation by Search in Probabilistic Programs

We introduce an approximate search algorithm for fast maximum a posteriori probability estimation in probabilistic programs, which we call Bayesian ascent Monte Carlo (BaMC). Probabilistic programs represent probabilistic models with varying number of mutually dependent finite, countable, and continuous random variables. BaMC is an anytime MAP search algorithm applicable to any combination of random variables and dependencies. We compare BaMC to other MAP estimation algorithms and show that BaMC is faster and more robust on a range of probabilistic models.

💡 Research Summary

This paper introduces a novel approximate search algorithm for performing fast Maximum a Posteriori (MAP) estimation within probabilistic programs, named Bayesian Ascent Monte Carlo (BaMC). The authors address the challenge of finding the single most probable assignment to all random variables in complex models, a problem central to many AI tasks like planning and abductive reasoning. While probabilistic programming offers unparalleled expressiveness for defining such models, existing inference methods like PMCMC are geared towards characterizing the full posterior distribution rather than efficiently finding its mode.

BaMC is presented as an anytime search algorithm inspired by Monte Carlo Tree Search (MCTS), but critically adapted for the MAP estimation context. Its core innovation lies in two areas. First, it re-purposes the MCTS philosophy for a “simple reward” setting, where the goal is the quality of the final output (the MAP estimate) rather than the cumulative reward of all samples. Second, it introduces a novel selection mechanism called Open Randomized Probability Matching (ORPM). ORPM generalizes Thompson Sampling to handle variables with infinite or unknown support (e.g., continuous distributions). It maintains Bayesian beliefs about the “reward” (contribution to the total log-weight of the trace) of each previously tried value for a variable. At each selection point, ORPM treats “drawing a new random value from the prior” as a viable action alongside the known choices. It uses Thompson Sampling twice: first to estimate the potential of this random draw, and then to select among all options (known and random). This elegantly balances exploration of new, high-potential regions of the variable’s domain with exploitation of values known to yield high rewards.

The algorithm works by repeatedly executing the probabilistic program. For each random variable encountered, it uses ORPM to select a value. After a full trace is generated, its total log-weight (proportional to unnormalized posterior probability) is computed. If this weight exceeds the previous best, the trace is output as an improved MAP estimate. Finally, the reward beliefs for each variable in the trace are updated in reverse order using the partial log-weight observed after that variable’s selection. Beliefs are modeled as normal distributions based on sample mean and variance, using an uninformative prior.

The empirical evaluation compares BaMC against two tuned versions of Simulated Annealing (with exponential and Lundy-Mees schedules) and a Lightweight Metropolis-Hastings baseline adapted for MAP search. Tests are conducted on two non-trivial models involving mixed discrete and continuous variables: a Hidden Markov Model with unknown transition probabilities and a Probabilistic Deterministic Infinite Automata applied to text modeling. Results over 50 runs show that BaMC consistently finds higher probability MAP estimates faster and with lower run-to-run variance than the alternatives, even when SA is given its optimally tuned annealing rate. A visualization of a single BaMC run demonstrates how the algorithm’s samples progressively concentrate around high-probability regions.

In conclusion, BaMC is a parameter-free, adaptive, and robust search algorithm for MAP estimation. It successfully transplants ideas from MCTS into the probabilistic programming inference domain, overcoming challenges related to variable types and the nature of the objective through the novel ORPM mechanism. The results demonstrate its practical superiority over established methods like Simulated Annealing for this specific task.