A* Sampling

Chris J. Maddison; Daniel Tarlow; Tom Minka

Source-linked AI summary

A* Sampling

Chris J. Maddison, Daniel Tarlow, Tom Minka

arXiv:1411.0030v2Oct 31, 2014stat.COstat.ML

TL;DR

Exact sampling from arbitrary continuous distributions is recast as optimization over a Gumbel-perturbed energy landscape. The paper constructs a Gumbel process and uses A∗ search to find its maximum, reporting more efficient use of bound and likelihood evaluations than closely related adaptive rejection methods. The approach remains limited in high dimensions and requires tractable initialization assumptions in its top-down construction.

Problem
The paper addresses how to obtain exact, generic samples from arbitrary probability distributions for reliable probabilistic inference.
Method
It constructs a Gumbel process and applies A∗ search to optimize perturbed continuous-space energies without instantiating all perturbations.
Results
A∗ sampling uses bound and likelihood evaluations more efficiently than the most closely related adaptive rejection sampling-based algorithms.
Takeaways & Limitations
The method provides a practical generic route to exact continuous sampling by combining Gumbel-process perturbations with region-based search bounds.
Takeaways & Limitations
High-dimensional problems remain unresolved, and the top-down construction assumes log μ(Ω) is efficiently computable.

Abstract

from arXiv · show

The problem of drawing samples from a discrete distribution can be converted into a discrete optimization problem. In this work, we show how sampling from a continuous distribution can be converted into an optimization problem over continuous space. Central to the method is a stochastic process recently described in mathematical statistics that we call the Gumbel process. We present a new construction of the Gumbel process and A* sampling, a practical generic sampling algorithm that searches for the maximum of a Gumbel process using A* search. We analyze the correctness and convergence time of A* sampling and demonstrate empirically that it makes more efficient use of bound and likelihood evaluations than the most closely related adaptive rejection sampling-based algorithms.

1 Introduction

The paper develops exact generic sampling by reframing continuous sampling as optimization over randomly perturbed energies. It introduces the Gumbel process and A∗ sampling to avoid instantiating infinitely many perturbations.

Motivation: Exact generic sampling is motivated by the need for reliable inference from arbitrary probability distributions.The paper situates sampling in statistics and machine learning, including probabilistic-model training, evaluation, and prediction.
Motivation: The Gumbel-Max trick samples discrete Gibbs distributions by adding independent Gumbel perturbations and returning the perturbed argmax.Prior structured-space approaches used approximate sampling because exact methods appeared to require exponentially many perturbations.
Core idea: The key observation is that irrelevant perturbations can be bounded while relevant ones are instantiated, yielding the argmax and an exact sample.This avoids instantiating all potentially exponentially many perturbations and extends the idea to continuous spaces with infinitely many perturbations.
Core idea: The Gumbel process generalizes independent Gumbel perturbations over space, enabling a continuous analogue of the Gumbel-Max trick.It addresses how to perturb space independently in a mathematically meaningful way.
Core idea: A∗ sampling uses a top-down Gumbel-process construction to sample arbitrary continuous distributions and studies its computational efficiency against adaptive rejection methods.The paper analyzes bound and likelihood evaluations across illustrative and challenging problems.

2 The Gumbel Process

The Gumbel process extends the discrete Gumbel-Max construction to measurable continuous spaces. Its marginal, independence, and consistency properties make the argmax distributed according to the underlying measure.

Discrete foundation: The discrete Gumbel-Max trick samples class i with probability proportional to exp(φ(i)) by maximizing φ(i) plus independent Gumbel noise.The maximum has a Gumbel distribution whose location is the log partition function, while max and argmax are independent.
Continuous generalization: For continuous densities p(x) ∝ exp(φ(x)), the desired process has maxima distributed by log ∫_B exp(φ(x)) and argmaxes distributed proportionally to exp(φ(x)).The Gumbel process is introduced to satisfy these continuous-space analogues.
Definition: A Gumbel process assigns each measurable set B a Gumbel random variable with location log μ(B).This is the process’s marginal-distribution condition.
Definition: Gumbel variables associated with disjoint sets are independent, providing the process’s spatial independence property.The definition states independence between a measurable set and its complement.
Consequences: Together, the process properties make the argmax restricted to B′ distributed according to μ normalized on B′.The marginal condition controls maxima, while consistency and independence support the argmax distribution.

3 Top-Down Construction for the Gumbel Process

The top-down construction samples a Gumbel process by generating maxima and locations recursively from a heap-like structure. Truncation and partitioning make this descending construction applicable even to infinite spaces.

Heap interpretation: For finite independent Gumbels, the heap can be instantiated from root to leaves by sampling each node conditional on its parent.The root provides the maximum and argmax, while recursive child sampling reconstructs the independent set.
Algorithm: Algorithm 1 initializes the space and queue, then repeatedly samples truncated child Gumbels and locations before yielding each pair.The procedure maintains a descending heap of sampled values and associated locations.
Algorithm: Child maxima can be sampled efficiently by truncated-Gumbel CDF inversion instead of rejection sampling.Independent child argmaxes are sampled from their respective partitions, and partitioning strategy does not alter the result.
Infinite spaces: The basic procedure therefore generates an infinite descending heap of Gumbel values and locations rather than explicitly representing the full continuous perturbed density.The heap perspective connects the finite construction to the continuous setting.
Top-down construction: The construction samples a global Gumbel maximum and its location, removes that location, partitions the remaining space, and recurses.Child maxima are sampled from Gumbel distributions truncated at the parent maximum.
Infinite spaces: The construction instantiates the perturbed function at countably many points while preserving each set’s maximum once its first point is encountered.This descending-path property supplies the process interpretation for an infinite space.

4 A∗Sampling

A∗ sampling turns the Gumbel-Max objective into a bounded search over regions, using a tractable component and bounds on the residual log density. The search returns exact samples without requiring a tractable total measure.

Algorithm: A∗ sampling removes the top-down construction’s requirement that log μ(Ω) be computed efficiently.It executes the Gumbel-Max trick by exploiting properties of the Gumbel process.
Process transformation: Adding a bounded log-likelihood difference to prior-process Gumbels transforms them into a posterior Gumbel process.This supports using a tractable prior as i(x) and the bounded likelihood contribution as o(x).
Algorithm: The target log density is decomposed into tractable i(x) and boundable residual o(x), with the sampler searching for argmax{G_k + o(X_k)}.The tractable component acts analogously to a rejection sampler’s proposal distribution.
Search bounds: A∗ search expands regions according to upper bounds while lower bounds come from sampled Gumbel maxima and residual evaluations.Queue priorities use region upper bounds to guide increasingly refined spatial search.
Correctness: When the current lower bound exceeds every queued upper bound, the selected location is an exact sample from p(x) ∝ exp(φ(x)).The same condition also yields an exact Gumbel sample for the global maximum; correctness and termination are proved in the Appendix.
Variants: Bounds can be reused across runs when multiple samples are desired, with lazy A∗ maximization over partitioned regions.Only regions needed to determine the global maximum are expanded.

5 Comparison to Rejection Samplers

A* sampling matches standard rejection sampling with a global bound, while refined bounds let it focus search and outperform OS* in computation. Its refinement is coupled to the regions most relevant for accepting the current sample.

With one global bound M across nodes, A* sampling has the same runtime distribution as standard rejection sampling.The iteration count follows a Geometric distribution with rate µ(R^d)/(exp(M)ν(R^d)).
OS* selects regions by proposal volume, whereas A* selects the region with the highest upper bound.This region-selection difference determines where refinement is concentrated during search.
A* couples refinement to the eventual accepted sample and can prune regions that cannot produce the returned sample.OS* remains blind to the eventual sample location and can waste computation refining irrelevant regions.
Across operating points, A* consistently uses fewer bound and likelihood evaluations than both OS* refinement strategies.The comparison reports computational costs for the algorithms across the operating points shown in Figure 3.

6 Experiments

The experiments evaluate A* sampling across controlled peakiness, dimensionality, bounding strategies, generic interval bounds, and robust regression. Across these settings, the method remains tractable in selected multimodal problems and often reduces computation relative to rejection sampling or OS*.

Scaling versus Peakiness and Dimension: As the target becomes peakier, A* sampling mainly pays the cost of locating the peak, which is essentially binary search.The experiment uses p(x) = exp(−x)/(1+x)^a with exp(−x) as the proposal and averages results over 1000 runs.
Scaling versus Peakiness and Dimension: 900 likelihood evaluations at D = 3 and 4000 at D = 4 keep the clutter problem reasonably tractable despite exponential growth with dimension.The analogous OS* algorithm requires 16% to 40% more computation on average.
Bounding Strategies: For large N, linear bounds multiply evaluations by 3, while constant bounds multiply evaluations by O(N), relative to tight bounds.The experiment compares constant, linear, and quadratic bounds for estimating a one-dimensional Gaussian mean.
Using Generic Interval Bounds: Automatically constructed interval bounds support A* sampling for nonlinear regression models with multimodal posteriors and varying observation noise.The models are expressed symbolically, and FuncDesigner computes the bounds used by the samplers.
Robust Bayesian Regression: In robust Bayesian regression, A* uses fewer bound and likelihood evaluations than both OS* refinement strategies across operating points.The single-sample experiment compares variants trading likelihood computations for bound computations.
Robust Bayesian Regression: For 200 posterior samples, OS* uses fewer likelihood evaluations only by incurring more bound computations, while A* reaches a computational operating point OS* cannot achieve.A slice sampler given ten times A*’s budget had difficulty mixing and rarely switched between posterior modes.

7 Discussion

The discussion frames A* sampling as a continuous-space extension of the Gumbel-Max idea, while acknowledging unresolved high-dimensional limitations and several future directions.

The paper asks whether the Gumbel-Max trick can be extended to continuous spaces and used for tractable continuous-distribution sampling.
Figure 3 compares A* and OS* computational costs across Cauchy regression dimensions, refinement strategies, and refinement rates.Circles denote A*, while squares and diamonds denote two OS* strategies; marker colors encode refinement rates.
Perturb-and-MAP methods motivate continuous analogs that may trade exactness for more efficient computation through principled approximate perturbations.
A* does not solve high-dimensional sampling: bounds can become uninformative, leaving little gain over vanilla rejection sampling.The paper identifies adapting the split between i(·) and o(·) to each node as an open question.
Future work includes exploiting conditional independence and search methods such as AND/OR search or branch-and-bound to scale to larger dimensions.

Appendix for “A∗Sampling”

The appendix introduces shorthand used in the theoretical development and identifies the cumulative distribution and density functions of a Gumbel variable.

The appendix defines shorthand before proving the paper’s main theoretical results and providing additional experimental details.
F_φ(g) and f_φ(g) denote the CDF and PDF, respectively, of a Gumbel(φ) distribution.

Joint Distribution of Gumbel Max and Argmax

The construction generates truncated Gumbel values and associated locations whose maxima and argmaxes encode Gibbs-distributed choices. The Top-Down and In-Order realizations differ by ordering and partitioning operations.

Joint distribution: The joint distribution of the maximum and argmax of independent truncated Gumbels yields a Gibbs distribution and a truncated Gumbel maximum.The maximum has location log Z, while the argmax follows the corresponding mass proportions.
Joint distribution: The maximum Gumbel value estimates the log partition function with variance π2/6N for N samples.
Joint distribution: Gumbel argmax probabilities satisfy Luce’s choice axiom and are independent of the truncation bound.The bound may therefore be set to infinity.
Constructions: Algorithm 3 produces Gumbel values in non-increasing order without a queue, whereas Algorithm 1 uses top-down partitioning.Figure 4 presents the two realizations as re-orderings of one another.

Correctness of the Top-Down Construction of the Gumbel Process

The Top-Down Construction is shown to generate a Gumbel process by proving that partitioning does not alter the distribution of the leading Gumbels or their induced maxima.

Proof strategy: The proof establishes that the Top-Down Construction constructs the Gumbel process.The argument uses a special In-Order case and then shows partitioning preserves the relevant distributions.
Proof strategy: With no subdivision, the construction becomes the In-Order Construction, which generates Gumbel values in non-increasing order.This special case needs no queue.
Partition invariance: The resulting collection of subset maxima satisfies the Gumbel-process construction requirements.The proof separately establishes the required conditions, with the third condition stated as immediate.
Partition invariance: The top n Gumbels have the In-Order distribution regardless of the partition function, provided it does not produce immeasurable sets.This makes the induced collection of subset maxima invariant to partition choice.

Equivalence Under partition

Partition invariance makes the Top-Down process analytically tractable, while A* termination and correctness follow from comparisons with global-bound search and the Gumbel-process linearity result.

Equivalence Under partition: For continuous µ, removals can be omitted while preserving the same distribution.Equal-measure successive regions permit the same simplification.
Equivalence Under partition: Boundary Gumbels are independent truncated Gumbels whose locations are log µ(Bi), and the boundary subsets partition the remaining space.This supplies the conditional structure used to derive the next maximum and location.
Equivalence Under partition: The joint distribution of max-Gumbels in B and Bc is invariant under the partition and equals that of independent Gumbels with locations log µ(B) and log µ(Bc).The proof conditions on which region contains the first sampled location and combines the two resulting events.
A* termination: A* cannot visit nodes that global-bound A* never visits, so almost-sure termination transfers from the global-bound procedure.Both searches use the same Gumbel-process realization and differ in explored nodes through their bounds.
A* termination: Global-bound A* is equivalent to rejection sampling and terminates after k iterations with a geometric distribution.The equivalence holds when one global bound M is reused at every node.
A* correctness: Given termination, A* returns the correct distribution because of the Gumbel-process linearity result.The proof introduces an auxiliary process to represent the bounded difference between tractable and target log densities.
Bound efficiency: Additional regions explored by a suboptimal bound are linear in the total bound suboptimality under the stated width-proportional assumption.Balanced splitting gives the corresponding depth relationship with high probability.
Bound efficiency: For n-term log-likelihoods, constant termwise bounds have linear suboptimality, whereas posterior concentration reduces significant-region suboptimality to O(√n).The concentration statement invokes regions of width O(n^-1/2).