Title: A * SAMPLING WITH PROBABILITY MATCHING

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.
Under review as a conference paper at ICLR 2019 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.
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) = 1 Z P exp{φ P (S)}, where Z P = 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).
GUMBEL PROCESS

1. INTRODUCTION:

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 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) = log P (S) + log Z P . Then G P = {G P (S)|S ⊆ Ω} is a Gumbel process induced from P , if:

• (marginal distributions) G P (S) ∼ G(φ P (S)). • (independence of disjoint sets) G P (S) ⊥ G P (S c ). • (consistency constraints) for measurable S 1 , S 2 ⊆ Ω, then G P (S 1 ∪ S 2 ) = max(G P (S 1 ), G P (S 2 )).

The Gumbel process is useful in sampling since arg max x∈Ω G P (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 G P (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 G P (x) = g(x) + φ P (dx) and compute G P (S) = max x∈S G P (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 )=S S = 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 G P (S) ∼ T G(φ P (S), L(S)) where L(S) := G P (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 G P (S). After that it divides S into subsets S 1 , S 2 and adds them into the collection.

2.2 A * SAMPLING 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 G Q (x) induces a Gumbel process of Q, then G P (x) := G Q (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 G Q (S) ⊥ arg max x∈S G Q (x) (Maddison et al., 2014) , A * sampling executes the top-down construction for Q, and for each subset, A * sampling computes

GP (S) = G Q (S) + φ P (x Q (S)) -φ Q (x Q (S)) where x Q (S) ∼ Q(•|S).

1 Suppose at some time point that A * sampling has processed n nodes in tree(Q), denoted by Done Q (n). It can be shown that there are n + 1 nodes in the to be processed collection, denoted by Collect Q (n). As introduced above, for each A ∈ Done Q (n), we have a pair (x Q (A), GP (A)), and each

S ∈ Collect Q (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) ≥ sup x∈S φ P (x) -φ Q (x). Let D(S) := L(S) + B(S)

. If for some n, max S∈Done(n) GP (S) ≥ max S∈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).

3. CHALLENGES OF A * SAMPLING:

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 G P (S) and it is possible that G P (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

]10 400 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, G P (S) being high for such S is with probability only about 0.00001.

4. A * SAMPLING WITH PROBABILITY MATCHING:

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 G P (S) in Section 4.2 and an efficient implementation with well-designed balanced trees in Section 4.3. 

4.1. PROBABILITY MATCHING AND AN EXAMPLE WITH TWO SUBSETS:

Algorithm 1 Probability matching with Monte-Carlo 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 = N one. 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 S 1 , S 2 . 9:

L(S 1 ) = L(S 2 ) = G Q (S * ). 10:

for S ∈ S 1 , S 2 do 11:

x(S) ∼ Q(•|S).

12: Collect.insert(S 1 ), Collect.insert(S 2 ). 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

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 = Gm (S), x = x(S).

G P (S) = max x∈S T G(φ Q (dx), L(S)) + φ P (x) -φ Q (x). Suppose the state space Ω is partitioned into {S 1 , • • • , S K }.

PM selects a subset according to the probability:

p i = arg max k∈[K] G P (S k ) .

(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) 

(T ) = E T t=1 1[i A,t = i * ] .

Intuitively, the smaller R A 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 :

R P M (T ) ≤ 2R opt(T ) (T ), ∀T

where R P M (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.

4.2. PROBABILITY MATCHING WITH A MONTE-CARLO ESTIMATOR:

Unfortunately, drawing samples from the probability in Eq. ( 1) is intractable when G P (S) is complex. So in this section, we present an efficient PM algorithm based on a Monte-Carlo estimator of G P (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 G P (S) as

EQUATION

We present the proof of Eq. ( 2) in Appendix B. Eq. ( 2) suggests that we can get a Monte-Carlo estimator of G P (S) by estimating Y . More specifically, let Y 1 , • • • , Y m be a sequence of random variables and w 1 , • • • , w m be the corresponding weights such that

m i=1 w i = 1, w i > 0. Suppose the random variable Y m : p(Y m = Y i ) = w i is an unbiased estimator of Y , then we can estimate G P (S) by: ĜP (S) = max i∈[m] T G(log(w i Q(S)), L(S)) + Y i = max i∈[m] T G(log(w i Q(S)) + Y i , L(S) + Y i ) (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 G P (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

EQUATION

)

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 y i ∼ ĜP (S i ), and take î = arg max i y i , then it can be shown that î 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] m i is the number of particles in total and m i is the number of particles in S i . So our selection algorithm executing m rounds would require running time O(m 2 ). It is relatively slow comparing with the O(m log m) time complexity for A * sampling (Maddison et al., 2014) .

4.3. AN EFFICIENT IMPLEMENTATION BY BALANCED TREES:

We now present a novel algorithm that only requires O(log m) 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 î = arg max j∈[n]

ĜP (S j ) = ĜP (S))dx.

Thus, sampling from the distribution in Eq. ( 4) equals to the following two sampling problems:

x ∼ max i∈[n] x i , x i ∼ ĜP (S i ), î ∼ p(i = arg max j∈[n]

x j ∼ ĜP (S i )|x = max x j )

Recall that Ĝ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 {v i } m i=1 with parameters (a i , b i ), i.e., v i ∼ T G(a i , b i ). We define two sampling problems:

EQUATION

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 U X (x), we first draw a sample s ∼ unif orm(0, 1), and then compute x such that U X (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 U a,b (g) denote the CDF of a truncated Gumbel distribution, T G(a, b).

According to Defn. 2, we have U a,b (g) = exp(-exp(-min(g,b)+a))

exp(-exp(-b+a))

.

Recall v = max i v i , then v has CDF: U (g) = i U ai,bi (g) = i exp(-exp(-min(g,bi)+ai)) exp(-exp(-bi+ai))

. Take logarithm on both sides, we get log

U (g) = i∈[m] (-exp(-min(g, b i ) + a i ) + exp(-b i + a i )).

Without loss of generality, we sort b i 's in a non-decreasing order, that is, b i ≤ b i+1 . Since U (g) is a monotonically increasing function, for g ∈ (b i , b i+1 ], we have:

log U (g) = - j>i exp(-g + a j ) + j>i exp(-b j + a j ) = -exp(-g) j>i exp(a j ) + j>i exp(-b j + a j )

Thus, given U (g) and suppose g ∈ (b i , b i+1 ], we can compute g by:

EQUATION

Thus, when we get s ∼ unif orm(0, 1), we need to find i such that U (b i ) ≤ s ≤ U (b i+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 b i , that is, for all j in the right subtree of node i, we have b j ≥ b i and for all j in the left subtree, we have b j ≤ b i . Suppose that from the balanced tree, we can query in O(1) time at each node i for terms: (1) exp(-b i ) j>i exp(a j );

(2) exp(a i ) j>i exp(-b j ) and (3) j>i exp(a j -b j ). 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 log U (b i ) = j>i (exp(-b j + a i ) -exp(a j -b i )), we can check out whether log s < log U (b i ) in O(1). Therefore, we can find the index i such that

U (b i ) ≤ s ≤ U (b i+1 ) in running time O(log m).

After that, we can compute U (g) = s via Eq. ( 7) in running time O(1).

Now we turn to sample î in Eq. ( 6). Without loss of generality, suppose g ∈ (b i , b i+1 ). Obviously, for j < i, p(j = arg max j v j |g = max j v j ) = 0. For j ≥ i, by Defn. 2 and with simple calculations, we have:

EQUATION

According to Eq. ( 8), we can sample î in O(log m) running time with a balanced tree from which we can query j>i exp(a j ) efficiently. Putting the previous results together, we get the algorithm as outlined in Alg. 1.

5. EXPERIMENTS:

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. 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.

5.2. THE COUNTER-EXAMPLE:

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(y i |x) where x is the variable we are going to sample, p(•) is the prior distribution over x, and {y i } n i=1 are observations, and p(y i |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, max x∈S i log p(y i |x). For the Monte-Carlo estimator in Alg. 1, we exploit the importance sampling over the trajectory of the same SGD algorithm as in the vanilla A * sampling2 .

5.3.1. EVALUATION ON THE CLUTTER PROBLEM:

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 distribution:p(y|x) = (1 -w)N (y; x, I) + wN (y; 0, σ 1 I), p(x) = N (x; 0, σ 2 I), where w is the ratio of outliers which is a known parameter, and N (µ, Σ) represents Gaussian distribution. Our goal is to inference x given data {y i } n i=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.

5.3.2. EVALUATION ON BAYESIAN LOGISTIC REGRESSION:

Our last experiment is on Bayesian Logistic Regression. Given a dataset {x i } n i=1 associated with label {y i } n i=1 where y i ∈ {0, 1}. We follow the setting in (Gershman et al., 2012) and define the Bayesian logistic regression: p(α) = Gamma(α; a, b), p(w k ) = N (w k ; 0, α -1 ), p(y i = 1; x i , w) = sigmoid(w T x i ). 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.

6. CONCLUSION AND FUTURE WORK:

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.

A PROBABILITY MATCHING FOR TWO SUBSETS:

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 S 1 , S 2 . For notational convenience, let k t,i denote the number of selected rounds of S i up to time step t. Since we do not split S 1 , S 2 , if we have processed S i for t rounds, we will have t observations

{(x i,j , u i,j )} t j=1 such that x i,1 ∼ Q(•|S i ), u i,1 ∼ G(φ Q (S i )) and x i,j ∼ Q(•|S i \X i,j-1 ), u i,j ∼ T G(φ Q (S i \X i,j-1 ), u i,j-1 ) for j > 1 where X i,j = {x i,j } j j =1 . Let V i,t = {(x i,j , u i,j )} kt,i j=1

denote the observations from subset S i until time t and V t = {V 1,t , V 2,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, i A,t+1 , as i A,t+1 := i A (V t ). 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 Ω = S 1 ∪ S 2 , a subset selection algorithm A, time horizon T . 2: Output: x: a sample from P . 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 R opt(T ) (T ) with time horizon T via dynamic programming. Let q i (V ) denote the probability of S i being the optimal subset given observations V . For convenience, let R A (T |V ) be the stochastic regret of algorithm A given observations V . Then: Definition 4. We can define R opt(T ) (T |V ) as follows:

3: maxgumbel = -∞, x = N one. 4: u 1,0 = u 2,0 = ∞, t = 0. 5: while t ≤ T do 6: t = t + 1. 7: i = i A (V t-1 ). 8: x ∼ Q(•|S i \X i,kt-1,i ). 9: u i,t ∼ T G(φ Q (S i \X i,kt-1,i ), u i,t-1 ). 10: u = u i,t + φ P (x(S)) -φ Q (x(S)).

R opt(1) (1|V ) = min i∈{1,2}

(1 -q i (V )).

∀T > 1, R opt(T ) (T |V ) = min i∈{1,2} 1 -q i (V ) + xi,ui R opt(T -1) (T -1|V ∪ (x, u))p(x, u|V )

where (x i , u i ) is the new observation from subset S i .

We now present the formal definition and some properties of q i (V ) which are useful in the sequel:

Corollary 1. Let V = {V 1 := {(x 1,j , u 1,j )} k1 j=1 , V 2 := {(x 2,j , u 2,j )} k2 j=1 } denote a set of observations. Let S i,\V = S i \ ∪ {x i,j } ki j=1 . Then let i = {1, 2}\i and m(x) = φ P (x) -φ Q (x), we have:

q i (V ) = p max x∈S i,\V T G(φ Q (dx), u i,ki ) + m(x) > max x∈S i ,\V T G(φ Q (dx), u i ,k i ) + m(x)

Suppose we process S i for the k i + 1 time, and then receive observation (x i , u i ), we have:

EQUATION

The first inequality follows the inductive hypothesis in Lemma 3 and q 1 (V ) + q 2 (V ) = 1. Above inequalities show that R P M (1, T ) -R P M (2, T ) ≤ 1. Reversing the roles of S 1 and S 2 , we can show that probability matching satisfies the condition in Lemma 3 which completes the proof.

B THE EQUIVALENT CHARACTERIZATION OF G P (S)

We now prove the equivalent characterization of G P (S) in Section 4.2. According to definitions, we have: C STATISTICS OF DATASETS PROPOSED BY (MIKA ET AL., 1999)

G P (S) = max

We present statistics of datasets proposed by (Mika et al., 1999) 1 : Some statistics of datasets proposed by (Mika et al., 1999) .

D A * SAMPLING WITHOUT THE BOUND FUNCTION 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).

:

In this paper, we use P (•|S) to denote the distribution P conditioned on state space S.

We first tune the parameters of SGD for vanilla A * sampling, and then apply them to Alg. 1 without further tuning.

Figures and Tables:

Figure 1: The counter example with (a) the log likelihood of P and Q; (b) the selected number on the intervals containing the optimal point; (c) the maximum Gumbel value found by A * sampling.

Carlo estimators for S 1 , S 2 and update balanced trees. 19:

Figure 2: Experiment results with (a) the toy experiment; (b) the counter example.

Figure 3: Experiment results on the clutter problem with 5 dimensions on the left, 15 dimensions in the middle and 20 dimensions on the right.

Figure 4: Experiments on logistic regression with (a) averaged Gumbel values; (b) averaged loglikelihoods.

x∈ST 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 thatmax x∈S T G(φ Q (dx), L) = T G(log( x∈S exp(φ Q (x))dx), L).

which is the expected number of selections on the suboptimal subset. Formally, suppose Ω is partitioned into S 1 , S 2 . Let i * = arg max i∈{1,2} G P (S i ) which is a random variable. Consider an algorithm A which selects a subset S i A,t at time step t. The stochastic regret of A at time T is: R A

Table content: <table/>

in Table1where n denotes the number of data points and dim denotes dimension.

Table content: <table><tr><td/><td colspan="4">Banana B.Cancer Diabetes German</td><td>Heart</td><td>Image</td><td>Ringnorm</td></tr><tr><td>dim</td><td>2</td><td>9</td><td>8</td><td>20</td><td>5</td><td>18</td><td>20</td></tr><tr><td>n</td><td>5300</td><td>263</td><td>768</td><td>1000</td><td>215</td><td>2086</td><td>7400</td></tr><tr><td/><td>F.Sonar</td><td>Splice</td><td>Thyroid</td><td colspan="3">Titanic Twonorm Waveform</td><td/></tr><tr><td>dim</td><td>9</td><td>60</td><td>5</td><td>3</td><td>20</td><td>21</td><td/></tr><tr><td>n</td><td>144</td><td>2991</td><td>215</td><td>24</td><td>7400</td><td>5000</td><td/></tr><tr><td/><td>Table</td><td/><td/><td/><td/><td/><td/></tr></table>

Peer review (iclr2019):
 Summary: 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. Quality: This is an interesting and natural idea, and the experiments support the author*s claim that it improves over A* Sampling in the regimes that they consider. The claims and proofs look correct to me. Originality: This idea has not been explored in the literature, and is a sensible one to consider when having an exact sample is not necessary. Clarity: The comparison to A* Sampling is a bit apples-to-oranges and the paper lacks some clarity overall. In particular, the authors should be much more clear that their termination condition necessarily differs from A* Sampling*s. In A* Sampling the termination condition guarantees that the algorithm has found the true optimum, but in PM-A* the authors consider only a fixed budget of time --- terminating whether or not the algorithm has found the true optimum. The difference is a necessary feature, because otherwise I*m fairly sure A* Sampling is optimal (It has to push the upper bound of every subset below the true maximum, so if the term. cond. is not currently satisfied the subset with the current max. upper bound must eventually be visited before termination. Since it is the only subset where that is necessarily true, taking it is the optimal choice.) More comments below. Significance: I think this is interesting work, but I would recommend the authors consider some examples that clarify where these ideas might impact that problems that the ICLR audience would be interested in. Overall: The paper*s clarity problems get in the way of its contribution. I think it is an interesting direction, but the authors are not clear enough about the correctness of the algorithm (it is not guaranteed to return an exact sample) to recommend acceptance. If this and other issues are fixed, it would be a good paper. Pros: - Solid contribution to the literature on Gumbels. - Beginnings of a regret analysis. - Uniform performance gains in comparison to A* Sampling (in the regime considered). Cons: - Lacks clarity in comparisons to A* Sampling. - Writing could clarify significance as well. Specific questions / suggestions: - One of the important omissions of the paper is the following question. What is the distribution of PM-A**s output at time T for a fixed T? It won*t be the target distribution in general, but it may be close. This applicability of this work will be greatly improved by that analysis. - In general the figures and experiments need more details. In particular, (1) In figure 2a, why are there two sets of lines? Different values of T? (2) You need to report T throughout. (3) What did you take as Y? This is a critical part of the method, and you should report the estimator used. - In A* Sampling for two subsets without splitting algorithm, shouldn*t you be removing ALL previous Xi from be S_i? i.e. S_i . {Xi_1, Xi_2, ... Xi_{k-1}} or something like that.