Year: 2018

Abstract:
Probabilistic methods often need to draw samples from a nontrivial distribution. A∗ sampling is a nice algorithm by building upon a top-down construction of a Gumbel process, where a large state space is divided into subsets and at each round A∗ sampling selects a subset to process. However, the selection rule depends on a bound function, which can be intractable. Moreover, we show that such a selection criterion can be inefficient. This paper aims to improve A∗ sampling by addressing these issues. To design a suitable selection rule, we apply Probability Matching, a widely used method for decision making, to A∗ sampling. We provide insights into the relationship between A∗ sampling and probability matching by analyzing a nontrivial special case in which the state space is partitioned into two subsets. We show that in this case probability matching is optimal within a constant gap. Furthermore, as directly applying probability matching to A∗ sampling is time consuming, we design an approximate version based on Monte-Carlo estimators. We also present an efficient implementation by leveraging special properties of Gumbel distributions and well-designed balanced trees. Empirical results show that our method saves a significantly amount of computational resources on suboptimal regions compared with A∗ sampling.

heading: 1 INTRODUCTION
text: Probabilistic methods provide an important family of tools in machine learning for modeling uncertainty of complex systems, performing probabilistic inference, revealing hidden factors (Ghahramani, 2015), and making decisions Kocsis & Szepesvári (2006). These methods usually involve a fundamental task of drawing samples from a nontrivial distribution.
There exists a lot of work approaching the sampling problems, including rejection sampling (Gilks & Wild, 1992), MCMC (Propp & Wilson, 1996), etc. Recently, sampling algorithms based on the Gumbel process have received increasing attentions (Malmberg, 2013; Hazan et al., 2013; Gane et al., 2014; Hazan & Jaakkola, 2012; Papandreou & Yuille, 2011; Tarlow et al., 2012; Kappes et al., 2015; Kim et al., 2016) since a Gumbel process can turn a sampling task to an optimization problem so that we can use optimization tools to finish the original sampling task.
In this work, we focus on A∗ sampling (Maddison et al., 2014) which is one of the most famous Gumbel process based sampling algorithm. The major advantage of A∗ sampling is that it can be applied to large state spaces, e.g., a continuous sample space or a discrete space whose size is exponentially large. The reason is that A∗ sampling divides the state space into disjoint subsets and takes each subset as a whole, so that it can avoid initializing a large number of states, which is often encountered by other Gumbel process based algorithms (Papandreou & Yuille, 2011). Furthermore, A∗ sampling adaptively selects subsets to process and the performance of A∗ sampling is highly dependent on the selection rule.
However, how to select subsets to process is very challenging. In each round, A∗ sampling processes the subset with maximum D(S) which is an upper bound of the maximum Gumbel value within a subset S (see Section 2 for more details of D(S)). But in general, it is difficult to compute D(S) since it is an instance of non-convex optimization. Another challenge is that even if we are able to compute D(S) efficiently, selecting a subset with the maximum D(S) may not be a good choice. This is because our target is to process subsets with larger Gumbel values, but D(S) only provides an upper bound. So it is possible that the Gumbel value of S is relatively small with high probability while D(S) is very large. In this case, A∗ sampling will waste many computational resources on
suboptimal regions. We’ll discuss more on how this inaccuracy of D(S) deteriorates the performance of A∗ sampling by analyzing a counter example in Section 3.
To address the above challenges, we improve the subset selecting procedure of A∗ sampling with probability matching (PM) which has been proven efficient in many settings of making decisions, including Bayesian bandits (Chapelle & Li, 2011), MDP (Osband & Van Roy, 2016), economic decisions (Vulkan, 2000), etc.
Contributions: Intuitively, PM randomly selects an option according to its probability of being the optimal, so that it won’t select a suboptimal option for too many rounds. To provide more insights into the efficiency of applying PM to A∗ sampling, we first analyze a simple but nontrivial special case in which the state space is partitioned into two subsets. As we’ll present in Section 4.1, in this case, PM is optimal within a constant gap in terms of the stochastic regret (Guha & Munagala, 2014) which measures the number of selected rounds on suboptimal options. Furthermore, as directly applying PM to A∗ sampling is time consuming, we design a novel approximate algorithm based on Monte-Carlo estimators. The approximate algorithm is computationally efficient since it utilizes special properties of Gumbel distributions and well-designed balanced trees. We empirically compare our method with popular baselines of A∗ sampling and Metropolis-Hastings algorithm. Experiments show that our algorithm works well.

heading: 2 PRELIMINARIES
text: In this section, we present some preliminary knowledge of the Gumbel process and A∗ sampling. Below, we first introduce basic definitions of probability distributions and Gumbel distributions.
Definition 1 (Probability distributions). In general, a distribution P on a state space Ω provided its potential function, φP : 2Ω → R, is a sigma-finite measure such that P (S) = 1ZP exp{φP (S)}, where ZP = exp(φP (Ω)) is normalizing constant.
Definition 2 (Gumbel and Truncated Gumbel distributions (Malmberg, 2013)). Let c denote the Euler constant. For convenience, define eψ(g) = exp(−g + ψ), Fψ(g) = exp(− exp(−g + ψ)) and fψ(g) = eφ(g)Fψ(g). Then (1), G(ψ): a Gumbel distribution with location ψ has PDF and CDF at state g: fψ+c(g), Fψ+c(g). (2), T G(ψ, b): a Truncated Gumbel distribution with location ψ and truncated value b has PDF and CDF at state g < b: fψ+c(g)/Fψ+c(b), Fψ+c(g)/Fψ+c(b).

heading: 2.1 GUMBEL PROCESS
text: Now we are ready to introduce the Gumbel process.
Definition 3 (Gumbel process (Malmberg, 2013)). Let P (S) be a sigma-finite measure on sample space Ω, S ⊆ Ω is a measurable subset. Let φP (·) denote the potential function of P such that φP (S) = logP (S) + logZP . Then GP = {GP (S)|S ⊆ Ω} is a Gumbel process induced from P , if:
• (marginal distributions) GP (S) ∼ G(φP (S)). • (independence of disjoint sets) GP (S) ⊥ GP (Sc). • (consistency constraints) for measurable S1, S2 ⊆ Ω, then GP (S1 ∪ S2) =
max(GP (S1), GP (S2)).
The Gumbel process is useful in sampling since arg maxx∈ΩGP (x) ∼ P (Malmberg, 2013). Therefore, we can draw a sample from P by constructing a Gumbel process for distribution P , and then finding the maximum one with some optimization techniques.
In the sequel, we will use P to denote the target distribution, and we call GP (S) the Gumbel value of subset S. According to (Malmberg, 2013), Defn. 3 is associated with a natural bottom-up construction: for any x ∈ Ω, we first perturb it with an independent Gumbel noise, i.e., g(x) ∼ G(0). After that we simply set GP (x) = g(x) + φP (dx) and compute GP (S) = maxx∈S GP (x) for all S ⊆ Ω according to the consistency constraints. However, when Ω is infinite, such a bottom-up construction is infeasible.
Top-down construction: (Maddison et al., 2014) presents a top-down construction, which partitions the state space into regions and resolves the problem caused by infinite spaces by considering each region as a whole. Formally, the top-down procedure constructs a top-down tree, tree(P ), with each
node corresponding to a subset of Ω. tree(P ) is rooted in Ω. Let par(S) denote the parent of subset S. For each S ∈ tree(P ), its children is a disjoint partition of S, that is, ∪S′:par(S′)=SS′ = S. The top-down construction computes Gumbel values for subsets in the order from the top to the bottom of tree(P ). Formally, according to the consistency constraints and marginal distributions, we compute GP (S) ∼ T G(φP (S), L(S)) where L(S) := GP (par(S)). In the algorithmic view of point, the top-down construction maintains a collection of subsets of Ω. Initially, the collection contains only Ω. At each round, the algorithm selects an element S from the collection and computes GP (S). After that it divides S into subsets S1, S2 and adds them into the collection.

heading: 2.2 A∗ SAMPLING
text: Obviously, if φP (S) is hard to compute, the top-down construction for P is computationally intractable. (Maddison et al., 2014) solves this problem by utilizing the linearity of Gumbel distribution. More specifically, given a distribution Q, if GQ(x) induces a Gumbel process of Q, then GP (x) := GQ(x) + φP (x) − φQ(x) induces a Gumbel process of distribution P . Based on this insight, (Maddison et al., 2014) proposes the A∗ sampling, which relies on a tractable proposal distribution Q. Furthermore, since GQ(S) ⊥ arg maxx∈S GQ(x) (Maddison et al., 2014), A∗ sampling executes the top-down construction for Q, and for each subset, A∗ sampling computes G̃P (S) = GQ(S) + φP (xQ(S)) − φQ(xQ(S)) where xQ(S) ∼ Q(·|S).1 Suppose at some time point that A∗ sampling has processed n nodes in tree(Q), denoted by DoneQ(n). It can be shown that there are n+ 1 nodes in the to be processed collection, denoted by CollectQ(n). As introduced above, for each A ∈ DoneQ(n), we have a pair (xQ(A), G̃P (A)), and each S ∈ CollectQ(n) is associated with a truncated Gumbel distribution T G(φQ(S), L(S)).
The subset selection and termination in A∗ sampling rely on a bound function B : 2Ω → R such that B(S) ≥ supx∈S φP (x) − φQ(x). Let D(S) := L(S) + B(S). If for some n, maxS∈Done(n) G̃P (S) ≥ maxS∈Collect(n)D(S), A∗ terminates and outputs the element with maximum value among the processed nodes. At round n, A∗ sampling selects the node S with maximum value of D(S) from Collect(n).

heading: 3 CHALLENGES OF A∗ SAMPLING
text: There are two challenges in A∗ sampling. The first one is about the function D on which A∗ sampling relies. Computing this function for every S can be intractable since it can be a non-convex optimization. If we simply remove the (possibly intractable) bound function or use a very loose bound, A∗ sampling will degenerate to an algorithm which is not efficiency (Maddison et al., 2014). We name the degenerated algorithm as A∗ sampling without a bound ( See Appendix D for details.).
The second challenge is that selecting the subset with maximum D(S) is not always a good choice. This is because D(S) is just an upper bound of GP (S) and it is possible that GP (S) is relatively small with high probability while D(S) is very large. We now present a simple counter example for A∗ sampling to intuitively explain the reason. In this example, Ω = (−10.0,+10.0), the target is a mixture distribution: P (x) ∝ (1.0− 10−5)N (x;−5.0, 1.0) + 1[|x| ≤ 0.5 ∗ 10−405]10400 and Q(x) ∝ N (x; 5.0, 1.0). The log likelihoods of P and Q are shown in Fig. 1(a). We first empirically evaluate A∗ sampling on this example. Fig. 1(b) shows the selected rounds on the optimal subsets and Fig. 1(c) shows the maximum Gumbel value found by A∗ sampling. Results are averaged over 100 runs. We can see that A∗ sampling has a poor performance on this example. In this case, D(S) is large if S covers points near x = 0. So A∗ sampling will allocate lots of computational resources into such intervals, however, GP (S) being high for such S is with probability only about 0.00001.

heading: 4 A∗ SAMPLING WITH PROBABILITY MATCHING
text: We now present how to use PM to improve A∗ sampling by addressing the above challenges. We first present an intuitive example in Section 4.1. Then, we present a practical PM algorithm based on Monte-Carlo estimators of GP (S) in Section 4.2 and an efficient implementation with well-designed balanced trees in Section 4.3.
1In this paper, we use P (·|S) to denote the distribution P conditioned on state space S.

heading: 4.1 PROBABILITY MATCHING AND AN EXAMPLE WITH TWO SUBSETS
text: Algorithm 1 Probability matching with MonteCarlo estimators.
1: Input: the target distribution P , proposal distribution Q, state space Ω, time horizon T . 2: Output: x: a sample from P . 3: maxgumbel = −∞, x = None. 4: Collect = {Ω}, L(Ω) =∞,∀S ⊆ Ω, t = 1. 5: while t ≤ T do 6: t = t+ 1. 7: Select S∗ according to Eq. (4). 8: Split S∗ into disjoint sets S1, S2. 9: L(S1) = L(S2) = GQ(S∗). 10: for S ∈ S1, S2 do 11: x(S) ∼ Q(·|S). 12: G(S) ∼ T G(φQ(S), L(S)). 13: G̃(S) = G(S) +φP (x(S))−φQ(x(S)).
14: if maxgumbel < G̃(S) then 15: maxgumbel = G̃m(S), x = x(S). 16: end if 17: end for 18: Compute Monte-Carlo estimators for S1, S2 and update balanced trees. 19: Collect.insert(S1), Collect.insert(S2). 20: end while
In general, when making a choice among a set of options, PM selects an option randomly according to its probability of being the optimal. More specifically, in our problem, the optimal option is the subset with the maximum Gumbel value. Formally, by definition, the maximum Gumbel value within region S is a random variable GP (S) = maxx∈S T G(φQ(dx), L(S)) + φP (x) − φQ(x). Suppose the state space Ω is partitioned into {S1, · · · , SK}. PM selects a subset according to the probability:
p ( i = arg max
k∈[K] GP (Sk)
) . (1)
Intuitively, PM has an excellent performance since it allocates computational resources into the options which are likely to have large outcomes. To provide more intuition into why PM suits A∗ sampling, we analyze a simple but nontrivial case in which we divide Ω into two sets.
In order to get a clean theoretical result, we additionally assume A∗ sampling does not further split a subset after processing it. We focus on the stochastic regret (Guha & Munagala, 2014) which is the expected number of selections on the suboptimal subset. Formally, suppose Ω is partitioned into S1, S2. Let i∗ = arg maxi∈{1,2}GP (Si) which is a random variable. Consider an algorithm A which selects a subset SiA,t at time step t. The stochastic regret of A at time T is: RA(T ) = E (∑T t=1 1[iA,t 6= i∗] ) .
Intuitively, the smaller RA is, the better A is, since A won’t waste many computational resources on the suboptimal subset. Moreover, we can prove that PM is optimal within a constant gap in terms of the stochastic regret: Lemma 1. Let opt(T ) denote the algorithm which minimizes the stochastic regret. Then :
RPM (T ) ≤ 2Ropt(T )(T ), ∀T where RPM (T ) is the stochastic regret of PM.
The proof of Lemma 1 is adapted from the proof in (Guha & Munagala, 2014) for Bayesian bandits with two arms, we defer the details in Appendix A.

heading: 4.2 PROBABILITY MATCHING WITH A MONTE-CARLO ESTIMATOR
text: Unfortunately, drawing samples from the probability in Eq. (1) is intractable whenGP (S) is complex. So in this section, we present an efficient PM algorithm based on a Monte-Carlo estimator of GP (S).
Consider a random variable Y = φP (x) − φQ(x), x ∼ Q(·|S) whose expectation is a constant plussing the KL-divergence between Q and P conditioned on the subset S. We can equally characterize GP (S) as
max y T G(log(Q(S) · p(Y = y)), L(S)) + y. (2)
We present the proof of Eq. (2) in Appendix B. Eq. (2) suggests that we can get a Monte-Carlo estimator of GP (S) by estimating Y . More specifically, let Y1, · · · , Ym be a sequence of random variables and w1, · · · , wm be the corresponding weights such that ∑m i=1 wi = 1, wi > 0. Suppose the random variable Ym : p(Ym = Yi) = wi is an unbiased estimator of Y , then we can estimate GP (S) by:
ĜP (S) = max i∈[m] T G(log(wiQ(S)), L(S)) + Yi = max i∈[m] T G(log(wiQ(S)) + Yi, L(S) + Yi) (3)
The second equality holds due to the linearity of the truncated Gumbel distribution (Maddison et al., 2014). According to Eq. (3), we can estimate GP (S) with existing Monte-Carlo estimators of Y , such as adaptive importance sampling (Gilks & Wild, 1992).
The corresponding PM with Monte-Carlo estimators is to draw samples from
p ( î = arg max
j∈[n] ĜP (Sj)
) . (4)
What remains is how to sample from the probability in Eq. (4) efficiently. The most popular execution of Eq. (4) is as in (Chapelle & Li, 2011): we draw yi ∼ ĜP (Si), and take î = arg maxi yi, then it can be shown that î is a sample from the probability in Eq. (4).
However, a direct implementation of the above ideas requires time complexity O(m) since we need to draw samples from m truncated Gumbel distributions, where m = ∑ i∈[n]mi is the number of particles in total and mi is the number of particles in Si. So our selection algorithm executing m rounds would require running time O(m2). It is relatively slow comparing with the O(m logm) time complexity for A∗ sampling (Maddison et al., 2014).

heading: 4.3 AN EFFICIENT IMPLEMENTATION BY BALANCED TREES
text: We now present a novel algorithm that only requires O(logm) running time to sample from the distribution in Eq. (4). Our algorithm is based on the properties of the truncated Gumbel distribution and under the help of well-designed balanced trees.
We first decompose sampling from the distribution in Eq. (4) into two steps which can be done efficiently. The decomposition is an immediate inference of Eq. (4):
p ( î = arg max
j∈[n] ĜP (Sj)
) = ∫ x p(x = max j∈[n] ĜP (S))p(̂i = arg max j∈[n] ĜP (S)|x = max j∈[n] ĜP (S))dx.
Thus, sampling from the distribution in Eq. (4) equals to the following two sampling problems:
x ∼ max i∈[n] xi, xi ∼ ĜP (Si), î ∼ p(i = arg max j∈[n] xj ∼ ĜP (Si)|x = maxxj)
Recall that ĜP (S) is the maximum one among a set of truncated Gumbels. Thus, the above two sampling problems are essentially sampling the maxima and the argument of the maxima among a set of truncated Gumbels. So our target can be converted into the following problem:
Problem 1. Given a set of truncated Gumbel variables {vi}mi=1 with parameters (ai, bi), i.e., vi ∼ T G(ai, bi). We define two sampling problems:
v = maxi∈[m] vi (5)
î ∼ p(i = arg maxj∈[m] vj |v = maxj∈[m] vj) (6)
We use the inverse transform sampling (Devroye, 1986) to sample v in Eq. (5). In inverse transform sampling, for a random variable X with CDF UX(x), we first draw a sample s ∼ uniform(0, 1), and then compute x such that UX(x) = s, it can be shown that x ∼ X . Thus, let U(g) denote the CDF of v, we only need an algorithm to compute g such that U(g) = s, s ∈ (0, 1). We now show how to compute such g efficiently with balanced trees.
For notational clarity, let Ua,b(g) denote the CDF of a truncated Gumbel distribution, T G(a, b). According to Defn. 2, we have Ua,b(g) = exp(− exp(−min(g,b)+a)) exp(− exp(−b+a)) . Recall v = maxi vi, then v has
CDF: U(g) = ∏ i Uai,bi(g) = ∏ i exp(− exp(−min(g,bi)+ai)) exp(− exp(−bi+ai)) . Take logarithm on both sides, we get
logU(g) = ∑ i∈[m] (− exp(−min(g, bi) + ai) + exp(−bi + ai)).
Without loss of generality, we sort bi’s in a non-decreasing order, that is, bi ≤ bi+1. Since U(g) is a monotonically increasing function, for g ∈ (bi, bi+1], we have: logU(g) = − ∑ j>i exp(−g + aj) + ∑ j>i exp(−bj + aj) = − exp(−g) ∑ j>i exp(aj) + ∑ j>i exp(−bj + aj)
Thus, given U(g) and suppose g ∈ (bi, bi+1], we can compute g by: g = − log ∑ j>i exp(aj − bj)− logU(g)∑
j>i exp(aj) (7)
Thus, when we get s ∼ uniform(0, 1), we need to find i such that U(bi) ≤ s ≤ U(bi+1), and then solve g according to Eq. (7) and inverse sampling . Both of above two steps can be done efficiently via a balanced tree.
Suppose we have a balanced tree such that each node in the tree corresponds to an index i ∈ [m], and the key of the balanced tree is bi, that is, for all j in the right subtree of node i, we have bj ≥ bi and for all j in the left subtree, we have bj ≤ bi. Suppose that from the balanced tree, we can query in O(1) time at each node i for terms: (1) exp(−bi) ∑ j>i exp(aj); (2)
exp(ai) ∑ j>i exp(−bj) and (3) ∑ j>i exp(aj − bj). We can query these terms efficiently in a balanced tree because they all are summations over an interval. And according to Defn. 2, we know that logU(bi) = ∑ j>i (exp(−bj + ai)− exp(aj − bi)), we can check out whether log s < logU(bi) in O(1). Therefore, we can find the index i such that U(bi) ≤ s ≤ U(bi+1) in running time O(logm). After that, we can compute U(g) = s via Eq. (7) in running time O(1).
Now we turn to sample î in Eq. (6). Without loss of generality, suppose g ∈ (bi, bi+1). Obviously, for j < i, p(j = arg maxj′ vj′ |g = maxj′ vj′) = 0. For j ≥ i, by Defn. 2 and with simple calculations, we have:
p ( j = arg max
j′ vj′ |g
) ∝ dUaj ,bj (g)
dg
∏ j′ 6=j,j′>i Uaj′ ,bj′ (g)
= exp(−g + aj) exp(− exp(−g + aj)) exp(− exp(−bj + aj)) ∏ j′ 6=j,j′>i Uaj′ ,bj′ (g) = exp(aj) exp(−g) ∏ j′>i Uaj′ ,bj′ (g) ∝ exp(aj)
(8)
According to Eq. (8), we can sample î in O(logm) running time with a balanced tree from which we can query ∑ j>i exp(aj) efficiently. Putting the previous results together, we get the algorithm as outlined in Alg. 1.

heading: 5 EXPERIMENTS
text: In this section, we present our empirical results. We first check the correctness on a simple toy experiment and the counter example in Section 3. After that, we evaluate the efficiency of Alg. 1 on two Bayesian posterior inference tasks, results show that our algorithm outperforms vanilla A∗ sampling significantly.

heading: 5.1 CORRECTNESS OF ALG. 1
text: We first verify the correctness of Alg. 1 on a greenhouse experiment. We consider sampling from a one-dimensional Gaussian mixture with potential function φP (x) = − log(N (x;−2.0, 1.0) + 2N (x; 2.0, 1.0)), which is a multi-mode target distribution. We set Q = N (0, 2). We present our result in Fig. 2(a) which shows the ground truth and the empirical sample distributions of Alg. 1 and baselines. From Fig. 2(a), Alg. 1 has a similar performance to A∗ sampling and outperforms the A∗ sampling without a bound function.

heading: 5.2 THE COUNTER-EXAMPLE
text: We empirically compare the performance of algorithms on the example in Section 3. The result is shown in Fig. 2(b). We can see that PM-A∗ outperforms baselines significantly since our algorithm wastes less resources on suboptimal subsets. 5.3 BAYESIAN POSTERIOR INFERENCE In this section, we evaluate our algorithm on two Bayesian posterior tasks: the clutter problem and the Bayesian logistic regression. More specifically, we focus on sampling tasks of formulations P (x) ∝ p(x) ∏n i=1 p(yi|x) where x is the variable we are going to sample, p(·) is the prior distribution over x, and {yi}ni=1 are observations, and p(yi|x) is the likelihood function. We simply set Q(x) := p(x) in both A∗ sampling and Alg. 1. For vanilla A∗ sampling, we exploit the standard stochastic gradient descent (SGD) algorithm to calculate the bound function, maxx∈S ∑ i log p(yi|x). For the MonteCarlo estimator in Alg. 1, we exploit the importance sampling over the trajectory of the same SGD algorithm as in the vanilla A∗ sampling 2.

heading: 5.3.1 EVALUATION ON THE CLUTTER PROBLEM
text: We now evaluate Alg. 1 on the Clutter problem proposed by (Minka, 2001). Clutter problem aims to inference the mean of an isotropic Gaussian with some data points are outliers. Consider a mixture
2We first tune the parameters of SGD for vanilla A∗ sampling, and then apply them to Alg. 1 without further tuning.
distribution:p(y|x) = (1 − w)N (y;x, I) + wN (y; 0, σ1I), p(x) = N (x; 0, σ2I), where w is the ratio of outliers which is a known parameter, andN (µ,Σ) represents Gaussian distribution. Our goal is to inference x given data {yi}ni=1. We do experiments on dimensions varying in 5, 15, 20, n = 20. We compare the Gumbel of these algorithms. We run 100 times and present the averaged results in Fig. 3. We can see that Alg. 1 outperforms A∗ sampling constantly.

heading: 5.3.2 EVALUATION ON BAYESIAN LOGISTIC REGRESSION
text: Our last experiment is on Bayesian Logistic Regression. Given a dataset {xi}ni=1 associated with label {yi}ni=1 where yi ∈ {0, 1}. We follow the setting in (Gershman et al., 2012) and define the Bayesian logistic regression: p(α) = Gamma(α; a, b), p(wk) = N (wk; 0, α−1), p(yi = 1;xi,w) = sigmoid(w
Txi). In this model, {w, α} are the hidden variables, where w denotes the regression coefficients and α is a precision parameter. We set a = b = 1. We do experiments on 13 binary classification datasets proposed by (Mika et al., 1999). The number of features of these data sets are in range from 2 to 60, and the number of points ranges from 24 to 7400 (See Appendix C for more statistics). We present our results in Fig. 4(a) where all results are averaged over 100 runs. Fig. 4(a) presents the summation of the maximum likelihood found by each algorithm on these datasets over time. From Fig. 4(a), we can see that PM-A∗ outperforms all baselines.
Furthermore, we compare our algorithm with standard Matropolis-Hastings algorithm (MH) and adaptive inference with exploration (AIE) (Rainforth et al., 2018) which also attempts to bridge the gap between sampling problems and decision-making techniques. For MH, the initial points are sampled from the prior. To make the comparison fair, we also evaluate Alg. 1 and AIE with the prior as the Monte-Carlo estimator instead of gradient-based methods. We compare the likelihoods in Fig. 4(b). We can see that Alg. 1 outperforms AIE even if they use the same Monte-Carlo estimator. This is AIE attempts to use UCB-like algorithm to make decisions, but UCB works only for those models in which concentration bounds hold which is not always valid in sampling problems.

heading: 6 CONCLUSION AND FUTURE WORK
text: In this work, we focus on improving the subset selection procedure in A∗ sampling with PM. We proved that in the special case of two subsets, PM is optimal within a constant gap in terms of the stochastic regret. Moreover, we proposed a practical algorithm based on Monte-Carlo estimators and well-designed balanced trees. Empirical results show that our methods saves a significantly amount of computational resources on suboptimal regions compared with A∗ sampling.
There exists several challenges in future work. The first one is on the analysis of PM. Though we proved PM is efficient in the case of two subsets, it is very challenging to prove the efficiency in general. The second one is that the performance of Alg. 1 relies on the accuracy of the Monte-Carlo estimator. However, it is time-consuming to compute an accurate Monte-Carlo estimator. So it is important to balance the accuracy of the Monte-Carlo estimator and the performance of PM. We hope our work is a starting point to address these problems.

heading: ACKNOWLEDGMENTS

heading: A PROBABILITY MATCHING FOR TWO SUBSETS
text: We present the proof of Lemma 1 in this section. This proof is adapted from the proof of Thompson sampling for Bayesian bandits with two arms in (Guha & Munagala, 2014).
We first make the simplified A∗ sampling for two subsets clearer. Recall Ω is divided into two subsets S1, S2. For notational convenience, let kt,i denote the number of selected rounds of Si up to time step t. Since we do not split S1, S2, if we have processed Si for t rounds, we will have t observations {(xi,j , ui,j)}tj=1 such that xi,1 ∼ Q(·|Si), ui,1 ∼ G(φQ(Si)) and xi,j ∼ Q(·|Si\Xi,j−1), ui,j ∼ T G(φQ(Si\Xi,j−1), ui,j−1) for j > 1 where Xi,j = {xi,j′}jj′=1. Let Vi,t = {(xi,j , ui,j)} kt,i j=1 denote the observations from subset Si until time t and Vt = {V1,t, V2,t} denote the collection of observations from both subsets. Obviously, a subset selection algorithm A selects a subset according to previous observations, that is, we can represent the subset selected by A at time t, iA,t+1, as iA,t+1 := iA(Vt). Now we present A∗ sampling for two subsets with A as the subset selection algorithm in Alg. 2.
Algorithm 2 A∗ sampling for two subsets. 1: Input: the target distribution P , proposal distribution Q, state space Ω = S1 ∪ S2, a subset
selection algorithm A, time horizon T . 2: Output: x: a sample from P . 3: maxgumbel = −∞, x = None. 4: u1,0 = u2,0 =∞, t = 0. 5: while t ≤ T do 6: t = t+ 1. 7: i = iA(Vt−1). 8: x′ ∼ Q(·|Si\Xi,kt−1,i). 9: ui,t ∼ T G(φQ(Si\Xi,kt−1,i), ui,t−1).
10: u′ = ui,t + φP (x(S))− φQ(x(S)). 11: if maxgumbel < u′ then 12: maxgumbel = u′, x = x′. 13: end if 14: end while
Recall opt(T ) denote the the algorithm which minimizes the stochastic regret up to time horizon T . We now explicitly define the optimal stochastic regret Ropt(T )(T ) with time horizon T via dynamic programming. Let qi(V ) denote the probability of Si being the optimal subset given observations V . For convenience, let RA(T |V ) be the stochastic regret of algorithm A given observations V . Then: Definition 4. We can define Ropt(T )(T |V ) as follows:
Ropt(1)(1|V ) = min i∈{1,2} (1− qi(V )).
∀T > 1, Ropt(T )(T |V ) = min i∈{1,2}
( 1− qi(V ) + ∫ xi,ui Ropt(T−1)(T − 1|V ∪ (x, u))p(x, u|V ) )
where (xi, ui) is the new observation from subset Si.
We now present the formal definition and some properties of qi(V ) which are useful in the sequel:
Corollary 1. Let V = {V1 := {(x1,j , u1,j)}k1j=1, V2 := {(x2,j , u2,j)} k2 j=1} denote a set of observations. Let Si,\V = Si\ ∪ {xi,j}kij=1. Then let i′ = {1, 2}\i and m(x) = φP (x) − φQ(x), we have:
qi(V ) = p
( max
x∈Si,\V T G(φQ(dx), ui,ki) +m(x) > max x∈Si′,\V T G(φQ(dx), ui′,ki′ ) +m(x) ) Suppose we process Si for the ki + 1 time, and then receive observation (x′i, u ′ i), we have:
qi(V ) ≥ Ex′i,u′i [qi(V ∪ (x ′ i, u ′ i))] (9)
Proof. Let E denote the event of the (ki + 1)-th observation from subset i is with the maximum Gumbel value within Si, that is, u′i + m(x ′ i) ≥ maxSi,\(V∪x′ i ) T G(φQ(dx), u′i) + m(x). Then we have qi(V ∪(x′i, u′i)|Ē) = qi(V |Ē); otherwise, qi will decrease, that is, qi(V ∪(x′i, u′i)|E) ≤ qi(V |E). Overall, we have
Ex′i,u′i [qi(V ∪ (x ′ i, u ′ i))] = p(E)qi(V ∪ (x′i, u′i)|E) + p(Ē)qi(V ∪ (x′i, u′i)|Ē ≤ qi(V )
For convenience, let RA(T |V ) = ∑ i p(iA,1 = i)RA(i, T |V ) where RA(i, T |V ) denote the stochastic regret of A if we select subset i at the first round. The proof of Lemma 1 relies on the following lemma: Lemma 2. Let P denote the target distribution andQ denote the proposal. If an algorithmA satisfies that for any observations V and time horizon T :
RA(T |V ) ≤ RA(i, T |V ) + c(1− qi(V )) (10) where c is a constant. Then we have RA(T |V ) ≤ (c+ 1)Ropt(T )(T |V ) for all T and V .
Proof. For T = 1, it is obvious that Eq. (10) holds. We use mathematical induction to prove the case of T > 1. According to definition and with straight-forward calculations, we have:
RA(T |V ) ≤ RA(i, T |V ) + c(1− qi(V )), ∀i = (c+ 1)(1− qi(V )) + E(x,u)RA(i, T − 1|V ∪ (x, u)), ∀i ≤ (c+ 1)(1− qi(V ) + E(x,u)Ropt(T−1)(i, T − 1|V ∪ (x, u))), ∀i ≤ (c+ 1) min
i (1− qi(V ) + E(x,u)Ropt(T−1)(i, T − 1|V ∪ (x, u)))
= (c+ 1)Ropt(T )(i, T |V )
According to Lemma 2, we can prove Lemma 1 by proving that probability matching satisfies precondition (10) for c = 1. The following Lemma provides an equivalent characterization of Eq. (10) in the context of probability matching. Lemma 3. For probability matching, we have:
RPM (T ) ≤ RPM (i, T ) + 1− qi,∀i ∈ {1, 2} ⇐⇒ |RPM (1, T )−RPM (2, T )| ≤ 1 (11)
Proof. Suppose RPM (T ) ≤ RPM (1, T ) + 1 − q1. By definition, we have RPM (T ) = q1RPM (1, T ) + q2RPM (2, T ). Observing q1 + q2 = 1, we have RPM (2, T ) − RPM (1, T ) ≤ 1. Conversely, ifRPM (1, T ) ≤ RPM (2, T )+1, we haveRPM (T ) = q1RPM (1, T )+q2RPM (2, T ) ≤ RPM (2, T ) + 1− q2. Swapping the roles of 1 and 2, we complete the proof.
Now, we can complete the proof of Lemma 1.
Proof. For T = 1, it is obvious that Lemma 3 is true. We use mathematical induction to prove Lemma 3 holds for T > 1. For convenience, let u1 ∼ T G(φQ(S1), L(S1)), x1 ∼ Q(·|S1), u2 ∼ T G(φQ(S2), L(S2)), x2 ∼ Q(·|S2). We have RPM (1, T |V ) = 1− q1(V ) + Ex1,u1 [RPM (T − 1|V ∪ (x1, u1))]
≤ 1− q1(V ) + Ex1,u1 [1− q2(V ∪ (x1, u1)) +RPM (2, T − 1|V ∪ (x1, u1))] = 1− q1(V ) + Ex1,u1 [2q1(V ∪ (x1, u1)) + Ex2,u2RPM (T − 2|V ∪ {(x1, u1), (x2, u2)})] ≤ 1 + q1(V ) + Ex1,u1 [Ex2,u2RPM (T − 2|V ∪ {(x1, u1), (x2, u2)})]
The first equation is according to the inductive hypothesis, the next equality follows q1(V )+q2(V ) = 1 and the last inequality is due to Corollary 1. Similarly, we have:
RPM (2, T |V ) = 1− q2(V ) + Ex2,u2 [RPM (T − 1|V ∪ (x2, u2))] = q1(V ) + Ex2,u2 [q1(V ∪ (x2, u2))RPM (1, T − 1|V ∪ (x2, u2)) + q2(V ∪ (x2, u2))RPM (2, T − 1|V ∪ (x2, u2))] ≥ q1(V ) + Ex2,u2 [RPM (1, T − 1|V ∪ (x2, u2))− q2(V ∪ (x2, u2))] = q1(V ) + Ex2,u2 [Ex1,u1RPM (T − 2|(x1, u1), (x2, u2))]
The first inequality follows the inductive hypothesis in Lemma 3 and q1(V ) + q2(V ) = 1. Above inequalities show that RPM (1, T ) − RPM (2, T ) ≤ 1. Reversing the roles of S1 and S2, we can show that probability matching satisfies the condition in Lemma 3 which completes the proof.
B THE EQUIVALENT CHARACTERIZATION OF GP (S)
We now prove the equivalent characterization of GP (S) in Section 4.2. According to definitions, we have:
GP (S) = max x∈S T G(φQ(dx), L(S)) + φP (dx)− φQ(x)
= max y max φP (x)−φQ(x)=y
T G(φQ(dx), L(S)) + y
= max y
(y + T G(log(Q(S)p(Y = y)), L(S)))
the first equation is by definition and the third equation follows the fact that maxx∈S T G(φQ(dx), L) = T G(log( ∫ x∈S exp(φQ(x))dx), L).

heading: C STATISTICS OF DATASETS PROPOSED BY (MIKA ET AL., 1999)
text: We present statistics of datasets proposed by (Mika et al., 1999) in Table 1 where n denotes the number of data points and dim denotes dimension.

heading: D A∗ SAMPLING WITHOUT THE BOUND FUNCTION
text: As presented in Section 3, when the bound B(S) is not available or very loose, the vanilla A∗ sampling degenerates to Alg .3 which is not very efficient since it selects subsets simply according to L(S).
Algorithm 3 A∗ sampling without a bound. 1: Input: the target distribution P , proposal distribution Q, state space Ω, time horizon T . 2: Output: x: a sample from P . 3: maxgumbel = −∞, x = None. 4: Collect = {Ω}, L(Ω) =∞,∀S ⊆ Ω, t = 1. 5: while t ≤ T do 6: t = t+ 1. 7: Select S∗ ∈ Collect with maximum L(Ω). 8: Split S∗ into disjoint sets S1, S2. 9: L(S1) = L(S2) = GQ(S∗). 10: for S ∈ S1, S2 do 11: x(S) ∼ Q(·|S). 12: G(S) ∼ T G(φQ(S), L(S)). 13: G̃(S) = G(S) + φP (x(S))− φQ(x(S)). 14: if maxgumbel < G̃(S) then 15: maxgumbel = G̃m(S), x = x(S). 16: end if 17: end for 18: Collect.insert(S1), Collect.insert(S2). 19: end while

Review summaries:
This paper introduces a probability matching approach for optimizing Gumbel processes, i.e. the extension of the Gumbel-Max trick to more general measure spaces. The basic idea is to use a more refined subset selection mechanism as compared to A* Sampling, but at the cost of being able to guarantee an exact sample. Instead, the authors study the algorithm's best Gumbel so far as a function of the time budget.