$A^*$ sampling with probability matching

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 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 . 2 PRELIMINARIES . 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 ) . 2.1 GUMBEL PROCESS . 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 . 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 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 ) .

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.